Add an option for dispersion corrections (#11)

* Add an option for dispersion corrections

* Minor formatting corrections

* Minor formatting correction

* Remove whitespace

Author: trdurrant

README.md

@@ -200,6 +200,9 @@ after optimisation
 │                                                                      [default: cpu]              │
 │ --model                                  TEXT                        ML model to use.            │
 │                                                                      [default: medium-omat-0]    │
+│ --dispersion            --no-dispersion                              Include a dispersion        │
+│                                                                      correction.                 │
+│                                                                      [default: no-dispersion]    │
 │ --arch                                   TEXT                        MLIP architecture to use.   │
 │                                                                      [default: mace_mp]          │
 │ --temperature                            FLOAT                       Temperature for the Monte   │

pack_mm/cli/packmm.py

@@ -113,6 +113,7 @@ def packmm(
         "cpu", help="Device to run calculations on (e.g., 'cpu' or 'cuda')."
     ),
     model: str = Option("medium-omat-0", help="ML model to use."),
+    dispersion: bool = Option(False, help="Include a dispersion correction."),
     arch: str = Option("mace_mp", help="MLIP architecture to use."),
     temperature: float = Option(
         300.0, help="Temperature for the Monte Carlo acceptance rule."
@@ -152,6 +153,7 @@ def packmm(
     print(f"{cell_c=}")
     print(f"{arch=}")
     print(f"{model=}")
+    print(f"{dispersion=}")
     print(f"{device=}")
     print(f"{temperature=}")
     print(f"{fmax=}")

pack_mm/core/core.py

@@ -191,6 +191,7 @@ def pack_molecules(
     nmols: int = -1,
     arch: str = "mace_mp",
     model: str = "medium-omat-0",
+    dispersion: bool = False,
     device: str = "cpu",
     where: str = "anywhere",
     center: tuple[float, float, float] = None,
@@ -226,6 +227,7 @@ def pack_molecules(
         nmols (int): Number of molecules to insert.
         arch (str): Architecture for the calculator.
         model (str): Path to the model file.
+        dispersion (bool): If a dispersion correction is added to MLIP calculator.
         device (str): Device to run calculations on (e.g., "cpu" or "cuda").
         where (str): Region to insert molecules ("anywhere",
                      "sphere", "cylinderZ", etc.).
@@ -297,7 +299,12 @@ def pack_molecules(
         cell, center, where, a, b, c, radius, height
     )
 
-    calc = choose_calculator(arch=arch, model_path=model, device=device)
+    calc = choose_calculator(
+        arch=arch,
+        model_path=model,
+        device=device,
+        dispersion=dispersion,
+    )
     sys.calc = calc
 
     e = sys.get_potential_energy() if len(sys) > 0 else 0.0
@@ -319,7 +326,14 @@ def pack_molecules(
             tsys = csys.copy() + mol.copy()
             if insert_strategy == "hmc":
                 tsys = run_md_nve(
-                    tsys, md_temperature, md_steps, md_timestep, arch, model, device
+                    tsys,
+                    md_temperature,
+                    md_steps,
+                    md_timestep,
+                    arch,
+                    model,
+                    dispersion,
+                    device,
                 )
 
             if every > 0 and _itry / every == 0:
@@ -328,6 +342,7 @@ def pack_molecules(
                     device=device,
                     arch=arch,
                     model=model,
+                    dispersion=dispersion,
                     fmax=fmax,
                     out_path=out_path,
                     md_temperature=md_temperature,
@@ -362,6 +377,7 @@ def pack_molecules(
                 device,
                 arch,
                 model,
+                dispersion,
                 fmax,
                 out_path,
                 md_temperature,
@@ -382,6 +398,7 @@ def pack_molecules(
             fmax=fmax,
             out_path=out_path,
             opt_cell=True,
+            dispersion=dispersion,
         )
     return (energy_final, csys)
 
@@ -430,6 +447,7 @@ def save_the_day(
     device: str = "",
     arch: str = "",
     model: str = "",
+    dispersion: bool = False,
     fmax: float = 0.01,
     out_path: str = ".",
     md_temperature: float = 100.0,
@@ -446,11 +464,19 @@ def save_the_day(
             model,
             fmax,
             out_path,
+            dispersion=dispersion,
         )
         return a
     if relax_strategy == "md":
         return run_md_nve(
-            struct, md_temperature, md_steps, md_timestep, arch, model, device
+            struct,
+            md_temperature,
+            md_steps,
+            md_timestep,
+            arch,
+            model,
+            dispersion,
+            device,
         )
     return None
 
@@ -462,6 +488,7 @@ def run_md_nve(
     timestep: float = 1.0,
     arch: str = "",
     model: str = "",
+    dispersion: bool = False,
     device: str = "",
 ) -> Atoms:
     """Run nve simulation."""
@@ -470,7 +497,7 @@ def run_md_nve(
         temp=temp,
         device=device,
         arch=arch,
-        calc_kwargs={"model_paths": model},
+        calc_kwargs={"model_paths": model, "dispersion": dispersion},
         stats_every=1,
         steps=steps,
         timestep=timestep,
@@ -487,14 +514,15 @@ def optimize_geometry(
     fmax: float,
     out_path: str = ".",
     opt_cell: bool = False,
+    dispersion: bool = False,
 ) -> tuple(float, Atoms):
     """Optimize the geometry of a structure."""
     geo = GeomOpt(
         struct=struct,
         device=device,
         arch=arch,
         fmax=fmax,
-        calc_kwargs={"model_paths": model},
+        calc_kwargs={"model_paths": model, "dispersion": dispersion},
         filter_kwargs={"hydrostatic_strain": opt_cell},
     )
     geo.run()

tests/test_cli.py

@@ -194,6 +194,13 @@ def test_packmm_custom_threshold():
     assert "threshold=0.9" in strip_ansi_codes(result.output)
 
 
+def test_packmm_with_dispersion():
+    """Check dispersion."""
+    result = runner.invoke(app, ["--dispersion"])
+    assert result.exit_code == 0
+    assert "dispersion=True" in strip_ansi_codes(result.output)
+
+
 def test_packmm_no_geometry_optimization():
     """Check optimisation."""
     result = runner.invoke(app, ["--no-geometry"])

tests/test_core.py

@@ -248,6 +248,26 @@ def test_optimize_geometry(tmp_path):
     assert optimized_energy == pytest.approx(-14.17098106193308, abs=err)
 
 
+def test_optimize_geometry_dispersion(tmp_path):
+    """Test go with dispersion."""
+    molecule = build_molecule("H2O")
+    molecule.set_cell([10, 10, 10])
+    molecule.set_pbc([True, True, True])
+    structure_file = tmp_path / "water.cif"
+    write(structure_file, molecule)
+    optimized_energy, _ = optimize_geometry(
+        str(structure_file),
+        device="cpu",
+        arch="mace_mp",
+        model="small-0b2",
+        fmax=0.01,
+        out_path=tmp_path,
+        opt_cell=True,
+        dispersion=True,
+    )
+    assert optimized_energy == pytest.approx(-14.180876622530265, abs=err)
+
+
 def test_pack_molecules(tmp_path):
     """Test pack molecule."""
     system = Atoms(
@@ -337,6 +357,34 @@ def test_pack_molecules_2(tmp_path, capsys):
     assert e == pytest.approx(-47.194755808249454, abs=err)
 
 
+def test_pack_molecules_atoms_dispersion(tmp_path):
+    """Test pack molecule with dispersion."""
+    system = Atoms(
+        "Ca", positions=[(2.5, 2.5, 2.5)], cell=[5, 5, 5], pbc=[True, True, True]
+    )
+
+    e, _ = pack_molecules(
+        system=system,
+        molecule="H2O",
+        nmols=3,
+        arch="mace_mp",
+        model="small-0b2",
+        dispersion=True,
+        device="cpu",
+        where="sphere",
+        center=(2.5, 2.5, 2.5),
+        radius=2.5,
+        seed=2042,
+        temperature=300,
+        ntries=10,
+        geometry=False,
+        fmax=0.1,
+        out_path=tmp_path,
+    )
+    # Without dispersion: -42.386622937612316
+    assert e == pytest.approx(-42.93461715053182, abs=err)
+
+
 def test_save_the_day(tmp_path):
     """Test save the day."""
     molecule = build_molecule("H2O")
@@ -377,6 +425,26 @@ def test_save_the_day_md(tmp_path):
     assert s[0].position == pytest.approx([4.99684244, 5.00440785, 5.2987255], abs=err)
 
 
+def test_save_the_day_dispersion(tmp_path):
+    """Test save the day with dispersion."""
+    molecule = build_molecule("H2O")
+    molecule.set_cell([10, 10, 10])
+    molecule.set_pbc([True, True, True])
+    molecule.center()
+    structure_file = tmp_path / "water.cif"
+    write(structure_file, molecule)
+    s = save_the_day(
+        str(structure_file),
+        device="cpu",
+        arch="mace_mp",
+        model="small-0b2",
+        dispersion=True,
+        fmax=0.01,
+        out_path=tmp_path,
+    )
+    assert s[0].position == pytest.approx([4.99998938, 4.99618389, 5.30704305], abs=err)
+
+
 def test_save_the_day_invalid(tmp_path):
     """Test save the day."""
     molecule = build_molecule("H2O")
@@ -414,3 +482,22 @@ def test_run_md(tmp_path):
         temp=100.0,
     )
     assert s[0].position == pytest.approx([4.99684244, 5.00440785, 5.2987255], abs=err)
+
+
+def test_run_md_dispersion(tmp_path):
+    """Test md with dispersion."""
+    molecule = build_molecule("H2O")
+    molecule.set_cell([10, 10, 10])
+    molecule.set_pbc([True, True, True])
+    molecule.center()
+    s = run_md_nve(
+        molecule,
+        device="cpu",
+        arch="mace_mp",
+        model="small-0b2",
+        dispersion=True,
+        timestep=1.0,
+        steps=10.0,
+        temp=100.0,
+    )
+    assert s[0].position == pytest.approx([4.99684244, 5.00440795, 5.29872694], abs=err)