taxonomyreport can write databases per query with --report-mode 2

milot-mirdita · milot-mirdita · commit 8284a8b1aa1f · 2025-02-22T18:35:59.000+09:00
diff --git a/src/commons/Parameters.cpp b/src/commons/Parameters.cpp
@@ -288,7 +288,7 @@ Parameters::Parameters():
         PARAM_PAIRING_DUMMY_MODE(PARAM_PAIRING_DUMMY_MODE_ID, "--pairing-dummy-mode", "Include dummy pairing", "0: dont include, 1: include - an entry that will cause result2msa to write a gap only line", typeid(int), (void *) &pairdummymode, "^[0-1]{1}$", MMseqsParameter::COMMAND_EXPERT),
         PARAM_PAIRING_MODE(PARAM_PAIRING_MODE_ID, "--pairing-mode", "Pairing mode", "0: pair maximal per species, 1: pair only if all chains are covered per species", typeid(int), (void *) &pairmode, "^[0-1]{1}$", MMseqsParameter::COMMAND_EXPERT),
         // taxonomyreport
-        PARAM_REPORT_MODE(PARAM_REPORT_MODE_ID, "--report-mode", "Report mode", "Taxonomy report mode 0: Kraken 1: Krona", typeid(int), (void *) &reportMode, "^[0-1]{1}$"),
+        PARAM_REPORT_MODE(PARAM_REPORT_MODE_ID, "--report-mode", "Report mode", "Taxonomy report mode 0: Kraken 1: Krona 2: Kraken per query database", typeid(int), (void *) &reportMode, "^[0-2]{1}$"),
         // createtaxdb
         PARAM_NCBI_TAX_DUMP(PARAM_NCBI_TAX_DUMP_ID, "--ncbi-tax-dump", "NCBI tax dump directory", "NCBI tax dump directory. The tax dump can be downloaded here \"ftp://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz\"", typeid(std::string), (void *) &ncbiTaxDump, ""),
         PARAM_TAX_MAPPING_FILE(PARAM_TAX_MAPPING_FILE_ID, "--tax-mapping-file", "Taxonomy mapping file", "File to map sequence identifier to taxonomical identifier", typeid(std::string), (void *) &taxMappingFile, ""),
diff --git a/src/taxonomy/NcbiTaxonomy.cpp b/src/taxonomy/NcbiTaxonomy.cpp
@@ -486,7 +486,7 @@ size_t NcbiTaxonomy::loadMerged(const std::string &mergedFile) {
     if (localMaxTaxID > maxTaxID) {
         int* newD = new int[localMaxTaxID + 1];
         std::copy(D, D + maxTaxID + 1, newD);
-        std::fill(newD + maxTaxID + 1, newD + (localMaxTaxID + 1), -1);
+        std::fill_n(newD + maxTaxID + 1, localMaxTaxID - maxTaxID, -1);
         delete[] D;
         D = newD;
         maxTaxID = localMaxTaxID;
@@ -501,8 +501,7 @@ size_t NcbiTaxonomy::loadMerged(const std::string &mergedFile) {
     return count;
 }
 
-std::unordered_map<TaxID, TaxonCounts> NcbiTaxonomy::getCladeCounts(std::unordered_map<TaxID, unsigned int>& taxonCounts) const {
-    Debug(Debug::INFO) << "Calculating clade counts ... ";
+std::unordered_map<TaxID, TaxonCounts> NcbiTaxonomy::getCladeCounts(const std::unordered_map<TaxID, unsigned int>& taxonCounts) const {
     std::unordered_map<TaxID, TaxonCounts> cladeCounts;
 
     for (std::unordered_map<TaxID, unsigned int>::const_iterator it = taxonCounts.begin(); it != taxonCounts.end(); ++it) {
@@ -525,7 +524,6 @@ std::unordered_map<TaxID, TaxonCounts> NcbiTaxonomy::getCladeCounts(std::unorder
         }
     }
 
-    Debug(Debug::INFO) << " Done\n";
     return cladeCounts;
 }
 
diff --git a/src/taxonomy/NcbiTaxonomy.h b/src/taxonomy/NcbiTaxonomy.h
@@ -107,7 +107,7 @@ class NcbiTaxonomy {
     TaxonNode const* taxonNode(TaxID taxonId, bool fail = true) const;
     bool nodeExists(TaxID taxId) const;
 
-    std::unordered_map<TaxID, TaxonCounts> getCladeCounts(std::unordered_map<TaxID, unsigned int>& taxonCounts) const;
+    std::unordered_map<TaxID, TaxonCounts> getCladeCounts(const std::unordered_map<TaxID, unsigned int>& taxonCounts) const;
 
     WeightedTaxResult weightedMajorityLCA(const std::vector<WeightedTaxHit> &setTaxa, const float majorityCutoff);
 
diff --git a/src/taxonomy/taxonomyreport.cpp b/src/taxonomy/taxonomyreport.cpp
@@ -15,50 +15,57 @@
 #include <omp.h>
 #endif
 
-template<typename K, typename V>
-V at(const std::unordered_map<K, V> &map, K key, V default_value = V()) {
-    typename std::unordered_map<K, V>::const_iterator it = map.find(key);
-    if (it == map.end()) {
-        return default_value;
-    } else {
-        return it->second;
-    }
-}
-
 unsigned int cladeCountVal(const std::unordered_map<TaxID, TaxonCounts> &map, TaxID key) {
-    typename std::unordered_map<TaxID, TaxonCounts>::const_iterator it = map.find(key);
+    std::unordered_map<TaxID, TaxonCounts>::const_iterator it = map.find(key);
     if (it == map.end()) {
         return 0;
     } else {
         return it->second.cladeCount;
     }
 }
 
-void taxReport(FILE *FP, const NcbiTaxonomy &taxDB, const std::unordered_map<TaxID, TaxonCounts> &cladeCounts, unsigned long totalReads, TaxID taxID = 0, int depth = 0) {
+void taxReport(
+    DBWriter& writer,
+    unsigned int thread_idx,
+    const NcbiTaxonomy &taxDB,
+    const std::unordered_map<TaxID, TaxonCounts> &cladeCounts,
+    unsigned long totalReads,
+    TaxID taxID = 0,
+    int depth = 0
+) {
     std::unordered_map<TaxID, TaxonCounts>::const_iterator it = cladeCounts.find(taxID);
     unsigned int cladeCount = it == cladeCounts.end() ? 0 : it->second.cladeCount;
     unsigned int taxCount = it == cladeCounts.end() ? 0 : it->second.taxCount;
+    char buffer[4096];
+    int len = 0;
     if (taxID == 0) {
         if (cladeCount > 0) {
-            fprintf(FP, "%.4f\t%i\t%i\tno rank\t0\tunclassified\n",
-                    100 * cladeCount / double(totalReads),
-                    cladeCount, taxCount);
+            len = snprintf(buffer, sizeof(buffer), "%.4f\t%i\t%i\tno rank\t0\tunclassified\n",
+                           100 * cladeCount / double(totalReads),
+                           cladeCount, taxCount);
+            writer.writeAdd(buffer, static_cast<size_t>(len), thread_idx);
         }
-        taxReport(FP, taxDB, cladeCounts, totalReads, 1);
+        taxReport(writer, thread_idx, taxDB, cladeCounts, totalReads, 1);
     } else {
         if (cladeCount == 0) {
             return;
         }
         const TaxonNode *taxon = taxDB.taxonNode(taxID);
-        fprintf(FP, "%.4f\t%i\t%i\t%s\t%i\t%s%s\n",
-                100 * cladeCount / double(totalReads), cladeCount, taxCount,
-                taxDB.getString(taxon->rankIdx), taxID, std::string(2 * depth, ' ').c_str(), taxDB.getString(taxon->nameIdx));
+        std::string indent = std::string(2 * depth, ' ');
+        len = snprintf(buffer, sizeof(buffer), "%.4f\t%i\t%i\t%s\t%i\t%s%s\n",
+                       100 * cladeCount / double(totalReads),
+                       cladeCount, taxCount,
+                       taxDB.getString(taxon->rankIdx),
+                       taxID,
+                       indent.c_str(),
+                       taxDB.getString(taxon->nameIdx));
+        writer.writeAdd(buffer, static_cast<size_t>(len), thread_idx);
         std::vector<TaxID> children = it->second.children;
         SORT_SERIAL(children.begin(), children.end(), [&](int a, int b) { return cladeCountVal(cladeCounts, a) > cladeCountVal(cladeCounts, b); });
         for (size_t i = 0; i < children.size(); ++i) {
             TaxID childTaxId = children[i];
             if (cladeCounts.count(childTaxId)) {
-                taxReport(FP, taxDB, cladeCounts, totalReads, childTaxId, depth + 1);
+                taxReport(writer, thread_idx, taxDB, cladeCounts, totalReads, childTaxId, depth + 1);
             } else {
                 break;
             }
@@ -94,32 +101,56 @@ std::string escapeAttribute(const std::string &data) {
     return buffer;
 }
 
-void kronaReport(FILE *FP, const NcbiTaxonomy &taxDB, const std::unordered_map<TaxID, TaxonCounts> &cladeCounts, unsigned long totalReads, TaxID taxID = 0, int depth = 0) {
+void kronaReport(
+    DBWriter& writer,
+    unsigned int thread_idx,
+    const NcbiTaxonomy &taxDB,
+    const std::unordered_map<TaxID, TaxonCounts> &cladeCounts,
+    unsigned long totalReads,
+    TaxID taxID = 0,
+    int depth = 0
+) {
     std::unordered_map<TaxID, TaxonCounts>::const_iterator it = cladeCounts.find(taxID);
     unsigned int cladeCount = it == cladeCounts.end() ? 0 : it->second.cladeCount;
+    char buffer[1024];
+    int len = 0;
     if (taxID == 0) {
         if (cladeCount > 0) {
-            fprintf(FP, "<node name=\"unclassified\"><magnitude><val>%d</val></magnitude></node>", cladeCount);
+            len = snprintf(buffer, sizeof(buffer), "<node name=\"unclassified\"><magnitude><val>%d</val></magnitude></node>", cladeCount);
+            writer.writeAdd(buffer, static_cast<size_t>(len), thread_idx);
         }
-        kronaReport(FP, taxDB, cladeCounts, totalReads, 1);
+        kronaReport(writer, thread_idx, taxDB, cladeCounts, totalReads, 1);
     } else {
         if (cladeCount == 0) {
             return;
         }
         const TaxonNode *taxon = taxDB.taxonNode(taxID);
         std::string escapedName = escapeAttribute(taxDB.getString(taxon->nameIdx));
-        fprintf(FP, "<node name=\"%s\"><magnitude><val>%d</val></magnitude>", escapedName.c_str(), cladeCount);
+        len = snprintf(buffer, sizeof(buffer), "<node name=\"%s\"><magnitude><val>%d</val></magnitude>", escapedName.c_str(), cladeCount);
+        writer.writeAdd(buffer, static_cast<size_t>(len), thread_idx);
         std::vector<TaxID> children = it->second.children;
         SORT_SERIAL(children.begin(), children.end(), [&](int a, int b) { return cladeCountVal(cladeCounts, a) > cladeCountVal(cladeCounts, b); });
         for (size_t i = 0; i < children.size(); ++i) {
             TaxID childTaxId = children[i];
             if (cladeCounts.count(childTaxId)) {
-                kronaReport(FP, taxDB, cladeCounts, totalReads, childTaxId, depth + 1);
+                kronaReport(writer, thread_idx, taxDB, cladeCounts, totalReads, childTaxId, depth + 1);
             } else {
                 break;
             }
         }
-        fprintf(FP, "</node>");
+        len = snprintf(buffer, sizeof(buffer), "</node>");
+        writer.writeAdd(buffer, static_cast<size_t>(len), thread_idx);
+    }
+}
+
+void mergeMaps(std::unordered_map<TaxID, unsigned int>& target, const std::unordered_map<TaxID, unsigned int>& source) {
+    for (std::unordered_map<TaxID, unsigned int>::const_iterator it = source.cbegin(); it != source.cend(); ++it) {
+        std::unordered_map<TaxID, unsigned int>::iterator found = target.find(it->first);
+        if (found != target.end()) {
+            found->second += it->second;
+        } else {
+            target.emplace(it->first, it->second);
+        }
     }
 }
 
@@ -133,6 +164,10 @@ int taxonomyreport(int argc, const char **argv, const Command &command) {
     const bool isSequenceDB = Parameters::isEqualDbtype(readerDbType, Parameters::DBTYPE_HMM_PROFILE)
                              || Parameters::isEqualDbtype(readerDbType, Parameters::DBTYPE_AMINO_ACIDS)
                              || Parameters::isEqualDbtype(readerDbType, Parameters::DBTYPE_NUCLEOTIDES);
+    if (par.reportMode == 2 && isSequenceDB == true) {
+        Debug(Debug::ERROR) << "Cannot use Kraken DB report mode with sequence db input\n";
+        EXIT(EXIT_FAILURE);
+    }
     int dataMode = DBReader<unsigned int>::USE_INDEX;
     if (isSequenceDB == false) {
         dataMode |= DBReader<unsigned int>::USE_DATA;
@@ -147,11 +182,21 @@ int taxonomyreport(int argc, const char **argv, const Command &command) {
         mapping = new MappingReader(par.db1);
     }
 
-    FILE *resultFP = FileUtil::openAndDelete(par.db3.c_str(), "w");
+    int mode = Parameters::DBTYPE_OMIT_FILE;
+    unsigned int localThreads = 1;
+    if (par.reportMode == 2) {
+        mode = Parameters::DBTYPE_GENERIC_DB;
+        localThreads = 1;
+#ifdef OPENMP
+        localThreads = std::max(std::min((size_t)par.threads, reader.getSize()), (size_t)1);
+#endif
+    }
+    DBWriter writer(par.db3.c_str(), par.db3Index.c_str(), localThreads, false, mode);
+    writer.open();
 
     std::unordered_map<TaxID, unsigned int> taxCounts;
     Debug::Progress progress(reader.getSize());
-#pragma omp parallel
+#pragma omp parallel num_threads(localThreads)
     {
         unsigned int thread_idx = 0;
 #ifdef OPENMP
@@ -172,6 +217,7 @@ int taxonomyreport(int argc, const char **argv, const Command &command) {
             }
 
             char *data = reader.getData(i, thread_idx);
+            size_t entryCount = 0;
             while (*data != '\0') {
                 if (isTaxonomyInput) {
                     TaxID taxon = Util::fast_atoi<int>(data);
@@ -183,40 +229,56 @@ int taxonomyreport(int argc, const char **argv, const Command &command) {
                         ++localTaxCounts[taxon];
                     }
                 }
+                entryCount++;
                 data = Util::skipLine(data);
             }
+            if (par.reportMode == 2) {
+                std::unordered_map<TaxID, TaxonCounts> cladeCounts = taxDB->getCladeCounts(localTaxCounts);
+                writer.writeStart(thread_idx);
+                taxReport(writer, thread_idx, *taxDB, cladeCounts, entryCount);
+                writer.writeEnd(reader.getDbKey(i), thread_idx);
+                localTaxCounts.clear();
+            }
         }
-
-        // merge maps again
+        if (par.reportMode != 2) {
 #pragma omp critical
-        for (std::unordered_map<TaxID, unsigned int>::const_iterator it = localTaxCounts.cbegin(); it != localTaxCounts.cend(); ++it) {
-            if (taxCounts[it->first]) {
-                taxCounts[it->first] += it->second;
-            } else {
-                taxCounts[it->first] = it->second;
-            }
+{
+            mergeMaps(taxCounts, localTaxCounts);
+}
         }
     }
-    Debug(Debug::INFO) << "Found " << taxCounts.size() << " different taxa for " << reader.getSize() << " different reads\n";
-    unsigned int unknownCnt = (taxCounts.find(0) != taxCounts.end()) ? taxCounts.at(0) : 0;
-    Debug(Debug::INFO) << unknownCnt << " reads are unclassified\n";
-    const size_t entryCount = reader.getSize();
-    reader.close();
-
-    std::unordered_map<TaxID, TaxonCounts> cladeCounts = taxDB->getCladeCounts(taxCounts);
-    if (par.reportMode == 0) {
-        taxReport(resultFP, *taxDB, cladeCounts, entryCount);
+    if (par.reportMode == 2) {
+        reader.close();
+        writer.close(true);
     } else {
-        fwrite(krona_prelude_html, krona_prelude_html_len, sizeof(char), resultFP);
-        fprintf(resultFP, "<node name=\"all\"><magnitude><val>%zu</val></magnitude>", entryCount);
-        kronaReport(resultFP, *taxDB, cladeCounts, entryCount);
-        fprintf(resultFP, "</node></krona></div></body></html>");
+        Debug(Debug::INFO) << "Found " << taxCounts.size() << " different taxa for " << reader.getSize() << " different reads\n";
+        unsigned int unknownCnt = (taxCounts.find(0) != taxCounts.end()) ? taxCounts.at(0) : 0;
+        Debug(Debug::INFO) << unknownCnt << " reads are unclassified\n";
+        const size_t entryCount = reader.getSize();
+        reader.close();
+
+        Debug(Debug::INFO) << "Calculating clade counts ... ";
+        std::unordered_map<TaxID, TaxonCounts> cladeCounts = taxDB->getCladeCounts(taxCounts);
+        Debug(Debug::INFO) << " Done\n";
+        if (par.reportMode == 0) {
+            writer.writeStart(0);
+            taxReport(writer, 0, *taxDB, cladeCounts, entryCount);
+            writer.writeEnd(0, 0, false, false);
+        } else {
+            writer.writeStart(0);
+            writer.writeAdd(reinterpret_cast<const char*>(krona_prelude_html), krona_prelude_html_len, 0);
+            char buffer[1024];
+            int len = snprintf(buffer, sizeof(buffer), "<node name=\"all\"><magnitude><val>%zu</val></magnitude>", entryCount);
+            writer.writeAdd(buffer, static_cast<size_t>(len), 0);
+            kronaReport(writer, 0, *taxDB, cladeCounts, entryCount);
+            len = snprintf(buffer, sizeof(buffer), "</node></krona></div></body></html>");
+            writer.writeAdd(buffer, static_cast<size_t>(len), 0);
+            writer.writeEnd(0, 0, false, false);
+        }
+        writer.close(true);
+        FileUtil::remove(writer.getIndexFileName());
     }
     delete taxDB;
-    if (fclose(resultFP) != 0) {
-        Debug(Debug::ERROR) << "Cannot close file " << par.db3 << "\n";
-        return EXIT_FAILURE;
-    }
     return EXIT_SUCCESS;
 }
 
