Fix MPI issues with per family rates

BenoitMorel · Aug 11, 2024 · ab7bc8a · ab7bc8a
1 parent c013727
commit ab7bc8a
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/src/ale/AleEvaluator.cpp b/src/ale/AleEvaluator.cpp
@@ -362,6 +362,7 @@ double AleEvaluator::optimizeModelRates(bool thorough)
       settings.optimizationMinImprovement = settings.lineSearchMinImprovement;
     }
     if (_info.perFamilyRates) {
+      _optimizer.enableCheckpoints(false); // to avoid MPI issues 
       for (unsigned int family = 0; family < _evaluations.size(); ++family) {
         DTLFamilyParametersOptimizer function(*this, family);
         auto categorizedParameters = getOptimizationClasses().getCompressedParameters(_modelParameters[family].getParameters());
@@ -371,6 +372,7 @@ double AleEvaluator::optimizeModelRates(bool thorough)
               settings);
         function.setParameters(bestParameters);
       }
+      _optimizer.enableCheckpoints(true); // to avoid MPI issues 
       ll = computeLikelihood();
       Logger::timed << "[Species search]   After model rate opt, ll=" << ll << std::endl;
     } else {

diff --git a/src/ale/AleOptimizer.cpp b/src/ale/AleOptimizer.cpp
@@ -47,7 +47,8 @@ AleOptimizer::AleOptimizer(
   _info(info),
   _outputDir(outputDir),
   _checkpointDir(getCheckpointDir(outputDir)),
-  _rootLikelihoods(_geneTrees.getTrees().size())
+  _rootLikelihoods(_geneTrees.getTrees().size()),
+  _enableCheckpoints(true)
 {
 
   if (checkpointExists()) {
@@ -564,6 +565,8 @@ bool AleOptimizer::checkpointExists(const std::string &outputDir)
 
 void AleOptimizer::onBetterParametersFoundCallback()
 {
-  saveCheckpoint();
+  if (_enableCheckpoints) {
+    saveCheckpoint();
+  }
 }
 
diff --git a/src/ale/AleOptimizer.hpp b/src/ale/AleOptimizer.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+
 #include "AleState.hpp"
 #include <search/SpeciesRootSearch.hpp>
 #include <trees/SpeciesTree.hpp>
@@ -56,6 +57,10 @@ class AleOptimizer: public SpeciesTree::Listener,
    */
   void optimize();
 
+  void enableCheckpoints(bool enable) {
+    _enableCheckpoints = enable;
+  }
+
   /**
    *  Optize the species tree root
    */
@@ -155,7 +160,7 @@ class AleOptimizer: public SpeciesTree::Listener,
   std::string _checkpointDir;
   std::unique_ptr<SpeciesSearchState> _speciesTreeSearchState;
   RootLikelihoods _rootLikelihoods;
-
+  bool _enableCheckpoints;  
   double sprSearch(unsigned int radius);
   double transferSearch();
   std::string saveCurrentSpeciesTreeId(std::string str = "inferred_species_tree.newick", bool masterRankOnly = true);