Merge pull request #75 from wiederm/absolute_solvation

absolute solvation

Author: JohannesKarwou

pytest.ini

@@ -1,10 +1,10 @@
 [pytest]
 markers =
+    asfe: part of the workflow for absolute solvation free energy calculation.
     rsfe: part of the workflow for a relative solvation free energy calculation.
     postprocessing: part of the postprocessing workflow. 
     rbfe: part of the workflow for a relative binding free energy calculation.
     full_workflow: performs a complete free energy calculation
-    requires_loeffler_systems: test that requires that the loeffler system are installed.
     requires_charmm_installation: test requires CHARMM installation.
     charmm: charmm test
     restraints: All tests for restraints.py

transformato/analysis.py

@@ -2,7 +2,7 @@
 import gc
 import logging
 import multiprocessing as mp
-import os
+import os,sys
 import pickle
 import subprocess
 from collections import defaultdict
@@ -21,7 +21,7 @@
 from tqdm import tqdm
 
 from transformato.constants import temperature
-from transformato.utils import get_structure_name
+from transformato.utils import get_structure_name, isnotebook
 
 logger = logging.getLogger(__name__)
 
@@ -76,7 +76,10 @@ def __init__(self, configuration: dict, structure_name: str):
         # decide if the name of the system corresponds to structure1 or structure2
         structure = get_structure_name(configuration, structure_name)
 
-        if configuration["simulation"]["free-energy-type"] == "rsfe":
+        if (
+            configuration["simulation"]["free-energy-type"] == "rsfe"
+            or configuration["simulation"]["free-energy-type"] == "asfe"
+        ):
             self.envs = ("vacuum", "waterbox")
             self.mbar_results = {"waterbox": None, "vacuum": None}
 
@@ -96,15 +99,15 @@ def __init__(self, configuration: dict, structure_name: str):
         self.save_results_to_path: str = f"{self.configuration['system_dir']}/results/"
         self.traj_files = defaultdict(list)
 
-    def load_trajs(self, nr_of_max_snapshots: int = 300):
+    def load_trajs(self, nr_of_max_snapshots: int = 300, multiple_runs: int = 0):
         """
         load trajectories, thin trajs and merge themn.
         Also calculate N_k for mbar.
         """
 
         assert type(nr_of_max_snapshots) == int
         self.nr_of_max_snapshots = nr_of_max_snapshots
-        self.snapshots, self.unitcell, self.nr_of_states, self.N_k = self._merge_trajs()
+        self.snapshots, self.unitcell, self.nr_of_states, self.N_k = self._merge_trajs(multiple_runs)
 
     def _generate_openMM_system(self, env: str, lambda_state: int) -> Simulation:
         # read in necessary files
@@ -156,7 +159,7 @@ def _thinning(self, any_list):
             further_thinning,
         )
 
-    def _merge_trajs(self) -> Tuple[dict, dict, int, dict]:
+    def _merge_trajs(self, multiple_runs: int) -> Tuple[dict, dict, int, dict]:
         """
         load trajectories, thin trajs and merge themn.
         Also calculate N_k for mbar.
@@ -179,7 +182,11 @@ def _merge_trajs(self) -> Tuple[dict, dict, int, dict]:
             unitcell_ = []
             conf_sub = self.configuration["system"][self.structure][env]
             for lambda_state in tqdm(range(1, nr_of_states + 1)):
-                dcd_path = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.dcd"
+                if multiple_runs:
+                    dcd_path = f"{self.base_path}/intst{lambda_state}/run_{multiple_runs}/{conf_sub['intermediate-filename']}.dcd"
+                else:
+                    dcd_path = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.dcd"
+                
                 psf_path = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.psf"
                 if not os.path.isfile(dcd_path):
                     raise RuntimeError(f"{dcd_path} does not exist.")
@@ -425,9 +432,8 @@ def _evaluate_e_with_openMM(
         return energies
 
     def energy_at_lambda(
-        self, lambda_state: int, env: str, nr_of_max_snapshots: int, in_memory: bool
+        self, lambda_state: int, env: str, nr_of_max_snapshots: int, in_memory: bool, multiple_runs: int,
     ) -> Tuple:
-
         gc.enable()
         logger.info(f"Analysing lambda state {lambda_state} of {self.nr_of_states}")
         conf_sub = self.configuration["system"][self.structure][env]
@@ -440,7 +446,11 @@ def energy_at_lambda(
 
         for lambda_state in range(1, self.nr_of_states + 1):
 
-            dcd_path = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.dcd"
+            if not multiple_runs:
+                dcd_path = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.dcd"
+            else:
+                dcd_path = f"{self.base_path}/intst{lambda_state}/run_{multiple_runs}/{conf_sub['intermediate-filename']}.dcd"
+
             if not os.path.isfile(dcd_path):
                 raise RuntimeError(f"{dcd_path} does not exist.")
 
@@ -497,10 +507,13 @@ def _analyse_results_using_mda(
         engine: str,
         num_proc: int,
         nr_of_max_snapshots: int,
+        multiple_runs: int,
         in_memory: bool = False,
     ):
 
         logger.info(f"Evaluating with {engine}, using {num_proc} CPUs")
+        if not os.path.isdir(self.base_path):
+            sys.exit(f"{self.base_path} does not exist")
         self.nr_of_states = len(next(os.walk(f"{self.base_path}"))[1])
 
         if engine == "openMM":
@@ -520,6 +533,7 @@ def _analyse_results_using_mda(
                         repeat(env),
                         repeat(nr_of_max_snapshots),
                         repeat(in_memory),
+                        repeat(multiple_runs),
                     ),
                 )
             )
@@ -535,12 +549,19 @@ def _analyse_results_using_mda(
             )
 
         if save_results:
-            file = f"{self.save_results_to_path}/mbar_data_for_{self.structure_name}_in_{env}.pickle"
+            if multiple_runs:
+                file = f"{self.save_results_to_path}/mbar_data_for_{self.structure_name}_in_{env}_run_{multiple_runs}.pickle"
+            else:
+                file = f"{self.save_results_to_path}/mbar_data_for_{self.structure_name}_in_{env}.pickle"
+
             logger.info(f"Saving results: {file}")
             results = {"u_kn": u_kn, "N_k": N_k}
             pickle.dump(results, open(file, "wb+"))
-
-        return self.calculate_dG_using_mbar(u_kn, N_k, env)
+            
+            return self.calculate_dG_using_mbar(u_kn, N_k, env)
+        
+        else:
+            return self.calculate_dG_using_mbar(u_kn, N_k, env)
 
     def _analyse_results_using_mdtraj(
         self,
@@ -611,6 +632,7 @@ def calculate_dG_to_common_core(
         num_proc: int = 1,
         in_memory: bool = False,
         nr_of_max_snapshots: int = -1,
+        multiple_runs: int = 0,
     ):
         """
         Calculate mbar results using either the python package mdtraj
@@ -636,13 +658,14 @@ def calculate_dG_to_common_core(
                     engine,
                     num_proc,
                     nr_of_max_snapshots,
+                    multiple_runs,
                 )
             elif analyze_traj_with == "mdtraj":
                 self.mbar_results[env] = self._analyse_results_using_mdtraj(
                     env, self.snapshots[env], self.unitcell[env], save_results, engine
                 )
             else:
-                raise RuntimeError("Either mda or mdtray")
+                raise RuntimeError("Either mda or mdtraj")
 
     def load_waterbox_results(self, file: str):
         self.mbar_results["waterbox"] = self._load_mbar_results(file)
@@ -765,8 +788,8 @@ def plot_free_energy_overlap(self, env: str):
         plt.savefig(
             f"{self.save_results_to_path}/ddG_to_common_core_overlap_{env}_for_{self.structure_name}.png"
         )
-
-        plt.show()
+        if isnotebook():
+            plt.show()
         plt.close()
 
     def plot_free_energy(self, env: str):
@@ -805,31 +828,17 @@ def plot_free_energy(self, env: str):
         plt.savefig(
             f"{self.save_results_to_path}/ddG_to_common_core_line_plot_{env}_for_{self.structure_name}.png"
         )
-        plt.show()
+        if isnotebook():
+            plt.show()
         plt.close()
 
-    def plot_vacuum_free_energy_overlap(self):
-        self.plot_free_energy_overlap("vacuum")
-
-    def plot_complex_free_energy_overlap(self):
-        self.plot_free_energy_overlap("complex")
-
-    def plot_waterbox_free_energy_overlap(self):
-        self.plot_free_energy_overlap("waterbox")
-
-    def plot_vacuum_free_energy(self):
-        self.plot_free_energy("vacuum")
-
-    def plot_complex_free_energy(self):
-        self.plot_free_energy("complex")
-
-    def plot_waterbox_free_energy(self):
-        self.plot_free_energy("waterbox")
-
     @property
     def end_state_free_energy_difference(self):
         """DeltaF[lambda=1 --> lambda=0]"""
-        if self.configuration["simulation"]["free-energy-type"] == "rsfe":
+        if (
+            self.configuration["simulation"]["free-energy-type"] == "rsfe"
+            or self.configuration["simulation"]["free-energy-type"] == "asfe"
+        ):
             return (
                 self.waterbox_free_energy_differences[0, -1]
                 - self.vacuum_free_energy_differences[0, -1],
@@ -848,26 +857,20 @@ def end_state_free_energy_difference(self):
             raise RuntimeError()
 
     def show_summary(self):
-        from transformato.utils import isnotebook
-
-        if self.configuration["simulation"]["free-energy-type"] == "rsfe":
-            if isnotebook:
-                # only show this if we are in a notebook
-                self.plot_vacuum_free_energy_overlap()
-                self.plot_waterbox_free_energy_overlap()
-            self.plot_vacuum_free_energy()
-            self.plot_waterbox_free_energy()
-            self.detailed_overlap("waterbox")
-            self.detailed_overlap("vacuum")
+
+        if (
+            self.configuration["simulation"]["free-energy-type"] == "rsfe"
+            or self.configuration["simulation"]["free-energy-type"] == "asfe"
+        ):
+            self.plot_free_energy_overlap("vacuum")
+            self.plot_free_energy_overlap("waterbox")
+            self.plot_free_energy("vacuum")
+            self.plot_free_energy("waterbox")
         else:
-            if isnotebook:
-                # only show this if we are in a notebook
-                self.plot_complex_free_energy_overlap()
-                self.plot_waterbox_free_energy_overlap()
-            self.plot_complex_free_energy()
-            self.plot_waterbox_free_energy()
-            self.detailed_overlap("complex")
-            self.detailed_overlap("waterbox")
+            self.plot_free_energy_overlap("waterbox")
+            self.plot_free_energy_overlap("complex")
+            self.plot_free_energy("waterbox")
+            self.plot_free_energy("complex")
 
         energy_estimate, uncertainty = self.end_state_free_energy_difference
         print(