Reduce MpiCallbacks complexity (#4767)

Follow-up to #4614
Follow-up to #3180 (comment)

src/core/communication.hpp

@@ -47,6 +47,7 @@
 #include "MpiCallbacks.hpp"
 
 #include <boost/mpi/communicator.hpp>
+#include <boost/mpi/environment.hpp>
 
 #include <memory>
 #include <utility>
@@ -97,38 +98,6 @@ void mpi_call_all(void (*fp)(Args...), ArgRef &&...args) {
   Communication::mpiCallbacks().call_all(fp, std::forward<ArgRef>(args)...);
 }
 
-/** @brief Call a local function.
- *  @tparam Tag    Any tag type defined in @ref Communication::Result
- *  @tparam R      Return type of the local function
- *  @tparam Args   Local function argument types
- *  @tparam ArgRef Local function argument types
- *  @param tag     Reduction strategy
- *  @param fp      Local function
- *  @param args    Local function arguments
- */
-template <class Tag, class R, class... Args, class... ArgRef>
-auto mpi_call(Tag tag, R (*fp)(Args...), ArgRef &&...args) {
-  return Communication::mpiCallbacks().call(tag, fp,
-                                            std::forward<ArgRef>(args)...);
-}
-
-/** @brief Call a local function.
- *  @tparam Tag    Any tag type defined in @ref Communication::Result
- *  @tparam TagArg Types of arguments to @p Tag
- *  @tparam R      Return type of the local function
- *  @tparam Args   Local function argument types
- *  @tparam ArgRef Local function argument types
- *  @param tag     Reduction strategy
- *  @param tag_arg Arguments to the reduction strategy
- *  @param fp      Local function
- *  @param args    Local function arguments
- */
-template <class Tag, class TagArg, class R, class... Args, class... ArgRef>
-auto mpi_call(Tag tag, TagArg &&tag_arg, R (*fp)(Args...), ArgRef &&...args) {
-  return Communication::mpiCallbacks().call(tag, std::forward<TagArg>(tag_arg),
-                                            fp, std::forward<ArgRef>(args)...);
-}
-
 /** Process requests from head node. Worker nodes main loop. */
 void mpi_loop();
 

src/core/observables/CylindricalDensityProfile.hpp

@@ -23,6 +23,7 @@
 
 #include "BoxGeometry.hpp"
 #include "grid.hpp"
+#include "utils_histogram.hpp"
 
 #include <utils/Histogram.hpp>
 #include <utils/math/coordinate_transformation.hpp>
@@ -53,9 +54,8 @@ class CylindricalDensityProfile : public CylindricalPidProfileObservable {
               transform_params->orientation()));
     }
 
-    std::vector<decltype(local_folded_positions)> global_folded_positions{};
-    boost::mpi::gather(comm, local_folded_positions, global_folded_positions,
-                       0);
+    auto const global_folded_positions =
+        detail::gather(comm, local_folded_positions);
 
     if (comm.rank() != 0) {
       return {};

src/core/observables/CylindricalFluxDensityProfile.hpp

@@ -60,21 +60,15 @@ class CylindricalFluxDensityProfile : public CylindricalPidProfileObservable {
               traits.velocity(p), transform_params->axis(), pos));
     }
 
-    auto const world_size = comm.size();
-    std::vector<decltype(local_folded_positions)> global_folded_positions{};
-    std::vector<decltype(local_velocities)> global_velocities{};
-    global_folded_positions.reserve(world_size);
-    global_velocities.reserve(world_size);
-    boost::mpi::gather(comm, local_folded_positions, global_folded_positions,
-                       0);
-    boost::mpi::gather(comm, local_velocities, global_velocities, 0);
+    auto const [global_folded_positions, global_velocities] =
+        detail::gather(comm, local_folded_positions, local_velocities);
 
     if (comm.rank() != 0) {
       return {};
     }
 
     Utils::CylindricalHistogram<double, 3> histogram(n_bins(), limits());
-    accumulate(histogram, global_folded_positions, global_velocities);
+    detail::accumulate(histogram, global_folded_positions, global_velocities);
     histogram.normalize();
     return histogram.get_histogram();
   }

src/core/observables/CylindricalLBFluxDensityProfileAtParticlePositions.cpp

@@ -61,17 +61,15 @@ CylindricalLBFluxDensityProfileAtParticlePositions::evaluate(
     local_flux_densities.emplace_back(flux_cyl);
   }
 
-  std::vector<decltype(local_folded_positions)> global_folded_positions{};
-  std::vector<decltype(local_flux_densities)> global_flux_densities{};
-  boost::mpi::gather(comm, local_folded_positions, global_folded_positions, 0);
-  boost::mpi::gather(comm, local_flux_densities, global_flux_densities, 0);
+  auto const [global_folded_positions, global_flux_densities] =
+      detail::gather(comm, local_folded_positions, local_flux_densities);
 
   if (comm.rank() != 0) {
     return {};
   }
 
   Utils::CylindricalHistogram<double, 3> histogram(n_bins(), limits());
-  accumulate(histogram, global_folded_positions, global_flux_densities);
-  return normalize_by_bin_size(histogram);
+  detail::accumulate(histogram, global_folded_positions, global_flux_densities);
+  return detail::normalize_by_bin_size(histogram);
 }
 } // namespace Observables

src/core/observables/CylindricalLBVelocityProfile.cpp

@@ -25,9 +25,6 @@
 #include <utils/Histogram.hpp>
 #include <utils/math/coordinate_transformation.hpp>
 
-#include <boost/mpi/collectives/gather.hpp>
-#include <boost/serialization/vector.hpp>
-
 #include <vector>
 
 namespace Observables {
@@ -55,21 +52,16 @@ std::vector<double> CylindricalLBVelocityProfile::operator()(
     }
   }
 
-  auto const world_size = comm.size();
-  std::vector<decltype(sampling_positions)> global_positions{};
-  std::vector<decltype(local_velocities)> global_velocities{};
-  global_positions.reserve(world_size);
-  global_velocities.reserve(world_size);
-  boost::mpi::gather(comm, local_positions, global_positions, 0);
-  boost::mpi::gather(comm, local_velocities, global_velocities, 0);
+  auto const [global_positions, global_velocities] =
+      detail::gather(comm, local_positions, local_velocities);
 
   if (comm.rank() != 0) {
     return {};
   }
 
   Utils::CylindricalHistogram<double, 3> histogram(n_bins(), limits());
-  accumulate(histogram, global_positions, global_velocities);
-  return normalize_by_bin_size(histogram);
+  detail::accumulate(histogram, global_positions, global_velocities);
+  return detail::normalize_by_bin_size(histogram);
 }
 
 } // namespace Observables

src/core/observables/CylindricalLBVelocityProfileAtParticlePositions.cpp

@@ -26,9 +26,6 @@
 #include <utils/Histogram.hpp>
 #include <utils/math/coordinate_transformation.hpp>
 
-#include <boost/mpi/collectives/gather.hpp>
-#include <boost/serialization/vector.hpp>
-
 #include <utility>
 #include <vector>
 
@@ -59,21 +56,16 @@ std::vector<double> CylindricalLBVelocityProfileAtParticlePositions::evaluate(
     local_velocities.emplace_back(vel_cyl);
   }
 
-  auto const world_size = comm.size();
-  std::vector<decltype(local_folded_positions)> global_folded_positions{};
-  std::vector<decltype(local_velocities)> global_velocities{};
-  global_folded_positions.reserve(world_size);
-  global_velocities.reserve(world_size);
-  boost::mpi::gather(comm, local_folded_positions, global_folded_positions, 0);
-  boost::mpi::gather(comm, local_velocities, global_velocities, 0);
+  auto const [global_folded_positions, global_velocities] =
+      detail::gather(comm, local_folded_positions, local_velocities);
 
   if (comm.rank() != 0) {
     return {};
   }
 
   Utils::CylindricalHistogram<double, 3> histogram(n_bins(), limits());
-  accumulate(histogram, global_folded_positions, global_velocities);
-  return normalize_by_bin_size(histogram);
+  detail::accumulate(histogram, global_folded_positions, global_velocities);
+  return detail::normalize_by_bin_size(histogram);
 }
 
 } // namespace Observables

src/core/observables/CylindricalVelocityProfile.hpp

@@ -28,7 +28,6 @@
 #include <utils/math/coordinate_transformation.hpp>
 
 #include <cstddef>
-#include <utility>
 #include <vector>
 
 namespace Observables {
@@ -59,22 +58,16 @@ class CylindricalVelocityProfile : public CylindricalPidProfileObservable {
               traits.velocity(p), transform_params->axis(), pos));
     }
 
-    auto const world_size = comm.size();
-    std::vector<decltype(local_folded_positions)> global_folded_positions{};
-    std::vector<decltype(local_velocities)> global_velocities{};
-    global_folded_positions.reserve(world_size);
-    global_velocities.reserve(world_size);
-    boost::mpi::gather(comm, local_folded_positions, global_folded_positions,
-                       0);
-    boost::mpi::gather(comm, local_velocities, global_velocities, 0);
+    auto const [global_folded_positions, global_velocities] =
+        detail::gather(comm, local_folded_positions, local_velocities);
 
     if (comm.rank() != 0) {
       return {};
     }
 
     Utils::CylindricalHistogram<double, 3> histogram(n_bins(), limits());
-    accumulate(histogram, global_folded_positions, global_velocities);
-    return normalize_by_bin_size(histogram);
+    detail::accumulate(histogram, global_folded_positions, global_velocities);
+    return detail::normalize_by_bin_size(histogram);
   }
 
   std::vector<std::size_t> shape() const override {

src/core/observables/DensityProfile.hpp

@@ -23,6 +23,7 @@
 #include "Particle.hpp"
 #include "PidProfileObservable.hpp"
 #include "grid.hpp"
+#include "utils_histogram.hpp"
 
 #include <utils/Histogram.hpp>
 
@@ -49,11 +50,8 @@ class DensityProfile : public PidProfileObservable {
           folded_position(traits.position(p), box_geo));
     }
 
-    auto const world_size = comm.size();
-    std::vector<decltype(local_folded_positions)> global_folded_positions{};
-    global_folded_positions.reserve(world_size);
-    boost::mpi::gather(comm, local_folded_positions, global_folded_positions,
-                       0);
+    auto const global_folded_positions =
+        detail::gather(comm, local_folded_positions);
 
     if (comm.rank() != 0) {
       return {};

src/core/observables/FluxDensityProfile.hpp

@@ -27,8 +27,6 @@
 
 #include <utils/Histogram.hpp>
 
-#include <boost/range/combine.hpp>
-
 #include <cstddef>
 #include <utility>
 #include <vector>
@@ -60,21 +58,15 @@ class FluxDensityProfile : public PidProfileObservable {
       local_velocities.emplace_back(traits.velocity(p));
     }
 
-    auto const world_size = comm.size();
-    std::vector<decltype(local_folded_positions)> global_folded_positions{};
-    std::vector<decltype(local_velocities)> global_velocities{};
-    global_folded_positions.reserve(world_size);
-    global_velocities.reserve(world_size);
-    boost::mpi::gather(comm, local_folded_positions, global_folded_positions,
-                       0);
-    boost::mpi::gather(comm, local_velocities, global_velocities, 0);
+    auto const [global_folded_positions, global_velocities] =
+        detail::gather(comm, local_folded_positions, local_velocities);
 
     if (comm.rank() != 0) {
       return {};
     }
 
     Utils::Histogram<double, 3> histogram(n_bins(), limits());
-    accumulate(histogram, global_folded_positions, global_velocities);
+    detail::accumulate(histogram, global_folded_positions, global_velocities);
     histogram.normalize();
     return histogram.get_histogram();
   }

src/core/observables/ForceDensityProfile.hpp

@@ -62,21 +62,15 @@ class ForceDensityProfile : public PidProfileObservable {
       local_forces.emplace_back(traits.force(p));
     }
 
-    auto const world_size = comm.size();
-    std::vector<decltype(local_folded_positions)> global_folded_positions{};
-    std::vector<decltype(local_forces)> global_forces{};
-    global_folded_positions.reserve(world_size);
-    global_forces.reserve(world_size);
-    boost::mpi::gather(comm, local_folded_positions, global_folded_positions,
-                       0);
-    boost::mpi::gather(comm, local_forces, global_forces, 0);
+    auto const [global_folded_positions, global_forces] =
+        detail::gather(comm, local_folded_positions, local_forces);
 
     if (comm.rank() != 0) {
       return {};
     }
 
     Utils::Histogram<double, 3> histogram(n_bins(), limits());
-    accumulate(histogram, global_folded_positions, global_forces);
+    detail::accumulate(histogram, global_folded_positions, global_forces);
     histogram.normalize();
     return histogram.get_histogram();
   }

src/core/observables/LBVelocityProfile.cpp

@@ -24,9 +24,6 @@
 #include <utils/Histogram.hpp>
 #include <utils/Vector.hpp>
 
-#include <boost/mpi/collectives/gather.hpp>
-#include <boost/serialization/vector.hpp>
-
 #include <stdexcept>
 #include <vector>
 
@@ -48,23 +45,18 @@ LBVelocityProfile::operator()(boost::mpi::communicator const &comm) const {
     }
   }
 
-  auto const world_size = comm.size();
-  std::vector<decltype(sampling_positions)> global_positions{};
-  std::vector<decltype(local_velocities)> global_velocities{};
-  global_positions.reserve(world_size);
-  global_velocities.reserve(world_size);
-  boost::mpi::gather(comm, local_positions, global_positions, 0);
-  boost::mpi::gather(comm, local_velocities, global_velocities, 0);
+  auto const [global_positions, global_velocities] =
+      detail::gather(comm, local_positions, local_velocities);
 
   if (comm.rank() != 0) {
     return {};
   }
 
   Utils::Histogram<double, 3> histogram(n_bins(), limits());
-  accumulate(histogram, global_positions, global_velocities);
+  detail::accumulate(histogram, global_positions, global_velocities);
   try {
-    return normalize_by_bin_size(histogram, allow_empty_bins);
-  } catch (empty_bin_exception const &) {
+    return detail::normalize_by_bin_size(histogram, allow_empty_bins);
+  } catch (detail::empty_bin_exception const &) {
     throw std::runtime_error(
         "Decrease sampling delta(s), some bins have no hit");
   }

src/core/observables/PidObservable.hpp

@@ -31,9 +31,9 @@
 
 #include <boost/mpi/collectives/gather.hpp>
 #include <boost/mpi/collectives/reduce.hpp>
-
 #include <boost/range/algorithm/copy.hpp>
 #include <boost/serialization/utility.hpp>
+#include <boost/serialization/vector.hpp>
 
 #include <cassert>
 #include <cstddef>

src/core/observables/utils_histogram.hpp

@@ -21,10 +21,40 @@
 
 #include <utils/Histogram.hpp>
 
+#include <boost/mpi/collectives/gather.hpp>
+#include <boost/mpi/communicator.hpp>
+#include <boost/serialization/vector.hpp>
+
 #include <cstddef>
+#include <utility>
 #include <vector>
 
-namespace Observables {
+namespace Observables::detail {
+
+/** @brief Gather data from all MPI ranks. */
+template <class Pos>
+auto gather(boost::mpi::communicator const &comm,
+            std::vector<Pos> const &local_pos) {
+  std::vector<std::vector<Pos>> global_pos{};
+  global_pos.reserve(comm.size());
+  boost::mpi::gather(comm, local_pos, global_pos, 0);
+  return global_pos;
+}
+
+/** @brief Gather data from all MPI ranks. */
+template <class Pos, class Val>
+auto gather(boost::mpi::communicator const &comm,
+            std::vector<Pos> const &local_pos,
+            std::vector<Val> const &local_val) {
+  auto const world_size = comm.size();
+  std::vector<std::vector<Pos>> global_pos{};
+  global_pos.reserve(world_size);
+  boost::mpi::gather(comm, local_pos, global_pos, 0);
+  std::vector<std::vector<Val>> global_val{};
+  global_val.reserve(world_size);
+  boost::mpi::gather(comm, local_val, global_val, 0);
+  return std::make_pair(global_pos, global_val);
+}
 
 /** @brief Accumulate histogram data gathered from multiple MPI ranks. */
 template <class T, std::size_t N, std::size_t M, class U, class Pos, class Val>
@@ -58,4 +88,4 @@ auto normalize_by_bin_size(Utils::Histogram<T, N, M, U> &histogram,
   return hist_data;
 }
 
-} // Namespace Observables
+} // Namespace Observables::detail