Loss

class mlip.models.loss.LossTerm

Abstract loss base class that defines the signature of the call function.

property_name

The name of the property that is being penalized in this loss. Needs to be overridden by subclasses. Mostly meant for use in metrics logging.

Type:

str | None

abstractmethod __call__(pred_graph: Graph, ref_graph: Graph) → Array

The call function that outputs the loss.

Parameters:

  • pred_graph – Graph holding the predicted data.

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

Returns:

The loss.

class mlip.models.loss.Loss(losses: list[LossTerm], schedules: list[Callable[[int], float]], extended_metrics: bool = False)

Loss combining multiple loss terms and schedules to form a new weighted loss.

__init__(losses: list[LossTerm], schedules: list[Callable[[int], float]], extended_metrics: bool = False)

Constructor.

Parameters:

  • losses – A list of loss terms.

  • schedules – A list of schedules. Has to be same order as losses.

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

__call__(pred_graph: Graph, ref_graph: Graph, epoch: int, eval_metrics: bool = False) → tuple[Array, dict[str, float | Array]]

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

Parameters:

  • pred_graph – Graph holding the predicted data.

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

  • epoch – The epoch number passed to the schedules to get the correct weights for each sub-loss.

  • 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 MSELoss.<lambda>>, forces_weight_schedule: ~typing.Callable[[int], float] = <function MSELoss.<lambda>>, stress_weight_schedule: ~typing.Callable[[int], float] = <function MSELoss.<lambda>>, hessian_weight_schedule: ~typing.Callable[[int], float] = <function MSELoss.<lambda>>, partial_charges_weight_schedule: ~typing.Callable[[int], float] = <function MSELoss.<lambda>>, charge_weight_schedule: ~typing.Callable[[int], float] = <function MSELoss.<lambda>>, dipole_moment_weight_schedule: ~typing.Callable[[int], float] = <function MSELoss.<lambda>>, extended_metrics: bool = False)

Mean-squared error loss for energy, forces and stress.

Weights are epoch-dependent. Custom schedules can be passed.

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

Constructor.

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.

  • hessian_weight_schedule – The schedule function for the Hessian weight. Default is a constant weight of 0.

  • partial_charges_weight_schedule – The schedule function for the partial charges weight. Default is a constant weight of 0.

  • charge_weight_schedule – The schedule function for the charge weight. Default is a constant weight of 0.

  • dipole_moment_weight_schedule – The schedule function for the dipole moment weight. Default is a constant weight of 0.

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

__call__(pred_graph: Graph, ref_graph: Graph, epoch: int, eval_metrics: bool = False) → tuple[Array, dict[str, float | Array]]

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

Parameters:

  • pred_graph – Graph holding the predicted data.

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

  • epoch – The epoch number passed to the schedules to get the correct weights for each sub-loss.

  • 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 HuberLoss.<lambda>>, forces_weight_schedule: ~typing.Callable[[int], float] = <function HuberLoss.<lambda>>, stress_weight_schedule: ~typing.Callable[[int], float] = <function HuberLoss.<lambda>>, hessian_weight_schedule: ~typing.Callable[[int], float] = <function HuberLoss.<lambda>>, partial_charges_weight_schedule: ~typing.Callable[[int], float] = <function HuberLoss.<lambda>>, charge_weight_schedule: ~typing.Callable[[int], float] = <function HuberLoss.<lambda>>, dipole_moment_weight_schedule: ~typing.Callable[[int], float] = <function HuberLoss.<lambda>>, extended_metrics: bool = False)

Huber loss for energy, forces and stress.

Weights are epoch-dependent. Custom schedules can be passed.

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

Constructor.

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 forces weight. Default is a constant weight of 25.

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

  • hessian_weight_schedule – The schedule function for the Hessian weight. Default is a constant weight of 0.

  • partial_charges_weight_schedule – The schedule function for the partial charges weight. Default is a constant weight of 0.

  • charge_weight_schedule – The schedule function for the charge weight. Default is a constant weight of 0.

  • dipole_moment_weight_schedule – The schedule function for the dipole moment weight. Default is a constant weight of 0.

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

__call__(pred_graph: Graph, ref_graph: Graph, epoch: int, eval_metrics: bool = False) → tuple[Array, dict[str, float | Array]]

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

Parameters:

  • pred_graph – Graph holding the predicted data.

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

  • epoch – The epoch number passed to the schedules to get the correct weights for each sub-loss.

  • 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.

mlip.models.loss.eval_metrics.compute_eval_metrics(pred_graph: Graph, ref_graph: Graph, extended_metrics: bool = False) → dict[str, float | Array]

Computes the evaluation metrics for a given prediction and reference graph.

Parameters:

  • pred_graph – The graph containing the predictions.

  • ref_graph – The reference graph with the ground truth values.

  • extended_metrics – Whether to compute extended metrics. Default is False. Without extended metrics means that “per-atom” metrics and stress metrics are omitted. The smaller set just includes MSE and MAE metrics for energies and forces.

Returns:

The metrics as a dictionary. These metrics are averaged over the given batch.