From 1306324bbaddef74dd82163178f07f060336f308 Mon Sep 17 00:00:00 2001
From: "clara.bayley" <yoctoyotta1024@yahoo.com>
Date: Sun, 17 Nov 2024 20:24:06 +0100
Subject: [PATCH] refactor: added MPI capabilities to all roughpaper programs

---
 roughpaper/scratch/main.cpp  | 16 ++++++++++++++--
 roughpaper/src/main.cpp      | 15 ++++++++++++++-
 roughpaper/src/main_impl.hpp |  6 +++---
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/roughpaper/scratch/main.cpp b/roughpaper/scratch/main.cpp
index 82b197cd..1406d124 100644
--- a/roughpaper/scratch/main.cpp
+++ b/roughpaper/scratch/main.cpp
@@ -25,6 +25,7 @@
 #include <iostream>
 
 #include "./cleotypes_sizes.hpp"
+#include "zarr/dataset.hpp"
 #include "cartesiandomain/cartesianmaps.hpp"
 #include "cartesiandomain/createcartesianmaps.hpp"
 #include "cartesiandomain/null_boundary_conditions.hpp"
@@ -54,7 +55,6 @@
 #include "runcleo/sdmmethods.hpp"
 #include "superdrops/microphysicalprocess.hpp"
 #include "superdrops/motion.hpp"
-#include "zarr/dataset.hpp"
 #include "zarr/fsstore.hpp"
 
 template <typename Store>
@@ -172,7 +172,17 @@ inline auto create_sdm(const Config &config, const Timesteps &tsteps, Dataset<St
 }
 
 int main(int argc, char *argv[]) {
-  print_type_sizes(argc, argv);
+  // print_type_sizes(argc, argv);
+
+  MPI_Init(&argc, &argv);
+
+  int comm_size;
+  MPI_Comm_size(MPI_COMM_WORLD, &comm_size);
+  if (comm_size > 1) {
+    std::cout << "ERROR: The current example is not prepared"
+              << " to be run with more than one MPI process" << std::endl;
+    MPI_Abort(MPI_COMM_WORLD, 1);
+  }
 
   Kokkos::Timer kokkostimer;
 
@@ -210,6 +220,8 @@ int main(int argc, char *argv[]) {
                "Total Program Duration: "
             << ttot << "s \n-------------------------------\n";
 
+  MPI_Finalize();
+
   return 0;
 }
 /* ----------------------------------------------------------------------------------------------*/
diff --git a/roughpaper/src/main.cpp b/roughpaper/src/main.cpp
index baf44938..544f0c97 100644
--- a/roughpaper/src/main.cpp
+++ b/roughpaper/src/main.cpp
@@ -28,6 +28,16 @@ int main(int argc, char *argv[]) {
     throw std::invalid_argument("configuration file(s) not specified");
   }
 
+  MPI_Init(&argc, &argv);
+
+  int comm_size;
+  MPI_Comm_size(MPI_COMM_WORLD, &comm_size);
+  if (comm_size > 1) {
+    std::cout << "ERROR: The current example is not prepared"
+              << " to be run with more than one MPI process" << std::endl;
+    MPI_Abort(MPI_COMM_WORLD, 1);
+  }
+
   Kokkos::Timer kokkostimer;
 
   /* Read input parameters from configuration file(s) */
@@ -40,7 +50,6 @@ int main(int argc, char *argv[]) {
   auto dataset = Dataset(store);
 
   /* Initial conditions for CLEO run */
-  const InitialConditions auto initconds = create_initconds(config);
 
   /* Initialise Kokkos parallel environment */
   Kokkos::initialize(argc, argv);
@@ -48,6 +57,8 @@ int main(int argc, char *argv[]) {
     /* CLEO Super-Droplet Model (excluding coupled dynamics solver) */
     const SDMMethods sdm(create_sdm(config, tsteps, dataset));
 
+    const InitialConditions auto initconds = create_initconds(config, sdm.gbxmaps);
+
     /* Solver of dynamics coupled to CLEO SDM */
     CoupledDynamics auto coupldyn(
         create_coupldyn(config, sdm.gbxmaps, tsteps.get_couplstep(), tsteps.get_t_end()));
@@ -66,5 +77,7 @@ int main(int argc, char *argv[]) {
                "Total Program Duration: "
             << ttot << "s \n-------------------------------\n";
 
+  MPI_Finalize();
+
   return 0;
 }
diff --git a/roughpaper/src/main_impl.hpp b/roughpaper/src/main_impl.hpp
index d06c1f2a..2ee55694 100644
--- a/roughpaper/src/main_impl.hpp
+++ b/roughpaper/src/main_impl.hpp
@@ -30,6 +30,7 @@
 #include <stdexcept>
 #include <string_view>
 
+#include "zarr/dataset.hpp"
 #include "cartesiandomain/add_supers_at_domain_top.hpp"
 #include "cartesiandomain/cartesianmaps.hpp"
 #include "cartesiandomain/cartesianmotion.hpp"
@@ -74,7 +75,6 @@
 #include "superdrops/microphysicalprocess.hpp"
 #include "superdrops/motion.hpp"
 #include "superdrops/terminalvelocity.hpp"
-#include "zarr/dataset.hpp"
 #include "zarr/fsstore.hpp"
 
 inline CoupledDynamics auto create_coupldyn(const Config &config, const CartesianMaps &gbxmaps,
@@ -88,9 +88,9 @@ inline CoupledDynamics auto create_coupldyn(const Config &config, const Cartesia
   return FromFileDynamics(config.get_fromfiledynamics(), couplstep, ndims, nsteps);
 }
 
-inline InitialConditions auto create_initconds(const Config &config) {
+inline InitialConditions auto create_initconds(const Config &config, const CartesianMaps &gbxmaps) {
   // const InitAllSupersFromBinary initsupers(config.get_initsupersfrombinary());
-  const InitSupersFromBinary initsupers(config.get_initsupersfrombinary());
+  const InitSupersFromBinary initsupers(config.get_initsupersfrombinary(), gbxmaps);
   const InitGbxsNull initgbxs(config.get_ngbxs());
 
   return InitConds(initsupers, initgbxs);