improve speed for progressive alignment

Author: y3tseng

src/alignment-helper.cpp

@@ -376,12 +376,11 @@ void msa::alignment_helper::addGappyColumnsBack(alnPath &aln_before, alnPath &al
 void msa::alignment_helper::updateAlignment(NodePair &nodes, SequenceDB *database, Option* option, alnPath &aln)
 {
     int totalLen = aln.size(), startLen = 0, endLen = 0;
-    bool updateSeq = (database->currentTask != 2);
     tbb::this_task_arena::isolate([&] { 
     tbb::parallel_for(tbb::blocked_range<int>(0, nodes.first->seqsIncluded.size()), [&](tbb::blocked_range<int> range) {
     for (int idx = range.begin(); idx < range.end(); ++idx) {
         int sIdx = nodes.first->seqsIncluded[idx];
-        if (updateSeq) {
+        if (database->currentTask != 2 && sIdx >= 0) {
             database->id_map[sIdx]->memCheck(totalLen);
             int storeFrom = database->id_map[sIdx]->storage;
             int storeTo = 1 - storeFrom;
@@ -420,7 +419,7 @@ void msa::alignment_helper::updateAlignment(NodePair &nodes, SequenceDB *databas
     tbb::parallel_for(tbb::blocked_range<int>(0, nodes.second->seqsIncluded.size()), [&](tbb::blocked_range<int> range) {
     for (int idx = range.begin(); idx < range.end(); ++idx) {
         int sIdx = nodes.second->seqsIncluded[idx];
-        if (updateSeq) {
+        if (database->currentTask != 2 && sIdx >= 0) {
             database->id_map[sIdx]->memCheck(totalLen);
             int orgIdx = 0;
             int storeFrom = database->id_map[sIdx]->storage;
@@ -455,11 +454,31 @@ void msa::alignment_helper::updateAlignment(NodePair &nodes, SequenceDB *databas
     } 
     }); 
     });
-    for (auto idx: nodes.second->seqsIncluded) nodes.first->seqsIncluded.push_back(idx);
     nodes.first->alnNum += nodes.second->alnNum;
-    nodes.second->seqsIncluded.clear();
     nodes.first->alnLen = totalLen;
     nodes.first->alnWeight += nodes.second->alnWeight;
+    for (auto idx: nodes.second->seqsIncluded) nodes.first->seqsIncluded.push_back(idx);
+    nodes.second->seqsIncluded.clear();
+    if (nodes.first->seqsIncluded.size() > _UPDATE_SEQ_TH && !nodes.first->msaFreq.empty() && database->currentTask != 2) {
+        int seqCount = 0, firstSeqID = 0;
+        for (auto idx: nodes.first->seqsIncluded) {
+            if (idx > 1) {
+               if (firstSeqID == 0) firstSeqID = -idx;
+               seqCount++;
+            }
+        }
+        if (seqCount >= _UPDATE_SEQ_TH) {
+            database->subtreeAln[firstSeqID] = alnPath(totalLen, 0);
+            std::vector<int> new_seqsIncluded;
+            new_seqsIncluded.push_back(firstSeqID);
+            for (auto idx: nodes.first->seqsIncluded) {
+                if (idx >= 0) database->id_map[idx]->subtreeIdx = firstSeqID;
+                else new_seqsIncluded.push_back(idx);
+            }
+            nodes.first->seqsIncluded = new_seqsIncluded;
+        }
+    }
+    
     return;
 }
 

src/msa.hpp

@@ -175,6 +175,7 @@ namespace msa
     namespace alignment_helper
     {
         constexpr int _CAL_PROFILE_TH = 1000;
+        constexpr int _UPDATE_SEQ_TH = 1000;
 
         void calculateProfile(float *profile, NodePair &nodes, SequenceDB *database, Option *option, int32_t memLen);
         void removeGappyColumns(float *hostFreq, NodePair &nodes, Option *option, std::pair<IntPairVec, IntPairVec> &gappyColumns, int32_t memLen, IntPair &lens, int currentTask);
@@ -199,6 +200,7 @@ namespace msa
         void updateNode(Tree *tree, NodePairVec &nodes, SequenceDB *database);
         void progressiveAlignment(Tree *T, SequenceDB *database, Option *option, std::vector<NodePairVec> &alnPairPerLevel, Params &param, alnFunction alignmentKernel);
         void msaOnSubtree(Tree *T, SequenceDB *database, Option *option, Params &param, alnFunction alignmentKernel, int subtree=-1);
+        void updateAlignment(Node* node, SequenceDB *database);
 
         namespace cpu 
         {

src/phylogeny.hpp

@@ -43,7 +43,8 @@ namespace phylogeny {
         int alnNum = {0};
         float alnWeight = {0};
         int getAlnNum(int currentTask) {
-            return (currentTask == 2) ? alnNum : seqsIncluded.size();
+            // return (currentTask == 2) ? alnNum : seqsIncluded.size();
+            return alnNum;
         };
         int getAlnLen(int currentTask) {
             return alnLen;

src/progressive.cpp

@@ -4,6 +4,7 @@
 
 #include <chrono>
 #include <functional>
+#include <tbb/parallel_for.h>
 
 
 void msa::progressive::getProgressivePairs(std::vector<std::pair<NodePair,int>>& alnOrder, std::stack<Node*> postStack, int grpID, int mode) {
@@ -190,6 +191,43 @@ void msa::progressive::progressiveAlignment(Tree *T, SequenceDB *database, Optio
     }
 }
 
+void msa::progressive::updateAlignment(Node* node, SequenceDB *database) {
+    std::vector<int> new_seqsIncludes;
+    tbb::this_task_arena::isolate([&] { 
+    tbb::parallel_for(tbb::blocked_range<int>(0, database->sequences.size()), [&](tbb::blocked_range<int> range) {
+    for (int idx = range.begin(); idx < range.end(); ++idx) {
+        auto seq = database->sequences[idx];
+        auto sIdx = seq->id;
+        if (seq->subtreeIdx < -1) {
+            auto aln = database->subtreeAln[seq->subtreeIdx];
+            database->id_map[sIdx]->memCheck(aln.size());
+            int orgIdx = 0;
+            int storeFrom = database->id_map[sIdx]->storage;
+            int storeTo = 1 - storeFrom;
+            for (int k = 0; k < aln.size(); ++k) {
+                if (aln[k] == 0) {
+                    database->id_map[sIdx]->alnStorage[storeTo][k] = database->id_map[sIdx]->alnStorage[storeFrom][orgIdx];
+                    orgIdx++;
+                }
+                else {
+                    database->id_map[sIdx]->alnStorage[storeTo][k] = '-';
+                }
+            }
+            database->id_map[sIdx]->len = aln.size();
+            database->id_map[sIdx]->changeStorage();
+        }
+    }
+    });
+    });
+    for (auto sIdx: node->seqsIncluded) {
+        if (sIdx >= 0) new_seqsIncludes.push_back(sIdx);
+    }
+    for (auto seq: database->sequences) {
+        if (seq->subtreeIdx < 0) new_seqsIncludes.push_back(seq->id);
+    }
+    node->seqsIncluded = new_seqsIncludes;
+    return;
+}
 
 void msa::progressive::msaOnSubtree(Tree *T, SequenceDB *database, Option *option, Params &param, alnFunction alignmentKernel, int subtree) {
     auto progressiveStart = std::chrono::high_resolution_clock::now();
@@ -228,8 +266,11 @@ void msa::progressive::msaOnSubtree(Tree *T, SequenceDB *database, Option *optio
         if (database->currentTask != 2) std::cerr << "Alignment (length: " << T->root->alnLen << ") completed in " << progressiveTime.count() / 1000000000 << " s\n";
         else std::cerr<< "Alignment on " << T->allNodes.size() << " subalignments (length: " << T->root->getAlnLen(database->currentTask) << ") in " << progressiveTime.count() / 1000000 << " ms\n";
     }
-    if (database->fallback_nodes.empty())
+    if (database->fallback_nodes.empty()) {
+        if (option->alnMode == DEFAULT_ALN || option->alnMode == PLACE_W_TREE) updateAlignment(T->root, database);
         return;
+    }
+        
 
     // Adding bad sequences back
     auto badStart = std::chrono::high_resolution_clock::now();
@@ -248,6 +289,7 @@ void msa::progressive::msaOnSubtree(Tree *T, SequenceDB *database, Option *optio
     std::cerr << "Realign profiles that have been deferred. Total profiles/sequences: " << database->fallback_nodes.size() << " / " << badSeqBefore << '\n';
     database->fallback_nodes.clear();
     progressiveAlignment(T, database, option, alnPairsPerLevel,param, cpu::alignmentKernel_CPU);
+    if (option->alnMode == DEFAULT_ALN || option->alnMode == PLACE_W_TREE) updateAlignment(T->root, database);
     // Reset currentTask
     database->currentTask = 0;
     auto badEnd = std::chrono::high_resolution_clock::now();