feat: use ASE for all minimizations (#73)

Author: chrbrunk

src/mlipaudit/benchmarks/small_molecule_minimization/small_molecule_minimization.py

@@ -51,16 +51,14 @@
     "simulation_type": "minimization",
     "num_steps": 1000,
     "snapshot_interval": 10,
-    "num_episodes": 10,
-    "timestep_fs": 0.1,
+    "max_force_convergence_threshold": 0.01,
 }
 
 SIMULATION_CONFIG_FAST = {
     "simulation_type": "minimization",
     "num_steps": 10,
     "snapshot_interval": 1,
-    "num_episodes": 1,
-    "timestep_fs": 0.1,
+    "max_force_convergence_threshold": 0.01,
 }
 
 RMSD_SCORE_THRESHOLD = 0.075

src/mlipaudit/utils/simulation.py

@@ -27,8 +27,6 @@
 
 logger = logging.getLogger("mlipaudit")
 
-DEFAULT_ASE_MAX_FORCE_CONV_THRESH = 0.01
-
 
 class ASESimulationEngineWithCalculator(ASESimulationEngine):
     """Class derived from mlip's ASE simulation engine but allowing for a passed
@@ -73,9 +71,15 @@ def __init__(
 
 def get_simulation_engine(
     atoms: ase.Atoms, force_field: ForceField | ASECalculator, **kwargs
-) -> JaxMDSimulationEngine | ASESimulationEngineWithCalculator:
+) -> JaxMDSimulationEngine | ASESimulationEngineWithCalculator | ASESimulationEngine:
     """Returns the correct simulation engine based on the input force field type.
 
+    For MD simulations with `mlip.models.ForceField` objects, we return a
+    `JaxMDSimulationEngine`. For energy minimizations with those objects, we return
+    a `ASESimulationEngine`. For any type of simulations with ASE calculator objects,
+    we return a `ASESimulationEngineWithCalculator`, which is a custom class of
+    the `mlipaudit` library.
+
     Args:
         atoms: The ASE atoms.
         force_field: The force field, either an `mlip.models.ForceField`
@@ -89,21 +93,26 @@ def get_simulation_engine(
     Raises:
         ValueError: If force field type is not compatible.
     """
-    if isinstance(force_field, ForceField):
+    # Case 1: MD simulations with ForceField objects -> use JAX-MD
+    if (
+        isinstance(force_field, ForceField)
+        and kwargs.get("simulation_type", "md") == "md"
+    ):
         md_config = JaxMDSimulationEngine.Config(**kwargs)
         return JaxMDSimulationEngine(atoms, force_field, md_config)
 
-    elif isinstance(force_field, ASECalculator):
-        kwargs_copy = deepcopy(kwargs)
-        kwargs_copy.pop("num_episodes", None)  # remove this if exists
+    kwargs_copy = deepcopy(kwargs)
+    kwargs_copy.pop("num_episodes", None)  # remove this if exists
 
-        # for minimization:
-        kwargs_copy["max_force_convergence_threshold"] = (
-            DEFAULT_ASE_MAX_FORCE_CONV_THRESH
-        )
+    # Case 2: Minimization with ForceField objects -> use ASE
+    if isinstance(force_field, ForceField):
+        minimization_config = ASESimulationEngine.Config(**kwargs_copy)
+        return ASESimulationEngine(atoms, force_field, minimization_config)
 
-        md_config = ASESimulationEngine.Config(**kwargs_copy)
-        return ASESimulationEngineWithCalculator(atoms, force_field, md_config)
+    # Case 3: MD or minimization with ASECalculator objects -> use ASE
+    if isinstance(force_field, ASECalculator):
+        sim_config = ASESimulationEngine.Config(**kwargs_copy)
+        return ASESimulationEngineWithCalculator(atoms, force_field, sim_config)
 
     raise ValueError(
         "Provided force field must be either a mlip-compatible "

tests/small_molecule_minimization/test_minimization.py

@@ -73,13 +73,13 @@ def _gen_sim_output(mol_name: str, num_atoms: int) -> MoleculeSimulationOutput:
     "small_mol_minimization_benchmark", [True, False], indirect=True
 )
 def test_full_run_with_mocked_engine(
-    small_mol_minimization_benchmark, mock_jaxmd_simulation_engine
+    small_mol_minimization_benchmark, mock_ase_simulation_engine
 ):
     """Integration test testing a full run of the benchmark."""
     benchmark = small_mol_minimization_benchmark
-    mock_engine = mock_jaxmd_simulation_engine()
+    mock_engine = mock_ase_simulation_engine()
     with patch(
-        "mlipaudit.utils.simulation.JaxMDSimulationEngine",
+        "mlipaudit.utils.simulation.ASESimulationEngine",
         return_value=mock_engine,
     ) as mock_engine_class:
         benchmark.run_model()