API Reference ¶

classmethod from_ase_list(atoms_list, *, neighbor_list=None, dtype=None, device=None, requires_grad=False, metadata=None, state=None, extract_state=True)[source]¶

Concatenate ASE structures into one batched input.

Parameters:

atoms_list (Sequence[Atoms])
neighbor_list (NeighborListData | None)
dtype (dtype | None)
device (device | None)
requires_grad (bool)
metadata (Dict[str, object] | None)
state (UFPInputState | None)
extract_state (bool)

Return type:

missing_state_fields(fields)[source]¶

Return required charge/spin state fields that are absent.

Parameters:: fields (Sequence[str])
Return type:: tuple[str, …]

property n_atoms: int¶: Return the total number of atoms across all systems.

property n_systems: int¶: Return the number of systems represented by this input.

pair_atomic_numbers(mask=None)[source]¶

Return atomic numbers for the first and second atom of each selected pair.

Parameters:: mask (ndarray | Tensor | None)
Return type:: tuple[Tensor, Tensor]

pair_category_indices(atomic_types, *, symmetric=True)[source]¶

Return pair-category indices for every neighbor-list row.

Pairs containing atomic numbers outside atomic_types are marked -1. When symmetric is true, category ordering matches unordered combinations with replacement over sorted unique atomic types.

Parameters:

atomic_types (Sequence[int])
symmetric (bool)

Return type:

Tensor

pair_distances(mask=None)[source]¶

Return pair distances for all or selected neighbor-list entries.

Parameters:: mask (ndarray | Tensor | None)
Return type:: Tensor

pair_indices(mask=None)[source]¶

Return atom indices for all or selected neighbor-list pairs.

Parameters:: mask (ndarray | Tensor | None)
Return type:: tuple[Tensor, Tensor]

pair_mask(first_atomic_number, second_atomic_number, *, symmetric=False)[source]¶

Build a mask selecting pairs with the requested atomic numbers.

Parameters:

first_atomic_number (int)
second_atomic_number (int)
symmetric (bool)

Return type:

Tensor

pair_shifts(mask=None)[source]¶

Return cell-shift vectors for selected neighbor-list entries.

Parameters:: mask (ndarray | Tensor | None)
Return type:: Tensor

pair_system_index(mask=None)[source]¶

Return the owning system index for all or selected pairs.

Parameters:: mask (ndarray | Tensor | None)
Return type:: Tensor

pair_vectors(mask=None)[source]¶

Return pair displacement vectors for all or selected pairs.

Parameters:: mask (ndarray | Tensor | None)
Return type:: Tensor

pin_memory()[source]¶

Pin stored tensors in place and return self for dataloader-style use.

Return type:: UFPInput

slice_neighbor_list(mask)[source]¶

Return a NeighborListData view restricted to selected pairs.

Parameters:: mask (ndarray | Tensor)
Return type:: NeighborListData

property system_sizes: list[int]¶: Return the atom count for each system in concatenated order.

to(*, device=None, dtype=None, requires_grad=False, neighbor_list=None)[source]¶

Return a copy moved to the requested device/dtype and gradient state.

Parameters:

device (device | None)
dtype (dtype | None)
requires_grad (bool)
neighbor_list (NeighborListData | None)

Return type:

class ufp.core.input.UFPInputState(atomic_charges=None, atomic_spin_moments=None, system_charges=None, system_spin_moments=None)[source]¶

Bases: object

Optional local and system-level state variables aligned with UFPInput.

The first supported state fields are local charge, collinear scalar magnetic moment, total system charge, and total system spin moment. Missing fields stay as None so legacy models and datasets keep their existing behavior.

Parameters:

atomic_charges (ndarray | Tensor | None)
atomic_spin_moments (ndarray | Tensor | None)
system_charges (ndarray | Tensor | None)
system_spin_moments (ndarray | Tensor | None)

as_torch(*, n_atoms, n_systems, dtype, device)[source]¶

Return a copy normalized to the requested input shape and tensor options.

Parameters:

n_atoms (int)
n_systems (int)
dtype (dtype)
device (device)

Return type:

UFPInputState

classmethod from_ase_list(atoms_list, *, dtype, device=None)[source]¶

Extract supported charge and spin arrays from ASE structures.

Parameters:

atoms_list (Sequence[Atoms])
dtype (dtype)
device (device | None)

Return type:

UFPInputState

missing_fields(fields)[source]¶

Return requested state field names whose values are absent.

Parameters:: fields (Sequence[str])
Return type:: tuple[str, …]

pin_memory()[source]¶

Pin any torch-backed state tensors in place and return self.

Return type:: UFPInputState

to(*, n_atoms, n_systems, dtype, device)[source]¶

Move state tensors to the requested device and dtype.

Parameters:

n_atoms (int)
n_systems (int)
dtype (dtype)
device (device)

Return type:

UFPInputState

Normalized output structures for UFP predictions.

This module defines the per-system and batched result containers used across calculators, training, and adapter integrations.

class ufp.core.output.UFPBatchOutput(energy=None, forces=None, stress=None, per_atom_energy=None, features=<factory>, atom_mask=None, system_sizes=None)[source]¶

Bases: object

Batch-oriented output for multiple ASE or metatomic systems.

Parameters:

energy (ndarray | Tensor | None)
forces (ndarray | Tensor | None)
stress (ndarray | Tensor | None)
per_atom_energy (ndarray | Tensor | None)
features (Dict[str, ndarray | Tensor])
atom_mask (ndarray | Tensor | None)
system_sizes (ndarray | Tensor | None)

energy¶

Per-system total energies.

Type:: numpy.ndarray | torch.Tensor | None

forces¶

Padded per-system forces.

Type:: numpy.ndarray | torch.Tensor | None

stress¶

Per-system stress tensors.

Type:: numpy.ndarray | torch.Tensor | None

per_atom_energy¶

Padded per-system per-atom energies.

Type:: numpy.ndarray | torch.Tensor | None

features¶

Named feature tensors, padded when they are atomwise.

Type:: Dict[str, numpy.ndarray | torch.Tensor]

atom_mask¶

Boolean mask identifying real atoms in padded atomwise outputs.

Type:: numpy.ndarray | torch.Tensor | None

system_sizes¶

Number of atoms in each system.

Type:: numpy.ndarray | torch.Tensor | None

classmethod from_output(output, inputs)[source]¶

Convert a flat UFPOutput into the padded batch representation.

Parameters:

output (UFPOutput)
inputs (UFPInput)

Return type:

UFPBatchOutput

class ufp.core.output.UFPOutput(energy=None, forces=None, stress=None, per_atom_energy=None, features=<factory>)[source]¶

Bases: object

Standard output object returned by UFPotential.

Parameters:

energy (ndarray | Tensor | None)
forces (ndarray | Tensor | None)
stress (ndarray | Tensor | None)
per_atom_energy (ndarray | Tensor | None)
features (Dict[str, ndarray | Tensor])

energy¶

Total energies with shape (n_systems,) or (n_systems, 1).

Type:: numpy.ndarray | torch.Tensor | None

forces¶

Atomic forces with shape (n_atoms, 3).

Type:: numpy.ndarray | torch.Tensor | None

stress¶

Stress with shape (n_systems, 3, 3), (3, 3), or Voigt (6,).

Type:: numpy.ndarray | torch.Tensor | None

per_atom_energy¶

Per-atom energies with shape (n_atoms,) or (n_atoms, 1).

Type:: numpy.ndarray | torch.Tensor | None

features¶

Optional named feature tensors.

Type:: Dict[str, numpy.ndarray | torch.Tensor]

as_ase_results(n_atoms=None)[source]¶

Convert normalized outputs into ASE calculator results.

Parameters:: n_atoms (int | None)
Return type:: Dict[str, object]

ufp.core.output.sum_outputs(outputs, inputs)[source]¶

Sum term outputs into one normalized result for the same input.

Parameters:

outputs (Sequence[UFPOutput])
inputs (UFPInput)

Return type:

Base potential interfaces and batch wrappers for UFP models.

Subclass UFPotential for model logic; wrap it in BatchedUFPotential when callers need padded batch-shaped atomwise outputs.

class ufp.core.potential.BatchedUFPotential(potential, *, independence_tolerance=1e-12)[source]¶

Bases: UFPotential

Batch-oriented wrapper around UFPotential.

UFPotential already supports evaluating a list of structures internally, but it returns flat per-atom arrays over the concatenated atoms. This wrapper keeps the same model logic and exposes a dedicated batch API returning per-system energies and padded per-atom outputs with an explicit batch axis.

When forces are derived from differentiable energies, this wrapper differentiates each system energy separately and checks that no gradient leaks onto atoms from any other system. That makes the concatenated-neighbor-list batching strategy explicit and safe for multi-structure evaluation.

Parameters:

potential (UFPotential) – Wrapped UFP potential.
independence_tolerance (float) – Tolerance used when checking that one system’s energy does not depend on another system’s atoms.

compute_batch(atoms, neighbor_list=None, backend=None, *, device=None, dtype=None, derive_forces=False, neighbor_list_cache=None, state=None, extract_state=True)[source]¶

Prepare batched ASE input, then delegate to compute_batch_input.

Parameters:

atoms (Atoms | Sequence[Atoms])
neighbor_list (NeighborListData | Sequence[NeighborListData] | None)
backend (str | NeighborListBackend | None)
device (device | None)
dtype (dtype | None)
derive_forces (bool)
neighbor_list_cache (GeometryNeighborListCache | None)
state (UFPInputState | None)
extract_state (bool)

Return type:

UFPBatchOutput

compute_batch_input(inputs, *, derive_forces=False)[source]¶

Run the wrapped model and repack the result into UFPBatchOutput.

Parameters:

inputs (UFPInput)
derive_forces (bool)

Return type:

UFPBatchOutput

compute_input(inputs, *, derive_forces=False)[source]¶

Differentiate each system energy independently for forces.

Parameters:

inputs (UFPInput)
derive_forces (bool)

Return type:

forward(inputs)[source]¶

Delegate the flat model evaluation to the wrapped potential.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

provides_forces()[source]¶

Mirror force availability from the wrapped potential.

Return type:: bool

class ufp.core.potential.GeometryNeighborListCache(max_size=128, min_atoms=8)[source]¶

Bases: object

Caller-owned LRU cache for ASE neighbor lists keyed by exact geometry.

The key includes positions, cell, periodic flags, atomic numbers, cutoff, backend, list options, dtype, and device assumptions. The cache stores the normalized NeighborListData returned by the builder and leaves tensor dtype/device coercion to UFPInput.

Parameters:

max_size (int)
min_atoms (int)

clear()[source]¶

Remove all cached neighbor lists.

Return type:: None

get_or_build(*, atoms, cutoff, backend, arrays='torch', full_list=True, sorted=True, dtype, device=None)[source]¶

Return a cached neighbor list or build and store one.

Parameters:

atoms (Atoms)
cutoff (float)
backend (NeighborListBackend)
arrays (str)
full_list (bool)
sorted (bool)
dtype (dtype)
device (device | None)

Return type:

class ufp.core.potential.UFPotential(cutoff=None, neighbor_backend=NeighborListBackend.AUTO)[source]¶

Bases: Module, ABC

Base class for UFP interatomic potentials.

Subclasses receive a torch-native UFPInput bundle and should return UFPOutput. The base class provides ASE conversion, optional neighbor-list construction, output validation, and automatic force derivation from differentiable energies for slow standalone workflows.

Parameters:

cutoff (Optional[float]) – Default cutoff used when building neighbor lists internally.
neighbor_backend (Union[str, NeighborListBackend]) – Default backend used for neighbor-list construction.

compute(atoms, neighbor_list=None, backend=None, *, device=None, dtype=None, derive_forces=False, neighbor_list_cache=None, state=None, extract_state=True)[source]¶

Prepare ASE input first, then delegate to compute_input.

Parameters:

atoms (Atoms | Sequence[Atoms])
neighbor_list (NeighborListData | Sequence[NeighborListData] | None)
backend (str | NeighborListBackend | None)
device (device | None)
dtype (dtype | None)
derive_forces (bool)
neighbor_list_cache (GeometryNeighborListCache | None)
state (UFPInputState | None)
extract_state (bool)

Return type:

compute_input(inputs, *, derive_forces=False)[source]¶

Run forward, validate the result, and derive forces when requested.

Parameters:

inputs (UFPInput)
derive_forces (bool)

Return type:

abstractmethod forward(inputs)[source]¶

Run the potential on one or more systems.

Parameters:: inputs (UFPInput) – Normalized input bundle.
Returns:: Model predictions for the input systems.
Return type:: UFPOutput

preferred_dtype()[source]¶

Return the parameter dtype or the torch default.

Return type:: dtype

prepare_input(atoms, neighbor_list=None, backend=None, *, device=None, dtype=None, requires_grad=False, neighbor_list_cache=None, state=None, extract_state=True)[source]¶

Convert ASE structures into UFPInput before evaluation.

Parameters:

atoms (Atoms | Sequence[Atoms])
neighbor_list (NeighborListData | Sequence[NeighborListData] | None)
backend (str | NeighborListBackend | None)
device (device | None)
dtype (dtype | None)
requires_grad (bool)
neighbor_list_cache (GeometryNeighborListCache | None)
state (UFPInputState | None)
extract_state (bool)

Return type:

DenseThreeBodyFeatureCache

provides_forces()[source]¶

Report whether the subclass returns forces directly.

Return type:: bool

Terms¶

Status: Stable user API for built-in terms and UFPModel; extension API for ufp.terms.contracts; expert diagnostics for documented three-body bucket, feature-cache, and evaluator helpers.

Composable energy terms and model containers.

Terms are small additive torch.nn.Module objects for one-body references, pair interactions, spline two-body and three-body models, higher-body terms, repulsion priors, and alchemical coefficient sharing. Combine them with UFPModel to evaluate complete interatomic potentials.

Composite UFP model built from additive term objects.

Use this module to combine pair and three-body contributions while keeping atomic-type checks and output summation centralized.

class ufp.terms.model.UFPModel(*, terms=None, onebody_terms=(), pair_terms=(), threebody_terms=(), atomic_types=None, neighbor_backend=NeighborListBackend.AUTO)[source]¶

Bases: UFPotential

Concrete UFP model assembled from pair and three-body terms.

Parameters:

terms (Sequence[UFPTerm] | None)
onebody_terms (Sequence[OneBodyTerm])
pair_terms (Sequence[PairTerm])
threebody_terms (Sequence[ThreeBodyTerm])
atomic_types (Optional[Sequence[int]])
neighbor_backend (str | NeighborListBackend)

forward(inputs)[source]¶

Validate atomic types, then sum child term outputs into one model result.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

provides_forces()[source]¶

Report whether every child term provides forces directly.

Return type:: bool

property terms: tuple[UFPTerm, ...]¶: Return all terms in evaluation order.

Public extension contracts for custom UFP terms.

The names in this module are safe for third-party or downstream terms to import when they need to expose fittable parameters, participate in input-cache warming, or provide optional least-squares assembly hooks. Runtime evaluation helpers, kernel dispatchers, and cache implementations remain internal.

class ufp.terms.contracts.LinearAssemblyOptions(blocks=(), threebody_lstsq_backend=None, threebody_bucket_backend=None, threebody_runtime_config=None)[source]¶

Bases: object

Options passed to optional term-level least-squares assemblers.

Parameters:

blocks (tuple[object, ...])
threebody_lstsq_backend (str | None)
threebody_bucket_backend (str | None)
threebody_runtime_config (object | None)

class ufp.terms.contracts.ParameterBlock(name, kind, shape, read, write, label=None, coefficient_provider=None, coefficient_index=None, regularization_group=None, fittable=True, frozen=False, assembler=None, cache_descriptor=None)[source]¶

Bases: object

Named parameter tensor exposed by a term for fitting or inspection.

Blocks are resolved once when a fitting layout is created. Term forward paths do not use this object, which keeps the contract out of runtime hot loops.

Parameters:

name (str)
kind (str)
shape (tuple[int, ...])
read (Callable[[], Tensor])
write (Callable[[Tensor], None])
label (str | None)
coefficient_provider (AlchemicalCoefficients | None)
coefficient_index (int | None)
regularization_group (str | None)
fittable (bool)
frozen (bool)
assembler (Any)
cache_descriptor (ParameterBlockCacheDescriptor | None)

name¶

Local parameter name within the owning term.

Type:: str

kind¶

Term-specific block kind used by layout and solver code.

Type:: str

shape¶

Logical tensor shape for this block.

Type:: tuple[int, …]

read¶

Callback returning the current tensor data.

Type:: collections.abc.Callable[[], torch.Tensor]

write¶

Callback that writes solved values back into the owning term.

Type:: collections.abc.Callable[[torch.Tensor], None]

label¶

Optional human-readable label.

Type:: str | None

coefficient_provider¶

Optional alchemical coefficient provider.

Type:: ufp.terms.alchemical.AlchemicalCoefficients | None

coefficient_index¶

Optional provider index for this block.

Type:: int | None

regularization_group¶

Optional group name used by regularization settings.

Type:: str | None

fittable¶

Whether linear fitting should include this block.

Type:: bool

frozen¶

Whether this block should be treated as fixed.

Type:: bool

assembler¶

Optional custom block assembler.

Type:: Any

cache_descriptor¶

Optional semantic descriptor for reusable assembled caches.

Type:: ufp.terms._parameters.ParameterBlockCacheDescriptor | None

class ufp.terms.contracts.ParameterBlockCacheChannel(kind, values, start, stop, key=None)[source]¶

Bases: object

Semantic channel metadata for reusable least-squares cache layouts.

Parameters:

kind (str)
values (tuple[int, ...])
start (int)
stop (int)
key (str | None)

class ufp.terms.contracts.ParameterBlockCacheDescriptor(family, channels, reusable=True)[source]¶

Bases: object

Semantic cache descriptor attached to a fittable parameter block.

Parameters:

family (Mapping[str, object] | str)
channels (tuple[ParameterBlockCacheChannel, ...])
reusable (bool)

class ufp.terms.contracts.TermCacheOptions(feature_cache_storage='cpu', feature_cache_mode='auto', feature_cache_dir=None, cache_prefix='batch', legacy_cache_prefixes=(), include_per_atom_energy=True)[source]¶

Bases: object

Options passed to optional term-level input cache warmers.

Parameters:

feature_cache_storage (str)
feature_cache_mode (str)
feature_cache_dir (object)
cache_prefix (str)
legacy_cache_prefixes (tuple[str, ...])
include_per_atom_energy (bool)

class ufp.terms.contracts.TermInputRequirements(neighbor_list=False, full_neighbor_list=False, state_fields=())[source]¶

Bases: object

Declarative input requirements for a term’s forward or assembly path.

Parameters:

neighbor_list (bool)
full_neighbor_list (bool)
state_fields (tuple[str, ...])

validate(inputs, *, term_name)[source]¶

Raise clear errors for missing geometry or charge/spin state.

Parameters:

inputs (UFPInput)
term_name (str)

Return type:

None

class ufp.terms.contracts.UFPTerm(*, cutoff=None, atomic_types=None)[source]¶

Bases: Module, ABC

Base class for pair and higher-order interaction terms.

Parameters:

cutoff (Optional[float])
atomic_types (Optional[Sequence[int]])

assemble_linear_blocks(batch, targets, options=None)[source]¶

Optionally assemble this term’s least-squares blocks in one call.

Parameters:: options (LinearAssemblyOptions | None)

cache_input(inputs, options=None, **kwargs)[source]¶

Optionally precompute reusable metadata for one input.

Parameters:

inputs (UFPInput)
options (TermCacheOptions | None)

Return type:

None

abstractmethod forward(inputs)[source]¶

Compute this term’s contribution to the total model output.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

property input_requirements: TermInputRequirements¶: Return optional input requirements declared by this term.

property optimizer_group: str | None¶: Return an optional workflow optimizer group name for this term.

parameter_blocks()[source]¶

Return fittable parameter blocks exposed by this term.

Return type:: tuple[ParameterBlock, …]

property provides_forces: bool¶: Report whether the term returns forces analytically.

validate_inputs(inputs)[source]¶

Validate inputs against this term’s declared requirements.

Parameters:: inputs (UFPInput)
Return type:: None

ufp.terms.contracts.copy_parameter_data(parameter, values)[source]¶

Copy reshaped values into a parameter without replacing the object.

Parameters:

parameter (Tensor)
values (Tensor)

Return type:

None

Reusable cutoff envelope helpers for analytic pair priors.

class ufp.terms.cutoffs.CutoffEnvelope(cutoff, onset=0.0, kind='smoothstep')[source]¶

Bases: object

Smooth multiplicative envelope that tapers pair energies to zero.

Parameters:

cutoff (float)
onset (float)
kind (str)

derivatives(distances)[source]¶

Evaluate derivatives d envelope / d distance.

Parameters:: distances (object)
Return type:: Tensor

to_dict()[source]¶

Return metadata suitable for projection diagnostics or checkpoints.

Return type:: dict[str, float | str]

values(distances)[source]¶

Evaluate envelope values at pair distances.

Parameters:: distances (object)
Return type:: Tensor

ufp.terms.cutoffs.CutoffEnvelopeKind¶: alias of str

ufp.terms.cutoffs.apply_cutoff_envelope(values, distances, envelope)[source]¶

Multiply values by a cutoff envelope evaluated at the same distances.

Parameters:

values (Tensor)
distances (object)
envelope (CutoffEnvelope)

Return type:

Tensor

ufp.terms.cutoffs.cutoff_envelope_derivatives(distances, *, cutoff, onset=0.0, kind='smoothstep')[source]¶

Evaluate derivative of a supported cutoff envelope.

Parameters:

distances (object)
cutoff (float)
onset (float)
kind (str)

Return type:

Tensor

ufp.terms.cutoffs.cutoff_envelope_values(distances, *, cutoff, onset=0.0, kind='smoothstep')[source]¶

Evaluate a supported cutoff envelope.

Parameters:

distances (object)
cutoff (float)
onset (float)
kind (str)

Return type:

Tensor

ufp.terms.cutoffs.normalize_cutoff_envelope(envelope, *, cutoff, default_kind='smoothstep', default_onset_fraction=0.8)[source]¶

Normalize a cutoff-envelope specification.

Parameters:

envelope (CutoffEnvelope | str | None)
cutoff (float)
default_kind (str)
default_onset_fraction (float)

Return type:

CutoffEnvelope

Autograd-friendly analytic pair prior terms.

class ufp.terms.analytic.DampedInversePowerPairPrior(*, cutoff, atomic_types, prefactors_by_pair, damping_rates_by_pair, power=6.0, active_pairs=None, symmetric=True, eps=1e-12, cutoff_envelope=None, trainable=True, dtype=None)[source]¶

Bases: _CategorizedAnalyticPairTerm

Inverse-power prior with a short-range exponential damping factor.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
power (float)
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
eps (float)
cutoff_envelope (CutoffEnvelope | str | None)
trainable (bool)
dtype (torch.dtype | None)

class ufp.terms.analytic.ExponentialRepulsionPairPrior(*, cutoff, atomic_types, amplitudes_by_pair, decay_rates_by_pair, active_pairs=None, symmetric=True, cutoff_envelope=None, trainable=True, dtype=None)[source]¶

Bases: _CategorizedAnalyticPairTerm

Pair-dependent exponential repulsion prior.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
cutoff_envelope (CutoffEnvelope | str | None)
trainable (bool)
dtype (torch.dtype | None)

class ufp.terms.analytic.InversePowerPairPrior(*, cutoff, atomic_types, prefactors_by_pair, power=2.0, active_pairs=None, symmetric=True, eps=1e-12, cutoff_envelope=None, trainable=True, dtype=None)[source]¶

Bases: _CategorizedAnalyticPairTerm

Pair-dependent inverse-power analytic prior.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
power (float)
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
eps (float)
cutoff_envelope (CutoffEnvelope | str | None)
trainable (bool)
dtype (torch.dtype | None)

class ufp.terms.analytic.MorsePairPrior(*, cutoff, atomic_types, depths_by_pair, equilibrium_distances_by_pair, widths_by_pair, active_pairs=None, symmetric=True, cutoff_envelope=None, trainable=True, dtype=None)[source]¶

Bases: _CategorizedAnalyticPairTerm

Pair-dependent Morse-like analytic interaction prior.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
cutoff_envelope (CutoffEnvelope | str | None)
trainable (bool)
dtype (torch.dtype | None)

class ufp.terms.analytic.ScaledZBLPairPrior(*, cutoff, atomic_types, scales_by_pair, active_pairs=None, symmetric=True, eps=1e-12, cutoff_envelope=None, trainable=True, dtype=None)[source]¶

Bases: _CategorizedAnalyticPairTerm

ZBL screened nuclear repulsion with trainable pair-channel scales.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
eps (float)
cutoff_envelope (CutoffEnvelope | str | None)
trainable (bool)
dtype (torch.dtype | None)

Element-wise one-body energy term.

Use this module when the model needs fixed or trainable reference energies per atomic species in addition to geometric interaction terms.

class ufp.terms.onebody.ElementOneBodyTerm(*, atomic_types, values=None, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: OneBodyTerm

Per-element one-body energy contribution.

Parameters:

atomic_types (Sequence[int])
trainable (bool)
fittable (bool)
frozen (bool)
dtype (torch.dtype | None)

assemble_linear_blocks(batch, targets, options=None)[source]¶

Assemble one-body least-squares blocks for this term.

Parameters:: options (LinearAssemblyOptions | None)

forward(inputs)[source]¶

Assign learned reference energies and sum by system.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

parameter_blocks()[source]¶

Return the element-reference coefficient block.

Return type:: tuple[ParameterBlock, …]

property provides_forces: bool¶: Report that this term provides the expected zero-force contribution.

class ufp.terms.onebody.OneBodyTerm(*, cutoff=None, atomic_types=None)[source]¶

Bases: UFPTerm

Base class for one-body contribution terms.

Parameters:

cutoff (Optional[float])
atomic_types (Optional[Sequence[int]])

Charge and collinear-spin state terms for UFP models.

class ufp.terms.state.ChargeScaledSplinePairTerm(*, cutoff, atomic_types, coeffs_by_pair=None, coefficient_provider=None, coefficient_index=None, active_pairs=None, symmetric=True, spline='cubic', full_support_start=0.0, eps=1e-12, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: _StateScaledSplinePairTerm

Short-range spline pair correction scaled by fixed local charges.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
coefficient_provider (AlchemicalCoefficients | None)
coefficient_index (int | None)
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
spline (SplineKind)
full_support_start (float)
eps (float)
trainable (bool)
fittable (bool)
frozen (bool)
dtype (torch.dtype | None)

class ufp.terms.state.ChargeSelfEnergyTerm(*, atomic_types, electronegativities=None, hardnesses=None, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: OneBodyTerm

Per-element local charge electronegativity and hardness energy.

Parameters:

atomic_types (Sequence[int])
trainable (bool)
fittable (bool)
frozen (bool)
dtype (torch.dtype | None)

assemble_linear_blocks(batch, targets, options=None)[source]¶

Assemble all requested charge self-energy blocks.

Parameters:: options (LinearAssemblyOptions | None)

forward(inputs)[source]¶

Evaluate local charge self energy and charge potential.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

property input_requirements: TermInputRequirements¶: Require fixed local charge state.

property optimizer_group: str | None¶: Group trainable state-term parameters for workflow optimizers.

parameter_blocks()[source]¶

Return linear charge self-energy coefficient blocks.

Return type:: tuple[ParameterBlock, …]

property provides_forces: bool¶: Report that this term contributes explicit zero forces.

class ufp.terms.state.CollinearSpinExchangeTerm(*, cutoff, atomic_types, coeffs_by_pair=None, coefficient_provider=None, coefficient_index=None, active_pairs=None, symmetric=True, spline='cubic', full_support_start=0.0, eps=1e-12, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: _StateScaledSplinePairTerm

Pairwise collinear exchange spline scaled by fixed spin moments.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
coefficient_provider (AlchemicalCoefficients | None)
coefficient_index (int | None)
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
spline (SplineKind)
full_support_start (float)
eps (float)
trainable (bool)
fittable (bool)
frozen (bool)
dtype (torch.dtype | None)

class ufp.terms.state.CollinearSpinLandauTerm(*, atomic_types, quadratic=None, quartic=None, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: OneBodyTerm

Per-element onsite Landau energy for fixed scalar spin moments.

Parameters:

atomic_types (Sequence[int])
trainable (bool)
fittable (bool)
frozen (bool)
dtype (torch.dtype | None)

assemble_linear_blocks(batch, targets, options=None)[source]¶

Assemble all requested Landau blocks.

Parameters:: options (LinearAssemblyOptions | None)

forward(inputs)[source]¶

Evaluate onsite spin energy and effective field.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

property input_requirements: TermInputRequirements¶: Require fixed local collinear spin moments.

property optimizer_group: str | None¶: Group trainable state-term parameters for workflow optimizers.

parameter_blocks()[source]¶

Return linear spin Landau coefficient blocks.

Return type:: tuple[ParameterBlock, …]

property provides_forces: bool¶: Report that this term contributes explicit zero forces.

class ufp.terms.state.LocalChargeCoulombTerm(*, cutoff, atomic_types, active_pairs=None, symmetric=True, softening=1e-06, scale=1.0, cutoff_envelope=None, eps=1e-12)[source]¶

Bases: PairTerm

Finite-cutoff softened Coulomb interaction for fixed local charges.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
softening (float)
scale (float)
cutoff_envelope (CutoffEnvelope | str | None)
eps (float)

property active_pair_categories: tuple[tuple[int, int], ...]¶: Return active pair categories.

forward(inputs)[source]¶

Evaluate softened Coulomb energy, forces, and charge potential.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

property input_requirements: TermInputRequirements¶: Require a neighbor list and fixed local charge state.

property optimizer_group: str | None¶: Return the common charge/spin optimizer group name.

property pair_categories: tuple[tuple[int, int], ...]¶: Return configured pair categories.

property provides_forces: bool¶: Report that this term provides analytic pair forces.

Spline-based two-body interaction term implementation.

Use this module when pair energies and forces should come from 1D spline coefficients, optionally shared through alchemical providers.

class ufp.terms.twobody.RepulsiveSplineTwoBodyTerm(*, cutoff, atomic_types, coeff_size, transition_span, active_pairs=None, symmetric=True, spline='cubic', full_support_start=0.0, eps=1e-12, trainable=True, outer_coeffs_by_pair=None, curvature_floors_by_pair=None, initial_curvatures_by_pair=None, transition_force_floors_by_pair=None, initial_transition_forces_by_pair=None, initial_transition_values_by_pair=None, dtype=None)[source]¶

Bases: PairTerm

Cubic two-body spline family with a constrained repulsive inner wall.

The term owns one generated cubic spline coefficient row per pair category. Low-distance coefficients are solved from positive curvature coefficients, a positive transition force, and a trainable transition value. The generated rows are evaluated through the same uniform cubic stencil path as SplineTwoBodyTerm.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
coeff_size (int)
transition_span (int)
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
spline (SplineKind)
full_support_start (float)
eps (float)
trainable (bool)
dtype (torch.dtype | None)

property active_pair_categories: tuple[tuple[int, int], ...]¶: Return the subset of pair categories enabled for evaluation.

canonical_pair(first, second)[source]¶

Normalize a pair key using this term’s symmetry convention.

Parameters:

first (int)
second (int)

Return type:

tuple[int, int]

property curvature_coeffs_by_pair: Tensor¶: Return positive inner curvature coefficients by pair category.

forward(inputs)[source]¶

Accumulate generated spline blocks for all active pair categories.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

is_pair_active(first, second)[source]¶

Report whether a canonical pair category is enabled.

Parameters:

first (int)
second (int)

Return type:

bool

property pair_categories: tuple[tuple[int, int], ...]¶: Return the full ordered list of pair categories owned by this term.

pair_category_index(first, second)[source]¶

Return the block index for one canonical pair category.

Parameters:

first (int)
second (int)

Return type:

int

property provides_forces: bool¶: Report that this term produces analytic forces directly.

to_spline_twobody_term(*, trainable=True, fittable=True)[source]¶

Return an ordinary spline term with the current generated coefficients.

Parameters:

trainable (bool)
fittable (bool)

Return type:

SplineTwoBodyTerm

property transition_forces_by_pair: Tensor¶: Return positive transition force magnitudes by pair category.

property true_coeffs_by_pair: Tensor¶: Return generated cubic spline coefficients by pair category.

wall_diagnostics(*, n_samples=200)[source]¶

Sample value, gradient, and curvature on the constrained interval.

Parameters:: n_samples (int)
Return type:: dict[tuple[int, int], object]

class ufp.terms.twobody.SplinePairTerm(*, cutoff, pair, coeffs=None, coefficient_provider=None, coefficient_index=None, symmetric=True, spline='cubic', full_support_start=0.0, eps=1e-12, enabled=True, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: PairTerm

One-dimensional spline two-body interaction term for a specific element pair.

The term evaluates a uniformly spaced spline over pair distances in the same length unit as the input coordinates, normally angstroms for ASE structures. Coefficients store interaction energies, normally in electron volts, and the analytic force contribution is derived from the spline gradient.

Direct coefficients are supplied through coeffs. For alchemical fitting, pass coefficient_provider and coefficient_index instead; the provider supplies the true coefficient vector used during evaluation.

Parameters:

cutoff (float) – Maximum pair distance included in the interaction.
pair (tuple[int, int]) – Two atomic numbers identifying the element pair handled by this term.
coeffs – One-dimensional coefficient vector with shape (n_knots,).
coefficient_provider (AlchemicalCoefficients | None) – Optional shared provider that projects proxy coefficients into true pair-specific coefficients.
coefficient_index (int | None) – Index into coefficient_provider for this pair channel.
symmetric (bool) – If True, treat (a, b) and (b, a) as the same pair category.
spline (SplineKind) – Spline family name. Supported values are "quadratic", "cubic", and "quartic".
full_support_start (float) – Lower distance where the spline has full support. The upper full-support boundary is cutoff.
eps (float) – Small distance threshold used to avoid division by zero when forming force directions.
enabled (bool) – If False, keep the term parameters but skip evaluation.
trainable (bool) – If True, direct coeffs are stored as trainable parameters.
fittable (bool) – Whether this term exposes a coefficient block to linear fitters.
frozen (bool) – Whether this term’s parameter block is fixed during fitting.
dtype (torch.dtype | None) – Optional dtype used when converting direct coefficients to a tensor.

Examples

>>> import torch
>>> term = SplinePairTerm(
...     cutoff=2.0,
...     pair=(1, 1),
...     coeffs=torch.zeros(6),
... )
>>> term.covers_pair(1, 1)
True

assemble_linear_blocks(batch, targets, options=None)[source]¶

Assemble pair least-squares blocks for this term.

Parameters:: options (LinearAssemblyOptions | None)

assemble_selected_linear_block(block, inputs, targets, selected_indices)[source]¶

Assemble only requested coefficient columns for this pair block.

Parameters:

inputs (UFPInput)
selected_indices (Sequence[int])

Return type:

Tensor | None

property coefficient_provider: AlchemicalCoefficients | None¶: Return the shared coefficient provider, if any.

covered_pair_mask(inputs)[source]¶

Select neighbor-list rows handled by this pair-specific spline block.

Parameters:: inputs (UFPInput)
Return type:: Tensor

covers_pair(first_atomic_number, second_atomic_number)[source]¶

Report whether this term handles the requested pair.

Parameters:

first_atomic_number (int)
second_atomic_number (int)

Return type:

bool

forward(inputs)[source]¶

Evaluate this spline block over covered pairs.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

parameter_blocks()[source]¶

Return the pair spline coefficient block.

Return type:: tuple[ParameterBlock, …]

property provides_forces: bool¶: Report that this term produces analytic forces directly.

property true_coeffs: Tensor¶: Return direct or provider-projected true coefficients.

class ufp.terms.twobody.SplineTwoBodyTerm(*, cutoff, atomic_types, coeffs_by_pair=None, coefficient_provider=None, coefficient_index=None, active_pairs=None, symmetric=True, spline='cubic', full_support_start=0.0, eps=1e-12, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: PairTerm

One-dimensional spline family covering all configured chemical pair channels.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
coefficient_provider (AlchemicalCoefficients | None)
coefficient_index (int | None)
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
spline (SplineKind)
full_support_start (float)
eps (float)
trainable (bool)
fittable (bool)
frozen (bool)
dtype (torch.dtype | None)

property active_pair_categories: tuple[tuple[int, int], ...]¶: Return the subset of pair categories still enabled for evaluation.

assemble_linear_blocks(batch, targets, options=None)[source]¶

Assemble categorized pair least-squares blocks for this term.

Parameters:: options (LinearAssemblyOptions | None)

assemble_selected_linear_block(block, inputs, targets, selected_indices)[source]¶

Assemble only requested coefficient columns for this two-body block.

Parameters:

inputs (UFPInput)
selected_indices (Sequence[int])

Return type:

Tensor | None

canonical_pair(first, second)[source]¶

Normalize a pair key using this term’s symmetry convention.

Parameters:

first (int)
second (int)

Return type:

tuple[int, int]

property coefficient_provider: AlchemicalCoefficients | None¶: Return the shared coefficient provider, when one supplies the pair blocks.

forward(inputs)[source]¶

Accumulate spline blocks for all active pair categories.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

is_pair_active(first, second)[source]¶

Report whether a canonical pair category remains enabled.

Parameters:

first (int)
second (int)

Return type:

bool

property pair_categories: tuple[tuple[int, int], ...]¶: Return the full ordered list of pair categories owned by this term.

pair_category_index(first, second)[source]¶

Return the block index for one canonical pair category.

Parameters:

first (int)
second (int)

Return type:

int

parameter_blocks()[source]¶

Return the categorized pair spline coefficient block.

Return type:: tuple[ParameterBlock, …]

property provides_forces: bool¶: Report that this term produces analytic forces directly.

property true_coeffs_by_pair: Tensor¶: Return coefficients indexed by pair category.

ufp.terms.twobody.get_eval_1d_with_grads(spline)[source]¶

Return the two-body evaluator for a spline family.

Parameters:: spline (Literal['quadratic', 'cubic', 'quartic'] | str)
Return type:: Callable[[float, Tensor, Tensor], tuple[Tensor, Tensor]]

Spline-based three-body interaction term implementation.

SplineThreeBodyTerm is the stable user-facing term. Bucket containers, feature caches, and evaluator helpers exported from this module are expert diagnostics for benchmarks, tests, and performance investigations.

class ufp.terms.threebody.Buckets(src_ids, row_ptr, nbr_ids, pair_vectors, pair_distances, patterns, pattern_ptr, pattern_plans=<factory>, tensor_pattern_plans=None)[source]¶

Bases: object

Flattened CSR-like neighborhood batches grouped by shared source patterns.

Parameters:

src_ids (Tensor)
row_ptr (Tensor)
nbr_ids (Tensor)
pair_vectors (Tensor)
pair_distances (Tensor)
patterns (Tensor)
pattern_ptr (Tensor)
pattern_plans (tuple[BucketPatternPlan, ...])
tensor_pattern_plans (TensorBucketPatternPlans | None)

to_input_device(*, device, dtype)[source]¶

Return a copy on the requested input device and floating dtype.

Parameters:

device (device)
dtype (dtype)

Return type:

Buckets

with_pattern_plans(device=None)[source]¶

Return buckets with reusable CPU-side pattern metadata attached.

Parameters:: device (device | None)
Return type:: Buckets

class ufp.terms.threebody.DenseThreeBodyFeatureCache(blocks)[source]¶

Bases: object

Coalesced dense output features for one fixed input.

Parameters:: blocks (tuple[DenseTripletFeatureBlock, ...])

to_input_device(*, device, dtype)[source]¶

Return this dense feature cache on an input device and dtype.

Parameters:

device (device)
dtype (dtype)

Return type:

class ufp.terms.threebody.DenseTripletFeatureBlock(triplet_index, coeff_start, coeff_shape, energy, per_atom_energy, per_atom_indices, forces, force_atom_indices)[source]¶

Bases: object

Coalesced dense output features for one triplet category.

Parameters:

triplet_index (int)
coeff_start (tuple[int, int, int])
coeff_shape (tuple[int, int, int])
energy (Tensor)
per_atom_energy (Tensor | None)
per_atom_indices (Tensor | None)
forces (Tensor)
force_atom_indices (Tensor | None)

to_input_device(*, device, dtype)[source]¶

Return this dense feature block on an input device and dtype.

Parameters:

device (device)
dtype (dtype)

Return type:

DenseTripletFeatureBlock

class ufp.terms.threebody.MemmapDenseThreeBodyFeatureCache(blocks)[source]¶

Bases: object

Disk-backed dense output feature blocks for one fixed input.

Parameters:: blocks (tuple[MemmapDenseTripletFeatureBlock, ...])

to_input_device(*, device, dtype)[source]¶

Materialize this disk-backed dense cache on an input device.

Parameters:

device (device)
dtype (dtype)

Return type:

DenseThreeBodyFeatureCache

class ufp.terms.threebody.MemmapDenseTripletFeatureBlock(triplet_index, coeff_start, coeff_shape, energy, per_atom_energy, per_atom_indices, forces, force_atom_indices)[source]¶

Bases: object

Disk-backed dense output features for one triplet category.

Parameters:

triplet_index (int)
coeff_start (tuple[int, int, int])
coeff_shape (tuple[int, int, int])
energy (_MemmapTensor)
per_atom_energy (_MemmapTensor | None)
per_atom_indices (_MemmapTensor | None)
forces (_MemmapTensor)
force_atom_indices (_MemmapTensor | None)

to_input_device(*, device, dtype)[source]¶

Materialize this disk-backed dense block on an input device.

Parameters:

device (device)
dtype (dtype)

Return type:

DenseTripletFeatureBlock

class ufp.terms.threebody.SplineThreeBodyTerm(*, cutoff, atomic_types, coeffs_by_triplet=None, coefficient_provider=None, coefficient_index=None, active_triplets=None, spline='cubic', full_support_start_xy=0.0, full_support_start_z=2.0, eps=1e-12, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: ThreeBodyTerm

Source-distinguished three-body spline interaction term.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
coefficient_provider (AlchemicalCoefficients | None)
coefficient_index (int | None)
active_triplets (Sequence[tuple[int, int, int]] | None)
spline (SplineKind)
full_support_start_xy (float)
full_support_start_z (float)
eps (float)
trainable (bool)
fittable (bool)
frozen (bool)
dtype (torch.dtype | None)

property active_triplet_categories: tuple[tuple[int, int, int], ...]¶: Return the subset of triplet categories that remain enabled.

assemble_linear_blocks(batch, targets, options=None)[source]¶

Assemble three-body least-squares blocks for this term.

Parameters:: options (LinearAssemblyOptions | None)

assemble_selected_linear_block(block, inputs, targets, selected_indices)[source]¶

Assemble only requested coefficient columns for this three-body block.

Parameters:

inputs (UFPInput)
selected_indices (Sequence[int])

Return type:

Tensor | None

cache_input(inputs, options=None, *, feature_cache_storage='cpu', feature_cache_mode='auto', feature_cache_dir=None, cache_prefix='threebody', legacy_cache_prefixes=(), include_per_atom_energy=True)[source]¶

Precompute static dense three-body feature blocks for a cached input.

Parameters:

inputs (UFPInput)
options (TermCacheOptions | None)
feature_cache_storage (Literal['none', 'cpu', 'disk'])
feature_cache_mode (Literal['auto', 'read', 'refresh'])
feature_cache_dir (Path | str | None)
cache_prefix (str)
legacy_cache_prefixes (Sequence[str])
include_per_atom_energy (bool)

Return type:

None

canonical_triplet(source, first_neighbor, second_neighbor)[source]¶

Normalize a triplet key using the term’s neighbor-ordering convention.

Parameters:

source (int)
first_neighbor (int)
second_neighbor (int)

Return type:

tuple[int, int, int]

property coefficient_provider: AlchemicalCoefficients | None¶: Return the shared coefficient provider for alchemical fitting.

dense_atom_features(inputs, atom_indices=None, *, force_scope='output', runtime_config=None)[source]¶

Return dense coefficient-space output rows for selected atoms.

This is intended for debugging fixed-geometry feature construction. The returned rows are the dense equivalent of the sparse cached operators used for this term: one per-atom energy row and one force row for each Cartesian component. With force_scope="output", force rows are the full model-output rows for the selected atoms. With force_scope="source", force rows include only interactions centered on the selected atoms.

Parameters:

inputs (UFPInput) – Normalized input bundle with a full neighbor list.
atom_indices (Sequence[int] | Tensor | None) – Optional atom indices to extract. If omitted, all atoms are returned in input order.
force_scope (Literal['output', 'source']) – Whether force rows should include all output contributions or only source-centered contributions.
runtime_config (ThreeBodyRuntimeConfig | None)

Returns:

Dense per-atom energy and force-component feature rows.

Raises:

RuntimeError – If the input lacks a full neighbor list.
ValueError – If force_scope is not "output" or "source".

Return type:

ThreeBodyDenseAtomFeatures

forward(inputs)[source]¶

Build local triplet buckets and return one three-body contribution.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

property input_requirements: TermInputRequirements¶: Declare the directed full-neighbor-list requirement.

is_triplet_active(source, first_neighbor, second_neighbor)[source]¶

Report whether a canonical triplet category remains enabled.

Parameters:

source (int)
first_neighbor (int)
second_neighbor (int)

Return type:

bool

property n_categories: int¶: Return the number of atomic categories tracked by this term.

property neighbor_neighbor_cutoff: float¶: Return the derived cutoff used when forming neighbor-neighbor triplets.

parameter_blocks()[source]¶

Return the three-body spline coefficient block.

Return type:: tuple[ParameterBlock, …]

property provides_forces: bool¶: Report that this term produces analytic forces directly.

property triplet_categories: tuple[tuple[int, int, int], ...]¶: Return the ordered triplet categories addressed by the coefficient tensor.

triplet_category_index(source, first_neighbor, second_neighbor)[source]¶

Return the coefficient-block index for one canonical triplet.

Parameters:

source (int)
first_neighbor (int)
second_neighbor (int)

Return type:

int

property true_coeffs_by_triplet: Tensor¶: Return the direct or provider-projected triplet coefficient tensor.

class ufp.terms.threebody.ThreeBodyDenseAtomFeatures(atom_indices, atomic_energy, force_x, force_y, force_z)[source]¶

Bases: object

Dense coefficient-space output rows for selected atoms.

Parameters:

atom_indices (Tensor)
atomic_energy (Tensor)
force_x (Tensor)
force_y (Tensor)
force_z (Tensor)

property forces: Tensor¶: Return force-component rows stacked as (n_atoms, 3, n_features).

class ufp.terms.threebody.ThreeBodyTerm(*, cutoff=None, atomic_types=None)[source]¶

Bases: UFPTerm

Base class for three-body interaction terms.

Parameters:

cutoff (Optional[float])
atomic_types (Optional[Sequence[int]])

ufp.terms.threebody.build_edge_category_table(n_cat, device=None)[source]¶

Map ordered category pairs to unordered edge-category indices.

Parameters:: n_cat (int)
Return type:: Tensor

ufp.terms.threebody.evaluate_bucketed_energy_forces(buckets, node_cat, coeffs_by_triplet, edge_cat_table, eval_3d_with_grads=None, *, spline='cubic', active_triplet_mask=None, n_nodes, n_cat=10, first_knot_xy=0.0, first_knot_z=0.0, knot_spacing_xy=0.25, knot_spacing_z=0.25, lower_support_xy=0.0, lower_support_z=0.0, eps=1e-12, runtime_config=None)[source]¶

Public wrapper around the bucket evaluator used by three-body forward paths.

Parameters:

buckets (Buckets)
node_cat (Tensor)
coeffs_by_triplet (Tensor)
edge_cat_table (Tensor)
eval_3d_with_grads (Callable[[float, float, float, float, Tensor, Tensor, Tensor, Tensor], tuple[Tensor, Tensor, Tensor, Tensor]] | None)
spline (Literal['quadratic', 'cubic', 'quartic'])
active_triplet_mask (Tensor | None)
n_nodes (int)
n_cat (int)
first_knot_xy (float)
first_knot_z (float)
knot_spacing_xy (float)
knot_spacing_z (float)
lower_support_xy (float)
lower_support_z (float)
eps (float)
runtime_config (ThreeBodyRuntimeConfig | None)

Return type:

tuple[Tensor, Tensor]

ufp.terms.threebody.get_eval_3d_with_grads(spline)[source]¶

Return the 3D spline evaluator for the requested basis family.

Parameters:: spline (Literal['quadratic', 'cubic', 'quartic'] | str)
Return type:: Callable[[float, float, float, float, Tensor, Tensor, Tensor, Tensor], tuple[Tensor, Tensor, Tensor, Tensor]]

ufp.terms.threebody.load_memmap_threebody_feature_cache(directory, *, prefix=None, mmap_mode='r', copy=False)[source]¶

Load disk-backed dense three-body feature blocks for debugging.

The returned dictionary is grouped as cache[prefix]. Each value is a list of block dictionaries with triplet_index, coeff_start, coeff_shape, energy, optional per_atom_energy, optional atom-index arrays, and forces entries.

Parameters:

directory (Path | str) – Directory containing dense feature block manifests and .npy arrays.
prefix (str | None) – Optional cache prefix to filter on.
mmap_mode (Literal['r', 'r+', 'c'] | None) – Mode passed to numpy.load. Use None to load regular arrays directly.
copy (bool) – Whether to copy loaded arrays into normal numpy.ndarray instances.

Returns:

Dense feature blocks grouped by cache prefix.

Raises:

ValueError – If directory is not an existing directory.

Return type:

dict[str, list[dict[str, object]]]

ufp.terms.threebody.make_bucketed_energy_forces_evaluator(*, spline='cubic', n_nodes, n_cat=10, first_knot_xy=0.0, first_knot_z=0.0, knot_spacing_xy=0.25, knot_spacing_z=0.25, lower_support_xy=0.0, lower_support_z=0.0, eps=1e-12, runtime_config=None)[source]¶

Bind spline hyperparameters into a reusable bucket evaluator callable.

Parameters:

spline (Literal['quadratic', 'cubic', 'quartic'])
n_nodes (int)
n_cat (int)
first_knot_xy (float)
first_knot_z (float)
knot_spacing_xy (float)
knot_spacing_z (float)
lower_support_xy (float)
lower_support_z (float)
eps (float)
runtime_config (ThreeBodyRuntimeConfig | None)

Return type:

Callable[[Buckets, Tensor, Tensor, Tensor], tuple[Tensor, Tensor]]

ufp.terms.threebody.num_edge_categories(n_cat)[source]¶

Return the number of unordered neighbor-category pairs.

Parameters:: n_cat (int)
Return type:: int

ufp.terms.threebody.preprocess_sources(i, j, node_cat, n_cat, pair_vectors, pair_distances=None)[source]¶

Group directed neighbor rows into source buckets for triplet assembly.

Parameters:

i (Tensor)
j (Tensor)
node_cat (Tensor)
n_cat (int)
pair_vectors (Tensor)
pair_distances (Tensor | None)

Return type:

Buckets

Two-distance spline three-body term.

class ufp.terms.triplet2d.SplineTriplet2DTerm(*, cutoff, atomic_types, coeffs_by_triplet, active_triplets=None, spline='cubic', full_support_start=0.0, eps=1e-12, trainable=True, fittable=True, frozen=False, dtype=None)[source]¶

Bases: ThreeBodyTerm

Three-body spline over the two center-neighbor distances r_ij and r_ik.

Parameters:

cutoff (float)
atomic_types (Sequence[int])
active_triplets (Sequence[tuple[int, int, int]] | None)
spline (SplineKind)
full_support_start (float)
eps (float)
trainable (bool)
fittable (bool)
frozen (bool)
dtype (torch.dtype | None)

property active_triplet_categories: tuple[tuple[int, int, int], ...]¶: Return the subset of triplet categories enabled for evaluation.

assemble_linear_block(block, inputs, targets)[source]¶

Assemble this term’s dense least-squares block.

Parameters:

inputs (UFPInput)
targets (Any)

Return type:

Tensor | None

assemble_linear_blocks(batch, targets, options=None)[source]¶

Assemble all requested 2D triplet blocks for this term.

Parameters:: options (LinearAssemblyOptions | None)

assemble_selected_linear_block(block, inputs, targets, selected_indices)[source]¶

Assemble only requested coefficient columns for this 2D triplet block.

Parameters:

inputs (UFPInput)
targets (Any)
selected_indices (Sequence[int])

Return type:

Tensor | None

forward(inputs)[source]¶

Evaluate the 2D triplet spline over full-neighbor-list triplets.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

property input_requirements: TermInputRequirements¶: Declare the directed full-neighbor-list requirement.

property n_categories: int¶: Return the number of atomic categories.

parameter_blocks()[source]¶

Return the categorized 2D triplet coefficient block.

Return type:: tuple[ParameterBlock, …]

property provides_forces: bool¶: Report that this term produces analytic forces.

Torch-only four-body spline energy term.

class ufp.terms.fourbody.SplineFourBody6DTerm(*, cutoff, coeffs, spline='cubic', full_support_start=0.0, neighbor_neighbor_full_support_start=0.0, atomic_types=None, trainable=True, fittable=False, frozen=False, dtype=None)[source]¶

Bases: UFPTerm

Torch-only four-body term over the six pair distances in a source-neighbor triple.

This term contributes energy directly. Forces can be derived through autograd by calling model APIs with derive_forces=True.

Parameters:

cutoff (float) – Maximum center-neighbor pair distance included in the term.
coeffs – Six-dimensional coefficient tensor.
spline (SplineKind) – Spline family name.
full_support_start (float) – Lower center-neighbor distance with full spline support.
neighbor_neighbor_full_support_start (float) – Lower neighbor-neighbor distance with full spline support.
atomic_types (tuple[int, ...] | None) – Optional atomic-number coverage for the term.
trainable (bool) – Whether coefficients require gradients.
fittable (bool) – Whether this term exposes a coefficient block to linear fitters.
frozen (bool) – Whether the coefficient block is fixed during fitting.
dtype (torch.dtype | None) – Optional dtype used when converting coefficients.

forward(inputs)[source]¶

Evaluate four-body energies for source-centered neighbor triples.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

property input_requirements: TermInputRequirements¶: Declare the directed full-neighbor-list requirement.

parameter_blocks()[source]¶

Return the 6D coefficient block when explicitly fittable.

Return type:: tuple[ParameterBlock, …]

ZBL screened nuclear repulsion term.

class ufp.terms.zbl.ZBLRepulsionTerm(*, cutoff, atomic_types=None, eps=1e-12)[source]¶

Bases: PairTerm

Analytic Ziegler-Biersack-Littmark screened nuclear repulsion.

Parameters:

cutoff (float)
atomic_types (Sequence[int] | None)
eps (float)

forward(inputs)[source]¶

Evaluate ZBL energy and force contributions over covered pairs.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

property provides_forces: bool¶: Report that this term produces analytic pair forces.

Trainable pair-dependent inverse-power repulsion terms.

class ufp.terms.powerlaw.PowerLawRepulsionTerm(*, cutoff, atomic_types, prefactors_by_pair, power=2.0, active_pairs=None, symmetric=True, eps=1e-12, trainable=True, dtype=None)[source]¶

Bases: PairTerm

Pair-dependent inverse-power repulsion prior.

The term evaluates

\[E_{ij} = a_{Z_i, Z_j} / \max(r_{ij}, \epsilon)^p\]

where a is a trainable pair-channel prefactor. Forces are intentionally not returned by this term; callers that request forces should use the normal autograd path via derive_forces=True.

Parameters:

cutoff (float) – Maximum pair distance included in the interaction.
atomic_types (Sequence[int]) – Atomic numbers used to enumerate pair categories.
prefactors_by_pair – Either a tensor with shape (n_pair_categories,) in the term’s category order, or a mapping from pair tuples to prefactor values for all active pair categories.
power (float) – Positive inverse-power exponent. 2.0 gives a / r**2.
active_pairs (Sequence[tuple[int, int]] | None) – Optional subset of pair categories that should be evaluated. Inactive categories may have zero prefactors.
symmetric (bool) – If True, treat (a, b) and (b, a) as the same pair category.
eps (float) – Minimum distance used before applying the inverse power.
trainable (bool) – Whether prefactors require gradients.
dtype (torch.dtype | None) – Optional dtype used when converting prefactors to a tensor.

Examples

>>> term = PowerLawRepulsionTerm(
...     cutoff=2.5,
...     atomic_types=[1],
...     prefactors_by_pair={(1, 1): 0.4},
... )
>>> term.covers_pair(1, 1)
True

property active_pair_categories: tuple[tuple[int, int], ...]¶: Return the subset of pair categories enabled for evaluation.

canonical_pair(first, second)[source]¶

Normalize a pair key using this term’s symmetry convention.

Parameters:

first (int)
second (int)

Return type:

tuple[int, int]

covers_pair(first_atomic_number, second_atomic_number)[source]¶

Report whether this term evaluates the requested pair category.

Parameters:

first_atomic_number (int)
second_atomic_number (int)

Return type:

bool

forward(inputs)[source]¶

Evaluate active inverse-power pair energies.

Parameters:: inputs (UFPInput)
Return type:: UFPOutput

is_pair_active(first, second)[source]¶

Report whether a canonical pair category remains enabled.

Parameters:

first (int)
second (int)

Return type:

bool

property pair_categories: tuple[tuple[int, int], ...]¶: Return the full ordered list of pair categories owned by this term.

pair_category_index(first, second)[source]¶

Return the category index for one pair.

Parameters:

first (int)
second (int)

Return type:

int

prefactor_for_pair(first, second)[source]¶

Return the prefactor tensor for one pair category.

Parameters:

first (int)
second (int)

Return type:

Tensor

property true_prefactors_by_pair: Tensor¶: Return prefactors indexed by pair category.

Alchemical coefficient sharing for related spline terms.

Use this module when several true coefficient blocks should be generated from a smaller proxy basis with optional trainable mixing weights.

class ufp.terms.alchemical.AlchemicalCoefficients(*, proxy_coeffs, n_true_terms, weights=None, proxy_trainable=True, weights_trainable=True, dtype=None)[source]¶

Bases: Module

Shared linear projection from proxy coefficient tensors to true term coefficients.

Parameters:

n_true_terms (int)
proxy_trainable (bool)
weights_trainable (bool)
dtype (torch.dtype | None)

property true_coeffs: Tensor¶: Return the current true coefficient tensor after optional proxy mixing.

true_coeffs_for(index)[source]¶

Return one true coefficient block by index.

Parameters:: index (int)
Return type:: Tensor

property uses_identity_weights: bool¶: Report whether proxy coefficients are used directly as true coefficients.

Neighbor Lists¶

Status: Stable user API.

Neighbor-list containers and builders.

The public helpers normalize ASE, vesin, and metatomic neighbor data into one NeighborListData representation. Pass that representation into UFPInput when a model needs explicit pair geometry, distances, vectors, and system indexing.

class ufp.neighbors.NeighborListBackend(*values)[source]¶

Bases: str, Enum

Known neighbor-list backends supported directly by UFP.

class ufp.neighbors.NeighborListData(pairs, shifts, distances=None, vectors=None, backend='unknown', cutoff=None, full_list=True, sorted=False, strict=None)[source]¶

Bases: object

Lightweight container for neighbor-list data used by UFP models.

Parameters:

pairs (ndarray | Tensor)
shifts (ndarray | Tensor)
distances (ndarray | Tensor | None)
vectors (ndarray | Tensor | None)
backend (str)
cutoff (float | None)
full_list (bool)
sorted (bool)
strict (bool | None)

pairs¶

Pair indices with shape (2, n_pairs) or (n_pairs, 2).

Type:: numpy.ndarray | torch.Tensor

shifts¶

Periodic cell shifts with shape (n_pairs, 3).

Type:: numpy.ndarray | torch.Tensor

distances¶

Optional pair distances with shape (n_pairs,).

Type:: numpy.ndarray | torch.Tensor | None

vectors¶

Optional pair displacement vectors with shape (n_pairs, 3).

Type:: numpy.ndarray | torch.Tensor | None

backend¶

Name of the backend that produced this neighbor list.

Type:: str

cutoff¶

Cutoff used to build the list, if known.

Type:: float | None

full_list¶

Whether the list contains both i-j and j-i entries.

Type:: bool

sorted¶

Whether the list is lexicographically sorted by pair indices.

Type:: bool

strict¶

Whether all pairs are guaranteed to be strictly within the cutoff.

Type:: bool | None

as_torch(*, dtype=None, device=None)[source]¶

Return a torch-backed copy of this neighbor list.

Parameters:

dtype (dtype | None)
device (device | None)

Return type:

masked(mask)[source]¶

Return a new neighbor list containing only the selected pairs.

Parameters:: mask (ndarray | Tensor)
Return type:: NeighborListData

property n_pairs: int¶: Return the number of pairs.

pin_memory()[source]¶

Pin any torch-backed arrays in place and return self.

Return type:: NeighborListData

shifted(atom_offset)[source]¶

Return a copy with all pair indices shifted by the requested atom offset.

Parameters:: atom_offset (int)
Return type:: NeighborListData

ufp.neighbors.available_neighbor_backends()[source]¶

List neighbor-list backends available in the current Python environment.

Returns:: Available backends ordered by preference.
Return type:: List[NeighborListBackend]

ufp.neighbors.build_neighbor_list(atoms, cutoff, backend=NeighborListBackend.AUTO, arrays='torch', full_list=True, sorted=True)[source]¶

Build a neighbor list for an ASE structure.

When backend="auto", UFP prefers vesin if it is importable and falls back to ase.neighborlist otherwise.

Parameters:

atoms (Atoms) – Structure for which the neighbor list should be built.
cutoff (float) – Spherical cutoff radius.
backend (str | NeighborListBackend) – One of "auto", "vesin", or "ase".
arrays (str) – Output array type, either "torch" or "numpy".
full_list (bool) – Whether both i-j and j-i entries should be present.
sorted (bool) – Whether to sort the output lexicographically by pair indices.

Returns:

Normalized neighbor-list data.

Raises:

TypeError – If atoms is not an ASE structure.
ValueError – If options are invalid or the selected backend is unsupported.

Return type:

UniformCubicSpline1DDiagnostics

ufp.neighbors.concatenate_neighbor_lists(neighbor_lists, *, atom_offsets=None)[source]¶

Concatenate per-system neighbor lists into one batch-level list.

Parameters:

neighbor_lists (Sequence[NeighborListData])
atom_offsets (Sequence[int] | None)

Return type:

NeighborListData | None

Splines¶

Status: Extension API.

ufp.splines exports low-level support helpers for expert callers that need to coordinate coefficient support with fitting, projection, uncertainty, or custom term code. uniform_support_parameters, spline_support_mask_1d, spline_support_mask_3d, and all_supported_uniform_stencil_3d are public extension APIs. They are intentionally direct exports rather than facade calls so hot paths can bind the exact helper they need before entering tensor loops. supported_uniform_stencil_3d remains available for compatibility, but new performance-sensitive code should filter with spline_support_mask_3d and call all_supported_uniform_stencil_3d on already-supported coordinates.

Uniform spline basis helpers.

This package exposes local stencil builders, derivative rows, diagnostics, and fitting helpers used by spline terms and least-squares assembly. The stencil objects make coefficient locality visible for model inspection and tests. Support-mask and unchecked-stencil helpers are expert APIs for callers that already manage spline support filtering in hot paths.

Stencil representations for uniform spline bases.

Use this module to inspect which coefficients a geometry touches and the basis weights or gradients that should be applied to those coefficients. The support helpers exported here are low-level expert APIs; hot callers may use them directly after doing their own setup-time dispatch and filtering.

class ufp.splines.representation.Stencil1D(indices, values, grads)[source]¶

Bases: object

Local 1D spline stencil with coefficient indices, values, and gradients.

Parameters:

indices (Tensor)
values (Tensor)
grads (Tensor)

class ufp.splines.representation.Stencil2D(indices, values, grad_x, grad_y)[source]¶

Bases: object

Local 2D spline stencil with flat indices, values, and gradients.

Parameters:

indices (Tensor)
values (Tensor)
grad_x (Tensor)
grad_y (Tensor)

class ufp.splines.representation.Stencil3D(indices, values, grad_x, grad_y, grad_z)[source]¶

Bases: object

Local 3D spline stencil with coefficient indices, values, and gradients.

Parameters:

indices (Tensor)
values (Tensor)
grad_x (Tensor)
grad_y (Tensor)
grad_z (Tensor)

class ufp.splines.representation.Stencil6D(indices, values, grads)[source]¶

Bases: object

Local 6D spline stencil with flat indices, values, and per-axis gradients.

Parameters:

indices (Tensor)
values (Tensor)
grads (tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor])

class ufp.splines.representation.SupportedStencil3D(mask, x, y, z, stencil)[source]¶

Bases: object

3D stencil data for coordinates inside the spline support.

Parameters:

mask (Tensor)
x (Tensor)
y (Tensor)
z (Tensor)
stencil (Stencil3D)

ufp.splines.representation.all_supported_uniform_stencil_3d(x, y, z, *, coeff_shape, first_knot_xy, first_knot_z, knot_spacing_xy, knot_spacing_z, spline)[source]¶

Return an unchecked 3D stencil for coordinates already inside support.

Callers must filter coordinates to the full support region first. This is the fast path used by three-body evaluation, where center-neighbor and neighbor-neighbor distances are non-negative by construction.

Parameters:

x (Tensor) – First coordinate tensor.
y (Tensor) – Second coordinate tensor.
z (Tensor) – Third coordinate tensor.
coeff_shape (tuple[int, int, int]) – Three-dimensional coefficient grid shape.
first_knot_xy (float) – First knot for the first two distance coordinates.
first_knot_z (float) – First knot for the neighbor-neighbor distance coordinate.
knot_spacing_xy (float) – Knot spacing for the first two distance coordinates.
knot_spacing_z (float) – Knot spacing for the neighbor-neighbor distance coordinate.
spline (str) – Spline family name.

Returns:

Dense stencil values and gradients for all provided coordinates.

Raises:

ValueError – If coordinate tensors have different shapes.

Return type:

Stencil3D

ufp.splines.representation.spline_support_mask_1d(x, *, coeff_size, first_knot, knot_spacing, spline)[source]¶

Return a mask selecting 1D coordinates inside the active spline support.

Parameters:

x (Tensor)
coeff_size (int)
first_knot (float)
knot_spacing (float)
spline (str)

Return type:

Tensor

ufp.splines.representation.spline_support_mask_2d(x, y, *, coeff_shape, first_knot_x, first_knot_y, knot_spacing_x, knot_spacing_y, spline)[source]¶

Return a mask selecting 2D coordinates inside active spline support.

Parameters:

x (Tensor)
y (Tensor)
coeff_shape (tuple[int, int])
first_knot_x (float)
first_knot_y (float)
knot_spacing_x (float)
knot_spacing_y (float)
spline (str)

Return type:

Tensor

ufp.splines.representation.spline_support_mask_3d(x, y, z, *, coeff_shape, first_knot_xy, first_knot_z, knot_spacing_xy, knot_spacing_z, spline)[source]¶

Return a mask selecting 3D coordinates inside the active spline support.

Parameters:

x (Tensor)
y (Tensor)
z (Tensor)
coeff_shape (tuple[int, int, int])
first_knot_xy (float)
first_knot_z (float)
knot_spacing_xy (float)
knot_spacing_z (float)
spline (str)

Return type:

Tensor

ufp.splines.representation.spline_support_mask_6d(coords, *, coeff_shape, first_knots, knot_spacings, spline)[source]¶

Return a mask selecting 6D coordinates inside active spline support.

Parameters:

coords (tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor])
coeff_shape (tuple[int, int, int, int, int, int])
first_knots (tuple[float, float, float, float, float, float])
knot_spacings (tuple[float, float, float, float, float, float])
spline (str)

Return type:

Tensor

ufp.splines.representation.supported_uniform_stencil_3d(x, y, z, *, coeff_shape, first_knot_xy, first_knot_z, knot_spacing_xy, knot_spacing_z, spline)[source]¶

Return the support mask and unchecked stencil for supported 3D coordinates.

Prefer filtering coordinates before calling all_supported_uniform_stencil_3d() in performance-sensitive paths.

Parameters:

x (Tensor) – First coordinate tensor.
y (Tensor) – Second coordinate tensor.
z (Tensor) – Third coordinate tensor.
coeff_shape (tuple[int, int, int]) – Three-dimensional coefficient grid shape.
first_knot_xy (float) – First knot for the first two distance coordinates.
first_knot_z (float) – First knot for the neighbor-neighbor distance coordinate.
knot_spacing_xy (float) – Knot spacing for the first two distance coordinates.
knot_spacing_z (float) – Knot spacing for the neighbor-neighbor distance coordinate.
spline (str) – Spline family name.

Returns:

Support mask and stencil for the supported coordinates.

Raises:

ValueError – If coordinate tensors have different shapes.

Return type:

SupportedStencil3D

ufp.splines.representation.uniform_stencil_1d(x, *, coeff_size, first_knot, knot_spacing, spline)[source]¶

Return coefficient indices, values, and gradients for a 1D uniform stencil.

Parameters:

x (Tensor)
coeff_size (int)
first_knot (float)
knot_spacing (float)
spline (str)

Return type:

Stencil1D

ufp.splines.representation.uniform_stencil_2d(x, y, *, coeff_shape, first_knot_x, first_knot_y, knot_spacing_x, knot_spacing_y, spline)[source]¶

Return coefficient indices, values, and gradients for a 2D uniform stencil.

Parameters:

x (Tensor)
y (Tensor)
coeff_shape (tuple[int, int])
first_knot_x (float)
first_knot_y (float)
knot_spacing_x (float)
knot_spacing_y (float)
spline (str)

Return type:

Stencil2D

ufp.splines.representation.uniform_stencil_3d(x, y, z, *, coeff_shape, first_knot_xy, first_knot_z, knot_spacing_xy, knot_spacing_z, spline)[source]¶

Return coefficient indices, values, and gradients for a 3D uniform stencil.

Parameters:

x (Tensor)
y (Tensor)
z (Tensor)
coeff_shape (tuple[int, int, int])
first_knot_xy (float)
first_knot_z (float)
knot_spacing_xy (float)
knot_spacing_z (float)
spline (str)

Return type:

Stencil3D

ufp.splines.representation.uniform_stencil_6d(coords, *, coeff_shape, first_knots, knot_spacings, spline)[source]¶

Return coefficient indices, values, and gradients for a 6D uniform stencil.

Parameters:

coords (tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor])
coeff_shape (tuple[int, int, int, int, int, int])
first_knots (tuple[float, float, float, float, float, float])
knot_spacings (tuple[float, float, float, float, float, float])
spline (str)

Return type:

Stencil6D

ufp.splines.representation.uniform_support_parameters(*, coeff_size, lower_full_support, upper_full_support, spline)[source]¶

Compute the first knot and spacing for a uniform spline support.

Parameters:

coeff_size (int)
lower_full_support (float)
upper_full_support (float)
spline (str)

Return type:

tuple[float, float]

Derivative helpers for uniform one-dimensional cubic splines.

class ufp.splines.derivatives.UniformCubicSpline1DDiagnostics(distances, values, gradients, curvatures)[source]¶

Bases: object

Sampled value, first derivative, and second derivative for one spline.

Parameters:

distances (Tensor)
values (Tensor)
gradients (Tensor)
curvatures (Tensor)

ufp.splines.derivatives.cubic_derivative_rows_1d(distances, *, coeff_size, first_knot, knot_spacing, dtype=None, device=None)[source]¶

Return dense first-derivative rows for supported cubic spline distances.

Parameters:

coeff_size (int)
first_knot (float)
knot_spacing (float)
dtype (dtype | None)
device (device | str | None)

Return type:

Tensor

ufp.splines.derivatives.cubic_second_derivative_operator(*, coeff_size, knot_spacing, dtype=None, device=None)[source]¶

Return the operator mapping cubic coefficients to curvature coefficients.

Parameters:

coeff_size (int)
knot_spacing (float)
dtype (dtype | None)
device (device | str | None)

Return type:

Tensor

ufp.splines.derivatives.cubic_second_derivative_values_1d(distances, coeffs, *, first_knot, knot_spacing)[source]¶

Evaluate the second derivative of a uniform cubic spline.

Parameters:

coeffs (Tensor)
first_knot (float)
knot_spacing (float)

Return type:

Tensor

ufp.splines.derivatives.cubic_spline_diagnostics_1d(distances, coeffs, *, first_knot, knot_spacing)[source]¶

Return sampled value, gradient, and curvature diagnostics for one spline.

Parameters:

coeffs (Tensor)
first_knot (float)
knot_spacing (float)

Return type:

ufp.splines.derivatives.cubic_value_rows_1d(distances, *, coeff_size, first_knot, knot_spacing, dtype=None, device=None)[source]¶

Return dense cubic B-spline value rows for supported distances.

Parameters:

coeff_size (int)
first_knot (float)
knot_spacing (float)
dtype (dtype | None)
device (device | str | None)

Return type:

Tensor

Projection helpers for uniform one-dimensional splines.

class ufp.splines.fitting.UniformSpline1DFitResult(coeffs, sample_distances, target_values, predicted_values, rmse, max_abs_error, target_gradients=None, predicted_gradients=None, gradient_rmse=None, gradient_max_abs_error=None)[source]¶

Bases: object

Diagnostics from fitting one callable onto a uniform 1D spline basis.

Parameters:

coeffs (Tensor)
sample_distances (Tensor)
target_values (Tensor)
predicted_values (Tensor)
rmse (float)
max_abs_error (float)
target_gradients (Tensor | None)
predicted_gradients (Tensor | None)
gradient_rmse (float | None)
gradient_max_abs_error (float | None)

ufp.splines.fitting.fit_uniform_spline_1d(function, *, coeff_size, lower_full_support, upper_full_support, spline='cubic', n_samples=None, sample_distances=None, derivative_weight=1.0, dtype=None, device=None, rcond=None)[source]¶

Fit a callable on a uniform 1D spline basis with optional derivative rows.

Parameters:

function (Callable[[Tensor], Tensor]) – Callable accepting a distance tensor and returning energy-like values with the same shape.
coeff_size (int) – Number of spline coefficients.
lower_full_support (float) – Lower physical distance with full spline support.
upper_full_support (float) – Upper physical distance with full spline support.
spline (str) – Spline family name.
n_samples (int | None) – Number of midpoint samples when sample_distances is not supplied. Defaults to max(4 * coeff_size, coeff_size + 8).
sample_distances (Tensor | None) – Optional explicit sample distances.
derivative_weight (float) – Weight applied to derivative rows. Set to 0.0 to fit only function values.
dtype (dtype | None) – Optional fitting dtype.
device (device | str | None) – Optional fitting device.
rcond (float | None) – Optional cutoff passed to torch.linalg.lstsq().

Returns:

Coefficients and projection diagnostics.

Return type:

UniformSpline1DFitResult

Deprecated Compatibility APIs¶

API	Status	Replacement guidance
`TwoBodySplineShapePenalty`	Deprecated compatibility behavior. It remains importable so older training and least-squares scripts keep running.	Prefer term-family ridge values, `RegularizationSearchConfig`, coefficient selection, explicit priors, validation data checks, or `TwoBodyCriticalWallPenalty` for physical slope/curvature wall regularization during Torch training.
`supported_uniform_stencil_3d`	Deprecated helper that combines support filtering and stencil construction.	Filter coordinates outside performance-sensitive loops, then call `ufp.splines.representation.all_supported_uniform_stencil_3d` on already-supported coordinates.

Least Squares¶

Status: Stable user and expert fitting API. Public imports cover fitting datasets, parameter layouts, solvers, regularization, cache IO, checkpoints, and diagnostic problem containers. Underscore helpers used by assembly and block-matrix internals are private even when they remain directly importable by UFP internals.

Direct linear fitting for coefficient-linear UFP models.

The least-squares package builds fixed-geometry design blocks from FitSample objects, streams prepared batches, and solves for coefficient vectors with dense or iterative backends. Diagnostic materialized matrices are available when callers need to inspect the assembled problem directly.

Fit-sample normalization and batch preparation for least-squares workflows.

Use this module to convert labeled structures into weighted target rows and the shared UFPInput batches consumed by the assembly layer.

class ufp.leastsquares.dataset.BatchTargets(values, sqrt_weights, row_scales, energy_rows, force_rows, per_atom_rows)[source]¶

Bases: object

Weighted target rows and row mappings for one prepared fit batch.

Parameters:

values (Tensor)
sqrt_weights (Tensor)
row_scales (Tensor)
energy_rows (Tensor)
force_rows (Tensor)
per_atom_rows (Tensor)

property n_rows: int¶: Return the number of active target rows in this batch.

property weighted_values: Tensor¶: Return target values after square-root weighting is applied.

class ufp.leastsquares.dataset.FitSample(atoms, neighbor_list=None, energy=None, forces=None, per_atom_energy=None, energy_weight=1.0, force_weight=1.0, per_atom_weight=1.0)[source]¶

Bases: object

One labeled structure for least-squares fitting over fixed geometry.

Parameters:

atoms (Atoms)
neighbor_list (NeighborListData | None)
energy (float | None)
forces (ndarray | Tensor | None)
per_atom_energy (ndarray | Tensor | None)
energy_weight (float)
force_weight (float)
per_atom_weight (float)

has_targets(*, fit_energy, fit_forces, fit_per_atom_energy)[source]¶

Report whether this sample contributes to the enabled target types.

Parameters:

fit_energy (bool)
fit_forces (bool)
fit_per_atom_energy (bool)

Return type:

bool

class ufp.leastsquares.dataset.PreparedBatch(samples, inputs, targets)[source]¶

Bases: object

One batch coupling input geometry with weighted least-squares targets.

Parameters:

samples (tuple[FitSample, ...])
inputs (UFPInput)
targets (BatchTargets)

ufp.leastsquares.dataset.prepare_batches(model, samples, *, batch_size, fit_energy, fit_forces, fit_per_atom_energy, dtype, device)[source]¶

Convert fit samples into prepared batches ready for matrix assembly.

Parameters:

model (UFPotential)
samples (Sequence[FitSample])
batch_size (int)
fit_energy (bool)
fit_forces (bool)
fit_per_atom_energy (bool)
dtype (dtype)
device (device | None)

Return type:

tuple[PreparedBatch, …]

Direct linear fitting over assembled spline coefficient blocks.

Use this module to build least-squares problems, materialize diagnostic matrices, and solve for direct coefficient vectors with dense or iterative backends.

class ufp.leastsquares.linear.BlockLinearProblem(layout, batches)[source]¶

Bases: object

Matrix-free linear least-squares problem assembled from block batches.

Parameters:

layout (BlockProblemLayout)
batches (Sequence[BlockSolveBatch])

accumulate_normal_equations()[source]¶

Materialize the Gram matrix and right-hand side from all batches.

Return type:: tuple[Tensor, Tensor]

design_trace_by_block()[source]¶

Return weighted design-matrix trace contributions by solve block.

Return type:: dict[object, float]

property device: device¶: Return the device shared by the assembled problem tensors.

property dtype: dtype¶: Return the dtype shared by the assembled problem tensors.

materialize_design_matrix()[source]¶

Build the explicit dense design matrix for debugging or tiny problems.

Return type:: tuple[Tensor, Tensor]

matvec(theta)[source]¶

Apply the design matrix to a flat parameter vector.

Parameters:: theta (Tensor)
Return type:: Tensor

property n_rows: int¶: Return the total number of target rows across all batches.

normal_equation_diagonal()[source]¶

Return the diagonal of the weighted design A.T @ A.

Return type:: Tensor

normal_matvec(theta)[source]¶

Apply the regularized normal-equation operator to a flat vector.

Parameters:: theta (Tensor)
Return type:: Tensor

objective(theta)[source]¶

Evaluate the regularized least-squares objective at theta.

Parameters:: theta (Tensor)
Return type:: Tensor

regularization_apply(theta)[source]¶

Apply all block regularizers to a flat parameter vector.

Parameters:: theta (Tensor)
Return type:: Tensor

regularization_diagonal()[source]¶

Return the summed diagonal preconditioner implied by block regularizers.

Return type:: Tensor

regularization_rhs()[source]¶

Return the summed RHS shifts implied by block regularizers.

Return type:: Tensor

rhs()[source]¶

Return the right-hand side of the normal equations.

Return type:: Tensor

rmatvec(residual)[source]¶

Apply the transpose design matrix to a residual vector.

Parameters:: residual (Tensor)
Return type:: Tensor

solve(*, solver, cg_tolerance, cg_max_iter, progress=False, progress_frequency=10, initial_theta=None, cg_checkpoint_path=None, cg_checkpoint_frequency=1, cg_resume=False, cg_checkpoint_metadata=None, fallback_theta=None, return_info=False)[source]¶

Solve the assembled problem with the selected dense or iterative backend.

Parameters:

solver (str)
cg_tolerance (float)
cg_max_iter (int | None)
progress (bool)
progress_frequency (int)
initial_theta (torch.Tensor | None)
cg_checkpoint_path (Path | str | None)
cg_checkpoint_frequency (int)
cg_resume (bool)
cg_checkpoint_metadata (dict[str, object] | None)
fallback_theta (torch.Tensor | None)
return_info (bool)

Return type:

torch.Tensor | ‘LinearSolveResult’

target_vector()[source]¶

Concatenate all batch targets into one right-hand-side vector.

Return type:: Tensor

class ufp.leastsquares.linear.BlockProblemLayout(blocks, slices, size)[source]¶

Bases: object

Ordered block layout used by an assembled linear problem.

Parameters:

blocks (tuple[SolveBlock, ...])
slices (dict[Any, slice])
size (int)

classmethod from_blocks(blocks)[source]¶

Assign slices to an ordered sequence of solve blocks.

Parameters:: blocks (Sequence[SolveBlock])
Return type:: BlockProblemLayout

theta_slice(key)[source]¶

Return the solve-vector slice occupied by the requested block key.

Parameters:: key (Any)
Return type:: slice

class ufp.leastsquares.linear.BlockRegularization(shape, ridge=0.0, third_difference_penalty=0.0, active_rows=None, third_difference_stencils=None)[source]¶

Bases: object

Regularization recipe applied blockwise inside the linear problem.

Parameters:

shape (tuple[int, ...])
ridge (float)
third_difference_penalty (float)
active_rows (tuple[int, ...] | None)
third_difference_stencils (tuple[RegularizationStencil, ...] | None)

apply(values)[source]¶

Apply ridge and optional third-difference penalties to one block.

Parameters:: values (Tensor)
Return type:: Tensor

constant(*, dtype, device)[source]¶

Return the regularization constant from fixed-coefficient targets.

Parameters:

dtype (dtype)
device (device)

Return type:

Tensor

diagonal(*, dtype, device)[source]¶

Return the diagonal of the block regularizer in vector form.

Parameters:

dtype (dtype)
device (device)

Return type:

Tensor

least_squares_rows(*, dtype, device)[source]¶

Materialize regularization as least-squares rows and targets.

Parameters:

dtype (dtype)
device (device)

Return type:

tuple[Tensor, Tensor]

materialize(*, dtype, device)[source]¶

Materialize the block regularizer as an explicit square matrix.

Parameters:

dtype (dtype)
device (device)

Return type:

Tensor

quadratic(values)[source]¶

Evaluate the quadratic penalty contributed by one block vector.

Parameters:: values (Tensor)
Return type:: Tensor

rhs(*, dtype, device)[source]¶

Return the regularization contribution to the normal-equation RHS.

Parameters:

dtype (dtype)
device (device)

Return type:

Tensor

semantics_metadata()[source]¶

Return metadata for cache/checkpoint invalidation.

Return type:: dict[str, object]

property size: int¶: Return the flattened size of the regularized block.

stencils_metadata()[source]¶

Return compact metadata for selected regularization rows.

Return type:: dict[str, object]

class ufp.leastsquares.linear.BlockSolveBatch(target, matrices)[source]¶

Bases: object

One assembled batch contributing rows to the block linear problem.

Parameters:

target (Tensor)
matrices (dict[Any, Tensor | RowIndexedBlockMatrix | ColumnRowIndexedBlockMatrix])

property n_rows: int¶: Return the number of target rows stored in this batch.

class ufp.leastsquares.linear.CGCheckpointState(x, residual, direction, rz_old, iteration, metadata)[source]¶

Bases: object

Restart state for one conjugate-gradient solve.

Parameters:

x (Tensor)
residual (Tensor)
direction (Tensor)
rz_old (Tensor)
iteration (int)
metadata (dict[str, object])

class ufp.leastsquares.linear.CachedBlockLinearProblem(*, layout, cache_dir, entries, dtype, device, row_weights, repulsion_batch=None, load_batch_entry=<function _load_assembled_batch_entry_memmap>, row_indexed_blocks=True)[source]¶

Bases: object

Matrix-free linear problem that streams cached batches from disk.

Parameters:

layout (BlockProblemLayout)
cache_dir (Path)
entries (Sequence[object])
dtype (torch.dtype)
device (torch.device)
row_weights (Sequence[torch.Tensor])
repulsion_batch (BlockSolveBatch | None)
load_batch_entry (Callable[..., AssembledBatch])
row_indexed_blocks (bool)

accumulate_normal_equations()[source]¶

Materialize the Gram matrix and right-hand side from streamed batches.

Return type:: tuple[Tensor, Tensor]

property batches: _CachedBlockBatchSequence¶: Return a lazy sequence that loads cached batches on demand.

design_trace_by_block()[source]¶

Return weighted design-matrix trace contributions by solve block.

Return type:: dict[object, float]

property device: device¶: Return the device used when cached batches are streamed.

property dtype: dtype¶: Return the dtype used when cached batches are streamed.

materialize_design_matrix()[source]¶

Build the explicit dense design matrix for debugging or tiny problems.

Return type:: tuple[Tensor, Tensor]

matvec(theta)[source]¶

Apply the design matrix, loading cached batches on demand.

Parameters:: theta (Tensor)
Return type:: Tensor

property n_rows: int¶: Return the total number of target rows across all batches.

normal_equation_diagonal()[source]¶

Return the diagonal of A.T @ A by streaming cached batches.

Return type:: Tensor

normal_matvec(theta)[source]¶

Apply the regularized normal operator without concatenating rows.

Parameters:: theta (Tensor)
Return type:: Tensor

objective(theta)[source]¶

Evaluate the regularized objective without materializing residuals.

Parameters:: theta (Tensor)
Return type:: Tensor

regularization_apply(theta)[source]¶

Apply all block regularizers to a flat parameter vector.

Parameters:: theta (Tensor)
Return type:: Tensor

regularization_diagonal()[source]¶

Return the summed diagonal preconditioner implied by regularizers.

Return type:: Tensor

regularization_rhs()[source]¶

Return the summed RHS shifts implied by block regularizers.

Return type:: Tensor

residual_norm(theta)[source]¶

Return ||A theta - b|| without materializing all residual rows.

Parameters:: theta (Tensor)
Return type:: Tensor

rhs()[source]¶

Return the right-hand side by streaming cached batches.

Return type:: Tensor

rmatvec(residual)[source]¶

Apply the transpose design matrix to a residual vector.

Parameters:: residual (Tensor)
Return type:: Tensor

Solve the streamed problem with the selected backend.

Parameters:

solver (str)
cg_tolerance (float)
cg_max_iter (int | None)
progress (bool)
progress_frequency (int)
initial_theta (Tensor | None)
cg_checkpoint_path (Path | str | None)
cg_checkpoint_frequency (int)
cg_resume (bool)
cg_checkpoint_metadata (dict[str, object] | None)
fallback_theta (Tensor | None)
return_info (bool)

Return type:

Tensor | LinearSolveResult

target_vector()[source]¶

Concatenate targets, loading cached batches on demand.

Return type:: Tensor

class ufp.leastsquares.linear.CoefficientSelector(block, channel=None, coeff_slice=None)[source]¶

Bases: object

Select a subset of one or more coefficient blocks.

block uses the existing block selector semantics: block index, block name, kind, label, or regularization group. channel optionally selects a pair or triplet channel inside categorized coefficient tensors. coeff_slice selects trailing coefficient dimensions inside the selected channel or channels.

Parameters:

block (int | str)
channel (tuple[int, ...] | None)
coeff_slice (Any)

class ufp.leastsquares.linear.ColumnRowIndexedBlockMatrix(chunks, n_rows, n_cols)[source]¶

Bases: object

Dense block values stored by contiguous column chunks and active rows.

Parameters:

chunks (tuple[ColumnRowIndexedChunk, ...])
n_rows (int)
n_cols (int)

property device: device¶: Return the value device.

property dtype: dtype¶: Return the value dtype.

materialize()[source]¶

Return a dense tensor with zero chunks restored.

Return type:: Tensor

matvec(theta)[source]¶

Apply this matrix to one block parameter vector.

Parameters:: theta (Tensor)
Return type:: Tensor

rmatvec(residual)[source]¶

Apply this matrix transpose to a residual vector.

Parameters:: residual (Tensor)
Return type:: Tensor

property shape: tuple[int, int]¶: Return the logical dense matrix shape.

class ufp.leastsquares.linear.LinearFitResult(theta, objective, residual_norm, solver, n_rows, n_parameters, layout, problem, interrupted=False, restored_checkpoint_path=None)[source]¶

Bases: object

Summary of one direct linear fit together with its assembled problem.

Parameters:

theta (Tensor)
objective (float)
residual_norm (float)
solver (str)
n_rows (int)
n_parameters (int)
layout (ParameterLayout)
problem (BlockLinearProblem)
interrupted (bool)
restored_checkpoint_path (str | None)

class ufp.leastsquares.linear.LinearFitter(model, *, fit_energy=True, fit_forces=True, fit_per_atom_energy=False, solver='cg', ridge=0.0, onebody_ridge=None, pair_ridge=None, twobody_ridge=None, threebody_ridge=None, twobody_shape_penalty=None, dtype=None, device=None, cg_tolerance=1e-10, cg_max_iter=None, threebody_lstsq_backend=None, threebody_bucket_backend=None, fit_blocks=None, freeze_blocks=(), assembly_contract='term', matrix_storage='auto')[source]¶

Bases: DenseNormalEquationMixin

High-level entry point for direct least-squares fitting of spline models.

Parameters:

model (UFPModel)
fit_energy (bool)
fit_forces (bool)
fit_per_atom_energy (bool)
solver (str)
ridge (float)
onebody_ridge (float | None)
pair_ridge (float | None)
twobody_ridge (float | None)
threebody_ridge (float | None)
twobody_shape_penalty (TwoBodySplineShapePenalty | None)
dtype (Optional[torch.dtype])
device (Optional[torch.device])
cg_tolerance (float)
cg_max_iter (int | None)
threebody_lstsq_backend (str | None)
threebody_bucket_backend (str | None)
fit_blocks (Sequence[int | str] | None)
freeze_blocks (Sequence[int | str])
assembly_contract (str)
matrix_storage (MatrixStorageMode)

accumulate_normal_equations(samples, *, batch_size=32, cache_directory=None, cache_mode='auto', normal_equation_cache_directory=None, normal_equation_cache_mode='auto', normal_equation_build_device=None, progress=False)[source]¶

Build the normal equations directly from the provided samples.

Parameters:

samples (Sequence[FitSample])
batch_size (int)
cache_directory (Path | str | None)
cache_mode (Literal['auto', 'read', 'write', 'refresh'])
normal_equation_cache_directory (Path | str | None)
normal_equation_cache_mode (Literal['auto', 'read', 'write', 'refresh'])
normal_equation_build_device (device | str | None)
progress (bool)

Return type:

tuple[Tensor, Tensor]

build_problem(samples, *, batch_size=32, progress=False, cache_directory=None, cache_mode='auto')[source]¶

Prepare batches, assemble block matrices, and return the linear problem.

Parameters:

samples (Sequence[FitSample])
batch_size (int)
progress (bool)
cache_directory (Path | str | None)
cache_mode (Literal['auto', 'read', 'write', 'refresh'])

Return type:

BlockLinearProblem

fit(samples, *, batch_size=32, write_back=True, progress=False, progress_frequency=10, cache_directory=None, cache_mode='auto', normal_equation_cache=False, normal_equation_cache_directory=None, normal_equation_cache_mode='auto', normal_equation_build_device=None, warm_start=False, cg_checkpoint_path=None, cg_checkpoint_frequency=1, cg_resume=False)[source]¶

Build, solve, optionally write back, and summarize one direct linear fit.

Parameters:

samples (Sequence[FitSample])
batch_size (int)
write_back (bool)
progress (bool)
progress_frequency (int)
cache_directory (Path | str | None)
cache_mode (Literal['auto', 'read', 'write', 'refresh'])
normal_equation_cache (bool)
normal_equation_cache_directory (Path | str | None)
normal_equation_cache_mode (Literal['auto', 'read', 'write', 'refresh'])
normal_equation_build_device (device | str | None)
warm_start (bool)
cg_checkpoint_path (Path | str | None)
cg_checkpoint_frequency (int)
cg_resume (bool)

Return type:

LinearFitResult

classmethod from_model(model, **kwargs)[source]¶

Convenience constructor mirroring __init__.

Parameters:: model (UFPModel)
Return type:: LinearFitter

make_linear_operator(samples, *, batch_size=32, cache_directory=None, cache_mode='auto')[source]¶

Alias build_problem for callers that want a matrix-free operator view.

Parameters:

samples (Sequence[FitSample])
batch_size (int)
cache_directory (Path | str | None)
cache_mode (Literal['auto', 'read', 'write', 'refresh'])

Return type:

BlockLinearProblem

materialize_design_matrix(samples, *, batch_size=32, cache_directory=None, cache_mode='auto')[source]¶

Build the explicit design matrix for the provided samples.

Parameters:

samples (Sequence[FitSample])
batch_size (int)
cache_directory (Path | str | None)
cache_mode (Literal['auto', 'read', 'write', 'refresh'])

Return type:

tuple[Tensor, Tensor]

problem_from_assembled_batches(assembled_batches)[source]¶

Wrap preassembled batches in the common block-problem interface.

Parameters:: assembled_batches (Sequence[AssembledBatch])
Return type:: BlockLinearProblem

write_back(theta)[source]¶

Write one solved direct vector back into the model coefficients.

Parameters:: theta (Tensor)
Return type:: None

write_checkpoint_to_model(checkpoint_path)[source]¶

Write the coefficient vector from a CG checkpoint into the model.

Parameters:: checkpoint_path (Path | str)
Return type:: None

class ufp.leastsquares.linear.LinearSolveResult(theta, interrupted=False, restored_checkpoint_path=None)[source]¶

Bases: object

Parameter vector returned by one linear solve with interrupt metadata.

Parameters:

theta (Tensor)
interrupted (bool)
restored_checkpoint_path (str | None)

class ufp.leastsquares.linear.RowIndexedBlockMatrix(rows, values, n_rows)[source]¶

Bases: object

Dense block values stored only for rows that can be nonzero.

Parameters:

rows (Tensor)
values (Tensor)
n_rows (int)

property device: device¶: Return the value device.

property dtype: dtype¶: Return the value dtype.

materialize()[source]¶

Return a dense tensor with zero rows restored.

Return type:: Tensor

matvec(theta)[source]¶

Apply this matrix to one block parameter vector.

Parameters:: theta (Tensor)
Return type:: Tensor

rmatvec(residual)[source]¶

Apply this matrix transpose to a residual vector.

Parameters:: residual (Tensor)
Return type:: Tensor

property shape: tuple[int, int]¶: Return the logical dense matrix shape.

class ufp.leastsquares.linear.SolveBlock(key, size, label, regularization=None)[source]¶

Bases: object

Solve-time metadata for one block in the assembled linear system.

Parameters:

key (Any)
size (int)
label (str)
regularization (BlockRegularization | None)

ufp.leastsquares.linear.load_assembled_batches_memmap(directory, *, manifest_name='assembled_batches.json', mmap_mode='r+', dtype=None, device=None, expected_metadata=None, row_indexed_blocks=False, row_weights=None)[source]¶

Load block-separated assembled batches from a disk cache.

Parameters:

directory (Path | str) – Directory containing a cache manifest and .npy files.
manifest_name (str) – JSON manifest filename inside directory.
mmap_mode (Literal['r', 'r+', 'c'] | None) – Mode passed to numpy.load. Use None to load ordinary arrays.
dtype (dtype | None) – Optional dtype conversion for loaded tensors.
device (device | None) – Optional device conversion for loaded tensors.
expected_metadata (dict[str, object] | None) – Optional metadata payload that must match the cache.
row_indexed_blocks (bool) – Whether to preserve sparse row-indexed block storage.
row_weights (Sequence[Tensor] | None) – Optional row weights applied to loaded batches.

Returns:

Assembled least-squares batches.

Raises:

ValueError – If cache metadata or row-weight counts do not match.

Return type:

tuple[AssembledBatch, …]

ufp.leastsquares.linear.load_cg_checkpoint(path, *, dtype=None, device=None, expected_metadata=None)[source]¶

Load and validate one conjugate-gradient restart state.

Parameters:

path (Path | str)
dtype (dtype | None)
device (device | None)
expected_metadata (dict[str, object] | None)

Return type:

CGCheckpointState

ufp.leastsquares.linear.save_assembled_batches_memmap(directory, assembled_batches, *, manifest_name='assembled_batches.json', metadata=None, column_chunk_sizes=None, compact=True)[source]¶

Persist assembled least-squares batches as block-separated .npy files.

The cache stores each batch target and each term-block matrix separately. This keeps the existing block layout visible on disk and avoids materializing one global dense design matrix.

Parameters:

directory (Path | str) – Directory where cache files should be written.
assembled_batches (Sequence[AssembledBatch]) – Batches returned by assemble_true_blocks.
manifest_name (str) – JSON manifest filename inside directory.
metadata (dict[str, object] | None) – Optional metadata copied into the manifest.
column_chunk_sizes (dict[object, int] | None) – Optional column chunk sizes keyed by cache block key.
compact (bool) – Whether dense block matrices should be compacted during writing.

Return type:

None

ufp.leastsquares.linear.save_cg_checkpoint(path, state)[source]¶

Persist one conjugate-gradient restart state as an .npz file.

Parameters:

path (Path | str)
state (CGCheckpointState)

Return type:

None

Blockwise regularization operators for linear least-squares fits.

class ufp.leastsquares.regularization.BlockRegularization(shape, ridge=0.0, third_difference_penalty=0.0, active_rows=None, third_difference_stencils=None)[source]¶

Bases: object

Regularization recipe applied blockwise inside the linear problem.

Parameters:

shape (tuple[int, ...])
ridge (float)
third_difference_penalty (float)
active_rows (tuple[int, ...] | None)
third_difference_stencils (tuple[RegularizationStencil, ...] | None)

apply(values)[source]¶

Apply ridge and optional third-difference penalties to one block.

Parameters:: values (Tensor)
Return type:: Tensor

constant(*, dtype, device)[source]¶

Return the regularization constant from fixed-coefficient targets.

Parameters:

dtype (dtype)
device (device)

Return type:

Tensor

diagonal(*, dtype, device)[source]¶

Return the diagonal of the block regularizer in vector form.

Parameters:

dtype (dtype)
device (device)

Return type:

Tensor

least_squares_rows(*, dtype, device)[source]¶

Materialize regularization as least-squares rows and targets.

Parameters:

dtype (dtype)
device (device)

Return type:

tuple[Tensor, Tensor]

materialize(*, dtype, device)[source]¶

Materialize the block regularizer as an explicit square matrix.

Parameters:

dtype (dtype)
device (device)

Return type:

Tensor

quadratic(values)[source]¶

Evaluate the quadratic penalty contributed by one block vector.

Parameters:: values (Tensor)
Return type:: Tensor

rhs(*, dtype, device)[source]¶

Return the regularization contribution to the normal-equation RHS.

Parameters:

dtype (dtype)
device (device)

Return type:

Tensor

semantics_metadata()[source]¶

Return metadata for cache/checkpoint invalidation.

Return type:: dict[str, object]

property size: int¶: Return the flattened size of the regularized block.

stencils_metadata()[source]¶

Return compact metadata for selected regularization rows.

Return type:: dict[str, object]

class ufp.leastsquares.regularization.RegularizationStencil(columns, weights, target=0.0)[source]¶

Bases: object

One least-squares regularization row over a compact block vector.

Parameters:

columns (tuple[int, ...])
weights (tuple[float, ...])
target (float)

Hybrid residualized linear fitting helpers.

class ufp.leastsquares.hybrid.HybridLinearFitter(model, *, frozen_terms=(), frozen_term_filter=None, **linear_options)[source]¶

Bases: object

Residualize frozen nonlinear terms, then delegate to LinearFitter.

Frozen terms are evaluated on the same ASE structures and target components as least-squares fitting. The resulting residual FitSample objects are passed to the ordinary linear fitter so dense, normal-equation, cache, and CG solve paths remain centralized.

Parameters:

model (UFPModel)
frozen_terms (Sequence[FrozenTermSelector])
frozen_term_filter (FrozenTermFilter | None)

accumulate_normal_equations(samples, **kwargs)[source]¶

Accumulate normal equations from residualized samples.

Parameters:: samples (Sequence[FitSample])

build_problem(samples, **kwargs)[source]¶

Build a residualized linear problem through the delegated fitter.

Parameters:: samples (Sequence[FitSample])

fit(samples, **kwargs)[source]¶

Fit selected linear coefficients against residualized targets.

Parameters:: samples (Sequence[FitSample])

property frozen_term_indices: tuple[int, ...]¶: Return model-order indices for the selected frozen residual terms.

property frozen_terms_signature: str¶

Return the hybrid cache state digest.

This alias is kept for callers that used the original name. New code should prefer frozen_terms_state_hash, matching residual materialization metadata.

property frozen_terms_state_hash: str¶: Return a strict state digest for selected frozen residual terms.

property layout¶: Return the delegated linear fitter parameter layout.

make_linear_operator(samples, **kwargs)[source]¶

Alias build_problem for matrix-free callers.

Parameters:: samples (Sequence[FitSample])

materialize_design_matrix(samples, **kwargs)[source]¶

Build the explicit residualized design matrix and target vector.

Parameters:: samples (Sequence[FitSample])

residualized_samples(samples)[source]¶

Return samples with frozen-term target contributions subtracted.

Parameters:: samples (Sequence[FitSample])
Return type:: tuple[FitSample, …]

write_back(theta)[source]¶

Write fitted coefficients through the delegated linear fitter.

Parameters:: theta (Tensor)
Return type:: None

Alternating least-squares fitting for alchemical coefficient providers.

Use this module when spline blocks are shared through proxy coefficients and the proxy tensors and mixing weights should be fit in alternating subproblems.

class ufp.leastsquares.alchemical.AlchemicalALSFitter(model, *, fit_energy=True, fit_forces=True, fit_per_atom_energy=False, solver='cg', ridge=0.0, onebody_ridge=None, pair_ridge=None, twobody_ridge=None, threebody_ridge=None, twobody_shape_penalty=None, weight_ridge=None, dtype=None, device=None, cg_tolerance=1e-10, cg_max_iter=None, threebody_lstsq_backend=None, threebody_bucket_backend=None, max_sweeps=10, tolerance=1e-08)[source]¶

Bases: object

Alternating least-squares driver for models with shared alchemical providers.

Parameters:

model (UFPModel)
fit_energy (bool)
fit_forces (bool)
fit_per_atom_energy (bool)
solver (str)
ridge (float)
onebody_ridge (float | None)
pair_ridge (float | None)
twobody_ridge (float | None)
threebody_ridge (float | None)
twobody_shape_penalty (TwoBodySplineShapePenalty | None)
weight_ridge (float | None)
dtype (Optional[torch.dtype])
device (Optional[torch.device])
cg_tolerance (float)
cg_max_iter (int | None)
threebody_lstsq_backend (str | None)
threebody_bucket_backend (str | None)
max_sweeps (int)
tolerance (float)

fit(samples, *, batch_size=32, cache_directory=None, cache_mode='auto', initialize='svd', checkpoint_directory=None, checkpoint_frequency=1, cg_checkpoint_directory=None, cg_checkpoint_frequency=1, cg_resume=False, progress=False, progress_frequency=10)[source]¶

Alternate proxy and weight solves until convergence.

Parameters:

samples (Sequence[FitSample])
batch_size (int)
cache_directory (Path | str | None)
cache_mode (Literal['auto', 'read', 'write', 'refresh'])
initialize (Literal['svd', 'current'])
checkpoint_directory (Path | str | None)
checkpoint_frequency (int)
cg_checkpoint_directory (Path | str | None)
cg_checkpoint_frequency (int)
cg_resume (bool)
progress (bool)
progress_frequency (int)

Return type:

AlchemicalALSResult

initialize_from_direct_cg_checkpoint(checkpoint_path)[source]¶

Initialize alchemical coefficients from a direct true-CG checkpoint.

Parameters:: checkpoint_path (Path | str)
Return type:: None

class ufp.leastsquares.alchemical.AlchemicalALSResult(theta, objective_history, converged, sweeps, layout, problem, interrupted=False, restored_checkpoint_path=None)[source]¶

Bases: object

Summary of one alternating least-squares fit over alchemical coefficients.

Parameters:

theta (Tensor)
objective_history (tuple[float, ...])
converged (bool)
sweeps (int)
layout (ParameterLayout)
problem (BlockLinearProblem)
interrupted (bool)
restored_checkpoint_path (str | None)

Uncertainty¶

Status: Stable user diagnostics API.

Predictive uncertainty helpers for coefficient-linear UFP models.

This module builds dense Bayesian posteriors from the same weighted design matrices used by ufp.leastsquares.LinearFitter, then evaluates epistemic variances from sparse prediction rows. Alchemical providers are handled by freezing the learned mixing weights and forming one fixed-weight linear problem over direct coefficients and proxy coefficients.

class ufp.uncertainty.AleatoricFeatureSpec(kind='log_num_atoms')[source]¶

Bases: object

Prediction-time feature contract for aleatoric spline heads.

Parameters:: kind (str)

feature_for_atoms(atoms)[source]¶

Return one scalar feature for a structure.

Parameters:: atoms (Atoms)
Return type:: float

classmethod from_payload(payload)[source]¶

Deserialize a feature spec.

Parameters:: payload (object)
Return type:: AleatoricFeatureSpec

to_payload()[source]¶

Serialize this feature spec to JSON-compatible metadata.

Return type:: dict[str, object]

class ufp.uncertainty.AleatoricRowMetadata(kinds, features)[source]¶

Bases: object

Target-row kinds and features used to train aleatoric heads.

Parameters:

kinds (tuple[str, ...])
features (Tensor)

class ufp.uncertainty.BayesianLinearPosterior(theta_mean, Sigma_theta, metadata=<factory>, layout=None)[source]¶

Bases: object

Dense coefficient posterior for a coefficient-linear UFP solve layout.

Parameters:

theta_mean (Tensor)
Sigma_theta (Tensor | ndarray)
metadata (dict[str, object])
layout (BlockProblemLayout | None)

covariance_tensor(*, dtype=None, device=None)[source]¶

Return Sigma_theta as a torch tensor.

Parameters:

dtype (dtype | None)
device (device | None)

Return type:

Tensor

property device: device¶: Return the posterior mean device.

property dtype: dtype¶: Return the posterior mean dtype.

classmethod load(path, *, expected_metadata=None, map_location='cpu')[source]¶

Load a posterior checkpoint and validate optional metadata.

Parameters:

path (Path | str)
expected_metadata (Mapping[str, object] | None)
map_location (object)

Return type:

classmethod load_memmap(directory, *, expected_metadata=None)[source]¶

Load a memory-mapped posterior saved by save_memmap().

Parameters:

directory (Path | str)
expected_metadata (Mapping[str, object] | None)

Return type:

SplineAleatoricNoiseBundle

property n_parameters: int¶: Return the posterior parameter dimension.

save(path)[source]¶

Save this posterior to a torch checkpoint file.

Parameters:: path (Path | str)
Return type:: None

save_memmap(directory)[source]¶

Save covariance as a memory-mapped .npy plus metadata files.

Parameters:: directory (Path | str)
Return type:: None

class ufp.uncertainty.SparseLinearRow(indices, values, size)[source]¶

Bases: object

One sparse row in a coefficient solve layout.

Parameters:

indices (Tensor)
values (Tensor)
size (int)

property device: device¶: Return the row value device.

property dtype: dtype¶: Return the row value dtype.

matvec(theta)[source]¶

Apply this row to one coefficient vector.

Parameters:: theta (Tensor)
Return type:: Tensor

to_dense(*, dtype=None, device=None)[source]¶

Materialize the sparse row as one dense vector.

Parameters:

dtype (dtype | None)
device (device | None)

Return type:

Tensor

class ufp.uncertainty.SparsePredictionRows(atomic_energy_rows, total_energy_row, force_rows=None)[source]¶

Bases: object

Sparse rows needed for diagonal predictive variances.

Parameters:

atomic_energy_rows (tuple[SparseLinearRow, ...])
total_energy_row (SparseLinearRow)
force_rows (tuple[tuple[SparseLinearRow, SparseLinearRow, SparseLinearRow], ...] | None)

property force_component_rows: tuple[SparseLinearRow, ...]¶: Return flattened force-component rows, if present.

property n_atoms: int¶: Return the number of atomic energy rows.

class ufp.uncertainty.SplineAleatoricNoiseBundle(feature_spec=<factory>, energy_per_atom=None, per_atom_energy=None, force_component=None)[source]¶

Bases: object

Optional spline variance heads for UFP prediction targets.

Parameters:

feature_spec (AleatoricFeatureSpec)
energy_per_atom (SplineAleatoricNoiseModel | None)
per_atom_energy (SplineAleatoricNoiseModel | None)
force_component (SplineAleatoricNoiseModel | None)

classmethod from_payload(payload, *, dtype=None)[source]¶

Load a fitted aleatoric bundle from to_payload() output.

Parameters:

payload (object)
dtype (dtype | None)

Return type:

has_heads()[source]¶

Return whether this bundle contains any active spline head.

Return type:: bool

head_for_kind(kind)[source]¶

Return the variance head associated with one target kind.

Parameters:: kind (str)
Return type:: SplineAleatoricNoiseModel | None

predict_for_atoms(atoms, *, include_forces, dtype, device)[source]¶

Return structure-dependent aleatoric variances for one structure.

Parameters:

atoms (Atoms)
include_forces (bool)
dtype (dtype)
device (device)

Return type:

tuple[Tensor | None, Tensor | None, Tensor | None]

to_payload()[source]¶

Return a torch-serializable payload for this aleatoric bundle.

Return type:: dict[str, object]

class ufp.uncertainty.SplineAleatoricNoiseModel(*, n_coefficients, lower_full_support, upper_full_support, spline='cubic', variance_floor=1e-12, initial_raw=-6.0, dtype=None)[source]¶

Bases: Module

Positive 1D spline variance head using softplus(raw) + floor.

Parameters:

n_coefficients (int)
lower_full_support (float)
upper_full_support (float)
spline (str)
variance_floor (float)
initial_raw (float | torch.Tensor)
dtype (torch.dtype | None)

fit_residuals(residuals, x, *, steps=500, lr=0.01)[source]¶

Optimize this variance head against residuals with Adam.

Parameters:

residuals (Tensor)
x (Tensor)
steps (int)
lr (float)

Return type:

tuple[float, …]

forward(x)[source]¶

Return positive variances at scalar features x.

Parameters:: x (Tensor)
Return type:: Tensor

classmethod from_payload(payload, *, dtype=None)[source]¶

Load a fitted noise head from to_payload() output.

Parameters:

payload (object)
dtype (dtype | None)

Return type:

SplineAleatoricNoiseModel

gaussian_nll(residuals, x, *, reduction='mean', include_constant=False)[source]¶

Return Gaussian negative log likelihood for residuals.

Parameters:

residuals (Tensor)
x (Tensor)
reduction (Literal['mean', 'sum', 'none'])
include_constant (bool)

Return type:

Tensor

raw(x)[source]¶

Evaluate the unconstrained spline head at scalar features x.

Parameters:: x (Tensor)
Return type:: Tensor

to_payload()[source]¶

Return a torch-serializable payload for this fitted noise head.

Return type:: dict[str, object]

class ufp.uncertainty.UFPUncertaintyOutput(means, energy_epistemic_variance=None, per_atom_energy_epistemic_variance=None, force_epistemic_variance=None, energy_aleatoric_variance=None, per_atom_energy_aleatoric_variance=None, force_aleatoric_variance=None, energy_total_variance=None, per_atom_energy_total_variance=None, force_total_variance=None, rows=None)[source]¶

Bases: object

UFP predictions with epistemic, aleatoric, and total variances.

Parameters:

means (UFPOutput)
energy_epistemic_variance (Tensor | None)
per_atom_energy_epistemic_variance (Tensor | None)
force_epistemic_variance (Tensor | None)
energy_aleatoric_variance (Tensor | None)
per_atom_energy_aleatoric_variance (Tensor | None)
force_aleatoric_variance (Tensor | None)
energy_total_variance (Tensor | None)
per_atom_energy_total_variance (Tensor | None)
force_total_variance (Tensor | None)
rows (SparsePredictionRows | None)

class ufp.uncertainty.UncertaintyPredictionBundle(model, posterior, aleatoric_variance=None, aleatoric_noise_model=None, aleatoric_prediction_feature=None, aleatoric_noise_bundle=None, energy_variance_scale=1.0, manifest=<factory>)[source]¶

Bases: object

Model, posterior, and metadata loaded from an uncertainty bundle.

Parameters:

model (UFPModel)
posterior (BayesianLinearPosterior)
aleatoric_variance (float | None)
aleatoric_noise_model (SplineAleatoricNoiseModel | None)
aleatoric_prediction_feature (float | None)
aleatoric_noise_bundle (SplineAleatoricNoiseBundle | None)
energy_variance_scale (float)
manifest (dict[str, object])

prediction_aleatoric_noise_bundle()[source]¶

Return the structure-dependent aleatoric noise bundle, if present.

Return type:: SplineAleatoricNoiseBundle | None

prediction_aleatoric_variance()[source]¶

Return the scalar aleatoric variance to use for generic predictions.

Return type:: float | None

ufp.uncertainty.build_aleatoric_row_metadata(samples, *, fit_energy, fit_forces, fit_per_atom_energy, dtype=None, device=None, feature_spec=None)[source]¶

Return target-row kinds and features in least-squares row order.

Parameters:

samples (Sequence[FitSample])
fit_energy (bool)
fit_forces (bool)
fit_per_atom_energy (bool)
dtype (dtype | None)
device (device | None)
feature_spec (AleatoricFeatureSpec | None)

Return type:

AleatoricRowMetadata

ufp.uncertainty.build_prediction_rows(model, atoms, *, posterior=None, fitter=None, neighbor_list=None, include_forces=True, dtype=None, device=None)[source]¶

Build sparse atomic, total-energy, and optional force prediction rows.

Parameters:

model (UFPModel)
atoms (Atoms)
posterior (BayesianLinearPosterior | None)
fitter (LinearFitter | None)
neighbor_list (NeighborListData | None)
include_forces (bool)
dtype (dtype | None)
device (device | None)

Return type:

SparsePredictionRows

ufp.uncertainty.combine_total_energy_rows(rows)[source]¶

Return a sparse linear combination of total-energy rows.

Parameters:: rows (Sequence[tuple[SparseLinearRow, float]])
Return type:: SparseLinearRow

ufp.uncertainty.describe_uncertainty_prediction_bundle(directory, *, dtype=None, validate=True)[source]¶

Return a compact human-readable summary of an uncertainty bundle.

Parameters:

directory (Path | str)
dtype (dtype | None)
validate (bool)

Return type:

dict[str, object]

ufp.uncertainty.fit_alchemical_uncertainty_model(model, samples, *, mean_source='checkpoint', checkpoint_path=None, fixed_weight_refit=False, row_variances=None, observation_noise_variance=1.0, aleatoric_noise_model=None, aleatoric_features=None, aleatoric_noise_bundle=None, aleatoric_row_metadata=None, aleatoric_steps=0, aleatoric_lr=0.01, batch_size=32, progress=False, cache_directory=None, cache_mode='auto', covariance_path=None, jitter=0.0, als_fitter=None, **fitter_kwargs)[source]¶

Build a fixed-weight Bayesian posterior for an alchemical UFP model.

Parameters:

model (UFPModel)
samples (Sequence[FitSample])
mean_source (Literal['checkpoint', 'als'])
checkpoint_path (Path | str | None)
fixed_weight_refit (bool)
row_variances (object)
observation_noise_variance (float)
aleatoric_noise_model (SplineAleatoricNoiseModel | None)
aleatoric_features (object)
aleatoric_noise_bundle (SplineAleatoricNoiseBundle | None)
aleatoric_row_metadata (AleatoricRowMetadata | None)
aleatoric_steps (int)
aleatoric_lr (float)
batch_size (int)
progress (bool)
cache_directory (Path | str | None)
cache_mode (Literal['auto', 'read', 'write', 'refresh'])
covariance_path (Path | str | None)
jitter (float)
als_fitter (AlchemicalALSFitter | None)
fitter_kwargs (Any)

Return type:

ufp.uncertainty.fit_linear_uncertainty_model(model, samples, *, fitter=None, problem=None, refit_mean=True, write_back=True, row_variances=None, observation_noise_variance=1.0, aleatoric_noise_model=None, aleatoric_features=None, aleatoric_noise_bundle=None, aleatoric_row_metadata=None, aleatoric_steps=0, aleatoric_lr=0.01, batch_size=32, progress=False, cache_directory=None, cache_mode='auto', covariance_path=None, jitter=0.0, **fitter_kwargs)[source]¶

Fit or load the mean and build a dense linear Bayesian posterior.

Parameters:

model (UFPModel)
samples (Sequence[FitSample])
fitter (LinearFitter | None)
problem (BlockLinearProblem | None)
refit_mean (bool)
write_back (bool)
row_variances (object)
observation_noise_variance (float)
aleatoric_noise_model (SplineAleatoricNoiseModel | None)
aleatoric_features (object)
aleatoric_noise_bundle (SplineAleatoricNoiseBundle | None)
aleatoric_row_metadata (AleatoricRowMetadata | None)
aleatoric_steps (int)
aleatoric_lr (float)
batch_size (int)
progress (bool)
cache_directory (Path | str | None)
cache_mode (Literal['auto', 'read', 'write', 'refresh'])
covariance_path (Path | str | None)
jitter (float)
fitter_kwargs (Any)

Return type:

UncertaintyPredictionBundle

ufp.uncertainty.load_uncertainty_prediction_bundle(directory, *, dtype=None, validate=True)[source]¶

Load a reusable model/posterior bundle for predictive uncertainties.

Parameters:

directory (Path | str)
dtype (dtype | None)
validate (bool)

Return type:

ufp.uncertainty.predict_with_uncertainty(model, atoms, posterior, *, rows=None, fitter=None, neighbor_list=None, include_forces=True, aleatoric_variance=None, aleatoric_noise_bundle=None, variance_chunk_size=256, return_rows=False)[source]¶

Predict means and diagonal epistemic/total variances for one structure.

Parameters:

model (UFPModel)
atoms (Atoms)
posterior (BayesianLinearPosterior)
rows (SparsePredictionRows | None)
fitter (LinearFitter | None)
neighbor_list (NeighborListData | None)
include_forces (bool)
aleatoric_variance (float | Tensor | None)
aleatoric_noise_bundle (SplineAleatoricNoiseBundle | None)
variance_chunk_size (int)
return_rows (bool)

Return type:

UFPUncertaintyOutput

ufp.uncertainty.save_energy_variance_scale_to_bundle(directory, scale, *, metadata=None)[source]¶

Save or replace the post-hoc energy variance scale in a bundle.

Parameters:

directory (Path | str)
scale (float)
metadata (Mapping[str, object] | None)

Return type:

dict[str, object]

ufp.uncertainty.save_uncertainty_prediction_bundle(directory, *, model, posterior, source_checkpoint=None, aleatoric_variance=None, aleatoric_noise_model=None, aleatoric_prediction_feature=None, aleatoric_noise_bundle=None, energy_variance_scale=None, calibration_metadata=None, metadata=None)[source]¶

Save a reusable model/posterior bundle for predictive uncertainties.

Parameters:

directory (Path | str)
model (UFPModel)
posterior (BayesianLinearPosterior)
source_checkpoint (Path | str | None)
aleatoric_variance (float | None)
aleatoric_noise_model (SplineAleatoricNoiseModel | None)
aleatoric_prediction_feature (float | None)
aleatoric_noise_bundle (SplineAleatoricNoiseBundle | None)
energy_variance_scale (float | None)
calibration_metadata (Mapping[str, object] | None)
metadata (Mapping[str, object] | None)

Return type:

dict[str, object]

ufp.uncertainty.validate_uncertainty_prediction_bundle(directory, *, dtype=None)[source]¶

Return True when a bundle loads and validates successfully.

Parameters:

directory (Path | str)
dtype (dtype | None)

Return type:

bool

ufp.uncertainty.variance_for_energy_row(row, posterior)[source]¶

Return row @ Sigma_theta @ row.T for one sparse or dense row.

Parameters:

row (SparseLinearRow | Tensor)
posterior (BayesianLinearPosterior)

Return type:

Tensor

ufp.uncertainty.variance_for_sparse_rows(rows, posterior, *, chunk_size=256)[source]¶

Return diag(A @ Sigma_theta @ A.T) for sparse/dense prediction rows.

Parameters:

rows (Sequence[SparseLinearRow | Tensor])
posterior (BayesianLinearPosterior)
chunk_size (int)

Return type:

Tensor

Coefficients¶

Status: Stable user API.

Public helpers for moving coefficient channels between compatible UFP models.

Coefficient-channel interchange helpers for compatible UFP models.

class ufp.coefficients.interchange.CoefficientChannel(block, family, channel, index, value_shape, grid, cutoff, category_order, active_channels, symmetric)[source]¶

Bases: object

Resolved physical channel inside one model coefficient block.

Parameters:

block (TermBlock)
family (str)
channel (tuple[int, ...])
index (tuple[int | slice, ...])
value_shape (tuple[int, ...])
grid (tuple[object, ...])
cutoff (tuple[object, ...])
category_order (tuple[tuple[int, ...], ...])
active_channels (tuple[tuple[int, ...], ...])
symmetric (bool | None)

property block_kind: str¶: Return the owning coefficient-block kind.

property block_label: str¶: Return the owning coefficient-block label.

property device: device¶: Return the channel tensor device.

property dtype: dtype¶: Return the channel tensor dtype.

property term_type: str¶: Return the owning term class name.

class ufp.coefficients.interchange.CoefficientCompatibility(source, target, reasons)[source]¶

Bases: object

Compatibility result for one source/target channel pair.

Parameters:

source (CoefficientChannel)
target (CoefficientChannel)
reasons (tuple[str, ...])

property compatible: bool¶: Return whether there are no incompatibility reasons.

exception ufp.coefficients.interchange.CoefficientCompatibilityError[source]¶

Bases: ValueError

Raised when coefficient channels can not be safely interchanged.

class ufp.coefficients.interchange.CoefficientCopy(family, channel, source_block_label, target_block_label, shape)[source]¶

Bases: object

Summary of one copied coefficient channel.

Parameters:

family (str)
channel (tuple[int, ...])
source_block_label (str)
target_block_label (str)
shape (tuple[int, ...])

class ufp.coefficients.interchange.CoefficientCopyReport(copied, skipped)[source]¶

Bases: object

Summary of a copy-matching-coefficients operation.

Parameters:

copied (tuple[CoefficientCopy, ...])
skipped (tuple[str, ...])

property copied_count: int¶: Return the number of copied channels.

ufp.coefficients.interchange.clone_model_with_copied_coefficients(model)[source]¶

Deep-copy a model architecture and coefficient values.

Parameters:: model (UFPModel)
Return type:: UFPModel

ufp.coefficients.interchange.clone_model_with_zeroed_coefficients(model)[source]¶

Deep-copy a model architecture and zero all copied coefficient blocks.

Parameters:: model (UFPModel)
Return type:: UFPModel

ufp.coefficients.interchange.copy_coefficient_channel(source_model, target_model, selector, *, target_selector=None)[source]¶

Copy one compatible source coefficient channel into a target model.

Parameters:

source_model (UFPModel)
target_model (UFPModel)
selector (CoefficientSelector)
target_selector (CoefficientSelector | None)

Return type:

CoefficientCopy

ufp.coefficients.interchange.copy_matching_coefficients(source_model, target_model, *, strict=True)[source]¶

Copy every matching compatible pair/triplet channel into target_model.

Parameters:

source_model (UFPModel)
target_model (UFPModel)
strict (bool)

Return type:

CoefficientCopyReport

ufp.coefficients.interchange.read_coefficient_channel(model, selector)[source]¶

Read one physical pair or triplet coefficient channel.

Parameters:

model (UFPModel)
selector (CoefficientSelector)

Return type:

Tensor

ufp.coefficients.interchange.validate_coefficient_compatibility(source_model, target_model, selector, *, target_selector=None)[source]¶

Validate that one source channel can be copied into one target channel.

Parameters:

source_model (UFPModel)
target_model (UFPModel)
selector (CoefficientSelector)
target_selector (CoefficientSelector | None)

Return type:

CoefficientCompatibility

ufp.coefficients.interchange.write_coefficient_channel(model, selector, values)[source]¶

Write one physical pair or triplet coefficient channel.

Parameters:

model (UFPModel)
selector (CoefficientSelector)
values (Tensor)

Return type:

None

ufp.coefficients.interchange.zero_model_coefficients(model)[source]¶

Zero all fittable coefficient blocks in-place and return model.

Parameters:: model (UFPModel)
Return type:: UFPModel

Projection¶

Status: Stable user API.

Offline coefficient projection helpers and diagnostics.

Offline projection of radial callables onto 1D spline coefficient rows.

class ufp.projection.radial.RadialProjectionResult(coeffs, spline, coeff_size, full_support_start, cutoff, first_knot, knot_spacing, sample_distances, sample_weights, design_rank, diagnostic_distances, diagnostic_weights, target_values, predicted_values, residuals, weighted_rmse, diagnostics, metadata=None)[source]¶

Bases: object

Projected spline coefficients and associated projection metadata.

Parameters:

coeffs (Tensor)
spline (str)
coeff_size (int)
full_support_start (float)
cutoff (float)
first_knot (float)
knot_spacing (float)
sample_distances (Tensor)
sample_weights (Tensor | None)
design_rank (int)
diagnostic_distances (Tensor)
diagnostic_weights (Tensor | None)
target_values (Tensor)
predicted_values (Tensor)
residuals (Tensor)
weighted_rmse (float | None)
diagnostics (ProjectionDiagnostics)
metadata (Mapping[str, object] | None)

ufp.projection.radial.project_pair_prior(prior, pair, *, coeff_size, full_support_start=0.0, cutoff=None, spline='cubic', n_samples=None, sample_distances=None, sample_weights=None, diagnostic_distances=None, diagnostic_weights=None, dtype=None, device=None, rcond=None, channel_label=None)[source]¶

Project one analytic pair prior channel onto a 1D spline coefficient row.

The prior must expose radial_values(pair, distances). If it also exposes projection_metadata(), that metadata is attached to the returned result alongside the projection grid settings and pair channel.

Parameters:

prior (object)
pair (Sequence[int])
coeff_size (int)
full_support_start (float)
cutoff (float | None)
spline (str)
n_samples (int | None)
sample_distances (object | None)
sample_weights (object | None)
diagnostic_distances (object | None)
diagnostic_weights (object | None)
dtype (dtype | None)
device (device | str | None)
rcond (float | None)
channel_label (str | None)

Return type:

RadialProjectionResult

ufp.projection.radial.project_radial_function(function, *, coeff_size, full_support_start, cutoff, spline='cubic', n_samples=None, sample_distances=None, sample_weights=None, diagnostic_distances=None, diagnostic_weights=None, dtype=None, device=None, rcond=None, channel_label='radial')[source]¶

Project a radial callable onto a uniform 1D pair-spline coefficient row.

Parameters:

function (Callable[[Tensor], object]) – Callable accepting a distance tensor and returning values with the same shape.
coeff_size (int) – Number of 1D spline coefficients in the target row.
full_support_start (float) – Lower distance with full spline support.
cutoff (float) – Upper distance with full spline support and pair-term cutoff.
spline (str) – Spline family name.
n_samples (int | None) – Number of deterministic midpoint samples when sample_distances is not supplied.
sample_distances (object | None) – Optional explicit fitting radii.
sample_weights (object | None) – Optional non-negative least-squares row weights.
diagnostic_distances (object | None) – Optional radii for error diagnostics. Defaults to the fitting radii.
diagnostic_weights (object | None) – Optional non-negative diagnostic weights. Defaults to sample_weights when diagnostics reuse the fitting radii.
dtype (dtype | None) – Floating-point dtype for projection tensors.
device (device | str | None) – Device for projection tensors.
rcond (float | None) – Optional cutoff passed to torch.linalg.lstsq().
channel_label (str) – Diagnostic label for the projected coefficient channel.

Returns:

Projected coefficients and diagnostic data.

Return type:

RadialProjectionResult

Offline projection helpers for uniform spline coefficient blocks.

exception ufp.projection.spline.SplineProjectionError[source]¶

Bases: ValueError

Raised when spline coefficients can not be projected safely.

class ufp.projection.spline.SplineProjectionResult(coeffs, sample_coordinates, source_values, projected_values, rmse, max_abs_error, diagnostics, source_gradients=None, projected_gradients=None, gradient_rmse=None, gradient_max_abs_error=None, source_shape=(), target_shape=(), channel=None)[source]¶

Bases: object

Projected coefficients and sample-space diagnostics.

Parameters:

coeffs (Tensor)
sample_coordinates (tuple[Tensor, ...])
source_values (Tensor)
projected_values (Tensor)
rmse (float)
max_abs_error (float)
diagnostics (ProjectionDiagnostics)
source_gradients (tuple[Tensor, ...] | None)
projected_gradients (tuple[Tensor, ...] | None)
gradient_rmse (float | None)
gradient_max_abs_error (float | None)
source_shape (tuple[int, ...])
target_shape (tuple[int, ...])
channel (tuple[int, ...] | None)

ufp.projection.spline.evaluate_uniform_spline_1d(coeffs, distances, *, first_knot, knot_spacing, spline)[source]¶

Evaluate 1D uniform spline values and distance derivatives.

Parameters:

coeffs (Tensor)
distances (Tensor)
first_knot (float)
knot_spacing (float)
spline (str)

Return type:

tuple[Tensor, Tensor]

ufp.projection.spline.evaluate_uniform_spline_2d(coeffs, x, y, *, first_knot_x, first_knot_y, knot_spacing_x, knot_spacing_y, spline)[source]¶

Evaluate 2D uniform spline values and coordinate derivatives.

Parameters:

coeffs (Tensor)
x (Tensor)
y (Tensor)
first_knot_x (float)
first_knot_y (float)
knot_spacing_x (float)
knot_spacing_y (float)
spline (str)

Return type:

tuple[Tensor, Tensor, Tensor]

ufp.projection.spline.evaluate_uniform_spline_3d(coeffs, x, y, z, *, first_knot_xy, first_knot_z, knot_spacing_xy, knot_spacing_z, spline)[source]¶

Evaluate 3D uniform spline values and coordinate derivatives.

Parameters:

coeffs (Tensor)
x (Tensor)
y (Tensor)
z (Tensor)
first_knot_xy (float)
first_knot_z (float)
knot_spacing_xy (float)
knot_spacing_z (float)
spline (str)

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

ufp.projection.spline.project_pair_to_twobody(source, target, *, pair=None, n_samples=None, sample_distances=None, derivative_weight=1.0, rcond=None, write=True)[source]¶

Project a pair-specific spline into one categorized two-body row.

Parameters:

source (SplinePairTerm)
target (SplineTwoBodyTerm)
pair (Sequence[int] | None)
n_samples (int | None)
sample_distances (Tensor | None)
derivative_weight (float)
rcond (float | None)
write (bool)

Return type:

ufp.projection.spline.project_threebody_channel(source, target, channel, *, n_samples_per_axis=None, sample_coordinates=None, derivative_weight=1.0, rcond=None, write=True)[source]¶

Project one same-family 3D triplet channel between compatible terms.

Parameters:

source (SplineThreeBodyTerm)
target (SplineThreeBodyTerm)
channel (Sequence[int])
n_samples_per_axis (int | Sequence[int] | None)
sample_coordinates (Sequence[Tensor] | None)
derivative_weight (float)
rcond (float | None)
write (bool)

Return type:

ufp.projection.spline.project_triplet2d_channel(source, target, channel, *, n_samples_per_axis=None, sample_coordinates=None, derivative_weight=1.0, rcond=None, write=True)[source]¶

Project one same-family 2D triplet channel between compatible terms.

Parameters:

source (SplineTriplet2DTerm)
target (SplineTriplet2DTerm)
channel (Sequence[int])
n_samples_per_axis (int | Sequence[int] | None)
sample_coordinates (Sequence[Tensor] | None)
derivative_weight (float)
rcond (float | None)
write (bool)

Return type:

ufp.projection.spline.project_uniform_spline_1d(source_coeffs, *, source_lower_full_support, source_upper_full_support, target_coeff_size, target_lower_full_support=None, target_upper_full_support=None, source_spline='cubic', target_spline=None, n_samples=None, sample_distances=None, derivative_weight=1.0, rcond=None)[source]¶

Project one 1D uniform spline onto another 1D uniform spline basis.

The target interval must be contained in the source interval. The result uses the source coefficient dtype and device.

Parameters:

source_coeffs (Tensor)
source_lower_full_support (float)
source_upper_full_support (float)
target_coeff_size (int)
target_lower_full_support (float | None)
target_upper_full_support (float | None)
source_spline (str)
target_spline (str | None)
n_samples (int | None)
sample_distances (Tensor | None)
derivative_weight (float)
rcond (float | None)

Return type:

ufp.projection.spline.project_uniform_spline_2d(source_coeffs, *, source_lower_full_support, source_upper_full_support, target_coeff_shape, target_lower_full_support=None, target_upper_full_support=None, source_spline='cubic', target_spline=None, n_samples_per_axis=None, sample_coordinates=None, derivative_weight=1.0, rcond=None)[source]¶

Project one square 2D uniform spline grid onto another grid.

Parameters:

source_coeffs (Tensor)
source_lower_full_support (float)
source_upper_full_support (float)
target_coeff_shape (Sequence[int])
target_lower_full_support (float | None)
target_upper_full_support (float | None)
source_spline (str)
target_spline (str | None)
n_samples_per_axis (int | Sequence[int] | None)
sample_coordinates (Sequence[Tensor] | None)
derivative_weight (float)
rcond (float | None)

Return type:

ufp.projection.spline.project_uniform_spline_3d(source_coeffs, *, source_lower_full_support_xy, source_upper_full_support_xy, source_lower_full_support_z, source_upper_full_support_z, target_coeff_shape, target_lower_full_support_xy=None, target_upper_full_support_xy=None, target_lower_full_support_z=None, target_upper_full_support_z=None, source_spline='cubic', target_spline=None, n_samples_per_axis=None, sample_coordinates=None, derivative_weight=1.0, rcond=None)[source]¶

Project one 3D uniform spline grid onto another grid.

Parameters:

source_coeffs (Tensor)
source_lower_full_support_xy (float)
source_upper_full_support_xy (float)
source_lower_full_support_z (float)
source_upper_full_support_z (float)
target_coeff_shape (Sequence[int])
target_lower_full_support_xy (float | None)
target_upper_full_support_xy (float | None)
target_lower_full_support_z (float | None)
target_upper_full_support_z (float | None)
source_spline (str)
target_spline (str | None)
n_samples_per_axis (int | Sequence[int] | None)
sample_coordinates (Sequence[Tensor] | None)
derivative_weight (float)
rcond (float | None)

Return type:

ProjectionChannelDiagnostic

Projection diagnostic containers.

The objects in this module are intentionally independent of projection algorithms and plotting libraries. Projection code can return them directly, and workflow/checkpoint code can serialize them through to_dict().

class ufp.projection.diagnostics.ProjectionChannelDiagnostic(channel_label, sample_count, support_coverage=None, value_error=None, derivative_error=None, warnings=<factory>, failure_reason=None)[source]¶

Bases: object

Projection diagnostics for one physical or logical coefficient channel.

Parameters:

channel_label (str)
sample_count (int)
support_coverage (ProjectionSupportCoverage | None)
value_error (ProjectionErrorSummary | None)
derivative_error (ProjectionErrorSummary | None)
warnings (tuple[str, ...])
failure_reason (str | None)

property failed: bool¶: Return whether projection failed for this channel.

classmethod failure(*, channel_label, failure_reason, sample_count=0, warnings=(), support_coverage=None)[source]¶

Build a failed channel diagnostic with a required failure reason.

Parameters:

channel_label (str)
failure_reason (str)
sample_count (int)
warnings (Sequence[str])
support_coverage (ProjectionSupportCoverage | None)

Return type:

property succeeded: bool¶: Return whether projection succeeded for this channel.

to_dict()[source]¶

Return JSON/checkpoint-friendly channel diagnostic metadata.

Return type:: dict[str, object]

class ufp.projection.diagnostics.ProjectionDiagnostics(channels=<factory>)[source]¶

Bases: object

Collection of channel diagnostics returned by one projection operation.

Parameters:: channels (tuple[ProjectionChannelDiagnostic, ...])

property failed_channels: tuple[ProjectionChannelDiagnostic, ...]¶: Return channels whose projection failed.

property summary: ProjectionDiagnosticsSummary¶: Return aggregate diagnostics across all channels.

to_dict()[source]¶

Return JSON/checkpoint-friendly diagnostics metadata.

Return type:: dict[str, object]

property warning_messages: tuple[str, ...]¶: Return channel-qualified warning messages.

class ufp.projection.diagnostics.ProjectionDiagnosticsSummary(channel_count, successful_channel_count, failed_channel_count, sample_count, covered_count, support_sample_count, support_coverage_fraction, max_value_rmse, max_value_max_absolute, max_derivative_rmse, max_derivative_max_absolute, warning_count, failed_channel_labels=<factory>, warning_messages=<factory>)[source]¶

Bases: object

Aggregate projection diagnostics across channels.

Parameters:

channel_count (int)
successful_channel_count (int)
failed_channel_count (int)
sample_count (int)
covered_count (int)
support_sample_count (int)
support_coverage_fraction (float | None)
max_value_rmse (float | None)
max_value_max_absolute (float | None)
max_derivative_rmse (float | None)
max_derivative_max_absolute (float | None)
warning_count (int)
failed_channel_labels (tuple[str, ...])
warning_messages (tuple[str, ...])

to_dict()[source]¶

Return JSON/checkpoint-friendly aggregate metadata.

Return type:: dict[str, object]

class ufp.projection.diagnostics.ProjectionErrorSummary(sample_count, rmse=None, mean_absolute=None, max_absolute=None)[source]¶

Bases: object

Scalar error metrics for projected values or derivatives.

Parameters:

sample_count (int)
rmse (float | None)
mean_absolute (float | None)
max_absolute (float | None)

classmethod from_residuals(residuals)[source]¶

Build error metrics from projection residuals.

Parameters:: residuals (Sequence[float])
Return type:: ProjectionErrorSummary

to_dict()[source]¶

Return JSON/checkpoint-friendly error metadata.

Return type:: dict[str, int | float | None]

class ufp.projection.diagnostics.ProjectionSupportCoverage(sample_count, covered_count, lower_bound=None, upper_bound=None, minimum_sample=None, maximum_sample=None)[source]¶

Bases: object

Coverage of sampled points by a projection support interval or domain.

Parameters:

sample_count (int)
covered_count (int)
lower_bound (float | None)
upper_bound (float | None)
minimum_sample (float | None)
maximum_sample (float | None)

property coverage_fraction: float | None¶: Return covered samples divided by total samples, or None if empty.

classmethod from_samples(samples, *, lower_bound=None, upper_bound=None)[source]¶

Build coverage metadata from sampled coordinates and optional bounds.

Parameters:

samples (Sequence[float])
lower_bound (float | None)
upper_bound (float | None)

Return type:

ProjectionSupportCoverage

property fully_covered: bool¶: Return whether every sampled point was inside support.

to_dict()[source]¶

Return JSON/checkpoint-friendly support coverage metadata.

Return type:: dict[str, int | float | None | bool]

ufp.projection.diagnostics.aggregate_channel_diagnostics(channels)[source]¶

Aggregate per-channel projection diagnostics into one summary.

Parameters:: channels (Sequence[ProjectionChannelDiagnostic])
Return type:: ProjectionDiagnosticsSummary

ufp.projection.diagnostics.ensure_projection_succeeded(diagnostics)[source]¶

Raise ValueError when any channel diagnostic reports a failure.

Parameters:: diagnostics (ProjectionDiagnostics | Sequence[ProjectionChannelDiagnostic])
Return type:: None

Training¶

Status: Stable user API.

ASE-native supervised training utilities.

The training package keeps structures intact in datasets, flattens them during batch collation, prepares UFPInput objects, and runs ordinary PyTorch training and evaluation loops for energy, force, and stress targets.

Dataset and split helpers for supervised ase.Atoms samples.

Use this module to normalize energies, forces, stresses, and metadata into a small PyTorch-friendly dataset layer.

class ufp.training.dataset.ASEAtomsDataset(samples)[source]¶

Bases: Dataset[ASEAtomsSample]

Thin torch.utils.data dataset for supervised ASE structures.

Parameters:: samples (Sequence[ASEAtomsSample])

cache_neighbor_lists(*, cutoff, backend=NeighborListBackend.AUTO)[source]¶

Warm neighbor-list caches for every sample for one cutoff/backend pair.

Parameters:

cutoff (float)
backend (str | NeighborListBackend)

Return type:

None

cache_tensorized_samples()[source]¶

Warm the per-sample tensor cache.

Return type:: None

classmethod from_atoms(atoms_list, *, energies=None, forces=None, force_masks=None, stresses=None, neighbor_lists=None, energy_key='energy', forces_key=None, stress_key=None, metadata=None, metadata_keys=())[source]¶

Build dataset samples directly from labeled ase.Atoms objects.

Parameters:

atoms_list (Sequence[Atoms])
energies (Sequence[object] | None)
forces (Sequence[object] | None)
force_masks (Sequence[object] | None)
stresses (Sequence[object] | None)
neighbor_lists (Sequence[NeighborListData | None] | None)
energy_key (str | None)
forces_key (str | None)
stress_key (str | None)
metadata (Sequence[Mapping[str, object]] | None)
metadata_keys (Sequence[str])

Return type:

ASEAtomsDataset

class ufp.training.dataset.ASEAtomsSample(atoms, energy=None, forces=None, force_mask=None, stress=None, neighbor_list=None, metadata=<factory>)[source]¶

Bases: object

Single supervised training sample backed by ase.Atoms.

Parameters:

atoms (Atoms)
energy (float | None)
forces (Tensor | None)
force_mask (Tensor | None)
stress (Tensor | None)
neighbor_list (NeighborListData | None)
metadata (Mapping[str, object])

cached_neighbor_list(*, cutoff, backend)[source]¶

Return an explicit or cached neighbor list.

Parameters:

cutoff (float)
backend (str | NeighborListBackend)

Return type:

tensorized()[source]¶

Return and cache the tensorized structural fields for this sample.

Return type:: _ASETensorizedSample

class ufp.training.dataset.ASEDataLoaders(train, validation, test)[source]¶

Bases: object

Grouped dataloaders aligned with an ASEDatasetSplit.

Parameters:

train (DataLoader)
validation (DataLoader)
test (DataLoader)

class ufp.training.dataset.ASEDatasetSplit(train, validation, test)[source]¶

Bases: object

Grouped train, validation, and test subsets over one ASE dataset.

Parameters:

train (Subset)
validation (Subset)
test (Subset)

ufp.training.dataset.split_ase_dataset(dataset, *, train_fraction=0.8, validation_fraction=0.1, test_fraction=0.1, shuffle=True, seed=0)[source]¶

Split one ASE dataset into train, validation, and test subsets.

Parameters:

dataset (Dataset[ASEAtomsSample] | Sequence[ASEAtomsSample])
train_fraction (float)
validation_fraction (float)
test_fraction (float)
shuffle (bool)
seed (int)

Return type:

ASEDatasetSplit

Batch containers and dataloader helpers for ASE-backed training.

This module keeps batched structures and targets aligned while exposing a prepare_input bridge back into the shared model input path.

class ufp.training.batch.ASEAtomsBatch(atoms, energy=None, forces=None, force_mask=None, stress=None, neighbor_lists=None, metadata=<factory>, samples=<factory>, cached_input=None)[source]¶

Bases: object

Mini-batch of structures and targets produced by ase_atoms_collate_fn().

Parameters:

atoms (Sequence[Atoms])
energy (Tensor | None)
forces (Tensor | None)
force_mask (Tensor | None)
stress (Tensor | None)
neighbor_lists (Sequence[NeighborListData] | None)
metadata (Sequence[Mapping[str, object]])
samples (Sequence[ASEAtomsSample])
cached_input (UFPInput | None)

property atom_slices: tuple[slice, ...]¶: Return per-system slices into the concatenated atom axis.

cache_targets_on_device(*, device, dtype)[source]¶

Move supervised targets to the training device in place.

Parameters:

device (device)
dtype (dtype)

Return type:

ASEAtomsBatch

property n_atoms: int¶: Return the total atom count across the batch.

property n_systems: int¶: Return the number of structures stored in the batch.

pin_memory()[source]¶

Pin stored tensors in place for faster host-to-device transfer.

Return type:: ASEAtomsBatch

prepare_input(model, *, backend=None, device=None, dtype=None, requires_grad=False)[source]¶

Delegate structure conversion to the model’s prepare_input helper.

Parameters:

model (UFPotential)
backend (str | NeighborListBackend | None)
device (device | str | None)
dtype (dtype | None)
requires_grad (bool)

Return type:

ForceComponentSelection

class ufp.training.batch.CachedASEBatchLoader(dataset, batches, *, shuffle=False, seed=0)[source]¶

Bases: object

Reusable iterable over pre-collated batches for static ASE datasets.

Parameters:

dataset (Dataset[ASEAtomsSample])
batches (Sequence[ASEAtomsBatch])
shuffle (bool)
seed (int)

ufp.training.batch.ase_atoms_collate_fn(samples)[source]¶

Collate normalized ASE samples into one ASEAtomsBatch.

Parameters:: samples (Sequence[ASEAtomsSample])
Return type:: ASEAtomsBatch

ufp.training.batch.build_ase_dataloader(dataset, *, batch_size=1, shuffle=False, drop_last=False, num_workers=0, pin_memory=False, persistent_workers=False, prefetch_factor=None)[source]¶

Build a dataloader that uses ase_atoms_collate_fn for batching.

Parameters:

dataset (Dataset[ASEAtomsSample] | Sequence[ASEAtomsSample])
batch_size (int)
shuffle (bool)
drop_last (bool)
num_workers (int)
pin_memory (bool)
persistent_workers (bool)
prefetch_factor (int | None)

Return type:

DataLoader

ufp.training.batch.build_ase_dataloaders(split, *, batch_size=1, num_workers=0, pin_memory=False, persistent_workers=False, prefetch_factor=None, shuffle_train=True, drop_last_train=False)[source]¶

Build train, validation, and test dataloaders from one split object.

Parameters:

split (ASEDatasetSplit)
batch_size (int)
num_workers (int)
pin_memory (bool)
persistent_workers (bool)
prefetch_factor (int | None)
shuffle_train (bool)
drop_last_train (bool)

Return type:

ASEDataLoaders

ufp.training.batch.build_ase_training_loader(model, dataset, *, batch_size=1, shuffle=False, drop_last=False, num_workers=0, pin_memory=False, persistent_workers=False, prefetch_factor=None, dtype=None, backend=None, cache_batches=True, cache_batches_on_device=False, feature_cache_storage='cpu', feature_cache_mode='auto', feature_cache_dir=None, feature_cache_prefix='batch', feature_cache_per_atom_energy=True, device=None, seed=0, progress=False, progress_description='Caching training batches')[source]¶

Build a training loader, caching static geometry by default.

Parameters:

model (UFPotential)
dataset (Dataset[ASEAtomsSample] | Sequence[ASEAtomsSample])
batch_size (int)
shuffle (bool)
drop_last (bool)
num_workers (int)
pin_memory (bool)
persistent_workers (bool)
prefetch_factor (int | None)
dtype (dtype | None)
backend (str | NeighborListBackend | None)
cache_batches (bool)
cache_batches_on_device (bool)
feature_cache_storage (str)
feature_cache_mode (str)
feature_cache_dir (Path | str | None)
feature_cache_prefix (str)
feature_cache_per_atom_energy (bool)
device (device | str | None)
seed (int)
progress (bool)
progress_description (str)

Return type:

DataLoader | CachedASEBatchLoader

ufp.training.batch.build_cached_ase_dataloader(model, dataset, *, batch_size=1, shuffle=False, drop_last=False, num_workers=0, pin_memory=False, persistent_workers=False, prefetch_factor=None, dtype=None, device=None, backend=None, cache_batches_on_device=False, feature_cache_storage='cpu', feature_cache_mode='auto', feature_cache_dir=None, feature_cache_prefix='batch', feature_cache_per_atom_energy=True, seed=0, progress=False, progress_description='Caching training batches')[source]¶

Pre-collate fixed-geometry batches and reuse them across epochs.

Cached loaders intentionally materialize batches in-process. Using PyTorch worker processes here would retain shared-memory handles for every cached tensor and can quickly exhaust the process file-descriptor limit.

Parameters:

model (UFPotential) – Potential used to infer cutoff and optional feature-cache settings.
dataset (Dataset[ASEAtomsSample] | Sequence[ASEAtomsSample]) – Supervised ASE samples or raw sample sequence.
batch_size (int) – Number of samples per cached batch.
shuffle (bool) – Whether to shuffle cached batch order on each iteration.
drop_last (bool) – Whether to drop the final incomplete batch before caching.
num_workers (int) – Worker count used only while building the initial cache.
pin_memory (bool) – Whether the initial data loader should pin memory.
persistent_workers (bool) – Whether worker processes should persist while caching.
prefetch_factor (int | None) – Optional prefetch factor for worker-based cache creation.
dtype (dtype | None) – Optional dtype for cached inputs.
device (device | str | None) – Optional device for cached inputs.
backend (str | NeighborListBackend | None) – Optional neighbor-list backend override.
cache_batches_on_device (bool) – Whether to retain cached batches on device.
feature_cache_storage (str) – Three-body feature-cache storage policy.
feature_cache_mode (str) – Disk feature-cache read/write policy.
feature_cache_dir (Path | str | None) – Optional directory for disk-backed feature caches.
feature_cache_prefix (str) – Prefix for feature-cache entries.
feature_cache_per_atom_energy (bool) – Whether feature caches include per-atom rows.
seed (int) – Seed used when shuffling cached batches.
progress (bool) – Whether to show cache-building progress.
progress_description (str) – Label for the progress bar.

Returns:

Reusable iterable over cached ASE batches.

Return type:

CachedASEBatchLoader

Training and evaluation loops for UFP models.

This module runs epoch-level optimization over ASE-backed dataloaders while reusing the same compute path as inference.

ufp.training.engine.evaluate_model(model, loader, *, split='validation', dtype=None, device=None, neighbor_backend=None, loss_weights=None, progress=False, progress_description=None)[source]¶

Run one evaluation pass without optimizer updates.

Parameters:

model (UFPotential)
loader (DataLoader | Iterable[ASEAtomsBatch])
split (str)
dtype (dtype | None)
device (device | str | None)
neighbor_backend (str | NeighborListBackend | None)
loss_weights (LossWeights | None)
progress (bool)
progress_description (str | None)

Return type:

EpochMetrics

ufp.training.engine.fit_model(model, train_loader, *, optimizer, epochs, validation_loader=None, test_loader=None, dtype=None, device=None, neighbor_backend=None, loss_weights=None, twobody_shape_penalty=None, twobody_wall_penalty=None, max_grad_norm=None, validation_frequency=1, evaluate_initial=False, evaluate_final_train=False, compile_model=False, progress=False)[source]¶

Repeat training and evaluation epochs, then package the full history.

Parameters:

model (UFPotential)
train_loader (DataLoader | Iterable[ASEAtomsBatch])
optimizer (Optimizer)
epochs (int)
validation_loader (DataLoader | Iterable[ASEAtomsBatch] | None)
test_loader (DataLoader | Iterable[ASEAtomsBatch] | None)
dtype (dtype | None)
device (device | str | None)
neighbor_backend (str | NeighborListBackend | None)
loss_weights (LossWeights | None)
twobody_shape_penalty (TwoBodySplineShapePenalty | None)
twobody_wall_penalty (TwoBodyCriticalWallPenalty | None)
max_grad_norm (float | None)
validation_frequency (int | None)
evaluate_initial (bool)
evaluate_final_train (bool)
compile_model (bool)
progress (bool)

Return type:

TrainingHistory

ufp.training.engine.test_model(model, loader, *, dtype=None, device=None, neighbor_backend=None, loss_weights=None, progress=False, progress_description='test')[source]¶

Evaluate the model on a test dataloader.

Parameters:

model (UFPotential)
loader (DataLoader | Iterable[ASEAtomsBatch])
dtype (dtype | None)
device (device | str | None)
neighbor_backend (str | NeighborListBackend | None)
loss_weights (LossWeights | None)
progress (bool)
progress_description (str)

Return type:

EpochMetrics

ufp.training.engine.train_one_epoch(model, loader, *, optimizer, dtype=None, device=None, neighbor_backend=None, loss_weights=None, twobody_shape_penalty=None, twobody_wall_penalty=None, max_grad_norm=None, progress=False, progress_description='train')[source]¶

Run one optimization epoch over a training dataloader.

Parameters:

model (UFPotential)
loader (DataLoader | Iterable[ASEAtomsBatch])
optimizer (Optimizer)
dtype (dtype | None)
device (device | str | None)
neighbor_backend (str | NeighborListBackend | None)
loss_weights (LossWeights | None)
twobody_shape_penalty (TwoBodySplineShapePenalty | None)
twobody_wall_penalty (TwoBodyCriticalWallPenalty | None)
max_grad_norm (float | None)
progress (bool)
progress_description (str)

Return type:

EpochMetrics

Force-component selection helpers for supervised training subsets.

class ufp.training.force_selection.ForceComponent(dataset_index, atom_index, component_index, force, abs_force)[source]¶

Bases: object

One selected Cartesian force component in a source dataset.

Parameters:

dataset_index (int)
atom_index (int)
component_index (int)
force (float)
abs_force (float)

class ufp.training.force_selection.ForceComponentSelection(selected_indices, force_masks, ranked_components, max_force, requested_configurations, requested_components)[source]¶

Bases: object

Selected structures, force masks, and ranked force components.

Parameters:

selected_indices (ndarray)
force_masks (tuple[ndarray, ...])
ranked_components (tuple[ForceComponent, ...])
max_force (float)
requested_configurations (int | None)
requested_components (int | None)

as_metadata(*, prefix='force_selection')[source]¶

Return JSON-friendly scalar metadata for checkpoints.

Parameters:: prefix (str)
Return type:: dict[str, object]

property n_selected_components: int¶: Return the number of selected Cartesian force components.

property n_selected_configurations: int¶: Return the number of selected structures.

class ufp.training.force_selection.LargestForceComponentSelector(*, max_force, n_configurations=None, n_components=None, allow_fewer=False)[source]¶

Bases: object

Select high-force Cartesian components for focused force training.

The selector ranks absolute force components from largest to smallest, excludes components above max_force, and returns boolean masks aligned with the selected structures.

Parameters:

max_force (float)
n_configurations (int | None)
n_components (int | None)
allow_fewer (bool)

select(forces, *, indices=None)[source]¶

Select high-force components from one force-array sequence.

Parameters:

forces (Sequence[object])
indices (Sequence[int] | ndarray | None)

Return type:

Coefficient freezing helpers for optimizer-based training.

class ufp.training.freezing.CoefficientFreezeState(masks, values, parameter_metadata=<factory>, selector_metadata=<factory>)[source]¶

Bases: object

Frozen coefficient masks and reference values for optimizer training.

Parameters:

masks (dict[Parameter, Tensor])
values (dict[Parameter, Tensor])
parameter_metadata (tuple[FrozenParameterMetadata, ...])
selector_metadata (tuple[FrozenSelectorMetadata, ...])

property affected_parameter_names: tuple[str, ...]¶: Return parameter names touched by this freeze state.

apply_gradient_masks()[source]¶

Zero gradients for frozen coefficient entries.

Return type:: None

clear_optimizer_state(optimizer)[source]¶

Clear tensor optimizer-state entries for frozen coefficient positions.

Adam-style optimizers keep per-parameter moment buffers. Clearing the frozen entries prevents stale moments from reappearing when the same optimizer is later unfrozen or rewrapped with a different mask.

Parameters:: optimizer (Optimizer)
Return type:: int

property frozen_count: int¶: Return the total number of frozen coefficient entries.

property frozen_counts: dict[str, int]¶: Return frozen-entry counts keyed by parameter name.

restore()[source]¶

Restore frozen coefficient entries after an optimizer step.

Return type:: None

property trainable_count: int¶: Return the total number of unfrozen entries in affected parameters.

property trainable_counts: dict[str, int]¶: Return unfrozen-entry counts keyed by parameter name.

unwrap_optimizer(optimizer)[source]¶

Restore an optimizer previously patched by this freeze state.

Parameters:: optimizer (Optimizer)
Return type:: Optimizer

wrap_optimizer(optimizer)[source]¶

Patch one optimizer so frozen entries remain unchanged after step.

Parameters:: optimizer (Optimizer)
Return type:: Optimizer

class ufp.training.freezing.FrozenParameterMetadata(name, shape, frozen_count, trainable_count, block_labels)[source]¶

Bases: object

Inspection summary for one optimizer parameter touched by a freeze mask.

Parameters:

name (str)
shape (tuple[int, ...])
frozen_count (int)
trainable_count (int)
block_labels (tuple[str, ...])

class ufp.training.freezing.FrozenSelectorMetadata(selector, block_labels, frozen_count)[source]¶

Bases: object

Inspection summary for one user selector after layout resolution.

Parameters:

selector (str)
block_labels (tuple[str, ...])
frozen_count (int)

ufp.training.freezing.freeze_model_coefficients(model, selectors)[source]¶

Return masks that freeze selected coefficients during optimizer training.

Parameters:

model (UFPModel)
selectors (Sequence[int | str | CoefficientSelector])

Return type:

CoefficientFreezeState

Workflows¶

Status: Experimental workflow API.

Reusable workflow helpers for examples and small supervised UFP studies.

class ufp.workflows.data.DatasetSplits(training_indices, validation_indices, holdout_indices)[source]¶

Bases: object

Train, validation, and holdout indices for supervised examples.

Parameters:

training_indices (ndarray)
validation_indices (ndarray)
holdout_indices (ndarray)

property testing_indices: ndarray¶: Return the holdout split under the older testing alias.

class ufp.workflows.data.SupervisedAtomsDataset(frames, energies, forces, sizes, training_indices, validation_indices, holdout_indices)[source]¶

Bases: object

ASE structures with total-energy/force labels and canonical splits.

Parameters:

frames (tuple[Atoms, ...])
energies (ndarray)
forces (tuple[ndarray, ...])
sizes (ndarray)
training_indices (ndarray)
validation_indices (ndarray)
holdout_indices (ndarray)

property testing_indices: ndarray¶: Return the holdout split under the older testing alias.

ufp.workflows.data.load_or_create_split_indices(split_path, *, dataset, n_frames, seed, n_train, n_validation, sort=True)[source]¶

Load split indices, creating them deterministically when absent.

Parameters:

split_path (Path)
dataset (str)
n_frames (int)
seed (int)
n_train (int)
n_validation (int)
sort (bool)

Return type:

DatasetSplits

ufp.workflows.data.make_split_indices(n_frames, *, seed, n_train, n_validation, sort=True)[source]¶

Create deterministic train, validation, and holdout split indices.

Parameters:

n_frames (int)
seed (int)
n_train (int)
n_validation (int)
sort (bool)

Return type:

DatasetSplits

ufp.workflows.data.make_supervised_atoms_dataset(frames, energies, forces, *, splits)[source]¶

Build a labeled ASE dataset object from arrays and split indices.

Parameters:

frames (Sequence[Atoms])
energies (Sequence[float])
forces (Sequence[object])
splits (DatasetSplits)

Return type:

SupervisedAtomsDataset

ufp.workflows.data.minimum_pair_distances_from_dataset(dataset, *, atomic_types=None, active_pairs=None, symmetric=True, indices=None)[source]¶

Return shortest observed pair distance for each requested pair channel.

For SupervisedAtomsDataset inputs, the training split is used by default. For ASEAtomsDataset or a sequence-like dataset, all entries are used unless indices is provided.

Parameters:

atomic_types (Sequence[int] | None)
active_pairs (Sequence[tuple[int, int]] | None)
symmetric (bool)
indices (Sequence[int] | ndarray | None)

Return type:

dict[tuple[int, int], float]

ufp.workflows.data.read_split_indices(split_path)[source]¶

Read train, validation, and holdout indices from JSON.

Parameters:: split_path (Path)
Return type:: DatasetSplits

ufp.workflows.data.validate_split_indices(*, n_frames, training_indices, validation_indices, holdout_indices, n_train=None, n_validation=None, allow_empty_validation=True, allow_empty_holdout=True)[source]¶

Validate split indices and return normalized integer arrays.

Parameters:

n_frames (int)
training_indices (Sequence[int] | ndarray)
validation_indices (Sequence[int] | ndarray)
holdout_indices (Sequence[int] | ndarray)
n_train (int | None)
n_validation (int | None)
allow_empty_validation (bool)
allow_empty_holdout (bool)

Return type:

DatasetSplits

ufp.workflows.data.write_split_indices(split_path, *, dataset, n_frames, seed, splits)[source]¶

Write split indices to JSON so all example scripts share them.

Parameters:

split_path (Path)
dataset (str)
n_frames (int)
seed (int)
splits (DatasetSplits)

Return type:

None

Reusable workflow helpers for examples and small supervised UFP studies.

class ufp.workflows.models.TermSchemaCodec(term_type, class_name, encode, decode)[source]¶

Bases: object

Encoder/decoder pair for one checkpoint-reconstructable term class.

Parameters:

term_type (type)
class_name (str)
encode (Callable[[object, int | None], dict[str, object]])
decode (Callable[[Mapping[str, object], TermSchemaContext], object])

class ufp.workflows.models.TermSchemaContext(providers, dtype, atomic_types, provider, coefficient_index)[source]¶

Bases: object

Shared reconstruction context passed to registered term schema decoders.

Parameters:

providers (Sequence[AlchemicalCoefficients])
dtype (dtype)
atomic_types (tuple[int, ...] | None)
provider (AlchemicalCoefficients | None)
coefficient_index (int | None)

ufp.workflows.models.add_element_reference_term(interaction_model, atomic_type, onebody_energy, *, dtype=None, trainable=False)[source]¶

Return a model with a one-element one-body term prepended.

Parameters:

interaction_model (UFPModel)
atomic_type (int)
onebody_energy (float)
dtype (dtype | None)
trainable (bool)

Return type:

ufp.workflows.models.demonstrate_calculator(model, frame)[source]¶

Evaluate one structure with the ASE calculator wrapper.

Parameters:

model (UFPModel)
frame (Atoms)

Return type:

dict[str, object]

ufp.workflows.models.extract_pair_component(interaction_model, *, atomic_types=None)[source]¶

Copy only the pair terms from a fitted interaction model.

Parameters:

interaction_model (UFPModel)
atomic_types (Sequence[int] | None)

Return type:

ufp.workflows.models.load_model_from_checkpoint(checkpoint, *, dtype=torch.float64, strict=True)[source]¶

Build a model from a checkpoint model_schema and load its state.

Parameters:

checkpoint (Mapping[str, object])
dtype (dtype)
strict (bool)

Return type:

ufp.workflows.models.model_from_schema(schema, *, dtype=torch.float64)[source]¶

Reconstruct a supported UFP example model from a checkpoint schema.

Parameters:

schema (Mapping[str, object])
dtype (dtype)

Return type:

RegularizationCandidate

ufp.workflows.models.model_schema(model)[source]¶

Return a reconstructable model schema for supported UFP example terms.

Parameters:: model (UFPModel)
Return type:: dict[str, object]

ufp.workflows.models.register_term_schema_codec(term_type, encode, decode, *, class_name=None)[source]¶

Parameters:

term_type (type)
encode (Callable[[object, int | None], dict[str, object]])
decode (Callable[[Mapping[str, object], TermSchemaContext], object])
class_name (str | None)

Return type:

None

ufp.workflows.models.relax_structure(model, frame, *, fmax=0.05, steps=50)[source]¶

Run a short ASE relaxation using UFP forces.

Parameters:

model (UFPModel)
frame (Atoms)
fmax (float)
steps (int)

Return type:

dict[str, object]

ufp.workflows.models.save_checkpoint(filename, *, interaction_model, metadata, onebody_energy=None)[source]¶

Save a simple torch checkpoint for a fitted model.

Parameters:

filename (Path)
interaction_model (UFPModel)
metadata (dict[str, object])
onebody_energy (float | None)

Return type:

None

Reusable workflow helpers for examples and small supervised UFP studies.

ufp.workflows.onebody.fit_element_onebody_energies(atoms_list, energies=None, *, atomic_types=None, energy_key=None, rcond=None)[source]¶

Fit composition-only per-element reference energies by least squares.

Parameters:

atoms_list (Sequence[Atoms]) – ASE structures whose elemental compositions define the linear system.
energies (Sequence[float] | None) – Optional total energies, one per structure. If omitted, energies are read with atoms.get_potential_energy() unless energy_key is set.
atomic_types (Sequence[int] | None) – Atomic numbers to fit. When omitted, all elements observed in atoms_list are fitted.
energy_key (str | None) – Optional key in atoms.info or atoms.arrays used when energies is omitted.
rcond (float | None) – Cutoff passed to numpy.linalg.lstsq.

Returns:

Mapping from atomic number to fitted reference energy in eV.

Raises:

KeyError – If energy_key is set but missing from a structure.
TypeError – If atoms_list contains non-ASE objects.
ValueError – If the inputs are empty or have inconsistent lengths.

Return type:

dict[int, float]

ufp.workflows.onebody.initialize_onebody_terms_from_atoms(model, atoms_list, energies=None, *, energy_key=None, rcond=None)[source]¶

Initialize trainable element one-body terms from composition least squares.

Parameters:

model (UFPModel)
atoms_list (Sequence[Atoms])
energies (Sequence[float] | None)
energy_key (str | None)
rcond (float | None)

Return type:

dict[int, float]

ufp.workflows.onebody.initialize_onebody_terms_from_dataset(model, dataset, *, rcond=None)[source]¶

Initialize trainable element one-body terms from an ASE training dataset.

Parameters:

model (UFPModel)
dataset (ASEAtomsDataset)
rcond (float | None)

Return type:

dict[int, float]

ufp.workflows.onebody.onebody_reference_energies(model)[source]¶

Return all element one-body values keyed by atomic number.

Parameters:: model (UFPModel)
Return type:: dict[int, float]

ufp.workflows.onebody.onebody_reference_energy(model)[source]¶

Return the scalar value from a single ElementOneBodyTerm model.

Parameters:: model (UFPModel)
Return type:: float

Reusable workflow helpers for examples and small supervised UFP studies.

ufp.workflows.predictions.ase_training_dataset_from_frames(frames, energies, forces, *, indices=None, force_masks=None)[source]¶

Build an ASEAtomsDataset from frame-level labels.

Parameters:

frames (Sequence[Atoms])
energies (Sequence[float])
forces (Sequence[object])
indices (Sequence[int] | None)
force_masks (Sequence[object] | None)

Return type:

ASEAtomsDataset

ufp.workflows.predictions.fit_samples_from_dataset(dataset, *, indices=None, sample_weights=None, energy_weight=1.0, force_weight=1.0, per_atom_weight=1.0)[source]¶

Convert selected labeled frames into least-squares samples.

Parameters:

dataset (SupervisedAtomsDataset)
indices (Sequence[int] | ndarray | None)
sample_weights (dict[int, float] | None)
energy_weight (float)
force_weight (float)
per_atom_weight (float)

Return type:

list[FitSample]

ufp.workflows.predictions.prediction_metrics_for_split(model, dataset, indices, *, batch_size, device=None, dtype=None, progress=False)[source]¶

Predict a split and return standard energy/force metrics.

Parameters:

model (UFPModel)
dataset (SupervisedAtomsDataset)
indices (Sequence[int] | ndarray)
batch_size (int)
device (device | str | None)
dtype (dtype | None)
progress (bool)

Return type:

dict[str, ndarray | float]

ufp.workflows.predictions.print_dataset_summary(dataset)[source]¶

Print split sizes and atom-count range for an example dataset.

Parameters:: dataset (SupervisedAtomsDataset)
Return type:: None

ufp.workflows.predictions.print_metrics(prefix, metrics)[source]¶

Print compact energy and force RMSE metrics.

Parameters:

prefix (str)
metrics (dict[str, ndarray | float])

Return type:

None

ufp.workflows.predictions.save_prediction_split(model, dataset, indices, *, filename, batch_size, device=None, dtype=None, progress=False, uncertainty_posterior=None, uncertainty_aleatoric_variance=None, uncertainty_aleatoric_noise_bundle=None, uncertainty_energy_variance_scale=1.0, uncertainty_forces=False)[source]¶

Predict one split, write a standard .npz, and return metrics.

Parameters:

model (UFPModel)
dataset (SupervisedAtomsDataset)
indices (Sequence[int] | ndarray)
filename (Path)
batch_size (int)
device (device | str | None)
dtype (dtype | None)
progress (bool)
uncertainty_posterior (object | None)
uncertainty_aleatoric_variance (float | None)
uncertainty_aleatoric_noise_bundle (object | None)
uncertainty_energy_variance_scale (float)
uncertainty_forces (bool)

Return type:

dict[str, ndarray | float]

ufp.workflows.predictions.save_prediction_splits(model, dataset, *, output_directory, prefix, batch_size, device=None, dtype=None, progress=False, uncertainty_posterior=None, uncertainty_aleatoric_variance=None, uncertainty_aleatoric_noise_bundle=None, uncertainty_energy_variance_scale=1.0, uncertainty_forces=False)[source]¶

Save train, validation, and holdout prediction arrays.

Parameters:

model (UFPModel)
dataset (SupervisedAtomsDataset)
output_directory (Path)
prefix (str)
batch_size (int)
device (device | str | None)
dtype (dtype | None)
progress (bool)
uncertainty_posterior (object | None)
uncertainty_aleatoric_variance (float | None)
uncertainty_aleatoric_noise_bundle (object | None)
uncertainty_energy_variance_scale (float)
uncertainty_forces (bool)

Return type:

dict[str, dict[str, ndarray | float]]

ufp.workflows.predictions.split_frames(dataset, indices)[source]¶

Return frame copies for one selected split.

Parameters:

dataset (SupervisedAtomsDataset)
indices (Sequence[int] | ndarray)

Return type:

list[Atoms]

Compatibility aliases for public cache naming helpers.

New code should import cache identity helpers from ufp.cache or from the top-level ufp convenience exports. This module remains importable for older workflow scripts.

class ufp.workflows.cache.CacheIdentity(cache_kind, prefix, settings, digest_length=16)[source]¶

Bases: object

Stable identity for a settings-addressed cache family.

Parameters:

cache_kind (str)
prefix (str)
settings (Mapping[str, object])
digest_length (int)

property digest: str¶: Return the stable digest for this cache identity.

directory(parent)[source]¶

Return this identity’s cache directory under parent.

Parameters:: parent (Path | str)
Return type:: Path

property name: str¶: Return the deterministic directory name for this cache identity.

property normalized_settings: dict[str, object]¶: Return JSON-friendly settings used to compute this cache identity.

summary(*, extra=None)[source]¶

Return a JSON-friendly summary suitable for cache_settings.json.

Parameters:: extra (Mapping[str, object] | None)
Return type:: dict[str, object]

ufp.workflows.cache.cache_settings_digest(settings, *, length=16)[source]¶

Return a stable digest for JSON-friendly cache identity settings.

Parameters:

settings (object)
length (int)

Return type:

str

ufp.workflows.cache.cache_settings_from_names(namespace, names, *, example, cache_kind, extra=None)[source]¶

Build a JSON-friendly cache identity payload from named parameters.

Parameters:

namespace (Mapping[str, object])
names (Sequence[str])
example (str)
cache_kind (str)
extra (Mapping[str, object] | None)

Return type:

dict[str, object]

ufp.workflows.cache.normalize_cache_settings(value)[source]¶

Return a JSON-friendly representation for cache identity settings.

Parameters:: value (object)
Return type:: object

ufp.workflows.cache.suggest_cache_digest(settings, *, length=16)[source]¶

Return a stable digest for JSON-friendly cache identity settings.

Parameters:

settings (object)
length (int)

Return type:

str

ufp.workflows.cache.suggest_cache_directory(parent, prefix, settings, *, length=16)[source]¶

Return a deterministic cache directory path from settings.

Parameters:

parent (Path | str)
prefix (str)
settings (object)
length (int)

Return type:

Path

ufp.workflows.cache.suggest_cache_name(prefix, settings, *, length=16)[source]¶

Return a deterministic cache directory name from settings.

Parameters:

prefix (str)
settings (object)
length (int)

Return type:

str

Workflow helpers for physical pair priors and two-body warm starts.

ufp.workflows.prior.copy_twobody_coefficients(source, target, *, strict=True)[source]¶

Copy matching two-body spline coefficient rows from source to target.

source and target may be models, individual two-body terms, or sequences of such terms. Only active target rows are copied. With strict=True, every active target row must exist in the source and must have matching spline metadata.

Returns:: Number of coefficient rows copied.
Parameters:: strict (bool)
Return type:: int

ufp.workflows.prior.project_pair_prior_to_twobody(prior_term, *, coeff_size, spline='cubic', full_support_start, n_samples=None, derivative_weight=1.0, trainable=True, fittable=True)[source]¶

Project one pair-dependent analytic prior onto a spline two-body term.

Parameters:

prior_term (PowerLawRepulsionTerm) – Fitted analytic prior term to project.
coeff_size (int) – Number of 1D spline coefficients per pair category.
spline (str) – Spline family for the projected term.
full_support_start (float) – Lower distance where the projected spline has full support.
n_samples (int | None) – Optional number of midpoint samples per pair category.
derivative_weight (float) – Weight for radial-derivative projection rows.
trainable (bool) – Whether projected spline coefficients require gradients.
fittable (bool) – Whether the projected term exposes coefficients to least-squares fitters.

Returns:

The projected categorized two-body term and per-active-pair diagnostics.

Return type:

tuple[SplineTwoBodyTerm, dict[tuple[int, int], UniformSpline1DFitResult]]

Progressive regularization tuning for linear least-squares workflows.

class ufp.workflows.regularization.RegularizationCandidate(ridge=0.0, onebody_ridge=0.0, twobody_ridge=0.0, threebody_ridge=0.0)[source]¶

Bases: object

Concrete ridge settings passed to LinearFitter.

Parameters:

ridge (float)
onebody_ridge (float)
twobody_ridge (float)
threebody_ridge (float)

as_fitter_kwargs()[source]¶

Return ridge keyword arguments accepted by LinearFitter.

Return type:: dict[str, float]

group_value(group)[source]¶

Return the ridge value for one canonical group.

Parameters:: group (str)
Return type:: float

to_dict()[source]¶

Return a JSON-friendly representation.

Return type:: dict[str, float]

total_ridge()[source]¶

Return a scalar used for deterministic tie-breaking.

Return type:: float

with_group_value(group, value)[source]¶

Return this candidate with one group changed.

Parameters:

group (str)
value (float)

Return type:

class ufp.workflows.regularization.RegularizationSearchConfig(seed=0, alpha=1e-06, estimate_subset_size=64, stage_subset_sizes=(64, 256), candidate_multipliers=(0.001, 0.01, 0.1, 1.0, 10.0, 100.0, 1000.0), refinement_multipliers=(0.1, 1.0, 10.0), top_k_per_stage=5, validation_fraction=0.2, minimum_validation_size=1, energy_score_weight=1.0, force_score_weight=1.0, refit_full=True, cache_directory=None, cache_mode='auto', dense_cache_parameter_limit=20000, prediction_batch_size=64, progress=False)[source]¶

Bases: object

Options for progressive regularization tuning.

Parameters:

seed (int)
alpha (float)
estimate_subset_size (int | None)
stage_subset_sizes (tuple[int, ...])
candidate_multipliers (tuple[float, ...])
refinement_multipliers (tuple[float, ...])
top_k_per_stage (int)
validation_fraction (float)
minimum_validation_size (int)
energy_score_weight (float)
force_score_weight (float)
refit_full (bool)
cache_directory (Path | str | None)
cache_mode (str)
dense_cache_parameter_limit (int)
prediction_batch_size (int)
progress (bool)

to_dict()[source]¶

Return JSON-friendly config metadata.

Return type:: dict[str, object]

class ufp.workflows.regularization.RegularizationSearchResult(estimates, trials, best_candidate, best_fitter_kwargs, final_model=None, final_fit_result=None)[source]¶

Bases: object

Result of a progressive regularization search.

Parameters:

estimates (RidgeScaleEstimate)
trials (tuple[RegularizationTrial, ...])
best_candidate (RegularizationCandidate)
best_fitter_kwargs (Mapping[str, object])
final_model (UFPModel | None)
final_fit_result (LinearFitResult | None)

property metadata: dict[str, object]¶: Return JSON-friendly search metadata.

class ufp.workflows.regularization.RegularizationTrial(stage, candidate, training_size, validation_size, metrics=<factory>, score=None, status='ok', error=None)[source]¶

Bases: object

One candidate fit and validation attempt.

Parameters:

stage (str)
candidate (RegularizationCandidate)
training_size (int)
validation_size (int)
metrics (Mapping[str, float])
score (float | None)
status (str)
error (str | None)

to_dict()[source]¶

Return JSON-friendly trial metadata.

Return type:: dict[str, object]

class ufp.workflows.regularization.RidgeGroupEstimate(group, n_parameters, design_trace, trace_per_parameter, suggested_ridge)[source]¶

Bases: object

Data-scale ridge estimate for one coefficient regularization group.

Parameters:

group (str)
n_parameters (int)
design_trace (float)
trace_per_parameter (float)
suggested_ridge (float)

to_dict()[source]¶

Return a JSON-friendly representation.

Return type:: dict[str, object]

class ufp.workflows.regularization.RidgeScaleEstimate(groups, alpha, sample_count, diagnostics=<factory>)[source]¶

Bases: object

Block-scale ridge suggestions derived from least-squares design statistics.

Parameters:

groups (tuple[RidgeGroupEstimate, ...])
alpha (float)
sample_count (int)
diagnostics (Mapping[str, float])

property by_group: dict[str, RidgeGroupEstimate]¶: Return group estimates keyed by group name.

candidate()[source]¶

Return the direct ridge candidate implied by this estimate.

Return type:: RegularizationCandidate

to_dict()[source]¶

Return JSON-friendly estimate metadata.

Return type:: dict[str, object]

ufp.workflows.regularization.estimate_linear_ridge_scales(model, samples, *, fitter_kwargs=None, fit_kwargs=None, subset_size=None, seed=0, alpha=1e-06)[source]¶

Estimate block-scale ridge weights from weighted least-squares design traces.

The suggested group weight is alpha * trace(A_g.T @ A_g) / n_params_g.

Parameters:

model (UFPModel)
samples (Sequence[FitSample])
fitter_kwargs (Mapping[str, object] | None)
fit_kwargs (Mapping[str, object] | None)
subset_size (int | None)
seed (int)
alpha (float)

Return type:

RidgeScaleEstimate

ufp.workflows.regularization.tune_linear_regularization(model_factory, dataset, *, training_indices=None, validation_indices=None, sample_weights=None, energy_weight=1.0, force_weight=1.0, config=None, fitter_kwargs=None, fit_kwargs=None)[source]¶

Tune linear ridge weights with deterministic progressive subsets.

Parameters:

model_factory (Callable[[], UFPModel])
dataset (SupervisedAtomsDataset)
training_indices (Sequence[int] | ndarray | None)
validation_indices (Sequence[int] | ndarray | None)
sample_weights (dict[int, float] | None)
energy_weight (float)
force_weight (float)
config (RegularizationSearchConfig | None)
fitter_kwargs (Mapping[str, object] | None)
fit_kwargs (Mapping[str, object] | None)

Return type:

RegularizationSearchResult

Workflow helpers for choosing efficient three-body runtime defaults.

class ufp.workflows.threebody.ThreeBodyLeastSquaresRuntimeStatus(requested_backend, selected_backend, requested_bucket_backend, fit_device, prediction_device, assembled_cache_mode, normal_equation_cache_mode, normal_equation_build_device, native_cxx_assembly_available, native_cuda_assembly_available, native_cxx_bucketing_available, native_cxx_assembly_used, native_cuda_assembly_used, native_cxx_bucketing_used)[source]¶

Bases: object

Resolved three-body runtime configuration for least-squares workflows.

Parameters:

requested_backend (str)
selected_backend (str)
requested_bucket_backend (str)
fit_device (device)
prediction_device (device)
assembled_cache_mode (str)
normal_equation_cache_mode (str)
normal_equation_build_device (device | None)
native_cxx_assembly_available (bool)
native_cuda_assembly_available (bool)
native_cxx_bucketing_available (bool)
native_cxx_assembly_used (bool)
native_cuda_assembly_used (bool)
native_cxx_bucketing_used (bool)

as_metadata()[source]¶

Return JSON/checkpoint-friendly runtime metadata.

Return type:: dict[str, object]

class ufp.workflows.threebody.ThreeBodyRuntimeStatus(requested_backend, requested_bucket_backend, selected_device, native_cxx_evaluator_available, native_cuda_evaluator_available, native_cxx_bucketing_available, native_cxx_dense_cache_available, native_cxx_bucketing_used, native_cxx_dense_cache_used, native_cxx_dynamic_used, native_cuda_dynamic_used, dynamic_note)[source]¶

Bases: object

Resolved three-body runtime configuration for a training workflow.

Parameters:

requested_backend (str)
requested_bucket_backend (str)
selected_device (device)
native_cxx_evaluator_available (bool)
native_cuda_evaluator_available (bool)
native_cxx_bucketing_available (bool)
native_cxx_dense_cache_available (bool)
native_cxx_bucketing_used (bool)
native_cxx_dense_cache_used (bool)
native_cxx_dynamic_used (bool)
native_cuda_dynamic_used (bool)
dynamic_note (str)

as_metadata()[source]¶

Return JSON/checkpoint-friendly runtime metadata.

Return type:: dict[str, object]

ufp.workflows.threebody.configure_threebody_leastsquares_runtime(*, backend='auto', bucket_backend='python', fit_device_selection='auto', fit_device=None, prediction_device_selection='auto', prediction_device=None, assembled_cache_mode='auto', normal_equation_cache_mode='auto', normal_equation_build_device=None)[source]¶

Apply three-body least-squares runtime options and report dispatch status.

Least-squares auto follows the current assembly benchmarks: native CUDA, native CPU C++, CUDA Torch, then CPU Torch. Native source bucketing is not part of the default least-squares path because the benchmarked workloads do not justify its overhead.

Parameters:

backend (str) – Requested least-squares assembly backend.
bucket_backend (str) – Requested source-bucketing backend.
fit_device_selection (str | device) – Automatic or explicit fit-device selection.
fit_device (str | device | None) – Explicit fit device override.
prediction_device_selection (str | device) – Automatic or explicit prediction-device selection.
prediction_device (str | device | None) – Explicit prediction device override.
assembled_cache_mode (str) – Assembled-batch cache mode label.
normal_equation_cache_mode (str) – Normal-equation cache mode label.
normal_equation_build_device (str | device | None) – Optional device used to build normal equations.

Returns:

Fit device, prediction device, and selected runtime status.

Return type:

tuple[device, device, ThreeBodyLeastSquaresRuntimeStatus]

ufp.workflows.threebody.configure_threebody_runtime_defaults()[source]¶

Set conservative three-body runtime defaults unless already configured.

auto prefers native dynamic kernels when available. Bucket auto uses the Python/Torch path for CUDA inputs to avoid CPU-GPU transfer overhead, and may use native CPU source preprocessing for CPU inputs. Explicit UFP_THREEBODY_BACKEND and UFP_THREEBODY_BUCKET_BACKEND environment settings are preserved.

Return type:: None

ufp.workflows.threebody.configure_threebody_training_runtime(*, backend='auto', bucket_backend='auto', device_selection='auto', device=None, cache_batches=False, cache_batches_on_device=False, feature_cache_storage='none')[source]¶

Apply three-body runtime options and return the selected training status.

Native dynamic three-body kernels are inference-only today, so gradient training always uses the PyTorch evaluator. The native CPU bucket path can still be used when its inputs remain on CPU.

Parameters:

backend (str) – Requested dynamic three-body evaluator backend.
bucket_backend (str) – Requested source-bucketing backend.
device_selection (str | device) – Automatic or explicit training-device selection.
device (str | device | None) – Explicit training device override.
cache_batches (bool) – Whether training will reuse cached batches.
cache_batches_on_device (bool) – Whether cached batches stay on the training device.
feature_cache_storage (str) – Feature-cache storage policy.

Returns:

Training device and selected runtime status.

Return type:

tuple[device, ThreeBodyRuntimeStatus]

ufp.workflows.threebody.preferred_threebody_inference_device()[source]¶

Return the preferred dynamic three-body inference device.

The order follows current dynamic-inference benchmarks: native CUDA, native CPU C++, GPU PyTorch, then CPU PyTorch.

Returns:: Preferred torch device for dynamic three-body inference.
Return type:: device

ufp.workflows.threebody.print_threebody_leastsquares_runtime_status(status)[source]¶

Print native three-body least-squares availability and dispatch.

Parameters:: status (ThreeBodyLeastSquaresRuntimeStatus)
Return type:: None

ufp.workflows.threebody.print_threebody_runtime_status(status)[source]¶

Print native three-body availability and expected training dispatch.

Parameters:: status (ThreeBodyRuntimeStatus)
Return type:: None

Reusable workflow helpers for examples and small supervised UFP studies.

class ufp.workflows.training.TorchTrainingResult(history, initial_train_metrics, initial_validation_metrics, final_train_metrics, final_validation_metrics, loss_weights, device, dtype, batch_size, validation_batch_size, epochs, completed_epochs, learning_rate, weight_decay, onebody_weight_decay, pair_weight_decay, threebody_weight_decay, embedding_weight_decay, state_weight_decay, charge_spin_weight_decay, term_weight_decays, max_grad_norm, validation_frequency, early_stopping_patience, early_stopping_min_delta, restore_best, stopped_early, interrupted, restored_checkpoint_path, best_validation_epoch, best_validation_metrics, best_checkpoint_path, num_workers, pin_memory, cache_batches, cache_batches_on_device, feature_cache_storage, feature_cache_mode, validation_feature_cache_storage, validation_feature_cache_mode, feature_cache_dir, validation_feature_cache_dir, feature_cache_per_atom_energy, compile_model, evaluate_initial, evaluate_final_train, initial_onebody_energies, n_parameters)[source]¶

Bases: object

Summary of one torch-native supervised training run.

Parameters:

history (TrainingHistory)
initial_train_metrics (EpochMetrics | None)
initial_validation_metrics (EpochMetrics | None)
final_train_metrics (EpochMetrics)
final_validation_metrics (EpochMetrics | None)
loss_weights (LossWeights)
device (str)
dtype (str)
batch_size (int)
validation_batch_size (int | None)
epochs (int)
completed_epochs (int)
learning_rate (float)
weight_decay (float)
onebody_weight_decay (float)
pair_weight_decay (float)
threebody_weight_decay (float)
embedding_weight_decay (float)
state_weight_decay (float)
charge_spin_weight_decay (float)
term_weight_decays (dict[str, float])
max_grad_norm (float | None)
validation_frequency (int | None)
early_stopping_patience (int | None)
early_stopping_min_delta (float)
restore_best (bool)
stopped_early (bool)
interrupted (bool)
restored_checkpoint_path (str | None)
best_validation_epoch (int | None)
best_validation_metrics (EpochMetrics | None)
best_checkpoint_path (str | None)
num_workers (int)
pin_memory (bool)
cache_batches (bool)
cache_batches_on_device (bool)
feature_cache_storage (str)
feature_cache_mode (str)
validation_feature_cache_storage (str)
validation_feature_cache_mode (str)
feature_cache_dir (str | None)
validation_feature_cache_dir (str | None)
feature_cache_per_atom_energy (bool)
compile_model (bool)
evaluate_initial (bool)
evaluate_final_train (bool)
initial_onebody_energies (dict[int, float])
n_parameters (int)

ufp.workflows.training.train_interaction_model(model, train_dataset, *, batch_size, epochs, learning_rate, force_loss_weight=5.0, weight_decay=0.0, onebody_weight_decay=None, pair_weight_decay=None, threebody_weight_decay=None, embedding_weight_decay=None, state_weight_decay=None, charge_spin_weight_decay=None, term_weight_decays=None, twobody_shape_penalty=None, twobody_wall_penalty=None, max_grad_norm=None, validation_dataset=None, validation_batch_size=None, validation_frequency=10, num_workers=0, pin_memory=None, persistent_workers=True, prefetch_factor=2, cache_batches=True, cache_batches_on_device=False, feature_cache_storage='cpu', feature_cache_mode='auto', feature_cache_dir=None, validation_feature_cache_storage='none', validation_feature_cache_mode=None, validation_feature_cache_dir=None, feature_cache_per_atom_energy=False, compile_model='auto', evaluate_initial=False, evaluate_final_train=False, initialize_onebody=True, onebody_rcond=None, early_stopping_patience=None, early_stopping_min_delta=0.0, restore_best=True, best_checkpoint_path=None, best_checkpoint_metadata=None, device=None, dtype=torch.float32, seed=7, progress=False, progress_frequency=10)[source]¶

Optimize an interaction model with the torch-native training stack.

Pass progress="summary" to print validation-epoch summaries without per-batch progress bars.

Parameters:

model (UFPModel)
train_dataset (ASEAtomsDataset)
batch_size (int)
epochs (int)
learning_rate (float)
force_loss_weight (float)
weight_decay (float)
onebody_weight_decay (float | None)
pair_weight_decay (float | None)
threebody_weight_decay (float | None)
embedding_weight_decay (float | None)
state_weight_decay (float | None)
charge_spin_weight_decay (float | None)
term_weight_decays (Mapping[str, float] | None)
twobody_shape_penalty (TwoBodySplineShapePenalty | None)
twobody_wall_penalty (TwoBodyCriticalWallPenalty | None)
max_grad_norm (float | None)
validation_dataset (ASEAtomsDataset | None)
validation_batch_size (int | None)
validation_frequency (int | None)
num_workers (int)
pin_memory (bool | None)
persistent_workers (bool)
prefetch_factor (int | None)
cache_batches (bool)
cache_batches_on_device (bool)
feature_cache_storage (str)
feature_cache_mode (str)
feature_cache_dir (Path | str | None)
validation_feature_cache_storage (str)
validation_feature_cache_mode (str | None)
validation_feature_cache_dir (Path | str | None)
feature_cache_per_atom_energy (bool)
compile_model (bool | Literal['auto'])
evaluate_initial (bool)
evaluate_final_train (bool)
initialize_onebody (bool)
onebody_rcond (float | None)
early_stopping_patience (int | None)
early_stopping_min_delta (float)
restore_best (bool)
best_checkpoint_path (Path | str | None)
best_checkpoint_metadata (dict[str, object] | None)
device (device | str | None)
dtype (dtype)
seed (int)
progress (bool | Literal['summary'])
progress_frequency (int)

Return type:

TorchTrainingResult

Explicit workflow stage objects for composing existing UFP helpers.

class ufp.workflows.stages.LinearFitStage(name='linear_fit', model_key='model', samples_key='fit_samples', result_key='linear_fit_result', fitter_kwargs=<factory>, fit_kwargs=<factory>)[source]¶

Bases: object

Selector-aware wrapper around ufp.leastsquares.LinearFitter.

Parameters:

name (str)
model_key (str)
samples_key (str)
result_key (str)
fitter_kwargs (Mapping[str, object])
fit_kwargs (Mapping[str, object])

property metadata: Mapping[str, object]¶: Return JSON-friendly stage metadata.

property produced_outputs: tuple[str, ...]¶: Context keys produced by this stage.

property required_inputs: tuple[str, ...]¶: Context keys required by this stage.

run(context)[source]¶

Build a LinearFitter and execute its fit method.

Parameters:: context (Mapping[str, object])
Return type:: StageResult

class ufp.workflows.stages.ProjectStage(projector, name='project', input_bindings=<factory>, projector_kwargs=<factory>, result_key='projection_result')[source]¶

Bases: object

Projection stage that delegates to an explicit projection callable.

Parameters:

projector (Callable[[...], object])
name (str)
input_bindings (Mapping[str, str])
projector_kwargs (Mapping[str, object])
result_key (str)

property metadata: Mapping[str, object]¶: Return JSON-friendly stage metadata.

property produced_outputs: tuple[str, ...]¶: Context keys produced by this stage.

property required_inputs: tuple[str, ...]¶: Context keys required by this stage.

run(context)[source]¶

Call the configured projection helper.

Parameters:: context (Mapping[str, object])
Return type:: StageResult

class ufp.workflows.stages.ResidualizeStage(name='residualize', model_key='model', dataset_key='dataset', result_key='residual_result', dataset_output_key='residual_dataset', residual_kwargs=<factory>)[source]¶

Bases: object

Wrapper around nonlinear frozen-component residual materialization.

Parameters:

name (str)
model_key (str)
dataset_key (str)
result_key (str)
dataset_output_key (str)
residual_kwargs (Mapping[str, object])

property metadata: Mapping[str, object]¶: Return JSON-friendly stage metadata.

property produced_outputs: tuple[str, ...]¶: Context keys produced by this stage.

property required_inputs: tuple[str, ...]¶: Context keys required by this stage.

run(context)[source]¶

Materialize residual labels for a dataset.

Parameters:: context (Mapping[str, object])
Return type:: StageResult

class ufp.workflows.stages.StageResult(outputs, metadata)[source]¶

Bases: object

Outputs and metadata produced by one explicit workflow stage.

Parameters:

outputs (Mapping[str, object])
metadata (Mapping[str, object])

update_context(context)[source]¶

Update a user-owned context mapping with this stage’s outputs.

Parameters:: context (MutableMapping[str, object])
Return type:: MutableMapping[str, object]

class ufp.workflows.stages.TrainStage(name='train', model_key='model', train_loader_key='train_loader', history_key='training_history', optimizer_key='optimizer', freeze_state_key='freeze_state', freeze_selectors=(), optimizer_factory=<class 'torch.optim.adam.Adam'>, optimizer_kwargs=<factory>, fit_kwargs=<factory>)[source]¶

Bases: object

Wrapper around optimizer training with optional coefficient freeze masks.

Parameters:

name (str)
model_key (str)
train_loader_key (str)
history_key (str)
optimizer_key (str)
freeze_state_key (str)
freeze_selectors (Sequence[int | str])
optimizer_factory (Callable[[...], Optimizer])
optimizer_kwargs (Mapping[str, object])
fit_kwargs (Mapping[str, object])

property metadata: Mapping[str, object]¶: Return JSON-friendly stage metadata.

optimizer_factory¶: alias of Adam

property produced_outputs: tuple[str, ...]¶: Context keys produced by this stage.

property required_inputs: tuple[str, ...]¶: Context keys required by this stage.

run(context)[source]¶

Train a model with optional coefficient freeze masks.

Parameters:: context (Mapping[str, object])
Return type:: StageResult

class ufp.workflows.stages.ValidateStage(name='validate', model_key='model', loader_key='validation_loader', metrics_key='validation_metrics', evaluate_kwargs=<factory>)[source]¶

Bases: object

Validation/metrics stage that delegates to evaluate_model.

Parameters:

name (str)
model_key (str)
loader_key (str)
metrics_key (str)
evaluate_kwargs (Mapping[str, object])

property metadata: Mapping[str, object]¶: Return JSON-friendly stage metadata.

property produced_outputs: tuple[str, ...]¶: Context keys produced by this stage.

property required_inputs: tuple[str, ...]¶: Context keys required by this stage.

run(context)[source]¶

Evaluate a model and return aggregate metrics.

Parameters:: context (Mapping[str, object])
Return type:: StageResult

class ufp.workflows.stages.WorkflowStage(*args, **kwargs)[source]¶

Bases: Protocol

Protocol implemented by lightweight workflow stages.

property metadata: Mapping[str, object]¶: JSON-friendly stage metadata.

property produced_outputs: tuple[str, ...]¶: Context keys produced by the stage.

property required_inputs: tuple[str, ...]¶: Context keys required by the stage.

run(context)[source]¶

Execute the stage against a user-owned context mapping.

Parameters:: context (Mapping[str, object])
Return type:: StageResult

ufp.workflows.stages.workflow_stage_metadata(stages_or_results, *, name='workflow')[source]¶

Return checkpoint-ready metadata for an explicit stage sequence.

Parameters:

stages_or_results (Sequence[WorkflowStage | StageResult | Mapping[str, object]])
name (str)

Return type:

dict[str, object]

Self-describing workflow checkpoint helpers.

exception ufp.workflows.checkpoints.WorkflowCheckpointError[source]¶

Bases: ValueError

Raised when a workflow checkpoint is missing or incompatible.

ufp.workflows.checkpoints.build_workflow_checkpoint(model, *, fit_blocks=None, freeze_blocks=(), stage_metadata=None, projection_diagnostics=None, validation_metrics=None, metadata=None)[source]¶

Build a plain dictionary workflow checkpoint payload.

Parameters:

model (UFPModel)
fit_blocks (Sequence[int | str | CoefficientSelector] | None)
freeze_blocks (Sequence[int | str | CoefficientSelector])
stage_metadata (Mapping[str, object] | None)
projection_diagnostics (object | None)
validation_metrics (object | None)
metadata (Mapping[str, object] | None)

Return type:

dict[str, object]

ufp.workflows.checkpoints.coefficient_layout_metadata(layout)[source]¶

Return checkpoint metadata for a model coefficient layout.

Parameters:: layout (ParameterLayout)
Return type:: dict[str, object]

ufp.workflows.checkpoints.fixed_coefficient_hashes(layout, *, fit_blocks=None, freeze_blocks=())[source]¶

Return value hashes for coefficients outside the selected solve layout.

Parameters:

layout (ParameterLayout)
fit_blocks (Sequence[int | str | CoefficientSelector] | None)
freeze_blocks (Sequence[int | str | CoefficientSelector])

Return type:

dict[str, object]

ufp.workflows.checkpoints.load_workflow_checkpoint(path, model=None, *, map_location='cpu', load_state_dict=True, strict=True, validate_model=True, validate_fixed_coefficients=False)[source]¶

Load and optionally apply a self-describing workflow checkpoint.

Parameters:

path (Path | str)
model (UFPModel | None)
map_location (Callable[[Storage, str], Storage] | device | str | dict[str, str] | None)
load_state_dict (bool)
strict (bool)
validate_model (bool)
validate_fixed_coefficients (bool)

Return type:

Mapping[str, object]

ufp.workflows.checkpoints.model_metadata(model)[source]¶

Return model architecture and term metadata without parameter values.

Parameters:: model (UFPModel)
Return type:: dict[str, object]

ufp.workflows.checkpoints.normalize_checkpoint_metadata(value)[source]¶

Return value converted into workflow-checkpoint metadata values.

Parameters:: value (Any)
Return type:: Any

ufp.workflows.checkpoints.save_workflow_checkpoint(path, model, *, fit_blocks=None, freeze_blocks=(), stage_metadata=None, projection_diagnostics=None, validation_metrics=None, metadata=None)[source]¶

Save a self-describing workflow checkpoint and return the payload.

Parameters:

path (Path | str)
model (UFPModel)
fit_blocks (Sequence[int | str | CoefficientSelector] | None)
freeze_blocks (Sequence[int | str | CoefficientSelector])
stage_metadata (Mapping[str, object] | None)
projection_diagnostics (object | None)
validation_metrics (object | None)
metadata (Mapping[str, object] | None)

Return type:

dict[str, object]

ufp.workflows.checkpoints.selector_metadata(layout, *, fit_blocks=None, freeze_blocks=())[source]¶

Resolve and serialize fit/freeze selector metadata for one layout.

Parameters:

layout (ParameterLayout)
fit_blocks (Sequence[int | str | CoefficientSelector] | None)
freeze_blocks (Sequence[int | str | CoefficientSelector])

Return type:

dict[str, object]

ufp.workflows.checkpoints.validate_fixed_coefficient_hashes(model, fixed_hashes)[source]¶

Validate stored fixed-coefficient hashes against the current model.

Parameters:

model (UFPModel)
fixed_hashes (Mapping[str, object])

Return type:

None

ufp.workflows.checkpoints.validate_model_metadata(model, checkpoint_model_metadata)[source]¶

Raise if a current model is incompatible with checkpoint metadata.

Parameters:

model (UFPModel)
checkpoint_model_metadata (Mapping[str, object])

Return type:

None

ufp.workflows.checkpoints.validate_workflow_checkpoint(payload)[source]¶

Validate the outer schema and required fields of a checkpoint payload.

Parameters:: payload (object)
Return type:: Mapping[str, object]

Residual-label materialization for frozen model components.

class ufp.workflows.residuals.ResidualDatasetResult(dataset, metadata)[source]¶

Bases: object

Residualized ASE dataset and the metadata used to materialize it.

Parameters:

dataset (ASEAtomsDataset)
metadata (Mapping[str, object])

property metadata_hash: str¶: Return the deterministic residual metadata hash.

ufp.workflows.residuals.materialize_residual_dataset(model, dataset, *, selectors=None, term_filter=None, targets=None, target_keys=None, target_weights=None, units=None, require_frozen=True, dtype=None, device=None, backend=None, metadata_key='residuals')[source]¶

Subtract selected frozen model contributions from supervised ASE labels.

Parameters:

model (UFPotential) – Model that owns the frozen terms to subtract.
dataset (SampleDataset) – Input ASE training samples.
selectors (Sequence[int | str | type[UFPTerm]] | None) – Optional term selectors. Integers match evaluation-order term indices; strings match indices, term families, class names, or parameter-block labels/kinds/groups; term classes match by isinstance.
term_filter (Callable[[int, UFPTerm], bool] | None) – Optional callable receiving (index, term).
targets (Sequence[str] | None) – Labels to residualize. If omitted, all labels present on every sample are residualized.
target_keys (Mapping[str, str] | None) – Label keys stored into copied ASE objects and metadata.
target_weights (Mapping[str, float] | None) – Optional downstream target weights to record in metadata. These do not scale residual labels.
units (Mapping[str, str] | None) – Target unit strings recorded in metadata; values are not converted.
require_frozen (bool) – If True, reject selected terms with trainable parameters.
dtype (dtype | None) – Evaluation dtype. Defaults to model.preferred_dtype().
device (device | str | None) – Optional evaluation device.
backend (str | NeighborListBackend | None) – Optional neighbor-list backend override.
metadata_key (str) – Per-sample metadata key that stores residual metadata.

Returns:

Residualized dataset plus deterministic metadata.

Return type:

ResidualDatasetResult

ufp.workflows.residuals.residualize_ase_dataset(model, dataset, **kwargs)[source]¶

Alias for materialize_residual_dataset().

Parameters:

model (UFPotential)
dataset (SampleDataset)
kwargs (Any)

Return type:

ResidualDatasetResult

Advanced prepared-geometry helpers for workflow-level cache reuse.

class ufp.workflows.prepared.PairCategorySpec(atomic_types, symmetric=True)[source]¶

Bases: object

Pair-category request stored by a prepared geometry object.

Parameters:

atomic_types (tuple[int, ...])
symmetric (bool)

class ufp.workflows.prepared.PreparedGeometry(positions, cell, pbc, atomic_numbers, system_index, neighbor_list, pair_system_index, pair_vectors, pair_distances, pair_category_indices, metadata, input_metadata=<factory>, triplet_caches=<factory>, source_atoms=None)[source]¶

Bases: object

Reusable tensorized geometry for advanced workflow caching.

This object intentionally stays outside term forward dispatch. Convert it back to UFPInput when a runtime model path is needed.

Parameters:

positions (Tensor)
cell (Tensor)
pbc (Tensor)
atomic_numbers (Tensor)
system_index (Tensor)
neighbor_list (NeighborListData | None)
pair_system_index (Tensor | None)
pair_vectors (Tensor | None)
pair_distances (Tensor | None)
pair_category_indices (Mapping[PairCategorySpec, Tensor])
metadata (PreparedGeometryMetadata)
input_metadata (Mapping[str, object])
triplet_caches (Mapping[str, object])
source_atoms (Sequence[Atoms] | None)

assert_valid_for(inputs, *, species_ordering=None, cutoff=None)[source]¶

Raise if inputs no longer match this prepared geometry metadata.

Parameters:

inputs (UFPInput)
species_ordering (Sequence[int] | None)
cutoff (float | None)

Return type:

None

is_valid_for(inputs, *, species_ordering=None, cutoff=None)[source]¶

Return whether inputs match this prepared geometry metadata.

Parameters:

inputs (UFPInput)
species_ordering (Sequence[int] | None)
cutoff (float | None)

Return type:

bool

property memory_nbytes: int¶: Return an estimated memory footprint for owned tensor payloads.

property n_atoms: int¶: Return the number of atoms in the prepared geometry.

property n_pairs: int¶: Return the number of neighbor-list rows.

property n_systems: int¶: Return the number of systems in the prepared geometry.

property pair_geometry_nbytes: int¶: Return the estimated memory footprint of reusable pair geometry.

reuse_estimate(*, n_reuses, memory_budget_bytes=None)[source]¶

Estimate when prepared pair geometry reuse is worth its memory cost.

Parameters:

n_reuses (int)
memory_budget_bytes (int | None)

Return type:

PreparedGeometryReuseEstimate

to_input(*, device=None, dtype=None, requires_grad=False, include_triplet_caches=True)[source]¶

Convert prepared geometry back to the runtime UFPInput container.

Parameters:

device (device | str | None)
dtype (dtype | None)
requires_grad (bool)
include_triplet_caches (bool)

Return type:

PreparedGeometryMetadata

property triplet_cache_nbytes: int¶: Return the estimated memory footprint of opaque triplet caches.

class ufp.workflows.prepared.PreparedGeometryMetadata(schema_version, species_ordering, cutoff, neighbor_backend, neighbor_full_list, neighbor_sorted, neighbor_strict, device, dtype, periodicity, n_atoms, n_systems, system_sizes, source_structure_ids, source_structure_hashes, geometry_hash, metadata_hash)[source]¶

Bases: object

Strict validity metadata for one prepared geometry object.

Parameters:

schema_version (int)
species_ordering (tuple[int, ...])
cutoff (float | None)
neighbor_backend (str | None)
neighbor_full_list (bool | None)
neighbor_sorted (bool | None)
neighbor_strict (bool | None)
device (str)
dtype (str)
periodicity (tuple[tuple[bool, bool, bool], ...])
n_atoms (int)
n_systems (int)
system_sizes (tuple[int, ...])
source_structure_ids (tuple[str | None, ...])
source_structure_hashes (tuple[str, ...])
geometry_hash (str)
metadata_hash (str)

classmethod from_input(inputs, *, species_ordering=None, cutoff=None)[source]¶

Build strict prepared-geometry metadata from a normalized input.

Parameters:

inputs (UFPInput)
species_ordering (Sequence[int] | None)
cutoff (float | None)

Return type:

to_dict()[source]¶

Return JSON-friendly metadata.

Return type:: dict[str, object]

class ufp.workflows.prepared.PreparedGeometryReuseEstimate(n_reuses, n_atoms, n_pairs, total_nbytes, pair_geometry_nbytes, triplet_cache_nbytes, avoided_pair_geometry_nbytes, reuse_helpful, memory_dominated)[source]¶

Bases: object

Memory and reuse estimate for a prepared geometry object.

Parameters:

n_reuses (int)
n_atoms (int)
n_pairs (int)
total_nbytes (int)
pair_geometry_nbytes (int)
triplet_cache_nbytes (int)
avoided_pair_geometry_nbytes (int)
reuse_helpful (bool)
memory_dominated (bool)

ufp.workflows.prepared.prepare_geometry_from_ase(atoms, *, cutoff, backend=NeighborListBackend.AUTO, full_list=True, sorted=True, dtype=None, device=None, species_ordering=None, pair_category_specs=None, triplet_terms=(), feature_cache_storage='none', feature_cache_dir=None, cache_prefix='prepared', include_per_atom_energy=True)[source]¶

Build prepared geometry directly from one or more ASE structures.

Parameters:

atoms (Atoms | Sequence[Atoms])
cutoff (float)
backend (str | NeighborListBackend)
full_list (bool)
sorted (bool)
dtype (dtype | None)
device (device | str | None)
species_ordering (Sequence[int] | None)
pair_category_specs (Sequence[PairCategorySpec | Sequence[int]] | None)
triplet_terms (Sequence[object])
feature_cache_storage (str)
feature_cache_dir (Path | str | None)
cache_prefix (str)
include_per_atom_energy (bool)

Return type:

ufp.workflows.prepared.prepare_geometry_from_batch(batch, model, *, backend=None, dtype=None, device=None, species_ordering=None, pair_category_specs=None, triplet_terms=(), feature_cache_storage='none', feature_cache_dir=None, cache_prefix='prepared', include_per_atom_energy=True)[source]¶

Prepare geometry from a training ASEAtomsBatch.

Parameters:

batch (ASEAtomsBatch)
model (UFPotential)
backend (str | NeighborListBackend | None)
dtype (dtype | None)
device (device | str | None)
species_ordering (Sequence[int] | None)
pair_category_specs (Sequence[PairCategorySpec | Sequence[int]] | None)
triplet_terms (Sequence[object])
feature_cache_storage (str)
feature_cache_dir (Path | str | None)
cache_prefix (str)
include_per_atom_energy (bool)

Return type:

ufp.workflows.prepared.prepare_geometry_from_input(inputs, *, species_ordering=None, pair_category_specs=None, triplet_terms=(), feature_cache_storage='none', feature_cache_dir=None, cache_prefix='prepared', include_per_atom_energy=True)[source]¶

Materialize reusable geometry tensors from an existing UFPInput.

Parameters:

inputs (UFPInput)
species_ordering (Sequence[int] | None)
pair_category_specs (Sequence[PairCategorySpec | Sequence[int]] | None)
triplet_terms (Sequence[object])
feature_cache_storage (str)
feature_cache_dir (Path | str | None)
cache_prefix (str)
include_per_atom_energy (bool)

Return type:

ufp.workflows.prepared.prepare_geometry_from_prepared_batch(batch, *, species_ordering=None, pair_category_specs=None, triplet_terms=(), feature_cache_storage='none', feature_cache_dir=None, cache_prefix='prepared', include_per_atom_energy=True)[source]¶

Prepare geometry from a least-squares PreparedBatch.

Parameters:

batch (PreparedBatch)
species_ordering (Sequence[int] | None)
pair_category_specs (Sequence[PairCategorySpec | Sequence[int]] | None)
triplet_terms (Sequence[object])
feature_cache_storage (str)
feature_cache_dir (Path | str | None)
cache_prefix (str)
include_per_atom_energy (bool)

Return type:

Adapters¶

Status: Optional integration API.

Public adapter surface for ASE, metatomic, and torch-sim integrations.

Use this namespace when selecting an engine-specific wrapper around a UFP model.

ASE calculator integration for UFP potentials.

Use this adapter when an ase.Atoms workflow should consume UFP energies, forces, and optional per-atom outputs.

class ufp.adapters.ase.UFPASECalculator(potential, *, neighbor_backend=NeighborListBackend.AUTO, **kwargs)[source]¶

Bases: Calculator

Wrap a UFP potential as an ASE calculator.

Parameters:

potential (UFPotential) – Potential to evaluate.
neighbor_backend (Union[str, NeighborListBackend]) – Backend used when the potential builds neighbor lists.

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

Run the wrapped potential and populate the ASE results dictionary.

Parameters:

properties (Sequence[str] | None)
system_changes (Iterable[str] | None)

Return type:

None

implemented_properties: list[str] = ['energy', 'free_energy', 'forces', 'stress', 'energies']¶: Properties calculator can handle (energy, forces, …)

Metatomic conversion and wrapper utilities for UFP models.

This module translates metatomic systems and outputs into the shared UFP input and output structures without changing model logic.

class ufp.adapters.metatomic.UFPMetatomicModule(potential, *, full_neighbor_list=True, strict_neighbor_list=True)[source]¶

Bases: Module

Wrap a UFP potential for metatomic execution.

The wrapped UFP potential still owns all chemistry and neighbor-list consumption. This adapter only converts metatomic systems to the shared UFP input bundle and returns metatensor outputs.

Parameters:

potential (UFPotential) – Wrapped UFP potential.
full_neighbor_list (bool) – Whether the requested metatomic neighbor list should contain both directions for every pair.
strict_neighbor_list (bool) – Whether the requested neighbor list should be strictly within the cutoff.

forward(systems, outputs, selected_atoms=None)[source]¶

Evaluate the wrapped UFP potential inside a metatomic workflow.

Parameters:

systems – Metatomic systems passed by the engine.
outputs – Requested metatomic outputs.
selected_atoms – Optional atom selection. This is not implemented.

Returns:

Dictionary mapping output names to metatensor tensor maps.

Raises:

NotImplementedError – If selected_atoms is provided.

Return type:

Dict[str, object]

requested_neighbor_lists()[source]¶

Declare the neighbor list needed by the wrapped UFP model.

Returns:: A singleton list of NeighborListOptions or an empty list.

ufp.adapters.metatomic.make_energy_output(energies)[source]¶

Create a metatensor TensorMap for batched total energies.

Parameters:: energies (Tensor) – Tensor with shape (n_systems,) or (n_systems, 1).
Returns:: Metatensor tensor map storing total energies.

ufp.adapters.metatomic.make_per_atom_output(values, system_sizes, property_name='feature')[source]¶

Create a metatensor TensorMap for batched per-atom outputs.

Parameters:

values (Tensor) – Tensor with shape (n_atoms_total,) or (n_atoms_total, n_properties).
system_sizes (Sequence[int]) – Atom counts for each system in the batch.
property_name (str) – Name to use for the output properties axis.

Returns:

Metatensor tensor map storing the per-atom values.

Raises:

ImportError – If metatensor-torch is not installed.
ValueError – If values has an unsupported shape or atom count.

ufp.adapters.metatomic.make_system_output(values, property_name)[source]¶

Create a metatensor TensorMap for per-system outputs.

Parameters:

values (Tensor) – Tensor with shape (n_systems,) or (n_systems, n_properties).
property_name (str) – Name of the property axis.

Returns:

Metatensor tensor map storing the values.

Raises:

ImportError – If metatensor-torch is not installed.
ValueError – If values does not have one or two dimensions.

ufp.adapters.metatomic.neighbor_list_to_data(neighbors, options=None, *, atom_offset=0)[source]¶

Convert a metatomic neighbor-list block into UFP neighbor-list data.

The expected metatomic convention is a block with sample names ["first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"] and distance vectors stored in neighbors.values.

Parameters:

neighbors – Metatensor-like tensor block holding a neighbor list.
options – Optional object exposing cutoff, full_list, and strict.
atom_offset (int) – Offset applied to local atom indices.

Returns:

Normalized neighbor-list data.

Raises:

ValueError – If sample names or vector shapes do not match metatomic neighbor-list conventions.

Return type:

ufp.adapters.metatomic.system_to_ase_atoms(system)[source]¶

Convert a metatomic system-like object into ASE atoms.

This helper intentionally uses duck typing so it can operate on scripted or Python metatomic systems without importing metatomic at module import time.

Parameters:: system – Object with positions, cell, pbc, and types fields.
Returns:: ASE structure with the same geometry.
Raises:: TypeError – If system does not expose the required geometry fields.
Return type:: Atoms

ufp.adapters.metatomic.systems_to_ase_atoms(systems)[source]¶

Convert metatomic system-like objects into ASE atoms.

Parameters:: systems – Iterable of system-like objects.
Returns:: ASE structures in input order.
Return type:: list[Atoms]

ufp.adapters.metatomic.systems_to_input(systems, *, neighbor_options=None, metadata=None)[source]¶

Convert metatomic systems into a shared UFP input bundle.

Parameters:

systems – Sequence of system-like objects.
neighbor_options – Optional neighbor-list request used to fetch neighbors.
metadata (Dict[str, object] | None) – Optional metadata dictionary copied into the input.

Returns:

Normalized UFP input bundle.

Raises:

ValueError – If no systems are provided.
TypeError – If any system does not expose the required geometry fields.

Return type: