Add new search algorithm

CMakeLists.txt

@@ -762,6 +762,7 @@ target_link_libraries(
   planewave
   surchem
   neighbor
+  neighbor_search
   io_input
   io_basic
   io_advanced

source/Makefile.Objects

@@ -16,6 +16,7 @@ VPATH=./src_global:\
 ./source_basis/module_ao:\
 ./source_basis/module_nao:\
 ./source_cell/module_neighbor:\
+./source_cell/module_neighbor_search:\
 ./source_cell/module_symmetry:\
 ./source_cell:\
 ./source_base:\
@@ -86,6 +87,7 @@ ${OBJS_HAMILT}\
 ${OBJS_HSOLVER}\
 ${OBJS_MD}\
 ${OBJS_NEIGHBOR}\
+${OBJS_NEIGHBOR_SEARCH}\
 ${OBJS_PSI}\
 ${OBJS_PSI_INITIALIZER}\
 ${OBJS_PW}\
@@ -411,6 +413,10 @@ OBJS_NEIGHBOR=sltk_atom.o\
     sltk_grid.o\
     sltk_grid_driver.o\
 
+OBJS_NEIGHBOR_SEARCH=neighbor_search.o\
+    bin_manager.o\
+
+
 OBJS_ORBITAL=ORB_atomic.o\
       ORB_atomic_lm.o\
       ORB_gaunt_table.o\

source/source_cell/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_subdirectory(module_symmetry)
 add_subdirectory(module_neighbor)
+add_subdirectory(module_neighbor_search)
 
 add_library(
     cell

source/source_cell/module_neighbor_search/CMakeLists.txt

@@ -0,0 +1,16 @@
+add_library(
+    neighbor_search
+    OBJECT
+    bin_manager.cpp
+    neighbor_search.cpp
+)
+
+if(ENABLE_COVERAGE)
+  add_coverage(neighbor_search)
+endif()
+
+if(BUILD_TESTING)
+  if(ENABLE_MPI)
+    add_subdirectory(test)
+  endif()
+endif()

source/source_cell/module_neighbor_search/bin_manager.cpp

@@ -0,0 +1,191 @@
+#include <limits>
+#include <cmath>
+#include <algorithm>
+#include "bin_manager.h"
+
+
+
+void BinManager::init_bins(
+    double sr,
+    const std::vector<NeighborAtom>& inside_atoms,
+    const std::vector<NeighborAtom>& ghost_atoms
+)
+{
+    sradius = sr;
+
+    x_min = y_min = z_min = std::numeric_limits<double>::max();
+
+    x_max = y_max = z_max = std::numeric_limits<double>::lowest();
+
+    auto update_bounds = [&](const std::vector<NeighborAtom>& atoms)
+    {
+        for (const auto& atom : atoms)
+        {
+            x_min = std::min(x_min, atom.position_x);
+            x_max = std::max(x_max, atom.position_x);
+
+            y_min = std::min(y_min, atom.position_y);
+            y_max = std::max(y_max, atom.position_y);
+
+            z_min = std::min(z_min, atom.position_z);
+            z_max = std::max(z_max, atom.position_z);
+        }
+    };
+
+    update_bounds(inside_atoms);
+    update_bounds(ghost_atoms);
+
+    bin_sizex = bin_sizey = bin_sizez = sradius;
+
+    nbinx = std::ceil((x_max - x_min) / bin_sizex);
+    nbiny = std::ceil((y_max - y_min) / bin_sizey);
+    nbinz = std::ceil((z_max - z_min) / bin_sizez);
+
+    nbinx = std::max(1, nbinx);
+    nbiny = std::max(1, nbiny);
+    nbinz = std::max(1, nbinz);
+
+    int nbins = nbinx * nbiny * nbinz;
+
+    bins.clear();
+
+    bins.resize(nbins);
+
+    for (int ix = 0; ix < nbinx; ++ix)
+    {
+        for (int iy = 0; iy < nbiny; ++iy)
+        {
+            for (int iz = 0; iz < nbinz; ++iz)
+            {
+                int idx = ix * nbiny * nbinz + iy * nbinz + iz;
+
+                bins[idx].id_x = ix;
+                bins[idx].id_y = iy;
+                bins[idx].id_z = iz;
+
+                bins[idx].atoms.clear();
+            }
+        }
+    }
+}
+
+void BinManager::do_binning(
+    const std::vector<NeighborAtom>& inside_atoms,
+    const std::vector<NeighborAtom>& ghost_atoms
+)
+{
+    auto bin_atom = [&](const NeighborAtom& atom)
+    {
+        int ix = std::min(
+            std::max(int((atom.position_x - x_min) / bin_sizex), 0),
+            nbinx - 1
+        );
+
+        int iy = std::min(
+            std::max(int((atom.position_y - y_min) / bin_sizey), 0),
+            nbiny - 1
+        );
+
+        int iz = std::min(
+            std::max(int((atom.position_z - z_min) / bin_sizez), 0),
+            nbinz - 1
+        );
+
+        int idx = ix * nbiny * nbinz + iy * nbinz + iz;
+
+        bins[idx].atoms.push_back(atom);
+    };
+
+    for (const auto& atom : inside_atoms) bin_atom(atom);
+
+    for (const auto& atom : ghost_atoms) bin_atom(atom);
+}
+
+void BinManager::build_atom_neighbors(
+    NeighborList& neighbor_list,
+    std::vector<NeighborAtom>& atoms
+)
+{
+    assert(atoms.size() == neighbor_list.numneigh.size());
+
+    double sradius2 = sradius * sradius;
+
+    neighbor_list.reset();
+
+    for (int i = 0; i < atoms.size(); i++)
+    {
+        std::vector<int> neigh_tmp;
+
+        int ix = std::min(
+            std::max(int((atoms[i].position_x - x_min) / bin_sizex), 0),
+            nbinx - 1
+        );
+
+        int iy = std::min(
+            std::max(int((atoms[i].position_y - y_min) / bin_sizey), 0),
+            nbiny - 1
+        );
+
+        int iz = std::min(
+            std::max(int((atoms[i].position_z - z_min) / bin_sizez), 0),
+            nbinz - 1
+        );
+
+        for (int dx = -1; dx <= 1; dx++)
+        {
+            for (int dy = -1; dy <= 1; dy++)
+            {
+                for (int dz = -1; dz <= 1; dz++)
+                {
+                    int jx = ix + dx;
+                    int jy = iy + dy;
+                    int jz = iz + dz;
+
+                    if (jx < 0 || jx >= nbinx ||
+                        jy < 0 || jy >= nbiny ||
+                        jz < 0 || jz >= nbinz)
+                        continue;
+
+                    int nidx = jx * nbiny * nbinz + jy * nbinz + jz;
+
+                    for (auto natom : bins[nidx].atoms)
+                    {
+                        double dx = atoms[i].position_x - natom.position_x;
+                        double dy = atoms[i].position_y - natom.position_y;
+                        double dz = atoms[i].position_z - natom.position_z;
+
+                        double dist2 = dx * dx + dy * dy + dz * dz;
+
+                        if (dist2 <= sradius2 && dist2 != 0)
+                        {
+                            neigh_tmp.push_back(natom.atom_id);
+                        }
+                    }
+                }
+            }
+        } 
+        int n = neigh_tmp.size();
+
+        //std::cout<<n<<std::endl;
+
+        int* ptr = neighbor_list.allocator.allocate(n);
+
+        for (int k = 0; k < n; k++)
+        {
+            ptr[k] = neigh_tmp[k];
+        }
+
+        neighbor_list.firstneigh[i] = ptr;
+        neighbor_list.numneigh[i] = n;
+    }
+}
+
+void BinManager::clear()
+{
+    for (auto& bin : bins)
+    {
+        bin.atoms.clear();
+    }
+
+    bins.clear();
+}

source/source_cell/module_neighbor_search/bin_manager.h

@@ -0,0 +1,66 @@
+#ifndef BIN_MANAGER_H
+#define BIN_MANAGER_H
+
+#include <vector>
+#include "source_cell/module_neighbor_search/neighbor_atom.h"
+#include "source_cell/module_neighbor_search/neighbor_list.h"
+
+
+
+class Bin
+{
+public:
+    Bin()=default;
+    ~Bin()=default;
+
+    int id_x;
+    int id_y;
+    int id_z;
+
+    std::vector<NeighborAtom> atoms;
+};
+
+
+class BinManager
+{
+public:
+
+    void init_bins(
+        double sr,
+        const std::vector<NeighborAtom>& inside_atoms,
+        const std::vector<NeighborAtom>& ghost_atoms
+    );
+
+
+    void do_binning(
+        const std::vector<NeighborAtom>& inside_atoms,
+        const std::vector<NeighborAtom>& ghost_atoms
+    );
+
+
+    void build_atom_neighbors(
+        NeighborList& neighbor_list,
+        std::vector<NeighborAtom>& atoms
+    );
+
+
+    void clear();
+
+
+    double sradius;
+
+    double x_min, y_min, z_min;
+    double x_max, y_max, z_max;
+
+    double bin_sizex;
+    double bin_sizey;
+    double bin_sizez;
+
+    int nbinx;
+    int nbiny;
+    int nbinz;
+
+    std::vector<Bin> bins;
+};
+
+#endif // BIN_MANAGER_H

source/source_cell/module_neighbor_search/neighbor_atom.h

@@ -0,0 +1,34 @@
+#ifndef NEIGHBOR_ATOM_H
+#define NEIGHBOR_ATOM_H
+
+#include <vector>
+
+class NeighborAtom
+{
+public:
+    double position_x;
+    double position_y;
+    double position_z;
+    int atom_type;
+    int atom_index;
+    int atom_id;
+    //bool isghost;
+    bool is_inside;
+
+    NeighborAtom(double x, double y, double z, int type, int index, int id)
+        : position_x(x), position_y(y), position_z(z),
+          atom_type(type), atom_index(index), atom_id(id) {}
+};
+
+class InputAtoms
+{
+public:
+    std::vector<NeighborAtom> InputAtom;
+    double x_low, x_high, y_low, y_high, z_low, z_high;
+    int n_atoms;
+
+    InputAtoms()
+        : x_low(0), x_high(0), y_low(0), y_high(0), z_low(0), z_high(0), n_atoms(0) {}
+};
+
+#endif // NEIGHBOR_ATOM_H