Loss

class mlip.models.loss.Loss

Very simple loss base class that defines the signature of the call function.

abstract __call__(prediction: Prediction, ref_graph: GraphsTuple, epoch: int, eval_metrics: bool = False) → tuple[float, dict[str, float]]

The call function that outputs the loss and metrics (auxiliary data).

Parameters:

  • prediction – The force field predictor’s outputs.

  • ref_graph – The reference graph holding the ground truth data.

  • epoch – The epoch number.

  • eval_metrics – Switch deciding whether to include additional evaluation metrics to the returned dictionary. Default is False.

Returns:

The loss and the auxiliary metrics dictionary.

class mlip.models.loss.WeightedEFSLoss(energy_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, forces_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, stress_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, extended_metrics: bool = False)

Loss base class for scheduled average of energy, forces and stress errors.

__init__(energy_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, forces_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, stress_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, extended_metrics: bool = False) → None

Loss averaging energy, forces and stress errors with epoch-dependent weights.

Parameters:

  • energy_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 1.

  • forces_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 25.

  • stress_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 0.

  • extended_metrics – Whether to include an extended list of metrics. Defaults to False.

__call__(prediction: Prediction, ref_graph: GraphsTuple, epoch: int, eval_metrics: bool = False) → tuple[float, dict[str, float]]

The call function that outputs the loss and metrics (auxiliary data).

The metrics returned by this class if eval_metrics=False:

  • average loss per structure

  • energy, forces, and stress weighting factors

The metrics returned by this class if eval_metrics=True:

  • average loss per structure

  • MAE and MAE per atom (for energies, forces, and stress)

  • MSE and MSE per atom (for energies, forces, and stress)

Important note 1: we provide MSE instead of RMSE, because MSE and MAE metrics allow for downstream reweighting by number of real graphs per batch to obtain the correct metrics over the whole dataset. This reweighting is necessary as not every batch has the same number of real (not dummy) graphs and is therefore done as part of the training loop. Feel free to take the square root of the final MSE metric before logging it. The default loggers provided with this library also report RMSE instead of MSE during training.

Important note 2: we use per-component errors for forces instead of computing force error vectors per atom and then computing their norm.

Parameters:

  • prediction – The force field predictor’s outputs.

  • ref_graph – The reference graph holding the ground truth data.

  • epoch – The epoch number.

  • eval_metrics – Switch deciding whether to include additional evaluation metrics to the returned dictionary. Default is False.

Returns:

The loss and the auxiliary metrics dictionary.

class mlip.models.loss.MSELoss(energy_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, forces_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, stress_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, extended_metrics: bool = False)

Mean squared-error loss for scheduled average of energy, forces and stress errors.

__init__(energy_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, forces_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, stress_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, extended_metrics: bool = False) → None

Loss averaging energy, forces and stress errors with epoch-dependent weights.

Parameters:

  • energy_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 1.

  • forces_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 25.

  • stress_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 0.

  • extended_metrics – Whether to include an extended list of metrics. Defaults to False.

__call__(prediction: Prediction, ref_graph: GraphsTuple, epoch: int, eval_metrics: bool = False) → tuple[float, dict[str, float]]

The call function that outputs the loss and metrics (auxiliary data).

The metrics returned by this class if eval_metrics=False:

  • average loss per structure

  • energy, forces, and stress weighting factors

The metrics returned by this class if eval_metrics=True:

  • average loss per structure

  • MAE and MAE per atom (for energies, forces, and stress)

  • MSE and MSE per atom (for energies, forces, and stress)

Important note 1: we provide MSE instead of RMSE, because MSE and MAE metrics allow for downstream reweighting by number of real graphs per batch to obtain the correct metrics over the whole dataset. This reweighting is necessary as not every batch has the same number of real (not dummy) graphs and is therefore done as part of the training loop. Feel free to take the square root of the final MSE metric before logging it. The default loggers provided with this library also report RMSE instead of MSE during training.

Important note 2: we use per-component errors for forces instead of computing force error vectors per atom and then computing their norm.

Parameters:

  • prediction – The force field predictor’s outputs.

  • ref_graph – The reference graph holding the ground truth data.

  • epoch – The epoch number.

  • eval_metrics – Switch deciding whether to include additional evaluation metrics to the returned dictionary. Default is False.

Returns:

The loss and the auxiliary metrics dictionary.

class mlip.models.loss.HuberLoss(energy_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, forces_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, stress_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, extended_metrics: bool = False)

Huber loss for scheduled average of energy, forces and stress errors.

__init__(energy_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, forces_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, stress_weight_schedule: ~typing.Callable[[int], float] = <function WeightedEFSLoss.<lambda>>, extended_metrics: bool = False) → None

Loss averaging energy, forces and stress errors with epoch-dependent weights.

Parameters:

  • energy_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 1.

  • forces_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 25.

  • stress_weight_schedule – The schedule function for the energy weight. Default is a constant weight of 0.

  • extended_metrics – Whether to include an extended list of metrics. Defaults to False.

__call__(prediction: Prediction, ref_graph: GraphsTuple, epoch: int, eval_metrics: bool = False) → tuple[float, dict[str, float]]

The call function that outputs the loss and metrics (auxiliary data).

The metrics returned by this class if eval_metrics=False:

  • average loss per structure

  • energy, forces, and stress weighting factors

The metrics returned by this class if eval_metrics=True:

  • average loss per structure

  • MAE and MAE per atom (for energies, forces, and stress)

  • MSE and MSE per atom (for energies, forces, and stress)

Important note 1: we provide MSE instead of RMSE, because MSE and MAE metrics allow for downstream reweighting by number of real graphs per batch to obtain the correct metrics over the whole dataset. This reweighting is necessary as not every batch has the same number of real (not dummy) graphs and is therefore done as part of the training loop. Feel free to take the square root of the final MSE metric before logging it. The default loggers provided with this library also report RMSE instead of MSE during training.

Important note 2: we use per-component errors for forces instead of computing force error vectors per atom and then computing their norm.

Parameters:

  • prediction – The force field predictor’s outputs.

  • ref_graph – The reference graph holding the ground truth data.

  • epoch – The epoch number.

  • eval_metrics – Switch deciding whether to include additional evaluation metrics to the returned dictionary. Default is False.

Returns:

The loss and the auxiliary metrics dictionary.