ASE calculator

FlashCartCalculator evaluates energies, forces, and stress with ASE. Provide a trained checkpoint and a structure containing only elements supported by that model. Reference labels are not required for prediction. The calculator does not convert units. The examples below use models trained with energies in eV and positions in Å.

from ase.io import read
from flashcart.calculators import FlashCartCalculator

atoms = read("water_256.extxyz")
atoms.calc = FlashCartCalculator(checkpoint="models/my_run/best", skin=1.0)
energy = atoms.get_potential_energy()
forces = atoms.get_forces()

Calculator options

Checkpoints are loaded on CPU by default. Set device="cuda" to use a GPU. The calculator accepts PyTorch device names. In training configurations, device: gpu selects GPU execution.

• add_atomic_offsets=True restores the per-element energy shifts subtracted from the reference energies during training. It changes the total energy but not forces or stress. The default is False.

• skin adds a buffer to the neighbor-list radius without changing the model cutoff. For a fixed cell, the list can be reused until an atom moves by skin / 2 from its position when the list was built. Changes to the cell or other atomic properties can also require rebuilding.

• compile_mode="reduce-overhead" enables torch.compile and graph padding. Padding keeps input shapes unchanged while the graph fits within the allocated capacities. Larger graphs increase those capacities and can require recompilation. Compilation adds an initial cost but can reduce the time required for repeated evaluations.

Molecular dynamics

The repository includes examples/md/ase_md.py, a Langevin molecular dynamics script. From the repository root, provide a checkpoint directory and an initial structure:

python examples/md/ase_md.py models/my_run/best structure.extxyz \
    --temperature 300 --timestep 0.5 --steps 2000

This runs 1 ps with a time step of 0.5 fs and a target temperature of 300 K. The script uses the calculator’s CPU default. Every 100 steps, it prints potential, kinetic, and total energies per atom and writes a frame to md.traj.

Initial velocities are sampled at the requested temperature. The script does not relax the structure before starting dynamics.

To use a GPU or compilation, adjust the calculator construction in the script using the options above.

"""Run Langevin molecular dynamics with ASE and a trained FlashCart model.

    python ase_md.py models/i3c3f64l2_s0/best water/water_256.extxyz --steps 2000

Prints the potential, kinetic, and total energy every 100 steps and writes
every 100th frame to md.traj (view it with ``ase gui md.traj``). The water
box comes from water/build_box.py.
"""

import argparse

from ase import units
from ase.io import read
from ase.io.trajectory import Trajectory
from ase.md.langevin import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution

from flashcart.calculators import FlashCartCalculator


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("checkpoint", help="checkpoint directory, e.g. <output_path>/best")
    parser.add_argument("structure", help="any file ASE can read")
    parser.add_argument("--steps", type=int, default=2000)
    parser.add_argument("--temperature", type=float, default=300.0, help="K")
    parser.add_argument("--timestep", type=float, default=0.5, help="fs")
    args = parser.parse_args()

    atoms = read(args.structure)
    atoms.calc = FlashCartCalculator(checkpoint=args.checkpoint, skin=1.0)
    MaxwellBoltzmannDistribution(atoms, temperature_K=args.temperature)
    dyn = Langevin(
        atoms, args.timestep * units.fs, temperature_K=args.temperature, friction=0.01 / units.fs, fixcm=False
    )
    traj = Trajectory("md.traj", "w", atoms)

    def report():
        epot, ekin = atoms.get_potential_energy(), atoms.get_kinetic_energy()
        n = len(atoms)
        print(f"step {dyn.nsteps:6d}  Epot {epot / n:9.4f}  Ekin {ekin / n:7.4f}  Etot {(epot + ekin) / n:9.4f} eV/atom")

    dyn.attach(report, interval=100)
    dyn.attach(traj.write, interval=100)
    dyn.run(args.steps)


if __name__ == "__main__":
    main()