nvalchemi.training.StressMSELoss

class nvalchemi.training.StressMSELoss(*, target_key='stress', prediction_key='predicted_stress', ignore_nonfinite=False, dtype_policy='strict')

Mean-squared-error loss on the per-graph stress tensor.

Both prediction and target are \(3 \times 3\) tensors of shape (B, 3, 3). Each graph contributes the mean of its nine squared component residuals – equivalently its squared Frobenius norm divided by the number of valid components – and these are averaged over graphs:

\[L = \frac{1}{B} \sum_{i=1}^{B} \frac{1}{9} \sum_{p=1}^{3} \sum_{q=1}^{3} \left(\hat{\sigma}_{ipq} - \sigma_{ipq}\right)^2 = \frac{1}{B} \sum_{i=1}^{B} \frac{\bigl\|\hat{\sigma}_i - \sigma_i\bigr\|_F^2}{9},\]

where the fraction shown assumes all nine components are valid. Concretely, reduce() sums the squared component residuals of each graph and divides by the number of valid components for that graph (clamped to at least 1), giving a per-graph component-mean; these per-graph values are then averaged over graphs. When ignore_nonfinite drops components the per-graph denominator is the number of remaining valid components rather than 9, and a graph whose entire stress tensor is non-finite contributes 0.0. A hat denotes the prediction; see the module docstring for the shared notation.

Tensor Contract

pred, targetStress

Per-graph stress tensors of shape (B, 3, 3). Shape validation requires exact equality.

param target_key:

Target container key for the target tensor.

type target_key:

str, default “stress”

param prediction_key:

Prediction container key for the model output.

type prediction_key:

str, default “predicted_stress”

param ignore_nonfinite:

When True, target stress components that are NaN or infinite are excluded from both loss value and gradient using torch.isfinite(). Intended for inputs that mix samples with and without stress labels. Implemented with branch-free tensor ops for torch.compile compatibility. A graph whose entire stress tensor is non-finite contributes 0.0 to the loss.

type ignore_nonfinite:

bool, default False

param dtype_policy:

How to handle prediction/target dtype mismatches before validation. "strict" raises; the other policies cast one tensor to match the other.

type dtype_policy:

{“strict”, “prediction_to_target”, “target_to_prediction”}, default “strict”

compute_residual(pred, target, valid)

Return squared stress residuals, zeroing invalid entries.

Parameters:

• pred (Tensor)

• target (Tensor)

• valid (Tensor)

Return type:

Tensor

extra_repr()

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

Return type:

str

mask(pred, target, ctx, **kwargs)

Exclude non-finite stress components when ignore_nonfinite=True.

Parameters:

• pred (Tensor)

• target (Tensor)

• ctx (ReductionContext)

• kwargs (Any)

Return type:

Tensor

reduce(residual, valid, ctx, **kwargs)

Reduce per-component stress residuals to a per-graph mean scalar.

Parameters:

• residual (Tensor)

• valid (Tensor)

• ctx (ReductionContext)

• kwargs (Any)

Return type:

Tensor

Parameters:

• target_key (str)

• prediction_key (str)

• ignore_nonfinite (bool)

• dtype_policy (DTypePolicy)