xnn.common.train.losses.weighted_loss#

xnn.common.train.losses.weighted_loss(pred, data, energy_weight, force_weight, stress_weight, *, huber_delta=0.0, huber_delta_energy=None, huber_delta_forces=None, huber_delta_stress=None)[source]#

Compute a weighted sum of energy, force and stress loss terms.

Each term is included only when the target is present in data and its weight is positive (and, for forces and stress, when the corresponding key is present in pred). The individual terms are:

Energy: per-atom energy error, i.e. the difference between the predicted and target total energy divided by the number of atoms in each structure, then averaged over structures. Dividing by the atom count makes the term size-extensive and well-scaled across structures of different sizes.
Forces: per-atom force error, averaged over every atom and Cartesian component in the batch.
Stress: error between predicted and target stress tensors.

Per-structure weighting. When data.weight is set (shape (B,)), every term becomes a weighted mean instead of a plain one: structure b contributes in proportion to weight[b], and for the force term that weight is spread to each of its atoms. Uniform weights reproduce the unweighted loss exactly, so this is a no-op unless the dataset asks for it. It exists because the two terms otherwise disagree about what a sample is: energy counts each structure once while forces count each atom once, so a set mixing small and large structures, or dense scans with sparse sampling, silently allocates the fit. Set weight in the structure dicts (see to_graph()) to allocate it on purpose.

Huber tails. With huber_delta > 0 the squared error is replaced by a Huber-like function that is quadratic up to delta and linear beyond, which caps the pull of a few large residuals. The three residuals have very different natural scales (per-atom energies in eV, forces in eV/A), so a single huber_delta is rarely right for all of them; the per-term arguments override it. See _residual_sq() for the exact form and how it compares with MACE’s.

Parameters:

pred (dict of str to Tensor) – Model outputs. Must contain "energy"; may also contain "forces" and "stress".
data (AtomicGraph) – The (possibly batched) target graph. Reads energy, forces, stress, n_atoms (per-structure atom counts), batch (the structure index of every atom) and the optional weight.
energy_weight (float) – Weight applied to the energy term. The term is skipped when this is not positive or data.energy is None.
force_weight (float) – Weight applied to the force term. The term is skipped when this is not positive, data.forces is None, or pred lacks "forces".
stress_weight (float) – Weight applied to the stress term. The term is skipped when this is not positive, data.stress is None, or pred lacks "stress".
huber_delta (float, optional) – Default crossover from quadratic to linear for every term. 0.0 (the default) means plain squared error throughout.
huber_delta_energy (float, optional) – Per-term overrides of huber_delta. None falls back to it.
huber_delta_forces (float, optional) – Per-term overrides of huber_delta. None falls back to it.
huber_delta_stress (float, optional) – Per-term overrides of huber_delta. None falls back to it.

Returns:

A pair (loss, logs) where loss is the scalar weighted total loss (a 0-dim tensor carrying gradients) and logs maps "loss" to the detached total plus, for each included term, its detached value under "energy_mse", "force_mse" and/or "stress_mse". Those keys keep their names when Huber tails are on, where they hold the Huber value rather than a mean square.

Return type:

tuple of (Tensor, dict of str to float)