Merge pull request #84 from wiederm/82-request-for-macos-compatibilit…

…y-and-optional-dependency-remove-conda-package-nnops Allow EndstateCorrectionAMBER class to run only with CPU
wiederm · Oct 2, 2023 · e075570 · e075570
2 parents 86e9493 + 62ffd16
commit e075570
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/endstate_correction/simulation/base.py b/endstate_correction/simulation/base.py
@@ -33,7 +33,7 @@ def __init__(
         name: str = "endstate_correction",
         work_dir: str = "./",
         potential: str = "ani2x",
-        implementation: str = "nnpops",
+        implementation: Literal["nnpops", "torchani"] = "nnpops",
         interpolate: bool = True,
     ):
         self.logger = logging.getLogger("EndstateCorrectionBase")
@@ -52,7 +52,7 @@ def __init__(
             mm_system,
             ml_atoms,
             interpolate=interpolate,
-            implementation="nnpops",
+            implementation=implementation,
         )
         integrator = self.get_integrator()
         platform = get_fastest_platform(minimum_precision="mixed")

diff --git a/endstate_correction/tests/test_perform_endstate_correction.py b/endstate_correction/tests/test_perform_endstate_correction.py
@@ -65,6 +65,7 @@ def ec(tmp_path_factory, bss_protocol):
             protocol=bss_protocol,
             name=system_name,
             work_dir=str(output_base),
+            implementation="torchani",
         )
         simulation.set_trajectory(str(parameter_base / f"{system_name}.h5"))
         return simulation