flashcart.data

Represent atomic configurations and neighbor graphs, construct batches, and compute the statistics used to initialize a model. AtomicConfig stores a configuration and its reference properties as NumPy arrays. AtomicData holds the PyTorch tensors used for model evaluation, including atom types and neighbor indices. Graph construction connects these representations using the model cutoff and an ordered list of chemical elements.

See Data preparation for data preparation and reference-property keys. The sections below describe the individual representations and processing steps.

flashcart.data.data

Represent atomic configurations and tensor graphs.

class flashcart.data.data.AtomicConfig(atomic_numbers, positions, cell=None, pbc=None, energy=None, forces=None, stress=None)

Bases: object

Store one atomic configuration and its optional reference properties.

The configuration stores NumPy arrays and attaches a neighbor list when compute_neighbors is called. Neighbor-list reuse depends on the stored cutoff and skin. Changes to the geometry do not invalidate the list automatically.

Parameters:

• atomic_numbers (np.ndarray) – Per-atom atomic numbers.

• positions (np.ndarray) – Cartesian positions, shape (n, 3).

• cell (np.ndarray, optional) – 3x3 cell matrix. Default: zeros.

• pbc (list[bool] | bool, optional) – Periodic flags. Default: all False.

• energy (float, optional) – Reference total energy. Default: None.

• forces (np.ndarray, optional) – Reference forces, with shape (n_atoms, 3). Default: None.

• stress (np.ndarray, optional) – Reference stress, with shape (3, 3) or (6,) in the order xx, yy, zz, yz, xz, xy. Matrix inputs are symmetrized and converted to this order. Default: None.

compute_neighbors(r_max, skin=0.0, backend='matscipy')

Build or reuse the neighbor list at r_max + skin.

An existing list is reused when its cutoff and skin match the requested values. This check does not compare positions, cell, periodicity, or backend. Clear neighbors before calling this method if these have changed and a new list is required.

Parameters:

• r_max (float) – Cutoff radius.

• skin (float, optional) – Extra radius for reuse. Default: 0.0.

• backend (str, optional) – Neighbor backend. Default: “matscipy”.

classmethod from_ase(atoms, energy_key='REF_energy', forces_key='REF_forces', stress_key='REF_stress')

Read an ASE configuration and its stored reference properties.

Reference energies and stresses are read from atoms.info and forces from atoms.arrays. Missing properties are stored as None.

Parameters:

• atoms (ase.Atoms) – Atomic configuration to convert.

• energy_key (str, optional) – Key for the reference energy in atoms.info. Default: “REF_energy”.

• forces_key (str, optional) – Key for the reference forces in atoms.arrays. Default: “REF_forces”.

• stress_key (str, optional) – Key for the reference stress in atoms.info. Default: “REF_stress”.

Returns:

Configuration containing the geometry and available reference

properties.

Return type:

AtomicConfig

to_ase(energy_key='REF_energy', forces_key='REF_forces', stress_key='REF_stress')

Convert the configuration to ASE and store its reference properties.

Only properties present in the configuration are written.

Parameters:

• energy_key (str, optional) – Key for the reference energy in atoms.info. Default: “REF_energy”.

• forces_key (str, optional) – Key for the reference forces in atoms.arrays. Default: “REF_forces”.

• stress_key (str, optional) – Key for the reference stress in atoms.info. Default: “REF_stress”.

Returns:

Atomic configuration with the available reference properties

stored under the specified keys.

Return type:

ase.Atoms

class flashcart.data.data.AtomicData(atomic_numbers, positions, cell, pbc, edge_index, unit_shifts, atom_types, energy=None, forces=None, stress=None)

Bases: Data

Represent one atomic configuration as a graph of PyTorch tensors.

Floating-point fields use torch.get_default_dtype(). Connectivity, atomic numbers, atom types, and periodic shifts use integer tensors. The shifts specify translations in units of the cell vectors. Cartesian translations are formed when edge vectors are evaluated.

Parameters:

• atomic_numbers (np.ndarray) – Per-atom atomic numbers.

• positions (np.ndarray) – Cartesian positions, shape (n, 3).

• cell (np.ndarray) – 3x3 cell matrix.

• pbc (np.ndarray) – Periodic flags.

• edge_index (np.ndarray) – Pair indices, shape (2, n_pairs). Row 0 contains receivers and row 1 contains senders.

• unit_shifts (np.ndarray) – Integer translations applied to the sender, with shape (n_pairs, 3).

• atom_types (np.ndarray or None) – Zero-based element indices, with shape (n_atoms,), or None when atom types are unavailable.

• energy (float, optional) – Reference total energy. Any desired energy offset must already have been subtracted. Default: None.

• forces (np.ndarray, optional) – Reference forces, with shape (n_atoms, 3). Default: None.

• stress (np.ndarray, optional) – Reference stress, with shape (3, 3) or (6,) in ASE Voigt order. Stored as a tensor of shape (1, 6), together with virials = -stress * volume. Default: None.

classmethod from_config(config, elements, energy_offset=0.0)

Convert an atomic configuration to a tensor graph.

If no neighbor list is attached, the graph has no edges. An attached neighbor list must already have been built. The reference energy is reduced by energy_offset when it is present.

Parameters:

• config (AtomicConfig) – Source structure.

• elements (list[str | int]) – Element order defining atom types.

• energy_offset (float, optional) – Subtracted from the reference energy (the summed per-element shifts). Default: 0.0.

Returns:

Graph containing the geometry, connectivity, atom types, and

available reference properties.

Return type:

AtomicData

to_config()

Convert a single graph to a NumPy configuration.

Tensor fields are detached and transferred to CPU. Connectivity and integer periodic shifts are retained when present. The stored energy is copied without restoring any offset subtracted by from_config.

Returns:

Configuration containing the graph’s geometry, available

reference properties, and neighbor list.

Return type:

AtomicConfig

flashcart.data.dataset

Load datasets, build graph batches, and split data for training and evaluation.

For training, make_loader() combines dataset loading, graph construction, and batching. Its batch_size argument is a target for one rank. The training command divides the global target among ranks before calling this function.

flashcart.data.dataset.make_loader(path, elements, r_max, batch_size, energy_key='REF_energy', forces_key='REF_forces', shuffle=False, dynamic_batching=True, n_replicas=1, rank=0, drop_last=False, balance_batches=False, data_seed=0, atomic_shifts=None, n_workers=0)

Construct an extxyz data loader with fixed or variable batch sizes.

Dynamic batching groups structures within atom and edge budgets computed from batch_size. For fixed distributed batching, ranks use DistributedSampler when balance_batches or drop_last is set. Otherwise, rank-strided sampling assigns each structure to one rank without duplication.

Parameters:

• path (str | Path) – extxyz file.

• elements (list[str]) – Element order defining atom types.

• r_max (float) – Neighbor cutoff radius.

• batch_size (int) – Structures per batch on this rank. With dynamic batching, this value determines the atom and edge budgets. The caller must divide a global target among ranks before calling this function.

• energy_key (str, optional) – Key for the reference energy in atoms.info. Default: “REF_energy”.

• forces_key (str, optional) – Key for the reference forces in atoms.arrays. Default: “REF_forces”.

• shuffle (bool, optional) – Shuffle per epoch. Default: False.

• dynamic_batching (bool, optional) – Group structures by atom and edge budgets. Default: True.

• n_replicas (int, optional) – Number of distributed ranks. Default: 1.

• rank (int, optional) – Index of this rank. Default: 0.

• drop_last (bool, optional) – For fixed batching, discard incomplete batches. For dynamic distributed batching with balance_batches, truncate all ranks to the smallest rank batch count. Otherwise, dynamic batching retains the final batch. Default: False.

• balance_batches (bool, optional) – Equalize the number of batches across distributed ranks. This may repeat samples or discard batches, depending on the sampler and drop_last. Default: False.

• data_seed (int, optional) – Seed used by shuffling and sampling. Default: 0.

• atomic_shifts (Sequence[float], optional) – Per-element energy offsets, ordered as elements, whose sum is subtracted from each reference energy during graph construction. Default: None.

• n_workers (int, optional) – DataLoader worker processes. Default: 0.

Returns:

Loader yielding batches of atomic graphs.

Return type:

DataLoader

class flashcart.data.dataset.AtomicDataset(configs, r_max, elements, skin=0.0, atomic_shifts=None)

Bases: Dataset

In-memory dataset of AtomicConfig structures.

Neighbor lists are built when needed and stored on each configuration for reuse. When atomic_shifts is set, the summed per-element shifts are subtracted from each structure’s reference energy when its graph is built. The model then learns energies relative to these offsets. Calculators can add the shifts back.

Parameters:

• configs (list[AtomicConfig]) – Structures.

• r_max (float) – Neighbor cutoff radius.

• elements (list[str | int]) – Element order defining atom types.

• skin (float, optional) – Extra neighbor radius. Default: 0.0.

• atomic_shifts (Sequence[float], optional) – Per-element energy offsets, ordered as elements, whose sum is subtracted from each reference energy during graph construction. Default: None.

classmethod from_extxyz(file_path, r_max, elements, skin=0.0, index=':', energy_key='REF_energy', forces_key='REF_forces', stress_key='REF_stress', atomic_shifts=None)

Read selected extxyz frames into an atomic dataset.

Parameters:

• file_path (str | Path) – Input extxyz file.

• r_max (float) – Neighbor-list cutoff.

• elements (list[str | int]) – Element order defining atom types.

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

• index (str | int | slice, optional) – ASE frame selection. Default: “:”, which selects all frames.

• energy_key (str, optional) – Key for the reference energy in atoms.info. Default: “REF_energy”.

• forces_key (str, optional) – Key for the reference forces in atoms.arrays. Default: “REF_forces”.

• stress_key (str, optional) – Key for the reference stress in atoms.info. Default: “REF_stress”.

• atomic_shifts (Sequence[float], optional) – Per-element energy offsets, ordered as elements, whose sum is subtracted from each reference energy during graph construction. Default: None.

Returns:

Dataset containing the selected configurations.

Return type:

AtomicDataset

len()

Return the underlying dataset size before applying an index selection.

Returns:

Number of examples supplied by the subclass.

Return type:

int

get(idx)

Construct the graph for one configuration.

The neighbor list is built or reused, and the summed per-element energy offset is subtracted when offsets have been provided.

Parameters:

idx (int) – Configuration index.

Returns:

Graph containing the configuration and its reference properties.

Return type:

AtomicData

get_sizes()

Count atoms and neighbor-list edges in each configuration.

Neighbor lists are built or reused before their edges are counted.

Returns:

Atom and edge counts for each configuration, in

dataset order.

Return type:

list[tuple[int, int]]

class flashcart.data.dataset.HDF5Dataset(file_path, r_max, elements, skin=0.0)

Bases: Dataset

Read atomic configurations from HDF5 when they are requested.

The file is opened on first access, and the open handle is excluded when the dataset is pickled. Each requested configuration is read and converted to a graph with a newly constructed neighbor list.

Parameters:

• file_path (str | Path) – HDF5 file written by write_hdf5.

• r_max (float) – Neighbor cutoff radius.

• elements (list[str | int]) – Element order defining atom types.

• skin (float, optional) – Extra neighbor radius. Default: 0.0.

len()

Return the underlying dataset size before applying an index selection.

Returns:

Number of examples supplied by the subclass.

Return type:

int

get(idx)

Read one configuration and construct its atomic graph.

Parameters:

idx (int) – Index in the sorted HDF5 group names.

Returns:

Graph containing the stored reference properties and a neighbor

list constructed at r_max + skin.

Return type:

AtomicData

flashcart.data.dataset.random_split_indices(n_samples, sizes, seed=None)

Partition a random permutation of sample indices into named splits.

Splits are filled in the order given by sizes. Any remaining indices go to "test". Sizes should be nonnegative and sum to at most n_samples. If a remainder is expected, reserve the name "test" for it. An existing entry with that name is overwritten.

Parameters:

• n_samples (int) – Total number of samples.

• sizes (dict[str, int]) – Number of samples for each split, in insertion order.

• seed (int, optional) – Random-number generator seed. Default: None.

Returns:

Sample indices for each named split.

Return type:

dict[str, np.ndarray]

flashcart.data.dataset.split_dataset(dataset, sizes, seed=None)

Split a dataset into subsets using random_split_indices.

Split sizes and the reserved "test" remainder follow random_split_indices.

Parameters:

• dataset – Dataset to split.

• sizes (dict[str, int]) – Number of samples for each split, in insertion order.

• seed (int, optional) – Random-number generator seed. Default: None.

Returns:

Dataset subsets keyed by split name.

Return type:

dict[str, Subset]

flashcart.data.graph

Construct neighbor graphs from atomic configurations.

flashcart.data.graph.graph_from_config(config, elements, skin=0.0, device=None)

Construct a model graph from a configuration with a built neighbor list.

Parameters:

• config (AtomicConfig) – Configuration whose neighbor list has already been built.

• elements (list[str | int]) – Element order defining atom types.

• skin (float, optional) – Unused by this function. The graph retains the skin recorded in config.neighbors. Default: 0.0.

• device (torch.device, optional) – Device to which the graph is moved. Default: None, which leaves the graph on CPU.

Returns:

Graph for one configuration, with every atom assigned to graph index

zero.

Return type:

AtomicData

Raises:

ValueError – The configuration has no built neighbor list.

flashcart.data.graph.graph_from_ase(atoms, elements, r_max, skin=0.0, device=None)

Build a model graph directly from an ase.Atoms object.

Parameters:

• atoms (ase.Atoms) – Structure to convert.

• elements (list[str | int]) – Element order defining atom types.

• r_max (float) – Cutoff radius for the neighbor list.

• skin (float, optional) – Additional neighbor-list radius. The list is built at r_max + skin. Default: 0.0.

• device (torch.device, optional) – Device to which the graph is moved. Default: None, which leaves the graph on CPU.

Returns:

Graph containing the atomic configuration, neighbor list, and

available reference properties.

Return type:

AtomicData

flashcart.data.graph.update_graph_positions(graph, atoms)

Update graph geometry in place while retaining its connectivity.

Positions, cell, periodicity flags, and use_shifts are updated from the ASE configuration. Connectivity and integer periodic shifts are unchanged. The caller must ensure that the retained neighbor list remains valid for the new geometry.

Parameters:

• graph (AtomicData) – Graph to update.

• atoms (ase.Atoms) – Configuration providing the new geometry.

Returns:

The input graph with updated geometry.

Return type:

AtomicData

flashcart.data.neighbors

Build and update neighbor lists.

class flashcart.data.neighbors.NeighborList(r_max, skin=0.0, backend='matscipy')

Bases: object

Store atom pairs and their integer periodic translations.

The list is constructed at r_max + skin. The additional radius allows connectivity to be reused while it still covers all neighbors within r_max. In FlashCart evaluation, the cutoff envelope removes contributions from pairs outside the model cutoff.

Parameters:

• r_max (float) – Cutoff radius.

• skin (float, optional) – Extra radius for reuse across steps. Default: 0.0.

• backend (str, optional) – Neighbor backend. Default: “matscipy”.

property is_built

Whether pair indices have been assigned.

copy()

Copy the neighbor list and its arrays.

Returns:

Independent copies of the pair-index and shift arrays, or an

unbuilt list if the original has not been built.

Return type:

NeighborList

build(positions, cell, pbc)

Compute edge_index and unit_shifts for one structure.

Parameters:

• positions (np.ndarray) – Cartesian positions, shape (n, 3).

• cell (np.ndarray) – 3x3 cell matrix (rows are lattice vectors).

• pbc (bool | list) – Periodic flags.

Returns:

This instance with edge_index and unit_shifts updated.

Pair-index row 0 contains receivers and row 1 contains senders. Shifts translate the senders by integer cell vectors.

Return type:

NeighborList

classmethod from_arrays(edge_index, unit_shifts, r_max=nan, skin=0.0, backend='matscipy')

Wrap precomputed arrays (e.g. restored from a serialized graph).

Parameters:

• edge_index (np.ndarray) – Pair indices, shape (2, n_pairs).

• unit_shifts (np.ndarray) – Integer image shifts, shape (n_pairs, 3).

• r_max (float, optional) – Model cutoff associated with the list. The construction radius is r_max + skin. Default: NaN.

• skin (float, optional) – Additional radius used when constructing the list. Default: 0.0.

• backend (str, optional) – Backend to use if the list is rebuilt. Default: “matscipy”.

Returns:

List containing the provided arrays, converted to the required

integer dtypes. Arrays may share memory with the inputs when conversion does not require a copy.

Return type:

NeighborList

flashcart.data.padding

Pad graph batches for compiled evaluation using reusable tensor buffers. Complete evaluation and backward passes before passing another batch to the same padder.

class flashcart.data.padding.PadAtomicData(r_max, atom_budget=0, edge_budget=0, graph_budget=0, extra_atoms=None, atom_headroom=1.0, atom_multiple=1, edge_headroom=1.0, edge_multiple=128, graph_headroom=1.0, graph_multiple=1)

Bases: object

Pad atomic graphs to reusable atom, edge, and graph capacities.

Real atoms, edges, and graphs occupy the leading entries. Padding atoms belong to the final graph and lie on a cubic grid with spacing 2 * r_max. Padding edges connect distinct padding atoms, so their lengths exceed the model cutoff.

Padding edges are grouped by receiver and distributed across the padding atoms. This retains receiver ordering and avoids assigning all padding edges to one receiver in the tensor-product kernels. Real edges must already have the ordering required by those kernels.

Capacities grow when an input exceeds them and are rounded to the specified multiples. The atom capacity includes extra atoms for padding edges, and the graph capacity includes one padding graph. Tensor shapes and storage are reused while capacities, device, and dtypes remain unchanged.

Returned graphs reference reusable buffers. Complete any evaluation and backward pass that uses these buffers before the next call.

Parameters:

• r_max (float) – Model cutoff. Padding atoms sit on a 2 * r_max grid.

• atom_budget (int, optional) – Initial atom budget. Default: 0.

• edge_budget (int, optional) – Initial edge budget. Default: 0.

• graph_budget (int, optional) – Initial graph budget. Default: 0.

• extra_atoms (int, optional) – Extra atoms reserved for padding edges, with a minimum of two. Default: None, which sets the count from atom_budget and edge_headroom.

• atom_headroom (float, optional) – Factor applied to the required atom count when increasing its capacity. Default: 1.0.

• atom_multiple (int, optional) – Multiple to which the total atom count is rounded. Default: 1.

• edge_headroom (float, optional) – Factor applied to the required edge count when increasing its capacity. Default: 1.0.

• edge_multiple (int, optional) – Multiple to which the total edge count is rounded. Default: 128.

• graph_headroom (float, optional) – Factor applied to the required graph count when increasing its capacity. Default: 1.0.

• graph_multiple (int, optional) – Multiple to which the total graph count is rounded. Default: 1.

property extra_atoms

Number of atoms reserved as endpoints of padding edges.

property total_atoms

Total atom capacity, including reserved atoms and rounding.

property total_edges

Total edge capacity after rounding.

property total_graphs

Total graph capacity, including the padding graph and rounding.

seed_budgets(atom_budget=0, edge_budget=0, graph_budget=0)

Increase the initial capacities without reducing existing values.

Parameters:

• atom_budget (int, optional) – Minimum capacity for real atoms. Default: 0.

• edge_budget (int, optional) – Minimum capacity for real edges. Default: 0.

• graph_budget (int, optional) – Minimum capacity for real graphs. Default: 0.

__call__(graph)

Copy graph inputs into the reusable padding buffers.

Only fields required for model evaluation are retained. Reference energies, forces, and stresses are not copied.

Parameters:

graph (Data) – Atomic graph or batch to pad.

Returns:

Padded graph, number of real atoms, and number of

real graphs. The graph references reusable storage and must no longer be in use when the padder is called again.

Return type:

tuple[Data, int, int]

flashcart.data.padding.slice_padded_outputs(out, n_real_atoms, n_real_graphs=1)

Remove padding entries from model predictions.

Parameters:

• out (dict) – Output of model.predict on a padded graph.

• n_real_atoms (int) – Real-atom count returned by the padder.

• n_real_graphs (int, optional) – Real-graph count. Default: 1.

Returns:

Dictionary containing predictions for real atoms and graphs.

Tensor slices share storage with the inputs. Unrecognized entries are retained unchanged.

Return type:

dict[str, Any]

flashcart.data.samplers

Assign structures to ranks and group them into batches. Batches can be constrained by atom and edge counts.

class flashcart.data.samplers.DistributedStridedSampler(dataset, n_replicas, rank, shuffle=False, seed=0)

Bases: Sampler

Assign each dataset index to one distributed rank.

Rank r receives indices[r::n_replicas] after optional shuffling. Samples are not duplicated to equalize rank lengths, so the numbers of samples and batches may differ across ranks. This supports evaluation that reduces accumulated statistics after iteration. Training with synchronized gradient reductions requires compatible iteration counts.

Parameters:

• dataset (Sized) – Dataset providing its length.

• n_replicas (int) – World size.

• rank (int) – Rank of this process.

• shuffle (bool, optional) – Reshuffle globally each epoch. Default: False.

• seed (int, optional) – Base seed. The epoch is added to it. Default: 0.

set_epoch(epoch)

Set the shuffle epoch (called by the training loop).

Parameters:

epoch (int) – Epoch number mixed into the shuffle seed.

class flashcart.data.samplers.DynamicBatchSampler(sizes, max_nodes, max_edges, shuffle=True, seed=0, n_replicas=1, rank=0, drop_last=False, balance_batches=False)

Bases: Sampler

Group variable-size graphs within atom and edge budgets.

Indices are optionally shuffled and then distributed by rank. Each rank adds graphs in order until the next graph would exceed a budget. A graph that exceeds a budget forms a batch by itself.

With distributed balance_batches=True, all ranks use the same number of batches. drop_last=True reduces this to the smallest batch count. Otherwise, ranks with fewer batches repeat their final batch. A rank with no samples uses a sample from the full index list when needed. Balanced sampling can therefore discard or repeat samples.

The shuffle epoch advances after a complete iteration. If iteration stops early, the shuffle epoch does not advance.

Parameters:

• sizes (list[tuple[int, int]]) – (n_nodes, n_edges) per structure.

• max_nodes (int) – Node budget per batch.

• max_edges (int) – Edge budget per batch.

• shuffle (bool, optional) – Reshuffle each epoch. Default: True.

• seed (int, optional) – Base seed. The epoch is added to it. Default: 0.

• n_replicas (int, optional) – World size. Default: 1.

• rank (int, optional) – Rank of this process. Default: 0.

• drop_last (bool, optional) – Truncate balanced distributed ranks to the smallest batch count. Has no effect without distributed batch balancing. Default: False.

• balance_batches (bool, optional) – Equal batch counts across ranks. Default: False.

property max_graphs

Upper bound on structures per batch, over every possible shuffle.

Node and edge counts are sorted separately. Adding the smallest counts first can fit at least as many structures as any actual batch, so the result is an upper bound. A structure that exceeds a budget can still form its own batch.

Returns:

Upper bound on the number of graphs in a batch, with a minimum of one,

including for an empty dataset.

Return type:

int

classmethod from_target_batch_size(sizes, target_batch_size, factor=1.1, shuffle=True, seed=0, n_replicas=1, rank=0, drop_last=False, balance_batches=False)

Derive atom and edge budgets from a target batch size.

Budgets are computed separately as factor * target_batch_size * mean_size, converted to integers with a minimum of one. sizes must be nonempty. The target controls the budgets. Actual batch sizes depend on graph sizes and their order.

Parameters:

• sizes (list[tuple[int, int]]) – (n_nodes, n_edges) per structure.

• target_batch_size (int) – Average structures per batch to aim for.

• factor (float, optional) – Multiplier for the budgets computed from average graph sizes. Default: 1.1.

• shuffle (bool, optional) – Reshuffle each epoch. Default: True.

• seed (int, optional) – Base seed. Default: 0.

• n_replicas (int, optional) – World size. Default: 1.

• rank (int, optional) – Rank of this process. Default: 0.

• drop_last (bool, optional) – Truncate balanced distributed ranks to the smallest batch count. Has no effect without distributed batch balancing. Default: False.

• balance_batches (bool, optional) – Equal batch counts across ranks. Default: False.

Returns:

Sampler initialized with the derived budgets.

Return type:

DynamicBatchSampler

set_epoch(epoch)

Set the shuffle epoch (called by the training loop).

Parameters:

epoch (int) – Epoch number mixed into the shuffle seed.

flashcart.data.statistics

Compute dataset statistics used to initialize a potential.

flashcart.data.statistics.get_atomic_energy_shifts(data_loader, n_elements, atomic_shifts=None, lam=1.0, fit=True)

Estimate per-element energy offsets by regularized least squares.

The provided offsets define a prior. Their summed contribution is subtracted from each reference energy before fitting a correction. The residual energy per atom, averaged across all atoms, supplies a common offset. Ridge regularization is applied to the remaining element-specific corrections.

Batches without energies or atom types are skipped. If no usable batches remain, the prior is returned. With fit=False, provided offsets are returned as a float64 array without fitting.

Parameters:

• data_loader (DataLoader) – Batches with reference energies and atom types.

• n_elements (int) – Number of elements.

• atomic_shifts (np.ndarray, optional) – Prior offsets, with shape (n_elements,). Required when fit=False. Default: None, which uses zeros when fitting.

• lam (float, optional) – Ridge regularization strength for the fitted corrections. Default: 1.0.

• fit (bool, optional) – Fit residual shifts on top of the prior. Default: True.

Returns:

Per-element energy offsets, with shape (n_elements,) and dtype

float64.

Return type:

np.ndarray

flashcart.data.statistics.get_force_rms(data_loader, n_elements=None)

Compute the root mean square of the reference force components.

All Cartesian components and atoms contribute equally. Batches without forces are skipped.

Parameters:

• data_loader (DataLoader) – Batches with reference forces.

• n_elements (int, optional) – If given, return an array of this length with the same value for every element, replacing zero by 1.0. Default: None, which returns a scalar.

Returns:

Scalar force-component RMS when n_elements is None, or a

float64 array containing the same value for every element. Zero is replaced by one only in the array form. With no reference forces, the scalar is zero and the array contains ones.

Return type:

float | np.ndarray

flashcart.data.statistics.get_avg_neighbors(data_loader, r_max=None)

Average neighbor count per atom over the loader.

Parameters:

• data_loader (DataLoader) – Batches with edge indices.

• r_max (float, optional) – Recount edges within this radius (needed when the lists were built with a skin). Default: None (raw edge count).

Returns:

Total included edge count divided by the total atom count, or zero when

the loader contains no atoms.

Return type:

float

flashcart.data.utils

Convert atomic configurations and read or write structure data.

flashcart.data.utils.numbers_to_types(atomic_numbers, elements)

Map atomic numbers to 0-based indices into elements.

The returned indices are the model’s internal atom types. Their order is fixed by elements (the training element order).

Parameters:

• atomic_numbers (np.ndarray) – Per-atom atomic numbers.

• elements (list[str | int]) – Element symbols or atomic numbers. Raises on duplicates or on atomic numbers not covered by the list.

Returns:

Zero-based element indices with the same shape as atomic_numbers

and dtype int64.

Return type:

np.ndarray

flashcart.data.utils.types_to_one_hot(atom_types, n_elements)

One-hot encode atom types.

Parameters:

• atom_types (np.ndarray) – 0-based type indices.

• n_elements (int) – Number of columns of the one-hot matrix.

Returns:

One-hot matrix with shape (n_atoms, n_elements) and dtype

float64.

Return type:

np.ndarray

flashcart.data.utils.get_elements(file_path, energy_key='REF_energy', forces_key='REF_forces')

Read an extxyz file and list its elements in increasing atomic-number order.

Parameters:

• file_path (str | Path) – extxyz file to scan.

• energy_key (str, optional) – Key for the reference energy in atoms.info. Default: “REF_energy”.

• forces_key (str, optional) – Key for the reference forces in atoms.arrays. Default: “REF_forces”.

Returns:

Element symbols ordered by increasing atomic number.

Return type:

list[str]

flashcart.data.utils.read_extxyz(file_path, index=':', energy_key='REF_energy', forces_key='REF_forces', stress_key='REF_stress')

Read extxyz frames into AtomicConfig objects.

Parameters:

• file_path (str | Path) – extxyz file.

• index (str | int | slice, optional) – ase-style frame selection. Default: “:” (all frames).

• energy_key (str, optional) – Key for the reference energy in atoms.info. Default: “REF_energy”.

• forces_key (str, optional) – Key for the reference forces in atoms.arrays. Default: “REF_forces”.

• stress_key (str, optional) – Key for the reference stress in atoms.info. Default: “REF_stress”.

Returns:

Selected configurations in the requested order, including a

one-element list when one frame is selected.

Return type:

list[AtomicConfig]

flashcart.data.utils.write_extxyz(file_path, configs, energy_key='REF_energy', forces_key='REF_forces', stress_key='REF_stress', append=False)

Write atomic configurations and their reference properties as extxyz frames.

Parameters:

• file_path (str | Path) – Output file.

• configs (list[AtomicConfig]) – Structures to write.

• energy_key (str, optional) – Key for the reference energy in atoms.info. Default: “REF_energy”.

• forces_key (str, optional) – Key for the reference forces in atoms.arrays. Default: “REF_forces”.

• stress_key (str, optional) – Key for the reference stress in atoms.info. Default: “REF_stress”.

• append (bool, optional) – Append instead of overwrite. Default: False.

flashcart.data.utils.read_hdf5(file_path)

Read all structure groups of an HDF5 file written by write_hdf5.

Parameters:

file_path (str | Path) – Input file.

Returns:

Configurations read in sorted group-name order.

Return type:

list[AtomicConfig]

flashcart.data.utils.write_hdf5(file_path, configs)

Write atomic configurations to HDF5, replacing any existing file.

Each frame is stored in a structure_%08d group containing its geometry and available reference properties. Neighbor lists are not stored.

Parameters:

• file_path (str | Path) – Output file.

• configs (list[AtomicConfig]) – Structures to write.