From ab7bc8a6c73a63a75c158446811f23d483eace00 Mon Sep 17 00:00:00 2001
From: BenoitMorel <benomorel@gmail.com>
Date: Sun, 11 Aug 2024 22:13:15 +0200
Subject: [PATCH] Fix MPI issues with per family rates

---
 src/ale/AleEvaluator.cpp | 2 ++
 src/ale/AleOptimizer.cpp | 7 +++++--
 src/ale/AleOptimizer.hpp | 7 ++++++-
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/ale/AleEvaluator.cpp b/src/ale/AleEvaluator.cpp
index b4ec33c..2c7d428 100644
--- 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
index e2accb7..b47b771 100644
--- 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
index 3e79db7..2225c3f 100644
--- 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);