Ctrl+K

nvalchemi.training.ForceL2NormLoss#

class nvalchemi.training.ForceL2NormLoss(*, target_key='forces', prediction_key='predicted_forces', normalize_by_atom_count=True, ignore_nonfinite=True, dtype_policy='strict')[source]#

Mean per-atom force-vector L2 loss.

Unlike the component-wise force terms, the per-atom residual here is the Euclidean norm of the force-vector error,

\[\rho_{ia} = \bigl\|\hat{\mathbf{F}}_{ia} - \mathbf{F}_{ia}\bigr\|_2 = \sqrt{\sum_{\alpha=1}^{3} \left(\hat{F}_{ia\alpha} - F_{ia\alpha}\right)^2},\]

reduced over atoms according to normalize_by_atom_count:

  • normalize_by_atom_count=True (default): the mean atom norm per graph is averaged over graphs (graph-balanced):

    \[L = \frac{1}{B} \sum_{i=1}^{B} \frac{1}{N_i} \sum_{a=1}^{N_i} \rho_{ia}.\]

  • normalize_by_atom_count=False: one global mean over all valid atoms, \(L = \tfrac{1}{V} \sum_{i=1}^{B} \sum_{a=1}^{N_i} \rho_{ia}\).

Dense (V, 3) inputs can be graph-balanced with batch_idx and num_graphs. Padded (B, V_max, 3) inputs require num_nodes_per_graph counts or a node mask so padding can be excluded from the atom-level reduction. The reduction is over atoms (one norm per atom), not over individual components.

Parameters:

  • target_key (str, default "forces") – Target container key for the target tensor.

  • prediction_key (str, default "predicted_forces") – Prediction container key for the model output.

  • normalize_by_atom_count (bool, default True) – When True, compute a mean atom L2 norm per graph, then mean over graphs. When False, compute one global mean over valid atom L2 norms.

  • ignore_nonfinite (bool, default True) – When True, atoms whose target vector contains NaN or infinity are excluded from both loss value and gradient using torch.isfinite().

  • dtype_policy ({"strict", "prediction_to_target", "target_to_prediction"}, default "strict") – How to handle prediction/target dtype mismatches before validation. "strict" raises; the other policies cast one tensor to match the other.

compute_residual(pred, target, valid)[source]#

Return per-atom L2 norm of force residuals, zeroing invalid atoms.

Parameters:
Return type:

Tensor

extra_repr()[source]#

Human-readable hyperparameter summary for nn.Module’s repr.

Return type:

str

mask(pred, target, ctx, **kwargs)[source]#

Return atom-level validity mask (not component-level) for forces.

The mask has shape (V,) for dense or (B, V_max) for padded forces — one validity flag per atom, not per component.

Parameters:
Return type:

Tensor

reduce(residual, valid, ctx, **kwargs)[source]#

Reduce per-atom L2 norms to a scalar loss.

Parameters:
Return type:

Tensor

On this page
Show Source