flashcart.model

Construct, save, and load interatomic potentials, and evaluate their properties. FlashCartPotential implements the energy model. AtomisticModel provides the prediction and checkpoint interfaces shared by atomistic models.

Call forward() for site energies or predict() for total energies and requested derivatives. Computing energy derivatives requires PyTorch gradient tracking, even during model evaluation. When training on these derivatives, create_graph=True keeps them differentiable with respect to model parameters. Direct model predictions omit the stored per-element energy shifts. Calculators can restore these shifts when reporting energies.

For evaluation from ASE structures, start with flashcart.calculators and ASE calculator. The interfaces below are useful when constructing models or supplying atomic graphs directly.

For LAMMPS, use the interface described in LAMMPS interface. When supplying LAMMPS graph dictionaries directly, set lammps_exchange=True, provide edge vectors, and place owned atoms before ghost atoms. batch covers only owned atoms, while atom_types covers both. The model infers the counts from these tensor shapes and returns site energies only for owned atoms. With more than one interaction layer, keep evaluation and differentiation inside lammps_data_slot(), including on ranks without ghost atoms. The LAMMPS calculator manages this context automatically.

flashcart.model.flashcart

Construct a FlashCart potential from equivariant interaction and product layers.

FlashCartPotential inherits its prediction and checkpoint interfaces from AtomisticModel. Use predict() to evaluate a graph, from_checkpoint() to load a saved potential, and save_inference_checkpoint() to save its configuration and weights. The separate load_weights() and save_weights() methods operate on weights alone.

class flashcart.model.flashcart.FlashCartPotential(elements, r_max, r_min=0.0, n_hidden_feats=32, l_max_hidden_feats=1, l_max_edge_attrs=3, n_radial=8, n_interactions=2, correlation=3, path_reduced_product=True, element_agnostic_interaction=True, element_agnostic_product=True, n_element_interaction_features=16, nonlinearity=False, gate_activation='silu', gate_bias=True, layer_norm=False, hidden_radial=None, hidden_readout=None, radial_mapping='affine', cutoff_poly_order=6, avg_neighbors=None, atomic_shifts=None, fit_atomic_shifts=True, recompute_radial=False, use_triton=False)

Bases: AtomisticModel

Equivariant potential that predicts site energies from atomic graphs.

Each interaction layer combines neighboring node features with the edge directions through a tensor-product convolution. Optional gated nonlinearities and repeated node-local products transform the aggregated features. The scalar channels from every layer contribute to the energy readout, followed by per-element scaling.

The model predicts energies relative to the stored atomic shifts. Training loaders subtract these shifts from the reference energies and calculators can restore them when reporting energies.

Parameters:

• elements (list[str]) – Chemical symbols in the order used by the integer atom-type indices.

• r_max (float) – Neighbor cutoff radius.

• r_min (float, optional) – Lower bound of the radial mapping interval. Default: 0.0.

• n_hidden_feats (int, optional) – Channels per tensor rank in the hidden node features. Default: 32.

• l_max_hidden_feats (int, optional) – Maximum tensor rank of the hidden node features carried between layers. Default: 1.

• l_max_edge_attrs (int, optional) – Maximum tensor rank of the irreducible Cartesian edge attributes (also the interaction output tensor rank when product layers are used). Default: 3.

• n_radial (int, optional) – Number of Chebyshev radial basis functions. Default: 8.

• n_interactions (int, optional) – Number of message-passing layers. Default: 2.

• correlation (int, optional) – Maximum polynomial degree of the repeated products in the aggregated node features. A value of 1 disables product layers. Default: 3.

• path_reduced_product (bool, optional) – Path-reduced product coupling with learnable per-path weights (see ProductLayer). Default: True.

• element_agnostic_interaction (bool, optional) – If False, interaction layers condition on atom types. Default: True.

• element_agnostic_product (bool, optional) – If False, product layers use element-dependent weights. Default: True.

• n_element_interaction_features (int, optional) – Element-embedding width for element-conditioned interactions. Default: 16.

• nonlinearity (bool, optional) – Gated nonlinearity after each interaction. Default: False.

• gate_activation (str, optional) – Activation name for the gates. Default: “silu”.

• gate_bias (bool, optional) – Bias in the gate-producing linear. Default: True.

• layer_norm (bool, optional) – RMS layer norm inside interaction and product layers. Default: False.

• hidden_radial (list[int], optional) – Radial-MLP hidden widths. Default: [64, 64].

• hidden_readout (list[int], optional) – Readout-MLP hidden widths. Default: [64, 64].

• radial_mapping (str, optional) – Distance mapping of the Chebyshev basis, “cosine” or “affine”. Default: “affine”.

• cutoff_poly_order (int, optional) – Order of the polynomial cutoff envelope. Default: 6.

• avg_neighbors (float, optional) – Average neighbor count used for normalization. Default: None, estimated during pre_fit.

• atomic_shifts (list[float], optional) – Initial per-element energy shifts, in elements order. Default: None, using a zero prior when fitting shifts. Required when fit_atomic_shifts=False.

• fit_atomic_shifts (bool, optional) – Fit the shifts from the training set during pre_fit. Default: True.

• recompute_radial (bool, optional) – Recompute the final radial linear output during differentiation instead of retaining that tensor. Hidden activations remain stored. This option is omitted from saved model configurations. Default: False.

• use_triton (bool, optional) – Use the Triton kernels. Default: False.

forward(graph)

Evaluate site energies from a graph dictionary.

When edge vectors are absent, they are constructed from positions, cell, and periodic shifts. For LAMMPS graphs, the supplied edge vectors are used directly, and node features are exchanged before each interaction layer after the first. Ranks without ghost atoms also take part, since their atoms may be ghosts on neighboring ranks.

For LAMMPS graphs with more than one interaction layer, evaluation and differentiation must run inside lammps_data_slot(). The LAMMPS calculator manages this context automatically.

Parameters:

graph (dict) – Graph tensors: edge_index, batch, atom_types, and either positions/cell/shifts (plus optional use_shifts) or vectors. LAMMPS graphs set lammps_exchange=True. Their batch covers owned atoms and atom_types covers owned atoms followed by ghosts. The model infers nlocal and ntotal from these tensor shapes.

Returns:

Site energies with shape (n_atoms,). For LAMMPS graphs,

only the nlocal owned atoms are included.

Return type:

torch.Tensor

broadcast_pre_fit_state(fabric)

Broadcast the fitted shifts, scales, and neighbor normalization from rank 0.

For a single process, the method returns without communication.

Parameters:

fabric (Fabric) – Launched Fabric handle providing the broadcast.

pre_fit(train_loader)

Initialize atomic energy shifts, energy scales, and neighbor normalization.

Energy scales are set from the root-mean-square reference force component. The same scale is assigned to every element. The average neighbor count is estimated only when it was not supplied to the constructor.

Parameters:

train_loader (DataLoader) – Training data the statistics are read from.

to_model_config()

Return the configuration needed to reconstruct the model architecture.

The configuration excludes recompute_radial, which controls memory use during a run. Model weights are saved separately.

Returns:

Architecture and normalization settings used to reconstruct

the model.

Return type:

dict[str, Any]

classmethod from_config(cfg)

Construct a potential from the supplied settings.

This method does not load the packaged training defaults.

Parameters:

cfg (dict) – Constructor settings. elements and r_max are required. Omitted settings use the Python constructor defaults.

Returns:

Newly constructed potential.

Return type:

FlashCartPotential

flashcart.model.atomistic

Predict atomic properties and save or load model checkpoints.

flashcart.model.atomistic.save_inference_metadata(ckpt_dir, model_config, meta=None)

Write the model configuration and optional metadata to a checkpoint directory.

Parameters:

• ckpt_dir (Path) – Checkpoint directory.

• model_config (dict) – Constructor config as returned by to_model_config.

• meta (dict, optional) – Additional metadata written to meta.yaml. Default: None.

flashcart.model.atomistic.load_model_config(ckpt_dir)

Read model_config.yaml from a checkpoint directory (raises if absent).

Parameters:

ckpt_dir (Path) – Checkpoint directory.

Returns:

Model constructor settings read from model_config.yaml.

Return type:

dict[str, Any]

class flashcart.model.atomistic.AtomisticModel

Bases: Module

Base class for potentials that map atomic graphs to site energies.

Subclasses implement forward and the methods from_config and to_model_config. This class sums site energies over each structure and obtains forces, stress, and edge forces by automatic differentiation. It also provides optional compilation, saving and loading checkpoints, and parameter selection for optimizers.

forward(graph)

Return per-atom energies for a graph given as a tensor dict.

Parameters:

graph (dict) – Graph tensors (see FlashCartPotential.forward).

Returns:

Site energies with shape (n_atoms,).

Return type:

torch.Tensor

predict(graph, compute_forces=False, compute_stress=False, compute_edge_forces=False, create_graph=False, use_compile=False, compile_mode='reduce-overhead', fullgraph=True, dynamic=False)

Predict energies and requested derivatives for a batched graph.

Forces are -dE/dpositions. Stress comes from a symmetrized virtual cell strain and is returned in Voigt-6 form together with virials = -stress * volume. Edge forces are dE/dvectors for graphs carrying precomputed edge vectors (LAMMPS). When derivatives are requested without create_graph, the returned energies are detached.

Gradient recording must remain enabled when requesting energy derivatives. Calling eval() is compatible with derivative prediction, but torch.no_grad() and torch.inference_mode() are not.

Position forces and stress require energies computed from positions and cell data. For graphs with precomputed vectors, request edge forces instead: the supplied vectors are not reconstructed from positions or transformed by the virtual strain.

Parameters:

• graph (Data) – Batched atomic graph.

• compute_forces (bool, optional) – Return forces. Default: False.

• compute_stress (bool, optional) – Return stress and virials. Requires a nonsingular cell for every structure. Default: False.

• compute_edge_forces (bool, optional) – Return edge_forces. Requires graph.vectors. Default: False.

• create_graph (bool, optional) – Construct a differentiable graph for the energy derivatives, allowing training on forces or stress. Default: False.

• use_compile (bool, optional) – Evaluate a compiled implementation, cached by the prediction and compilation options. Default: False.

• compile_mode (str, optional) – torch.compile mode. Default: “reduce-overhead”.

• fullgraph (bool, optional) – Compile without graph breaks. Default: True.

• dynamic (bool, optional) – Allow dynamic shapes when compiling. Default: False.

Returns:

For ordinary batched graphs, node_energies has

shape (n_atoms,) and energy has shape (n_graphs,). Requested derivatives are returned as forces with shape (n_atoms, 3), stress and virials with shape (n_graphs, 6), and edge_forces with shape (n_edges, 3). Stress and virials use the order xx, yy, zz, yz, xz, xy. For LAMMPS graphs, site energies cover only the owned atoms.

Return type:

dict[str, torch.Tensor]

clear_compile_cache()

Clear cached prediction functions so they recompile on next use.

pre_fit(train_loader)

Hook for fitting data statistics before training (default: no-op).

Parameters:

train_loader (DataLoader) – Training data.

classmethod from_config(cfg)

Construct a model from the subclass’s constructor settings.

Parameters:

cfg (dict) – Keys as accepted by the subclass constructor.

Returns:

Model constructed from the supplied settings.

Return type:

AtomisticModel

to_model_config()

Return the constructor settings needed to reconstruct the model.

Returns:

Model constructor settings, excluding trained weights.

Return type:

dict[str, Any]

save_weights(path, filename='model.pt')

Save the state dict.

Parameters:

• path (str | Path) – Directory, created if missing.

• filename (str, optional) – Weights file name. Default: “model.pt”.

load_weights(path, filename='model.pt', map_location='cpu')

Load a state dict saved by save_weights.

Parameters:

• path (str | Path) – Directory containing the weights file.

• filename (str, optional) – Weights file name. Default: “model.pt”.

• map_location (str | torch.device, optional) – Where tensors are loaded. Default: “cpu”.

save_inference_checkpoint(path, meta=None)

Save the weights and configuration needed to reconstruct the model.

The checkpoint contains model.pt and model_config.yaml, with optional run metadata in meta.yaml.

Parameters:

• path (str | Path) – Checkpoint directory, created if missing.

• meta (dict, optional) – Additional metadata written to meta.yaml. Default: None.

classmethod load_model_config(ckpt_dir)

Read model_config.yaml from a checkpoint directory.

Parameters:

ckpt_dir (str | Path) – Checkpoint directory.

Returns:

Model constructor settings read from the checkpoint.

Return type:

dict[str, Any]

classmethod from_checkpoint(ckpt_dir, device=None, map_location='cpu')

Rebuild a model from a checkpoint directory.

Parameters:

• ckpt_dir (str | Path) – Directory with model_config.yaml and model.pt.

• device (str | torch.device, optional) – Device to which the reconstructed model is moved. Default: None, retaining the device used by the constructor.

• map_location (str | torch.device, optional) – Load location for the state dict. Default: “cpu”.

Returns:

Reconstructed model with loaded weights in evaluation mode.

Return type:

AtomisticModel