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.

Initialize the accumulator.

update(predictions: Dict[str, TensorMap], targets: Dict[str, TensorMap])

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.

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.

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.

Initialize the accumulator.

update(predictions: Dict[str, TensorMap], targets: Dict[str, TensorMap])

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.

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.

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

Return type:

float