flashcart.utils

These functions handle configuration, distributed execution, graph geometry, compilation, logging, and reductions. They support the interfaces described in flashcart.model, flashcart.calculators, and flashcart.cli.

Use the configuration helpers to combine defaults, YAML files, and overrides in a Python script. The geometry helpers construct edge vectors and convert stress tensors to the Voigt convention used by model predictions. Compilation settings affect the PyTorch process. The distributed helpers manage execution on individual ranks, and the LAMMPS helpers exchange atomic features between ranks.

For user-facing settings and examples, see Configuration files, Environment variables, and Multi-GPU training.

flashcart.utils.config

Read configuration files and combine defaults with supplied settings.

flashcart.utils.config.save_yaml(path, data)

Write a dict as YAML, creating parent directories.

Parameters:

• path (str | Path) – Output file.

• data (dict) – Mapping to dump (key order preserved).

flashcart.utils.config.load_yaml(path)

Read a YAML file.

Empty files and files containing only null produce an empty dictionary. Other YAML values are returned as parsed. A top-level mapping is not required.

Parameters:

path (str | Path) – Input file.

Returns:

Parsed YAML value, or an empty dictionary if the file is empty or

contains only null.

Return type:

Any

flashcart.utils.config.load_config(overrides=None, required=['train_path', 'valid_path', 'output_path'])

Load the packaged defaults and replace settings with the supplied overrides.

Merging replaces values at the top level. Nested mappings are replaced rather than recursively merged. Defaults come from flashcart/configs/default.yaml.

Parameters:

• overrides (str | Path | dict | argparse.Namespace, optional) – YAML file, mapping, or namespace containing overrides. None-valued entries are skipped only for namespaces. Default: None.

• required (list[str], optional) – Keys that must be non-None in the result. Pass [] to skip the check. Default: REQUIRED_KEYS.

Returns:

Default configuration updated with the supplied overrides and

checked for required values.

Return type:

dict[str, Any]

Raises:

ValueError – A required setting is missing or None.

flashcart.utils.config.parse_override(s)

Parse one KEY=VALUE command-line override into a key and value.

Values are interpreted first as integers or floats, then as explicit boolean or null values. Remaining values are parsed with yaml.safe_load. The original string is retained if YAML parsing fails.

Parameters:

s (str) – Override of the form “KEY=VALUE”.

Returns:

Override key and parsed value.

Return type:

tuple[str, Any]

Raises:

ValueError – The override has no key or does not contain an equals sign.

flashcart.utils.config.parse_config_args(args)

Parse the CLI convention [CONFIG.yaml] [KEY=VALUE ...] into overrides.

An optional leading YAML file is read first. KEY=VALUE pairs then replace the corresponding values.

Parameters:

args (Sequence[str]) – CLI arguments after the program name.

Returns:

Overrides read from the optional file and subsequent

assignments. Later assignments replace earlier values for the same key.

Return type:

dict[str, Any]

Raises:

ValueError – The loaded configuration is not a mapping, or an assignment does not have the form KEY=VALUE.

flashcart.utils.config.torch_device_from_config(device)

Translate the configuration value "gpu" to the PyTorch value "cuda".

Parameters:

device (Any) – Device spec from the run config.

Returns:

"cuda" for a case-insensitive "gpu" string. Other input values are

returned unchanged.

Return type:

Any

flashcart.utils.distributed

Manage batch sizes and loader iteration across distributed ranks.

flashcart.utils.distributed.fabric_from_config(cfg)

Build an unlaunched Fabric from the run config’s device keys.

Parameters:

cfg (Mapping) – Configuration providing device (default "gpu"), devices and strategy (default "auto"), n_nodes (default 1), and precision (default None).

Returns:

Configured Fabric instance, which has not yet been launched.

Return type:

Fabric

flashcart.utils.distributed.per_rank_batch_size(global_batch_size, world_size)

Split a global batch size across ranks by rounding upward.

Rounding upward can make the combined batch size exceed the requested global size. A world size below one is treated as one.

Parameters:

• global_batch_size (int) – Requested global batch size, before division across ranks.

• world_size (int) – Number of ranks.

Returns:

Ceiling of the global batch size divided by the world size, with a minimum

of one.

Return type:

int

flashcart.utils.distributed.is_global_zero(fabric)

Determine whether execution is on the global rank zero process.

Parameters:

fabric (Fabric or None) – Fabric instance, or None for execution without Fabric.

Returns:

True when fabric is None, its global_rank is zero, or it has no

global_rank attribute.

Return type:

bool

flashcart.utils.compile

Configure PyTorch compilation for model evaluation.

flashcart.utils.compile.configure_autograd_for_compile(allow_autograd=True)

Configure Dynamo to trace prediction functions that differentiate energy.

The function controls whether torch.autograd.grad is permitted in compiled graphs, raises the recompilation limit to at least eight, and disables Dynamo’s DDP optimization when that option is available. These settings affect the process globally. If Dynamo is unavailable, the function returns without making changes.

Parameters:

allow_autograd (bool, optional) – Allow (True) or forbid (False) autograd.grad inside compiled graphs. Default: True.

flashcart.utils.env

Read Boolean and integer settings from environment variables.

flashcart.utils.env.env_bool(name, default=False)

Interpret an environment variable as a boolean.

The values "1", "true", "yes", and "t" are true, ignoring case. Other values are false.

Parameters:

• name (str) – Environment variable to read.

• default (bool, optional) – Value used when the variable is unset. Default: False.

Returns:

Parsed value, or the default when the variable is unset.

Return type:

bool

flashcart.utils.env.env_int(name, default)

Read an integer environment variable with a fallback value.

Parameters:

• name (str) – Environment variable to read.

• default (int) – Value used when the variable is unset or cannot be parsed as an integer.

Returns:

Parsed value, or the provided default.

Return type:

int

flashcart.utils.geometry

Compute edge vectors and convert stress tensors.

flashcart.utils.geometry.uses_periodic_shifts(pbc)

Whether edge vectors must include periodic shift contributions.

Parameters:

pbc (object) – Per-axis periodicity flags (tensor/array/None). None is treated as periodic.

Returns:

Whether any periodicity flag is true, or True when the flags are

unavailable.

Return type:

bool

flashcart.utils.geometry.get_edge_vectors(positions, cell, shifts, edge_index, batch, use_shifts=True)

Compute receiver-to-sender edge vectors with optional periodic translations.

Parameters:

• positions (torch.Tensor) – Atom positions [n_atoms, 3].

• cell (torch.Tensor) – Cells [n_graphs, 3, 3].

• shifts (torch.Tensor) – Integer unit-cell shifts per edge [n_edges, 3].

• edge_index (torch.Tensor) – Shape (2, n_edges). Row 0 contains receivers and row 1 contains senders.

• batch (torch.Tensor) – Graph index per atom.

• use_shifts (bool, optional) – Include periodic translations in the edge vectors. Set False for fully non-periodic batches. Default: True.

Returns:

Receiver-to-sender vectors with shape (n_edges, 3), including

the specified periodic translations when use_shifts is True.

Return type:

torch.Tensor

flashcart.utils.geometry.to_voigt6(t)

Symmetrize a 3x3 tensor into Voigt order (xx, yy, zz, yz, xz, xy).

Parameters:

t (np.ndarray | torch.Tensor) – Tensor(s) of shape […, 3, 3].

Returns:

Symmetrized components with shape (..., 6),

preserving the input array type. Off-diagonal components are averaged, without an additional engineering-shear factor.

Return type:

np.ndarray | torch.Tensor

flashcart.utils.lammps

Exchange atomic features and their gradients between LAMMPS ranks. Forward and reverse exchanges are PyTorch custom operators that support compiled predictions. When calling the exchange helpers directly, keep both evaluation and differentiation inside lammps_data_slot(), including on ranks without ghost atoms. The LAMMPS calculator manages this context automatically.

flashcart.utils.lammps.lammps_exchange_features(feats, nlocal, ntotal)

Refresh ghost-atom features from their owning ranks between layers.

Features of shape (nlocal, C) are padded with zeros to (ntotal, C) before the exchange fills the ghost rows. All ranks take part, including ranks without ghost atoms, because their owned atoms may be ghosts on neighboring ranks. Evaluation and differentiation must run inside lammps_data_slot().

During differentiation, reverse communication accumulates gradients from ghost rows onto their owners.

Parameters:

• feats (torch.Tensor) – Node features with shape (nlocal, C) or (ntotal, C), where C is the number of features. Owned rows come first.

• nlocal (int) – Number of atoms owned by this rank.

• ntotal (int) – Total number of owned and ghost atoms on this rank.

Returns:

Features with shape (ntotal, C) after ghost rows have been

filled.

Return type:

torch.Tensor

Raises:

RuntimeError – The feature row count is incompatible with nlocal and ntotal, or no communication object with both exchange hooks is active.

flashcart.utils.lammps.lammps_data_slot(lammps_data)

Expose an ML-IAP data object to the exchange operators during evaluation.

Both forward evaluation and differentiation must run inside this context. The communication object is stored process-wide, and the previous object is restored on exit. Concurrent evaluations using different LAMMPS objects are not supported.

Parameters:

lammps_data (Any) – ML-IAP data object providing forward_exchange and reverse_exchange.

flashcart.utils.lammps.lammps_forward_exchange(feats)

Copy features from owning ranks to ghost atoms.

Requires an active lammps_data_slot() context.

Parameters:

feats (torch.Tensor) – Features with shape (ntotal, C), where ntotal counts owned and ghost atoms and C is the number of features. Owned rows come first. Noncontiguous inputs are accepted. The input is not modified.

Returns:

Newly allocated contiguous features with the same shape, dtype,

and device as feats. Owned rows are preserved, and ghost rows hold the features of their owners.

Return type:

torch.Tensor

Raises:

RuntimeError – No ML-IAP data object with both exchange hooks is active.

flashcart.utils.lammps.lammps_reverse_exchange(grads)

Accumulate ghost-row gradients on their owning ranks.

Requires an active lammps_data_slot() context.

Parameters:

grads (torch.Tensor) – Gradients with shape (ntotal, C), where ntotal counts owned and ghost atoms and C is the number of features. Owned rows come first. Noncontiguous inputs are accepted. The input is not modified.

Returns:

Newly allocated contiguous gradients with the same shape,

dtype, and device as grads. Ghost rows are zero and their contributions are added to the owner rows.

Return type:

torch.Tensor

Raises:

RuntimeError – No ML-IAP data object with both exchange hooks is active.

flashcart.utils.logging

Configure log output and report training progress.

flashcart.utils.logging.setup_logger(name, log_file, enabled=True)

Configure a named logger to write to standard error and a file.

Existing handlers are removed and closed before the new handlers are installed. Disabled logging uses a NullHandler and does not create the output file.

Parameters:

• name (str) – Logger name.

• log_file (Path) – Log file, parents created.

• enabled (bool, optional) – False installs only a NullHandler. Default: True.

Returns:

Configured logger with INFO level and propagation disabled.

Return type:

logging.Logger

flashcart.utils.parameter_groups

Collect parameters assigned to specific optimizer groups.

flashcart.utils.parameter_groups.iter_child_special_parameters(module, method_name)

Yield each parameter once from child modules at all nesting levels.

Call the named method on each nested child module that has it. Skip the supplied module itself. If several child modules return the same parameter object, yield it only once.

Parameters:

• module (nn.Module) – Module containing the child modules to inspect.

• method_name (str) – Name of the generator method to call.

Yields:

nn.Parameter – Each reported parameter once, in the order found.

flashcart.utils.scatter

Sum tensor values by index.

Trimmed-down pytorch_scatter: https://github.com/rusty1s/pytorch_scatter. Copyright (c) 2020 Matthias Fey, MIT License (see NOTICE).

flashcart.utils.scatter.scatter_sum(src, index, dim=-1, out=None, dim_size=None)

Sum src entries with equal index along dim (scatter-add).

Parameters:

• src (torch.Tensor) – Values to reduce.

• index (torch.Tensor) – Target bin per entry, broadcast to the shape of src.

• dim (int, optional) – Dimension to reduce along. Default: -1.

• out (torch.Tensor, optional) – Tensor into which values are added in place. Default: None, which allocates a zero-initialized output.

• dim_size (int, optional) – Output size along dim when allocating an output. Default: None, which uses the largest index plus one, or zero for an empty index tensor.

Returns:

Sum for each output index. If out is provided, returns that

tensor after accumulation. Otherwise, returns the newly allocated tensor.

Return type:

torch.Tensor