Loss Function and Error Metrics

For practical usage and configuration guidance, see the Loss and Metrics guide. This page provides technical API details.

Simplified Wrappers

The following MetricsManager wrappers can be used for common force field training scenarios:

class nequip.train.EnergyForceLoss(coeffs: Dict[str, float] = {'forces': 1.0, 'total_energy': 1.0}, per_atom_energy: bool = True, type_names: List[str] | None = None)[source]

Simplified MetricsManager wrapper for a loss term containing energy and forces mean squared errors (MSEs).

The loss component names are per_atom_energy_mse OR total_energy_mse (depending on whether per_atom_energy is True or False), and forces_mse, which are the names to refer to when neeeded, e.g. when scheduling loss component coefficients.

Example usage in config:

training_module:
  _target_: nequip.train.NequIPLightningModule

  loss:
    _target_: nequip.train.EnergyForceLoss
    per_atom_energy: true
    coeffs:
      total_energy: 1.0
      forces: 1.0
Parameters:

  • coeffs (Dict[str, float]) – dict that stores the relative weight of energy and forces to the overall loss (default {'total_energy': 1.0, 'forces': 1.0})

  • per_atom_energy (bool, optional) – whether to normalize the total energy by the number of atoms (default True)

class nequip.train.EnergyForceStressLoss(coeffs: Dict[str, float] = {'forces': 1.0, 'stress': 1.0, 'total_energy': 1.0}, per_atom_energy: bool = True, type_names: List[str] | None = None, ignore_nan: Dict[str, bool] | None = None)[source]

Simplified MetricsManager wrapper for a loss term containing energy, forces and stress mean squared errors (MSEs).

The loss component names are per_atom_energy_mse OR total_energy_mse (depending on whether per_atom_energy is True or False), forces_mse, and stress_mse, which are the names to refer to when neeeded, e.g. when scheduling loss component coefficients.

Example usage in config:

training_module:
  _target_: nequip.train.NequIPLightningModule

  loss:
    _target_: nequip.train.EnergyForceStressLoss
    per_atom_energy: true
    coeffs:
      total_energy: 1.0
      forces: 1.0
      stress: 1.0
    # if not all frames have stresses, one can populate the stress labels with NaN and set ignore_nan here:
    # ignore_nan:
    #   stress: true
Parameters:

  • coeffs (Dict[str, float]) – dict that stores the relative weight of energy and forces to the overall loss (default {'total_energy': 1.0, 'forces': 1.0, 'stress': 1.0})

  • per_atom_energy (bool, optional) – whether to normalize the total energy by the number of atoms (default True)

  • ignore_nan (Dict[str, bool], optional) – dict that specifies whether to ignore NaN values for each field (default: all False)

class nequip.train.EnergyForceMetrics(coeffs: Dict[str, float | None] = {'forces_mae': None, 'forces_maxabserr': None, 'forces_rmse': 1.0, 'per_atom_energy_mae': None, 'per_atom_energy_maxabserr': None, 'per_atom_energy_rmse': None, 'total_energy_mae': None, 'total_energy_maxabserr': None, 'total_energy_rmse': 1.0}, type_names: List[str] | None = None)[source]

Simplified MetricsManager wrapper for a metric term containing energy and force mean absolute errors (MAEs), root mean squared errors (RMSEs), and maximum absolute errors (MaxAbsErrs).

Example usage in config:

training_module:
  _target_: nequip.train.NequIPLightningModule

  val_metrics:
    _target_: nequip.train.EnergyForceMetrics
    coeffs:
      total_energy_rmse: 1.0
      per_atom_energy_rmse: null
      forces_rmse: 1.0
      total_energy_mae: null
      per_atom_energy_mae: null
      forces_mae: null
      total_energy_maxabserr: null
      per_atom_energy_maxabserr: null
      forces_maxabserr: null
Parameters:

coeffs (Dict[str, float]) – dict that stores the relative contribution of the different energy and forces metrics to the weighted_sum version of the metric as in nequip.train.MetricsManager (default {'total_energy_rmse': 1.0, 'per_atom_energy_rmse': None, 'forces_rmse': 1.0, 'total_energy_mae': None, 'per_atom_energy_mae': None, 'forces_mae': None, 'total_energy_maxabserr': None, 'per_atom_energy_maxabserr': None, 'forces_maxabserr': None})

class nequip.train.EnergyForceStressMetrics(coeffs: Dict[str, float | None] = {'forces_mae': None, 'forces_maxabserr': None, 'forces_rmse': 1.0, 'per_atom_energy_mae': None, 'per_atom_energy_maxabserr': None, 'per_atom_energy_rmse': None, 'stress_mae': None, 'stress_maxabserr': None, 'stress_rmse': 1.0, 'total_energy_mae': None, 'total_energy_maxabserr': None, 'total_energy_rmse': 1.0}, type_names: List[str] | None = None, ignore_nan: Dict[str, bool] | None = None)[source]

Simplified MetricsManager wrapper for a metric term containing energy, force and stress mean absolute errors (MAEs), root mean squared errors (RMSEs), and maximum absolute errors (MaxAbsErrs).

Example usage in config:

training_module:
  _target_: nequip.train.NequIPLightningModule

  val_metrics:
    _target_: nequip.train.EnergyForceStressMetrics
    coeffs:
      total_energy_rmse: 1.0
      per_atom_energy_rmse: null
      forces_rmse: 1.0
      stress_rmse: 1.0
      total_energy_mae: null
      per_atom_energy_mae: null
      forces_mae: null
      stress_mae: null
      total_energy_maxabserr: null
      per_atom_energy_maxabserr: null
      forces_maxabserr: null
      stress_maxabserr: null
    # if not all frames have stresses, one can populate the stress labels with NaN and set ignore_nan here:
    # ignore_nan:
    #   stress: true
Parameters:

  • coeffs (Dict[str, float]) – dict that stores the relative contribution of the different energy and forces metrics to the weighted_sum version of the metric as in nequip.train.MetricsManager (default {'total_energy_rmse': 1.0, 'per_atom_energy_rmse': None, 'forces_rmse': 1.0, 'stress_rmse': 1.0, 'total_energy_mae': None, 'per_atom_energy_mae': None, 'forces_mae': None, 'stress_mae': None, 'total_energy_maxabserr': None, 'per_atom_energy_maxabserr': None, 'forces_maxabserr': None, 'stress_maxabserr': None})

  • ignore_nan (Dict[str, bool], optional) – dict that specifies whether to ignore NaN values for each field (default: all False)

The following can be used for energy-only datasets (without forces):

class nequip.train.EnergyOnlyLoss(per_atom_energy: bool = True, type_names: List[str] | None = None)[source]

Simplified MetricsManager wrapper for a loss term containing only energy mean squared error (MSE).

The loss component name is per_atom_energy_mse OR total_energy_mse (depending on whether per_atom_energy is True or False).

Example usage in config:

training_module:
  _target_: nequip.train.NequIPLightningModule

  loss:
    _target_: nequip.train.EnergyOnlyLoss
    per_atom_energy: true
Parameters:

per_atom_energy (bool, optional) – whether to normalize the total energy by the number of atoms (default True)

class nequip.train.EnergyOnlyMetrics(coeffs: Dict[str, float | None] = {'per_atom_energy_mae': None, 'per_atom_energy_maxabserr': None, 'per_atom_energy_rmse': None, 'total_energy_mae': None, 'total_energy_maxabserr': None, 'total_energy_rmse': 1.0}, type_names: List[str] | None = None)[source]

Simplified MetricsManager wrapper for a metric term containing only energy mean absolute errors (MAEs), root mean squared errors (RMSEs), and maximum absolute errors (MaxAbsErrs).

Example usage in config:

training_module:
  _target_: nequip.train.NequIPLightningModule

  val_metrics:
    _target_: nequip.train.EnergyOnlyMetrics
    coeffs:
      total_energy_rmse: 1.0
      per_atom_energy_rmse: null
      total_energy_mae: null
      per_atom_energy_mae: null
      total_energy_maxabserr: null
      per_atom_energy_maxabserr: null
Parameters:

coeffs (Dict[str, float]) – dict that stores the relative contribution of the different energy metrics to the weighted_sum version of the metric as in nequip.train.MetricsManager (default {'total_energy_rmse': 1.0, 'per_atom_energy_rmse': None, 'total_energy_mae': None, 'per_atom_energy_mae': None, 'total_energy_maxabserr': None, 'per_atom_energy_maxabserr': None})

Advanced Configuration: MetricsManager

For users who need custom configurations beyond the simplified wrappers, the full MetricsManager API is available.

class nequip.train.MetricsManager(metrics: List[Dict[str, float | str | Dict[str, str | Callable] | Metric]], type_names: List[str] | None = None)[source]

Manages metrics for AtomicDataDict objects in loss functions and monitored error metrics.

This class handles both loss computation and monitored error metrics through a unified interface. Each metric is defined by a dictionary with mandatory and optional keys.

The metrics parameter is a list of dictionaries defining individual metrics. Each dictionary can contain keys: field, metric, coeff, name, per_type, ignore_nan. The type_names parameter is a list of atom type names, required if ANY metric uses per_type=True. When this class is instantiated in NequIPLightningModule (which is the case when using nequip-train), type_names is automatically provided from the dataset configuration and users need not specify it explicitly.

Mandatory Keys

metrictorchmetrics.Metric

The metric class to compute (e.g., MeanSquaredError, MeanAbsoluteError, RootMeanSquaredError).

Optional Keys

fieldstr or Callable, optional

Data field to extract for metric computation. Can be:

  • str: A nequip field name (e.g., 'total_energy', 'forces', 'stress')

  • Callable: A field modifier like PerAtomModifier

  • None: For custom metrics that need full AtomicDataDict objects

When None, the metric receives (preds, target) as complete data dictionaries instead of extracted field values. The metric must handle all data processing, type filtering, and NaN handling internally.

coefffloat, optional

Weight for weighted_sum calculation. If any metric has a coefficient, weighted_sum is computed as the weighted average of all metrics with coefficients. Coefficients are automatically normalized to sum to 1.

namestr, optional

Custom name for logging. Auto-generated from field and metric if not provided. Must be unique across all metrics.

per_typebool, default=False

Compute separate metrics for each atom type, then average them. Only valid for node fields (like forces). Requires type_names parameter.

ignore_nanbool, default=False

Handle NaN values in target data by masking them out. Useful for datasets with partial labels (e.g., stress data available only for some structures).

Key Behaviors

Coefficient Normalization: Coefficients are automatically normalized to sum to 1. For example, {energy: 3.0, forces: 1.0} becomes {energy: 0.75, forces: 0.25}.

Per-Type Requirements: If ANY metric uses per_type=True, the type_names parameter is mandatory for the entire MetricsManager instance.

Per-Type Averaging: During batch steps, per-type metrics may be NaN if that atom type is absent from the batch. The effective metric averages only over types present in that batch. During epoch computation, all configured types contribute to the average.

Special Field=None Mode: When field=None, the metric receives the complete AtomicDataDict objects (preds, target) directly. This enables custom metrics that need access to multiple fields or geometric information. However, per_type and ignore_nan features are disabled—the custom metric must handle type filtering and NaN processing itself if needed.

Weighted Sum: A weighted_sum metric is automatically computed when any metric has a coefficient. This serves as the effective loss (for training) or monitoring metric (for validation). Metrics without coefficients are computed but excluded from the weighted sum.

E.g., custom MetricsManager equivalent to EnergyForceLoss:

_target_: nequip.train.MetricsManager
metrics:
    - name: per_atom_energy_mse
    field:
        _target_: nequip.data.PerAtomModifier
        field: total_energy
    coeff: 1
    metric:
        _target_: nequip.train.MeanSquaredError
    - name: forces_mse
    field: forces
    coeff: 1
    metric:
        _target_: nequip.train.MeanSquaredError
forward(preds: Dict[str, Tensor], target: Dict[str, Tensor], prefix: str = '', suffix: str = '')[source]

Computes and accumulates metrics (intended for use at batch steps).

compute(prefix: str = '', suffix: str = '')[source]

Aggregates accumulated metrics (intended for use at the end of an epoch).

set_coeffs(coeff_dict: Dict[str, float | None]) → None[source]

Sanity checks and normalizes coefficients to one before setting the new coefficients. If some metrics are unspecified, the coeff will be assumed to be None.

Error Metrics

class nequip.train.MeanSquaredError(**kwargs)[source]

Mean squared error.

class nequip.train.RootMeanSquaredError(**kwargs)[source]

Root mean squared error.

class nequip.train.MeanAbsoluteError(**kwargs)[source]

Mean absolute error.

class nequip.train.MaximumAbsoluteError(**kwargs)[source]

Maximum absolute error.

class nequip.train.HuberLoss(reduction='mean', delta=1.0, **kwargs)[source]

Huber loss (see torch.nn.HuberLoss)

Note that delta takes on the units of the target and prediction tensors.

class nequip.train.StratifiedHuberForceLoss(delta_dict, reduction='mean', **kwargs)[source]

Stratified Huber loss for vectors (forces) (see torch.nn.HuberLoss).

This metrics class implements a stratified/conditional Huber loss, where the Huber delta parameter is scaled based on the magnitude of the reference vector (i.e. force), by providing a delta_dict of {lower bound: delta parameter} where the loss contributions for all vectors with a magnitude between lower bound and the next lower bound are computed as a Huber loss with the corresponding delta parameter.

Note that delta values take on the units of the target and prediction tensors.

If the first lower bound in delta_dict is not 0 (typically recommended), then a MSELoss (divided by 2; matching Huber loss in the L2 regime (|x| < delta), see torch.nn.HuberLoss) is used for vectors with a magnitude smaller than the first lower bound.