aimnet 0.0.1__py3-none-any.whl

aimnet/modules/core.py

@@ -0,0 +1,237 @@
+from typing import Any, Callable, Dict, List, Optional
+
+import torch
+from torch import Tensor, nn
+
+from aimnet import constants, nbops, ops
+from aimnet.config import get_init_module, get_module
+
+
+def MLP(
+    n_in: int,
+    n_out: int,
+    hidden: Optional[List[int]] = None,
+    activation_fn: Callable | str = "torch.nn.GELU",
+    activation_kwargs: Optional[Dict[str, Any]] = None,
+    weight_init_fn: Callable | str = "torch.nn.init.xavier_normal_",
+    bias: bool = True,
+    last_linear: bool = True,
+):
+    """Convenience function to build MLP from config"""
+    if hidden is None:
+        hidden = []
+    if activation_kwargs is None:
+        activation_kwargs = {}
+    # hp search hack
+    hidden = [x for x in hidden if x > 0]
+    if isinstance(activation_fn, str):
+        activation_fn = get_init_module(activation_fn, kwargs=activation_kwargs)
+    if isinstance(weight_init_fn, str):
+        weight_init_fn = get_module(weight_init_fn)
+    sizes = [n_in, *hidden, n_out]
+    layers = []
+    for i in range(1, len(sizes)):
+        n_in, n_out = sizes[i - 1], sizes[i]
+        layer = nn.Linear(n_in, n_out, bias=bias)
+        with torch.no_grad():
+            weight_init_fn(layer.weight)
+            if bias:
+                nn.init.zeros_(layer.bias)
+        layers.append(layer)
+        if not (last_linear and i == len(sizes) - 1):
+            layers.append(activation_fn)
+    return nn.Sequential(*layers)
+
+
+class Embedding(nn.Embedding):
+    def __init__(self, init: Optional[Dict[int, Any]] = None, **kwargs):
+        super().__init__(**kwargs)
+        with torch.no_grad():
+            if init is not None:
+                for i in range(self.weight.shape[0]):
+                    if self.padding_idx is not None and i == self.padding_idx:
+                        continue
+                    if i in init:
+                        self.weight[i] = init[i]
+                    else:
+                        self.weight[i].fill_(float("nan"))
+                for k, v in init.items():
+                    self.weight[k] = v
+
+    def reset_parameters(self) -> None:
+        nn.init.orthogonal_(self.weight)
+        if self.padding_idx is not None:
+            with torch.no_grad():
+                self.weight[self.padding_idx].fill_(0)
+
+
+class DSequential(nn.Module):
+    def __init__(self, *modules):
+        super().__init__()
+        self.module = nn.ModuleList(modules)
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        for m in self.module:
+            data = m(data)
+        return data
+
+
+class AtomicShift(nn.Module):
+    def __init__(
+        self,
+        key_in: str,
+        key_out: str,
+        num_types: int = 64,
+        dtype: torch.dtype = torch.float,
+        requires_grad: bool = True,
+        reduce_sum=False,
+    ):
+        super().__init__()
+        shifts = nn.Embedding(num_types, 1, padding_idx=0, dtype=dtype)
+        shifts.weight.requires_grad_(requires_grad)
+        self.shifts = shifts
+        self.key_in = key_in
+        self.key_out = key_out
+        self.reduce_sum = reduce_sum
+
+    def extra_repr(self) -> str:
+        return f"key_in: {self.key_in}, key_out: {self.key_out}"
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        shifts = self.shifts(data["numbers"]).squeeze(-1)
+        if self.reduce_sum:
+            shifts = nbops.mol_sum(shifts, data)
+        data[self.key_out] = data[self.key_in] + shifts
+        return data
+
+
+class AtomicSum(nn.Module):
+    def __init__(self, key_in: str, key_out: str):
+        super().__init__()
+        self.key_in = key_in
+        self.key_out = key_out
+
+    def extra_repr(self) -> str:
+        return f"key_in: {self.key_in}, key_out: {self.key_out}"
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        data[self.key_out] = nbops.mol_sum(data[self.key_in], data)
+        return data
+
+
+class Output(nn.Module):
+    def __init__(self, mlp: Dict | nn.Module, n_in: int, n_out: int, key_in: str, key_out: str):
+        super().__init__()
+        self.key_in = key_in
+        self.key_out = key_out
+        if not isinstance(mlp, nn.Module):
+            mlp = MLP(n_in=n_in, n_out=n_out, **mlp)
+        self.mlp = mlp
+
+    def extra_repr(self) -> str:
+        return f"key_in: {self.key_in}, key_out: {self.key_out}"
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        v = self.mlp(data[self.key_in]).squeeze(-1)
+        if data["_input_padded"].item():
+            v = nbops.mask_i_(v, data, mask_value=0.0)
+        data[self.key_out] = v
+        return data
+
+
+class Forces(nn.Module):
+    def __init__(self, module: nn.Module, x: str = "coord", y: str = "energy", key_out: str = "forces"):
+        super().__init__()
+        self.module = module
+        self.x = x
+        self.y = y
+        self.key_out = key_out
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        prev = torch.is_grad_enabled()
+        torch.set_grad_enabled(True)
+        data[self.x].requires_grad_(True)
+        data = self.module(data)
+        y = data[self.y]
+        g = torch.autograd.grad([y.sum()], [data[self.x]], create_graph=self.training)[0]
+        assert g is not None
+        data[self.key_out] = -g
+        torch.set_grad_enabled(prev)
+        return data
+
+
+class Dipole(nn.Module):
+    def __init__(self, key_in: str = "charges", key_out: str = "dipole", center_coord: bool = False):
+        super().__init__()
+        self.center_coord = center_coord
+        self.key_out = key_out
+        self.key_in = key_in
+        self.register_buffer("mass", constants.get_masses())
+
+    def extra_repr(self) -> str:
+        return f"key_in: {self.key_in}, key_out: {self.key_out}, center_coord: {self.center_coord}"
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        q = data[self.key_in]
+        r = data["coord"]
+        if self.center_coord:
+            r = ops.center_coordinates(r, data, self.mass[data["numbers"]])
+        data[self.key_out] = nbops.mol_sum(q.unsqueeze(-1) * r, data)
+        return data
+
+
+class Quadrupole(Dipole):
+    def __init__(self, key_in: str = "charges", key_out: str = "quadrupole", center_coord: bool = False):
+        super().__init__(key_in=key_in, key_out=key_out, center_coord=center_coord)
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        q = data[self.key_in]
+        r = data["coord"]
+        if self.center_coord:
+            r = ops.center_coordinates(r, data, self.mass[data["numbers"]])
+        _x = torch.cat([r.pow(2), r * r.roll(-1, -1)], dim=-1)
+        quad = nbops.mol_sum(q.unsqueeze(-1) * _x, data)
+        _x1, _x2 = quad.split(3, dim=-1)
+        _x1 = _x1 - _x1.mean(dim=-1, keepdim=True)
+        quad = torch.cat([_x1, _x2], dim=-1)
+        data[self.key_out] = quad
+        return data
+
+
+class SRRep(nn.Module):
+    """GFN1-stype short range repulsion function"""
+
+    def __init__(self, key_out="e_rep", cutoff_fn="none", rc=5.2, reduce_sum=True):
+        super().__init__()
+        from aimnet.constants import get_gfn1_rep
+
+        self.key_out = key_out
+        self.cutoff_fn = cutoff_fn
+        self.reduce_sum = reduce_sum
+
+        self.register_buffer("rc", torch.tensor(rc))
+        gfn1_repa, gfn1_repb = get_gfn1_rep()
+        weight = torch.stack([gfn1_repa, gfn1_repb], dim=-1)
+        self.params = nn.Embedding(87, 2, padding_idx=0, _weight=weight)
+        self.params.weight.requires_grad_(False)
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        p = self.params(data["numbers"])
+        p_i, p_j = nbops.get_ij(p, data)
+        p_ij = p_i * p_j
+        alpha_ij, zeff_ij = p_ij.unbind(-1)
+        d_ij = data["d_ij"]
+        e = torch.exp(-alpha_ij * d_ij.pow(1.5)) * zeff_ij / d_ij
+        e = nbops.mask_ij_(e, data, 0.0)
+        if self.cutoff_fn == "exp_cutoff":
+            e = e * ops.exp_cutoff(d_ij, self.rc)
+        elif self.cutoff_fn == "cosine_cutoff":
+            e = e * ops.cosine_cutoff(d_ij, self.rc)
+        e = e.sum(-1)
+        if self.reduce_sum:
+            e = nbops.mol_sum(e, data)
+        if self.key_out in data:
+            data[self.key_out] = data[self.key_out] + e
+        else:
+            data[self.key_out] = e
+        return data

aimnet/modules/lr.py

@@ -0,0 +1,243 @@
+from typing import Dict, Optional
+
+import torch
+from torch import Tensor, nn
+
+from aimnet import constants, nbops, ops
+
+
+class LRCoulomb(nn.Module):
+    def __init__(
+        self,
+        key_in: str = "charges",
+        key_out: str = "e_h",
+        rc: float = 4.6,
+        method: str = "simple",
+        dsf_alpha: float = 0.2,
+        dsf_rc: float = 15.0,
+    ):
+        super().__init__()
+        self.key_in = key_in
+        self.key_out = key_out
+        self._factor = constants.half_Hartree * constants.Bohr
+        self.register_buffer("rc", torch.tensor(rc))
+        self.dsf_alpha = dsf_alpha
+        self.dsf_rc = dsf_rc
+        if method in ("simple", "dsf", "ewald"):
+            self.method = method
+        else:
+            raise ValueError(f"Unknown method {method}")
+
+    def coul_simple(self, data: Dict[str, Tensor]) -> Tensor:
+        data = ops.lazy_calc_dij_lr(data)
+        d_ij = data["d_ij_lr"]
+        q = data[self.key_in]
+        q_i, q_j = nbops.get_ij(q, data, suffix="_lr")
+        q_ij = q_i * q_j
+        fc = 1.0 - ops.exp_cutoff(d_ij, self.rc)
+        e_ij = fc * q_ij / d_ij
+        e_ij = nbops.mask_ij_(e_ij, data, 0.0, suffix="_lr")
+        e_i = e_ij.sum(-1)
+        e = self._factor * nbops.mol_sum(e_i, data)
+        return e
+
+    def coul_simple_sr(self, data: Dict[str, Tensor]) -> Tensor:
+        d_ij = data["d_ij"]
+        q = data[self.key_in]
+        q_i, q_j = nbops.get_ij(q, data)
+        q_ij = q_i * q_j
+        fc = ops.exp_cutoff(d_ij, self.rc)
+        e_ij = fc * q_ij / d_ij
+        e_ij = nbops.mask_ij_(e_ij, data, 0.0)
+        e_i = e_ij.sum(-1)
+        e = self._factor * nbops.mol_sum(e_i, data)
+        return e
+
+    def coul_dsf(self, data: Dict[str, Tensor]) -> Tensor:
+        data = ops.lazy_calc_dij_lr(data)
+        d_ij = data["d_ij_lr"]
+        q = data[self.key_in]
+        q_i, q_j = nbops.get_ij(q, data, suffix="_lr")
+        J = ops.coulomb_matrix_dsf(d_ij, self.dsf_rc, self.dsf_alpha, data)
+        e = (q_i * q_j * J).sum(-1)
+        e = self._factor * nbops.mol_sum(e, data)
+        e = e - self.coul_simple_sr(data)
+        return e
+
+    def coul_ewald(self, data: Dict[str, Tensor]) -> Tensor:
+        J = ops.coulomb_matrix_ewald(data["coord"], data["cell"])
+        q_i, q_j = data["charges"].unsqueeze(-1), data["charges"].unsqueeze(-2)
+        e = self._factor * (q_i * q_j * J).flatten(-2, -1).sum(-1)
+        e = e - self.coul_simple_sr(data)
+        return e
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        if self.method == "simple":
+            e = self.coul_simple(data)
+        elif self.method == "dsf":
+            e = self.coul_dsf(data)
+        elif self.method == "ewald":
+            e = self.coul_ewald(data)
+        else:
+            raise ValueError(f"Unknown method {self.method}")
+        if self.key_out in data:
+            data[self.key_out] = data[self.key_out] + e
+        else:
+            data[self.key_out] = e
+        return data
+
+
+class DispParam(nn.Module):
+    def __init__(
+        self,
+        ref_c6: Optional[Dict[int, Tensor] | Tensor] = None,
+        ref_alpha: Optional[Dict[int, Tensor] | Tensor] = None,
+        ptfile: Optional[str] = None,
+        key_in: str = "disp_param",
+        key_out: str = "disp_param",
+    ):
+        super().__init__()
+        if (
+            ptfile is None
+            and (ref_c6 is None or ref_alpha is None)
+            or ptfile is not None
+            and (ref_c6 is not None or ref_alpha is not None)
+        ):
+            raise ValueError("Either ptfile or ref_c6 and ref_alpha should be supplied.")
+        # load data
+        ref = torch.load(ptfile) if ptfile is not None else torch.zeros(87, 2)
+        for i, p in enumerate([ref_c6, ref_alpha]):
+            if p is not None:
+                if isinstance(p, Tensor):
+                    ref[: p.shape[0], i] = p
+                else:
+                    for k, v in p.items():
+                        ref[k, i] = v
+        # c6=0 and alpha=1 for dummy atom
+        ref[0, 0] = 0.0
+        ref[0, 1] = 1.0
+        self.register_buffer("disp_param0", ref)
+        self.key_in = key_in
+        self.key_out = key_out
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        disp_param_mult = data[self.key_in].clamp(min=-4, max=4).exp()
+        disp_param = self.disp_param0[data["numbers"]]
+        vals = disp_param * disp_param_mult
+        data[self.key_out] = vals
+        return data
+
+
+class D3TS(nn.Module):
+    """DFT-D3-like pairwise dispersion with TS combination rule"""
+
+    def __init__(self, a1: float, a2: float, s8: float, s6: float = 1.0, key_in="disp_param", key_out="energy"):
+        super().__init__()
+        self.register_buffer("r4r2", constants.get_r4r2())
+        self.a1 = a1
+        self.a2 = a2
+        self.s6 = s6
+        self.s8 = s8
+        self.key_in = key_in
+        self.key_out = key_out
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        disp_param = data[self.key_in]
+        disp_param_i, disp_param_j = nbops.get_ij(disp_param, data, suffix="_lr")
+        c6_i, alpha_i = disp_param_i.unbind(dim=-1)
+        c6_j, alpha_j = disp_param_j.unbind(dim=-1)
+
+        # TS combination rule
+        c6ij = 2 * c6_i * c6_j / (c6_i * alpha_j / alpha_i + c6_j * alpha_i / alpha_j).clamp(min=1e-4)
+
+        rr = self.r4r2[data["numbers"]]
+        rr_i, rr_j = nbops.get_ij(rr, data, suffix="_lr")
+        rrij = 3 * rr_i * rr_j
+        rrij = nbops.mask_ij_(rrij, data, 1.0, suffix="_lr")
+        r0ij = self.a1 * rrij.sqrt() + self.a2
+
+        ops.lazy_calc_dij_lr(data)
+        d_ij = data["d_ij_lr"] * constants.Bohr_inv
+        e_ij = c6ij * (self.s6 / (d_ij.pow(6) + r0ij.pow(6)) + self.s8 * rrij / (d_ij.pow(8) + r0ij.pow(8)))
+        e = -constants.half_Hartree * nbops.mol_sum(e_ij.sum(-1), data)
+
+        if self.key_out in data:
+            data[self.key_out] = data[self.key_out] + e
+        else:
+            data[self.key_out] = e
+
+        return data
+
+
+class DFTD3(nn.Module):
+    """DFT-D3 implementation.
+    BJ dumping, C6 and C8 terms, without 3-body term.
+    """
+
+    def __init__(self, s8: float, a1: float, a2: float, s6: float = 1.0, key_out="energy"):
+        super().__init__()
+        self.key_out = key_out
+        # BJ damping parameters
+        self.s6 = s6
+        self.s8 = s8
+        self.s9 = 4.0 / 3.0
+        self.a1 = a1
+        self.a2 = a2
+        self.a3 = 16.0
+        # CN parameters
+        self.k1 = -16.0
+        self.k3 = -4.0
+        # data
+        self.register_buffer("c6ab", torch.zeros(95, 95, 5, 5, 3))
+        self.register_buffer("r4r2", torch.zeros(95))
+        self.register_buffer("rcov", torch.zeros(95))
+        self.register_buffer("cnmax", torch.zeros(95))
+        sd = constants.get_dftd3_param()
+        self.load_state_dict(sd)
+
+    def _calc_c6ij(self, data: Dict[str, Tensor]) -> Tensor:
+        # CN part
+        # short range for CN
+        # d_ij = data["d_ij"] * constants.Bohr_inv
+        data = ops.lazy_calc_dij_lr(data)
+        d_ij = data["d_ij_lr"] * constants.Bohr_inv
+
+        numbers = data["numbers"]
+        numbers_i, numbers_j = nbops.get_ij(numbers, data, suffix="_lr")
+        rcov_i, rcov_j = nbops.get_ij(self.rcov[numbers], data, suffix="_lr")
+        rcov_ij = rcov_i + rcov_j
+        cn_ij = 1.0 / (1.0 + torch.exp(self.k1 * (rcov_ij / d_ij - 1.0)))
+        cn_ij = nbops.mask_ij_(cn_ij, data, 0.0, suffix="_lr")
+        cn = cn_ij.sum(-1)
+        cn = torch.clamp(cn, max=self.cnmax[numbers]).unsqueeze(-1).unsqueeze(-1)
+        cn_i, cn_j = nbops.get_ij(cn, data, suffix="_lr")
+        c6ab = self.c6ab[numbers_i, numbers_j]
+        c6ref, cnref_i, cnref_j = torch.unbind(c6ab, dim=-1)
+        c6ref = nbops.mask_ij_(c6ref, data, 0.0, suffix="_lr")
+        l_ij = torch.exp(self.k3 * ((cn_i - cnref_i).pow(2) + (cn_j - cnref_j).pow(2)))
+        w = l_ij.flatten(-2, -1).sum(-1)
+        z = torch.einsum("...ij,...ij->...", c6ref, l_ij)
+        _w = w < 1e-5
+        z[_w] = 0.0
+        c6_ij = z / w.clamp(min=1e-5)
+        return c6_ij
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        c6ij = self._calc_c6ij(data)
+
+        rr = self.r4r2[data["numbers"]]
+        rr_i, rr_j = nbops.get_ij(rr, data, suffix="_lr")
+        rrij = 3 * rr_i * rr_j
+        rrij = nbops.mask_ij_(rrij, data, 1.0, suffix="_lr")
+        r0ij = self.a1 * rrij.sqrt() + self.a2
+
+        ops.lazy_calc_dij_lr(data)
+        d_ij = data["d_ij_lr"] * constants.Bohr_inv
+        e_ij = c6ij * (self.s6 / (d_ij.pow(6) + r0ij.pow(6)) + self.s8 * rrij / (d_ij.pow(8) + r0ij.pow(8)))
+        e = -constants.half_Hartree * nbops.mol_sum(e_ij.sum(-1), data)
+
+        if self.key_out in data:
+            data[self.key_out] = data[self.key_out] + e
+        else:
+            data[self.key_out] = e
+        return data

aimnet/nbops.py

@@ -0,0 +1,151 @@
+from typing import Dict, Tuple
+
+import torch
+from torch import Tensor
+
+
+def set_nb_mode(data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+    """Logic to guess and set the neighbor model."""
+    if "nbmat" in data:
+        if data["nbmat"].ndim == 2:
+            data["_nb_mode"] = torch.tensor(1)
+        elif data["nbmat"].ndim == 3:
+            data["_nb_mode"] = torch.tensor(2)
+        else:
+            raise ValueError(f"Invalid neighbor matrix shape: {data['nbmat'].shape}")
+    else:
+        data["_nb_mode"] = torch.tensor(0)
+    return data
+
+
+def get_nb_mode(data: Dict[str, Tensor]) -> int:
+    """Get the neighbor model."""
+    return int(data["_nb_mode"].item())
+
+
+def calc_masks(data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+    """Calculate neighbor masks"""
+    nb_mode = get_nb_mode(data)
+    if nb_mode == 0:
+        data["mask_i"] = data["numbers"] == 0
+        data["mask_ij"] = torch.eye(
+            data["numbers"].shape[1], device=data["numbers"].device, dtype=torch.bool
+        ).unsqueeze(0)
+        if data["mask_i"].any():
+            data["_input_padded"] = torch.tensor(True)
+            data["_natom"] = data["mask_i"].logical_not().sum(-1)
+            data["mol_sizes"] = (~data["mask_i"]).sum(-1)
+            data["mask_ij"] = data["mask_ij"] | (data["mask_i"].unsqueeze(-2) + data["mask_i"].unsqueeze(-1))
+        else:
+            data["_input_padded"] = torch.tensor(False)
+            data["_natom"] = torch.tensor(data["numbers"].shape[1], device=data["numbers"].device)
+            data["mol_sizes"] = torch.tensor(data["numbers"].shape[1], device=data["numbers"].device)
+        data["mask_ij_lr"] = data["mask_ij"]
+    elif nb_mode == 1:
+        # padding must be the last atom
+        data["mask_i"] = torch.zeros(data["numbers"].shape[0], device=data["numbers"].device, dtype=torch.bool)
+        data["mask_i"][-1] = True
+        for suffix in ("", "_lr"):
+            if f"nbmat{suffix}" in data:
+                data[f"mask_ij{suffix}"] = data[f"nbmat{suffix}"] == data["numbers"].shape[0] - 1
+        data["_input_padded"] = torch.tensor(True)
+        data["mol_sizes"] = torch.bincount(data["mol_idx"])
+        # last atom is padding
+        data["mol_sizes"][-1] -= 1
+    elif nb_mode == 2:
+        data["mask_i"] = data["numbers"] == 0
+        w = torch.where(data["mask_i"])
+        pad_idx = w[0] * data["numbers"].shape[1] + w[1]
+        for suffix in ("", "_lr"):
+            if f"nbmat{suffix}" in data:
+                data[f"mask_ij{suffix}"] = torch.isin(data[f"nbmat{suffix}"], pad_idx)
+        data["_input_padded"] = torch.tensor(True)
+        data["mol_sizes"] = (~data["mask_i"]).sum(-1)
+    else:
+        raise ValueError(f"Invalid neighbor mode: {nb_mode}")
+
+    return data
+
+
+def mask_ij_(
+    x: Tensor,
+    data: Dict[str, Tensor],
+    mask_value: float = 0.0,
+    inplace: bool = True,
+    suffix: str = "",
+) -> Tensor:
+    mask = data[f"mask_ij{suffix}"]
+    for _i in range(x.ndim - mask.ndim):
+        mask = mask.unsqueeze(-1)
+    if inplace:
+        x.masked_fill_(mask, mask_value)
+    else:
+        x = x.masked_fill(mask, mask_value)
+    return x
+
+
+def mask_i_(x: Tensor, data: Dict[str, Tensor], mask_value: float = 0.0, inplace: bool = True) -> Tensor:
+    nb_mode = get_nb_mode(data)
+    if nb_mode == 0:
+        if data["_input_padded"].item():
+            mask = data["mask_i"]
+            for _i in range(x.ndim - mask.ndim):
+                mask = mask.unsqueeze(-1)
+            if inplace:
+                x.masked_fill_(mask, mask_value)
+            else:
+                x = x.masked_fill(mask, mask_value)
+    elif nb_mode == 1:
+        if inplace:
+            x[-1] = mask_value
+        else:
+            x = torch.cat([x[:-1], torch.zeros_like(x[:1])], dim=0)
+    elif nb_mode == 2:
+        if inplace:
+            x[:, -1] = mask_value
+        else:
+            x = torch.cat([x[:, :-1], torch.zeros_like(x[:, :1])], dim=1)
+    else:
+        raise ValueError(f"Invalid neighbor mode: {nb_mode}")
+    return x
+
+
+def get_ij(x: Tensor, data: Dict[str, Tensor], suffix: str = "") -> Tuple[Tensor, Tensor]:
+    nb_mode = get_nb_mode(data)
+    if nb_mode == 0:
+        x_i = x.unsqueeze(2)
+        x_j = x.unsqueeze(1)
+    elif nb_mode == 1:
+        x_i = x.unsqueeze(1)
+        idx = data[f"nbmat{suffix}"]
+        x_j = torch.index_select(x, 0, idx.flatten()).unflatten(0, idx.shape)
+    elif nb_mode == 2:
+        x_i = x.unsqueeze(2)
+        idx = data[f"nbmat{suffix}"]
+        x_j = torch.index_select(x.flatten(0, 1), 0, idx.flatten()).unflatten(0, idx.shape)
+    else:
+        raise ValueError(f"Invalid neighbor mode: {nb_mode}")
+    return x_i, x_j
+
+
+def mol_sum(x: Tensor, data: Dict[str, Tensor]) -> Tensor:
+    nb_mode = get_nb_mode(data)
+    if nb_mode in (0, 2):
+        res = x.sum(dim=1)
+    elif nb_mode == 1:
+        assert x.ndim in (
+            1,
+            2,
+        ), "Invalid tensor shape for mol_sum, ndim should be 1 or 2"
+        idx = data["mol_idx"]
+        # assuming mol_idx is sorted, replace with max if not
+        out_size = int(idx[-1].item()) + 1
+        if x.ndim == 1:
+            res = torch.zeros(out_size, device=x.device, dtype=x.dtype)
+        else:
+            idx = idx.unsqueeze(-1).expand(-1, x.shape[1])
+            res = torch.zeros(out_size, x.shape[1], device=x.device, dtype=x.dtype)
+        res.scatter_add_(0, idx, x)
+    else:
+        raise ValueError(f"Invalid neighbor mode: {nb_mode}")
+    return res