From 245b046ba276223fe81bcc4c85dc4ba3a5b443e8 Mon Sep 17 00:00:00 2001
From: SimonBoothroyd <simon.boothroyd@hotmail.com>
Date: Sun, 14 Jan 2024 19:28:26 -0500
Subject: [PATCH] Add EQ arg to save samples

---
 absolv/runner.py     | 22 +++++++++++++++++-----
 regression/README.md |  2 ++
 regression/run.py    | 20 +++++++++++++++-----
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/absolv/runner.py b/absolv/runner.py
index d6bb88d..a744345 100644
--- a/absolv/runner.py
+++ b/absolv/runner.py
@@ -178,6 +178,7 @@ def _run_phase_hremd(
     temperature: openmm.unit.Quantity,
     prepared_system: PreparedSystem,
     platform: femto.md.constants.OpenMMPlatform,
+    output_dir: pathlib.Path | None = None,
 ) -> tuple[dict[str, float], dict[str, numpy.ndarray]]:
     platform = (
         femto.md.constants.OpenMMPlatform.REFERENCE
@@ -222,10 +223,11 @@ def _run_phase_hremd(
     )
 
     with tempfile.TemporaryDirectory() as tmp_dir:
-        tmp_dir = pathlib.Path(tmp_dir)
+        output_dir = pathlib.Path(tmp_dir) if output_dir is None else output_dir
+        output_dir.mkdir(parents=True, exist_ok=True)
 
-        femto.md.hremd.run_hremd(simulation, states, hremd_config, tmp_dir)
-        u_kn, n_k = femto.fe.ddg.load_u_kn(tmp_dir / "samples.arrow")
+        femto.md.hremd.run_hremd(simulation, states, hremd_config, output_dir)
+        u_kn, n_k = femto.fe.ddg.load_u_kn(output_dir / "samples.arrow")
 
     return femto.fe.ddg.estimate_ddg(u_kn, n_k, temperature)
 
@@ -235,6 +237,7 @@ def run_eq(
     prepared_system_a: PreparedSystem,
     prepared_system_b: PreparedSystem,
     platform: femto.md.constants.OpenMMPlatform = femto.md.constants.OpenMMPlatform.CUDA,
+    output_dir: pathlib.Path | None = None,
 ) -> absolv.config.Result:
     """Perform a simulation at each lambda window and for each solvent.
 
@@ -243,17 +246,26 @@ def run_eq(
         prepared_system_a: The prepared system a. See ``setup`` for more details.
         prepared_system_b: The prepared system b. See ``setup`` for more details.
         platform: The OpenMM platform to run using.
+        output_dir: The (optional) directory to save HREMD samples to.
     """
 
     results_a, overlap_a = _run_phase_hremd(
-        config.alchemical_protocol_a, config.temperature, prepared_system_a, platform
+        config.alchemical_protocol_a,
+        config.temperature,
+        prepared_system_a,
+        platform,
+        None if output_dir is None else output_dir / "solvant-a",
     )
 
     dg_a, dg_a_std = results_a["ddG_kcal_mol"], results_a["ddG_error_kcal_mol"]
     # overlap_a = overlap_a["overlap_0"]
 
     results_b, overlap_b = _run_phase_hremd(
-        config.alchemical_protocol_b, config.temperature, prepared_system_b, platform
+        config.alchemical_protocol_b,
+        config.temperature,
+        prepared_system_b,
+        platform,
+        None if output_dir is None else output_dir / "solvant-b",
     )
 
     dg_b, dg_b_std = results_b["ddG_kcal_mol"], results_b["ddG_error_kcal_mol"]
diff --git a/regression/README.md b/regression/README.md
index 09959c4..9bd5889 100644
--- a/regression/README.md
+++ b/regression/README.md
@@ -27,3 +27,5 @@ python run.py --solute methanol \
               --replica 0       \
               --replica 1
 ```
+
+The results will be written to `results/{timestamp}/{method}-{solute}-{replica}/results.json`.
diff --git a/regression/run.py b/regression/run.py
index cb01730..e322cad 100644
--- a/regression/run.py
+++ b/regression/run.py
@@ -182,8 +182,14 @@ def run_replica(
         parmed.openmm.load_topology(prepared_system_a.topology.to_openmm()),
     )
 
-    run_fn = absolv.runner.run_neq if method == "neq" else absolv.runner.run_eq
-    result = run_fn(config, prepared_system_a, prepared_system_b, "CUDA")
+    if method == "neq":
+        result = absolv.runner.run_neq(
+            config, prepared_system_a, prepared_system_b, "CUDA"
+        )
+    else:
+        result = absolv.runner.run_eq(
+            config, prepared_system_a, prepared_system_b, "CUDA", output_dir
+        )
 
     (output_dir / "result.json").write_text(result.model_dump_json(indent=2))
 
@@ -220,9 +226,13 @@ def main(solutes: list[str], methods: list[str], replicas: list[str]):
                 logging.info(f"running {method} {solute} {replica}")
 
                 replica_dir = output_dir / f"{method}-{solute}-{replica}"
-                run_replica(
-                    DEFAULT_SYSTEMS[solute], system_generator, method, replica_dir
-                )
+
+                try:
+                    run_replica(
+                        DEFAULT_SYSTEMS[solute], system_generator, method, replica_dir
+                    )
+                except BaseException as e:
+                    logging.exception(f"failed to run {method} {solute} {replica}", e)
 
 
 if __name__ == "__main__":