Neighbors convert cpp

pyproject.toml

@@ -0,0 +1,7 @@
+[build-system]
+requires = [
+    "setuptools >= 68",
+    "wheel",
+    "torch",
+]
+build-backend = "setuptools.build_meta"

requirements.txt

@@ -6,3 +6,4 @@ jupyter
 scikit-learn
 torch
 torch_geometric
+matscipy

setup.py

@@ -1,9 +1,30 @@
-from setuptools import setup
-
+from setuptools import setup, Extension
+from torch.utils.cpp_extension import BuildExtension, include_paths, library_paths
 
 with open("requirements.txt") as f:
     requirements = f.read().splitlines()
 
+# Collecting include and library paths
+include_dirs = include_paths()
+library_dirs = library_paths()
+
+libraries = []
+
+libraries.append('c10')
+libraries.append('torch')
+libraries.append('torch_cpu')
+
+
+# Defining the extension module without specifying the unwanted libraries
+neighbors_convert_extension = Extension(
+    name="pet.neighbors_convert",
+    sources=["src/neighbors_convert.cpp"],
+    include_dirs=include_dirs,
+    library_dirs=library_dirs,
+    libraries=libraries,
+    language='c++',
+)
+
 setup(
     name="pet",
     version="0.0.0",
@@ -18,4 +39,12 @@
         ],
     },
     install_requires=requirements,
+    ext_modules=[neighbors_convert_extension],
+    cmdclass={
+        "build_ext": BuildExtension.with_options(no_python_abi_suffix=True)
+    },
+    package_data={
+        'pet': ['neighbors_convert.so'],  # Ensure the shared object file is included
+    },
+    include_package_data=True,
 )

src/__init__.py

@@ -1 +1,15 @@
 from .single_struct_calculator import SingleStructCalculator
+
+import torch
+import importlib.resources as pkg_resources
+
+def load_neighbors_convert():
+    try:
+        # Locate the shared object file in the package
+        with pkg_resources.files(__name__).joinpath('neighbors_convert.so') as lib_path:
+            # Load the shared object file
+            torch.ops.load_library(str(lib_path))
+    except Exception as e:
+        print(f"Failed to load neighbors_convert.so: {e}")
+
+load_neighbors_convert()

src/molecule.py

@@ -3,7 +3,7 @@
 import numpy as np
 from torch_geometric.data import Data
 from .long_range import get_reciprocal, get_all_k, get_volume
-
+from matscipy.neighbours import neighbour_list as neighbor_list
 
 class Molecule:
     def __init__(
@@ -188,6 +188,73 @@ def get_graph(self, max_num, all_species, max_num_k):
         return result
 
 
+class MoleculeCPP:
+    def __init__(
+        self, atoms, r_cut, use_additional_scalar_attributes, use_long_range, k_cut
+    ):
+
+        self.use_additional_scalar_attributes = use_additional_scalar_attributes
+        self.atoms = atoms
+        self.r_cut = r_cut
+        self.use_long_range = use_long_range
+        self.k_cut = k_cut
+
+        if self.use_long_range:
+            raise NotImplementedError("Long range is not implemented in cpp")
+        if self.use_additional_scalar_attributes:
+            raise NotImplementedError("Additional scalar attributes are not implemented in cpp")
+
+        def is_3d_crystal(atoms):
+            pbc = atoms.get_pbc()
+            if isinstance(pbc, bool):
+                return pbc
+            return all(pbc)
+
+        if is_3d_crystal(atoms):
+            i_list, j_list, D_list, S_list = neighbor_list('ijDS', atoms, r_cut)
+        else:
+            i_list, j_list, D_list, S_list = ase.neighborlist.neighbor_list(
+                "ijDS", atoms, r_cut
+            )
+
+        self.i_list = torch.tensor(i_list, dtype=torch.int64).contiguous()
+        self.j_list = torch.tensor(j_list, dtype=torch.int64).contiguous()
+        self.D_list = torch.tensor(D_list, dtype=torch.get_default_dtype()).contiguous()
+        self.S_list = torch.tensor(S_list, dtype=torch.int64).contiguous()
+        self.species = torch.tensor(atoms.get_atomic_numbers(), dtype=torch.int64).contiguous()
+        if len(self.i_list) == 0:
+            self.max_num = 0
+        else:
+            self.max_num = torch.max(torch.bincount(self.i_list))
+
+    def get_num_k(self):
+        raise NotImplementedError("Long range is not implemented in cpp")
+
+    def get_max_num(self):
+        return self.max_num
+
+    def get_graph(self, max_num, all_species, max_num_k):
+        n_atoms = len(self.atoms.get_atomic_numbers())
+        all_species = torch.tensor(all_species, dtype=torch.int64).contiguous()
+
+        # torch.ops.my_extension.process(i_list, j_list, S_list, D_list, max_size, n_atoms, species, None)
+        neighbors_index, relative_positions, nums, mask, neighbor_species, neighbors_pos, species_mapped = torch.ops.neighbors_convert.process(self.i_list, self.j_list, self.S_list, self.D_list, max_num, n_atoms, self.species, all_species)
+
+        kwargs = {
+            "central_species": species_mapped,
+            "x": relative_positions,
+            "neighbor_species": neighbor_species,
+            "neighbors_pos": neighbors_pos,
+            "neighbors_index": neighbors_index.transpose(0, 1),
+            "nums": nums,
+            "mask": mask,
+            "n_atoms": len(self.atoms.positions),
+        }
+
+        result = Data(**kwargs)
+
+        return result
+
 def batch_to_dict(batch):
     batch_dict = {
         "x": batch.x,