Checkpoint and restart an MD simulation¶

restart_md continues the integrator from a saved snapshot and writes a new trajectory segment. Use it when stopping a long calculation, recovering from an interrupted run, or splitting a calculation across jobs. A geometry file retains atoms and momenta but does not contain all the random-generator, thermostat/barostat, or bias state required for faithful continuation.

This tutorial uses only the md extra and ASE's small EMT calculator. It downloads no ML model and needs no PLUMED installation. For AIMNet runs, construct the restart interface with the same model files, calculator options, and device setup you used originally.

1. Save a checkpoint and continue in another process¶

From a current source checkout, install the small simulation extra:

python -m pip install -e '.[md]'

Save this as first_segment.py and run python first_segment.py:

from ase.build import bulk
from ase.calculators.emt import EMT
from matmalib import mlmd

simulation = mlmd(calculator=EMT())
simulation.run_md(
    bulk("Cu", cubic=True),
    mlmd.MDConfig(
        steps=12, timestep_fs=0.2, temperature_K=100,
        ensemble="NVT", thermostat="langevin", seed=31,
        trajectory="restart_run/part1.xyz", trajectory_interval=2,
        log_file="restart_run/part1.log", log_interval=2,
        checkpoint_interval=5,
        checkpoint_file="restart_run/part1.checkpoint.npz",
    ),
)
print("Last step:", simulation.data.last_step)
print("Snapshot step:", simulation.data.checkpoint_step)
print("Checkpoint:", simulation.data.checkpoint_file)

Snapshots are written after steps 5 and 10, then at the final step 12. The same checkpoint path is atomically replaced each time; it is not an archive of all earlier checkpoints. Expect both printed step values to be 12.

Save the next script as second_segment.py and run it in a fresh Python process from the same working directory:

from ase.calculators.emt import EMT
from matmalib import mlmd

simulation = mlmd(calculator=EMT())
trajectory = simulation.restart_md(
    "restart_run/part1.checkpoint.npz",
    steps=8,
    trajectory="restart_run/part2.xyz",
    log_file="restart_run/part2.log",
)
report = simulation.data
print("Additional steps:", report.steps)
print("Step range:", report.start_step, report.last_step)
print("Absolute final time (fs):", report.time_fs)
print("Segment duration (fs):", report.simulation_time_fs)
print("New trajectory:", trajectory)
print("New checkpoint:", report.checkpoint_file)
print("Restarted from:", report.restarted_from)

Expect 8 additional steps, step range 12–20, absolute time 4 fs, and segment duration 1.6 fs. Restart keeps the saved timestep, ensemble, temperature, coupling parameters, and output/checkpoint intervals. This API does not accept a replacement MDConfig: a changed protocol is a new run, not an exact continuation.

The new segment includes the checkpoint geometry as its first frame, with step 12 and time 2.4 fs. The original segment also ends at step 12. When combining segments for analysis, discard one copy of that shared boundary frame. Do not blindly generate a new time axis starting at zero for each segment.

2. Understand the saved state¶

Checkpoints use a ZIP-based .npz container with JSON metadata and NumPy arrays, without pickle. They include:

Saved information	Why it matters
Atomic numbers, coordinates, cell, and PBC	The current molecular geometry and periodic box.
Masses, momenta, and velocities	Preserve mass scaling and the instantaneous dynamical state.
Charge, multiplicity, and constraints	Preserve the system specification, including initialized SHAKE target lengths.
Configuration and step/time counters	Continue the same protocol and label output continuously.
NumPy random-generator state	Continue Langevin random forces without redrawing the velocity distribution.
ASE integrator variables	Restore extended thermostat/barostat coordinates and momenta.
Generated PLUMED input and bias snapshots, when applicable	Reconstruct supported metadynamics history.

Model weights and the calculator object are not embedded. Recreate the same calculator and use the same inference weights and settings. Custom calculators with their own mutable state beyond the saved atomic/integrator state must manage that state separately. The loader requires exactly the ASE version recorded in the checkpoint and rejects incompatible integrator fields or inconsistent arrays. Preserve your environment and model versions alongside a research run; arbitrary cross-version or cross-hardware bitwise identity is not promised.

You can inspect a checkpoint using standard libraries and NumPy without importing the private checkpoint implementation. Run this after the first segment:

import json
from zipfile import ZipFile
import numpy as np

path = "restart_run/part1.checkpoint.npz"
with ZipFile(path) as archive:
    metadata = json.loads(archive.read("metadata.json"))
print("Format/version:", metadata["format"], metadata["version"])
print("ASE version:", metadata["ase_version"])
print("Integrator:", metadata["integrator"])
print("Saved step/time:", metadata["step"], metadata["time_fs"])
with np.load(path, allow_pickle=False) as arrays:
    print("Positions:", arrays["positions"].shape)
    print("Velocities (angstrom/fs):", arrays["velocities_angstrom_fs"].shape)

Treat the snapshot as a restart artifact rather than editing its JSON/arrays. Restart validates and restores constraints and integrator state together; manually changing one field can make the stored state inconsistent.

3. Configure the schedule and file names¶

from dataclasses import replace
from matmalib import mlmd

config = mlmd.MDConfig(
    steps=2300,
    trajectory="production.xyz",
    checkpoint_interval=1000,
)
# An ordinary new run begins at step zero.
scheduled = list(range(config.checkpoint_interval, config.steps + 1,
                       config.checkpoint_interval))
if not scheduled or scheduled[-1] != config.steps:
    scheduled.append(config.steps)
print("Planned snapshot steps:", scheduled)
print("Checkpoint selection:", config.checkpoint_file)
print("Disable checkpoints:", replace(config, checkpoint_file=None).checkpoint_file)

For a fresh 2,300-step run this plans snapshots at 1,000, 2,000, and 2,300. All MD and equilibration ensembles use the same default schedule of every 1,000 steps and at completion, including unbiased MD and generated or custom PLUMED input. A successful zero-step run also writes a final snapshot. No initial snapshot is written for a nonzero-length run, so a failure before its first scheduled save may leave no checkpoint.

Option	Behavior
`checkpoint_file="auto"`	Derive the path from the trajectory basename, replacing its suffix with `.checkpoint.npz`.
Equilibration without a trajectory	Derive the path from its log basename, or use `equilibration.checkpoint.npz` if no log is configured.
`checkpoint_file="my_snapshot.npz"`	Use an explicit checkpoint path.
`checkpoint_file=None`	Disable snapshot writes.
`checkpoint_interval=N`	Save at absolute step multiples of positive integer `N` and at completion.

On restart, the checkpoint interval is inherited; for a snapshot at step 12 with interval 5, the next periodic save is step 15. The restart method's checkpoint_file argument can choose a new destination or disable new snapshots. Its default "auto" uses the new trajectory's basename.

If trajectory is omitted from restart_md, the filename includes the original output basename and saved step, such as part1.restart-12.xyz, beside the checkpoint. log_file=None is the restart default: it disables a new text log rather than appending to the original log. Reports and the new trajectory are still produced.

Restart always protects existing outputs and does not expose an overwrite flag. Use new trajectory/log/checkpoint paths for each continuation. It also rejects a checkpoint destination equal to its input snapshot. The original snapshot remains unchanged, allowing independent branches with different new output names.

4. Know what is resumed and what is skipped¶

Restart restores velocities instead of thermalizing them, restores the random-generator state, and does not apply hydrogen/heavy-atom mass changes again. The saved configuration removes preparation phases, so minimization and equilibration are not rerun. This is true even when the original production run embedded those phases.

run_md("previous.xyz", config) is useful when deliberately starting a new protocol from its last geometry, but it creates new integrator/RNG state and new step/time labels. Extended XYZ preserves momenta, allowing initialize_velocities=False, but that alone is not a thermostat/barostat or biased-dynamics restart. equilibrate also writes checkpoints; its snapshot can be continued using restart_md with the saved ensemble/settings and a new trajectory destination.

If integration raises, output already flushed to the trajectory/log remains available. The most recent successfully written checkpoint is the restart point, which may be older than the last output frame. A failed snapshot write leaves the preceding snapshot intact. The API does not guarantee a new checkpoint at the exact failing step, and an abruptly terminated process may not update in-memory diagnostics.

5. Restart metadynamics with the correct history¶

Automatic biased restart supports input generated from MetadynamicsConfig fields, where neither plumed_file nor input_lines is supplied. The checkpoint bundles flushed HILLS and COLVAR history, the resolved PLUMED input, and cached bias results. Restoring from those snapshots does not depend on the original HILLS/COLVAR files still being present.

For example, after a generated metadynamics run:

from matmalib import mlmd

# Use the same AIMNet inference models and calculator settings as the first run.
simulation = mlmd(device="cpu")
simulation.restart_md(
    "biased_run/production.checkpoint.npz",
    steps=10000,
    trajectory="biased_run/production_part2.xyz",
    log_file="biased_run/production_part2.log",
)

The example assumes a previously created checkpoint; it is not an independent simulation. Restart creates biased_run/production_part2.plumed/ containing restored HILLS, COLVAR, and resolved plumed.dat. That directory must not already exist and its path must contain no whitespace. The step count and bias deposition schedule continue; growing histories are copied in chunks instead of loaded wholesale into memory. Small numerical differences from uninterrupted biased dynamics can arise from the finite precision of PLUMED history files.

Custom plumed_file or input_lines runs still write molecular checkpoints on the normal schedule, even when their input creates no HILLS or COLVAR. However, restart_md rejects their automatic continuation because arbitrary PLUMED actions may require additional external restart state. Restoring their molecular arrays alone would not reconstruct that bias. Use a workflow that explicitly handles all state required by the custom actions rather than treating such a checkpoint as an automatically portable biased restart.

Return to running simulations for ensemble selection, output diagnostics, and velocity initialization, or follow the metadynamics tutorial to create a generated bias configuration that supports this restart workflow.