diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 16770a9e..e19ee0f3 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,2 +1,2 @@ # people autorhized to approve a pull request -* @sayeg84 @kbno @aguljas @jacopoventurin @yaoyic \ No newline at end of file +* @sayeg84 @kbno @aguljas @jacopoventurin \ No newline at end of file diff --git a/src/mlcg/neighbor_list/nvalchemi_impl.py b/src/mlcg/neighbor_list/nvalchemi_impl.py index 99047f8d..0fa9e9ab 100644 --- a/src/mlcg/neighbor_list/nvalchemi_impl.py +++ b/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: diff --git a/src/mlcg/nn/prior/repulsion.py b/src/mlcg/nn/prior/repulsion.py index a4373b9a..e1f5db1a 100644 --- a/src/mlcg/nn/prior/repulsion.py +++ b/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 + ) diff --git a/src/mlcg/nn/schnet.py b/src/mlcg/nn/schnet.py index 9bb981fb..490f74c6 100644 --- a/src/mlcg/nn/schnet.py +++ b/src/mlcg/nn/schnet.py @@ -139,6 +139,7 @@ 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( @@ -146,7 +147,6 @@ def forward(self, data: AtomicData) -> AtomicData: edge_index, neighbor_list["cell_shifts"], ) - rbf_expansion = self.rbf_layer(distances) num_batch = data.batch[-1] + 1 for block in self.interaction_blocks: diff --git a/src/mlcg/scripts/mlcg_add_exclusion_list.py b/src/mlcg/scripts/mlcg_add_exclusion_list.py new file mode 100644 index 00000000..132ffe00 --- /dev/null +++ b/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() diff --git a/tests/unit/neighbor_list/assets/circular_with_nonbonded_exclusion.pt b/tests/unit/neighbor_list/assets/circular_with_nonbonded_exclusion.pt new file mode 100644 index 00000000..9f178786 Binary files /dev/null and b/tests/unit/neighbor_list/assets/circular_with_nonbonded_exclusion.pt differ diff --git a/tests/unit/neighbor_list/assets/dissolved_with_nonbonded_exclusion.pt b/tests/unit/neighbor_list/assets/dissolved_with_nonbonded_exclusion.pt new file mode 100644 index 00000000..12d98181 Binary files /dev/null and b/tests/unit/neighbor_list/assets/dissolved_with_nonbonded_exclusion.pt differ diff --git a/tests/unit/neighbor_list/assets/exp_rep_model.pt b/tests/unit/neighbor_list/assets/exp_rep_model.pt new file mode 100644 index 00000000..88120217 Binary files /dev/null and b/tests/unit/neighbor_list/assets/exp_rep_model.pt differ diff --git a/tests/unit/neighbor_list/assets/ordered_with_nonbonded_exclusion.pt b/tests/unit/neighbor_list/assets/ordered_with_nonbonded_exclusion.pt new file mode 100644 index 00000000..cdf4d9fe Binary files /dev/null and b/tests/unit/neighbor_list/assets/ordered_with_nonbonded_exclusion.pt differ diff --git a/tests/unit/neighbor_list/assets/rep_model.pt b/tests/unit/neighbor_list/assets/rep_model.pt new file mode 100644 index 00000000..60633119 Binary files /dev/null and b/tests/unit/neighbor_list/assets/rep_model.pt differ diff --git a/tests/unit/neighbor_list/test_neighbor_list_nvalchemi.py b/tests/unit/neighbor_list/test_neighbor_list_nvalchemi.py new file mode 100644 index 00000000..d2ad73f3 --- /dev/null +++ b/tests/unit/neighbor_list/test_neighbor_list_nvalchemi.py @@ -0,0 +1,307 @@ +import pytest +import ase +from ase.build import bulk, molecule +from torch_geometric.loader import DataLoader +import numpy as np +import torch + +from mlcg.neighbor_list.utils import ase2data +from mlcg.neighbor_list.ase_impl import ase_neighbor_list +from mlcg.geometry.internal_coordinates import compute_distances + +try: + from mlcg.neighbor_list.nvalchemi_impl import ( + nvalchemi_naive_neighbor_list, + nvalchemi_cell_neighbor_list, + nvalchemi_cell_neighbor_list_raw, + ) + + NVALCH_AVAILABLE = True +except ImportError: + print( + "nalchemi is not installed. Please install with " + + "pip install nvalchemi-toolkit-ops" + ) + NVALCH_AVAILABLE = False + + +def sort_edges(edge_index, *tensors): + if edge_index.numel() == 0: + return (edge_index,) + tuple(t for t in tensors) + stride = edge_index.max().item() + 1 + key = edge_index[0] * stride + edge_index[1] + order = torch.argsort(key) + return (edge_index[:, order],) + tuple(t[order] for t in tensors) + + +def bulk_metal(): + a = 4.0 + b = a / 2 + frames = [ + ase.Atoms( + "Ag", + cell=[(0, b, b), (b, 0, b), (b, b, 0)], + pbc=True, + ), + bulk("Cu", "fcc", a=3.6), + ] + return frames + + +def atomic_structures(): + frames = [ + molecule("CH3CH2NH2"), + molecule("H2O"), + molecule("methylenecyclopropane"), + ] + bulk_metal() + for frame in frames: + yield (frame.get_chemical_symbols(), frame) + + +nvalchemi_test_set = [ + (name, frame, rc, self_interaction) + for (name, frame) in atomic_structures() + for rc in range(2, 7, 2) + for self_interaction in [False] +] + +# resolved lazily by name so that collection works without nvalchemi installed +NVALCHEMI_CELL_METHODS = ( + { + "cell": nvalchemi_cell_neighbor_list, + "raw": nvalchemi_cell_neighbor_list_raw, + } + if NVALCH_AVAILABLE + else {} +) + +nvalchemi_cell_method_names = ["cell", "raw"] + + +@pytest.mark.skipif( + not NVALCH_AVAILABLE, + reason="nvalchemi is not available (install nvalchemi-toolkit-ops)", +) +@pytest.mark.parametrize( + "name, frame, cutoff, self_interaction", + nvalchemi_test_set, +) +def test_neighborlist_nvalchemi(name, frame, cutoff, self_interaction): + """Check that nvalchemi_neighbor_list gives the same NL as ASE by comparing + the resulting sorted list of distances between neighbors.""" + data_list = [ase2data(frame)] + dataloader = DataLoader(data_list, batch_size=1) + distance_results = {} + neighs_results = {} + method_list = ["current_nls_method", "ase_ref"] + for met_name in method_list: + dds = [] + for data in dataloader: + data.cell = data.cell.to(data.pos.dtype) + if met_name == "ase_ref": + met = ase_neighbor_list + else: + met = nvalchemi_naive_neighbor_list + idx_i, idx_j, cell_shifts, _ = met( + data, cutoff, self_interaction=self_interaction + ) + + dd = (data.pos[idx_j] - data.pos[idx_i] + cell_shifts).norm(dim=1) + dds.extend(dd.numpy()) + dds = np.sort(dds) + edge_index = torch.stack([idx_i, idx_j], dim=0) + edge_index = sort_edges(edge_index) + distance_results[met_name] = dds + neighs_results[met_name] = edge_index + assert np.allclose( + distance_results["current_nls_method"], distance_results["ase_ref"] + ) + assert np.allclose( + neighs_results["current_nls_method"], neighs_results["ase_ref"] + ) + + +@pytest.mark.skipif( + not NVALCH_AVAILABLE, + reason="nvalchemi is not available (install nvalchemi-toolkit-ops)", +) +@pytest.mark.parametrize("nls_name", nvalchemi_cell_method_names) +def test_neighborlist_pbc_nvalchemi(nls_name): + """Test that neighbor list with PBC correctly handles periodic images + and produces the same results as ASE reference implementation.""" + nls_method = NVALCHEMI_CELL_METHODS[nls_name] + + # Create test structures with PBC + structures = [ + bulk("Cu", "fcc", a=3.6), + ] + + cutoffs = [ + 3.0, + ] + + for structure in structures: + for cutoff in cutoffs: + for self_interaction in [False]: + # Convert to data format + data_list = [ase2data(structure)] + dataloader = DataLoader(data_list, batch_size=1) + + # Get nvalchemi neighbor list distances + nvalchemi_distances = [] + for data in dataloader: + data.cell = data.cell.to(data.pos.dtype) + if "cell" in data: + print("Cell:\n", data.cell) + idx_i, idx_j, cell_shifts, _ = nls_method( + data, cutoff, self_interaction=self_interaction + ) + if nls_name == "raw": + cell = data.cell.reshape(-1, 3, 3) + cell_shifts = ( + cell_shifts.to(cell.dtype) + .to(cell.dtype) + .unsqueeze(-1) + * cell[data.batch[idx_i]] + ).sum(dim=1) + mapping = torch.stack([idx_i, idx_j], dim=0) + dd = compute_distances(data.pos, mapping, cell_shifts) + nvalchemi_distances.extend(dd.numpy()) + + nvalchemi_distances = np.sort(nvalchemi_distances) + + # Get ASE reference distances + ase_distances = [] + for data in dataloader: + data.cell = data.cell.to(data.pos.dtype) + idx_i, idx_j, ase_cell_shifts, _ = ase_neighbor_list( + data, cutoff, self_interaction=self_interaction + ) + dd = ( + data.pos[idx_j] - data.pos[idx_i] + ase_cell_shifts + ).norm(dim=1) + ase_distances.extend(dd.numpy()) + + ase_distances = np.sort(ase_distances) + + assert np.allclose( + ase_distances, nvalchemi_distances, rtol=1e-5, atol=1e-6 + ) + + assert np.all(nvalchemi_distances <= cutoff + 1e-6) + + +@pytest.mark.skipif( + not NVALCH_AVAILABLE, + reason="nvalchemi is not available (install nvalchemi-toolkit-ops)", +) +@pytest.mark.parametrize("nls_name", nvalchemi_cell_method_names) +def test_pbc_minimum_image_convention_nvalchemi(nls_name): + """Test that PBC neighbor list correctly applies minimum image convention. + Neighbors should be found across periodic boundaries at the shortest distance. + """ + nls_method = NVALCHEMI_CELL_METHODS[nls_name] + + # Create a simple cubic cell with one atom + atoms = ase.Atoms( + "Ar", + positions=[[0.1, 0.1, 0.1]], + cell=[10.0, 10.0, 10.0], + pbc=True, + ) + + cutoff = 3.0 + data_list = [ase2data(atoms)] + dataloader = DataLoader(data_list, batch_size=1) + + for data in dataloader: + data.cell = data.cell.to(data.pos.dtype) + idx_i, idx_j, cell_shifts, _ = nls_method( + data, cutoff, self_interaction=False + ) + + assert len(idx_i) == 0, "Single isolated atom should have no neighbors" + + atoms = ase.Atoms( + "Ar2", + positions=[[0.1, 0.1, 0.1], [9.9, 0.1, 0.1]], + cell=[10.0, 10.0, 10.0], + pbc=True, + ) + + data_list = [ase2data(atoms)] + dataloader = DataLoader(data_list, batch_size=1) + + distances = [] + for data in dataloader: + data.cell = data.cell.to(data.pos.dtype) + idx_i, idx_j, cell_shifts, _ = nls_method( + data, cutoff, self_interaction=False + ) + if nls_name == "raw": + cell = data.cell.reshape(-1, 3, 3) + cell_shifts = ( + cell_shifts.to(cell.dtype).unsqueeze(-1) + * cell[data.batch[idx_i]] + ).sum(dim=1) + + mapping = torch.stack([idx_i, idx_j], dim=0) + dd = compute_distances(data.pos, mapping, cell_shifts) + distances.extend(dd.numpy()) + + distances = np.sort(distances) + assert np.all(distances < 1.0), ( + f"Minimum image convention not applied correctly. " + f"Distances: {distances}" + ) + + +@pytest.mark.skipif( + not NVALCH_AVAILABLE, + reason="nvalchemi is not available (install nvalchemi-toolkit-ops)", +) +@pytest.mark.parametrize("nls_name", nvalchemi_cell_method_names) +def test_mixed_pbc_nvalchemi(nls_name): + """Test neighbor list with partial periodic boundary conditions.""" + nls_method = NVALCHEMI_CELL_METHODS[nls_name] + + atoms = ase.Atoms( + "C4", + positions=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 5.0], + ], + cell=[5.0, 5.0, 10.0], + pbc=[True, True, False], # Periodic in x,y but not z + ) + + cutoff = 2.0 + data_list = [ase2data(atoms)] + dataloader = DataLoader(data_list, batch_size=1) + distances = [] + for data in dataloader: + data.cell = data.cell.to(data.pos.dtype) + idx_i, idx_j, cell_shifts, _ = nls_method( + data, cutoff, self_interaction=False + ) + + cell = data.cell.reshape(-1, 3, 3) + if nls_name == "raw": + cell = data.cell.reshape(-1, 3, 3) + cell_shifts = ( + cell_shifts.to(cell.dtype).unsqueeze(-1) + * cell[data.batch[idx_i]] + ).sum(dim=1) + + mapping = torch.stack([idx_i, idx_j], dim=0) + dd = compute_distances(data.pos, mapping, cell_shifts) + distances.extend(dd.numpy()) + distances = np.sort(distances) + + assert np.all(distances <= cutoff + 1e-6) + + atoms_involved = torch.cat([idx_i, idx_j]).unique() + + assert len(atoms_involved) >= 3 diff --git a/tests/unit/neighbor_list/test_neighbor_list.py b/tests/unit/neighbor_list/test_neighbor_list_torch.py similarity index 91% rename from tests/unit/neighbor_list/test_neighbor_list.py rename to tests/unit/neighbor_list/test_neighbor_list_torch.py index 05f9a116..89b0784f 100644 --- a/tests/unit/neighbor_list/test_neighbor_list.py +++ b/tests/unit/neighbor_list/test_neighbor_list_torch.py @@ -10,16 +10,14 @@ from mlcg.neighbor_list.torch_impl import torch_neighbor_list from mlcg.geometry.internal_coordinates import compute_distances -try: - from mlcg.neighbor_list.nvalchemi_impl import nvalchemi_neighbor_list - - NVALCH_AVAILABLE = True -except ImportError: - print( - "nalchemiis not installed. Please install with " - + "pip install nvalchemi-toolkit-ops" - ) - NVALCH_AVAILABLE = False + +def sort_edges(edge_index, *tensors): + if edge_index.numel() == 0: + return (edge_index,) + tuple(t for t in tensors) + stride = edge_index.max().item() + 1 + key = edge_index[0] * stride + edge_index[1] + order = torch.argsort(key) + return (edge_index[:, order],) + tuple(t[order] for t in tensors) def bulk_metal(): @@ -53,13 +51,12 @@ def atomic_structures(): for self_interaction in [False, True] ] -if NVALCH_AVAILABLE: - test_set += [ - (nvalchemi_neighbor_list, name, frame, rc, self_interaction) - for (name, frame) in atomic_structures() - for rc in range(2, 7, 2) - for self_interaction in [False] - ] +nvalchemi_test_set = [ + (name, frame, rc, self_interaction) + for (name, frame) in atomic_structures() + for rc in range(2, 7, 2) + for self_interaction in [False] +] @pytest.mark.parametrize( @@ -87,10 +84,10 @@ def test_neighborlist(nls_method, name, frame, cutoff, self_interaction): dd = (data.pos[idx_j] - data.pos[idx_i] + cell_shifts).norm(dim=1) dds.extend(dd.numpy()) dds = np.sort(dds) + edge_index = torch.stack([idx_i, idx_j], dim=0) + edge_index = sort_edges(edge_index) distance_results[met_name] = dds - neighs_results[met_name] = torch.stack([idx_i, idx_j], dim=0).sort( - dim=1 - ) + neighs_results[met_name] = edge_index assert np.allclose( distance_results["current_nls_method"], distance_results["ase_ref"] ) diff --git a/tests/unit/neighbor_list/test_nls_sharing.py b/tests/unit/neighbor_list/test_nls_sharing.py new file mode 100644 index 00000000..362999f1 --- /dev/null +++ b/tests/unit/neighbor_list/test_nls_sharing.py @@ -0,0 +1,195 @@ +import os + +import pytest +import torch +from torch_geometric.data.collate import collate + +from mlcg.data import AtomicData +from mlcg.nn.prior.repulsion import FastCutoffRepulsion, FastCutoffExpRepulsion + +try: + import nvalchemiops # noqa: F401 + + NVALCHEMI_AVAILABLE = True +except ImportError: + NVALCHEMI_AVAILABLE = False + +ASSETS_DIR = os.path.join(os.path.dirname(__file__), "assets") + +pytestmark = pytest.mark.skipif( + not (torch.cuda.is_available() and NVALCHEMI_AVAILABLE), + reason="requires a CUDA device and the nvalchemi-toolkit-ops package", +) + + +def _load_collated_data(asset_name: str, n_replicas: int) -> AtomicData: + confs = [] + for _ in range(n_replicas): + confs += torch.load( + os.path.join(ASSETS_DIR, f"{asset_name}.pt"), + weights_only=False, + ) + col_data, _, _ = collate(AtomicData, confs) + col_data = col_data.to("cuda") + col_data.pos = col_data.pos.to(torch.float32) + if "cell" in col_data: + col_data.cell = col_data.cell.to(torch.float32) + return col_data + + +@pytest.fixture( + params=[ + "dissolved_with_nonbonded_exclusion", + "ordered_with_nonbonded_exclusion", + "circular_with_nonbonded_exclusion", + ] +) +def col_data(request): + n_replicas = 3 + return _load_collated_data(request.param, n_replicas) + + +@pytest.mark.parametrize( + "nls_distance_method", ["torch", "nvalchemi_naive", "nvalchemi_cell"] +) +def test_fast_cutoff_exp_repulsion_and_network_nls_match( + nls_distance_method, col_data +): + """ + A model using `FastCutoffExpRepulsion` (a cutoff-table based + reimplementation of `CutoffExpRepulsion`) together with the + `nls_distance_method` neighbor list method should produce the same + energies and forces as the reference model, which uses the base + repulsion prior and the default `torch` neighbor list. Consistency + is checked over many steps of a randomly perturbed trajectory so + that both implementations are exercised over an evolving set of + neighbor lists. + """ + n_steps = 20 + n_replicas = col_data.n_atoms.shape[0] + + model = torch.load( + os.path.join(ASSETS_DIR, "exp_rep_model.pt"), weights_only=False + ) + new_model = torch.load( + os.path.join(ASSETS_DIR, "exp_rep_model.pt"), weights_only=False + ) + new_model.models["non_bonded"].model = FastCutoffExpRepulsion.from_base( + model.models["non_bonded"].model + ) + new_model.models["SchNet"].model.nls_distance_method = nls_distance_method + + model = model.to("cuda") + new_model = new_model.to("cuda") + + old_eners = torch.zeros(n_steps, n_replicas) + old_forces = torch.zeros(n_steps, *col_data.pos.shape) + new_eners = torch.zeros(n_steps, n_replicas) + new_forces = torch.zeros(n_steps, *col_data.pos.shape) + + for i in range(n_steps): + model(col_data) + old_eners[i] = col_data.out["non_bonded"]["energy"].detach().cpu() + old_forces[i] = col_data.out["non_bonded"]["forces"].detach().cpu() + col_data.out = {} + + new_model(col_data) + new_eners[i] = col_data.out["non_bonded"]["energy"].detach().cpu() + new_forces[i] = col_data.out["non_bonded"]["forces"].detach().cpu() + + # Advance positions along the current forces plus noise so that + # consistency is checked across an evolving set of neighbor lists, + # rather than just the initial configuration. + col_data.pos += 1e-5 * col_data.out["non_bonded"][ + "forces" + ] + 1e-1 * torch.randn(col_data.pos.shape, device=col_data.pos.device) + col_data.out = {} + + atol_energy = torch.mean(old_eners) * 1e-5 + torch.testing.assert_close( + new_eners, old_eners, atol=atol_energy, rtol=1e-5 + ) + + atol_forces = torch.mean(torch.abs(old_forces)) * 1e-3 + torch.testing.assert_close( + new_forces, old_forces, atol=atol_forces, rtol=1e-3 + ) + + +@pytest.fixture( + params=[ + "dissolved_with_nonbonded_exclusion", + "ordered_with_nonbonded_exclusion", + "circular_with_nonbonded_exclusion", + ] +) +def col_data(request): + n_replicas = 3 + return _load_collated_data(request.param, n_replicas) + + +@pytest.mark.parametrize( + "nls_distance_method", ["torch", "nvalchemi_naive", "nvalchemi_cell"] +) +def test_fast_cutoff_repulsion_and_network_nls_match( + nls_distance_method, col_data +): + """ + A model using `FastCutoffRepulsion` (a cutoff-table based + reimplementation of `CutoffRepulsion`) together with the + `nls_distance_method` neighbor list method should produce the same + energies and forces as the reference model, which uses the base + repulsion prior and the default `torch` neighbor list. Consistency + is checked over many steps of a randomly perturbed trajectory so + that both implementations are exercised over an evolving set of + neighbor lists. + """ + n_steps = 20 + n_replicas = col_data.n_atoms.shape[0] + + model = torch.load( + os.path.join(ASSETS_DIR, "rep_model.pt"), weights_only=False + ) + new_model = torch.load( + os.path.join(ASSETS_DIR, "rep_model.pt"), weights_only=False + ) + new_model.models["non_bonded"].model = FastCutoffRepulsion.from_base( + model.models["non_bonded"].model + ) + new_model.models["SchNet"].model.nls_distance_method = nls_distance_method + + model = model.to("cuda") + new_model = new_model.to("cuda") + + old_eners = torch.zeros(n_steps, n_replicas) + old_forces = torch.zeros(n_steps, *col_data.pos.shape) + new_eners = torch.zeros(n_steps, n_replicas) + new_forces = torch.zeros(n_steps, *col_data.pos.shape) + + for i in range(n_steps): + model(col_data) + old_eners[i] = col_data.out["non_bonded"]["energy"].detach().cpu() + old_forces[i] = col_data.out["non_bonded"]["forces"].detach().cpu() + col_data.out = {} + + new_model(col_data) + new_eners[i] = col_data.out["non_bonded"]["energy"].detach().cpu() + new_forces[i] = col_data.out["non_bonded"]["forces"].detach().cpu() + + # Advance positions along the current forces plus noise so that + # consistency is checked across an evolving set of neighbor lists, + # rather than just the initial configuration. + col_data.pos += 1e-5 * col_data.out["non_bonded"][ + "forces" + ] + 1e-1 * torch.randn(col_data.pos.shape, device=col_data.pos.device) + col_data.out = {} + + atol_energy = torch.mean(old_eners) * 1e-5 + torch.testing.assert_close( + new_eners, old_eners, atol=atol_energy, rtol=1e-5 + ) + + atol_forces = torch.mean(torch.abs(old_forces)) * 1e-3 + torch.testing.assert_close( + new_forces, old_forces, atol=atol_forces, rtol=1e-3 + )