aimnet 0.0.1__py3-none-any.whl

aimnet/ops.py

@@ -0,0 +1,208 @@
+import math
+from typing import Dict, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from aimnet import nbops
+
+
+def lazy_calc_dij_lr(data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+    if "d_ij_lr" not in data:
+        nb_mode = nbops.get_nb_mode(data)
+        if nb_mode == 0:
+            data["d_ij_lr"] = data["d_ij"]
+        else:
+            data["d_ij_lr"] = calc_distances(data, suffix="_lr")[0]
+    return data
+
+
+def calc_distances(data: Dict[str, Tensor], suffix: str = "", pad_value: float = 1.0) -> Tuple[Tensor, Tensor]:
+    coord_i, coord_j = nbops.get_ij(data["coord"], data, suffix)
+    if f"shifts{suffix}" in data:
+        assert "cell" in data, "cell is required if shifts are provided"
+        nb_mode = nbops.get_nb_mode(data)
+        if nb_mode == 2:
+            shifts = torch.einsum("bnmd,bdh->bnmh", data[f"shifts{suffix}"], data["cell"])
+        else:
+            shifts = data[f"shifts{suffix}"] @ data["cell"]
+        coord_j = coord_j + shifts
+    r_ij = coord_j - coord_i
+    d_ij = torch.norm(r_ij, p=2, dim=-1)
+    d_ij = nbops.mask_ij_(d_ij, data, mask_value=pad_value, inplace=False, suffix=suffix)
+    return d_ij, r_ij
+
+
+def center_coordinates(coord: Tensor, data: Dict[str, Tensor], masses: Optional[Tensor] = None) -> Tensor:
+    if masses is not None:
+        masses = masses.unsqueeze(-1)
+        center = nbops.mol_sum(coord * masses, data) / nbops.mol_sum(masses, data) / data["mol_sizes"].unsqueeze(-1)
+    else:
+        center = nbops.mol_sum(coord, data) / data["mol_sizes"]
+    nb_mode = nbops.get_nb_mode(data)
+    if nb_mode in (0, 2):
+        center = center.unsqueeze(-2)
+    coord = coord - center
+    return coord
+
+
+def cosine_cutoff(d_ij: Tensor, rc: float) -> Tensor:
+    fc = 0.5 * (torch.cos(d_ij.clamp(min=1e-6, max=rc) * (math.pi / rc)) + 1.0)
+    return fc
+
+
+def exp_cutoff(d: Tensor, rc: Tensor) -> Tensor:
+    fc = torch.exp(-1.0 / (1.0 - (d / rc).clamp(0, 1.0 - 1e-6).pow(2))) / 0.36787944117144233
+    return fc
+
+
+def exp_expand(d_ij: Tensor, shifts: Tensor, eta: float) -> Tensor:
+    # expand on axis -1, e.g. (b, n, m) -> (b, n, m, shifts)
+    return torch.exp(-eta * (d_ij.unsqueeze(-1) - shifts) ** 2)
+
+
+# pylint: disable=invalid-name
+def nse(
+    Q: Tensor,
+    q_u: Tensor,
+    f_u: Tensor,
+    data: Dict[str, Tensor],
+    epsilon: float = 1.0e-6,
+) -> Tensor:
+    # Q and q_u and f_u must have last dimension size 1 or 2
+    F_u = nbops.mol_sum(f_u, data) + epsilon
+    Q_u = nbops.mol_sum(q_u, data)
+    dQ = Q - Q_u
+    # for loss
+    data["_dQ"] = dQ
+
+    nb_mode = nbops.get_nb_mode(data)
+    if nb_mode in (0, 2):
+        F_u = F_u.unsqueeze(-2)
+        dQ = dQ.unsqueeze(-2)
+    elif nb_mode == 1:
+        data["mol_sizes"][-1] += 1
+        F_u = torch.repeat_interleave(F_u, data["mol_sizes"], dim=0)
+        dQ = torch.repeat_interleave(dQ, data["mol_sizes"], dim=0)
+        data["mol_sizes"][-1] -= 1
+    else:
+        raise ValueError(f"Invalid neighbor mode: {nb_mode}")
+    f = f_u / F_u
+    q = q_u + f * dQ
+    return q
+
+
+def coulomb_matrix_dsf(d_ij: Tensor, Rc: float, alpha: float, data: Dict[str, Tensor]) -> Tensor:
+    _c1 = (alpha * d_ij).erfc() / d_ij
+    _c2 = math.erfc(alpha * Rc) / Rc
+    _c3 = _c2 / Rc
+    _c4 = 2 * alpha * math.exp(-((alpha * Rc) ** 2)) / (Rc * math.pi**0.5)
+    J = _c1 - _c2 + (d_ij - Rc) * (_c3 + _c4)
+    # mask for d_ij > Rc
+    mask = data["mask_ij_lr"] & (d_ij > Rc)
+    J.masked_fill_(mask, 0.0)
+    return J
+
+
+def coulomb_matrix_sf(q_j: Tensor, d_ij: Tensor, Rc: float, data: Dict[str, Tensor]) -> Tensor:
+    _c1 = 1.0 / d_ij
+    _c2 = 1.0 / Rc
+    _c3 = _c2 / Rc
+    J = _c1 - _c2 + (d_ij - Rc) * _c3
+    mask = data["mask_ij_lr"] & (d_ij > Rc)
+    J.masked_fill_(mask, 0.0)
+    return J
+
+
+def get_shifts_within_cutoff(cell: Tensor, cutoff: Tensor) -> Tensor:
+    assert cell.shape == (3, 3), "Batch cell is not supported"
+    cell_inv = torch.inverse(cell).mT
+    inv_distances = cell_inv.norm(p=2, dim=-1)
+    num_repeats = torch.ceil(cutoff * inv_distances).to(torch.long)
+    device = cell.device
+    shifts = torch.cartesian_prod(
+        torch.arange(-num_repeats[0], num_repeats[0] + 1, device=device),  # type: ignore[attr-defined]
+        torch.arange(-num_repeats[1], num_repeats[1] + 1, device=device),  # type: ignore[attr-defined]
+        torch.arange(-num_repeats[2], num_repeats[2] + 1, device=device),  # type: ignore[attr-defined]
+    ).to(torch.float)
+    return shifts
+
+
+def coulomb_matrix_ewald(coord: Tensor, cell: Tensor) -> Tensor:
+    # single molecule implementation. nb_mode == 1
+    assert coord.ndim == 2 and cell.ndim == 2, "Only single molecule is supported"
+    accuracy = 1e-8
+    N = coord.shape[0]
+    volume = torch.det(cell)
+    eta = ((volume**2 / N) ** (1 / 6)) / math.sqrt(2.0 * math.pi)
+    cutoff_real = math.sqrt(-2.0 * math.log(accuracy)) * eta
+    cutoff_recip = math.sqrt(-2.0 * math.log(accuracy)) / eta
+
+    # real space
+    _grad_mode = torch.is_grad_enabled()
+    torch.set_grad_enabled(False)
+    shifts = get_shifts_within_cutoff(cell, cutoff_real)  # (num_shifts, 3)
+    torch.set_grad_enabled(_grad_mode)
+    disps_ij = coord[None, :, :] - coord[:, None, :]
+    disps = disps_ij[None, :, :, :] + torch.matmul(shifts, cell)[:, None, None, :]
+    distances_all = disps.norm(p=2, dim=-1)  # (num_shifts, num_atoms, num_atoms)
+    within_cutoff = (distances_all > 0.1) & (distances_all < cutoff_real)
+    distances = distances_all[within_cutoff]
+    e_real_matrix_aug = torch.zeros_like(distances_all)
+    e_real_matrix_aug[within_cutoff] = torch.erfc(distances / (math.sqrt(2) * eta)) / distances
+    e_real_matrix = e_real_matrix_aug.sum(dim=0)
+
+    # reciprocal space
+    recip = 2 * math.pi * torch.transpose(torch.linalg.inv(cell), 0, 1)
+    _grad_mode = torch.is_grad_enabled()
+    torch.set_grad_enabled(False)
+    shifts = get_shifts_within_cutoff(recip, cutoff_recip)
+    torch.set_grad_enabled(_grad_mode)
+    ks_all = torch.matmul(shifts, recip)
+    length_all = ks_all.norm(p=2, dim=-1)
+    within_cutoff = (length_all > 0.1) & (length_all < cutoff_recip)
+    ks = ks_all[within_cutoff]
+    length = length_all[within_cutoff]
+    # disps_ij[i, j, :] is displacement vector r_{ij}, (num_atoms, num_atoms, 3)
+    # disps_ij = coord[None, :, :] - coord[:, None, :] # computed above
+    phases = torch.sum(ks[:, None, None, :] * disps_ij[None, :, :, :], dim=-1)
+    e_recip_matrix_aug = (
+        torch.cos(phases)
+        * torch.exp(-0.5 * torch.square(eta * length[:, None, None]))
+        / torch.square(length[:, None, None])
+    )
+    e_recip_matrix = 4.0 * math.pi / volume * torch.sum(e_recip_matrix_aug, dim=0)
+    # self interaction
+    device = coord.device
+    diag = -math.sqrt(2.0 / math.pi) / eta * torch.ones(N, device=device)
+    e_self_matrix = torch.diag(diag)
+
+    J = e_real_matrix + e_recip_matrix + e_self_matrix
+    return J
+
+
+def huber(x: Tensor, delta: float = 1.0) -> Tensor:
+    return torch.where(x.abs() < delta, 0.5 * x**2, delta * (x.abs() - 0.5 * delta))
+
+
+def bumpfn(x: Tensor, low: float = 0.0, high: float = 1.0) -> Tensor:
+    """For x > 0, return smooth transition function which is 0 for x <= low and 1 for x >= high"""
+    x = (x - low) / (high - low)
+    x = x.clamp(min=1e-6, max=1 - 1e-6)
+    a = (-1 / x).exp()
+    b = (-1 / (1 - x)).exp()
+    return a / (a + b)
+
+
+def smoothstep(x: Tensor, low: float = 0.0, high: float = 1.0) -> Tensor:
+    """For x > 0, return smooth transition function which is 0 for x <= low and 1 for x >= high"""
+    x = (x - low) / (high - low)
+    x = x.clamp(min=0, max=1)
+    return x.pow(3) * (x * (x * 6 - 15) + 10)
+
+
+def expstep(x: Tensor, low: float = 0.0, high: float = 1.0) -> Tensor:
+    """For x > 0, return smooth transition function which is 0 for x <= low and 1 for x >= high"""
+    x = (x - low) / (high - low)
+    x = x.clamp(min=1e-6, max=1 - 1e-6)
+    return (-1 / (1 - x.pow(2))).exp() / 0.36787944117144233

aimnet/train/calc_sae.py

@@ -0,0 +1,43 @@
+import logging
+
+import click
+import numpy as np
+
+from aimnet.data import SizeGroupedDataset
+
+
+@click.command(short_help="Calculate SAE for a dataset.")
+@click.option("--samples", type=int, default=100000, help="Max number of samples to consider.")
+@click.argument("ds", type=str)
+@click.argument("output", type=str)
+def calc_sae(ds, output, samples=100000):
+    """Script to calculate energy SAE for a dataset DS. Writes SAE to OUTPUT file."""
+    logging.info(f"Loading dataset from {ds}")
+    ds = SizeGroupedDataset(ds, keys=["numbers", "energy"])
+    logging.info(f"Loaded dataset with {len(ds)} samples")
+    if samples > 0 and len(ds) > samples:
+        ds = ds.random_split(samples / len(ds))[0]
+    logging.info(f"Using {len(ds)} samples to calculate SAE")
+    sae = ds.apply_peratom_shift("energy", "_energy")
+    # remove up 2 percentiles from right and left
+    energy = ds.concatenate("_energy")
+    pct1, pct2 = np.percentile(energy, [2, 98])
+    for _n, g in ds.items():
+        mask = (g["_energy"] > pct1) & (g["_energy"] < pct2)
+        g = g.sample(mask)
+        if not len(g):
+            ds._data.pop(_n)
+        else:
+            ds[_n] = g
+    # now re-compute SAE
+    sae = ds.apply_peratom_shift("energy", "_energy")
+    # save
+    with open(output, "w") as f:
+        for k, v in sae.items():
+            s = f"{k}: {v}\n"
+            f.write(s)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    calc_sae()

aimnet/train/default_train.yaml

@@ -0,0 +1,166 @@
+# required to be set
+run_name: ???
+
+data:
+    # path to train and validation datasets
+    train: ???
+    val: null # same data file by default
+    # sae files
+    sae:
+        energy:
+            file: "???"
+            mode: linreg
+
+    # fraction of training data to use for validation is separate val file is not provided
+    val_fraction: 0.1
+    separate_val: true
+    # in DDP mode, will load only a shard of the dataset on each worker
+    ddp_load_full_dataset: false
+
+    # data keys
+    x: [coord, numbers, charge]
+    y: [energy, forces, charges]
+
+    # dataset class definition
+    datasets:
+        train:
+            class: aimnet.data.SizeGroupedDataset
+            kwargs: {}
+        val:
+            class: aimnet.data.SizeGroupedDataset
+            kwargs: {}
+
+    # sampler class definition
+    samplers:
+        train:
+            class: aimnet.data.SizeGroupedSampler
+            kwargs:
+                # this value if for each of DDP worker, total batch size is `batch_size*world_size`
+                batch_size: 512
+                # could be set to 'atoms', than 'batch_size' could be around 16384
+                batch_mode: molecules
+                shuffle: True
+                # for extra large datasets we want to do evaluation more often then onece per full epoch
+                # this value sets the size the epoch. `batches_per_epoch*batch_size` could be smaller or larger than the dataset size
+                batches_per_epoch: 10000
+        val:
+            class: aimnet.data.SizeGroupedSampler
+            kwargs:
+                batch_size: 1024
+                batch_mode: molecules
+                shuffle: False
+                # full dataset
+                batches_per_epoch: -1
+
+    # any additional torch.util.data.DataLoader options
+    # num_workers=0 and pin_memory=True are recommended
+    loaders:
+        train:
+            num_workers: 0
+            pin_memory: true
+        val:
+            num_workers: 0
+            pin_memory: true
+
+# definition for loss function class. Modify if trainig on different targets
+loss:
+    class: aimnet.train.loss.MTLoss
+    kwargs:
+        components:
+            energy:
+                fn: aimnet.train.loss.energy_loss_fn
+                weight: 1.0
+            forces:
+                fn: aimnet.train.loss.peratom_loss_fn
+                weight: 0.2
+                kwargs:
+                    key_true: forces
+                    key_pred: forces
+            charges:
+                fn: aimnet.train.loss.peratom_loss_fn
+                weight: 0.05
+                kwargs:
+                    key_true: charges
+                    key_pred: charges
+
+optimizer:
+    # lists of regular expressions for parameter names to enable or disable gradients
+    # force_no_grad will be processed first
+    force_no_train: []
+    force_train: []
+    class: torch.optim.RAdam
+    kwargs:
+        lr: 0.0004
+        weight_decay: 1e-8
+    # parameters with non-default optimizer settings
+    param_groups:
+        shifts:
+            re: ".*.atomic_shift.shifts.weight$"
+            weight_decay: 0.0
+
+# class definition for learning rate scheduler
+scheduler:
+    class: ignite.handlers.param_scheduler.ReduceLROnPlateauScheduler
+    kwargs:
+        metric_name: loss
+        factor: 0.75
+        patience: 10
+    # terminate training if learning rate is lower than this value
+    # useful for ReduceLROnPlateauScheduler
+    terminate_on_low_lr: 1.0e-5
+
+trainer:
+    # function that define trainig and validation loops
+    trainer: aimnet.train.utils.default_trainer
+    evaluator: aimnet.train.utils.default_evaluator
+    # total number of epochs to train
+    epochs: 100
+
+# perdicaly save chechpoints, set to null to disable
+checkpoint:
+    dirname: checkpoints
+    filename_prefix: ${run_name}
+    kwargs:
+        n_saved: 1
+        require_empty: False
+
+# wandb logger
+wandb:
+    init:
+        name: ${run_name}
+        mode: offline
+        entity: null
+        project: null
+        notes: null
+    watch_model:
+        log: all
+        log_freq: 1000
+        log_graph: true
+
+# standard set of metrics. Add an entry if training on different targets
+metrics:
+    class: aimnet.train.metrics.RegMultiMetric
+    kwargs:
+        cfg:
+            energy:
+                abbr: E
+                scale: 23.06 # eV to kcal/mol
+            dipole:
+                abbr: D
+                scale: 1.0
+                mult: 3
+            quadrupole:
+                abbr: Q
+                scale: 1.0
+                mult: 6
+            charges:
+                abbr: q
+                peratom: True
+            volumes:
+                abbr: V
+                peratom: True
+            forces:
+                abbr: F
+                peratom: True
+                mult: 3
+                scale: 23.06

aimnet/train/loss.py

@@ -0,0 +1,83 @@
+from functools import partial
+from typing import Any, Dict
+
+import torch
+from torch import Tensor
+
+from aimnet.config import get_module
+
+
+class MTLoss:
+    """Multi-target loss function with fixed weights.
+
+    This class allows for the combination of multiple loss functions, each with a specified weight.
+    The weights are normalized to sum to 1. The loss functions are applied to the predictions and
+    true values, and the weighted sum of the losses is computed.
+
+    Loss functions definition must contain keys:
+        name (str): The name of the loss function.
+        fn (str): The loss function (e.g. `aimnet2.train.loss.mse_loss_fn`).
+        weight (float): The weight of the loss function.
+        kwargs (Dict): Optional, additional keyword arguments for the loss function.
+
+    Methods:
+        __call__(y_pred, y_true):
+            Computes the weighted sum of the losses from the individual loss functions.
+            Args:
+                y_pred (Dict[str, Tensor]): Predicted values.
+                y_true (Dict[str, Tensor]): True values.
+            Returns:
+                Dict[str, Tensor]: total loss under key 'loss' and values for individual components.
+    """
+
+    def __init__(self, components: Dict[str, Any]):
+        w_sum = sum(c["weight"] for c in components.values())
+        self.components = {}
+        for name, c in components.items():
+            kwargs = c.get("kwargs", {})
+            fn = partial(get_module(c["fn"]), **kwargs)
+            self.components[name] = (fn, c["weight"] / w_sum)
+
+    def __call__(self, y_pred: Dict[str, Tensor], y_true: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        loss = {}
+        for name, (fn, w) in self.components.items():
+            _l = fn(y_pred=y_pred, y_true=y_true)
+            loss[name] = _l * w
+        # special name for the total loss
+        loss["loss"] = sum(loss.values())
+        return loss
+
+
+def mse_loss_fn(y_pred: Dict[str, Tensor], y_true: Dict[str, Tensor], key_pred: str, key_true: str) -> Tensor:
+    """General MSE loss function"""
+    x = y_true[key_true]
+    y = y_pred[key_pred]
+    loss = torch.nn.functional.mse_loss(x, y)
+    return loss
+
+
+def peratom_loss_fn(y_pred: Dict[str, Tensor], y_true: Dict[str, Tensor], key_pred: str, key_true: str) -> Tensor:
+    """MSE loss function with per-atom normalization correction.
+    Suitable when some of the values are zero both in y_pred and y_true due to padding of inputs.
+    """
+    x = y_true[key_true]
+    y = y_pred[key_pred]
+
+    if y_pred["_natom"].numel() == 1:
+        loss = torch.nn.functional.mse_loss(x, y)
+    else:
+        diff2 = (x - y).pow(2).view(x.shape[0], -1)
+        dim = diff2.shape[-1]
+        loss = (diff2 * (y_pred["_natom"].unsqueeze(-1) / dim)).mean()
+    return loss
+
+
+def energy_loss_fn(
+    y_pred: Dict[str, Tensor], y_true: Dict[str, Tensor], key_pred: str = "energy", key_true: str = "energy"
+) -> Tensor:
+    """MSE loss normalized by the number of atoms."""
+    x = y_true[key_true]
+    y = y_pred[key_pred]
+    s = y_pred["_natom"].sqrt()
+    loss = ((x - y).pow(2) / s).mean() if y_pred["_natom"].numel() > 1 else torch.nn.functional.mse_loss(x, y) / s
+    return loss

aimnet/train/metrics.py

@@ -0,0 +1,188 @@
+import logging
+from collections import defaultdict
+from typing import Dict, List, Tuple
+
+import ignite.distributed as idist
+import numpy as np
+import torch
+from ignite.exceptions import NotComputableError
+from ignite.metrics import Metric
+from ignite.metrics.metric import reinit__is_reduced
+from torch import Tensor
+
+
+def regression_stats(pred: Tensor, true: Tensor) -> Dict[str, Tensor]:
+    diff = true - pred
+    diff2 = diff.pow(2)
+    mae = diff.abs().mean(-1)
+    rmse = diff2.mean(-1).sqrt()
+    true_mean = true.mean()
+    tot = (true - true_mean).pow(2).to(torch.double).sum()
+    res = diff2.to(torch.double).sum(-1)
+    r2 = 1 - (res / tot)
+    return {"mae": mae, "rmse": rmse, "r2": r2}
+
+
+def cat_flatten(y_pred: Tensor, y_true: Tensor) -> Tuple[Tensor, Tensor]:
+    if isinstance(y_true, (list, tuple)):
+        y_true = torch.cat([x.view(-1) for x in y_true])
+    if isinstance(y_pred, (list, tuple)):
+        _n = sum(x.numel() for x in y_pred)
+        _npass = _n // y_true.numel()
+        y_pred = torch.cat([x.view(_npass, -1) for x in y_pred], dim=1)
+    y_true = y_true.view(-1)
+    if y_pred.ndim > y_true.ndim:
+        if y_pred.shape[1] != y_true.shape[0]:
+            y_pred = y_pred.view(-1, y_true.shape[0])
+        _npass = y_pred.shape[0]
+        y_pred = y_pred.view(_npass, -1)
+    else:
+        y_pred = y_pred.view(-1)
+    return y_pred, y_true
+
+
+def _iqr(a):
+    a = a.view(-1)
+    k1 = 1 + round(0.25 * (a.numel() - 1))
+    k2 = 1 + round(0.75 * (a.numel() - 1))
+    v1 = a.kthvalue(k1).values.item()
+    v2 = a.kthvalue(k2).values.item()
+    return v2 - v1
+
+
+def _freedman_diaconis_bins(a, max_bins=50):
+    """Calculate number of hist bins using Freedman-Diaconis rule."""
+    # From https://stats.stackexchange.com/questions/798/
+    if a.numel() < 2:
+        return 1
+    h = 2 * _iqr(a) / (a.numel() ** (1 / 3))
+    # fall back to sqrt(a) bins if iqr is 0
+    n_bins = int(np.sqrt(a.numel())) if h == 0 else int(np.ceil((a.max().item() - a.min().item()) / h))
+    return min(n_bins, max_bins)
+
+
+def calculate_metrics(result, histogram=False, corrplot=False):
+    keys = [k[:-5] for k in result if k.endswith("_pred")]
+    for k in keys:
+        y_pred = result.pop(k + "_pred")
+        y_true = result.pop(k + "_true")
+        y_pred, y_true = cat_flatten(y_pred, y_true)
+        stats = regression_stats(y_pred, y_true)
+        npass = stats["mae"].numel()
+        if k.split(".")[-1] in ("energy", "forces"):  # noqa: SIM108
+            f = 23.06  # eV to kcal/mol
+        else:
+            f = 1.0
+        for i in range(npass):
+            for m, v in stats.items():
+                if m in ("mae", "rmse"):
+                    v[i] = v[i] * f
+                result.log(f"{k}_{m}_{i}", v[i])
+        if histogram:
+            err = y_pred - y_true
+            for i in range(npass):
+                result[f"{k}_{i}_hist"] = torch.histc(err, bins=_freedman_diaconis_bins(err))
+    return result
+
+
+class RegMultiMetric(Metric):
+    def __init__(self, cfg: List[Dict], loss_fn=None):
+        super().__init__()
+        self.cfg = cfg
+        self.loss_fn = loss_fn
+
+    def attach_loss(self, loss_fn):
+        self.loss_fn = loss_fn
+
+    @reinit__is_reduced
+    def reset(self):
+        super().reset()
+        self.data = defaultdict(lambda: defaultdict(float))
+        self.atoms = 0.0
+        self.samples = 0.0
+        self.loss = defaultdict(float)
+
+    def _update_one(self, key: str, pred: Tensor, true: Tensor) -> None:
+        e = true - pred
+        e = e.view(pred.shape[0], -1) if pred.ndim > true.ndim else e.view(-1)
+        d = self.data[key]
+        d["sum_abs_err"] += e.abs().sum(-1).to(dtype=torch.double, device="cpu")  # type: ignore[attr-defined]
+        d["sum_sq_err"] += e.pow(2).sum(-1).to(dtype=torch.double, device="cpu")  # type: ignore[attr-defined]
+        d["sum_true"] += true.sum().to(dtype=torch.double, device="cpu")  # type: ignore[attr-defined]
+        d["sum_sq_true"] += true.pow(2).sum().to(dtype=torch.double, device="cpu")  # type: ignore[attr-defined]
+
+    @reinit__is_reduced
+    def update(self, output) -> None:
+        y_pred, y_true = output
+        if y_pred is None:
+            return
+        for k in y_pred:
+            if k not in y_true:
+                continue
+            with torch.no_grad():
+                self._update_one(k, y_pred[k].detach(), y_true[k].detach())
+            b = y_true[k].shape[0]
+        self.samples += b
+
+        _n = y_pred["_natom"]
+        if _n.numel() > 1:
+            self.atoms += _n.sum().item()
+        else:
+            self.atoms += y_pred["numbers"].shape[0] * y_pred["numbers"].shape[1]
+        if self.loss_fn is not None:
+            with torch.no_grad():
+                loss_d = self.loss_fn(y_pred, y_true)
+                for k, loss in loss_d.items():
+                    if isinstance(loss, Tensor):
+                        if loss.numel() > 1:
+                            loss = loss.mean()
+                        loss = loss.item()
+                    self.loss[k] += loss * b
+
+    def compute(self) -> Dict[str, float]:
+        if self.samples == 0:
+            raise NotComputableError
+        # Use custom reduction
+        if idist.get_world_size() > 1:
+            self.atoms = idist.all_reduce(self.atoms)
+            self.samples = idist.all_reduce(self.samples)
+            for k, loss in self.loss.items():
+                self.loss[k] = idist.all_reduce(loss)  # type: ignore[attr-defined]
+            for k1, v1 in self.data.items():
+                for k2, v2 in v1.items():
+                    self.data[k1][k2] = idist.all_reduce(v2)  # type: ignore[attr-defined]
+        self._is_reduced = True
+
+        # compute
+        ret = {}
+        for k in self.data:
+            if k not in self.cfg:
+                continue
+            cfg = self.cfg[k]
+            _n = self.atoms if cfg.get("peratom", False) else self.samples
+            _n *= cfg.get("mult", 1.0)
+            name = k
+            abbr = cfg["abbr"]
+            v = self.data[name]
+            m = {}
+            m["mae"] = v["sum_abs_err"] / _n
+            m["rmse"] = (v["sum_sq_err"] / _n).sqrt()
+            m["r2"] = 1.0 - v["sum_sq_err"] / (v["sum_sq_true"] - (v["sum_true"].pow(2)) / _n)  # type: ignore[attr-defined]
+            for k, v in m.items():
+                if k in ("mae", "rmse"):
+                    v *= cfg.get("scale", 1.0)
+                v = v.tolist()
+                if isinstance(v, list):
+                    for ii, vv in enumerate(v):
+                        ret[f"{abbr}_{k}_{ii}"] = vv
+                else:
+                    ret[f"{abbr}_{k}"] = v
+        if len(self.loss):
+            for k, loss in self.loss.items():
+                if not k.endswith("loss"):
+                    k = k + "_loss"
+                ret[k] = loss / self.samples
+
+        logging.info(str(ret))
+
+        return ret