Refactor module_neighlist (#7498)

* update module_neighlist

* remove magic numbers

* update

* update

* update

* update, add page_allocator

* update

---------

Co-authored-by: abacus_fixer <mohanchen@pku.eud.cn>

Author: mohanchen

source/source_cell/module_neighlist/CMakeLists.txt

@@ -3,6 +3,8 @@ add_library(
     OBJECT
     bin_manager.cpp
     neighbor_search.cpp
+    page_allocator.cpp
+    unitcell_lite.cpp
 )
 
 if(ENABLE_COVERAGE)

source/source_cell/module_neighlist/atom_provider.h

@@ -0,0 +1,74 @@
+#ifndef ATOM_PROVIDER_H
+#define ATOM_PROVIDER_H
+
+#include "source_base/vector3.h"
+#include "source_base/matrix3.h"
+
+/**
+ * @brief Interface for providing atom and lattice information.
+ *
+ * This abstract interface defines the minimum set of methods needed by
+ * the neighbor search module to access atom positions and lattice parameters.
+ * Any class implementing this interface can be used with NeighborSearch.
+ *
+ * @see UnitCell
+ * @see UnitCellLite
+ */
+class AtomProvider
+{
+public:
+    /**
+     * @brief Default destructor.
+     */
+    virtual ~AtomProvider() = default;
+
+    /**
+     * @brief Get the lattice constant.
+     * @return Lattice constant in Bohr.
+     */
+    virtual double get_lat0() const = 0;
+
+    /**
+     * @brief Get the volume of the unit cell.
+     * @return Unit cell volume in Bohr^3.
+     */
+    virtual double get_omega() const = 0;
+
+    /**
+     * @brief Get the lattice vectors.
+     * @return Const reference to the 3x3 lattice vector matrix.
+     */
+    virtual const ModuleBase::Matrix3& get_latvec() const = 0;
+
+    /**
+     * @brief Get the total number of atoms.
+     * @return Total atom count.
+     */
+    virtual int get_natom() const = 0;
+
+    /**
+     * @brief Get the number of atoms of a specific type.
+     * @param i Type index.
+     * @return Number of atoms of type i.
+     */
+    virtual int get_na(int i) const = 0;
+
+    /**
+     * @brief Get the number of atom types.
+     * @return Number of atom types.
+     */
+    virtual int get_ntype() const = 0;
+
+    /**
+     * @brief Get the Cartesian coordinates of a specific atom.
+     *
+     * Returns the position of the j-th atom of type i.
+     *
+     * @param i Type index.
+     * @param j Atom index within type i.
+     * @return Cartesian position vector.
+     */
+    virtual ModuleBase::Vector3<double> get_tau(int i, int j) const = 0;
+};
+
+#endif // ATOM_PROVIDER_H

source/source_cell/module_neighlist/bin_manager.cpp

@@ -1,79 +1,131 @@
 #include <limits>
 #include <cmath>
 #include <algorithm>
+#include <cassert>
 #include "bin_manager.h"
 
+// ========== Bin class implementation ==========
 
+int Bin::get_id_x() const {
+    return id_x_;
+}
+
+int Bin::get_id_y() const {
+    return id_y_;
+}
+
+int Bin::get_id_z() const {
+    return id_z_;
+}
+
+const std::vector<NeighborAtom>& Bin::get_atoms() const {
+    return atoms_;
+}
+
+void Bin::set_id(int ix, int iy, int iz) {
+    id_x_ = ix;
+    id_y_ = iy;
+    id_z_ = iz;
+}
+
+void Bin::clear_atoms() {
+    atoms_.clear();
+}
+
+void Bin::add_atom(const NeighborAtom& atom) {
+    atoms_.push_back(atom);
+}
+
+// ========== BinManager getter methods ==========
+
+int BinManager::get_nbinx() const {
+    return nbinx_;
+}
+
+int BinManager::get_nbiny() const {
+    return nbiny_;
+}
+
+int BinManager::get_nbinz() const {
+    return nbinz_;
+}
+
+int BinManager::get_total_bins() const {
+    return static_cast<int>(bins_.size());
+}
+
+int BinManager::get_bin_atom_count(int bin_index) const {
+    if (bin_index < 0 || bin_index >= static_cast<int>(bins_.size())) {
+        return 0;
+    }
+    return static_cast<int>(bins_[bin_index].get_atoms().size());
+}
+
+// ========== BinManager main methods ==========
 
 void BinManager::init_bins(
     double sr,
     const std::vector<NeighborAtom>& inside_atoms,
     const std::vector<NeighborAtom>& ghost_atoms
 )
 {
-    sradius = sr;
+    sradius_ = sr;
     if(inside_atoms.empty() && ghost_atoms.empty())
     {
-        x_min=y_min=z_min=0;
-        x_max=y_max=z_max=0;
-        nbinx=nbiny=nbinz=1;
-        bins.clear();
-        bins.resize(1);
+        x_min_ = y_min_ = z_min_ = 0;
+        x_max_ = y_max_ = z_max_ = 0;
+        nbinx_ = nbiny_ = nbinz_ = 1;
+        bins_.clear();
+        bins_.resize(1);
         return;
     }
 
-    x_min = y_min = z_min = std::numeric_limits<double>::max();
-
-    x_max = y_max = z_max = std::numeric_limits<double>::lowest();
-    
+    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);
+            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);
+            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);
+            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;
+    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::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);
+    nbinx_ = std::max(1, nbinx_);
+    nbiny_ = std::max(1, nbiny_);
+    nbinz_ = std::max(1, nbinz_);
 
-    int nbins = nbinx * nbiny * nbinz;
+    int nbins = nbinx_ * nbiny_ * nbinz_;
 
-    bins.clear();
+    bins_.clear();
+    bins_.resize(nbins);
 
-    bins.resize(nbins);
-
-    for (int ix = 0; ix < nbinx; ++ix)
+    for (int ix = 0; ix < nbinx_; ++ix)
     {
-        for (int iy = 0; iy < nbiny; ++iy)
+        for (int iy = 0; iy < nbiny_; ++iy)
         {
-            for (int iz = 0; iz < nbinz; ++iz)
+            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;
+                int idx = bin_index(ix, iy, iz);
 
-                bins[idx].atoms.clear();
+                bins_[idx].set_id(ix, iy, iz);
+                bins_[idx].clear_atoms();
             }
         }
     }
@@ -87,58 +139,63 @@ void BinManager::do_binning(
     auto bin_atom = [&](const NeighborAtom& atom)
     {
         int ix = std::min(
-            std::max(int((atom.position_x - x_min) / bin_sizex), 0),
-            nbinx - 1
+            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
+            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
+            std::max(int((atom.position_z - z_min_) / bin_sizez_), 0),
+            nbinz_ - 1
         );
 
-        int idx = ix * nbiny * nbinz + iy * nbinz + iz;
+        int idx = bin_index(ix, iy, iz);
 
-        bins[idx].atoms.push_back(atom);
+        bins_[idx].add_atom(atom);
     };
 
     for (const auto& atom : inside_atoms) bin_atom(atom);
-
     for (const auto& atom : ghost_atoms) bin_atom(atom);
 }
 
+int BinManager::bin_index(int ix, int iy, int iz) const {
+    return ix * nbiny_ * nbinz_ + iy * nbinz_ + iz;
+}
+
 void BinManager::build_atom_neighbors(
     NeighborList& neighbor_list,
     std::vector<NeighborAtom>& atoms
 )
 {
-    assert(atoms.size() == neighbor_list.numneigh.size());
+    assert(atoms.size() == static_cast<size_t>(neighbor_list.get_nlocal()));
 
-    double sradius2 = sradius * sradius;
+    double sradius2 = sradius_ * sradius_;
 
     neighbor_list.reset();
 
+    std::vector<int> neigh_tmp;
+
     for (int i = 0; i < atoms.size(); i++)
     {
-        std::vector<int> neigh_tmp;
+        neigh_tmp.clear();
 
         int ix = std::min(
-            std::max(int((atoms[i].position_x - x_min) / bin_sizex), 0),
-            nbinx - 1
+            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
+            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
+            std::max(int((atoms[i].position_z - z_min_) / bin_sizez_), 0),
+            nbinz_ - 1
         );
 
         for (int dx = -1; dx <= 1; dx++)
@@ -151,14 +208,14 @@ void BinManager::build_atom_neighbors(
                     int jy = iy + dy;
                     int jz = iz + dz;
 
-                    if (jx < 0 || jx >= nbinx ||
-                        jy < 0 || jy >= nbiny ||
-                        jz < 0 || jz >= nbinz)
+                    if (jx < 0 || jx >= nbinx_ ||
+                        jy < 0 || jy >= nbiny_ ||
+                        jz < 0 || jz >= nbinz_)
                         continue;
 
-                    int nidx = jx * nbiny * nbinz + jy * nbinz + jz;
+                    int nidx = bin_index(jx, jy, jz);
 
-                    for (const NeighborAtom& natom : bins[nidx].atoms)
+                    for (const NeighborAtom& natom : bins_[nidx].get_atoms())
                     {
                         double dx = atoms[i].position_x - natom.position_x;
                         double dy = atoms[i].position_y - natom.position_y;
@@ -173,30 +230,29 @@ void BinManager::build_atom_neighbors(
                     }
                 }
             }
-        } 
+        }
+
         int n = neigh_tmp.size();
 
-        //std::cout<<n<<std::endl;
+        int* ptr = neighbor_list.allocator_.allocate(n);
 
-        int* ptr = neighbor_list.allocator.allocate(n);
-    
         for (int k = 0; k < n; k++)
         {
             assert(ptr != nullptr);
             ptr[k] = neigh_tmp[k];
         }
 
-        neighbor_list.firstneigh[i] = ptr;
-        neighbor_list.numneigh[i] = n;
+        neighbor_list.firstneigh_[i] = ptr;
+        neighbor_list.numneigh_[i] = n;
     }
 }
 
 void BinManager::clear()
 {
-    for (auto& bin : bins)
+    for (auto& bin : bins_)
     {
-        bin.atoms.clear();
+        bin.clear_atoms();
     }
 
-    bins.clear();
+    bins_.clear();
 }