modified analysis for amber FF

transformato/analysis.py

@@ -18,7 +18,7 @@
 from pymbar import mbar
 from openmm import unit
 from openmm import Platform, XmlSerializer, vec3
-from openmm.app import CharmmPsfFile, Simulation
+from openmm.app import CharmmPsfFile, Simulation, AmberPrmtopFile
 from tqdm import tqdm
 
 from transformato.constants import temperature
@@ -99,6 +99,9 @@ def __init__(self, configuration: dict, structure_name: str):
         self.thinning: int = 0
         self.save_results_to_path: str = f"{self.configuration['system_dir']}/results/"
         self.traj_files = defaultdict(list)
+        self.forcefield = "charmm"
+        if self.configuration["simulation"]["forcefield"] == "amber":
+            self.forcefield = "amber"
 
     def load_trajs(self, nr_of_max_snapshots: int = 300, multiple_runs: int = 0):
         """
@@ -113,14 +116,19 @@ def load_trajs(self, nr_of_max_snapshots: int = 300, multiple_runs: int = 0):
         )
 
     def _generate_openMM_system(self, env: str, lambda_state: int) -> Simulation:
-        # read in necessary files
+        ### read in necessary xml (_system.xml and _integrator.xml) and topology files (charmm: psf and amber:parm7)
+
         conf_sub = self.configuration["system"][self.structure][env]
         file_name = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}_system.xml"
         system = XmlSerializer.deserialize(open(file_name).read())
         file_name = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}_integrator.xml"
         integrator = XmlSerializer.deserialize(open(file_name).read())
-        psf_file_path = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.psf"
-        psf = CharmmPsfFile(psf_file_path)
+        if self.forcefield == "charmm":
+            psf_file_path = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.psf"
+            psf = CharmmPsfFile(psf_file_path)
+        elif self.forcefield == "amber":
+            psf_file_path = f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.parm7"
+            psf = AmberPrmtopFile(psf_file_path)
 
         # generate simulations object and set states
         if self.configuration["simulation"]["GPU"] == True:
@@ -460,11 +468,18 @@ def energy_at_lambda(
             if not os.path.isfile(dcd_path):
                 raise RuntimeError(f"{dcd_path} does not exist.")
 
-            traj = MDAnalysis.Universe(
-                f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.psf",
-                f"{dcd_path}",
-                in_memory=in_memory,
-            )
+            if self.forcefield == "charmm":
+                traj = MDAnalysis.Universe(
+                    f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.psf",
+                    f"{dcd_path}",
+                    in_memory=in_memory,
+                )
+            elif self.forcefield == "amber":
+                traj = MDAnalysis.Universe(
+                    f"{self.base_path}/intst{lambda_state}/{conf_sub['intermediate-filename']}.parm7",
+                    f"{dcd_path}",
+                    in_memory=in_memory,
+                )
 
             # simple thinning of the Trajectory
             start = int(0.25 * len(traj.trajectory))

transformato/bin/openmm_run.py

@@ -158,7 +158,9 @@
 
 # needed for later analysis
 file_name = f"lig_in_{env}"
-print(file_name)
+state = simulation.context.getState(getPositions=True, getVelocities=True)
+with open(file_name + ".rst", "w") as f:
+    f.write(XmlSerializer.serialize(state))
 with open(file_name + "_integrator.xml", "w") as outfile:
     outfile.write(XmlSerializer.serialize(integrator))
 with open(file_name + "_system.xml", "w") as outfile:

transformato/tests/test_amber.py

@@ -5,6 +5,7 @@
     ProposeMutationRoute,
 )
 from transformato.mutate import perform_mutations
+from transformato.utils import postprocessing
 from transformato.tests.paths import get_test_output_dir
 from transformato_testsystems.testsystems import get_testsystems_dir
 import pytest
@@ -50,7 +51,6 @@ def test_full_amber_test():
     i = IntermediateStateFactory(
         system=s2,
         configuration=configuration,
-        multiple_runs=3,
     )
 
     perform_mutations(
@@ -59,3 +59,63 @@ def test_full_amber_test():
         mutation_list=mutation_list,
         nr_of_mutation_steps_charge=2,
     )
+
+
+@pytest.mark.rbfe
+@pytest.mark.charmm
+@pytest.mark.skipif(
+    os.getenv("CI") == "true",
+    reason="Skipping tests that cannot pass in github actions",
+)
+def test_asfe_amber_test():
+    molecule = "ethane"
+
+    configuration = load_config_yaml(
+        config=f"/site/raid3/johannes/amber_tests/data/config/{molecule}_asfe.yaml",
+        input_dir="/site/raid3/johannes/amber_tests/data/",
+        output_dir="/site/raid3/johannes/amber_tests/",
+    )
+
+    s1 = SystemStructure(configuration, "structure1")
+    s1_to_s2 = ProposeMutationRoute(s1)
+    s1_to_s2.propose_common_core()
+    s1_to_s2.finish_common_core()
+
+    mutation_list = s1_to_s2.generate_mutations_to_common_core_for_mol1()
+    i = IntermediateStateFactory(
+        system=s1,
+        configuration=configuration,
+    )
+    perform_mutations(
+        configuration=configuration,
+        i=i,
+        mutation_list=mutation_list,
+        nr_of_mutation_steps_charge=3,
+    )
+
+
+@pytest.mark.rbfe
+@pytest.mark.charmm
+@pytest.mark.skipif(
+    os.getenv("CI") == "true",
+    reason="Skipping tests that cannot pass in github actions",
+)
+def test_amber_analysis():
+    molecule = "ethane"
+    configuration = load_config_yaml(
+        config=f"/site/raid3/johannes/amber_tests/data/config/{molecule}_asfe.yaml",
+        input_dir="/site/raid3/johannes/amber_tests/data/",
+        output_dir="/site/raid3/johannes/amber_tests/",
+    )
+
+    ddG_openMM, dddG, f_openMM = postprocessing(
+        configuration,
+        name="ethane_amber",
+        engine="openMM",
+        max_snapshots=10000,
+        num_proc=6,
+        analyze_traj_with="mda",
+        show_summary=True,
+    )
+
+    print(f"Free energy is {ddG_openMM} +- {dddG} kT")