Add Symplectic Integrator

examples/al/.gitignore

@@ -0,0 +1,9 @@
+*.xyz*
+*.extxyz*
+*.out*
+*RESTART*
+outputs
+*.pt
+wandb
+*.ckpt
+RESTART

examples/al/create-datasets.py

@@ -0,0 +1,67 @@
+import copy
+
+import ase
+import ase.build
+import ase.io
+import ase.units
+from ase.calculators.emt import EMT
+from ase.md import VelocityVerlet
+from ase.md.langevin import Langevin
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+
+
+# We start by creating a simple system (a small box of aluminum).
+atoms = ase.build.bulk("Al", "fcc", cubic=True) * (2, 2, 2)
+
+# We first equilibrate the system at 300K using a Langevin thermostat.
+MaxwellBoltzmannDistribution(atoms, temperature_K=300)
+atoms.calc = EMT()
+dyn = Langevin(
+    atoms, 2 * ase.units.fs, temperature_K=300, friction=1 / (100 * ase.units.fs)
+)
+dyn.run(1000)  # 2 ps equilibration (around 10 ps is better in practice)
+
+# Then, we run a production simulation in the NVE ensemble.
+trajectory = []
+
+
+def store_trajectory():
+    trajectory.append(copy.deepcopy(atoms))
+
+
+dyn = VelocityVerlet(atoms, 1 * ase.units.fs)
+dyn.attach(store_trajectory, interval=1)
+dyn.run(2000)  # 2 ps NVE run
+
+time_lag = 32
+spacing = 200
+
+def get_structure_for_dataset_m2d(frame_now, frame_ahead):
+    s = copy.deepcopy(frame_now)
+    s.arrays["delta_positions"] = (
+        frame_ahead.get_positions() - frame_now.get_positions()
+    )
+    s.arrays["delta_momenta"] = frame_ahead.get_momenta() - frame_now.get_momenta()
+    s.set_positions(0.5 * (frame_now.get_positions() + frame_ahead.get_positions()))
+    s.set_momenta(0.5 * (frame_now.get_momenta() + frame_ahead.get_momenta()))
+    return s
+
+def get_structure_for_dataset_s2e(frame_now, frame_ahead):
+    s = copy.deepcopy(frame_now)
+    s.arrays["future_positions"] = frame_ahead.get_positions()
+    s.arrays["future_momenta"] = frame_ahead.get_momenta()
+    return s
+
+
+structures_for_dataset_m2d = []
+structures_for_dataset_s2e = []
+for i in range(0, len(trajectory) - time_lag, spacing):
+    frame_now = trajectory[i]
+    frame_ahead = trajectory[i + time_lag]
+    s_m2d = get_structure_for_dataset_m2d(frame_now, frame_ahead)
+    s_s2e = get_structure_for_dataset_s2e(frame_now, frame_ahead)
+    structures_for_dataset_m2d.append(s_m2d)
+    structures_for_dataset_s2e.append(s_s2e)
+
+ase.io.write("data/midpoint-to-delta.xyz", structures_for_dataset_m2d)
+ase.io.write("data/start-to-end.xyz", structures_for_dataset_s2e)

examples/al/input.xml

@@ -0,0 +1,33 @@
+<simulation verbosity='low' threading='false'>
+   <total_steps>100</total_steps>
+   <output prefix='md'>
+      <trajectory stride='1' filename='pos' format='ase'> positions </trajectory>
+      <trajectory stride='1' filename='vel' format='xyz'> velocities </trajectory>
+      <properties stride='1'> [ step, time{picosecond}, conserved, temperature{kelvin} ] </properties> 
+   </output>
+   <prng>
+      <seed>32123</seed>
+   </prng>
+   <ffdirect name='mlip'>
+      <pes>metatomic</pes>
+      <parameters>{model: ../models/mlip_pet-omatpes-v2.pt, template: ../data/equilibrated.xyz, device: cuda}</parameters>
+    </ffdirect>
+   <system>
+      <forces> 
+         <force forcefield='mlip'></force>
+      </forces>
+      <initialize nbeads='1'>
+         <file mode='ase'>../data/equilibrated.xyz</file>
+         <velocities mode='thermal' units='kelvin'>300</velocities>
+      </initialize>
+      <ensemble>
+         <temperature units='kelvin'>300</temperature>
+      </ensemble>
+      <motion mode='dynamics'>
+         <dynamics mode='nvt'>
+            <timestep units='femtosecond'>32</timestep>
+            <thermostat mode='svr'><tau units='femtosecond'>2</tau></thermostat>
+         </dynamics>
+      </motion>
+  </system>
+</simulation>

examples/al/options-flashmd-symplectic.yaml

@@ -0,0 +1,55 @@
+seed: 42
+base_precision: 32
+
+architecture:
+  name: experimental.flashmd_symplectic
+  training:
+    timestep: 32  # in this case 30 (time lag) * 1 fs (timestep of reference MD)
+    batch_size: 8  # to be increased in a production scenario
+    num_epochs: 100  # to be increased (at least 1000-10000) in a production scenario
+    log_interval: 1
+    learning_rate: 3e-4
+    fixed_scaling_weights:
+      positions: 1.0
+      momenta: 1.0
+    loss:
+      positions:
+        type: mse
+        weight: 1.0
+        reduction: mean
+      momenta:
+        type: mse
+        weight: 1.0
+        reduction: mean
+
+training_set:
+  systems:
+    read_from: data/midpoint-to-delta.xyz
+    length_unit: A
+  targets:
+    positions:
+      key: delta_positions
+      quantity: length
+      unit: A
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true
+    momenta:
+      key: delta_momenta
+      quantity: momentum
+      unit: (eV*u)^(1/2)
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true
+
+validation_set: 0.1
+test_set: 0.0
+
+wandb:
+  project: flashmd-variants
+  name: symplectic-flashmd 
+  tags:
+    - al
+    - symplectic-flashmd

examples/al/options-flashmd.yaml

@@ -0,0 +1,50 @@
+seed: 42
+
+architecture:
+  name: experimental.flashmd
+  training:
+    timestep: 32  # in this case 32 (time lag) * 1 fs (timestep of reference MD)
+    batch_size: 8  # to be increased in a production scenario
+    num_epochs: 100  # to be increased (at least 1000-10000) in a production scenario
+    log_interval: 1
+    loss:
+      positions:
+        type: mse
+        weight: 1.0
+        reduction: mean
+      momenta:
+        type: mse
+        weight: 1.0
+        reduction: mean
+
+training_set:
+  systems:
+    read_from: data/start-to-end.xyz
+    length_unit: A
+  targets:
+    positions:
+      key: future_positions
+      quantity: length
+      unit: A
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true
+    momenta:
+      key: future_momenta
+      quantity: momentum
+      unit: (eV*u)^(1/2)
+      type:
+        cartesian:
+          rank: 1
+      per_atom: true
+
+validation_set: 0.1
+test_set: 0.1
+
+wandb:
+  project: flashmd-variants
+  name: flashmd-baseline
+  tags:
+    - al
+    - flashmd

examples/al/simulation-baseline/baseline.xml

@@ -0,0 +1,33 @@
+<simulation verbosity='low' threading='false'>
+   <total_steps>3200</total_steps>
+   <output prefix='md'>
+      <trajectory stride='1' filename='pos' format='ase'> positions </trajectory>
+      <trajectory stride='1' filename='vel' format='xyz'> velocities </trajectory>
+      <properties stride='1'> [ step, time{picosecond}, conserved, temperature{kelvin} ] </properties> 
+   </output>
+   <prng>
+      <seed>32123</seed>
+   </prng>
+   <ffdirect name='mlip'>
+      <pes>metatomic</pes>
+      <parameters>{model: ../models/mlip_pet-omatpes-v2.pt, template: ../data/equilibrated.xyz, device: cuda}</parameters>
+    </ffdirect>
+   <system>
+      <forces> 
+         <force forcefield='mlip'></force>
+      </forces>
+      <initialize nbeads='1'>
+         <file mode='ase'>../data/equilibrated.xyz</file>
+         <velocities mode='thermal' units='kelvin'>300</velocities>
+      </initialize>
+      <ensemble>
+         <temperature units='kelvin'>300</temperature>
+      </ensemble>
+      <motion mode='dynamics'>
+         <dynamics mode='nvt'>
+            <timestep units='femtosecond'>1</timestep>
+            <thermostat mode='svr'><tau units='femtosecond'>2</tau></thermostat>
+         </dynamics>
+      </motion>
+  </system>
+</simulation>

examples/al/simulation-baseline/run.sh

@@ -0,0 +1 @@
+pixi run i-pi baseline.xml

examples/al/simulation-flashmd-omatpes/run.py

@@ -0,0 +1,20 @@
+import torch
+from ipi.utils.scripting import InteractiveSimulation
+from flashmd import get_pretrained
+from flashmd.steppers import FlashMDStepper
+from flashmd.wrappers import wrap_nvt
+from flashmd.vv import flashmd_vv
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+with open("../input.xml", "r") as input_xml:
+  sim = InteractiveSimulation(input_xml)
+
+# replace the motion step with a FlashMD stepper
+_, flashmd_model_32 = get_pretrained("pet-omatpes", 32)
+stepper = FlashMDStepper(flashmd_model_32, device=device)
+step_fn = flashmd_vv(sim, stepper, device=device, dtype=torch.float32, rescale_energy=False)
+step_fn = wrap_nvt(sim, step_fn)
+sim.set_motion_step(step_fn)
+
+sim.run(100)

examples/al/simulation-flashmd-symplectic/run.py

@@ -0,0 +1,42 @@
+from typing import Callable
+import torch
+from metatomic.torch import load_atomistic_model
+from ipi.utils.scripting import InteractiveSimulation
+from flashmd.steppers import SymplecticStepper, FlashMDStepper
+from flashmd.vv import flashmd_vv
+from flashmd.wrappers import wrap_nvt
+from flashmd.fpi import anderson_solver
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+with open("../input.xml", "r") as input_xml:
+  sim = InteractiveSimulation(input_xml)
+
+# load FlashMD model for initial guess
+flashmd_model_32 = load_atomistic_model("../models/flashmd.pt")
+flashmd_model_32.to(device)
+initial_guess = FlashMDStepper(flashmd_model_32, device=device)
+
+# load FlashMD symplectic model for corrector
+flashmd_symplectic_model_32 = load_atomistic_model("../models/flashmd-symplectic.pt")
+flashmd_symplectic_model_32.to(device)
+
+# create a fixed-point solver and attach a logger to see the convergence behavior
+solver_kwargs = dict(m=0, max_iter=100, tol=1e-3, beta=0.5)
+def solver_with_log(
+    g: Callable[[torch.Tensor], torch.Tensor],
+    x0: torch.Tensor,
+) -> torch.Tensor:
+    x_star, norms = anderson_solver(g, x0, return_residual_norms=True, **solver_kwargs)  # type: ignore
+    print("l2 accuracies (converged in %d steps):" % len(norms))
+    for i, n in enumerate(norms):
+       print("iteration", i, "residual norm:", n)
+    return x_star
+
+# replace the motion step with a FlashMD stepper
+stepper = SymplecticStepper(initial_guess, flashmd_symplectic_model_32, solver_with_log)
+step_fn = flashmd_vv(sim, stepper, device=device, dtype=torch.float32, rescale_energy=False, random_rotation=False)
+step_fn = wrap_nvt(sim, step_fn)
+sim.set_motion_step(step_fn)
+
+sim.run(100)

examples/al/simulation-flashmd/run.py

@@ -0,0 +1,23 @@
+import torch
+from metatomic.torch import load_atomistic_model
+from ipi.utils.scripting import InteractiveSimulation
+from flashmd.steppers import FlashMDStepper
+from flashmd.vv import flashmd_vv
+from flashmd.wrappers.nvt import wrap_nvt
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+with open("../input.xml", "r") as input_xml:
+  sim = InteractiveSimulation(input_xml)
+
+# load FlashMD model
+flashmd_model_32 = load_atomistic_model("../models/flashmd.pt")
+flashmd_model_32.to(device)
+
+# replace the motion step with a FlashMD stepper
+stepper = FlashMDStepper(flashmd_model_32, device=device)
+step_fn = flashmd_vv(sim, stepper, device=device, dtype=torch.float32, rescale_energy=False)
+step_fn = wrap_nvt(sim, step_fn)
+sim.set_motion_step(step_fn)
+
+sim.run(100)

src/flashmd/ase/velocity_verlet.py

@@ -8,7 +8,7 @@
 from metatomic.torch.ase_calculator import _ase_to_torch_data
 from scipy.spatial.transform import Rotation
 
-from ..stepper import FlashMDStepper
+from ..steppers.flashmd import FlashMDStepper
 
 
 class VelocityVerlet(MolecularDynamics):