Add dpnblist package with CUDA support

dmff/common/nblist.py

@@ -7,6 +7,73 @@
     freud = None
     import warnings
     warnings.warn("WARNING: freud not installed, users need to create neighbor list by themselves.")
+try:
+   import dpnblist
+except ImportError:
+    dpnblist = None
+    import warnings
+    warnings.warn("WARNING: dpdpnblist not installed, users need to create neighbor list by themselves.")
+
+class NeighborListDp:
+    def __init__(self, alg_type, box, rcut, cov_map, padding=True):
+        if dpnblist is None:
+            raise ImportError("dpnblist not installed.")
+        self.box = dpnblist.Box([box[0][0], box[1][1], box[2][2]], [90.0, 90.0, 90.0])
+        self.nb = dpnblist.NeighborList(alg_type)
+        self.flag = False
+        self.rcut = rcut
+        self.capacity_multiplier = None
+        self.padding = padding
+        self.cov_map = cov_map
+
+    def _do_cov_map(self, pairs):
+        nbond = self.cov_map[pairs[:, 0], pairs[:, 1]]
+        pairs = jnp.concatenate([pairs, nbond[:, None]], axis=1)
+        return pairs
+
+    def allocate(self, coords, box=None):
+        self._positions = coords  # cache it
+        dbox = dpnblist.Box([box[0][0], box[1][1], box[2][2]], [90.0, 90.0, 90.0]) if box is not None else self.box
+        self.nb.build(dbox, coords, self.rcut)
+        pair = self.nb.get_neighbor_pair()
+        nlist = np.vstack((pair[:, 0], pair[:, 1])).T
+        nlist = nlist.astype(np.int32)
+        msk = (nlist[:, 0] - nlist[:, 1]) < 0
+        nlist = nlist[msk]
+        if self.capacity_multiplier is None:
+            self.capacity_multiplier = int(nlist.shape[0] * 1.3)
+
+        if not self.padding:
+            self._pairs = self._do_cov_map(nlist)
+            return self._pairs
+
+        self.capacity_multiplier = max(self.capacity_multiplier, nlist.shape[0])
+        padding_width = self.capacity_multiplier - nlist.shape[0]
+        if padding_width == 0:
+            self._pairs = self._do_cov_map(nlist)
+            return self._pairs
+        elif padding_width > 0:
+            padding = np.ones((self.capacity_multiplier - nlist.shape[0], 2), dtype=np.int32) * coords.shape[0]
+            nlist = np.vstack((nlist, padding))
+            self._pairs = self._do_cov_map(nlist)
+            return self._pairs
+        else:
+            raise ValueError("padding width < 0")
+
+    def update(self, positions, box=None):
+        self.allocate(positions, box)
+
+    @property
+    def pairs(self):
+        return self._pairs
+
+    @property
+    def scaled_pairs(self):
+        return self._pairs
+
+    @property
+    def positions(self):
+        return self._positions
 
 
 class NeighborListFreud:
@@ -69,10 +136,10 @@ def scaled_pairs(self):
     def positions(self):
         return self._positions
 
-
 class NeighborList(NeighborListFreud):
     ...
 
+
 class NoCutoffNeighborList:
 
     def __init__(self, cov_map, padding=True):
@@ -88,8 +155,8 @@ def _do_cov_map(self, pairs):
     def allocate(self, coords, box=None):
         self._positions = coords  # cache it
         natoms = coords.shape[0]
-        nblist = np.fromiter(permutations(range(natoms), 2), dtype=np.dtype((int, 2)))
-        nlist = nblist[nblist[:, 0] < nblist[:, 1]]
+        dpnblist = np.fromiter(permutations(range(natoms), 2), dtype=np.dtype(int, 2))
+        nlist = dpnblist[dpnblist[:, 0] < dpnblist[:, 1]]
         if self.capacity_multiplier is None:
             self.capacity_multiplier = int(nlist.shape[0] * 1.3)
 
@@ -135,8 +202,8 @@ def __init__(self, rcut, cov_map, padding=True):
     def allocate(self, coords):
         self._positions = coords  # cache it
         natoms = coords.shape[0]
-        nblist = np.fromiter(permutations(range(natoms), 2), dtype=np.dtype(int, 2))
-        nlist = nblist[nblist[:, 0] < nblist[:, 1]]
+        dpnblist = np.fromiter(permutations(range(natoms), 2), dtype=np.dtype(int, 2))
+        nlist = dpnblist[dpnblist[:, 0] < dpnblist[:, 1]]
         distances = np.linalg.norm(coords[nlist[:, 0]] - coords[nlist[:, 1]], axis=1)
         nlist = nlist[distances < self.rcut]
         if self.capacity_multiplier is None:
@@ -157,4 +224,4 @@ def allocate(self, coords):
             self._pairs = self._do_cov_map(nlist)
             return self._pairs
         else:
-            raise ValueError("padding width < 0")
+            raise ValueError("padding width < 0")

dmff/dpnblist/CMakeLists.txt

@@ -0,0 +1,85 @@
+cmake_minimum_required(VERSION 3.15)
+# project(dpnblist LANGUAGES CUDA CXX)
+project(dpnblist LANGUAGES CXX)
+
+# Set C++ standard
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+# Find CUDA
+find_package(CUDA QUIET)
+if (CUDA_FOUND)
+    message(STATUS "CUDA found, compiling CUDA version")
+    enable_language(CUDA)
+    set(CUDA_NVCC_FLAGS 61)
+    set(SOURCES
+        src/pywrap.cpp
+        src/nbList.cpp
+        src/box.cpp
+        src/nbSchAlg.cpp
+        src/SchAlg/hashSchAlgGPU.cu
+        src/SchAlg/hashSchAlgCPU.cpp
+        src/SchAlg/octreeSchAlgCPU.cpp 
+        src/SchAlg/octreeSchAlgGPU.cu 
+        src/SchAlg/cellSchAlgCPU.cpp
+        src/SchAlg/cellSchAlgGPU.cu
+    )
+else()
+    set(SOURCES
+        src/pywrap_CPU.cpp
+        src/nbList_CPU.cpp
+        src/box.cpp
+        src/nbSchAlg.cpp
+        src/SchAlg/hashSchAlgCPU.cpp
+        src/SchAlg/octreeSchAlgCPU.cpp 
+        src/SchAlg/cellSchAlgCPU.cpp
+    )
+endif()
+
+# Find pybind11
+find_package(pybind11 QUIET)
+if (NOT pybind11_FOUND)
+    add_subdirectory(${CMAKE_SOURCE_DIR}/external/pybind11-2.11.1)
+endif()
+
+# Add the pybind11 module
+pybind11_add_module(dpnblist SHARED)
+target_sources(dpnblist PRIVATE ${SOURCES})
+
+# Include directories
+target_include_directories(dpnblist PRIVATE
+    ${CMAKE_SOURCE_DIR}/src
+    ${CMAKE_SOURCE_DIR}/src/SchAlg
+)
+
+# Find and link OpenMP
+find_package(OpenMP REQUIRED)
+if(OpenMP_CXX_FOUND)
+    target_link_libraries(dpnblist PUBLIC OpenMP::OpenMP_CXX)
+endif()
+
+# Find and link Python3
+find_package(Python3 REQUIRED COMPONENTS Development)
+if (Python3_FOUND)
+    target_link_libraries(dpnblist PUBLIC Python3::Python)
+endif()
+
+# Add subdirectory for tests
+if(CMAKE_BUILD_TYPE MATCHES Debug)
+    add_library(dpnblist ${SOURCES})
+    target_include_directories(dpnblist PRIVATE
+        ${CMAKE_SOURCE_DIR}/src
+        ${CMAKE_SOURCE_DIR}/src/SchAlg
+    )
+    target_link_libraries(dpnblist PUBLIC OpenMP::OpenMP_CXX Python3::Python)
+    add_subdirectory(tests)
+endif()
+
+# Set position independent code
+set_target_properties(dpnblist PROPERTIES POSITION_INDEPENDENT_CODE ON)
+
+# Set compiler flags
+if (CUDA_FOUND)
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -fopenmp -O3 -arch=sm_${CUDA_NVCC_FLAGS}")
+endif()
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp -O3")

dmff/dpnblist/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Zhu Benhao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

dmff/dpnblist/README.md

@@ -0,0 +1,47 @@
+# dpnblist
+ Resable and Modular Neighborlist Lib
+
+### Installation
+follow the below commands to build the project
+```
+pip install .
+```
+then a dpnblist package will be compiled and installed into the python environment.
+
+### Usage
+In our neighbor list library, we provide a full type of neighbor list and include three different types of algorithms: cell list, octree, and hash methods, and support CPU and GPU.  
+
+1. First, it is necessary to define a simulation system through the `box` class, mainly including two parameters, the `length` and `angle` of the simulation system, as shown in the table below:
+
+    |Box Parameter | Data Type |
+    |---|---|
+    |length  |  list |
+    |amgle   |   list |
+
+2. Then, specify different algorithms and platforms through the `NeighborList` class, with parameters as follows
+
+    |Parameter| Data Type |Value|
+    |---|---|---|
+    |mode|string|"Linked_Cell-CPU"  "Linked_Cell-GPU"  "Octree-CPU"  "Octree-GPU"  "Hash-CPU"  "Hash-GPU"|
+
+3. Finally, call the `build` method to generate the neighbor list, passing the `box`, `points`, and `cutoff` parameters. Then, use the `get_neighbor_pair` method to retrieve the neighbor list.
+
+Here's an example of usage:  
+```python
+import dpnblist
+import numpy as np
+
+dpnblist.set_num_threads(4)
+print(dpnblist.get_max_threads())
+num_particle = 30000
+shape = (num_particle, 3)
+points = np.random.random(shape) * domain_size
+domain_size = [50.0, 50.0, 50.0]
+angle = [90.0, 90.0, 90.0]
+cutoff = 6.0
+
+box = dpnblist.Box(domain_size, angle)
+nb = dpnblist.NeighborList("Linked_Cell-GPU")    # Linked_Cell-GPU  Linked_Cell-CPU  Octree-GPU  Octree-CPU  Hash-GPU  Hash-CPU
+nb.build(box, points, cutoff)
+pairs = nb.get_neighbor_pair()
+```

dmff/dpnblist/examples/nblist_test.py

@@ -0,0 +1,22 @@
+import dpnblist
+import numpy as np
+import time
+import freud
+
+dpnblist.set_num_threads(4)
+print(dpnblist.get_max_threads())
+
+num_particle = 3000
+domain_size = [50.0, 50.0, 50.0]
+angle = [90.0, 90.0, 90.0]
+shape = (num_particle, 3)
+inputs = np.random.random(shape) * domain_size
+cutoff = 6.0
+
+box = dpnblist.Box(domain_size, angle)
+nb = dpnblist.NeighborList("Octree-CPU")    # Linked_Cell-GPU  Linked_Cell-CPU  Octree-GPU  Octree-CPU  Hash-GPU  Hash-CPU
+nb.build(box, inputs, cutoff)
+pairs = nb.get_neighbor_pair()
+lists = nb.get_neighbor_list()
+print(pairs[0:9])
+print(lists[0:3])

dmff/dpnblist/external/doctest-2.4.11/.clang-format

@@ -0,0 +1,82 @@
+# https://clang.llvm.org/docs/ClangFormatStyleOptions.html
+
+DisableFormat: false
+Language: Cpp
+Standard: Cpp11
+IndentWidth: 4
+TabWidth: 4
+UseTab: Never
+ColumnLimit: 100
+ReflowComments: false
+SortIncludes: false
+
+AlignConsecutiveAssignments: true
+AlignConsecutiveDeclarations: true
+AlignEscapedNewlinesLeft: false
+AlignOperands: true
+AlignTrailingComments: true
+AlignAfterOpenBracket: true
+DerivePointerAlignment: false
+PointerAlignment: Left
+IndentCaseLabels: true
+ContinuationIndentWidth: 8
+NamespaceIndentation: Inner
+CompactNamespaces: true
+FixNamespaceComments: true
+AccessModifierOffset: -4
+
+SpaceAfterControlStatementKeyword: false
+SpaceAfterCStyleCast: false
+SpaceBeforeAssignmentOperators: true
+SpaceBeforeParens: Never
+SpaceInEmptyParentheses: false
+SpacesBeforeTrailingComments: 1
+SpacesInAngles: false
+SpacesInCStyleCastParentheses: false
+SpacesInContainerLiterals: false
+SpacesInParentheses: false
+SpacesInSquareBrackets: false
+Cpp11BracedListStyle: true
+
+KeepEmptyLinesAtTheStartOfBlocks: false
+MaxEmptyLinesToKeep: 1
+BinPackArguments: true
+BinPackParameters: true
+AlwaysBreakAfterReturnType: None
+AlwaysBreakAfterDefinitionReturnType: None
+AlwaysBreakTemplateDeclarations: true
+BreakConstructorInitializersBeforeComma: true
+ConstructorInitializerAllOnOneLineOrOnePerLine: false
+ConstructorInitializerIndentWidth: 8
+
+AllowShortIfStatementsOnASingleLine: false
+AllowShortLoopsOnASingleLine: false
+AllowShortBlocksOnASingleLine: true
+AllowShortCaseLabelsOnASingleLine: true
+AllowAllParametersOfDeclarationOnNextLine: true
+AllowShortFunctionsOnASingleLine: All
+
+BreakBeforeBinaryOperators: false
+BreakBeforeTernaryOperators: false
+BreakStringLiterals: true
+BreakBeforeBraces: Custom
+BraceWrapping:
+    AfterClass: true
+    AfterEnum: true
+    AfterNamespace: false
+    AfterStruct: true
+    AfterUnion: true
+
+    BeforeCatch: false
+    BeforeElse: false
+    IndentBraces: false
+    AfterFunction: false
+    AfterControlStatement: false
+
+# penalties not thought of yet
+PenaltyBreakBeforeFirstCallParameter: 19
+PenaltyBreakComment: 60
+PenaltyBreakString: 1000
+PenaltyBreakFirstLessLess: 120
+PenaltyExcessCharacter: 1000000
+PenaltyReturnTypeOnItsOwnLine: 1000