Add Al example for symplectic integrator

examples/al/.gitignore

@@ -0,0 +1,10 @@
+# ignore generated dataset and simulation files
+*.*xyz
+RESTART
+*#
+*.out
+
+# ignore training artifacts
+outputs/
+*.pt
+*.ckpt

examples/al/al.py

@@ -0,0 +1,130 @@
+"""
+This is a self-contained example to show how to use different integrators.
+
+First, we run a simple MD simulation for Al and log the trajectory. Then, we process
+the trajectory into two datasets: one for training a FlashMD model and one for training
+a symplectic FlashMD model. Then, we train both models and run dynamics with them with
+i-PI.
+
+NOTE: This example is designed to run quickly on a CPU. For a real experiment, make sure
+to run longer simulations and train for more epochs.
+"""
+
+# %%
+import shutil
+import subprocess
+from ase import Atoms
+import ase.io
+from ase.build import bulk
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary
+from ase.md.verlet import VelocityVerlet
+from ase.md.langevin import Langevin
+from ase import units
+from upet.calculator import UPETCalculator
+from upet import save_upet
+from tqdm import trange
+from metatomic.torch import load_atomistic_model
+from ipi.scripting import InteractiveSimulation
+from flashmd.ipi import get_nve_stepper
+
+# %%
+# Create a bulk Al system for demonstation.
+atoms = bulk("Al", "fcc", cubic=True) * (3, 3, 3)
+ase.io.write("al.xyz", atoms)
+len(atoms)
+
+# %%
+# Attach a UPET calculator
+calc = UPETCalculator(model="pet-mad-xs", version="1.5.0", device="cpu")
+save_upet(model="pet-mad", size="xs", version="1.5.0", output="mlip.pt")
+atoms.calc = calc
+atoms.get_potential_energy()
+
+# %%
+# Set up a simulation and equilibrate
+MaxwellBoltzmannDistribution(atoms, temperature_K=400)
+Stationary(atoms)
+gamma = 1 / (200 * units.fs)
+Langevin(atoms, 2 * units.fs, temperature_K=400, friction=gamma, fixcm=False).run(100)
+
+# %%
+# Run NVE MD with ASE for an Al system.
+mlip_integrator = VelocityVerlet(atoms, 2 * units.fs)
+structures = []
+for _ in trange(100):
+  mlip_integrator.run(1)
+  structures.append(atoms.copy())
+
+# %%
+# Write the trajectory to an easy-to-use format.
+ase.io.write("al.xyz", structures)
+
+# %%
+# Preprocess the trajectories to be readable for both versions of FlashMD. This code is
+# largely equal to the code in metatrain showing how to train the various models.
+
+structures: list[Atoms] = ase.io.read("al.xyz", index=":")  # type: ignore
+i = 0
+num_step_frames = 4
+num_decorrelation_frames = 10
+assert num_decorrelation_frames > 1
+flashmd_structures = []
+symplectic_flashmd_structures = []
+while i < len(structures) - num_step_frames + 1:
+  # Extract the current and future positions and momenta.
+  current_q = structures[i].get_positions()
+  current_p = structures[i].get_momenta()
+  future_q = structures[i + num_step_frames - 1].get_positions()
+  future_p = structures[i + num_step_frames - 1].get_momenta()
+
+  # For FlashMD, take a frame and frame + num_step_frames ahead.
+  flashmd_structure = structures[i].copy()
+  flashmd_structure.arrays["future_positions"] = future_q
+  flashmd_structure.arrays["future_momenta"] = future_p
+  flashmd_structures.append(flashmd_structure)
+
+  # For symplectic FlashMD, the input is a midpoint and the target is the delta between
+  # the start and the end configuration.
+  symplectic_flashmd_structure = structures[i].copy()
+  symplectic_flashmd_structure.set_positions((current_q + future_q) / 2)
+  symplectic_flashmd_structure.set_momenta((current_p + future_p) / 2)
+  symplectic_flashmd_structure.arrays["delta_positions"] = future_q - current_q
+  symplectic_flashmd_structure.arrays["delta_momenta"] = future_p - current_p
+  symplectic_flashmd_structures.append(symplectic_flashmd_structure)
+
+  i += num_decorrelation_frames
+print(f"{len(flashmd_structures)=}, {len(symplectic_flashmd_structures)=}")
+
+# %%
+# Write the processed frames to two dataset files.
+ase.io.write("start-to-end.xyz", flashmd_structures)
+ase.io.write("midpoint-to-delta.xyz", symplectic_flashmd_structures)
+
+# %%
+# Train models with the datasets.
+subprocess.run(["mtt", "train", "options-flashmd.yaml"], check=True)
+shutil.move("model.pt", "flashmd.pt")
+subprocess.run(["mtt", "train", "options-symplectic-flashmd.yaml"], check=True)
+shutil.move("model.pt", "symplectic-flashmd.pt")
+
+# %%
+# Load the input file template for i-PI. We replace the motion step later with an
+# appropriate FlashMD step function.
+with open("simulation-template.xml") as f:
+  input_template = f.read()
+
+# %%
+# Run NVE dynamics with i-PI and FlashMD
+flashmd = load_atomistic_model("flashmd.pt")
+simulation = InteractiveSimulation(input_template.replace("PREFIX", "flashmd"))
+step_fn = get_nve_stepper(simulation, flashmd, "cpu", rescale_energy=False)
+simulation.set_motion_step(step_fn)
+simulation.run(100)
+
+# %%
+# Run NVE dynamics with i-PI and symplectic FlashMD
+symplectic_flashmd = load_atomistic_model("symplectic-flashmd.pt")
+symplectic_simulation = InteractiveSimulation(input_template.replace("PREFIX", "symplectic-flashmd"))
+symplectic_step_fn = get_nve_stepper(symplectic_simulation, (flashmd, (symplectic_flashmd, {})), "cpu", rescale_energy=False)
+symplectic_simulation.set_motion_step(symplectic_step_fn)
+symplectic_simulation.run(100)

examples/al/options-flashmd.yaml

@@ -0,0 +1,45 @@
+seed: 42
+base_precision: 32
+
+architecture:
+  name: experimental.flashmd
+  training:
+    timestep: 8 # in this case 4 (time lag) * 2 fs (timestep of reference MD)
+    batch_size: 8  # to be increased in a production scenario
+    num_epochs: 25  # to be increased (at least 1000-10000) in a production scenario
+    log_interval: 1
+    learning_rate: 3e-4
+    loss:
+      positions:
+        type: mse
+        weight: 1.0
+        reduction: mean
+      momenta:
+        type: mse
+        weight: 1.0
+        reduction: mean
+
+training_set:
+  systems:
+    read_from: 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.0

examples/al/options-symplectic-flashmd.yaml

@@ -0,0 +1,48 @@
+seed: 42
+base_precision: 32
+
+architecture:
+  name: experimental.flashmd_symplectic
+  training:
+    timestep: 8  # in this case 4 (time lag) * 2 fs (timestep of reference MD)
+    batch_size: 8  # to be increased in a production scenario
+    num_epochs: 25  # 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: 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

examples/al/simulation-template.xml

@@ -0,0 +1,33 @@
+<simulation verbosity='low' threading='false'>
+   <total_steps>100</total_steps>
+   <output prefix='PREFIX'>
+      <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>
+   <!-- Declare a "dummy" forcefield. We're not going to use it because the FlashMD flavors replace the `motion_step`. -->
+   <ffdirect name='mlip'>
+      <pes>metatomic</pes>
+      <parameters>{model: ./mlip.pt, template: ./al.xyz, device: cpu}</parameters>
+    </ffdirect>
+   <system>
+      <!-- Use the dummy force field here. -->
+      <forces> 
+         <force forcefield='mlip'></force>
+      </forces>
+
+      <initialize nbeads='1'>
+         <file mode='ase'>./al.xyz</file>
+         <velocities mode='thermal' units='kelvin'>400</velocities>
+      </initialize>
+      <motion mode='dynamics'>
+         <dynamics mode='nve'>
+            <!-- This needs to match the FlashMD time step! -->
+            <timestep units='femtosecond'>8</timestep>
+         </dynamics>
+      </motion>
+  </system>
+</simulation>

pyproject.toml

@@ -92,6 +92,8 @@ filterwarnings = [
     "ignore:`torch.jit.load` is not supported:DeprecationWarning",
     "ignore:The 'nvalchemiops.neighborlist' module has been renamed to 'nvalchemiops.neighbors':DeprecationWarning",
     "ignore:From version 0.3.0 onwards, PyTorch is now an optional dependency:DeprecationWarning",
+    "ignore:currentThread\\(\\) is deprecated:DeprecationWarning",
+    "ignore::ResourceWarning",
     "ignore:`torch.jit.interface` is deprecated:DeprecationWarning",
     "ignore:`torch.jit.set_fusion_strategy` is deprecated:DeprecationWarning",
     "ignore:the '.*' output name is deprecated:UserWarning",

tests/test_example_al.py

@@ -0,0 +1,98 @@
+"""Test that the Al example runs end-to-end on CPU with a tiny model."""
+
+import os
+import tempfile
+from pathlib import Path
+
+import yaml
+
+
+EXAMPLE_DIR = Path(__file__).parent.parent / "examples" / "al"
+_MINIMAL_MODEL_HYPERS = {
+    "d_pet": 1,
+    "d_head": 1,
+    "d_node": 1,
+    "d_feedforward": 1,
+    "num_heads": 1,
+    "num_attention_layers": 1,
+    "num_gnn_layers": 1,
+}
+
+
+def _modify_al_py(code: str) -> str:
+    # Swap out the MLIP calculator for EMT (no GPU, no download required).
+    code = code.replace(
+        "from upet.calculator import UPETCalculator\nfrom upet import save_upet",
+        "from ase.calculators.emt import EMT",
+    )
+    code = code.replace(
+        'calc = UPETCalculator(model="pet-mad-xs", version="1.5.0", device="cpu")\n'
+        'save_upet(model="pet-mad", size="xs", version="1.5.0", output="mlip.pt")\n'
+        "atoms.calc = calc",
+        "atoms.calc = EMT()",
+    )
+
+    # Reduce the number of MD steps so the test finishes quickly.
+    code = code.replace(
+        "Langevin(atoms, 2 * units.fs, temperature_K=400, friction=gamma, fixcm=False).run(100)",
+        "Langevin(atoms, 2 * units.fs, temperature_K=400, friction=gamma, fixcm=False).run(5)",
+    )
+    code = code.replace("trange(100)", "trange(30)")
+    code = code.replace("num_decorrelation_frames = 10", "num_decorrelation_frames = 2")
+
+    # Reduce i-PI simulation steps.
+    code = code.replace("simulation.run(100)", "simulation.run(5)")
+    code = code.replace(
+        "symplectic_simulation.run(100)", "symplectic_simulation.run(5)"
+    )
+
+    return code
+
+
+def _modify_training_yaml(path: Path, architecture_name: str) -> str:
+    with open(path) as f:
+        hypers = yaml.safe_load(f)
+
+    hypers["architecture"]["model"] = _MINIMAL_MODEL_HYPERS.copy()
+    hypers["architecture"]["training"]["num_epochs"] = 2
+    hypers["architecture"]["training"]["batch_size"] = 2
+
+    return yaml.dump(hypers)
+
+
+def _modify_simulation_xml(xml: str) -> str:
+    # Replace the metatomic force field with a dummy PES so no model file is
+    # needed (FlashMD replaces the motion step entirely anyway).
+    xml = xml.replace("<pes>metatomic</pes>", "<pes>dummy</pes>")
+    xml = xml.replace(
+        "<parameters>{model: ./mlip.pt, template: ./al.xyz, device: cpu}</parameters>",
+        "<parameters>{}</parameters>",
+    )
+    xml = xml.replace("<total_steps>100</total_steps>", "<total_steps>5</total_steps>")
+    return xml
+
+
+def test_example_al():
+    code = _modify_al_py((EXAMPLE_DIR / "al.py").read_text())
+    flashmd_yaml = _modify_training_yaml(
+        EXAMPLE_DIR / "options-flashmd.yaml", "experimental.flashmd"
+    )
+    symplectic_yaml = _modify_training_yaml(
+        EXAMPLE_DIR / "options-symplectic-flashmd.yaml",
+        "experimental.flashmd_symplectic",
+    )
+    simulation_xml = _modify_simulation_xml(
+        (EXAMPLE_DIR / "simulation-template.xml").read_text()
+    )
+
+    original_dir = os.getcwd()
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        (Path(tmp_dir) / "options-flashmd.yaml").write_text(flashmd_yaml)
+        (Path(tmp_dir) / "options-symplectic-flashmd.yaml").write_text(symplectic_yaml)
+        (Path(tmp_dir) / "simulation-template.xml").write_text(simulation_xml)
+
+        try:
+            os.chdir(tmp_dir)
+            exec(code, {})  # noqa: S102
+        finally:
+            os.chdir(original_dir)

tox.ini

@@ -45,6 +45,9 @@ description = Run package tests with pytest
 passenv = *
 deps =
     pytest
+    ipi
+    tqdm
+    pyyaml
 changedir = tests
 commands =
     pytest \