Add NPT for ASE Molecular Dynamics

atomistics/calculators/__init__.py

@@ -1,4 +1,5 @@
 from atomistics.calculators.ase import (
+    calc_molecular_dynamics_npt_with_ase,
     calc_molecular_dynamics_langevin_with_ase,
     calc_static_with_ase,
     evaluate_with_ase,

atomistics/calculators/ase.py

@@ -2,6 +2,7 @@
 
 from ase import units
 from ase.md.langevin import Langevin
+from ase.md.npt import NPT
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.constraints import UnitCellFilter
 import numpy as np
@@ -114,6 +115,60 @@ def calc_static_with_ase(
     )
 
 
+def _calc_md_step_with_ase(
+    dyn, structure, ase_calculator, temperature, run, thermo, output
+):
+    structure.calc = ase_calculator
+    MaxwellBoltzmannDistribution(atoms=structure, temperature_K=temperature)
+    cache = {q: [] for q in output}
+    for i in range(int(run / thermo)):
+        dyn.run(thermo)
+        calc_dict = ASEOutputMolecularDynamics.get(
+            ASEExecutor(ase_structure=structure, ase_calculator=ase_calculator),
+            *output,
+        )
+        for k, v in calc_dict.items():
+            cache[k].append(v)
+    return {q: np.array(cache[q]) for q in output}
+
+
+def calc_molecular_dynamics_npt_with_ase(
+    structure,
+    ase_calculator,
+    run=100,
+    thermo=100,
+    timestep=1 * units.fs,
+    ttime=100 * units.fs,
+    pfactor=2e6 * units.GPa * (units.fs**2),
+    temperature=100,
+    externalstress=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
+    * units.bar,
+    output=ASEOutputMolecularDynamics.fields(),
+):
+    return _calc_md_step_with_ase(
+        dyn=NPT(
+            atoms=structure,
+            timestep=timestep,
+            temperature=None,
+            externalstress=externalstress,
+            ttime=ttime,
+            pfactor=pfactor,
+            temperature_K=temperature,
+            mask=None,
+            trajectory=None,
+            logfile=None,
+            loginterval=1,
+            append_trajectory=False,
+        ),
+        structure=structure,
+        ase_calculator=ase_calculator,
+        temperature=temperature,
+        run=run,
+        thermo=thermo,
+        output=output,
+    )
+
+
 def calc_molecular_dynamics_langevin_with_ase(
     structure,
     ase_calculator,
@@ -122,24 +177,22 @@ def calc_molecular_dynamics_langevin_with_ase(
     timestep=1 * units.fs,
     temperature=100,
     friction=0.002,
-    quantities=ASEOutputMolecularDynamics.fields(),
+    output=ASEOutputMolecularDynamics.fields(),
 ):
-    structure.calc = ase_calculator
-    MaxwellBoltzmannDistribution(atoms=structure, temperature_K=temperature)
-    dyn = Langevin(
-        atoms=structure, timestep=timestep, temperature_K=temperature, friction=friction
+    return _calc_md_step_with_ase(
+        dyn=Langevin(
+            atoms=structure,
+            timestep=timestep,
+            temperature_K=temperature,
+            friction=friction,
+        ),
+        structure=structure,
+        ase_calculator=ase_calculator,
+        temperature=temperature,
+        run=run,
+        thermo=thermo,
+        output=output,
     )
-    loops_to_execute = int(run / thermo)
-    cache = {q: [] for q in quantities}
-    for i in range(loops_to_execute):
-        dyn.run(thermo)
-        calc_dict = ASEOutputMolecularDynamics.get(
-            ASEExecutor(ase_structure=structure, ase_calculator=ase_calculator),
-            *quantities,
-        )
-        for k, v in calc_dict.items():
-            cache[k].append(v)
-    return {q: np.array(cache[q]) for q in quantities}
 
 
 def optimize_positions_with_ase(

tests/test_ase_md_emt.py

@@ -1,10 +1,18 @@
 from ase import units
 from ase.build import bulk
 from ase.calculators.emt import EMT
-from atomistics.calculators import calc_molecular_dynamics_langevin_with_ase
+from atomistics.calculators import (
+    calc_molecular_dynamics_langevin_with_ase,
+    calc_molecular_dynamics_npt_with_ase,
+)
+import numpy as np
 import unittest
 
 
+def get_volume(cell):
+    return np.abs(np.linalg.det(cell))
+
+
 class TestASEMD(unittest.TestCase):
     def test_ase_langevin(self):
         structure = bulk("Al", cubic=True).repeat([2, 2, 2])
@@ -26,4 +34,29 @@ def test_ase_langevin(self):
         self.assertEqual(result_dict["energy_tot"].shape, (10, ))
         self.assertEqual(result_dict["pressure"].shape, (10, 3, 3))
         self.assertTrue(result_dict["temperature"][-1] > 25)
-        self.assertTrue(result_dict["temperature"][-1] < 75)
+        self.assertTrue(result_dict["temperature"][-1] < 75)
+
+    def test_ase_npt(self):
+        structure = bulk("Al", a=3.5, cubic=True).repeat([2, 2, 2])
+        result_dict = calc_molecular_dynamics_npt_with_ase(
+            structure=structure,
+            ase_calculator=EMT(),
+            run=100,
+            thermo=10,
+            timestep=1 * units.fs,
+            ttime=100 * units.fs,
+            pfactor=2e6 * units.GPa * (units.fs ** 2),
+            temperature=100,
+            externalstress=np.array([0., 0., 0., 0., 0., 0.]) * units.bar,
+        )
+        self.assertEqual(result_dict["positions"].shape, (10, 32, 3))
+        self.assertEqual(result_dict["velocities"].shape, (10, 32, 3))
+        self.assertEqual(result_dict["cell"].shape, (10, 3, 3))
+        self.assertEqual(result_dict["forces"].shape, (10, 32, 3))
+        self.assertEqual(result_dict["temperature"].shape, (10,))
+        self.assertEqual(result_dict["energy_pot"].shape, (10,))
+        self.assertEqual(result_dict["energy_tot"].shape, (10,))
+        self.assertEqual(result_dict["pressure"].shape, (10, 3, 3))
+        self.assertTrue(result_dict["temperature"][-1] > 50)
+        self.assertTrue(result_dict["temperature"][-1] < 100)
+        self.assertTrue(get_volume(result_dict["cell"][0]) < get_volume(result_dict["cell"][-1]))