flashcart.calculators

Use trained potentials in atomistic simulations. The ASE calculator constructs neighbor graphs from ase.Atoms objects and returns predictions using ASE’s calculator interface. The LAMMPS interface receives neighbor and atom data from ML-IAP and exchanges features between owned and ghost atoms when needed.

Start with ASE calculator for Python simulations or LAMMPS interface for export, build requirements, and simulation input. The reference below describes the calculator options and the lower-level ML-IAP graph and export helpers.

flashcart.calculators.ase

Evaluate trained potentials with ASE.

class flashcart.calculators.ase.FlashCartCalculator(checkpoint=None, model=None, device=None, skin=0.0, compile_mode=None, compile_fullgraph=True, pad_n_atoms=0, pad_n_edges=0, pad_extra_atoms=None, pad_atom_multiple=1, pad_edge_headroom=1.1, pad_edge_multiple=128, add_atomic_offsets=False, **kwargs)[source]

Bases: Calculator

Evaluate energies, forces, and stress with a FlashCart model in ASE.

A positive skin extends the neighbor-list radius beyond the model cutoff, allowing graph connectivity to be reused between evaluations. The cutoff envelope removes contributions from edges outside the model cutoff.

Setting compile_mode enables compiled prediction and graph padding. Padding keeps tensor shapes unchanged while the graph fits within the allocated capacities, reducing recompilation as the number of edges varies. The capacities grow when required. Positive pad_n_atoms or pad_n_edges also enables padding without compilation.

Parameters:

  • checkpoint (Union[str, Path], optional) – Inference checkpoint directory. Required when model is not provided. Otherwise ignored.

  • model (FlashCartPotential, optional) – Model to use for prediction. If set, it overrides checkpoint.

  • device (Union[str, torch.device], optional) – Device used for prediction. If None, retain the device of a provided model. Checkpoints are loaded on CPU unless a device is specified.

  • skin (float, optional) – Additional neighbor-list radius. The list is constructed at r_max + skin. For a fixed cell, it can be reused while every atom has moved less than skin / 2 from its position at construction, provided no other atomic properties have changed. Cell changes are checked separately. Default: 0.0, which rebuilds the list for each calculation.

  • compile_mode (str, optional) – torch.compile mode used for prediction. Setting a mode also enables graph padding. Default: None, which evaluates the model without compilation.

  • compile_fullgraph (bool, optional) – Whether compilation requires a single graph without graph breaks. Used when compile_mode is set. Default: True.

  • pad_n_atoms (int, optional) – Initial capacity for atoms before adding reserved padding atoms. The capacity grows when required. Default: 0.

  • pad_n_edges (int, optional) – Initial capacity for edges before rounding to pad_edge_multiple. The capacity grows when required. Default: 0.

  • pad_extra_atoms (int, optional) – Number of additional padding atoms reserved as endpoints of padding edges, with a minimum of two. Default: None, which determines the number from the atom capacity and pad_edge_headroom.

  • pad_atom_multiple (int, optional) – Round the total number of atoms, including reserved padding atoms, up to this multiple. Default: 1.

  • pad_edge_headroom (float, optional) – Multiplicative factor applied to the required edge count when increasing the edge capacity. Default: 1.1.

  • pad_edge_multiple (int, optional) – Round the total number of edges, including padding edges, up to this multiple. Default: 128.

  • add_atomic_offsets (bool, optional) – Restore the per-element energy shifts subtracted from the reference energies during training. These shifts affect the returned energy but not forces or stress. Default: False.

  • **kwargs – Additional keyword arguments passed to ase.calculators.Calculator.

implemented_properties = ['energy', 'forces', 'stress']

Properties calculator can handle (energy, forces, …)

classmethod from_checkpoint(checkpoint, device=None, skin=0.0, **kwargs)[source]

Build a calculator from a saved inference checkpoint.

Parameters:

  • checkpoint (Union[str, Path]) – Inference checkpoint directory.

  • device (Union[str, torch.device], optional) – Device used for prediction. Default: None, which loads the model on CPU.

  • skin (float, optional) – Additional neighbor-list radius, as described in FlashCartCalculator. Default: 0.0.

  • **kwargs – Additional arguments passed to FlashCartCalculator.

Returns:

Calculator initialized with the saved model.

Return type:

FlashCartCalculator

calculate(atoms=None, properties=None, system_changes=['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'])[source]

Evaluate the requested properties and store them in self.results.

Graph connectivity is reused when the neighbor-list reuse conditions are satisfied. Otherwise, the graph is rebuilt. The results contain only the requested properties. Energy is stored as a scalar, forces as an array of shape (n_atoms, 3), and stress as a six-component array in ASE Voigt order: xx, yy, zz, yz, xz, xy.

Parameters:

  • atoms (ase.Atoms, optional) – Atomic configuration to evaluate. If None, use the configuration already stored by the calculator.

  • properties (list[str], optional) – Properties to evaluate, chosen from "energy", "forces", and "stress". Default: ["energy", "forces"].

  • system_changes (list[str], optional) – Changes reported by ASE since the preceding evaluation. Changes other than positions or cell require the neighbor list to be rebuilt. Default: all_changes.

flashcart.calculators.lammps_mliap

Build ML-IAP graphs, export models, and evaluate potentials in LAMMPS.

flashcart.calculators.lammps_mliap.graph_from_mliap(data, element_types, r_max, device=None, dtype=None, cell=None, pbc=None, validate=False)[source]

Build a FlashCart graph from a LAMMPS ML-IAP data object.

LAMMPS supplies edge vectors through rij, including periodic translations. The graph stores these vectors directly and sets use_shifts=False, so positions and cell data are not used to reconstruct them. Nodes include both owned and ghost atoms, while batch and n_atoms describe only the nlocal owned atoms. Site energies are therefore summed over owned atoms. lammps_exchange=True selects this layout in the model. When ghost features must be exchanged between layers, evaluation and differentiation require an active lammps_data_slot() context.

Parameters:

  • data – LAMMPS ML-IAP data object (ntotal, nlocal, npairs, pair_i, pair_j, rij, elems).

  • element_types (list[str]) – Chemical symbols in ML-IAP element order, indexed by data.elems.

  • r_max (float) – Model cutoff radius, stored as graph metadata. The neighbor list provided by LAMMPS is not rebuilt or filtered.

  • device (torch.device, optional) – Requested tensor device. CUDA input arrays require a CUDA device. Default: inferred from the input arrays.

  • dtype (torch.dtype, optional) – Floating-point dtype. Default: torch.get_default_dtype().

  • cell (np.ndarray, optional) – Cell matrix stored as graph metadata. Default: a diagonal matrix with entries of 1000.

  • pbc (np.ndarray, optional) – Periodic boundary flags stored as graph metadata. Default: all False.

  • validate (bool, optional) – Check that atom and element-type indices lie within their allowed ranges. Default: False.

Returns:

Graph containing owned and ghost atoms and the edge vectors provided by

LAMMPS. n_atoms and batch cover only owned atoms.

Return type:

Data

flashcart.calculators.lammps_mliap.graph_dict_from_mliap(data, element_types, device=None, dtype=None, validate=False)[source]

Build a graph dictionary for energy and edge-force evaluation.

The dictionary uses the same atom and edge representation as graph_from_mliap but omits positions, cell data, and periodic shifts. LAMMPS supplies the edge vectors directly, and edge forces are obtained by differentiating the energy with respect to these vectors. The dictionary contains only tensors and flags. With lammps_exchange=True, the model infers owned and total atom counts from the shapes of batch and atom_types for dynamic compilation. When ghost features must be exchanged between layers, evaluation and differentiation require an active lammps_data_slot() context.

Parameters:

  • data – LAMMPS ML-IAP data object, as in graph_from_mliap.

  • element_types (list[str]) – Chemical symbols in ML-IAP element order, indexed by data.elems.

  • device (torch.device, optional) – Requested tensor device. CUDA input arrays require a CUDA device. Default: inferred from the input arrays.

  • dtype (torch.dtype, optional) – Floating-point dtype. Default: torch.get_default_dtype().

  • validate (bool, optional) – Check that atom and element-type indices lie within their allowed ranges. Default: False.

Returns:

Graph dictionary containing atom types, connectivity, and

edge vectors, without positions, cell data, or periodic shifts.

Return type:

dict[str, object]

class flashcart.calculators.lammps_mliap.FlashCartLAMMPSUnified(*args, **kwargs)[source]

Bases: MLIAPUnified

Evaluate a FlashCart potential through the LAMMPS ML-IAP interface.

compute_forces evaluates site energies for owned atoms and obtains edge forces by differentiating the energy with respect to the edge vectors provided by LAMMPS. Models with multiple interaction layers exchange ghost-atom features between layers through ML-IAP, with reverse communication during differentiation. Every rank participates in feature exchange, including ranks without ghost atoms, because other ranks may need features from its atoms. Ranks without owned atoms or neighbor pairs are not supported.

Construction moves the provided model to CPU, sets evaluation mode, and disables gradients for its parameters. During evaluation, the model moves to the device selected for the ML-IAP arrays.

Setting compile_mode enables compilation with torch.compile and dynamic shapes. Compilation starts on the first model evaluation, after the model has moved to the device used by the ML-IAP arrays. Dynamic shapes allow graph reuse as atom and neighbor-pair counts change, although some changes can still require recompilation. The compiled function is not saved with the interface. Each LAMMPS process creates its own compiled function. The PyTorch Inductor cache (TORCHINDUCTOR_CACHE_DIR) can shorten later compilations.

Parameters:

  • model (FlashCartPotential) – Potential to evaluate.

  • add_atomic_offsets (bool, optional) – Restore the per-element energy shifts subtracted from the reference energies during training. These shifts affect the reported energies but not the forces. Default: False.

  • compile_mode (str, optional) – torch.compile mode: "default" or "max-autotune-no-cudagraphs". Modes using CUDA graphs are not supported by this interface. Default: None, which evaluates the model without compilation.

  • compile_fullgraph (bool, optional) – Whether compilation requires a single graph without graph breaks. Used when compile_mode is set. Default: True.

Raises:

ValueError – compile_mode is not supported.

compute_forces(data)[source]

Evaluate forces and write the requested energies to the ML-IAP data.

Edge forces are passed to data.update_pair_forces when available. Otherwise, their contributions are accumulated for owned and ghost atoms and written to data.f. LAMMPS returns ghost-atom contributions to their owners through reverse communication.

When data.eflag is true or absent, the method writes the summed energy of the owned atoms and their site energies when an eatoms buffer is available. Each rank must own at least one atom and have at least one neighbor pair. With several interaction layers, the ghost-feature exchange is matched between neighboring ranks, and a rank that cannot evaluate the model would leave them waiting. Unsupported configurations terminate an attached LAMMPS process with a nonzero status. Standalone calls raise RuntimeError.

Atom and element-type indices are checked for the first evaluated graph. Setting FLASHCART_LAMMPS_VALIDATE=1 or FLASHCART_LAMMPS_DEBUG=1 enables these checks for every evaluated graph.

Parameters:

data – LAMMPS ML-IAP data object for the current step.

flashcart.calculators.lammps_mliap.build_lammps_unified(checkpoint, compile_mode=None)[source]

Load an inference checkpoint on CPU and wrap it for LAMMPS.

Setting compile_mode records how to evaluate the model in LAMMPS. It does not compile the model here.

Parameters:

  • checkpoint (Union[str, Path]) – Inference checkpoint directory.

  • compile_mode (str, optional) – torch.compile mode, as described in FlashCartLAMMPSUnified. Default: None.

Returns:

Interface initialized with the saved model.

Return type:

FlashCartLAMMPSUnified

flashcart.calculators.lammps_mliap.save_lammps_unified(unified, path)[source]

Serialize a FlashCart ML-IAP interface for loading by LAMMPS.

The file contains the interface and its model, saved with torch.save, for use with the LAMMPS mliap unified command. Parent directories are created if needed. Compilation settings are saved, but cached compiled functions are omitted. They are recreated on evaluation after loading when compilation is enabled.

Parameters:
Returns:

Path to the saved interface.

Return type:

Path