fix(esolver): sort DeePMD MD neighbors by distance

source/source_cell/module_neighlist/CMakeLists.txt

@@ -1,5 +1,6 @@
 set(neighbor_search_sources
     bin_manager.cpp
+    neighbor_list.cpp
     neighbor_search.cpp
     page_allocator.cpp
 )

source/source_cell/module_neighlist/neighbor_list.cpp

@@ -0,0 +1,31 @@
+#include "neighbor_list.h"
+
+#include <algorithm>
+#include <utility>
+
+std::vector<int> NeighborList::get_neighbors_sorted_by_distance(const int central_atom,
+                                                                const std::vector<double>& coordinates) const
+{
+    const int neighbor_count = get_numneigh(central_atom);
+    std::vector<std::pair<double, int> > distance_neighbors;
+    distance_neighbors.reserve(static_cast<std::size_t>(neighbor_count));
+    for (int ineighbor = 0; ineighbor < neighbor_count; ++ineighbor)
+    {
+        const int neighbor_atom = get_firstneigh(central_atom)[ineighbor];
+        const double dx = coordinates[3 * central_atom] - coordinates[3 * neighbor_atom];
+        const double dy = coordinates[3 * central_atom + 1] - coordinates[3 * neighbor_atom + 1];
+        const double dz = coordinates[3 * central_atom + 2] - coordinates[3 * neighbor_atom + 2];
+        const double distance2 = dx * dx + dy * dy + dz * dz;
+        distance_neighbors.push_back(std::make_pair(distance2, neighbor_atom));
+    }
+
+    std::sort(distance_neighbors.begin(), distance_neighbors.end());
+
+    std::vector<int> sorted_neighbors;
+    sorted_neighbors.reserve(static_cast<std::size_t>(neighbor_count));
+    for (const std::pair<double, int>& distance_neighbor : distance_neighbors)
+    {
+        sorted_neighbors.push_back(distance_neighbor.second);
+    }
+    return sorted_neighbors;
+}

source/source_cell/module_neighlist/neighbor_list.h

@@ -29,6 +29,9 @@ class NeighborList
     int get_numneigh(int i) const { return numneigh_[i]; }
     int* get_firstneigh(int i) { return firstneigh_[i]; }
     const int* get_firstneigh(int i) const { return firstneigh_[i]; }
+    /// Return a distance-sorted copy of one atom's neighbors without changing the stored order.
+    std::vector<int> get_neighbors_sorted_by_distance(int central_atom,
+                                                      const std::vector<double>& coordinates) const;
     PageAllocator& get_allocator() { return allocator_; }
     const PageAllocator& get_allocator() const { return allocator_; }
 

source/source_cell/module_neighlist/test/CMakeLists.txt

@@ -31,6 +31,7 @@ AddTest(
   LIBS parameter base device
   SOURCES
   test_neighbor_list.cpp
+  ../neighbor_list.cpp
   ../page_allocator.cpp
 )
 

source/source_cell/module_neighlist/test/test_neighbor_list.cpp

@@ -76,3 +76,24 @@ TEST(NeighborList_Getters, Accessors)
     EXPECT_EQ(nl.get_numneigh(0), 0);
     EXPECT_EQ(nl.get_firstneigh(0), nullptr);
 }
+
+TEST(NeighborList_Sort, ByDistanceWithIndexTieBreak)
+{
+    NeighborList nl;
+    nl.initialize(1, 16);
+    const std::vector<int> neighbors = {3, 1, 4, 2};
+    nl.set_neighbors(0, neighbors);
+
+    const std::vector<double> coordinates = {
+        0.0, 0.0, 0.0,
+        2.0, 0.0, 0.0,
+        1.0, 0.0, 0.0,
+        0.0, 3.0, 0.0,
+        -1.0, 0.0, 0.0};
+
+    const std::vector<int> sorted = nl.get_neighbors_sorted_by_distance(0, coordinates);
+
+    const std::vector<int> expected = {2, 4, 1, 3};
+    EXPECT_EQ(sorted, expected);
+    EXPECT_EQ(nl.get_firstneigh(0)[0], 3);
+}

source/source_esolver/esolver_dp.cpp

@@ -26,7 +26,6 @@
 #include "source_io/module_output/output_log.h"
 #include "source_io/module_parameter/parameter.h"
 
-#include <algorithm>
 #include <iomanip>
 #include <array>
 #include <sstream>
@@ -131,11 +130,18 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep)
         std::vector<int> ilist(static_cast<std::size_t>(nowned_atoms), 0);
         std::vector<int> numneigh(static_cast<std::size_t>(nowned_atoms), 0);
         std::vector<int*> firstneigh(static_cast<std::size_t>(nowned_atoms), NULL);
+        std::vector<std::vector<int> > sorted_neighbors(static_cast<std::size_t>(nowned_atoms));
         for (int iat = 0; iat < nowned_atoms; ++iat)
         {
             ilist[static_cast<std::size_t>(iat)] = iat;
-            numneigh[static_cast<std::size_t>(iat)] = neighbor_list.get_numneigh(iat);
-            firstneigh[static_cast<std::size_t>(iat)] = const_cast<int*>(neighbor_list.get_firstneigh(iat));
+            sorted_neighbors[static_cast<std::size_t>(iat)]
+                = neighbor_list.get_neighbors_sorted_by_distance(iat, coord);
+            numneigh[static_cast<std::size_t>(iat)]
+                = static_cast<int>(sorted_neighbors[static_cast<std::size_t>(iat)].size());
+            firstneigh[static_cast<std::size_t>(iat)]
+                = sorted_neighbors[static_cast<std::size_t>(iat)].empty()
+                      ? NULL
+                      : sorted_neighbors[static_cast<std::size_t>(iat)].data();
         }
 #ifdef __DPMDC
         deepmd::hpp::InputNlist nlist(nowned_atoms,

source/source_esolver/test/CMakeLists.txt

@@ -24,5 +24,6 @@ AddTest(
     ../esolver_dp.cpp
     ../../source_cell/basecell.cpp
     ../../source_cell/cif_io.cpp
+    ../../source_cell/module_neighlist/neighbor_list.cpp
     ../../source_io/module_output/output_log.cpp
 )