Merge pull request #24 from michellab/bugfix-calcset-calc-attrs

Ensure SimulationRunnerIterator saves states when tearing down

Author: fjclark

a3fe/_version.py

@@ -1 +1 @@
-__version__ = "0.3.0"
+__version__ = "0.3.1"

a3fe/read/_process_somd_files.py

@@ -137,7 +137,7 @@ def read_overlap_mat(outfile: str) -> _np.ndarray:
     overlap_mat = []
     in_overlap_mat = False
     for line in lines:
-        if line.startswith("#Overlap matrix"):
+        if "#Overlap matrix" in line:
             in_overlap_mat = True
             continue
         if line.startswith("#"):

a3fe/run/_simulation_runner.py

@@ -235,15 +235,6 @@ def delta_g_er(self) -> _np.ndarray:
 
         return self._delta_g_er
 
-    @property
-    def is_complete(self) -> bool:
-        f"""Whether the {self.__class__.__name__} has completed."""
-        # Check if the overall_stats.dat file exists
-        if _pathlib.Path(f"{self.output_dir}/overall_stats.dat").is_file():
-            return True
-        else:
-            return False
-
     def __str__(self) -> str:
         return self.__class__.__name__
 
@@ -439,7 +430,6 @@ def analyse(
     def get_results_df(
         self,
         save_csv: bool = True,
-        get_normalised_costs: bool = False,
         add_sub_sim_runners: bool = True,
     ) -> _pd.DataFrame:
         """

a3fe/run/_utils.py

@@ -153,18 +153,19 @@ def __iter__(self):
         return self
 
     def __next__(self) -> _T:
-        if self.i >= len(self.base_dirs):
-            # Tear down the current simulation runner
-            if self.current_sim_runner is not None:
-                self.current_sim_runner._close_logging_handlers()
-                del self.current_sim_runner
-                self.current_sim_runner = None
-            raise StopIteration
+        end_of_list = self.i >= len(self.base_dirs)
 
         # Tear down the current simulation runner
         if self.current_sim_runner is not None:
+            self.current_sim_runner._dump()
             self.current_sim_runner._close_logging_handlers()
             del self.current_sim_runner
+            if end_of_list:
+                self.current_sim_runner = None
+
+        if end_of_list:
+            raise StopIteration
+
         # Set up the next simulation runner
         self.current_sim_runner = self.subclass(
             **self.kwargs, base_dir=self.base_dirs[self.i]

a3fe/run/calc_set.py

@@ -452,3 +452,36 @@ def analyse(
                 offset=offset,
                 stats=stats,
             )
+
+    # Avoid base class methods which aren't valid being called
+    @property
+    def delta_g(self):
+        raise AttributeError("CalcSet objects do not have a delta_g attribute.")
+
+    @property
+    def delta_g_err(self):
+        raise AttributeError("CalcSet objects do not have a delta_g_err attribute.")
+
+    def get_results_df(self, save_csv: bool = True, add_sub_sim_runners: bool = True):
+        # TODO: Implement this method
+        raise NotImplementedError(
+            "This method is not implemented for CalcSet objects. Use the analyse method to get free energy changes."
+        )
+
+    def analyse_convergence(
+        self,
+        slurm: bool = False,
+        run_nos: _Optional[_List[int]] = None,
+        mode: str = "cumulative",
+        fraction: float = 1,
+        equilibrated: bool = True,
+    ):
+        """
+        Not implemented for CalcSet objects as convergence analysis is expensive.
+        Call the analyse_convergence method on individual calculations instead (or
+        run analyse convergence on each calculation in the set, if you're determined).
+        """
+        raise NotImplementedError(
+            "This method is not implemented for CalcSet objects (due to high computational ",
+            "expense. Call the analyse_convergence method on individual calculations instead.",
+        )

a3fe/run/calculation.py

@@ -8,6 +8,7 @@
 import shutil as _shutil
 from typing import List as _List
 from typing import Optional as _Optional
+from pathlib import Path as _Path
 
 from ._simulation_runner import SimulationRunner as _SimulationRunner
 from .enums import LegType as _LegType
@@ -154,6 +155,15 @@ def prep_stage(self) -> _PreparationStage:
 
         return self._prep_stage
 
+    @property
+    def is_complete(self) -> bool:
+        f"""Whether the {self.__class__.__name__} has completed."""
+        # Check if the overall_stats.dat file exists
+        if _Path(f"{self.output_dir}/overall_stats.dat").is_file():
+            return True
+
+        return False
+
     def setup(
         self,
         bound_leg_sysprep_config: _Optional[_SystemPreparationConfig] = None,

a3fe/run/lambda_window.py

@@ -405,7 +405,3 @@ def analyse_convergence(self) -> None:
 
     def setup(self) -> None:
         raise NotImplementedError("LamWindows are set up when they are created")
-
-    @property
-    def is_complete(self) -> bool:
-        raise NotImplementedError()

a3fe/run/simulation.py

@@ -658,7 +658,3 @@ def analyse_convergence(self) -> None:
 
     def setup(self) -> None:
         raise NotImplementedError("Simulations are set up when they are created")
-
-    @property
-    def is_complete(self) -> bool:
-        raise NotImplementedError()

a3fe/tests/test_calc_set.py

@@ -52,3 +52,28 @@ def test_calc_set_analysis(calc_set):
     assert results_exp.loc["mdm2_pip2_short", "calc_er"] == pytest.approx(
         0.0935, abs=1e-2
     )
+
+
+def test_calc_set_sub_sim_attrs(calc_set):
+    """Test that the attributes of the sub-simulation runners are correctly saved."""
+
+    def check_attr(sim_runner, attr_name="test_attr", value=42, force=True):
+        """Recursively check that the attribute is set for all sub-simulation runners."""
+        if force:
+            assert hasattr(sim_runner, attr_name)
+            assert getattr(sim_runner, attr_name) == value
+        else:
+            if hasattr(sim_runner, attr_name):
+                assert getattr(sim_runner, attr_name) == value
+
+        for sub_sim_runner in sim_runner._sub_sim_runners:
+            check_attr(sub_sim_runner)
+
+    # Try forcing an attribute to be set
+    calc_set.recursively_set_attr("test_attr", force=True, value=42)
+    check_attr(calc_set, attr_name="test_attr", value=42, force=True)
+
+    # Also check that we can set the equilibration time attribute
+    calc_set.set_equilibration_time(0.77)
+    check_attr(calc_set, attr_name="_equilibration_time", value=0.77, force=False)
+    check_attr(calc_set, attr_name="_equilibrated", value=True, force=False)

docs/CHANGELOG.rst

@@ -2,6 +2,11 @@
 Change Log
 ===============
 
+0.3.1
+====================
+- Modified the SimulationRunnerIterator to ensure that the state of sub-simulation runners is saved before they are torn down. This means that setting the equilibration time at the CalcSet level will now work as expected (previously the state of the calculations was not saved, causing the analysis to fail).
+- Updated read_overlap_mat so that it works when "#Overlap" is not printed to a new line (but is included at the end of the previous line). This seems to happen when MBAR requires many iterations to converge.
+
 0.3.0
 ====================