diff --git a/source/source_cell/module_neighlist/CMakeLists.txt b/source/source_cell/module_neighlist/CMakeLists.txt index 119eba4dfd4..0c5d9bebffa 100644 --- a/source/source_cell/module_neighlist/CMakeLists.txt +++ b/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 ) diff --git a/source/source_cell/module_neighlist/neighbor_list.cpp b/source/source_cell/module_neighlist/neighbor_list.cpp new file mode 100644 index 00000000000..920a32fa6d4 --- /dev/null +++ b/source/source_cell/module_neighlist/neighbor_list.cpp @@ -0,0 +1,31 @@ +#include "neighbor_list.h" + +#include +#include + +std::vector NeighborList::get_neighbors_sorted_by_distance(const int central_atom, + const std::vector& coordinates) const +{ + const int neighbor_count = get_numneigh(central_atom); + std::vector > distance_neighbors; + distance_neighbors.reserve(static_cast(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 sorted_neighbors; + sorted_neighbors.reserve(static_cast(neighbor_count)); + for (const std::pair& distance_neighbor : distance_neighbors) + { + sorted_neighbors.push_back(distance_neighbor.second); + } + return sorted_neighbors; +} diff --git a/source/source_cell/module_neighlist/neighbor_list.h b/source/source_cell/module_neighlist/neighbor_list.h index 017d23fb328..c460deb11c3 100644 --- a/source/source_cell/module_neighlist/neighbor_list.h +++ b/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 get_neighbors_sorted_by_distance(int central_atom, + const std::vector& coordinates) const; PageAllocator& get_allocator() { return allocator_; } const PageAllocator& get_allocator() const { return allocator_; } diff --git a/source/source_cell/module_neighlist/test/CMakeLists.txt b/source/source_cell/module_neighlist/test/CMakeLists.txt index 80270b213d3..d2f6c85845a 100644 --- a/source/source_cell/module_neighlist/test/CMakeLists.txt +++ b/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 ) diff --git a/source/source_cell/module_neighlist/test/test_neighbor_list.cpp b/source/source_cell/module_neighlist/test/test_neighbor_list.cpp index d40d5e03c21..2e949859296 100644 --- a/source/source_cell/module_neighlist/test/test_neighbor_list.cpp +++ b/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 neighbors = {3, 1, 4, 2}; + nl.set_neighbors(0, neighbors); + + const std::vector 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 sorted = nl.get_neighbors_sorted_by_distance(0, coordinates); + + const std::vector expected = {2, 4, 1, 3}; + EXPECT_EQ(sorted, expected); + EXPECT_EQ(nl.get_firstneigh(0)[0], 3); +} diff --git a/source/source_esolver/esolver_dp.cpp b/source/source_esolver/esolver_dp.cpp index b5a84727b9c..3de1522e053 100644 --- a/source/source_esolver/esolver_dp.cpp +++ b/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 #include #include #include @@ -131,11 +130,18 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep) std::vector ilist(static_cast(nowned_atoms), 0); std::vector numneigh(static_cast(nowned_atoms), 0); std::vector firstneigh(static_cast(nowned_atoms), NULL); + std::vector > sorted_neighbors(static_cast(nowned_atoms)); for (int iat = 0; iat < nowned_atoms; ++iat) { ilist[static_cast(iat)] = iat; - numneigh[static_cast(iat)] = neighbor_list.get_numneigh(iat); - firstneigh[static_cast(iat)] = const_cast(neighbor_list.get_firstneigh(iat)); + sorted_neighbors[static_cast(iat)] + = neighbor_list.get_neighbors_sorted_by_distance(iat, coord); + numneigh[static_cast(iat)] + = static_cast(sorted_neighbors[static_cast(iat)].size()); + firstneigh[static_cast(iat)] + = sorted_neighbors[static_cast(iat)].empty() + ? NULL + : sorted_neighbors[static_cast(iat)].data(); } #ifdef __DPMDC deepmd::hpp::InputNlist nlist(nowned_atoms, diff --git a/source/source_esolver/test/CMakeLists.txt b/source/source_esolver/test/CMakeLists.txt index 80ca8f061b1..2b7a04a122c 100644 --- a/source/source_esolver/test/CMakeLists.txt +++ b/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 )