Torch Model

apax/config/train_config.py

@@ -171,7 +171,7 @@ class ModelConfig(BaseModel, extra="forbid"):
     n_radial: PositiveInt = 5
     r_min: NonNegativeFloat = 0.5
     r_max: PositiveFloat = 6.0
-    n_contr: int = -1
+    n_contr: int = 8
     emb_init: Optional[str] = "uniform"
 
     nn: List[PositiveInt] = [512, 512]

apax/layers/scaling.py

@@ -24,6 +24,7 @@ def setup(self):
             scale = einops.repeat(scale, "species -> species 1")
         if len(shift.shape) == 1:
             shift = einops.repeat(shift, "species -> species 1")
+
         scale_init = nn.initializers.constant(scale)
         shift_init = nn.initializers.constant(shift)
 

apax/md/torch_ase_calc.py

@@ -0,0 +1,84 @@
+from pathlib import Path
+from typing import Callable, Union
+
+import numpy as np
+import torch
+from ase.calculators.calculator import Calculator, all_changes
+from matscipy.neighbours import neighbour_list
+
+
+class TorchASECalculator(Calculator):
+    """
+    ASE Calculator for apax models.
+    """
+
+    implemented_properties = [
+        "energy",
+        "forces",
+    ]
+
+    def __init__(
+        self,
+        model_path: Union[Path, list[Path]],
+        dr_threshold: float = 0.5,
+        transformations: Callable = [],
+        **kwargs
+    ):
+        Calculator.__init__(self, **kwargs)
+        self.skin = dr_threshold
+
+        self.model = torch.jit.load(model_path)
+        self.r_max = (
+            self.model.energy_model.atomistic_model.descriptor.radial_fn.basis_fn.r_max
+        )
+
+        self.step = None
+        self.neighbor_fn = None
+        self.neighbors = None
+        self.offsets = None
+        self.pos0 = 0
+        self.Z = [0, 0]
+        self.pbc = False
+
+    def set_neighbours_and_offsets(self, atoms, box):
+        condition = (
+            np.any(self.pbc != atoms.pbc)
+            or len(self.Z) != len(atoms.numbers)
+            or np.max(np.sum(((self.pos0 - atoms.positions) ** 2), axis=1))
+            > self.skin**2 / 4.0
+        )
+        if condition:
+            idxs_i, idxs_j, offsets = neighbour_list(
+                "ijS", positions=atoms.positions, pbc=atoms.pbc, cutoff=self.r_max
+            )
+
+            self.neighbors = np.array([idxs_i, idxs_j], dtype=np.int32)
+            self.offsets = np.matmul(offsets, box)  # np.zeros_like(self.neighbors) #
+            self.pos0 = atoms.positions
+            self.Z = atoms.numbers
+            self.pbc = atoms.pbc
+
+    def calculate(self, atoms, properties=["energy"], system_changes=all_changes):
+        Calculator.calculate(self, atoms, properties, system_changes)
+        positions = atoms.positions
+        box = atoms.cell.array
+        # if np.any(atoms.pbc):
+        #     positions = atoms.positions @ np.linalg.inv(box)
+
+        # predict
+        self.set_neighbours_and_offsets(atoms, box)
+
+        inputt = (
+            torch.from_numpy(positions),
+            torch.from_numpy(atoms.numbers),
+            torch.from_numpy(np.asarray(self.neighbors, dtype=np.int64)),
+            torch.from_numpy(np.asarray(box, dtype=np.float64)),
+            torch.from_numpy(np.asarray(self.offsets, dtype=np.float64)),
+        )
+
+        results = self.model(*inputt)
+
+        self.results = {
+            k: np.array(v.detach().numpy(), dtype=np.float64) for k, v in results.items()
+        }
+        self.results["energy"] = self.results["energy"].item()

apax/model/gmnn.py

@@ -66,7 +66,6 @@ def __call__(
         gm = self.descriptor(dr_vec, Z, idx)
         h = jax.vmap(self.readout)(gm)
         output = self.scale_shift(h, Z)
-
         if self.mask_atoms:
             output = mask_by_atom(output, Z)
         return output

apax/nn/impl/triangular_indices.py

@@ -0,0 +1,49 @@
+import numpy as np
+
+# import numpy as np
+
+# def tril_2d_indices(n: int):
+#     indices = np.zeros((int(n * (n + 1) / 2), 2), dtype=int)
+#     sparse_idx = 0
+#     for i in range(0, n):
+#         for j in range(i, n):
+#             indices[sparse_idx] = i, j
+#             sparse_idx += 1
+
+#     return jnp.asarray(indices)
+
+
+# def tril_3d_indices(n: int):
+#     indices = np.zeros((int(n * (n + 1) * (n + 2) / 6), 3), dtype=int)
+#     sparse_idx = 0
+#     for i in range(0, n):
+#         for j in range(i, n):
+#             for k in range(j, n):
+#                 indices[sparse_idx] = i, j, k
+#                 sparse_idx += 1
+
+#     return jnp.asarray(indices)
+
+
+def tril_2d_indices(n_radial):
+    tril_idxs = []
+    for i in range(n_radial):
+        tril_idxs.append([i, i])
+        for j in range(i + 1, n_radial):
+            tril_idxs.append([i, j])
+    tril_idxs = np.array(tril_idxs)
+    return tril_idxs
+
+
+def tril_3d_indices(n_radial):
+    tril_idxs = []
+    for i in range(n_radial):
+        tril_idxs.append([i, i, i])
+        for j in range(n_radial):
+            if j != i:
+                tril_idxs.append([i, j, j])
+        for j in range(i + 1, n_radial):
+            for k in range(j + 1, n_radial):
+                tril_idxs.append([i, j, k])
+    tril_idxs = np.array(tril_idxs)
+    return tril_idxs

apax/nn/torch/layers/__init__.py

@@ -0,0 +1,3 @@
+from apax.nn.torch.layers import descriptor, ntk_linear, scaling
+
+__all__ = ["descriptor", "ntk_linear", "scaling"]

apax/nn/torch/layers/activation.py

@@ -0,0 +1,11 @@
+import torch
+import torch.nn as nn
+
+
+class SwishT(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x):
+        h = 1.6765324703310907 * torch.nn.functional.silu(x)
+        return h

apax/nn/torch/layers/descriptor/__init__.py

@@ -0,0 +1,5 @@
+from apax.nn.torch.layers.descriptor.gaussian_moment_descriptor import (
+    GaussianMomentDescriptorT,
+)
+
+__all__ = ["GaussianMomentDescriptorT"]

apax/nn/torch/layers/descriptor/basis.py

@@ -0,0 +1,114 @@
+from typing import Any
+
+import einops
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+def cosine_cutoff(dr, r_max):
+    # shape: neighbors
+    dr_clipped = torch.clamp(dr, max=r_max)
+    cos_cutoff = 0.5 * (torch.cos(np.pi * dr_clipped / r_max) + 1.0)
+    cutoff = cos_cutoff[:, None]
+    return cutoff
+
+
+class GaussianBasisT(nn.Module):
+    def __init__(
+        self,
+        n_basis: int = 7,
+        r_min: float = 0.5,
+        r_max: float = 6.0,
+        dtype: Any = torch.float32,
+    ) -> None:
+        super().__init__()
+        self.n_basis = n_basis
+        self.r_min = r_min
+        self.r_max = r_max
+        self.dtype = dtype
+
+        self.betta = self.n_basis**2 / self.r_max**2
+        self.rad_norm = (2.0 * self.betta / np.pi) ** 0.25
+        shifts = self.r_min + (self.r_max - self.r_min) / self.n_basis * np.arange(
+            self.n_basis
+        )
+
+        self.betta = torch.tensor(self.betta)
+        self.rad_norm = torch.tensor(self.rad_norm)
+
+        # shape: 1 x n_basis
+        shifts = einops.repeat(shifts, "n_basis -> 1 n_basis")
+        self.shifts = torch.tensor(shifts, dtype=self.dtype)
+
+    def forward(self, dr: torch.Tensor) -> torch.Tensor:
+        # dr shape: neighbors
+        # neighbors -> neighbors x 1
+        dr = dr[:, None].type(self.dtype)
+        # 1 x n_basis, neighbors x 1 -> neighbors x n_basis
+        distances = self.shifts - dr
+
+        # shape: neighbors x n_basis
+        basis = torch.exp(-self.betta * (distances**2))
+        basis = self.rad_norm * basis
+        return basis
+
+
+class RadialFunctionT(nn.Module):
+    def __init__(
+        self,
+        n_radial: int = 5,
+        basis_fn: nn.Module = GaussianBasisT(),
+        emb_init: str = "uniform",
+        n_species: int = 119,
+        params=None,
+        dtype: Any = torch.float32,
+    ) -> None:
+        super().__init__()
+        self.n_radial = n_radial
+        self.basis_fn = basis_fn
+        self.n_species = n_species
+        self.emb_init = emb_init
+        self.dtype = dtype
+
+        self.r_max = torch.tensor(self.basis_fn.r_max)
+        norm = 1.0 / np.sqrt(self.basis_fn.n_basis)
+        self.embed_norm = torch.tensor(norm, dtype=self.dtype)
+        self.embeddings = None
+
+        if params:
+            emb = params["atomic_type_embedding"]
+            emb = torch.from_numpy(np.array(emb))
+            self.embeddings = nn.Parameter(emb)
+        elif self.emb_init is not None:
+            self.n_radial = n_radial
+            emb = torch.rand(
+                (self.n_species, self.n_species, self.n_radial, self.basis_fn.n_basis)
+            )
+            self.embeddings = nn.Parameter(emb)
+        else:
+            self.n_radial = self.basis_fn.n_basis
+
+    def forward(self, dr, Z_i, Z_j):
+        dr = dr.type(self.dtype)
+        # basis shape: neighbors x n_basis
+        basis = self.basis_fn(dr)
+
+        if self.embeddings is None:
+            radial_function = basis
+        else:
+            # coeffs shape: n_neighbors x n_radialx n_basis
+            # reverse convention to match original
+            species_pair_coeffs = self.embeddings[Z_j, Z_i, ...]
+            species_pair_coeffs = self.embed_norm * species_pair_coeffs
+
+            radial_function = torch.einsum(
+                "nrb, nb -> nr",
+                species_pair_coeffs,
+                basis,
+            )
+        cutoff = cosine_cutoff(dr, self.r_max)
+        radial_function = radial_function * cutoff
+
+        assert radial_function.dtype == self.dtype
+        return radial_function