Merge pull request #88 from Andrew-S-Rosen/mtk

Add IsotropicMTKNPT to `_md.py`

Author: shyuep

src/matcalc/_md.py

@@ -9,6 +9,7 @@
 from ase.md import Langevin
 from ase.md.andersen import Andersen
 from ase.md.bussi import Bussi
+from ase.md.nose_hoover_chain import IsotropicMTKNPT
 from ase.md.npt import NPT
 from ase.md.nptberendsen import Inhomogeneous_NPTBerendsen, NPTBerendsen
 from ase.md.nvtberendsen import NVTBerendsen
@@ -52,6 +53,7 @@ def __init__(
             "npt_nose_hoover",
             "npt_berendsen",
             "npt_inhomogeneous",
+            "npt_isotropic_mtk",
         ] = "nvt",
         temperature: int = 300,
         timestep: float = 1.0,
@@ -65,6 +67,10 @@ def __init__(
         pfactor: float = 75.0**2.0,
         external_stress: float | np.ndarray | None = None,
         compressibility_au: float | None = None,
+        tchain: int = 3,
+        pchain: int = 3,
+        tloop: int = 1,
+        ploop: int = 1,
         trajfile: Any = None,
         logfile: str | None = None,
         loginterval: int = 1,
@@ -83,20 +89,31 @@ def __init__(
             calculator (Calculator): The calculator used for energy, force, and stress evaluations.
                 Default to the provided calculator.
             ensemble (str): Ensemble for MD simulation. Options include "nve", "nvt_langevin",
-                "nvt_andersen", "nvt_bussi", "npt", "npt_berendsen", "npt_nose_hoover". Default to "nvt".
+                "nvt_andersen", "nvt_bussi", "npt", "npt_berendsen", "npt_nose_hoover", "npt_mtk".
+                Default to "nvt".
             temperature (int): Simulation temperature in Kelvin. Default to 300.
             timestep (float): Time step in femtoseconds. Default to 1.0.
             steps (int): Number of MD simulation steps. Default to 100.
             pressure (float): External pressure for NPT simulations (in eV/ų). Default to 1.01325 * units.bar.
             taut (float | None): Time constant for temperature coupling. If None, defaults to 100 * timestep * fs.
+                For npt_isotropic_mtk, this is the time constant for temperature damping.
             taup (float | None): Time constant for pressure coupling. If None, defaults to 1000 * timestep * fs.
+                For npt_isotropic_mtk, this is the time constant for pressure damping.
             friction (float): Friction coefficient for Langevin dynamics. Default to 1.0e-3.
             andersen_prob (float): Collision probability for Andersen thermostat. Default to 1.0e-2.
             ttime (float): Characteristic time scale for the thermostat in ASE units (fs). Default to 25.0.
             pfactor (float): Barostat differential equation constant. Default to 75.0**2.0.
             external_stress (float | np.ndarray | None): External stress applied to the system.
                 If not provided, defaults to 0.0.
             compressibility_au (float | None): Material compressibility in ų/eV. Default to None.
+            tchain (int): The number of thermostat variables in the Nose-Hoover thermostat. Default to 3.
+                Only used by IsotropicMTKNPT.
+            pchain (int): The number of barostat variables in the Nose-Hoover barostat. Default to 3.
+                Only used by IsotropicMTKNPT.
+            tloop (int): The number of sub-steps in thermostat integration. Default to 1.
+                Only used by IsotropicMTKNPT.
+            ploop (int): T The number of sub-steps in barostat integration. Default to 1.
+                Only used by IsotropicMTKNPT.
             trajfile (Any): Trajectory object or file for storing simulation data. Default to None.
             logfile (str | None): Filename for simulation logs. Default to None.
             loginterval (int): Interval (in steps) for logging simulation data. Default to 1.
@@ -125,6 +142,10 @@ def __init__(
         self.pfactor = pfactor
         self.external_stress = external_stress
         self.compressibility_au = compressibility_au
+        self.tchain = tchain
+        self.pchain = pchain
+        self.tloop = tloop
+        self.ploop = ploop
         self.trajfile = trajfile
         self.logfile = logfile
         self.loginterval = loginterval
@@ -266,11 +287,30 @@ def _initialize_md(self, atoms: Atoms) -> Any:
                 loginterval=self.loginterval,
                 append_trajectory=self.append_trajectory,
             )
+        elif self.ensemble.lower() == "npt_isotropic_mtk":
+            md = IsotropicMTKNPT(
+                atoms,
+                timestep=timestep_fs,
+                temperature_K=self.temperature,
+                pressure_au=self.pressure,
+                tdamp=taut,
+                pdamp=taup,
+                tchain=self.tchain,
+                pchain=self.pchain,
+                tloop=self.tloop,
+                ploop=self.ploop,
+                trajectory=self.trajfile,
+                logfile=self.logfile,
+                loginterval=self.loginterval,
+                append_trajectory=self.append_trajectory,
+            )
+
         else:
             raise ValueError(
                 "The specified ensemble is not supported, choose from 'nve', 'nvt',"
                 " 'nvt_nose_hoover', 'nvt_berendsen', 'nvt_langevin', 'nvt_andersen',"
-                " 'nvt_bussi', 'npt', 'npt_nose_hoover', 'npt_berendsen', 'npt_inhomogeneous'."
+                " 'nvt_bussi', 'npt', 'npt_nose_hoover', 'npt_berendsen', 'npt_inhomogeneous',"
+                " 'npt_isotropic_mtk'."
             )
         return md
 

tests/test_md.py

@@ -19,15 +19,16 @@
 @pytest.mark.parametrize(
     ("ensemble", "expected_energy"),
     [
-        ("nve", -10.74606),
-        ("nvt", -10.81289),
-        ("nvt_berendsen", -10.73112),
-        ("nvt_langevin", -10.78238),
-        ("nvt_andersen", -10.82750),
-        ("nvt_bussi", -10.77756),
-        ("npt_inhomogeneous", -10.74737),
-        ("npt_berendsen", -10.74714),
-        ("npt_nose_hoover", -10.86120),
+        ("nve", -10.839367595419139),
+        ("nvt", -10.752140577997002),
+        ("nvt_berendsen", -10.79344712397928),
+        ("nvt_langevin", -10.719885845311552),
+        ("nvt_andersen", -10.838280482248559),
+        ("nvt_bussi", -10.83345229703825),
+        ("npt_inhomogeneous", -10.778117238180538),
+        ("npt_berendsen", -10.797141692994748),
+        ("npt_nose_hoover", -10.773028687488921),
+        ("npt_isotropic_mtk", -10.817933454243002),
     ],
 )
 def test_md_calc(
@@ -46,8 +47,6 @@ def test_md_calc(
         calculator=matpes_calculator,
         ensemble=ensemble,
         temperature=300,
-        taut=0.1,
-        taup=0.1,
         steps=10,
         frames=5,
         compressibility_au=1,
@@ -63,7 +62,7 @@ def test_md_calc(
     assert "kinetic_energy" in results
     assert "total_energy" in results
 
-    assert results["total_energy"] == pytest.approx(expected_energy, rel=1e-1)
+    assert results["total_energy"] == pytest.approx(expected_energy, rel=1e-2)
 
     energies = np.array(results["trajectory"].total_energies)