Speed-up getCladeCounts

milot-mirdita · milot-mirdita · commit 033d5f5f4cb1 · 2025-02-22T18:35:59.000+09:00
diff --git a/src/taxonomy/NcbiTaxonomy.cpp b/src/taxonomy/NcbiTaxonomy.cpp
@@ -501,7 +501,23 @@ size_t NcbiTaxonomy::loadMerged(const std::string &mergedFile) {
     return count;
 }
 
-std::unordered_map<TaxID, TaxonCounts> NcbiTaxonomy::getCladeCounts(const std::unordered_map<TaxID, unsigned int>& taxonCounts) const {
+std::unordered_map<TaxID, std::vector<TaxID>> NcbiTaxonomy::getParentToChildren() const {
+    std::unordered_map<TaxID, std::vector<TaxID>> result;
+    result.reserve(maxNodes);
+    
+    // Build the adjacency (parent -> children)
+    for (size_t i = 0; i < maxNodes; ++i) {
+        const TaxonNode& tn = taxonNodes[i];
+        if (tn.parentTaxId == tn.taxId) {
+            continue;
+        }
+        result[tn.parentTaxId].push_back(tn.taxId);
+    }
+
+    return result;
+}
+
+std::unordered_map<TaxID, TaxonCounts> NcbiTaxonomy::getCladeCounts(const std::unordered_map<TaxID, unsigned int>& taxonCounts, const std::unordered_map<TaxID, std::vector<TaxID>>& parentToChildren) const {
     std::unordered_map<TaxID, TaxonCounts> cladeCounts;
 
     for (std::unordered_map<TaxID, unsigned int>::const_iterator it = taxonCounts.begin(); it != taxonCounts.end(); ++it) {
@@ -516,11 +532,12 @@ std::unordered_map<TaxID, TaxonCounts> NcbiTaxonomy::getCladeCounts(const std::u
         }
     }
 
-    for (size_t i = 0; i < maxNodes; ++i) {
-        TaxonNode& tn = taxonNodes[i];
-        if (tn.parentTaxId != tn.taxId && cladeCounts.count(tn.taxId)) {
-            std::unordered_map<TaxID, TaxonCounts>::iterator itp = cladeCounts.find(tn.parentTaxId);
-            itp->second.children.push_back(tn.taxId);
+   for (std::unordered_map<TaxID, TaxonCounts>::iterator it = cladeCounts.begin(); it != cladeCounts.end(); ++it) {
+        TaxID parentTaxId = it->first;
+        TaxonCounts& taxCounts = it->second;
+        std::unordered_map<TaxID, std::vector<TaxID>>::const_iterator ptcIt = parentToChildren.find(parentTaxId);
+        if (ptcIt != parentToChildren.end()) {
+            taxCounts.children = ptcIt->second;
         }
     }
 
diff --git a/src/taxonomy/NcbiTaxonomy.h b/src/taxonomy/NcbiTaxonomy.h
@@ -107,7 +107,8 @@ class NcbiTaxonomy {
     TaxonNode const* taxonNode(TaxID taxonId, bool fail = true) const;
     bool nodeExists(TaxID taxId) const;
 
-    std::unordered_map<TaxID, TaxonCounts> getCladeCounts(const std::unordered_map<TaxID, unsigned int>& taxonCounts) const;
+    std::unordered_map<TaxID, std::vector<TaxID>> getParentToChildren() const;
+    std::unordered_map<TaxID, TaxonCounts> getCladeCounts(const std::unordered_map<TaxID, unsigned int>& taxonCounts, const std::unordered_map<TaxID, std::vector<TaxID>>& parentToChildren) const;
 
     WeightedTaxResult weightedMajorityLCA(const std::vector<WeightedTaxHit> &setTaxa, const float majorityCutoff);
 
diff --git a/src/taxonomy/taxonomyreport.cpp b/src/taxonomy/taxonomyreport.cpp
@@ -194,6 +194,8 @@ int taxonomyreport(int argc, const char **argv, const Command &command) {
     DBWriter writer(par.db3.c_str(), par.db3Index.c_str(), localThreads, false, mode);
     writer.open();
 
+    std::unordered_map<TaxID, std::vector<TaxID>> parentToChildren = taxDB->getParentToChildren();
+
     std::unordered_map<TaxID, unsigned int> taxCounts;
     Debug::Progress progress(reader.getSize());
 #pragma omp parallel num_threads(localThreads)
@@ -233,7 +235,7 @@ int taxonomyreport(int argc, const char **argv, const Command &command) {
                 data = Util::skipLine(data);
             }
             if (par.reportMode == 2) {
-                std::unordered_map<TaxID, TaxonCounts> cladeCounts = taxDB->getCladeCounts(localTaxCounts);
+                std::unordered_map<TaxID, TaxonCounts> cladeCounts = taxDB->getCladeCounts(localTaxCounts, parentToChildren);
                 writer.writeStart(thread_idx);
                 taxReport(writer, thread_idx, *taxDB, cladeCounts, entryCount);
                 writer.writeEnd(reader.getDbKey(i), thread_idx);
@@ -258,7 +260,7 @@ int taxonomyreport(int argc, const char **argv, const Command &command) {
         reader.close();
 
         Debug(Debug::INFO) << "Calculating clade counts ... ";
-        std::unordered_map<TaxID, TaxonCounts> cladeCounts = taxDB->getCladeCounts(taxCounts);
+        std::unordered_map<TaxID, TaxonCounts> cladeCounts = taxDB->getCladeCounts(taxCounts, parentToChildren);
         Debug(Debug::INFO) << " Done\n";
         if (par.reportMode == 0) {
             writer.writeStart(0);
