core: Parallel search for reaction methods

espressomd · May 16, 2022 · f5bba5b · f5bba5b
1 parent b60c0d6
commit f5bba5b
Show file tree

Hide file tree

Showing 10 changed files with 118 additions and 54 deletions.
diff --git a/src/core/reaction_methods/ReactionAlgorithm.cpp b/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;
     }
   }
 }

diff --git a/src/core/reaction_methods/ReactionAlgorithm.hpp b/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;
 
 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();

diff --git a/src/python/espressomd/reaction_methods.py b/src/python/espressomd/reaction_methods.py
@@ -72,6 +72,13 @@ class ReactionAlgorithm(ScriptInterfaceHelper):
         Initial counter value (or seed) of the Mersenne Twister RNG.
     exclusion_radius_per_type : :obj:`dict`, optional
          Mapping of particle types to exclusion radii.
+    search_algorithm : :obj:`str`
+        Pair search algorithm. Default is ``"order_n"``, which evaluates the
+        distance between the inserted particle and all other particles in the
+        system, which scales with O(N). For MPI-parallel simulations, the
+        ``"parallel"`` method is faster. The ``"parallel"`` method is not
+        recommended for simulations on 1 MPI rank, since it comes with the
+        overhead of a ghost particle update.
 
     Methods
     -------
@@ -286,7 +293,8 @@ def __init__(self, **kwargs):
             utils.check_valid_keys(self.valid_keys(), kwargs.keys())
 
     def valid_keys(self):
-        return {"kT", "exclusion_range", "seed", "exclusion_radius_per_type"}
+        return {"kT", "exclusion_range", "seed",
+                "exclusion_radius_per_type", "search_algorithm"}
 
     def required_keys(self):
         return {"kT", "exclusion_range", "seed"}
@@ -409,7 +417,7 @@ class ConstantpHEnsemble(ReactionAlgorithm):
 
     def valid_keys(self):
         return {"kT", "exclusion_range", "seed",
-                "constant_pH", "exclusion_radius_per_type"}
+                "constant_pH", "exclusion_radius_per_type", "search_algorithm"}
 
     def required_keys(self):
         return {"kT", "exclusion_range", "seed", "constant_pH"}

diff --git a/src/script_interface/reaction_methods/ConstantpHEnsemble.hpp b/src/script_interface/reaction_methods/ConstantpHEnsemble.hpp
@@ -56,6 +56,9 @@ class ConstantpHEnsemble : public ReactionAlgorithm {
         get_value<double>(params, "constant_pH"),
         get_value_or<std::unordered_map<int, double>>(
             params, "exclusion_radius_per_type", {}));
+    do_set_parameter("search_algorithm",
+                     Variant{get_value_or<std::string>(
+                         params, "search_algorithm", "order_n")});
   }
 
 private:

diff --git a/src/script_interface/reaction_methods/ReactionAlgorithm.hpp b/src/script_interface/reaction_methods/ReactionAlgorithm.hpp
@@ -67,6 +67,24 @@ class ReactionAlgorithm : public AutoParameters<ReactionAlgorithm> {
             return out;
           }},
          {"kT", AutoParameter::read_only, [this]() { return RE()->get_kT(); }},
+         {"search_algorithm",
+          [this](Variant const &v) {
+            auto const key = get_value<std::string>(v);
+            if (key == "order_n") {
+              RE()->neighbor_search_order_n = true;
+            } else if (key == "parallel") {
+              RE()->neighbor_search_order_n = false;
+            } else {
+              throw std::invalid_argument("Unknown search algorithm '" + key +
+                                          "'");
+            }
+          },
+          [this]() {
+            if (RE()->neighbor_search_order_n) {
+              return std::string("order_n");
+            }
+            return std::string("parallel");
+          }},
          {"exclusion_range", AutoParameter::read_only,
           [this]() { return RE()->get_exclusion_range(); }},
          {"exclusion_radius_per_type",

diff --git a/src/script_interface/reaction_methods/ReactionEnsemble.hpp b/src/script_interface/reaction_methods/ReactionEnsemble.hpp
@@ -45,6 +45,9 @@ class ReactionEnsemble : public ReactionAlgorithm {
         get_value<double>(params, "exclusion_range"),
         get_value_or<std::unordered_map<int, double>>(
             params, "exclusion_radius_per_type", {}));
+    do_set_parameter("search_algorithm",
+                     Variant{get_value_or<std::string>(
+                         params, "search_algorithm", "order_n")});
   }
 
 private:

diff --git a/src/script_interface/reaction_methods/WidomInsertion.hpp b/src/script_interface/reaction_methods/WidomInsertion.hpp
@@ -40,6 +40,15 @@ class WidomInsertion : public ReactionAlgorithm {
     return m_re;
   }
 
+  WidomInsertion() {
+    add_parameters({{"search_algorithm",
+                     [](Variant const &) {
+                       throw std::runtime_error(
+                           "No search algorithm for WidomInsertion");
+                     },
+                     []() { return none; }}});
+  }
+
   void do_construct(VariantMap const &params) override {
     m_re = std::make_shared<::ReactionMethods::WidomInsertion>(
         get_value<int>(params, "seed"), get_value<double>(params, "kT"), 0.,

diff --git a/testsuite/python/constant_pH_stats.py b/testsuite/python/constant_pH_stats.py
@@ -23,7 +23,7 @@
 import espressomd.reaction_methods
 
 
-class ReactionEnsembleTest(ut.TestCase):
+class Test(ut.TestCase):
 
     """Test the core implementation of the constant pH reaction ensemble."""
 
@@ -39,20 +39,21 @@ class ReactionEnsembleTest(ut.TestCase):
         types["A-"]: -1,
         types["H+"]: +1,
     }
-    temperature = 1.0
+    temperature = 1.
     # choose target alpha not too far from 0.5 to get good statistics in a
     # small number of steps
     pKa_minus_pH = -0.2
-    pH = 2
+    pH = 2.
     pKa = pKa_minus_pH + pH
-    Ka = 10**(-pKa)
-    box_l = (N0 / c0)**(1.0 / 3.0)
+    Ka = 10.**(-pKa)
+    box_l = np.cbrt(N0 / c0)
     system = espressomd.System(box_l=[box_l, box_l, box_l])
     np.random.seed(69)  # make reaction code fully deterministic
     system.cell_system.skin = 0.4
     system.time_step = 0.01
     RE = espressomd.reaction_methods.ConstantpHEnsemble(
-        kT=1.0, exclusion_range=1, seed=44, constant_pH=pH)
+        kT=1., exclusion_range=1., seed=44, constant_pH=pH,
+        search_algorithm="parallel")
 
     @classmethod
     def setUpClass(cls):
@@ -68,13 +69,13 @@ def setUpClass(cls):
 
     @classmethod
     def ideal_alpha(cls, pH):
-        return 1.0 / (1 + 10**(cls.pKa - pH))
+        return 1. / (1. + 10.**(cls.pKa - pH))
 
     def test_ideal_titration_curve(self):
-        N0 = ReactionEnsembleTest.N0
-        types = ReactionEnsembleTest.types
-        system = ReactionEnsembleTest.system
-        RE = ReactionEnsembleTest.RE
+        RE = self.RE
+        N0 = self.N0
+        types = self.types
+        system = self.system
 
         # Set the hidden particle type to the lowest possible number to speed
         # up the simulation
@@ -100,9 +101,9 @@ def test_ideal_titration_curve(self):
         # note you cannot calculate the pH via -log10(<NH>/volume) in the
         # constant pH ensemble, since the volume is totally arbitrary and does
         # not influence the average number of protons
-        pH = ReactionEnsembleTest.pH
-        pKa = ReactionEnsembleTest.pKa
-        target_alpha = ReactionEnsembleTest.ideal_alpha(pH)
+        pH = self.pH
+        pKa = self.pKa
+        target_alpha = Test.ideal_alpha(pH)
         rel_error_alpha = abs(average_alpha - target_alpha) / target_alpha
         # relative error
         self.assertLess(

diff --git a/testsuite/python/reaction_ensemble.py b/testsuite/python/reaction_ensemble.py
@@ -58,7 +58,7 @@ class ReactionEnsembleTest(ut.TestCase):
     product_types = [types["A-"], types["H+"]]
     product_coefficients = [1, 1]
     nubar = 1
-    system = espressomd.System(box_l=np.ones(3) * (N0 / c0)**(1.0 / 3.0))
+    system = espressomd.System(box_l=np.ones(3) * np.cbrt(N0 / c0))
     np.random.seed(69)  # make reaction code fully deterministic
     system.cell_system.skin = 0.4
     volume = system.volume()
@@ -68,8 +68,8 @@ class ReactionEnsembleTest(ut.TestCase):
     # degree of dissociation alpha = N_A / N_HA = N_H / N_0
     gamma = target_alpha**2 / (1. - target_alpha) * N0 / (volume**nubar)
     RE = espressomd.reaction_methods.ReactionEnsemble(
-        kT=temperature,
-        exclusion_range=exclusion_range, seed=12)
+        seed=12, kT=temperature,
+        exclusion_range=exclusion_range, search_algorithm="parallel")
 
     @classmethod
     def setUpClass(cls):

diff --git a/testsuite/python/reaction_methods_interface.py b/testsuite/python/reaction_methods_interface.py
@@ -34,7 +34,7 @@ def tearDown(self):
         self.system.part.clear()
 
     def check_interface(self, method, kT, exclusion_range,
-                        gamma, exclusion_radius_per_type):
+                        gamma, exclusion_radius_per_type, search_algorithm):
         def check_reaction_parameters(reactions, parameters):
             for reaction, params in zip(reactions, parameters):
                 for key in reaction.required_keys():
@@ -72,6 +72,7 @@ def check_reaction_parameters(reactions, parameters):
             method.exclusion_range,
             exclusion_range,
             delta=1e-10)
+        self.assertEqual(method.search_algorithm, search_algorithm)
         if not isinstance(method, espressomd.reaction_methods.WidomInsertion):
             self.assertEqual(
                 list(method.exclusion_radius_per_type.keys()), [1])
@@ -162,18 +163,28 @@ def test_interface(self):
 
         # reaction ensemble
         method = espressomd.reaction_methods.ReactionEnsemble(
-            kT=1.4, seed=12, **params)
-        self.check_interface(method, kT=1.4, gamma=1.2, **params)
+            kT=1.4, seed=12, search_algorithm="order_n", **params)
+        self.check_interface(
+            method,
+            kT=1.4,
+            gamma=1.2,
+            search_algorithm="order_n",
+            **params)
 
         # constant pH ensemble
         method = espressomd.reaction_methods.ConstantpHEnsemble(
-            kT=1.5, seed=14, constant_pH=10, **params)
-        self.check_interface(method, kT=1.5, gamma=1.2, **params)
+            kT=1.5, seed=14, search_algorithm="parallel", constant_pH=10., **params)
+        self.check_interface(
+            method,
+            kT=1.5,
+            gamma=1.2,
+            search_algorithm="parallel",
+            **params)
 
         # Widom insertion
         method = espressomd.reaction_methods.WidomInsertion(kT=1.6, seed=16)
         self.check_interface(method, kT=1.6, gamma=1., exclusion_range=0.,
-                             exclusion_radius_per_type={})
+                             exclusion_radius_per_type={}, search_algorithm=None)
 
     def test_exceptions(self):
         single_reaction_params = {
@@ -257,6 +268,8 @@ def test_exceptions(self):
             method.delete_particle(p_id=0)
         with self.assertRaisesRegex(RuntimeError, "Trying to remove some non-existing particles from the system via the inverse Widom scheme"):
             widom.calculate_particle_insertion_potential_energy(reaction_id=0)
+        with self.assertRaisesRegex(RuntimeError, "No search algorithm for WidomInsertion"):
+            widom.search_algorithm = "order_n"
 
         # check other exceptions
         with self.assertRaisesRegex(ValueError, "Invalid value for 'volume'"):
@@ -287,6 +300,9 @@ def test_exceptions(self):
         with self.assertRaisesRegex(ValueError, "Invalid excluded_radius value for type 1: radius -0.10"):
             espressomd.reaction_methods.ReactionEnsemble(
                 kT=1., seed=12, exclusion_range=1., exclusion_radius_per_type={1: -0.1})
+        with self.assertRaisesRegex(ValueError, "Unknown search algorithm 'unknown'"):
+            espressomd.reaction_methods.ReactionEnsemble(
+                kT=1., seed=12, exclusion_range=1., search_algorithm="unknown")
         method = espressomd.reaction_methods.ReactionEnsemble(
             kT=1., exclusion_range=1., seed=12, exclusion_radius_per_type={1: 0.1})
         with self.assertRaisesRegex(ValueError, "Invalid excluded_radius value for type 2: radius -0.10"):