aimnet 0.0.1__py3-none-any.whl

aimnet/models/aimnet2.py

@@ -0,0 +1,188 @@
+from typing import Dict, List, Mapping, Sequence, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+
+from aimnet import nbops, ops
+from aimnet.models.base import AIMNet2Base
+from aimnet.modules import AEVSV, MLP, ConvSV, Embedding
+
+
+# pylint: disable=too-many-arguments, too-many-instance-attributes
+class AIMNet2(AIMNet2Base):
+    def __init__(
+        self,
+        aev: Dict,
+        nfeature: int,
+        d2features: bool,
+        ncomb_v: int,
+        hidden: Tuple[List[int]],
+        aim_size: int,
+        outputs: Union[List[nn.Module], Dict[str, nn.Module]],
+        num_charge_channels: int = 1,
+    ):
+        super().__init__()
+
+        if num_charge_channels not in [1, 2]:
+            raise ValueError("num_charge_channels must be 1 (closed shell) or 2 (NSE for open-shell).")
+        self.num_charge_channels = num_charge_channels
+
+        self.aev = AEVSV(**aev)
+        nshifts_s = aev["nshifts_s"]
+        nshifts_v = aev.get("nshitfs_v") or nshifts_s
+        if d2features:
+            if nshifts_s != nshifts_v:
+                raise ValueError("nshifts_s must be equal to nshifts_v for d2features")
+            nfeature_tot = nshifts_s * nfeature
+        else:
+            nfeature_tot = nfeature
+        self.nfeature = nfeature
+        self.nshifts_s = nshifts_s
+        self.d2features = d2features
+
+        self.afv = Embedding(num_embeddings=64, embedding_dim=nfeature, padding_idx=0)
+
+        with torch.no_grad():
+            nn.init.orthogonal_(self.afv.weight[1:])
+            if d2features:
+                self.afv.weight = nn.Parameter(
+                    self.afv.weight.clone().unsqueeze(-1).expand(64, nfeature, nshifts_s).flatten(-2, -1)
+                )
+
+        conv_param = {"nshifts_s": nshifts_s, "nshifts_v": nshifts_v, "ncomb_v": ncomb_v, "do_vector": True}
+        self.conv_a = ConvSV(nchannel=nfeature, d2features=d2features, **conv_param)
+        self.conv_q = ConvSV(nchannel=num_charge_channels, d2features=False, **conv_param)
+
+        mlp_param = {"activation_fn": nn.GELU(), "last_linear": True}
+        mlps = [
+            MLP(
+                n_in=self.conv_a.output_size() + nfeature_tot,
+                n_out=nfeature_tot + 2 * num_charge_channels,
+                hidden=hidden[0],
+                **mlp_param,
+            )
+        ]
+        mlp_param = {"activation_fn": nn.GELU(), "last_linear": False}
+        for h in hidden[1:-1]:
+            mlps.append(
+                MLP(
+                    n_in=self.conv_a.output_size() + self.conv_q.output_size() + nfeature_tot + num_charge_channels,
+                    n_out=nfeature_tot + 2 * num_charge_channels,
+                    hidden=h,
+                    **mlp_param,
+                )
+            )
+        mlp_param = {"activation_fn": nn.GELU(), "last_linear": False}
+        mlps.append(
+            MLP(
+                n_in=self.conv_a.output_size() + self.conv_q.output_size() + nfeature_tot + num_charge_channels,
+                n_out=aim_size,
+                hidden=hidden[-1],
+                **mlp_param,
+            )
+        )
+        self.mlps = nn.ModuleList(mlps)
+
+        if isinstance(outputs, Sequence):
+            self.outputs = nn.ModuleList(outputs)
+        elif isinstance(outputs, Mapping):
+            self.outputs = nn.ModuleDict(outputs)
+        else:
+            raise TypeError("`outputs` is not either list or dict")
+
+    def _preprocess_spin_polarized_charge(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        if "mult" not in data:
+            raise ValueError("mult key is required for NSE if two channels for charge are not provided")
+        _half_spin = 0.5 * (data["mult"] - 1.0)
+        _half_q = 0.5 * data["charge"]
+        data["charge"] = torch.stack([_half_q + _half_spin, _half_q - _half_spin], dim=-1)
+        return data
+
+    def _postprocess_spin_polarized_charge(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        data["spin_charges"] = data["charges"][..., 0] - data["charges"][..., 1]
+        data["charges"] = data["charges"].sum(dim=-1)
+        data["charge"] = data["charge"].sum(dim=-1)
+        return data
+
+    def _prepare_in_a(self, data: Dict[str, Tensor]) -> Tensor:
+        a_i, a_j = nbops.get_ij(data["a"], data)
+        avf_a = self.conv_a(a_j, data["gs"], data["gv"])
+        if self.d2features:
+            a_i = a_i.flatten(-2, -1)
+        _in = torch.cat([a_i.squeeze(-2), avf_a], dim=-1)
+        return _in
+
+    def _prepare_in_q(self, data: Dict[str, Tensor]) -> Tensor:
+        q_i, q_j = nbops.get_ij(data["charges"], data)
+        avf_q = self.conv_q(q_j, data["gs"], data["gv"])
+        _in = torch.cat([q_i.squeeze(-2), avf_q], dim=-1)
+        return _in
+
+    def _update_q(self, data: Dict[str, Tensor], x: Tensor, delta_q: bool = True) -> Dict[str, Tensor]:
+        _q, _f, delta_a = x.split(
+            [
+                self.num_charge_channels,
+                self.num_charge_channels,
+                x.shape[-1] - 2 * self.num_charge_channels,
+            ],
+            dim=-1,
+        )
+        # for loss
+        data["_delta_Q"] = data["charge"] - nbops.mol_sum(_q, data)
+        q = data["charges"] + _q if delta_q else _q
+        f = _f.pow(2)
+        q = ops.nse(data["charge"], q, f, data, epsilon=1.0e-6)
+        data["charges"] = q
+        data["a"] = data["a"] + delta_a.view_as(data["a"])
+        return data
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        data = self.prepare_input(data)
+
+        # initial features
+        a: Tensor = self.afv(data["numbers"])
+        if self.d2features:
+            a = a.unflatten(-1, (self.nfeature, self.nshifts_s))
+        data["a"] = a
+
+        # NSE case
+        if self.num_charge_channels == 2:
+            data = self._preprocess_spin_polarized_charge(data)
+        else:
+            # make sure that charge has channel dimension
+            data["charge"] = data["charge"].unsqueeze(-1)
+
+        # AEV
+        data = self.aev(data)
+
+        # MP iterations
+        _npass = len(self.mlps)
+        for ipass, mlp in enumerate(self.mlps):
+            if ipass == 0:
+                _in = self._prepare_in_a(data)
+            else:
+                _in = torch.cat([self._prepare_in_a(data), self._prepare_in_q(data)], dim=-1)
+
+            _out = mlp(_in)
+            if data["_input_padded"].item():
+                _out = nbops.mask_i_(_out, data, mask_value=0.0)
+
+            if ipass == 0:
+                data = self._update_q(data, _out, delta_q=False)
+            elif ipass < _npass - 1:
+                data = self._update_q(data, _out, delta_q=True)
+            else:
+                data["aim"] = _out
+
+        # squeeze charges
+        if self.num_charge_channels == 2:
+            data = self._postprocess_spin_polarized_charge(data)
+        else:
+            data["charges"] = data["charges"].squeeze(-1)
+            data["charge"] = data["charge"].squeeze(-1)
+
+        # readout
+        for m in self.outputs.children():
+            data = m(data)
+
+        return data

aimnet/models/aimnet2.yaml

@@ -0,0 +1,44 @@
+class: aimnet.models.AIMNet2
+kwargs:
+    nfeature: 16
+    d2features: true
+    ncomb_v: 12
+    hidden:
+        - [512, 380]
+        - [512, 380]
+        - [512, 380, 380]
+    aim_size: 256
+    aev:
+        rc_s: 5.0
+        nshifts_s: 16
+    outputs:
+        energy_mlp:
+            class: aimnet.modules.Output
+            kwargs:
+                n_in: 256
+                n_out: 1
+                key_in: aim
+                key_out: energy
+                mlp:
+                    activation_fn: torch.nn.GELU
+                    last_linear: true
+                    hidden: [128, 128]
+
+        atomic_shift:
+            class: aimnet.modules.AtomicShift
+            kwargs:
+                key_in: energy
+                key_out: energy
+
+        atomic_sum:
+            class: aimnet.modules.AtomicSum
+            kwargs:
+                key_in: energy
+                key_out: energy
+
+        lrcoulomb:
+            class: aimnet.modules.LRCoulomb
+            kwargs:
+                rc: 4.6
+                key_in: charges
+                key_out: energy

aimnet/models/aimnet2_dftd3_wb97m.yaml

@@ -0,0 +1,51 @@
+class: aimnet.models.AIMNet2
+kwargs:
+    nfeature: 16
+    d2features: true
+    ncomb_v: 12
+    hidden:
+        - [512, 380]
+        - [512, 380]
+        - [512, 380, 380]
+    aim_size: 256
+    aev:
+        rc_s: 5.0
+        nshifts_s: 16
+    outputs:
+        energy_mlp:
+            class: aimnet.modules.Output
+            kwargs:
+                n_in: 256
+                n_out: 1
+                key_in: aim
+                key_out: energy
+                mlp:
+                    activation_fn: torch.nn.GELU
+                    last_linear: true
+                    hidden: [128, 128]
+
+        atomic_shift:
+            class: aimnet.modules.AtomicShift
+            kwargs:
+                key_in: energy
+                key_out: energy
+
+        atomic_sum:
+            class: aimnet.modules.AtomicSum
+            kwargs:
+                key_in: energy
+                key_out: energy
+
+        lrcoulomb:
+            class: aimnet.modules.LRCoulomb
+            kwargs:
+                rc: 4.6
+                key_in: charges
+                key_out: energy
+
+        dftd3:
+            class: aimnet.modules.DFTD3
+            kwargs:
+                s8: 0.3908
+                a1: 0.5660
+                a2: 3.1280

aimnet/models/base.py

@@ -0,0 +1,51 @@
+from typing import ClassVar, Dict, Final
+
+import torch
+from torch import Tensor, nn
+
+from aimnet import nbops
+
+
+class AIMNet2Base(nn.Module):  # pylint: disable=abstract-method
+    """Base class for AIMNet2 models. Implements pre-processing data:
+    converting to right dtype and device, setting nb mode, calculating masks.
+    """
+
+    __default_dtype = torch.get_default_dtype()
+
+    _required_keys: Final = ["coord", "numbers", "charge"]
+    _required_keys_dtype: Final = [__default_dtype, torch.int64, __default_dtype]
+    _optional_keys: Final = ["mult", "nbmat", "nbmat_lr", "mol_idx", "shifts", "shifts_lr", "cell"]
+    _optional_keys_dtype: Final = [
+        __default_dtype,
+        torch.int64,
+        torch.int64,
+        torch.int64,
+        __default_dtype,
+        __default_dtype,
+        __default_dtype,
+    ]
+    __constants__: ClassVar = ["_required_keys", "_required_keys_dtype", "_optional_keys", "_optional_keys_dtype"]
+
+    def __init__(self):
+        super().__init__()
+
+    def _prepare_dtype(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        for k, d in zip(self._required_keys, self._required_keys_dtype):
+            assert k in data, f"Key {k} is required"
+            data[k] = data[k].to(d)
+        for k, d in zip(self._optional_keys, self._optional_keys_dtype):
+            if k in data:
+                data[k] = data[k].to(d)
+        return data
+
+    def prepare_input(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Some sommon operations"""
+        data = self._prepare_dtype(data)
+        data = nbops.set_nb_mode(data)
+        data = nbops.calc_masks(data)
+
+        assert data["charge"].ndim == 1, "Charge should be 1D tensor."
+        if "mult" in data:
+            assert data["mult"].ndim == 1, "Mult should be 1D tensor."
+        return data

aimnet/modules/__init__.py

@@ -0,0 +1,3 @@
+from .aev import AEVSV, ConvSV  # noqa: F401
+from .core import MLP, AtomicShift, AtomicSum, Dipole, Embedding, Forces, Output, Quadrupole  # noqa: F401
+from .lr import D3TS, DFTD3, LRCoulomb  # noqa: F401

aimnet/modules/aev.py

@@ -0,0 +1,201 @@
+import math
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch import Tensor, nn
+
+from aimnet import nbops, ops
+
+
+class AEVSV(nn.Module):
+    """AEV module to expand distances and vectors toneighbors over shifted Gaussian basis functions.
+
+    Parameters:
+    -----------
+    rmin : float, optional
+        Minimum distance for the Gaussian basis functions. Default is 0.8.
+    rc_s : float, optional
+        Cutoff radius for scalar features. Default is 5.0.
+    nshifts_s : int, optional
+        Number of shifts for scalar features. Default is 16.
+    eta_s : Optional[float], optional
+        Width of the Gaussian basis functions for scalar features. Will estimate reasonable default.
+    rc_v : Optional[float], optional
+        Cutoff radius for vector features. Default is same as `rc_s`.
+    nshifts_v : Optional[int], optional
+        Number of shifts for vector features. Default is same as `nshifts_s`
+    eta_v : Optional[float], optional
+        Width of the Gaussian basis functions for vector features. Will estimate reasonable default.
+    shifts_s : Optional[List[float]], optional
+        List of shifts for scalar features. Default equidistant between `rmin` and `rc_s`
+    shifts_v : Optional[List[float]], optional
+        List of shifts for vector features. Default equidistant between `rmin` and `rc_v`
+    """
+
+    def __init__(
+        self,
+        rmin: float = 0.8,
+        rc_s: float = 5.0,
+        nshifts_s: int = 16,
+        eta_s: Optional[float] = None,
+        rc_v: Optional[float] = None,
+        nshifts_v: Optional[int] = None,
+        eta_v: Optional[float] = None,
+        shifts_s: Optional[List[float]] = None,
+        shifts_v: Optional[List[float]] = None,
+    ):
+        super().__init__()
+
+        self._init_basis(rc_s, eta_s, nshifts_s, shifts_s, rmin, mod="_s")
+        if rc_v is not None:
+            if rc_v > rc_s:
+                raise ValueError("rc_v must be less than or equal to rc_s")
+            if nshifts_v is None:
+                raise ValueError("nshifts_v must not be None")
+            self._init_basis(rc_v, eta_v, nshifts_v, shifts_v, rmin, mod="_v")
+            self._dual_basis = True
+        else:
+            # dummy init
+            self._init_basis(rc_s, eta_s, nshifts_s, shifts_s, rmin, mod="_v")
+            self._dual_basis = False
+
+        self.dmat_fill = rc_s
+
+    def _init_basis(self, rc, eta, nshifts, shifts, rmin, mod="_s"):
+        self.register_parameter(
+            "rc" + mod,
+            nn.Parameter(torch.tensor(rc, dtype=torch.float), requires_grad=False),
+        )
+        if eta is None:
+            eta = (1 / ((rc - rmin) / nshifts)) ** 2
+        self.register_parameter(
+            "eta" + mod,
+            nn.Parameter(torch.tensor(eta, dtype=torch.float), requires_grad=False),
+        )
+        if shifts is None:
+            shifts = torch.linspace(rmin, rc, nshifts + 1)[:nshifts]
+        else:
+            shifts = torch.as_tensor(shifts, dtype=torch.float)
+        self.register_parameter("shifts" + mod, nn.Parameter(shifts, requires_grad=False))
+
+    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        # shapes (..., m) and (..., m, 3)
+        d_ij, r_ij = ops.calc_distances(data)
+        data["d_ij"] = d_ij
+        # shapes (..., nshifts, m) and (..., nshifts, 3, m)
+        u_ij, gs, gv = self._calc_aev(r_ij, d_ij, data)  # pylint: disable=unused-variable
+        # for now, do not save u_ij
+        data["gs"], data["gv"] = gs, gv
+        return data
+
+    def _calc_aev(self, r_ij: Tensor, d_ij: Tensor, data: Dict[str, Tensor]) -> Tuple[Tensor, Tensor, Tensor]:
+        fc_ij = ops.cosine_cutoff(d_ij, self.rc_s)  # (..., m)
+        fc_ij = nbops.mask_ij_(fc_ij, data, 0.0)
+        gs = ops.exp_expand(d_ij, self.shifts_s, self.eta_s) * fc_ij.unsqueeze(
+            -1
+        )  # (..., m, nshifts) * (..., m, 1) -> (..., m, shitfs)
+        u_ij = r_ij / d_ij.unsqueeze(-1)  # (..., m, 3) / (..., m, 1) -> (..., m, 3)
+        if self._dual_basis:
+            fc_ij = ops.cosine_cutoff(d_ij, self.rc_v)
+            gsv = ops.exp_expand(d_ij, self.shifts_v, self.eta_v) * fc_ij.unsqueeze(-1)
+            gv = gsv.unsqueeze(-2) * u_ij.unsqueeze(-1)
+        else:
+            # (..., m, 1,  shifts), (..., m, 3, 1) -> (..., m, 3, shifts)
+            gv = gs.unsqueeze(-2) * u_ij.unsqueeze(-1)
+        return u_ij, gs, gv
+
+
+class ConvSV(nn.Module):
+    """AIMNet2 type convolution: encoding of local environment which combines geometry of local environment and atomic features.
+
+    Parameters:
+    -----------
+    nshifts_s : int
+        Number of shifts (gaussian basis functions) for scalar convolution.
+    nchannel : int
+        Number of feature channels for atomic features.
+    d2features : bool, optional
+        Flag indicating whether to use 2D features. Default is False.
+    do_vector : bool, optional
+        Flag indicating whether to perform vector convolution. Default is True.
+    nshifts_v : Optional[int], optional
+        Number of shifts for vector convolution. If not provided, defaults to the value of nshifts_s.
+    ncomb_v : Optional[int], optional
+        Number of linear combinations for vector features. If not provided, defaults to the value of nshifts_v.
+    """
+
+    def __init__(
+        self,
+        nshifts_s: int,
+        nchannel: int,
+        d2features: bool = False,
+        do_vector: bool = True,
+        nshifts_v: Optional[int] = None,
+        ncomb_v: Optional[int] = None,
+    ):
+        super().__init__()
+        nshifts_v = nshifts_v or nshifts_s
+        ncomb_v = ncomb_v or nshifts_v
+        agh = _init_ahg(nchannel, nshifts_v, ncomb_v)
+        self.register_parameter("agh", nn.Parameter(agh, requires_grad=True))
+        self.do_vector = do_vector
+        self.nchannel = nchannel
+        self.d2features = d2features
+        self.nshifts_s = nshifts_s
+        self.nshifts_v = nshifts_v
+        self.ncomb_v = ncomb_v
+
+    def output_size(self):
+        n = self.nchannel * self.nshifts_s
+        if self.do_vector:
+            n += self.nchannel * self.ncomb_v
+        return n
+
+    def forward(self, a: Tensor, gs: Tensor, gv: Optional[Tensor] = None) -> Tensor:
+        avf = []
+        if self.d2features:
+            avf_s = torch.einsum("...mag,...mg->...ag", a, gs)
+        else:
+            avf_s = torch.einsum("...mg,...ma->...ag", gs, a)
+        avf.append(avf_s.flatten(-2, -1))
+        if self.do_vector:
+            assert gv is not None
+            agh = self.agh
+            if self.d2features:
+                avf_v = torch.einsum("...mag,...mdg,agh->...ahd", a, gv, agh)
+            else:
+                avf_v = torch.einsum("...ma,...mdg,agh->...ahd", a, gv, agh)
+            avf.append(avf_v.pow(2).sum(-1).flatten(-2, -1))
+        return torch.cat(avf, dim=-1)
+
+
+def _init_ahg(b: int, m: int, n: int):
+    ret = torch.zeros(b, m, n)
+    for i in range(b):
+        ret[i] = _init_ahg_one(m, n)  # pylinit: disable-arguments-out-of-order
+    return ret
+
+
+def _init_ahg_one(m: int, n: int):
+    # make x8 times more vectors to select most diverse
+    x = torch.arange(m).unsqueeze(0)
+    a1, a2, a3, a4 = torch.randn(8 * n, 4).unsqueeze(-2).unbind(-1)
+    y = a1 * torch.sin(a2 * 2 * x * math.pi / m) + a3 * torch.cos(a4 * 2 * x * math.pi / m)
+    y -= y.mean(dim=-1, keepdim=True)
+    y /= y.std(dim=-1, keepdim=True)
+
+    dmat = torch.cdist(y, y)
+    # most distant point
+    ret = torch.zeros(n, m)
+    mask = torch.ones(y.shape[0], dtype=torch.bool)
+    i = dmat.sum(-1).argmax()
+    ret[0] = y[i]
+    mask[i] = False
+
+    # simple maxmin impementation
+    for j in range(1, n):
+        mindist, _ = torch.cdist(ret[:j], y).min(dim=0)
+        maxidx = torch.argsort(mindist)[mask][-1]
+        ret[j] = y[maxidx]
+        mask[maxidx] = False
+    return ret.t()