Merge pull request #24 from iqtree/piyumal

Fix issue for MCMCTree dating when there is a leaf at first child of the super tree and Formatted code for Hessain calculation.

Author: bqminh

main/phyloanalysis.cpp

@@ -2862,38 +2862,31 @@ void printMCMCFileFormat(PhyloTree *tree, MatrixXd &hessian, stringstream &tree_
                 counter++;
             }
             // here we check whether the first child of the partition tree is a leaf in the case of missing taxa.
-            if (leftSingle && counter==1 && (nei->link_neighbors[part_id]->node->neighbors.size()==1 || nei1->link_neighbors[part_id]->node->neighbors.size()==1))
-            {
+            if (leftSingle && counter == 1 && (nei->link_neighbors[part_id]->node->neighbors.size() == 1 || nei1->
+                link_neighbors[part_id]->node->neighbors.size() == 1)){
                 leftSingleRoot = true;
                 counter++;
             }
         }
         // Note: if we get an unrooted tree with left child is a leaf, we need to push it to back as the last branch
         // This is a special case in MCMCTree
-        if (leftSingleRoot)
-        {
+        if (leftSingleRoot){
             auto partBrM = partBrMmap.front();
             vector<pair<int, int>> partBrMapleftSingle;
             vector<pair<int, int>> partBrMmap2;
             int leftSinglePartId = partBrM.second;
 
-            for (auto it : partBrMmap)
-            {
-                if (it.second != leftSinglePartId)
-                {
+            // Note: multiples of same id could be mapped to the original tree
+            for (auto it : partBrMmap){
+                if (it.second != leftSinglePartId){
                     partBrMmap2.emplace_back(it.first, it.second);
-                }
-                else
-                {
+                }else{
                     partBrMapleftSingle.emplace_back(it.first, it.second);
                 }
             }
-
-            for (auto it : partBrMapleftSingle)
-            {
+            for (auto it : partBrMapleftSingle){
                 partBrMmap2.emplace_back(it.first, it.second);
             }
-
             partBrMmap = partBrMmap2;
         }
         int numStates = tree->getModel()->num_states;
@@ -2910,21 +2903,17 @@ void printMCMCFileFormat(PhyloTree *tree, MatrixXd &hessian, stringstream &tree_
         auto *hessian_diagonal_part = aligned_alloc<double>(branchNum);
         memset(hessian_diagonal_part, 0, branchNum * sizeof(double));
 
-//        double *G_matrix_sub_tree = tree->G_matrix;
+        double *G_matrix_sub_tree = tree->G_matrix;
         if (superTree->params->partition_type != BRLEN_OPTIMIZE) {
             for (auto mapping: partBrMmap) {
                 int stree_branch_id = mapping.first;
                 int part_branch_id = mapping.second;
-                for (int i = 0; i < nPtn; i++) {
-                    G_matrix_part[stree_branch_id * nPtn + i] = tree->G_matrix[part_branch_id * nPtn + i];
-                }
-                //memcpy(G_matrix_part+(sizeof(double)*stree_branch_id*nPtn), G_matrix_sub_tree+sizeof(double)*(part_branch_id*nPtn), sizeof(double)*nPtn);
+                memmove(G_matrix_part+stree_branch_id*nPtn, G_matrix_sub_tree+part_branch_id*nPtn, sizeof(double)*nPtn);
                 gradient_vector_part[stree_branch_id] = tree->gradient_vector[part_branch_id];
                 hessian_diagonal_part[stree_branch_id] = tree->hessian_diagonal[part_branch_id];
             }
             saveTreeMCMCTree(branchLengths, branch_lengths_vector, tree, tree_stream);
         } else {
-
             DoubleVector branchLengths2;
             int branchCounter = 0;
             std::unordered_map<int, int> brVisitedMap;
@@ -2935,29 +2924,65 @@ void printMCMCFileFormat(PhyloTree *tree, MatrixXd &hessian, stringstream &tree_
                     continue;
                 }
                 brVisitedMap[part_branch_id] = 1;
-                for (int i = 0; i < nPtn; i++) {
-                    G_matrix_part[branchCounter * nPtn + i] = tree->G_matrix[part_branch_id * nPtn + i];
-                }
-                //memcpy(G_matrix_part+(sizeof(double)*stree_branch_id*nPtn), G_matrix_sub_tree+sizeof(double)*(part_branch_id*nPtn), sizeof(double)*nPtn);
+                memmove(G_matrix_part+branchCounter*nPtn, G_matrix_sub_tree+part_branch_id*nPtn, sizeof(double)*nPtn);
                 gradient_vector_part[branchCounter] = tree->gradient_vector[part_branch_id];
                 hessian_diagonal_part[branchCounter] = tree->hessian_diagonal[part_branch_id];
                 branchLengths2.push_back(branchLengths[part_branch_id]);
                 branchCounter++;
             }
             saveTreeMCMCTree(branchLengths2, branch_lengths_vector, tree, tree_stream);
         }
-
         processDervMCMCTree(gradient_vector_part, branchNum, nPtn, G_matrix_part, ptn_freq_diagonal,
                             gradient_vector_eigen, hessian, hessian_diagonal_part);
         aligned_free(G_matrix_part);
         aligned_free(gradient_vector_part);
         aligned_free(hessian_diagonal_part);
 
     } else {
-        processDervMCMCTree(tree->gradient_vector, nBranches, nPtn, tree->G_matrix, ptn_freq_diagonal,
-                            gradient_vector_eigen, hessian, tree->hessian_diagonal);
-        saveTreeMCMCTree(branchLengths, branch_lengths_vector, tree, tree_stream);
 
+        int numStates = tree->getModel()->num_states;
+        size_t mem_size = get_safe_upper_limit(tree->getAlnNSite()) + max(get_safe_upper_limit(numStates),
+                                                                          get_safe_upper_limit(tree->getModelFactory()->unobserved_ptns.size()));
+        BranchVector singleAlnBranches;
+        tree->getBranches(singleAlnBranches);
+        vector<int> branch_ids;
+        for (Branch branch: singleAlnBranches){
+            auto nei = branch.first->findNeighbor(branch.second);
+            branch_ids.push_back(nei->id);
+        }
+        if (leftSingle){
+            int single_root = branch_ids.front();
+            branch_ids.erase(branch_ids.begin());
+            branch_ids.push_back(single_root);
+        }
+        size_t branchNum = tree->branchNum;
+        size_t g_matrix_size = branchNum * mem_size;
+        auto *G_matrix_new = aligned_alloc<double>(g_matrix_size);
+        memset(G_matrix_new, 0, g_matrix_size * sizeof(double));
+        auto *gradient_vector_new = aligned_alloc<double>(branchNum);
+        memset(gradient_vector_new, 0, branchNum * sizeof(double));
+        auto *hessian_diagonal_new = aligned_alloc<double>(branchNum);
+        memset(hessian_diagonal_new, 0, branchNum * sizeof(double));
+
+        double *G_matrix_tree = tree->G_matrix;
+        int br_counter = 0;
+        DoubleVector branchLengthsNew;
+        for (int  br_id: branch_ids) {
+            memmove(G_matrix_new+br_counter*nPtn, G_matrix_tree+br_id*nPtn, sizeof(double)*nPtn);
+            gradient_vector_new[br_counter] = tree->gradient_vector[br_id];
+            hessian_diagonal_new[br_counter] = tree->hessian_diagonal[br_id];
+            branchLengthsNew.push_back(branchLengths[br_id]);
+            br_counter++;
+        }
+
+
+        processDervMCMCTree(gradient_vector_new, nBranches, nPtn, G_matrix_new, ptn_freq_diagonal,
+                            gradient_vector_eigen, hessian, hessian_diagonal_new);
+        saveTreeMCMCTree(branchLengthsNew, branch_lengths_vector, tree, tree_stream);
+
+        aligned_free(G_matrix_new);
+        aligned_free(gradient_vector_new);
+        aligned_free(hessian_diagonal_new);
     }
 }
 
@@ -2968,7 +2993,7 @@ void generateDummyAlignment(PhyloTree* tree, ofstream &dummyAlignment){
     // auto alignment = tree->aln->getPattern(0);
     tree->getOrderedTaxa(nodeVector);
     dummyAlignment << "   " << nodeVector.size() << "   " << 1 << endl;
-    for (Node *node: nodeVector) {
+    for (const Node *node: nodeVector) {
         dummyAlignment << node->name << "              " << "A" << endl;
     }
     dummyAlignment << endl;
@@ -2982,7 +3007,6 @@ void printMCMCTreeCtlFile(IQTree *iqtree, ofstream &ctl, ofstream &dummyAlignmen
     } else {
         ndata = 1;
     }
-
     StrVector mcmc_iter_vec;
     convert_string_vec(Params::getInstance().mcmc_iter.c_str(), mcmc_iter_vec, ',');
     ctl << "seed = -1" << endl
@@ -3094,38 +3118,10 @@ void printHessian(IQTree *iqtree, int partition_type) {
             RowVectorXd partition_branch_lengths_vector;
             RowVectorXd partition_gradient_vector_eigen;
 
-            if (it->traversal_starting_node)
-            {
+            if (it->traversal_starting_node){
                 auto part_traversal_starting_nei = (Neighbor*)(((Node*)it->traversal_starting_node)->neighbors[0]->node)
                 ->findNeighbor((Node*)it->traversal_starting_node);
                 it->root = part_traversal_starting_nei->node;
-
-                // This is special case in MCMCtree: if the fist son of root has only one taxa, then select the next neighbour.
-                // if (part_traversal_starting_nei->node->neighbors.size() == 1 && it->leftSingleRoot)
-                // {
-                //     // leftSingle = true;
-                //     part_traversal_starting_nei = (PhyloNeighbor*)((part_traversal_starting_nei->node)->
-                //         findNeighbor((part_traversal_starting_nei->node->neighbors[0]->node)));
-                //     it->root = part_traversal_starting_nei->node;
-                //
-                //     NeighborVec neighbors = part_traversal_starting_nei->node->neighbors;
-                //     auto nei1 = part_traversal_starting_nei->node->neighbors[0];
-                //     auto nei2 = part_traversal_starting_nei->node->neighbors[1];
-                //     auto nei3 = part_traversal_starting_nei->node->neighbors[2];
-                //     if (part_traversal_starting_nei->node->neighbors[0]->node == it->traversal_starting_node)
-                //     {
-                //         part_traversal_starting_nei->node->neighbors[0] = nei2;
-                //         part_traversal_starting_nei->node->neighbors[1] = nei3;
-                //         part_traversal_starting_nei->node->neighbors[2] = nei1;
-                //     }
-                //     else if (part_traversal_starting_nei->node->neighbors[1]->node == it->traversal_starting_node)
-                //     {
-                //         part_traversal_starting_nei->node->neighbors[0] = nei1;
-                //         part_traversal_starting_nei->node->neighbors[1] = nei3;
-                //         part_traversal_starting_nei->node->neighbors[2] = nei2;
-                //     }
-                //
-                // }
             }
 
             printMCMCFileFormat(it, partition_hessian, partition_tree_stream, partition_branch_lengths_vector,
@@ -3179,36 +3175,20 @@ void printHessian(IQTree *iqtree, int partition_type) {
         size_t nBranches = iqtree->branchNum;
 
         // iqtree->sortTaxa();
-
-        if (iqtree->traversal_starting_node)
-        {
+        if (iqtree->traversal_starting_node){
             auto traversal_starting_nei = (Neighbor*)(((Node*)iqtree->traversal_starting_node)->neighbors[0]->node)
                 ->findNeighbor((Node*)iqtree->traversal_starting_node);
             iqtree->root = traversal_starting_nei->node;
-            iqtree->sortTaxa();
+            // iqtree->sortTaxa();
             iqtree->initializeTree();
 
             // This is special case in MCMCtree: if the fist son of root has only one taxa, then select the next neighbour.
-            if (traversal_starting_nei->node->neighbors.size() == 1)
-            {
-                traversal_starting_nei = (PhyloNeighbor*)((traversal_starting_nei->node)->findNeighbor(
-                    (traversal_starting_nei->node->neighbors[0]->node)));
-                iqtree->root = traversal_starting_nei->node;
-                // iqtree->sortTaxa();
-                // iqtree->initializeTree();
-                NeighborVec neighbors = traversal_starting_nei->node->neighbors;
-                auto nei1 = traversal_starting_nei->node->neighbors[0];
-                auto nei2 = traversal_starting_nei->node->neighbors[1];
-                auto nei3 = traversal_starting_nei->node->neighbors[2];
-                traversal_starting_nei->node->neighbors[0] = nei2;
-                traversal_starting_nei->node->neighbors[1] = nei3;
-                traversal_starting_nei->node->neighbors[2] = nei1;
-                // iqtree->sortTaxa();
-                iqtree->initializeTree();
+            if (traversal_starting_nei->node->neighbors.size() == 1){
+                iqtree->leftSingleRoot = true;
             }
         }
 
-        printMCMCFileFormat(iqtree, hessian, tree_stream, branch_lengths_vector, gradient_vector_eigen);
+        printMCMCFileFormat(iqtree, hessian, tree_stream, branch_lengths_vector, gradient_vector_eigen, NULL, 0, iqtree->leftSingleRoot);
 
         outfile << endl << iqtree->aln->getNSeq() << endl << endl;
         outfile << tree_stream.str() << endl << endl;
@@ -3225,7 +3205,7 @@ void printHessian(IQTree *iqtree, int partition_type) {
     alnFile.close();
 
     cout << endl << "Gradients and Hessians written to: " << outFileName << endl;
-    cout << "Ctl file for MCMCTree written to: " << ctlFileName << endl;
+    cout << "Ctl file for MCMCTree written to: " << ctlFileName << endl << endl;
     // cout << "Add time records calibrations to: " << iqtree->params->user_file << endl;
 }
 
@@ -3260,12 +3240,10 @@ SuperNeighbor* findRootedNeighbour(SuperNeighbor* super_root, int part_id) {
     if (!super_root) {
         return nullptr;
     }
-
     // Check if the root itself satisfies the condition
     if (super_root->link_neighbors[part_id]) {
         return super_root;
     }
-
     // Use a queue for level-order traversal
     std::queue<SuperNeighbor*> q;
     q.push(super_root);
@@ -3274,19 +3252,16 @@ SuperNeighbor* findRootedNeighbour(SuperNeighbor* super_root, int part_id) {
     while (!q.empty()) {
         auto current = q.front();
         q.pop();
-
         // Iterate over the current node's neighbors
         for (auto nei : current->node->neighbors) {
             auto super_nei = dynamic_cast<SuperNeighbor*>(nei);
             if (!super_nei) {
                 continue; // Skip invalid neighbors
             }
-
             // Check if this neighbor satisfies the condition
             if (super_nei->link_neighbors[part_id]) {
                 return super_nei;
             }
-
             // Add the neighbor to the queue for further exploration
             q.push(super_nei);
         }
@@ -3355,25 +3330,23 @@ void startTreeReconstruction(Params &params, IQTree* &iqtree, ModelCheckpoint &m
     // if users want to perform tree dating (with mcmc)
     // and if ModelFinder was run, the traversal starting node was incidently deleted (after copyTree and restoreCheckpoint)
     // we have to delete tree nodes to force IQ-TREE to re-read the tree from the treefile
-    if (params.dating_method == "mcmctree" && params.dating_mf)
-    {
+
+    if (params.dating_method == "mcmctree" && params.dating_mf){
         // if it's a supertree, delete all tree members
-        if (iqtree->isSuperTree())
-        {
+        if (iqtree->isSuperTree()){
             PhyloSuperTree* stree = (PhyloSuperTree*) iqtree;
+
             // delete member trees one by one
-            for (PhyloSuperTree::iterator it = stree->begin(); it != stree->end(); it++)
-            {
-                if ((*it)->root)
-                {
+            for (PhyloSuperTree::iterator it = stree->begin(); it != stree->end(); it++){
+                if ((*it)->root){
+
                     (*it)->freeNode();
                     (*it)->root = NULL;
                 }
             }
         }
         // delete the tree itself
-        if (iqtree->root)
-        {
+        if (iqtree->root){
             iqtree->freeNode();
             iqtree->root = NULL;
         }
@@ -3960,15 +3933,14 @@ void runTreeReconstruction(Params &params, IQTree* &iqtree) {
     //check for dating with MCMCTree
     if (params.dating_method == "mcmctree")
     {
-
         cout << endl << "--- Generating the gradients and the Hessian for MCMCTree ---" << endl;
         if (iqtree->isSuperTree())
         {
             auto* stree = (PhyloSuperTree*)iqtree;
             int part_id = 0;
             for (auto& it : *stree)
             {
-                auto* partition_tree = (PhyloTree*)it;
+                auto* partition_tree = it;
                 bool leftSingleRoot = false;
 
                 // If we memorized the traversal starting node -> find the corresponding traversal starting node for the current partition
@@ -3982,13 +3954,17 @@ void runTreeReconstruction(Params &params, IQTree* &iqtree) {
                     // identify whether there is missing taxa at the root. This is needed if we have a clade of 2 taxa at
                     // the root in the rooted tree. When unrooted, they become two children for root which is are leaves.
                     // if leftSingleRoot == true, that means the leaves at the root are not from a 2 leaves clade.
-                    for(auto linkedNei: nei->node->neighbors){
+                    for (auto linkedNei : nei->node->neighbors)
+                    {
                         auto superLinkedNei = (SuperNeighbor*)linkedNei;
-                        if(!superLinkedNei->link_neighbors[part_id]){
+                        if (!superLinkedNei->link_neighbors[part_id])
+                        {
                             leftSingleRoot = true;
+                        }else
+                        {
+                            it->root_available = false;
                         }
                     }
-
                     // this is the other case where we don't miss any nodes at the root but a leaf at the root
                     if (!leftSingleRoot && linked_super_nei->node->neighbors.size() == 1)
                     {

main/timetree.cpp

@@ -367,9 +367,8 @@ void computeHessian(PhyloTree *tree) {
     if (tree->traversal_starting_node && tree->root != tree->traversal_starting_node){
         tree->root = (Node *) tree->traversal_starting_node;
     }
-
     // sort the internal nodes according to their smallest taxon id
-//    tree->sortTaxa();
+    // tree->sortTaxa();
     tree->clearBranchDirection();
     tree->initializeTree();
     tree->computeBranchDirection();
@@ -388,7 +387,6 @@ void computeHessian(PhyloTree *tree) {
             cout << "lh: " << lh << " df: " << df << " ddf: " << ddf << endl;
         }
     }
-
 }
 
 void runMCMCTree(PhyloTree *tree) {

tree/phylotree.h

@@ -841,7 +841,7 @@ class PhyloTree : public MTree, public Optimization, public CheckpointFactory {
     int num_packets;
 
     /** flag to identify partition-trees with missing root for MCMCTree branch traversal order*/
-    bool leftSingleRoot;
+    bool leftSingleRoot = false;
 
     /****************************************************************************
             helper functions for computing tree traversal
@@ -1586,6 +1586,11 @@ class PhyloTree : public MTree, public Optimization, public CheckpointFactory {
      * */
     double *hessian_diagonal;
 
+    /**
+    Flag to check root node availability in case of missing data in partitions
+    * */
+    bool root_available = true;
+
 
     /****************************************************************************
             Nearest Neighbor Interchange by maximum likelihood