Cartesian losses

Cartesian losses compare sample-level quantities such as total energies or atom-level quantities such as forces. Componentwise absolute and squared errors apply to scalars, vectors, and tensors; vector-norm errors treat the final Cartesian axis as a single geometric vector.

For extensive sample-level quantities, the per-atom variants divide the residual by the sample's atom count before measuring the error. Atom-level losses instead receive one row per atom and can give equal weight to atoms or to samples when taking the mean.

Sample-level values

mae_loss

mae_loss(input: Tensor, target: Tensor, *, reduction: _Reduction = 'mean') -> Tensor

Return the componentwise absolute error of sample-level values.

With reduction="none", the elementary errors are

\[ \ell_{bc} = \left|x_{bc} - y_{bc}\right|. \]

reduction="mean" returns

\[ L = \frac{1}{BC}\sum_{b,c}\ell_{bc}, \]

where \(C\) is the product of value_shape; \(C=1\) for one scalar per sample. reduction="sum" returns the corresponding unnormalized sum. Because the absolute errors are taken component by component, this loss is generally not invariant under rotations when the value axes contain Cartesian vectors or tensors.

Parameters:

input (Tensor) –

Predicted sample-level values with shape [B, *value_shape].
target (Tensor) –

Target values with shape [B, *value_shape].
reduction (_Reduction, default: 'mean' ) –

"none" returns elementary errors with shape [B, *value_shape]; "mean" averages all elementary errors; "sum" sums them.

Returns:

Tensor –

Loss tensor with the shape determined by reduction.

mse_loss

mse_loss(input: Tensor, target: Tensor, *, reduction: _Reduction = 'mean') -> Tensor

Return the componentwise squared error of sample-level values.

With reduction="none", the elementary errors are

\[ \ell_{bc} = \left(x_{bc} - y_{bc}\right)^2. \]

reduction="mean" returns

\[ L = \frac{1}{BC}\sum_{b,c}\ell_{bc}, \]

where \(C\) is the product of value_shape; \(C=1\) for one scalar per sample. reduction="sum" returns the corresponding unnormalized sum. No factor of \(1/2\) is applied.

Parameters:

input (Tensor) –

Predicted sample-level values with shape [B, *value_shape].
target (Tensor) –

Target values with shape [B, *value_shape].
reduction (_Reduction, default: 'mean' ) –

"none" returns elementary errors with shape [B, *value_shape]; "mean" averages all elementary errors; "sum" sums them.

Returns:

Tensor –

Loss tensor with the shape determined by reduction.

vector_norm_loss

vector_norm_loss(
    input: Tensor, target: Tensor, *, reduction: _Reduction = "mean"
) -> Tensor

Return Euclidean errors of sample-level Cartesian vectors.

The final axis is interpreted as one Cartesian vector and is consumed by the Euclidean norm. With reduction="none", the elementary errors are

\[ \ell_{bc} = \sqrt{\sum_{\alpha=1}^{3} \left(x_{bc\alpha} - y_{bc\alpha}\right)^2}. \]

reduction="mean" returns

\[ L = \frac{1}{BC}\sum_{b,c}\ell_{bc}, \]

where \(C\) is the product of extra_shape; \(C=1\) for one vector per sample. reduction="sum" returns the corresponding unnormalized sum. Unlike componentwise MAE, the elementary vector errors are invariant under simultaneous orthogonal transformations of input and target.

Parameters:

input (Tensor) –

Predicted sample-level Cartesian vectors with shape [B, *extra_shape, 3].
target (Tensor) –

Target vectors with shape [B, *extra_shape, 3].
reduction (_Reduction, default: 'mean' ) –

"none" returns vector errors with shape [B, *extra_shape]; "mean" averages all vector errors; "sum" sums them.

Returns:

Tensor –

Loss tensor with the shape determined by reduction.

Extensive values per atom

These functions still take one row per sample. For squared error, dividing the residual by the atom count introduces the square of that count in the denominator.

per_atom_mae_loss

per_atom_mae_loss(
    input: Tensor, target: Tensor, num_atoms: Tensor, *, reduction: _Reduction = "mean"
) -> Tensor

Return the per-atom absolute error of extensive sample-level values.

This loss is for an extensive quantity such as total energy: the residual is divided by the number of atoms in its sample before the absolute error is formed. With reduction="none",

\[ \ell_{bc} = \left|\frac{x_{bc} - y_{bc}} {N_{\mathrm{atom}}^{(b)}}\right|. \]

reduction="mean" averages these errors over all samples and value components. For one scalar per sample,

\[ L = \frac{1}{B}\sum_b \frac{\left|x_b-y_b\right|}{N_{\mathrm{atom}}^{(b)}}. \]

This differs from an atom-level loss: input and target have one row per sample, not one row per atom.

Parameters:

input (Tensor) –

Predicted extensive sample-level values with shape [B, *value_shape].
target (Tensor) –

Target values with shape [B, *value_shape].
num_atoms (Tensor) –

Number of atoms in each sample, with shape [B].
reduction (_Reduction, default: 'mean' ) –

"none" returns per-atom elementary errors with shape [B, *value_shape]; "mean" averages them; "sum" sums them.

Returns:

Tensor –

Loss tensor with the shape determined by reduction.

per_atom_mse_loss

per_atom_mse_loss(
    input: Tensor, target: Tensor, num_atoms: Tensor, *, reduction: _Reduction = "mean"
) -> Tensor

Return the per-atom squared error of extensive sample-level values.

This loss is for an extensive quantity such as total energy: the residual is divided by the number of atoms in its sample before it is squared. With reduction="none",

\[ \ell_{bc} = \left(\frac{x_{bc} - y_{bc}} {N_{\mathrm{atom}}^{(b)}}\right)^2. \]

reduction="mean" averages these errors over all samples and value components. For one scalar per sample,

\[ L = \frac{1}{B}\sum_b \left(\frac{x_b-y_b}{N_{\mathrm{atom}}^{(b)}}\right)^2. \]

The division therefore contributes \((N_{\mathrm{atom}}^{(b)})^{-2}\), not \((N_{\mathrm{atom}}^{(b)})^{-1}\), to the squared loss. No factor of \(1/2\) is applied.

Parameters:

input (Tensor) –

Predicted extensive sample-level values with shape [B, *value_shape].
target (Tensor) –

Target values with shape [B, *value_shape].
num_atoms (Tensor) –

Number of atoms in each sample, with shape [B].
reduction (_Reduction, default: 'mean' ) –

"none" returns per-atom elementary errors with shape [B, *value_shape]; "mean" averages them; "sum" sums them.

Returns:

Tensor –

Loss tensor with the shape determined by reduction.

Atom-level values

mean_over chooses whether each atom or each sample has equal weight in the mean reduction.

atom_mae_loss

atom_mae_loss(
    input: Tensor,
    target: Tensor,
    num_atoms: Tensor,
    *,
    mean_over: _MeanOver = "atoms",
    reduction: _Reduction = "mean",
) -> Tensor

Return the componentwise absolute error of packed atom-level values.

The elementary errors are

\[ \ell_{ic} = \left|x_{ic} - y_{ic}\right|. \]

With reduction="mean" and mean_over="atoms", every atom and value component in the batch has equal weight:

\[ L_{\mathrm{atoms}} = \frac{1}{N_{\mathrm{atom}}C} \sum_{i,c}\ell_{ic}. \]

With mean_over="samples", every sample instead has equal weight, independently of its number of atoms:

\[ L_{\mathrm{samples}} = \frac{1}{B}\sum_b \frac{1}{N_{\mathrm{atom}}^{(b)}C} \sum_{i\in b,c}\ell_{ic}. \]

Here \(C\) is the product of value_shape; \(C=1\) for one scalar per atom. mean_over affects only the "mean" reduction. "none" returns every elementary error and "sum" returns their unnormalized sum. This componentwise loss is generally not invariant under rotations when the value axes contain Cartesian vectors or tensors.

Parameters:

input (Tensor) –

Predicted packed atom-level values with shape [N_atom, *value_shape].
target (Tensor) –

Target values with shape [N_atom, *value_shape].
num_atoms (Tensor) –

Number of consecutive atoms in each sample, with shape [B] and sum N_atom.
mean_over (_MeanOver, default: 'atoms' ) –

"atoms" gives every atom equal weight; "samples" gives every sample equal weight.
reduction (_Reduction, default: 'mean' ) –

"none" returns elementary errors with shape [N_atom, *value_shape]; "mean" averages them according to mean_over; "sum" sums them.

Returns:

Tensor –

Loss tensor with the shape determined by reduction.

atom_mse_loss

atom_mse_loss(
    input: Tensor,
    target: Tensor,
    num_atoms: Tensor,
    *,
    mean_over: _MeanOver = "atoms",
    reduction: _Reduction = "mean",
) -> Tensor

Return the componentwise squared error of packed atom-level values.

The elementary errors are

\[ \ell_{ic} = \left(x_{ic} - y_{ic}\right)^2. \]

With reduction="mean" and mean_over="atoms", every atom and value component in the batch has equal weight:

\[ L_{\mathrm{atoms}} = \frac{1}{N_{\mathrm{atom}}C} \sum_{i,c}\ell_{ic}. \]

With mean_over="samples", every sample instead has equal weight:

\[ L_{\mathrm{samples}} = \frac{1}{B}\sum_b \frac{1}{N_{\mathrm{atom}}^{(b)}C} \sum_{i\in b,c}\ell_{ic}. \]

Here \(C\) is the product of value_shape; \(C=1\) for one scalar per atom. mean_over affects only the "mean" reduction. "none" returns every elementary error and "sum" returns their unnormalized sum. No factor of \(1/2\) is applied.

Parameters:

input (Tensor) –

Predicted packed atom-level values with shape [N_atom, *value_shape].
target (Tensor) –

Target values with shape [N_atom, *value_shape].
num_atoms (Tensor) –

Number of consecutive atoms in each sample, with shape [B] and sum N_atom.
mean_over (_MeanOver, default: 'atoms' ) –

"atoms" gives every atom equal weight; "samples" gives every sample equal weight.
reduction (_Reduction, default: 'mean' ) –

"none" returns elementary errors with shape [N_atom, *value_shape]; "mean" averages them according to mean_over; "sum" sums them.

Returns:

Tensor –

Loss tensor with the shape determined by reduction.

atom_vector_norm_loss

atom_vector_norm_loss(
    input: Tensor,
    target: Tensor,
    num_atoms: Tensor,
    *,
    mean_over: _MeanOver = "atoms",
    reduction: _Reduction = "mean",
) -> Tensor

Return Euclidean errors of packed atom-level Cartesian vectors.

The final axis is interpreted as one Cartesian vector and is consumed by the Euclidean norm. The elementary errors are

\[ \ell_{ic} = \sqrt{\sum_{\alpha=1}^{3} \left(x_{ic\alpha} - y_{ic\alpha}\right)^2}. \]

With reduction="mean" and mean_over="atoms", every atom and extra vector in the batch has equal weight:

\[ L_{\mathrm{atoms}} = \frac{1}{N_{\mathrm{atom}}C} \sum_{i,c}\ell_{ic}. \]

With mean_over="samples", every sample instead has equal weight:

\[ L_{\mathrm{samples}} = \frac{1}{B}\sum_b \frac{1}{N_{\mathrm{atom}}^{(b)}C} \sum_{i\in b,c}\ell_{ic}. \]

Here \(C\) is the product of extra_shape; \(C=1\) for one vector per atom. mean_over affects only the "mean" reduction. "none" returns every vector error and "sum" returns their unnormalized sum. Each elementary error is invariant under simultaneous orthogonal transformations of input and target.

Parameters:

input (Tensor) –

Predicted packed atom-level Cartesian vectors with shape [N_atom, *extra_shape, 3].
target (Tensor) –

Target vectors with shape [N_atom, *extra_shape, 3].
num_atoms (Tensor) –

Number of consecutive atoms in each sample, with shape [B] and sum N_atom.
mean_over (_MeanOver, default: 'atoms' ) –

"atoms" gives every atom equal weight; "samples" gives every sample equal weight.
reduction (_Reduction, default: 'mean' ) –

"none" returns vector errors with shape [N_atom, *extra_shape]; "mean" averages them according to mean_over; "sum" sums them.

Returns:

Tensor –

Loss tensor with the shape determined by reduction.