Metrics

class metatrain.utils.metrics.RMSEAccumulator(separate_blocks: bool = False)

Bases: object

Accumulates the RMSE between predictions and targets for an arbitrary number of keys, each corresponding to one target.

Parameters:

separate_blocks (bool) – if true, the RMSE will be computed separately for each block in the target and prediction TensorMap objects.

information: Dict[str, Tuple[float, int]]

A dictionary mapping each target key to a tuple containing the sum of squared errors and the number of elements for which the error has been computed.

separate_blocks

Whether the RMSE should be computed separately for each block in the target and prediction TensorMap objects.

update(predictions: Dict[str, TensorMap], targets: Dict[str, TensorMap], extra_data: Dict[str, TensorMap] | None = None) → None

Updates the accumulator with new predictions and targets.

Parameters:

• predictions (Dict[str, TensorMap]) – A dictionary of predictions, where the keys correspond to the keys in the targets dictionary, and the values are the predictions.

• targets (Dict[str, TensorMap]) – A dictionary of targets, where the keys correspond to the keys in the predictions dictionary, and the values are the targets.

• extra_data (Dict[str, TensorMap] | None) – A dictionary of extra data, where the keys correspond to mask keys (i.e. “{target_key}_mask”), and the values are the masks to apply when computing the RMSE.

Return type:

None

finalize(not_per_atom: List[str], is_distributed: bool = False, device: device = None) → Dict[str, float]

Finalizes the accumulator and returns the RMSE for each key.

All keys will be returned as “{key} RMSE (per atom)” in the output dictionary, unless key contains one or more of the strings in not_per_atom, in which case “{key} RMSE” will be returned.

Parameters:

• not_per_atom (List[str]) – a list of strings. If any of these strings are present in a key, the RMSE key will not be labeled as “(per atom)”.

• is_distributed (bool) – if true, the RMSE will be computed across all ranks of the distributed system.

• device (device) – the local device to use for the computation. Only needed if is_distributed is True.

Returns:

The RMSE for each key.

Return type:

Dict[str, float]

class metatrain.utils.metrics.MAEAccumulator(separate_blocks: bool = False)

Bases: object

Accumulates the MAE between predictions and targets for an arbitrary number of keys, each corresponding to one target.

Parameters:

separate_blocks (bool) – if true, the RMSE will be computed separately for each block in the target and prediction TensorMap objects.

information: Dict[str, Tuple[float, int]]

A dictionary mapping each target key to a tuple containing the sum of absolute errors and the number of elements for which the error has been computed.

separate_blocks: bool

Whether the MAE should be computed separately for each block in the target and prediction TensorMap objects.

update(predictions: Dict[str, TensorMap], targets: Dict[str, TensorMap], extra_data: Dict[str, TensorMap] | None = None) → None

Updates the accumulator with new predictions and targets.

Parameters:

• predictions (Dict[str, TensorMap]) – A dictionary of predictions, where the keys correspond to the keys in the targets dictionary, and the values are the predictions.

• targets (Dict[str, TensorMap]) – A dictionary of targets, where the keys correspond to the keys in the predictions dictionary, and the values are the targets.

• extra_data (Dict[str, TensorMap] | None) – A dictionary of extra data, where the keys correspond to mask keys (i.e. “{target_key}_mask”), and the values are the masks to apply when computing the MAE.

Return type:

None

finalize(not_per_atom: List[str], is_distributed: bool = False, device: device = None) → Dict[str, float]

Finalizes the accumulator and returns the MAE for each key.

All keys will be returned as “{key} MAE (per atom)” in the output dictionary, unless key contains one or more of the strings in not_per_atom, in which case “{key} MAE” will be returned.

Parameters:

• not_per_atom (List[str]) – a list of strings. If any of these strings are present in a key, the MAE key will not be labeled as “(per atom)”.

• is_distributed (bool) – if true, the MAE will be computed across all ranks of the distributed system.

• device (device) – the local device to use for the computation. Only needed if is_distributed is True.

Returns:

The MAE for each key.

Return type:

Dict[str, float]

class metatrain.utils.metrics.EquivarianceAccumulator(separate_blocks: bool = False)

Bases: object

Accumulates the O(3) equivariance error of a model over a dataset.

This consumes the o3::variance::<name> outputs of a metatomic SymmetrizedModel, which hold, for every sample, the variance of the back-rotated predictions over a rotation grid, averaged over the components of each block. Pooling them with the same normalization as RMSEAccumulator gives an error in the same units as the accuracy RMSE of the target, and a lower bound on it.

Parameters:

separate_blocks (bool) – as in RMSEAccumulator. The blocks are then not weighted by 2 * o3_lambda + 1.

information: Dict[str, Tuple[float, int]]

A dictionary mapping each target key to a tuple containing the sum of the variances and the number of elements over which it was accumulated.

separate_blocks

Whether the equivariance error should be computed separately for each block of each target.

update(variances: Dict[str, TensorMap]) → None

Updates the accumulator with the variances of a new batch.

Parameters:

variances (Dict[str, TensorMap]) – a dictionary mapping each target name to the o3::variance::<name> output of a SymmetrizedModel, already normalized per atom like the predictions.

Return type:

None

finalize(not_per_atom: List[str], is_distributed: bool = False, device: device = None) → Dict[str, float]

Finalizes the accumulator and returns the equivariance error per key.

All keys will be returned as “{key} equivariance RMSE (per atom)” in the output dictionary, unless key contains one or more of the strings in not_per_atom, in which case “{key} equivariance RMSE” will be returned.

Parameters:

• not_per_atom (List[str]) – a list of strings. If any of these strings are present in a key, the metric will not be labeled as “(per atom)”.

• is_distributed (bool) – if true, the equivariance error will be computed across all ranks of the distributed system.

• device (device) – the local device to use for the computation. Only needed if is_distributed is True.

Returns:

The equivariance error for each key.

Return type:

Dict[str, float]

metatrain.utils.metrics.get_selected_metric(metric_dict: Dict[str, float], selected_metric: str) → float

Selects and/or calculates a (user-)selected metric from a dictionary of metrics.

This is useful when choosing the best model from a training run.

Parameters:

• metric_dict (Dict[str, float]) – A dictionary of metrics, where the keys are the names of the metrics and the values are the corresponding values.

• selected_metric (str) – The metric to return. This can be one of the following: - “loss”: return the loss value - “rmse_prod”: return the product of all RMSEs - “mae_prod”: return the product of all MAEs

Returns:

The value of the selected metric.

Return type:

float