Reduce usage of MPI callbacks (#4571)

Description of changes:
- make `Observable_stat` fully MPI-parallel instead of relying on MPI callbacks
- expose the `Context` MPI communicator to script interface objects to reduce usage of MPI callbacks

src/core/analysis/statistics.cpp

@@ -73,32 +73,21 @@ double mindist(PartCfg &partCfg, std::vector<int> const &set1,
   return std::sqrt(mindist_sq);
 }
 
-static Utils::Vector3d mpi_particle_momentum_local() {
-  auto const particles = cell_structure.local_particles();
-  auto const momentum =
-      std::accumulate(particles.begin(), particles.end(), Utils::Vector3d{},
-                      [](Utils::Vector3d &m, Particle const &p) {
-                        return m + p.mass() * p.v();
-                      });
-
-  return momentum;
-}
-
-REGISTER_CALLBACK_REDUCTION(mpi_particle_momentum_local,
-                            std::plus<Utils::Vector3d>())
-
 Utils::Vector3d calc_linear_momentum(bool include_particles,
                                      bool include_lbfluid) {
-  Utils::Vector3d linear_momentum{};
+  Utils::Vector3d momentum{};
   if (include_particles) {
-    linear_momentum +=
-        mpi_call(::Communication::Result::reduction,
-                 std::plus<Utils::Vector3d>(), mpi_particle_momentum_local);
+    auto const particles = cell_structure.local_particles();
+    momentum =
+        std::accumulate(particles.begin(), particles.end(), Utils::Vector3d{},
+                        [](Utils::Vector3d &m, Particle const &p) {
+                          return m + p.mass() * p.v();
+                        });
   }
   if (include_lbfluid) {
-    linear_momentum += lb_lbfluid_calc_fluid_momentum();
+    momentum += lb_lbfluid_calc_fluid_momentum();
   }
-  return linear_momentum;
+  return momentum;
 }
 
 Utils::Vector3d center_of_mass(PartCfg &partCfg, int p_type) {

src/core/energy.cpp

@@ -43,7 +43,7 @@
 
 #include <memory>
 
-static std::shared_ptr<Observable_stat> calculate_energy_local() {
+std::shared_ptr<Observable_stat> calculate_energy() {
 
   auto obs_energy_ptr = std::make_shared<Observable_stat>(1);
 
@@ -117,23 +117,19 @@ static std::shared_ptr<Observable_stat> calculate_energy_local() {
   // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks)
 }
 
-REGISTER_CALLBACK_MAIN_RANK(calculate_energy_local)
+REGISTER_CALLBACK_MAIN_RANK(calculate_energy)
 
-std::shared_ptr<Observable_stat> calculate_energy() {
-  return mpi_call(Communication::Result::main_rank, calculate_energy_local);
+double mpi_calculate_potential_energy() {
+  auto const obs = mpi_call(Communication::Result::main_rank, calculate_energy);
+  return obs->accumulate(-obs->kinetic[0]);
 }
 
-double calculate_current_potential_energy_of_system() {
-  auto const obs_energy = calculate_energy();
-  return obs_energy->accumulate(-obs_energy->kinetic[0]);
+double mpi_observable_compute_energy() {
+  auto const obs = mpi_call(Communication::Result::main_rank, calculate_energy);
+  return obs->accumulate(0);
 }
 
-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 particle_short_range_energy_contribution(int pid) {
   double ret = 0.0;
 
   if (cell_structure.get_resort_particles()) {
@@ -158,11 +154,3 @@ double particle_short_range_energy_contribution_local(int pid) {
   }
   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

@@ -34,16 +34,16 @@
 std::shared_ptr<Observable_stat> calculate_energy();
 
 /** Calculate the total energy of the system. */
-double calculate_current_potential_energy_of_system();
+double mpi_calculate_potential_energy();
 
 /** Helper function for @ref Observables::Energy. */
-double observable_compute_energy();
+double mpi_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.
+ * energy contribution for a specific particle.
  *
  * @param pid    Particle id
  * @return Non-bonded energy of the particle.

src/core/forces.cpp

@@ -211,9 +211,13 @@ void force_calc(CellStructure &cell_structure, double time_step, double kT) {
     // There are two global quantities that need to be evaluated:
     // object's surface and object's volume.
     for (int i = 0; i < max_oif_objects; i++) {
-      auto const area_volume = calc_oif_global(i, cell_structure);
-      if (fabs(area_volume[0]) < 1e-100 && fabs(area_volume[1]) < 1e-100)
+      auto const area_volume = boost::mpi::all_reduce(
+          comm_cart, calc_oif_global(i, cell_structure), std::plus<>());
+      auto const oif_part_area = std::abs(area_volume[0]);
+      auto const oif_part_vol = std::abs(area_volume[1]);
+      if (oif_part_area < 1e-100 and oif_part_vol < 1e-100) {
         break;
+      }
       add_oif_global_forces(area_volume, i, cell_structure);
     }
   }

src/core/grid_based_algorithms/lb.cpp

@@ -1316,15 +1316,15 @@ Utils::Vector3d lb_calc_local_momentum_density(Lattice::index_t index,
 }
 
 /** Calculate momentum of the LB fluid.
- *  @param[out] result         Fluid momentum in MD units
  *  @param[in]  lb_parameters  LB parameters
  *  @param[in]  lb_fields      Hydrodynamic fields of the fluid
  *  @param[in]  lb_lattice     The underlying lattice
  */
-void lb_calc_fluid_momentum(double *result, const LB_Parameters &lb_parameters,
-                            const std::vector<LB_FluidNode> &lb_fields,
-                            const Lattice &lb_lattice) {
-  Utils::Vector3d momentum_density{}, momentum{};
+Utils::Vector3d
+mpi_lb_calc_fluid_momentum_local(LB_Parameters const &lb_parameters,
+                                 std::vector<LB_FluidNode> const &lb_fields,
+                                 Lattice const &lb_lattice) {
+  Utils::Vector3d momentum_density{}, momentum{}, result{};
 
   for (int x = 1; x <= lb_lattice.grid[0]; x++) {
     for (int y = 1; y <= lb_lattice.grid[1]; y++) {
@@ -1338,7 +1338,8 @@ void lb_calc_fluid_momentum(double *result, const LB_Parameters &lb_parameters,
   }
 
   momentum *= lb_parameters.agrid / lb_parameters.tau;
-  boost::mpi::reduce(comm_cart, momentum.data(), 3, result, std::plus<>(), 0);
+  boost::mpi::reduce(::comm_cart, momentum, result, std::plus<>(), 0);
+  return result;
 }
 
 void lb_collect_boundary_forces(double *result) {

src/core/grid_based_algorithms/lb.hpp

@@ -250,9 +250,10 @@ void lb_bounce_back(LB_Fluid &lbfluid, const LB_Parameters &lb_parameters,
 
 #endif /* LB_BOUNDARIES */
 
-void lb_calc_fluid_momentum(double *result, const LB_Parameters &lb_parameters,
-                            const std::vector<LB_FluidNode> &lb_fields,
-                            const Lattice &lb_lattice);
+Utils::Vector3d
+mpi_lb_calc_fluid_momentum_local(LB_Parameters const &lb_parameters,
+                                 std::vector<LB_FluidNode> const &lb_fields,
+                                 Lattice const &lb_lattice);
 void lb_collect_boundary_forces(double *result);
 void lb_initialize_fields(std::vector<LB_FluidNode> &fields,
                           LB_Parameters const &lb_parameters,

src/core/grid_based_algorithms/lb_interface.cpp

@@ -1099,23 +1099,18 @@ void lb_lbnode_set_pop(const Utils::Vector3i &ind,
   }
 }
 
-static void mpi_lb_lbfluid_calc_fluid_momentum_local() {
-  lb_calc_fluid_momentum(nullptr, lbpar, lbfields, lblattice);
-}
-
-REGISTER_CALLBACK(mpi_lb_lbfluid_calc_fluid_momentum_local)
-
 Utils::Vector3d lb_lbfluid_calc_fluid_momentum() {
-  Utils::Vector3d fluid_momentum{};
+  Utils::Vector3d momentum{};
   if (lattice_switch == ActiveLB::GPU) {
 #ifdef CUDA
-    lb_calc_fluid_momentum_GPU(fluid_momentum.data());
+    if (::comm_cart.rank() == 0) {
+      lb_calc_fluid_momentum_GPU(momentum.data());
+    }
 #endif
   } else if (lattice_switch == ActiveLB::CPU) {
-    mpi_call(mpi_lb_lbfluid_calc_fluid_momentum_local);
-    lb_calc_fluid_momentum(fluid_momentum.data(), lbpar, lbfields, lblattice);
+    momentum = mpi_lb_calc_fluid_momentum_local(lbpar, lbfields, lblattice);
   }
-  return fluid_momentum;
+  return momentum;
 }
 
 const Utils::Vector3d

src/core/object-in-fluid/oif_global_forces.cpp

@@ -22,7 +22,6 @@
 #include "BoxGeometry.hpp"
 #include "Particle.hpp"
 #include "cell_system/CellStructure.hpp"
-#include "communication.hpp"
 #include "grid.hpp"
 
 #include "bonded_interactions/bonded_interaction_data.hpp"
@@ -32,12 +31,7 @@
 #include <utils/constants.hpp>
 #include <utils/math/triangle_functions.hpp>
 
-#include <boost/mpi/collectives.hpp>
-
-#include <functional>
-
-using Utils::area_triangle;
-using Utils::get_n_triangle;
+int max_oif_objects = 0;
 
 Utils::Vector2d calc_oif_global(int molType, CellStructure &cs) {
   // first-fold-then-the-same approach
@@ -60,20 +54,19 @@ Utils::Vector2d calc_oif_global(int molType, CellStructure &cs) {
       auto const p33 = p11 + box_geo.get_mi_vector(partners[1]->pos(), p11);
 
       // unfolded positions correct
-      auto const VOL_A = area_triangle(p11, p22, p33);
+      auto const VOL_A = Utils::area_triangle(p11, p22, p33);
       partArea += VOL_A;
 
-      auto const VOL_norm = get_n_triangle(p11, p22, p33);
+      auto const VOL_norm = Utils::get_n_triangle(p11, p22, p33);
       auto const VOL_dn = VOL_norm.norm();
       auto const VOL_hz = 1.0 / 3.0 * (p11[2] + p22[2] + p33[2]);
-      VOL_partVol += VOL_A * -1 * VOL_norm[2] / VOL_dn * VOL_hz;
+      VOL_partVol -= VOL_A * VOL_norm[2] / VOL_dn * VOL_hz;
     }
 
     return false;
   });
 
-  auto const area_volume_local = Utils::Vector2d{{partArea, VOL_partVol}};
-  return boost::mpi::all_reduce(comm_cart, area_volume_local, std::plus<>());
+  return {{partArea, VOL_partVol}};
 }
 
 void add_oif_global_forces(Utils::Vector2d const &area_volume, int molType,
@@ -96,8 +89,8 @@ void add_oif_global_forces(Utils::Vector2d const &area_volume, int molType,
 
       // unfolded positions correct
       // starting code from volume force
-      auto const VOL_norm = get_n_triangle(p11, p22, p33).normalize();
-      auto const VOL_A = area_triangle(p11, p22, p33);
+      auto const VOL_norm = Utils::get_n_triangle(p11, p22, p33).normalize();
+      auto const VOL_A = Utils::area_triangle(p11, p22, p33);
       auto const VOL_vv = (VOL_volume - iaparams->V0) / iaparams->V0;
 
       auto const VOL_force =
@@ -127,15 +120,3 @@ void add_oif_global_forces(Utils::Vector2d const &area_volume, int molType,
     return false;
   });
 }
-
-int max_oif_objects = 0;
-
-void mpi_set_max_oif_objects_local(int max_oif_objects) {
-  ::max_oif_objects = max_oif_objects;
-}
-
-REGISTER_CALLBACK(mpi_set_max_oif_objects_local)
-
-void mpi_set_max_oif_objects(int max_oif_objects) {
-  mpi_call_all(mpi_set_max_oif_objects_local, max_oif_objects);
-}

src/core/object-in-fluid/oif_global_forces.hpp

@@ -43,5 +43,4 @@ void add_oif_global_forces(Utils::Vector2d const &area_volume, int molType,
 
 extern int max_oif_objects;
 
-void mpi_set_max_oif_objects(int max_oif_objects);
 #endif

src/core/observables/EnergyObservable.hpp

@@ -32,7 +32,7 @@ class Energy : public Observable {
   std::vector<std::size_t> shape() const override { return {1}; }
   std::vector<double> operator()() const override {
     std::vector<double> res{1};
-    res[0] = observable_compute_energy();
+    res[0] = mpi_observable_compute_energy();
     return res;
   }
 };

src/core/observables/PressureObservable.hpp

@@ -30,7 +30,7 @@ class Pressure : public Observable {
 public:
   std::vector<std::size_t> shape() const override { return {1}; }
   std::vector<double> operator()() const override {
-    auto const ptensor = observable_compute_pressure_tensor();
+    auto const ptensor = mpi_observable_compute_pressure_tensor();
     std::vector<double> res{1};
     res[0] = (ptensor[0] + ptensor[4] + ptensor[8]) / 3.;
     return res;

src/core/observables/PressureTensor.hpp

@@ -30,7 +30,7 @@ class PressureTensor : public Observable {
 public:
   std::vector<std::size_t> shape() const override { return {3, 3}; }
   std::vector<double> operator()() const override {
-    return observable_compute_pressure_tensor().as_vector();
+    return mpi_observable_compute_pressure_tensor().as_vector();
   }
 };
 

src/core/pressure.cpp

@@ -52,7 +52,7 @@
 #include <memory>
 #include <utility>
 
-static std::shared_ptr<Observable_stat> calculate_pressure_local() {
+std::shared_ptr<Observable_stat> calculate_pressure() {
 
   auto obs_pressure_ptr = std::make_shared<Observable_stat>(9);
 
@@ -125,17 +125,14 @@ static std::shared_ptr<Observable_stat> calculate_pressure_local() {
   return obs_pressure_ptr;
 }
 
-REGISTER_CALLBACK_MAIN_RANK(calculate_pressure_local)
-
-std::shared_ptr<Observable_stat> calculate_pressure() {
-  return mpi_call(Communication::Result::main_rank, calculate_pressure_local);
-}
+REGISTER_CALLBACK_MAIN_RANK(calculate_pressure)
 
-Utils::Vector9d observable_compute_pressure_tensor() {
-  auto const obs_pressure = calculate_pressure();
+Utils::Vector9d mpi_observable_compute_pressure_tensor() {
+  auto const obs =
+      mpi_call(Communication::Result::main_rank, calculate_pressure);
   Utils::Vector9d pressure_tensor{};
   for (std::size_t j = 0; j < 9; j++) {
-    pressure_tensor[j] = obs_pressure->accumulate(0, j);
+    pressure_tensor[j] = obs->accumulate(0, j);
   }
   return pressure_tensor;
 }

src/core/pressure.hpp

@@ -35,6 +35,6 @@
 std::shared_ptr<Observable_stat> calculate_pressure();
 
 /** Helper function for @ref Observables::PressureTensor. */
-Utils::Vector9d observable_compute_pressure_tensor();
+Utils::Vector9d mpi_observable_compute_pressure_tensor();
 
 #endif

src/core/reaction_methods/ReactionAlgorithm.cpp

@@ -54,7 +54,7 @@ int ReactionAlgorithm::do_reaction(int reaction_steps) {
   // calculate potential energy; only consider potential energy since we
   // assume that the kinetic part drops out in the process of calculating
   // ensemble averages (kinetic part may be separated and crossed out)
-  auto current_E_pot = calculate_current_potential_energy_of_system();
+  auto current_E_pot = mpi_calculate_potential_energy();
   for (int i = 0; i < reaction_steps; i++) {
     int reaction_id = i_random(static_cast<int>(reactions.size()));
     generic_oneway_reaction(*reactions[reaction_id], current_E_pot);
@@ -295,7 +295,7 @@ void ReactionAlgorithm::generic_oneway_reaction(
 
   auto const E_pot_new = (particle_inside_exclusion_range_touched)
                              ? std::numeric_limits<double>::max()
-                             : calculate_current_potential_energy_of_system();
+                             : mpi_calculate_potential_energy();
 
   auto const bf = calculate_acceptance_probability(
       current_reaction, E_pot_old, E_pot_new, old_particle_numbers);
@@ -592,14 +592,14 @@ bool ReactionAlgorithm::do_global_mc_move_for_particles_of_type(
     return false;
   }
 
-  auto const E_pot_old = calculate_current_potential_energy_of_system();
+  auto const E_pot_old = mpi_calculate_potential_energy();
 
   auto const original_positions = generate_new_particle_positions(
       type, particle_number_of_type_to_be_changed);
 
   auto const E_pot_new = (particle_inside_exclusion_range_touched)
                              ? std::numeric_limits<double>::max()
-                             : calculate_current_potential_energy_of_system();
+                             : mpi_calculate_potential_energy();
 
   auto const beta = 1.0 / kT;