nvalchemi.training.EnergyHuberLoss#

class nvalchemi.training.EnergyHuberLoss(*, target_key='energy', prediction_key='predicted_energy', per_atom=True, delta=0.01, ignore_nonfinite=True, dtype_policy='strict')[source]#

Huber loss on total energy or energy per atom.

The elementwise Huber function \(H_\delta\) (quadratic within \(\delta\) of zero, linear beyond it; see the module docstring) is applied to each per-graph energy residual, then averaged over the \(B\) labeled graphs:

\[L = \frac{1}{B} \sum_{i=1}^{B} H_\delta\!\left(\hat{E}_i - E_i\right).\]

With per_atom=True (default) the prediction and target are first divided by each graph’s atom count \(N_i\), so the Huber function acts on energy-per-atom residuals \((\hat{E}_i - E_i) / N_i\). Unlike the per-atom MSE/MAE terms, the final reduction is an unweighted mean over labeled structures rather than an atom-count-weighted mean.

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

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

  • per_atom (bool, default True) – Divide prediction and target by num_nodes_per_graph before computing Huber residuals.

  • delta (float, default 0.01) – Positive transition point between quadratic and linear Huber regimes.

  • ignore_nonfinite (bool, default True) – When True, target entries that are NaN or infinite 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 elementwise Huber losses, zeroing invalid entries.

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]#

Exclude non-finite target entries when ignore_nonfinite=True.

Parameters:
Return type:

Tensor

normalize(pred, target, **kwargs)[source]#

Divide by atom counts when per_atom=True.

Parameters:
Return type:

tuple[Tensor, Tensor, ReductionContext]