TeraMS.cpp

#include "TeraIndex/TeraIndex.h"
#include "util/fasta.h"
#include "util/util.h"
#include <queue>

static constexpr size_t DEFAULT_BUFFER_SIZE = 64ULL * 1024ULL * 1024ULL; // 64MB

void printUsage() {
    std::cout << 
        "TeraMS computes matching statistics of a pattern.\n"
        "It requires as input the ms index generated by TeraIndexBuilder.\n"
        "\n"
        "Usage: TeraMS <arguments>\n"
        "Options:\n"
        "  Input:\n"
        "    -i          FILE                       REQUIRED       file name of input (ends with '" << ms_index_extension << "'), an index generated by TeraIndexBuilder.\n"
        "    -q          FILE                       REQUIRED       file name of query file.\n"
        "  Output:\n"
        "    -o          FILE                       optional       base name of output file for matching statistics (writes to FILE.[len|pos]), otherwise QUERY.[len|pos] is used.\n"
        "  Behavior:\n"
        "    -m          [phi,psi,dual]             optional       matching statistics mode. one of 'phi', 'psi', or 'dual'. dual is default.\n"
        "    -oracle     bool                       optional       output repositioning oracle. default is false.\n"
        "    -p          INT                        optional       Limit the program to (nonnegative) INT threads. By default uses maximum available. Maximum on this hardware is " << omp_get_max_threads() << "\n"
        "    -v          [quiet,time,verb]          optional       Verbosity, verb for most verbose output, time for timer info, and quiet for no output. time is default.\n"
        "    -h, --help                             optional       Print this help message.\n"
        ;
}

struct options{
    std::string indexFile, patternFile, outputFile = "", mode = "dual";
    unsigned numThreads = omp_get_max_threads();
    bool oracle = false;
    verbosity v = TIME;
}o;

void processOptions(const int argc, const char* argv[]) {
    std::vector<bool> used(argc);
    used[0] = true;

    if (argc == 1 || getArgument(argc, argv, used, "-h", false, false) != "" || getArgument(argc, argv, used, "--help", false, false) != "") {
        printUsage();
        exit(0);
    }

    o.indexFile = getArgument(argc, argv, used, "-i", true, true);
    o.patternFile = getArgument(argc, argv, used, "-q", true, true);
    std::string s = getArgument(argc, argv, used, "-o", false, true);
    if (s != "") {
        o.outputFile = s;
    } else {
        o.outputFile = o.patternFile;
    }
    s = getArgument(argc, argv, used, "-m", false, true);
    if (s != "" && s != "phi" && s != "psi" && s != "dual" && s != "oracle") {
        std::cout << "Invalid value passed to -m '" << s << "'\n";
        exit(1);
    }
    else if (s != "") {
        o.mode = s;
    }
    o.oracle = getArgument(argc, argv, used, "-oracle", false, false) != ""; // Default is false
    s = getArgument(argc, argv, used, "-p", false, true);
    if (s != "")
        o.numThreads = std::stoul(s); 
    s = getArgument(argc, argv, used, "-v", false, true);
    if (s == "quiet")
        o.v = QUIET;
    else if (s == "time" || s == "")  
        o.v = TIME;
    else if (s == "verb")
        o.v = VERB;
    for (int i = 0; i < argc; ++i) {
        if (!used[i]) {
            std::cout << "Argument " << i << ", '" << argv[i] << "' not recognized or used as an argument for another option. It might have been passed more than once (invalid).\n";
            exit(1);
        }
	}
    testInFile(o.indexFile);
    testInFile(o.patternFile);
}

// Structure to hold data for writing
struct WriteData {
    std::string seq_name;
    std::vector<uint32_t> len_data;
    std::vector<uint64_t> pos_data;
};

// Thread-safe queue for write operations
struct WriteQueue {
    std::queue<WriteData> queue;
    std::mutex mutex;
    std::condition_variable cv;
    std::atomic<bool> done{false};
};

// Write thread function
void write_thread_func(WriteQueue& write_queue, FILE* out_len, FILE* out_pos, double& total_write_time) {
    while (true) {
        std::unique_lock<std::mutex> lock(write_queue.mutex);
        write_queue.cv.wait(lock, [&] { return !write_queue.queue.empty() || write_queue.done.load(); });
        
        if (write_queue.queue.empty() && write_queue.done.load()) {
            break;
        }
        
        if (!write_queue.queue.empty()) {
            WriteData data = std::move(write_queue.queue.front());
            write_queue.queue.pop();
            lock.unlock();
            
            auto start_time = std::chrono::high_resolution_clock::now();
            fprintf(out_len, ">%s\n", data.seq_name.c_str());
            fprintf(out_pos, ">%s\n", data.seq_name.c_str());
            size_t len_data_size = data.len_data.size();
            fwrite(&len_data_size, sizeof(size_t), 1, out_len);
            fwrite(data.len_data.data(), sizeof(uint32_t), len_data_size, out_len);
            size_t pos_data_size = data.pos_data.size();
            fwrite(&pos_data_size, sizeof(size_t), 1, out_pos);
            fwrite(data.pos_data.data(), sizeof(uint64_t), pos_data_size, out_pos);
            fputc('\n', out_len);
            fputc('\n', out_pos);
            auto end_time = std::chrono::high_resolution_clock::now();
            double duration = std::chrono::duration<double>(end_time - start_time).count();
            total_write_time += duration;
        }
    }
}

int main(const int argc, const char*argv[]) {
	processOptions(argc, argv);

    #ifdef STATS
    if (o.numThreads > 1) {
        std::cerr << "Stats are not supported with multiple threads. Please run with -p 1." << std::endl;
        exit(1);
    }
    #endif
    if (o.oracle) {
        if (o.numThreads > 1) {
            std::cout << "Oracle computation is enabled. This is not supported for multiple threads." << std::endl;
            exit(1);
        }
        if (o.mode == "oracle") {
            std::cout << "Oracle computation is enabled. This is not supported with the 'oracle' mode." << std::endl;
            exit(1);
        }
        #ifndef WRITE_ORACLE
        std::cout << "Oracle computation is enabled. This is not supported without compiling with the WRITE_ORACLE macro." << std::endl;
        exit(1);
        #endif
    }

    omp_set_num_threads(o.numThreads);

    if (o.v >= TIME) { Timer.start("msComputer"); }
    if (o.v >= TIME) { Timer.start("Program Initialization"); }
	if (o.v >= TIME) { Timer.start("Loading index " + o.indexFile); }
	std::ifstream in(o.indexFile);
	TeraIndex msIndex;
	msIndex.load(in);
    #ifdef WRITE_ORACLE
    msIndex.set_oracle(o.oracle);
    #endif
	in.close();
	if (o.v >= TIME) { Timer.stop(); } //Loading " + o.indexFile

    Timer.start("Opening files for read (" + o.patternFile + ") and write (" + o.outputFile + ".[len|pos])");
    FILE *fp;
    kseq_t *seq = open_fasta(o.patternFile, &fp);
    
    // For maximum performance, use system buffer size
    std::string outputFile_len = o.outputFile + ".len";
    std::string outputFile_pos = o.outputFile + ".pos";
    FILE* out_len = fopen(outputFile_len.c_str(), "w");
    FILE* out_pos = fopen(outputFile_pos.c_str(), "w");
    int fd_len = fileno(out_len);
    int fd_pos = fileno(out_pos);

    long long buffer_size_len = fpathconf(fd_len, _PC_REC_XFER_ALIGN);
    long long buffer_size_pos = fpathconf(fd_pos, _PC_REC_XFER_ALIGN);

    if (buffer_size_len <= 0) {
        buffer_size_len = DEFAULT_BUFFER_SIZE;
    }
    if (buffer_size_pos <= 0) {
        buffer_size_pos = DEFAULT_BUFFER_SIZE;
    }
    setvbuf(out_len, nullptr, _IOFBF, buffer_size_len);
    setvbuf(out_pos, nullptr, _IOFBF, buffer_size_pos);
    Timer.stop(); //Opening files for read (" + o.patternFile + " and write (" + o.outputFile + ")
    Timer.stop(); //Program Initialization

    #ifdef STATS
    msIndex.reset_ms_stats();
    #endif

    auto total_ms_time = 0.0;
    double total_write_time = 0.0;
    uint64_t total_seq_len = 0;
    
    // Create write queue and start write thread
    WriteQueue write_queue;
    std::thread write_thread(write_thread_func, std::ref(write_queue), out_len, out_pos, std::ref(total_write_time));
    
    auto ms_step = [&](const SeqInfo& seq_info) {
        #pragma omp atomic
        total_seq_len += seq_info.seq_len;
        thread_local static std::pair<std::vector<uint32_t>, std::vector<uint64_t>> ms_result;

        // Load repositioning oracle if needed
        std::vector<uint32_t> repositioning_oracle;
        if (o.mode == "oracle") {
            std::ifstream in_oracle(o.patternFile + "." + seq_info.seq_name + ".oracle");
            repositioning_oracle = msIndex.load_oracle(in_oracle);
            in_oracle.close();
        }

        auto start_time = std::chrono::high_resolution_clock::now();
        ms_result.first.clear();
        ms_result.second.clear();
        if (o.mode == "phi") {
            ms_result = msIndex.ms_phi(seq_info.seq_content, seq_info.seq_len);
        } else if (o.mode == "psi") {
            ms_result = msIndex.ms_psi(seq_info.seq_content, seq_info.seq_len);
        } else if (o.mode == "dual") {
            ms_result = msIndex.ms_dual(seq_info.seq_content, seq_info.seq_len);
        } else if (o.mode == "oracle") {
            ms_result = msIndex.ms_oracle(seq_info.seq_content, seq_info.seq_len, repositioning_oracle);
        } else {
            std::cerr << "Invalid mode: " << o.mode << std::endl;
            exit(1);
        }
        auto end_time = std::chrono::high_resolution_clock::now();
        #pragma omp atomic
        total_ms_time += std::chrono::duration<double>(end_time - start_time).count();

        // Push result to write queue instead of writing directly
        WriteData write_data;
        write_data.seq_name = std::string(seq_info.seq_name);
        write_data.len_data = ms_result.first;
        write_data.pos_data = ms_result.second;
        
        {
            std::lock_guard<std::mutex> lock(write_queue.mutex);
            write_queue.queue.push(std::move(write_data));
        }
        write_queue.cv.notify_one();

        #ifdef WRITE_ORACLE
        if (o.oracle) {
            std::ofstream out_oracle(o.patternFile + "." + seq_info.seq_name + ".oracle");
            msIndex.serialize_oracle(out_oracle);
            out_oracle.close();
        }
        #endif
    };

    Timer.start("Processing patterns");
    process_sequences(seq, o.numThreads, ms_step);
    
    // Signal write thread to finish and wait for it
    write_queue.done = true;
    write_queue.cv.notify_one();
    write_thread.join();
    Timer.stop(); //Processing patterns
    
    std::ofstream out_stats(o.outputFile + ".stats");
    out_stats << "\tCPU query time: " << total_ms_time << " seconds" << std::endl;
    out_stats << "\t\tTime per base: " << (total_ms_time / total_seq_len) * 1e9 << " nanoseconds" << std::endl;
    out_stats << "\tCPU write time: " << total_write_time << " seconds" << std::endl << std::endl;
    #ifdef STATS
    out_stats << "\tTotal bases: " << total_seq_len << std::endl;
    msIndex.print_ms_stats(out_stats);
    #endif
    out_stats.close();

    fclose(out_len);
    fclose(out_pos);
    kseq_destroy(seq);
    fclose(fp);
    Timer.stop(); //msComputer
    return 0;
}