Add Molecular Dynamics support for ASE

atomistics/calculators/__init__.py

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

atomistics/calculators/ase.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
+from ase import units
+from ase.md.langevin import Langevin
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.constraints import UnitCellFilter
+import numpy as np
 from typing import TYPE_CHECKING
 
 from atomistics.calculators.interface import get_quantities_from_tasks
-from atomistics.shared.output import OutputStatic
+from atomistics.shared.output import OutputStatic, OutputMolecularDynamics
 from atomistics.calculators.wrapper import as_task_dict_evaluator
 
 if TYPE_CHECKING:
@@ -28,13 +32,41 @@ def get_energy(self):
     def get_stress(self):
         return self.structure.get_stress(voigt=False)
 
+    def get_total_energy(self):
+        return (
+            self.structure.get_potential_energy() + self.structure.get_kinetic_energy()
+        )
+
+    def get_cell(self):
+        return self.structure.get_cell()
+
+    def get_positions(self):
+        return self.structure.get_positions()
+
+    def get_velocities(self):
+        return self.structure.get_velocities()
+
+    def get_temperature(self):
+        return self.structure.get_temperature()
+
 
 ASEOutputStatic = OutputStatic(
     forces=ASEExecutor.get_forces,
     energy=ASEExecutor.get_energy,
     stress=ASEExecutor.get_stress,
 )
 
+ASEOutputMolecularDynamics = OutputMolecularDynamics(
+    positions=ASEExecutor.get_positions,
+    cell=ASEExecutor.get_cell,
+    forces=ASEExecutor.get_forces,
+    temperature=ASEExecutor.get_temperature,
+    energy_pot=ASEExecutor.get_energy,
+    energy_tot=ASEExecutor.get_total_energy,
+    pressure=ASEExecutor.get_stress,
+    velocities=ASEExecutor.get_velocities,
+)
+
 
 @as_task_dict_evaluator
 def evaluate_with_ase(
@@ -82,6 +114,34 @@ def calc_static_with_ase(
     )
 
 
+def calc_molecular_dynamics_langevin_with_ase(
+    structure,
+    ase_calculator,
+    run=100,
+    thermo=100,
+    timestep=1 * units.fs,
+    temperature=100,
+    friction=0.002,
+    quantities=ASEOutputMolecularDynamics.fields(),
+):
+    structure.calc = ase_calculator
+    MaxwellBoltzmannDistribution(atoms=structure, temperature_K=temperature)
+    dyn = Langevin(
+        atoms=structure, timestep=timestep, temperature_K=temperature, friction=friction
+    )
+    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(
     structure, ase_calculator, ase_optimizer, ase_optimizer_kwargs
 ):

tests/test_ase_md_emt.py

@@ -0,0 +1,30 @@
+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
+import unittest
+
+
+class TestASEMD(unittest.TestCase):
+    def test_ase_langevin(self):
+        structure = bulk("Al", cubic=True).repeat([2, 2, 2])
+        result_dict = calc_molecular_dynamics_langevin_with_ase(
+            structure=structure,
+            ase_calculator=EMT(),
+            run=100,
+            thermo=10,
+            timestep=1 * units.fs,
+            temperature=100,
+            friction=0.002,
+        )
+        print(result_dict)
+        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] > 25)
+        self.assertTrue(result_dict["temperature"][-1] < 75)