Fixes to rate-variation

Fix positional errors causing incorrect likelihood calculations when using rate-variation.

Author: SR-Martin

alignment/seqregions.h

@@ -564,6 +564,7 @@ auto updateLHwithModel(const ModelBase* model,
     RealNumType tot = 0;
     if (total_blength > 0)  // TODO: avoid
     {
+
       const RealNumType* mutation_mat_row = model->getMutationMatrixRow(i, pos);
       tot += dotProduct<num_states>(&(prior)[0], mutation_mat_row);
 
@@ -630,7 +631,6 @@ void merge_N_O(const RealNumType lower_plength,
     total_blength = reg_o.plength_observation2node +
                     (lower_plength > 0 ? lower_plength : 0);
   }
-
   auto new_lh =
       cmaple::make_unique<SeqRegion::LHType>();  // = new RealNumType[num_states];
   RealNumType sum_lh = updateLHwithModel<num_states>(model, *reg_o.likelihood,
@@ -765,7 +765,7 @@ void merge_RACGT_O(const SeqRegion& seq2_region,
   assert(seq2_region.type == TYPE_O);
   assert(model);
   assert(aln);
-    
+
   RealNumType sum_new_lh = updateMultLHwithMat<num_states>(
       model->getMutationMatrix(end_pos), *(seq2_region.likelihood), new_lh, total_blength_2);
 
@@ -1043,7 +1043,7 @@ auto merge_O_O_TwoLowers(const SeqRegion& seq2_region,
   assert(seq2_region.type == TYPE_O);
   assert(model);
   assert(aln);
-    
+
   RealNumType sum_lh = updateMultLHwithMat<num_states>(
       model->getMutationMatrix(end_pos), *seq2_region.likelihood, new_lh, total_blength_2);
 
@@ -1179,7 +1179,7 @@ auto merge_RACGT_O_TwoLowers(const SeqRegion& seq2_region,
   assert(seq2_region.type == TYPE_O);
   assert(model);
   assert(aln);
-    
+
   RealNumType sum_lh = updateMultLHwithMat<num_states>(
       model->getMutationMatrix(end_pos), *(seq2_region.likelihood), new_lh, total_blength_2);
 
@@ -1363,7 +1363,7 @@ auto merge_notN_notN_TwoLowers(const SeqRegion& seq1_region,
         aln, model, threshold_prob, log_lh, merged_regions, return_log_lh);
     return ret;
   }
-
+  
   // no error
   return true;
 }
@@ -1447,12 +1447,6 @@ RealNumType SeqRegions::mergeTwoLowers(
     }
     // neither seq1_entry nor seq2_entry = N
     else {
-      //RealNumType* vec = (&seq1_region->likelihood)[0];
-      /*if(!seq1_region->likelihood.get()) {
-        std::cout << "[mergeTwoLowers] Likelihood pointer: " << seq1_region->likelihood.get() << std::endl;
-        return MIN_NEGATIVE;
-      }*/
-      //std::cout << "[mergeTwoLowers] Prior: " << vec[0] << " " << vec[1] << " " << vec[2] << " " << vec[3] << std::endl;
       if (!merge_notN_notN_TwoLowers<num_states>(
               *seq1_region, *seq2_region, plength1, plength2, end_pos, pos, aln,
               model, cumulative_rate, threshold_prob, log_lh, merged_regions,
@@ -1480,7 +1474,6 @@ RealNumType SeqRegions::mergeTwoLowers(
          max_elements);  // ensure we did the correct reserve, otherwise it was
   // a pessimization
 #endif
-
   return log_lh;
 }
 
@@ -1505,7 +1498,7 @@ auto SeqRegions::computeAbsoluteLhAtRoot(
       for (StateType i = 0; i < num_states; ++i) {
         log_lh += model->getRootLogFreq(i) *
                   (cumulative_base[static_cast<size_t>(region.position) + 1][i] -
-                   cumulative_base[static_cast<size_t>(start_pos)][i]);
+                   cumulative_base[static_cast<size_t>(start_pos)][i]); 
       }
     }
     // type ACGT
@@ -1737,7 +1730,7 @@ bool calSiteLhs_O_O(std::vector<RealNumType>& site_lh_contributions,
   assert(seq2_region.type == TYPE_O);
   assert(aln);
   assert(model);
-    
+
   RealNumType sum_lh = updateMultLHwithMat<num_states>(
       model->getMutationMatrix(end_pos), *seq2_region.likelihood, new_lh, total_blength_2);
 
@@ -1869,7 +1862,6 @@ bool calSiteLhs_RACGT_O(std::vector<RealNumType>& site_lh_contributions,
   assert(seq2_region.type == TYPE_O);
   assert(aln);
   assert(model);
-    
   RealNumType sum_lh = updateMultLHwithMat<num_states>(
       model->getMutationMatrix(end_pos), *(seq2_region.likelihood), new_lh, total_blength_2);
 

model/model_dna_rate_variation.cpp

@@ -59,6 +59,9 @@ void ModelDNARateVariation::printCountsAndWaitingTimes(const RealNumType* counts
 }
 
 bool ModelDNARateVariation::updateMutationMatEmpirical() {
+    if(ratesEstimated) {
+        std::cout << "[ModelDNARateVariation] Warning: Overwriting estimated rate matrices with single empirical mutation matrix." << std::endl;
+    }
     bool val = ModelDNA::updateMutationMatEmpirical();
     for(int i = 0; i < genomeSize; i++) {
         for(int j = 0; j < matSize; j++) {
@@ -76,6 +79,7 @@ bool ModelDNARateVariation::updateMutationMatEmpirical() {
 }
 
 void ModelDNARateVariation::estimateRates(cmaple::Tree* tree) {
+    ratesEstimated = true;
     if(useSiteRates) {
         estimateRatePerSite(tree);
 
@@ -109,21 +113,22 @@ void ModelDNARateVariation::estimateRates(cmaple::Tree* tree) {
             outFile << "Rate Matrix: " << std::endl;
             printMatrix(getMutationMatrix(i), &outFile);
             outFile << std::endl;
+
         }
         outFile.close();
     } 
 }
 
 void ModelDNARateVariation::estimateRatePerSite(cmaple::Tree* tree){
     std::cout << "Estimating mutation rate per site..." << std::endl;
-    RealNumType** totals = new RealNumType*[num_states_];
+    RealNumType** waitingTimes = new RealNumType*[num_states_];
     for(int j = 0; j < num_states_; j++) {
-        totals[j] = new RealNumType[genomeSize];
+        waitingTimes[j] = new RealNumType[genomeSize];
     }
     RealNumType* numSubstitutions = new RealNumType[genomeSize];
     for(int i = 0; i < genomeSize; i++) {
         for(int j = 0; j < num_states_; j++) {
-            totals[j][i] = 0;
+            waitingTimes[j][i] = 0;
         }
         numSubstitutions[i] = 0;
     }
@@ -170,12 +175,12 @@ void ModelDNARateVariation::estimateRatePerSite(cmaple::Tree* tree){
                 // both states are type REF
                 for(int i = pos; i <= end_pos; i++) {
                     int state = tree->aln->ref_seq[static_cast<std::vector<cmaple::StateType>::size_type>(i)];
-                    totals[state][i] += blength;
+                    waitingTimes[state][i] += blength;
                 }
             }  else if(seq1_region->type == seq2_region->type && seq1_region->type < TYPE_R) {
                 // both states are equal but not of type REF
                  for(int i = pos; i <= end_pos; i++) {
-                    totals[seq1_region->type][i] += blength;
+                    waitingTimes[seq1_region->type][i] += blength;
                 }               
             } else if(seq1_region->type <= TYPE_R && seq2_region->type <= TYPE_R) {
                 // both states are not equal
@@ -194,7 +199,7 @@ void ModelDNARateVariation::estimateRatePerSite(cmaple::Tree* tree){
         } else {
             RealNumType expectedRateNoSubstitution = 0;
             for(int j = 0; j < num_states_; j++) {
-                RealNumType summand = totals[j][i] * abs(diagonal_mut_mat[j]);
+                RealNumType summand = waitingTimes[j][i] * abs(diagonal_mut_mat[j]);
                 expectedRateNoSubstitution += summand;
             }
             if(expectedRateNoSubstitution <= 0.01) {
@@ -210,7 +215,6 @@ void ModelDNARateVariation::estimateRatePerSite(cmaple::Tree* tree){
     RealNumType averageRate = rateCount / genomeSize;
     for(int i = 0; i < genomeSize; i++) {
         rates[i] /= averageRate;
-        //std::cout << rates[i] << " ";
         for(int stateA = 0; stateA < num_states_; stateA++) {
             RealNumType rowSum = 0;
             for(int stateB = 0; stateB < num_states_; stateB++) {
@@ -236,9 +240,9 @@ void ModelDNARateVariation::estimateRatePerSite(cmaple::Tree* tree){
     }
 
     for(int j = 0; j < num_states_; j++) {
-        delete[] totals[j];
+        delete[] waitingTimes[j];
     }
-    delete[] totals;
+    delete[] waitingTimes;
     delete[] numSubstitutions;
 }
 
@@ -556,4 +560,14 @@ void ModelDNARateVariation::setAllMatricesToDefault() {
 
         }
     }  
+}
+
+void ModelDNARateVariation::setMatrixAtPosition(RealNumType* matrix, PositionType i) {
+    for(int stateA = 0; stateA < num_states_; stateA++) {
+        diagonalMutationMatrices[i * num_states_ + stateA] = matrix[stateA + row_index[stateA]];
+        for(int stateB = 0; stateB < num_states_; stateB++) {
+            mutationMatrices[i * matSize + (stateB + row_index[stateA])] = matrix[stateB + row_index[stateA]];
+            transposedMutationMatrices[i * matSize + (stateB + row_index[stateA])] = matrix[stateA + row_index[stateB]];
+        }
+    }
 }

model/model_dna_rate_variation.h

@@ -72,6 +72,7 @@ class ModelDNARateVariation : public ModelDNA {
   virtual bool updateMutationMatEmpirical() override;
 
   void setAllMatricesToDefault();
+  void setMatrixAtPosition(RealNumType* matrix, PositionType i);
 
   void printMatrix(const RealNumType* matrix, std::ostream* outStream);
   void printCountsAndWaitingTimes(const RealNumType* counts, const RealNumType* waitingTImes, std::ostream* outStream);
@@ -94,6 +95,7 @@ class ModelDNARateVariation : public ModelDNA {
     cmaple::RealNumType* rates = nullptr;
     uint16_t matSize;
     bool useSiteRates = false;
+    bool ratesEstimated = false;
 
 };
 }