Feat/nls share

.github/CODEOWNERS

@@ -1,2 +1,2 @@
 # people autorhized to approve a pull request
-* @sayeg84 @kbno @aguljas @jacopoventurin @yaoyic
+* @sayeg84 @kbno @aguljas @jacopoventurin

src/mlcg/neighbor_list/nvalchemi_impl.py

@@ -1,7 +1,7 @@
 from typing import Tuple, Optional
 import torch
 from torch_geometric.data import Data
-from nvalchemiops.neighborlist import (
+from nvalchemiops.torch.neighbors import (
     batch_cell_list,
     batch_naive_neighbor_list,
 )
@@ -72,7 +72,7 @@ def nvalchemi_naive_neighbor_list(
     if "pbc" in data:
         pbc = data.pbc
         # the type casting has to be done otherwise the library complains
-        cell = data.cell.to(torch.float32)
+        cell = data.cell.reshape(-1, 3, 3)
         with_pbc = True
     else:
         pbc = None
@@ -98,7 +98,10 @@ def nvalchemi_naive_neighbor_list(
 
     if with_pbc:
         (idx_i, idx_j), _, idx_S = result
-        cell_shifts = torch.matmul(idx_S.to(cell.dtype), cell)
+        # cell_shifts = torch.einsum("ni,nij->nj", idx_S.to(cell.dtype), cell[data.batch[idx_i]])
+        cell_shifts = (
+            idx_S.to(cell.dtype).unsqueeze(-1) * cell[data.batch[idx_i]]
+        ).sum(dim=1)
         return idx_i, idx_j, cell_shifts, None
     else:
         (idx_i, idx_j), _ = result
@@ -152,13 +155,13 @@ def nvalchemi_cell_neighbor_list(
     if "pbc" in data:
         pbc = data.pbc
         # the type casting has to be done otherwise the library complains
-        cell = data.cell.to(torch.float32)
+        cell = data.cell.reshape(-1, 3, 3)
         with_pbc = True
 
     else:
         box_size = 70
         warnings.warn(no_pbc_warning(box_size), UserWarning)
-        # this is required as the method needs
+        # this is required as the method needs PBC
         cell = (
             torch.zeros(
                 data.batch[-1] + 1,
@@ -190,9 +193,11 @@ def nvalchemi_cell_neighbor_list(
     )
 
     (idx_i, idx_j), _, idx_S = result
-
     if with_pbc:
-        cell_shifts = torch.matmul(idx_S.to(cell.dtype), cell)
+        # cell_shifts = torch.einsum("ni,nij->nj", idx_S.to(cell.dtype), cell[data.batch[idx_i]])
+        cell_shifts = (
+            idx_S.to(cell.dtype).unsqueeze(-1) * cell[data.batch[idx_i]]
+        ).sum(dim=1)
     else:
         cell_shifts = torch.zeros(
             (idx_i.shape[0], 3), dtype=data.pos.dtype, device=data.pos.device
@@ -245,7 +250,7 @@ def nvalchemi_cell_neighbor_list_raw(
     if "pbc" in data:
         pbc = data.pbc
         # the type casting has to be done otherwise the library complains
-        cell = data.cell.to(torch.float32)
+        cell = data.cell.to(torch.float64).reshape(-1, 3, 3)
         with_pbc = True
 
     else:

src/mlcg/nn/prior/repulsion.py

@@ -504,3 +504,218 @@ def from_base(cls, repulsion: ExpRepulsion, cutoff: float):
             single_r_0 = repulsion.r_0[key]
             prior_stats[key] = {"alpha": single_alpha, "r_0": single_r_0}
         return cls(statistics=prior_stats, cutoff=cutoff, name=repulsion.name)
+
+
+class _BaseFastRepulsion:
+
+    @staticmethod
+    def build_exclusion_table(
+        exclusions_nls: torch.Tensor,
+        n_atoms: int,
+        dtype: torch.dtype = torch.int32,
+    ) -> torch.Tensor:
+        r"""Compile (2, m) exclusions into a padded (n_atoms, width) table.
+
+        Row `a` holds every `b` such that (a, b) is excluded and a < b, padded
+        with -1. Only the a < b direction is stored, because queries are
+        canonicalised to a < b before lookup.
+
+        `dtype` is int32 to halve the gather bandwidth in the hot path; the
+        comparison against int64 indices is a fused elementwise cast.
+        """
+        if exclusions_nls.numel() == 0:
+            return torch.full(
+                (n_atoms, 0), -1, dtype=dtype, device=exclusions_nls.device
+            )
+
+        lo, hi = torch.aminmax(exclusions_nls, dim=0)  # tolerate unsorted input
+        if int(hi.max()) >= torch.iinfo(dtype).max:
+            raise ValueError(f"atom index exceeds {dtype} range")
+
+        counts = torch.bincount(lo, minlength=n_atoms)
+        width = int(counts.max())
+
+        # Stable sort groups equal `lo` contiguously; slot = offset within group.
+        order = torch.argsort(lo, stable=True)
+        lo_s, hi_s = lo[order], hi[order]
+        starts = torch.cumsum(counts, 0) - counts
+        slot = torch.arange(lo_s.numel(), device=lo.device) - starts[lo_s]
+
+        table = torch.full((n_atoms, width), -1, dtype=dtype, device=lo.device)
+        table[lo_s, slot] = hi_s.to(dtype)
+        return table
+
+    @staticmethod
+    def pair_mask(
+        geometric_nls: torch.Tensor,
+        excl_table: torch.Tensor,
+    ) -> torch.Tensor:
+        r"""Boolean mask over columns: True = keep (unique, canonical, included).
+
+        Assumes `geometric_nls` is symmetric, i.e. every pair appears as both
+        (i, j) and (j, i). True for radius_graph unless max_num_neighbors
+        truncated the neighbor list.
+        """
+        i, j = geometric_nls[0], geometric_nls[1]
+        nbrs = excl_table[i]  # (n, width) gather
+        return (i < j) & (nbrs != j.unsqueeze(1).to(nbrs.dtype)).all(dim=1)
+
+
+class FastCutoffRepulsion(CutoffRepulsion, _BaseFastRepulsion):
+    def __init__(
+        self, statistics: Dict, cutoff: float, name: str = "repulsion"
+    ) -> None:
+        super().__init__(statistics=statistics, cutoff=cutoff)
+        self.name = name
+
+    @staticmethod
+    def compute_with_dev(x, sigma):
+        r"""Method defining the repulsion and its derivative"""
+        orig = CutoffRepulsion.compute(x, sigma)
+        deriv = -6 * orig / x
+        return orig, deriv
+
+    def forward(self, data: AtomicData) -> AtomicData:
+        """Forward pass through the repulsion interaction.
+
+        Parameters
+        ----------
+        data:
+            Input AtomicData instance that possesses an appropriate
+            neighbor list containing both an 'index_mapping'
+            field and a 'mapping_batch' field for accessing
+            beads relevant to the interaction and scattering
+            the interaction energies onto the correct example/structure
+            respectively.
+
+        Returns
+        -------
+        AtomicData:
+            Updated AtomicData instance with the 'out' field
+            populated with the predicted energies for each
+            example/structure
+        """
+
+        exc_table = data.neighbor_list[self.name].get("exc_table")
+        if exc_table is None:
+            exc_table = self.build_exclusion_table(
+                data.neighbor_list[self.name]["index_mapping_exclusions"],
+                data.pos.shape[0],
+            )
+            data.neighbor_list[self.name]["exc_table"] = exc_table
+        # get relevant information for network neighborlist
+        mask = self.pair_mask(
+            data.neighbor_list["mpnn"]["index_mapping"], exc_table
+        )
+        mapping = data.neighbor_list["mpnn"]["index_mapping"][:, mask]
+        cell_shifts = data.neighbor_list["mpnn"]["cell_shifts"][mask, :]
+        mapping_batch = data.neighbor_list["mpnn"]["mapping_batch"][mask]
+        # compute features
+        features = compute_distances(data.pos, mapping, cell_shifts)
+
+        cutoff = data.neighbor_list["mpnn"]["rcut"]
+
+        interaction_types = tuple(
+            data.atom_types[mapping[ii]] for ii in range(self.order)
+        )
+
+        y = Repulsion.compute(features, self.sigma[interaction_types])
+        yc, dyc = CutoffRepulsion.compute_with_dev(
+            cutoff, self.sigma[interaction_types]
+        )
+        # ensure that the cutoff is continupus
+        y = y - yc - (features - cutoff) * (dyc)
+        y = scatter(y, mapping_batch, dim=0, reduce="sum")
+        data.out[self.name] = {"energy": y}
+        data.neighbor_list.pop("mpnn")
+        return data
+
+    @classmethod
+    def from_base(cls, repulsion: CutoffRepulsion):
+        prior_stats = {}
+        for key in repulsion.allowed_interaction_keys:
+            single_sigma = repulsion.sigma[key]
+            prior_stats[key] = {"sigma": single_sigma}
+        return cls(
+            statistics=prior_stats, cutoff=repulsion.cutoff, name=repulsion.name
+        )
+
+
+class FastCutoffExpRepulsion(CutoffExpRepulsion, _BaseFastRepulsion):
+    def __init__(
+        self, statistics: Dict, cutoff: float, name: str = "repulsion"
+    ) -> None:
+        super().__init__(statistics=statistics, cutoff=cutoff)
+        self.name = name
+
+    def forward(self, data: AtomicData) -> AtomicData:
+        """Forward pass through the repulsion interaction.
+
+        Parameters
+        ----------
+        data:
+            Input AtomicData instance that possesses an appropriate
+            neighbor list containing both an 'index_mapping'
+            field and a 'mapping_batch' field for accessing
+            beads relevant to the interaction and scattering
+            the interaction energies onto the correct example/structure
+            respectively.
+
+        Returns
+        -------
+        AtomicData:
+            Updated AtomicData instance with the 'out' field
+            populated with the predicted energies for each
+            example/structure
+        """
+
+        exc_table = data.neighbor_list[self.name].get("exc_table")
+        if exc_table is None:
+            exc_table = self.build_exclusion_table(
+                data.neighbor_list[self.name]["index_mapping_exclusions"],
+                data.pos.shape[0],
+            )
+            data.neighbor_list[self.name]["exc_table"] = exc_table
+        # get relevant information for network neighborlist
+        mask = self.pair_mask(
+            data.neighbor_list["mpnn"]["index_mapping"], exc_table
+        )
+        mapping = data.neighbor_list["mpnn"]["index_mapping"][:, mask]
+        cell_shifts = data.neighbor_list["mpnn"]["cell_shifts"][mask, :]
+        mapping_batch = data.neighbor_list["mpnn"]["mapping_batch"][mask]
+        # compute features
+        features = compute_distances(data.pos, mapping, cell_shifts)
+
+        cutoff = data.neighbor_list["mpnn"]["rcut"]
+        interaction_types = tuple(
+            data.atom_types[mapping[ii]] for ii in range(self.order)
+        )
+
+        y = ExpRepulsion.compute(
+            features, self.alpha[interaction_types], self.r_0[interaction_types]
+        )
+        yc, dyc = CutoffExpRepulsion.compute_with_dev(
+            cutoff,
+            self.alpha[interaction_types],
+            self.r_0[interaction_types],
+        )
+        # ensure that the cutoff is continupus
+        y = y - yc - (features - cutoff) * dyc
+        # we can override the
+        y = scatter(y, mapping_batch, dim=0, reduce="sum")
+        data.out[self.name] = {"energy": y}
+        # remove unwanted network neighborlist
+        data.neighbor_list.pop("mpnn")
+        return data
+
+    @classmethod
+    def from_base(cls, repulsion: CutoffExpRepulsion):
+        """initialize a CutoffExpRepulsion from a normal ExpRepulsion"""
+        prior_stats = {}
+        for key in repulsion.allowed_interaction_keys:
+            single_alpha = repulsion.alpha[key]
+            single_r_0 = repulsion.r_0[key]
+            prior_stats[key] = {"alpha": single_alpha, "r_0": single_r_0}
+        return cls(
+            statistics=prior_stats, cutoff=repulsion.cutoff, name=repulsion.name
+        )

src/mlcg/nn/schnet.py

@@ -139,14 +139,14 @@ def forward(self, data: AtomicData) -> AtomicData:
                 self.rbf_layer.cutoff.cutoff_upper,
                 self.max_num_neighbors,
             )[self.name]
+            data.neighbor_list["mpnn"] = neighbor_list
 
         edge_index = neighbor_list["index_mapping"]
         distances = compute_distances(
             data.pos,
             edge_index,
             neighbor_list["cell_shifts"],
         )
-
         rbf_expansion = self.rbf_layer(distances)
         num_batch = data.batch[-1] + 1
         for block in self.interaction_blocks:

src/mlcg/scripts/mlcg_add_exclusion_list.py

@@ -0,0 +1,60 @@
+import argparse
+import os
+from itertools import combinations
+
+import torch
+
+
+def parse_cli():
+    parser = argparse.ArgumentParser(
+        description="Command line tool for adding a non-bonded exclusion "
+        "list to a set of configurations, computed as the complement of "
+        "the existing 'non_bonded' neighbor list within the fully "
+        "connected graph. The result is saved next to the input file."
+    )
+    parser.add_argument(
+        "conf_path",
+        type=str,
+        help="path to the input configurations. Must be a valid .pt file.",
+    )
+    parser.add_argument(
+        "-n",
+        "--nls_name",
+        type=str,
+        default="non_bonded",
+        help="name of the neighbor list to build an exclusion to ",
+    )
+    return parser
+
+
+def main():
+    parser = parse_cli()
+    args = parser.parse_args()
+
+    conf_path = args.conf_path
+    nls_name = args.nls_name
+    confs = torch.load(conf_path, weights_only=False)
+    for conf in confs:
+        actual_nls = conf.neighbor_list[nls_name]["index_mapping"]
+        fully_connected_nls = torch.tensor(
+            list(combinations(range(conf.pos.shape[0]), 2))
+        ).T
+        num_atoms = conf.pos.shape[0]
+        actual_codes = actual_nls[0] * num_atoms + actual_nls[1]
+        full_codes = fully_connected_nls[0] * num_atoms + fully_connected_nls[1]
+        mask = ~torch.isin(full_codes, actual_codes)
+        exclusion_nls = fully_connected_nls[:, mask]
+        conf.neighbor_list[nls_name]["index_mapping_exclusions"] = exclusion_nls
+
+    new_name = conf_path.replace(".pt", "_with_nonbonded_exclusion.pt")
+    if not os.path.isfile(new_name):
+        print(f"New configurations saved at {new_name}")
+        torch.save(confs, new_name)
+    else:
+        raise ValueError(
+            f"File {new_name} exists already, please rename it or move it"
+        )
+
+
+if __name__ == "__main__":
+    main()