Specifiable options for periodic neighbors calculations

matsciml/datasets/transforms/pbc.py

@@ -5,6 +5,7 @@
 import numpy as np
 import torch
 from pymatgen.core import Lattice, Structure
+from loguru import logger
 
 from matsciml.common.types import DataDict
 from matsciml.datasets.transforms.base import AbstractDataTransform
@@ -38,11 +39,23 @@ def __init__(
         cutoff_radius: float,
         adaptive_cutoff: bool = False,
         backend: Literal["pymatgen", "ase"] = "pymatgen",
+        max_neighbors: int = 1000,
+        allow_self_loops: bool = False,
+        convert_to_unit_cell: bool = False,
+        is_cartesian: bool | None = None,
     ) -> None:
         super().__init__()
         self.cutoff_radius = cutoff_radius
         self.adaptive_cutoff = adaptive_cutoff
         self.backend = backend
+        self.max_neighbors = max_neighbors
+        self.allow_self_loops = allow_self_loops
+        if is_cartesian is not None and backend == "ase":
+            logger.warning(
+                "`is_cartesian` passed but using `ase` backend; option will not affect anything."
+            )
+        self.is_cartesian = is_cartesian
+        self.convert_to_unit_cell = convert_to_unit_cell
 
     def __call__(self, data: DataDict) -> DataDict:
         """
@@ -84,7 +97,10 @@ def __call__(self, data: DataDict) -> DataDict:
             structure = data["structure"]
             if isinstance(structure, Structure):
                 graph_props = calculate_periodic_shifts(
-                    structure, self.cutoff_radius, self.adaptive_cutoff
+                    structure,
+                    self.cutoff_radius,
+                    self.adaptive_cutoff,
+                    max_neighbors=self.max_neighbors,
                 )
                 data.update(graph_props)
                 return data
@@ -123,16 +139,25 @@ def __call__(self, data: DataDict) -> DataDict:
             data["atomic_numbers"],
             data["pos"],
             lattice=lattice,
+            convert_to_unit_cell=self.convert_to_unit_cell,
+            is_cartesian=self.is_cartesian,
         )
         if self.backend == "pymatgen":
             graph_props = calculate_periodic_shifts(
-                structure, self.cutoff_radius, self.adaptive_cutoff
+                structure, self.cutoff_radius, self.adaptive_cutoff, self.max_neighbors
             )
         elif self.backend == "ase":
             graph_props = calculate_ase_periodic_shifts(
-                data, self.cutoff_radius, self.adaptive_cutoff
+                data, self.cutoff_radius, self.adaptive_cutoff, self.max_neighbors
             )
         else:
             raise RuntimeError(f"Requested backend f{self.backend} not available.")
         data.update(graph_props)
+        if not self.allow_self_loops:
+            mask = data["src_nodes"] == data["dst_nodes"]
+            # only mask out self-loops within the same image
+            mask &= data["unit_offsets"].sum(dim=-1) == 0
+            # apply mask to each of the tensors that depend on edges
+            for key in ["src_nodes", "dst_nodes", "images", "unit_offsets", "offsets"]:
+                data[key] = data[key][mask]
         return data

matsciml/datasets/transforms/tests/test_pbc.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from collections import Counter
+
+import torch
+import pytest
+import numpy as np
+from pymatgen.core import Structure, Lattice
+
+from matsciml.datasets.transforms import PeriodicPropertiesTransform
+
+"""
+This module uses reference Materials project structures and tests
+the edge calculation routines to ensure they at least work with
+various parameters.
+
+The key thing here is at least using feasible structures to perform
+this check, rather than using randomly generated coordinates and
+lattices, even if composing them isn't meaningful.
+"""
+
+hexa = Lattice.from_parameters(
+    4.81, 4.809999999999999, 13.12, 90.0, 90.0, 120.0, vesta=True
+)
+cubic = Lattice.from_parameters(6.79, 6.79, 12.63, 90.0, 90.0, 90.0, vesta=True)
+
+# mp-1143
+alumina = Structure(
+    hexa,
+    species=["Al", "O"],
+    coords=[[1 / 3, 2 / 3, 0.814571], [0.360521, 1 / 3, 0.583333]],
+    coords_are_cartesian=False,
+)
+# mp-1267
+nac = Structure(
+    cubic,
+    species=["Na", "C"],
+    coords=[[0.688819, 3 / 4, 3 / 8], [0.065833, 0.565833, 0.0]],
+    coords_are_cartesian=False,
+)
+
+
+@pytest.mark.parametrize(
+    "coords",
+    [
+        alumina.cart_coords,
+        nac.cart_coords,
+    ],
+)
+@pytest.mark.parametrize(
+    "cell",
+    [
+        hexa.matrix,
+        cubic.matrix,
+    ],
+)
+@pytest.mark.parametrize("self_loops", [True, False])
+@pytest.mark.parametrize("backend", ["pymatgen", "ase"])
+@pytest.mark.parametrize(
+    "cutoff_radius", [6.0, 9.0, 15.0]
+)  # TODO figure out why pmg fails on 3
+def test_periodic_generation(
+    coords: np.ndarray,
+    cell: np.ndarray,
+    self_loops: bool,
+    backend: str,
+    cutoff_radius: float,
+):
+    coords = torch.FloatTensor(coords)
+    cell = torch.FloatTensor(cell)
+    transform = PeriodicPropertiesTransform(
+        cutoff_radius=cutoff_radius,
+        adaptive_cutoff=False,
+        backend=backend,
+        max_neighbors=10,
+        allow_self_loops=self_loops,
+    )
+    num_atoms = coords.size(0)
+    atomic_numbers = torch.ones(num_atoms)
+    packed_data = {"pos": coords, "cell": cell, "atomic_numbers": atomic_numbers}
+    output = transform(packed_data)
+    # check to make sure no source node has more than 10 neighbors
+    src_nodes = output["src_nodes"].tolist()
+    counts = Counter(src_nodes)
+    for index, count in counts.items():
+        if not self_loops:
+            assert count < 10, print(f"Node {index} has too many counts. {src_nodes}")