"""Error metrics for benchmarking, plus prediction collection.
A name -> callable registry mirrors the model registry (see
``models.registry``) so new metrics are added without touching the core, and
user-defined metrics register the same way. The built-in metrics are MAE, MSE
and RMSE; each takes two flat tensors ``(pred, target)`` and returns a Python
float.
:func:`collect_predictions` runs a model over a data loader once and returns
the paired prediction/target tensors for each requested target quantity
(``energy`` / ``forces`` / ``stress``), which :func:`score` then reduces with
the selected metrics. Energy is scored per atom -- the same size-extensive
normalization the training loss uses (see ``train.losses.weighted_loss``).
"""
from __future__ import annotations
import importlib
from typing import Callable
import torch
from torch import Tensor
from ..data import AtomicGraph
from ..models import scatter_sum
Metric = Callable[[Tensor, Tensor], float]
_METRICS: dict[str, Metric] = {}
[docs]
def register_metric(name: str) -> Callable[[Metric], Metric]:
"""Return a decorator registering an error metric under ``name``.
Parameters
----------
name : str
Name under which to register the metric. Lookups are case-insensitive
(the name is lowercased internally).
Returns
-------
Callable[[Metric], Metric]
A decorator that registers the callable it wraps and returns it
unchanged.
Raises
------
KeyError
When applied, if ``name`` is already registered to a different callable.
Examples
--------
>>> @register_metric("maxae")
... def max_abs_error(pred, target):
... return float((pred - target).abs().max())
"""
def deco(fn: Metric) -> Metric:
key = name.lower()
if key in _METRICS and _METRICS[key] is not fn:
raise KeyError(f"metric '{name}' already registered")
_METRICS[key] = fn
return fn
return deco
[docs]
def get_metric(name: str) -> Metric:
"""Look up a registered metric by name.
Parameters
----------
name : str
Metric name (case-insensitive).
Returns
-------
Metric
The registered callable.
Raises
------
KeyError
If ``name`` is not registered.
"""
key = name.lower()
if key not in _METRICS:
raise KeyError(
f"unknown metric '{name}'. registered: {available_metrics()}")
return _METRICS[key]
[docs]
def available_metrics() -> list[str]:
"""List the names of all currently registered metrics.
Returns
-------
list of str
The registered metric names, sorted alphabetically.
"""
return sorted(_METRICS)
[docs]
def load_custom_metric(name: str, path: str) -> Metric:
"""Import a ``module:function`` callable and register it under ``name``.
Lets a benchmark config pull in user-defined metrics without any code
changes to xnn: give the dotted import path of a callable with the
``(pred, target) -> float`` signature (see :data:`Metric`).
Parameters
----------
name : str
Name to register the imported callable under.
path : str
Import path of the form ``"package.module:function"``.
Returns
-------
Metric
The imported (and now registered) callable.
Raises
------
ValueError
If ``path`` does not contain the ``module:function`` separator.
"""
if ":" not in path:
raise ValueError(
f"custom metric '{name}' path must be 'module:function', got {path!r}")
module_name, func_name = path.split(":", 1)
fn = getattr(importlib.import_module(module_name), func_name)
return register_metric(name)(fn)
[docs]
@register_metric("mae")
def mae(pred: Tensor, target: Tensor) -> float:
"""Mean absolute error, ``mean(|pred - target|)``."""
return float((pred - target).abs().mean())
[docs]
@register_metric("mse")
def mse(pred: Tensor, target: Tensor) -> float:
"""Mean squared error, ``mean((pred - target) ** 2)``."""
return float(((pred - target) ** 2).mean())
[docs]
@register_metric("rmse")
def rmse(pred: Tensor, target: Tensor) -> float:
"""Root mean squared error, ``sqrt(mean((pred - target) ** 2))``."""
return float(((pred - target) ** 2).mean().sqrt())
def _scored(data: AtomicGraph, mask: Tensor | None, term: str | None = None) -> Tensor | None:
"""Bool ``(B,)`` of the structures whose label is scored, or ``None`` for all.
A structure is left out when it lacks the label (``mask``) or when its
``weight`` or term weight (``term``, e.g. ``"energy_weight"``) is zero.
"""
keep = mask
for v in (data.weight, getattr(data, term) if term else None):
if v is not None:
keep = v != 0 if keep is None else keep & (v != 0)
return keep
[docs]
def prediction_pairs(out: dict, data: AtomicGraph, targets: list[str],
atomic_energies: Tensor | None = None,
energy_per_atom: bool = True) -> dict[str, tuple[Tensor, Tensor]]:
"""The scored ``(prediction, reference)`` pairs of one batch, as flat detached tensors.
The per-batch step of :func:`collect_predictions` (see there for the
energy handling and which labels are scored); the trainer accumulates
its error metrics from the same pairs.
Parameters
----------
out : dict of str to Tensor
The model output for ``data``.
data : AtomicGraph
The batch with its labels.
targets : list of str
Quantities to pair, as in :func:`collect_predictions`.
atomic_energies : torch.Tensor or None, optional
Per-element reference energies to subtract, as in :func:`collect_predictions`.
energy_per_atom : bool, optional
Divide energies by the atom count. Defaults to ``True``.
Returns
-------
dict of str to (Tensor, Tensor)
The targets present in the batch and in ``out``, on the batch's device.
"""
pairs: dict[str, tuple[Tensor, Tensor]] = {}
if "energy" in targets and data.energy is not None and "energy" in out:
e_pred, e_ref = out["energy"], data.energy
if atomic_energies is not None:
e0 = atomic_energies.to(e_pred.device, e_pred.dtype)
e0_sum = scatter_sum(
e0[data.atomic_numbers], data.batch, data.num_graphs)
e_pred = e_pred - e0_sum
e_ref = e_ref - e0_sum
if energy_per_atom:
n = data.n_atoms.to(e_pred.dtype)
e_pred = e_pred / n
e_ref = e_ref / n
mask = _scored(data, data.energy_mask, "energy_weight")
if mask is not None:
e_pred, e_ref = e_pred[mask], e_ref[mask]
pairs["energy"] = (e_pred.detach(), e_ref.detach())
# a batch mixing labelled and unlabelled structures scores the labelled ones
if "forces" in targets and data.forces is not None and "forces" in out:
f_pred, f_ref = out["forces"].detach(), data.forces.detach()
mask = _scored(data, data.forces_mask, "force_weight")
atoms = None if mask is None else mask[data.batch]
if data.atom_force_weight is not None:
nonzero = data.atom_force_weight != 0
atoms = nonzero if atoms is None else atoms & nonzero
if atoms is not None:
f_pred, f_ref = f_pred[atoms], f_ref[atoms]
pairs["forces"] = (f_pred.reshape(-1), f_ref.reshape(-1))
if "stress" in targets and data.stress is not None and "stress" in out:
s_pred, s_ref = out["stress"].detach(), data.stress.detach()
mask = _scored(data, data.stress_mask, "stress_weight")
if mask is not None:
s_pred, s_ref = s_pred[mask], s_ref[mask]
pairs["stress"] = (s_pred.reshape(-1), s_ref.reshape(-1))
for key, label in (("dipole", "dipole"), ("polarizability", "polarizability"),
("les_dipole", "dipole")):
ref = getattr(data, label, None)
if key in targets and ref is not None and key in out:
t_pred, t_ref = out[key].detach(), ref.detach()
mask = _scored(data, getattr(data, label + "_mask", None))
if mask is not None:
t_pred, t_ref = t_pred[mask], t_ref[mask]
pairs[key] = (t_pred.reshape(-1), t_ref.reshape(-1))
for key in ("charges", "les_charges"):
if key in targets and data.charges is not None and key in out:
q_pred, q_ref = out[key].detach(), data.charges.detach()
mask = _scored(data, data.charges_mask)
if mask is not None:
atoms = mask[data.batch]
q_pred, q_ref = q_pred[atoms], q_ref[atoms]
pairs[key] = (q_pred.reshape(-1), q_ref.reshape(-1))
return pairs
[docs]
def collect_predictions(model, loader, device, targets: list[str],
atomic_energies: Tensor | None = None,
energy_per_atom: bool = True,
) -> dict[str, tuple[Tensor, Tensor]]:
"""Run ``model`` over ``loader`` once, pairing predictions with targets.
The model is put in eval mode but gradients are left enabled, because
force predictions differentiate the energy with respect to positions (the
same reason :meth:`Trainer.evaluate` avoids ``torch.no_grad``). Each
prediction/target pair is detached and moved to CPU before being stacked.
Energy handling is controlled by ``atomic_energies`` and
``energy_per_atom``. When ``atomic_energies`` is given, the per-element
reference energy of every atom is subtracted from both the predicted and
reference total energy, turning them into *atomization* (interaction)
energies -- the physically meaningful quantity to report (see
:mod:`~xnn.common.benchmark.energy`). When ``energy_per_atom`` is set, the
energy is then divided by the per-structure atom count, the size-extensive
normalization the training loss uses, so energy metrics are comparable
across differently sized structures. Both are applied identically to the
prediction and the reference, so difference metrics (MAE/MSE/RMSE) are
invariant to the E0 offset while their reported values become meaningful.
Labels with a zero loss weight are not scored: a structure whose
``weight`` or term weight (``energy_weight`` / ``force_weight`` /
``stress_weight``) is zero, and an atom whose ``atom_force_weight`` is
zero. Nonzero weights do not scale the metrics.
Parameters
----------
model : torch.nn.Module
A model whose forward returns a dict with an ``"energy"`` key and,
when forces/stress are requested, ``"forces"`` / ``"stress"`` keys
(e.g. a :class:`~xnn.common.models.ForceStressOutput`).
loader : torch.utils.data.DataLoader
Loader yielding batched :class:`AtomicGraph` objects.
device : torch.device
Device to evaluate on.
targets : list of str
Which quantities to collect; a subset of ``"energy"``, ``"forces"``,
``"stress"``, ``"dipole"``, ``"polarizability"``, ``"charges"``
(per-atom partial charges), ``"les_dipole"`` and ``"les_charges"``
(the LES outputs, scored against the dipole and charge labels). A target is
skipped for a batch that lacks the reference value or the
corresponding prediction.
atomic_energies : torch.Tensor or None, optional
A ``Z``-indexed lookup of per-element reference energies (as built by
:func:`~xnn.common.benchmark.energy.build_e0_lookup`). When given,
energies are scored on an atomization basis. Defaults to ``None`` (raw
total energy).
energy_per_atom : bool, optional
Whether to divide the (atomization) energy by the atom count before
scoring. Defaults to ``True``.
Returns
-------
dict of str to (Tensor, Tensor)
Maps each target that had at least one value to a
``(prediction, reference)`` pair of flat CPU tensors of equal length.
Targets with no data are omitted.
"""
model.eval()
preds: dict[str, list[Tensor]] = {t: [] for t in targets}
refs: dict[str, list[Tensor]] = {t: [] for t in targets}
for data in loader:
data = data.to(device)
pairs = prediction_pairs(model(data), data, targets, atomic_energies, energy_per_atom)
for t, (p, r) in pairs.items():
preds[t].append(p.cpu())
refs[t].append(r.cpu())
out_pairs: dict[str, tuple[Tensor, Tensor]] = {}
for t in targets:
if preds[t]:
out_pairs[t] = (torch.cat(preds[t]), torch.cat(refs[t]))
return out_pairs
[docs]
def score(pairs: dict[str, tuple[Tensor, Tensor]],
metrics: dict[str, list[str]] | list[str]) -> dict[str, float]:
"""Reduce prediction/target pairs to a flat ``{target_metric: value}`` map.
Parameters
----------
pairs : dict of str to (Tensor, Tensor)
Prediction/target pairs keyed by target, as returned by
:func:`collect_predictions`.
metrics : dict of str to list of str, or list of str
Which metric is reported for which target: a mapping from target to
registered metric names (the canonical
:attr:`~xnn.common.benchmark.BenchmarkConfig.metrics` form; targets
absent from the mapping are skipped), or a plain list of names applied
to every target in ``pairs``.
Returns
-------
dict of str to float
One entry per scored ``(target, metric)`` combination, keyed
``f"{target}_{metric}"`` (e.g. ``"energy_mae"``, ``"forces_rmse"``).
"""
if not isinstance(metrics, dict):
metrics = {t: list(metrics) for t in pairs}
result: dict[str, float] = {}
for target, (pred, ref) in pairs.items():
for name in metrics.get(target, []):
result[f"{target}_{name}"] = get_metric(name)(pred, ref)
return result
