aimnet 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

aimnet/modules/__init__.py

@@ -1,3 +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
+from .lr import D3TS, DFTD3, LRCoulomb, SRCoulomb  # noqa: F401

aimnet/modules/aev.py

@@ -1,10 +1,10 @@
 import math
-from typing import Dict, List, Optional, Tuple
 
 import torch
 from torch import Tensor, nn
 
 from aimnet import nbops, ops
+from aimnet.kernels import conv_sv_2d_sp
 
 
 class AEVSV(nn.Module):
@@ -37,12 +37,12 @@ class AEVSV(nn.Module):
         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,
+        eta_s: float | None = None,
+        rc_v: float | None = None,
+        nshifts_v: int | None = None,
+        eta_v: float | None = None,
+        shifts_s: list[float] | None = None,
+        shifts_v: list[float] | None = None,
     ):
         super().__init__()
 
@@ -78,31 +78,31 @@ class AEVSV(nn.Module):
             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)
+    def forward(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        # d_ij: distances to neighbors (..., m)
+        # r_ij: displacement vectors to neighbors (..., 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
+        # Atomic environment vectors: (..., m, g, 4)
+        # 4 components = 1 scalar (radial) + 3 vector (directional)
+        g_sv = self._calc_aev(r_ij, d_ij, data)
+        data["g_sv"] = g_sv
         return data
 
-    def _calc_aev(self, r_ij: Tensor, d_ij: Tensor, data: Dict[str, Tensor]) -> Tuple[Tensor, Tensor, Tensor]:
+    def _calc_aev(self, r_ij: Tensor, d_ij: Tensor, data: dict[str, 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
+        # Apply cutoff envelope to Gaussian-expanded distances
+        # Shape: (..., m, nshifts) * (..., m, 1) -> (..., m, nshifts)
+        gs = ops.exp_expand(d_ij, self.shifts_s, self.eta_s) * fc_ij.unsqueeze(-1)
+        # Normalize displacement vectors to unit direction vectors
+        # Shape: (..., m, 3) / (..., m, 1) -> (..., m, 3)
+        u_ij = r_ij / d_ij.unsqueeze(-1)
+        # Combine radial basis with directional info for vector features
+        # Shape: (..., m, 1, nshifts) * (..., m, 3, 1) -> (..., m, nshifts, 3)
+        gv = gs.unsqueeze(-1) * u_ij.unsqueeze(-2)
+        g_sv = torch.cat([gs.unsqueeze(-1), gv], dim=-1)
+        return g_sv
 
 
 class ConvSV(nn.Module):
@@ -116,8 +116,6 @@ class ConvSV(nn.Module):
         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
@@ -129,16 +127,15 @@ class ConvSV(nn.Module):
         nshifts_s: int,
         nchannel: int,
         d2features: bool = False,
-        do_vector: bool = True,
-        nshifts_v: Optional[int] = None,
-        ncomb_v: Optional[int] = None,
+        nshifts_v: int | None = None,
+        ncomb_v: int | None = 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.do_vector = True
         self.nchannel = nchannel
         self.d2features = d2features
         self.nshifts_s = nshifts_s
@@ -151,33 +148,37 @@ class ConvSV(nn.Module):
             n += self.nchannel * self.ncomb_v
         return n
 
-    def forward(self, a: Tensor, gs: Tensor, gv: Optional[Tensor] = None) -> Tensor:
-        avf = []
+    def forward(self, data: dict[str, Tensor], a: Tensor) -> Tensor:
+        g_sv = data["g_sv"]
+        mode = nbops.get_nb_mode(data)
         if self.d2features:
-            avf_s = torch.einsum("...mag,...mg->...ag", a, gs)
+            if mode > 0 and a.device.type == "cuda":
+                avf_sv = conv_sv_2d_sp(a, data["nbmat"], g_sv)
+            elif mode > 0:
+                a_j = a.index_select(0, data["nbmat"].flatten()).unflatten(0, data["nbmat"].shape)
+                avf_sv = torch.einsum("...mag,...mgd->...agd", a_j, g_sv)
+            else:
+                avf_sv = torch.einsum("...mag,...mgd->...agd", a.unsqueeze(1), g_sv)
         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)
+            if mode > 0:
+                a_j = a.index_select(0, data["nbmat"].flatten()).unflatten(0, data["nbmat"].shape)
+                avf_sv = torch.einsum("...ma,...mgd->...agd", a_j, g_sv)
             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)
+                avf_sv = torch.einsum("...ma,...mgd->...agd", a.unsqueeze(1), g_sv)
+        avf_s, avf_v = avf_sv.split([1, 3], dim=-1)
+        avf_v = torch.einsum("agh,...agd->...ahd", self.agh, avf_v).pow(2).sum(-1)
+        return torch.cat([avf_s.squeeze(-1).flatten(-2, -1), avf_v.flatten(-2, -1)], 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
+        ret[i] = _init_ahg_one(m, n)  # pylint: disable=arguments-out-of-order
     return ret
 
 
 def _init_ahg_one(m: int, n: int):
-    # make x8 times more vectors to select most diverse
+    # Oversample by 8x to select most orthogonal vectors via maxmin algorithm
     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)
@@ -192,7 +193,7 @@ def _init_ahg_one(m: int, n: int):
     ret[0] = y[i]
     mask[i] = False
 
-    # simple maxmin impementation
+    # simple maxmin implementation
     for j in range(1, n):
         mindist, _ = torch.cdist(ret[:j], y).min(dim=0)
         maxidx = torch.argsort(mindist)[mask][-1]

aimnet/modules/core.py

@@ -1,4 +1,5 @@
-from typing import Any, Callable, Dict, List, Optional
+from collections.abc import Callable
+from typing import Any
 
 import torch
 from torch import Tensor, nn
@@ -10,9 +11,9 @@ from aimnet.config import get_init_module, get_module
 def MLP(
     n_in: int,
     n_out: int,
-    hidden: Optional[List[int]] = None,
+    hidden: list[int] | None = None,
     activation_fn: Callable | str = "torch.nn.GELU",
-    activation_kwargs: Optional[Dict[str, Any]] = None,
+    activation_kwargs: dict[str, Any] | None = None,
     weight_init_fn: Callable | str = "torch.nn.init.xavier_normal_",
     bias: bool = True,
     last_linear: bool = True,
@@ -44,7 +45,7 @@ def MLP(
 
 
 class Embedding(nn.Embedding):
-    def __init__(self, init: Optional[Dict[int, Any]] = None, **kwargs):
+    def __init__(self, init: dict[int, Any] | None = None, **kwargs):
         super().__init__(**kwargs)
         with torch.no_grad():
             if init is not None:
@@ -70,7 +71,7 @@ class DSequential(nn.Module):
         super().__init__()
         self.module = nn.ModuleList(modules)
 
-    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+    def forward(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
         for m in self.module:
             data = m(data)
         return data
@@ -97,7 +98,7 @@ class AtomicShift(nn.Module):
     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]:
+    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)
@@ -114,13 +115,13 @@ class AtomicSum(nn.Module):
     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]:
+    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):
+    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
@@ -131,7 +132,7 @@ class Output(nn.Module):
     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]:
+    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)
@@ -147,7 +148,7 @@ class Forces(nn.Module):
         self.y = y
         self.key_out = key_out
 
-    def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+    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)
@@ -171,7 +172,7 @@ class Dipole(nn.Module):
     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]:
+    def forward(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
         q = data[self.key_in]
         r = data["coord"]
         if self.center_coord:
@@ -184,7 +185,7 @@ 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]:
+    def forward(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
         q = data[self.key_in]
         r = data["coord"]
         if self.center_coord:
@@ -215,7 +216,7 @@ class SRRep(nn.Module):
         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]:
+    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