Add short range neighbor methods (#4401)

Fixes #4399, fixes #4438

Description of changes:
- `CellStructure.cpp`: Added the `CellStructure::run_on_particle_short_range_neighbors()` method which runs a kernel function over all particles inside the cell and its neighbors of a given particle
- `cells.hpp`: Added the `mpi_get_short_range_neighbors()` function to execute a parallel search
   - can be called from the Python interface via `system.cell_system.get_neighbors.get_neighbors(p, distance)`
- `energy.hpp`: Added the `compute_non_bonded_pair_energy()` function which returns both the short-range Coulomb interactions plus the non-bonded energy contributions  of two particles
   - can be called from the Python interface via `system.analysis.particle_energy(p)`
- Reaction methods can use the new neighbor search algorithm with constructor argument `search_algorithm="parallel"` (default is `search_algorithm="order_n"`); on 1 MPI rank the original order N algorithm is faster since the parallel algorithm introduces some overhead due to the ghost update (this overhead is negligible with 2 or more MPI ranks)

src/core/cell_system/CellStructure.hpp

@@ -692,6 +692,58 @@ struct CellStructure {
 
     return particle_to_cell(p);
   }
+
+  /**
+   * @brief Run kernel on all particles inside local cell and its neighbors.
+   *
+   * @param p      Particle to find cell for
+   * @param kernel Function with signature <tt>double(Particle const&,
+   *               Particle const&, Utils::Vector3d const&)</tt>
+   * @return false if cell is not found, otherwise true
+   */
+  template <class Kernel>
+  bool run_on_particle_short_range_neighbors(Particle const &p,
+                                             Kernel &kernel) {
+    auto const cell = find_current_cell(p);
+
+    if (cell == nullptr) {
+      return false;
+    }
+
+    auto const maybe_box = decomposition().minimum_image_distance();
+
+    if (maybe_box) {
+      auto const distance_function = detail::MinimalImageDistance{box_geo};
+      short_range_neighbor_loop(p, cell, kernel, distance_function);
+    } else {
+      auto const distance_function = detail::EuclidianDistance{};
+      short_range_neighbor_loop(p, cell, kernel, distance_function);
+    }
+    return true;
+  }
+
+private:
+  template <class Kernel, class DistanceFunc>
+  void short_range_neighbor_loop(Particle const &p1, Cell *const cell,
+                                 Kernel &kernel, DistanceFunc const &df) {
+    /* Iterate over particles inside cell */
+    for (auto const &p2 : cell->particles()) {
+      if (p1.id() != p2.id()) {
+        auto const vec = df(p1, p2).vec21;
+        kernel(p1, p2, vec);
+      }
+    }
+    /* Iterate over all neighbors */
+    for (auto const neighbor : cell->neighbors().all()) {
+      /* Iterate over particles in neighbors */
+      if (neighbor != cell) {
+        for (auto const &p2 : neighbor->particles()) {
+          auto const vec = df(p1, p2).vec21;
+          kernel(p1, p2, vec);
+        }
+      }
+    }
+  }
 };
 
 #endif

src/core/cells.cpp

@@ -33,6 +33,7 @@
 
 #include "Particle.hpp"
 #include "communication.hpp"
+#include "errorhandling.hpp"
 #include "event.hpp"
 #include "grid.hpp"
 #include "integrate.hpp"
@@ -113,8 +114,51 @@ static void search_distance_sanity_check(double const distance) {
                             std::to_string(range));
   }
 }
+static void search_neighbors_sanity_check(double const distance) {
+  search_distance_sanity_check(distance);
+  if (cell_structure.decomposition_type() ==
+      CellStructureType::CELL_STRUCTURE_HYBRID) {
+    throw std::runtime_error("Cannot search for neighbors in the hybrid "
+                             "decomposition cell system");
+  }
+}
 } // namespace detail
 
+boost::optional<std::vector<int>>
+mpi_get_short_range_neighbors_local(int const pid, double const distance,
+                                    bool run_sanity_checks) {
+
+  if (run_sanity_checks) {
+    detail::search_neighbors_sanity_check(distance);
+  }
+  on_observable_calc();
+
+  auto const p = cell_structure.get_local_particle(pid);
+  if (not p or p->is_ghost()) {
+    return {};
+  }
+
+  std::vector<int> ret;
+  auto const cutoff2 = distance * distance;
+  auto kernel = [&ret, cutoff2](Particle const &p1, Particle const &p2,
+                                Utils::Vector3d const &vec) {
+    if (vec.norm2() < cutoff2) {
+      ret.emplace_back(p2.id());
+    }
+  };
+  cell_structure.run_on_particle_short_range_neighbors(*p, kernel);
+  return ret;
+}
+
+REGISTER_CALLBACK_ONE_RANK(mpi_get_short_range_neighbors_local)
+
+std::vector<int> mpi_get_short_range_neighbors(int const pid,
+                                               double const distance) {
+  detail::search_neighbors_sanity_check(distance);
+  return mpi_call(::Communication::Result::one_rank,
+                  mpi_get_short_range_neighbors_local, pid, distance, false);
+}
+
 std::vector<std::pair<int, int>> get_pairs(double const distance) {
   detail::search_distance_sanity_check(distance);
   auto pairs =

src/core/cells.hpp

@@ -56,6 +56,8 @@
 
 #include "Particle.hpp"
 
+#include <boost/optional.hpp>
+
 #include <utility>
 #include <vector>
 
@@ -116,6 +118,15 @@ get_pairs_of_types(double distance, std::vector<int> const &types);
 /** Check if a particle resorting is required. */
 void check_resort_particles();
 
+/**
+ * @brief Get ids of particles that are within a certain distance
+ * of another particle.
+ */
+std::vector<int> mpi_get_short_range_neighbors(int pid, double distance);
+boost::optional<std::vector<int>>
+mpi_get_short_range_neighbors_local(int pid, double distance,
+                                    bool run_sanity_checks);
+
 /**
  * @brief Find the cell in which a particle is stored.
  *

src/core/energy.cpp

@@ -128,3 +128,36 @@ double observable_compute_energy() {
   auto const obs_energy = calculate_energy();
   return obs_energy->accumulate(0);
 }
+
+double particle_short_range_energy_contribution_local(int pid) {
+  double ret = 0.0;
+
+  if (cell_structure.get_resort_particles()) {
+    cells_update_ghosts(global_ghost_flags());
+  }
+
+  auto const p = cell_structure.get_local_particle(pid);
+
+  if (p) {
+    auto kernel = [&ret](Particle const &p, Particle const &p1,
+                         Utils::Vector3d const &vec) {
+#ifdef EXCLUSIONS
+      if (not do_nonbonded(p, p1))
+        return;
+#endif
+      auto const &ia_params = *get_ia_param(p.type(), p1.type());
+      // Add energy for current particle pair to result
+      ret += calc_non_bonded_pair_energy(p, p1, ia_params, vec, vec.norm());
+    };
+    cell_structure.run_on_particle_short_range_neighbors(*p, kernel);
+  }
+  return ret;
+}
+
+REGISTER_CALLBACK_REDUCTION(particle_short_range_energy_contribution_local,
+                            std::plus<double>())
+
+double particle_short_range_energy_contribution(int pid) {
+  return mpi_call(Communication::Result::reduction, std::plus<double>(),
+                  particle_short_range_energy_contribution_local, pid);
+}

src/core/energy.hpp

@@ -42,4 +42,15 @@ double calculate_current_potential_energy_of_system();
 /** Helper function for @ref Observables::Energy. */
 double observable_compute_energy();
 
+/**
+ * @brief Compute short-range energy of a particle.
+ *
+ * Iterates through particles inside cell and neighboring cells and computes
+ * energy contribution for a specificparticle.
+ *
+ * @param pid    Particle id
+ * @return Non-bonded energy of the particle.
+ */
+double particle_short_range_energy_contribution(int pid);
+
 #endif

src/core/reaction_methods/ReactionAlgorithm.cpp

@@ -21,6 +21,7 @@
 
 #include "reaction_methods/ReactionAlgorithm.hpp"
 
+#include "cells.hpp"
 #include "energy.hpp"
 #include "grid.hpp"
 #include "partCfg_global.hpp"
@@ -381,45 +382,46 @@ void ReactionAlgorithm::check_exclusion_range(int inserted_particle_id) {
 
   auto const &inserted_particle = get_particle_data(inserted_particle_id);
 
-  /* Check the excluded radius of the inserted particle */
-
+  /* Check the exclusion radius of the inserted particle */
   if (exclusion_radius_per_type.count(inserted_particle.type()) != 0) {
-    if (exclusion_radius_per_type[inserted_particle.type()] == 0) {
+    if (exclusion_radius_per_type[inserted_particle.type()] == 0.) {
       return;
     }
   }
 
-  auto particle_ids = get_particle_ids();
-  /* remove the inserted particle id*/
-  particle_ids.erase(std::remove(particle_ids.begin(), particle_ids.end(),
-                                 inserted_particle_id),
-                     particle_ids.end());
-
-  /* Check  if the inserted particle within the excluded_range of any other
-   * particle*/
-  double excluded_distance;
-  for (const auto &particle_id : particle_ids) {
-    auto const &already_present_particle = get_particle_data(particle_id);
+  std::vector<int> particle_ids;
+  if (neighbor_search_order_n) {
+    particle_ids = get_particle_ids();
+    /* remove the inserted particle id */
+    particle_ids.erase(std::remove(particle_ids.begin(), particle_ids.end(),
+                                   inserted_particle_id),
+                       particle_ids.end());
+  } else {
+    particle_ids = mpi_get_short_range_neighbors(inserted_particle.identity(),
+                                                 m_max_exclusion_range);
+  }
+
+  /* Check if the inserted particle within the exclusion radius of any other
+   * particle */
+  for (auto const &particle_id : particle_ids) {
+    auto const &p = get_particle_data(particle_id);
+    double excluded_distance;
     if (exclusion_radius_per_type.count(inserted_particle.type()) == 0 ||
-        exclusion_radius_per_type.count(inserted_particle.type()) == 0) {
+        exclusion_radius_per_type.count(p.type()) == 0) {
       excluded_distance = exclusion_range;
-    } else if (exclusion_radius_per_type[already_present_particle.type()] ==
-               0.) {
+    } else if (exclusion_radius_per_type[p.type()] == 0.) {
       continue;
     } else {
-      excluded_distance =
-          exclusion_radius_per_type[inserted_particle.type()] +
-          exclusion_radius_per_type[already_present_particle.type()];
+      excluded_distance = exclusion_radius_per_type[inserted_particle.type()] +
+                          exclusion_radius_per_type[p.type()];
     }
 
     auto const d_min =
-        box_geo
-            .get_mi_vector(already_present_particle.r.p, inserted_particle.r.p)
-            .norm();
+        box_geo.get_mi_vector(p.pos(), inserted_particle.pos()).norm();
 
     if (d_min < excluded_distance) {
       particle_inside_exclusion_range_touched = true;
-      return;
+      break;
     }
   }
 }

src/core/reaction_methods/ReactionAlgorithm.hpp

@@ -50,17 +50,16 @@ class ReactionAlgorithm {
 public:
   ReactionAlgorithm(
       int seed, double kT, double exclusion_range,
-      const std::unordered_map<int, double> &exclusion_radius_per_type)
-      : m_generator(Random::mt19937(std::seed_seq({seed, seed, seed}))),
+      std::unordered_map<int, double> const &exclusion_radius_per_type)
+      : kT{kT}, exclusion_range{exclusion_range},
+        m_generator(Random::mt19937(std::seed_seq({seed, seed, seed}))),
         m_normal_distribution(0.0, 1.0), m_uniform_real_distribution(0.0, 1.0) {
     if (kT < 0.) {
       throw std::domain_error("Invalid value for 'kT'");
     }
     if (exclusion_range < 0.) {
       throw std::domain_error("Invalid value for 'exclusion_range'");
     }
-    this->kT = kT;
-    this->exclusion_range = exclusion_range;
     set_exclusion_radius_per_type(exclusion_radius_per_type);
     update_volume();
   }
@@ -100,6 +99,7 @@ class ReactionAlgorithm {
   void update_volume();
   void
   set_exclusion_radius_per_type(std::unordered_map<int, double> const &map) {
+    auto max_exclusion_range = exclusion_range;
     for (auto const &item : map) {
       auto const type = item.first;
       auto const exclusion_radius = item.second;
@@ -108,8 +108,11 @@ class ReactionAlgorithm {
                                 std::to_string(type) + ": radius " +
                                 std::to_string(exclusion_radius));
       }
+      max_exclusion_range =
+          std::max(max_exclusion_range, 2. * exclusion_radius);
     }
     exclusion_radius_per_type = map;
+    m_max_exclusion_range = max_exclusion_range;
   }
 
   virtual int do_reaction(int reaction_steps);
@@ -134,6 +137,7 @@ class ReactionAlgorithm {
                                                int particle_number_of_type);
 
   bool particle_inside_exclusion_range_touched = false;
+  bool neighbor_search_order_n = true;
 
 protected:
   std::vector<int> m_empty_p_ids_smaller_than_max_seen_particle;
@@ -178,7 +182,6 @@ class ReactionAlgorithm {
   bool
   all_reactant_particles_exist(SingleReaction const &current_reaction) const;
 
-protected:
   virtual double
   calculate_acceptance_probability(SingleReaction const &, double, double,
                                    std::map<int, int> const &) const {
@@ -210,6 +213,7 @@ class ReactionAlgorithm {
   double m_cyl_y = -10.0;
   double m_slab_start_z = -10.0;
   double m_slab_end_z = -10.0;
+  double m_max_exclusion_range = 0.;
 
 protected:
   Utils::Vector3d get_random_position_in_box();

src/python/espressomd/analyze.pxd

@@ -40,6 +40,10 @@ cdef extern from "<array>" namespace "std" nogil:
         array2() except+
         double & operator[](size_t)
 
+cdef extern from "particle_data.hpp":
+    ctypedef struct particle "Particle"
+    const particle & get_particle_data(int pid) except +
+
 cdef extern from "PartCfg.hpp":
     cppclass PartCfg:
         pass
@@ -92,6 +96,7 @@ cdef extern from "pressure.hpp":
 cdef extern from "energy.hpp":
     cdef shared_ptr[Observable_stat] calculate_energy()
     double calculate_current_potential_energy_of_system()
+    double particle_short_range_energy_contribution(int pid)
 
 cdef extern from "dpd.hpp":
     Vector9d dpd_stress()

src/python/espressomd/analyze.pyx

@@ -29,6 +29,7 @@ from .utils cimport Vector3i, Vector3d, Vector9d
 from .utils cimport make_Vector3d
 from .utils cimport make_array_locked
 from .utils cimport create_nparray_from_double_array
+from .particle_data cimport particle, ParticleHandle
 
 
 def autocorrelation(time_series):
@@ -344,6 +345,24 @@ class Analysis:
         utils.handle_errors("calculate_energy() failed")
         return obs
 
+    def particle_energy(self, particle):
+        """
+        Calculate the non-bonded energy of a single given particle.
+
+        Parameters
+        ----------
+        particle : :class:`~espressomd.particle_data.ParticleHandle`
+
+        Returns
+        -------
+        :obj: `float`
+            non-bonded energy of that particle
+
+        """
+        energy_contribution = particle_short_range_energy_contribution(
+            particle.id)
+        return energy_contribution
+
     def calc_re(self, chain_start=None, number_of_chains=None,
                 chain_length=None):
         """