MPI norm calculation for multiple processes

src/orca-jedi/geometry/Geometry.cc

@@ -47,14 +47,14 @@ Geometry::Geometry(const eckit::Configuration & config,
                       grid_(config.getString("grid name"))
 {
     params_.validateAndDeserialize(config);
-    int64_t halo = params_.sourceMeshHalo.value().value_or(0);
+    int64_t halo = params_.sourceMeshHalo.value();
     auto meshgen_config = grid_.meshgenerator()
                           | atlas::option::halo(halo);
 
     atlas::MeshGenerator meshgen(meshgen_config);
     auto partitioner_config = grid_.partitioner();
     partitioner_config.set("type",
-        params_.partitioner.value().value_or("serial"));
+        params_.partitioner.value());
     partitioner_ = atlas::grid::Partitioner(partitioner_config);
     mesh_ = meshgen.generate(grid_, partitioner_);
     funcSpace_ = atlas::functionspace::NodeColumns(

src/orca-jedi/geometry/Geometry.h

@@ -64,7 +64,7 @@ class Geometry : public util::Printable {
     std::string variable_type) const;
   bool levelsAreTopDown() const {return true;}
   std::string distributionType() const {
-      return params_.partitioner.value().value_or("serial");}
+      return params_.partitioner.value();}
 
  private:
   void print(std::ostream &) const;

src/orca-jedi/geometry/GeometryParameters.h

@@ -36,8 +36,17 @@ class OrcaGeometryParameters : public oops::Parameters {
   oops::RequiredParameter<std::string> gridName
     {"grid name", this};
   oops::RequiredParameter<int> nLevels {"number levels", this};
-  oops::OptionalParameter<int> sourceMeshHalo {"source mesh halo", this};
-  oops::OptionalParameter<std::string> partitioner {"partitioner", this};
+  oops::Parameter<int> sourceMeshHalo {"source mesh halo",
+    "Size of the MPI halo when using a domain-distributed geometry."
+      " The default is 0 (no MPI halo)",
+    0,
+    this};
+  oops::Parameter<std::string> partitioner {
+    "partitioner",
+      "Name of the atlas partitioner to use to MPI distribute the model data"
+        " The default will not distribute the data ('serial').",
+      "serial",
+      this};
 };
 
 }  //  namespace orcamodel

src/orca-jedi/state/State.cc

@@ -235,9 +235,11 @@ double State::norm(const std::string & field_name) const {
     bool has_mv = static_cast<bool>(mv);
     double squares_TP = 0;
     size_t valid_points_TP = 0;
-    atlas_omp_for(atlas::idx_t j = 0; j < field_view.shape(0); ++j) {
+    atlas::idx_t num_h_locs = field_view.shape(0);
+    atlas::idx_t num_levels = field_view.shape(1);
+    atlas_omp_for(atlas::idx_t j = 0; j < num_h_locs; ++j) {
       if (!ghost(j)) {
-        for (atlas::idx_t k = 0; k < field_view.shape(1); ++k) {
+        for (atlas::idx_t k = 0; k < num_levels; ++k) {
           double pointValue = field_view(j, k);
           if (!has_mv || (has_mv && !mv(pointValue))) {
             squares_TP += pointValue*pointValue;
@@ -251,12 +253,28 @@ double State::norm(const std::string & field_name) const {
         valid_points += valid_points_TP;
     }
   }
-  // prevent divide by zero when there are no valid model points on the
-  // partition
-  if (!valid_points)
-    return 0;
 
-  return sqrt(squares)/valid_points;
+  // serial distributions have the entire model grid on each MPI rank
+  if (geom_->distributionType() == "serial") {
+    double local_norm = 0;
+    // prevent divide by zero when there are no valid model points on this
+    // MPI rank
+    if (valid_points) {
+      local_norm = sqrt(squares)/valid_points;
+    }
+    return local_norm;
+  }
+
+
+  // Accumulate values across MPI ranks.
+  geom_->getComm().allReduceInPlace(squares, eckit::mpi::sum());
+  geom_->getComm().allReduceInPlace(valid_points, eckit::mpi::sum());
+
+  if (valid_points) {
+    return sqrt(squares)/valid_points;
+  }
+
+  return 0;
 }
 
 }  // namespace orcamodel

src/tests/CMakeLists.txt

@@ -28,9 +28,19 @@ ecbuild_add_test( TARGET test_orcamodel_hofx_sst
                   ARGS testinput/hofx_nc_sst.yaml
                   COMMAND orcamodel_hofx.x )
 
-ecbuild_add_test( TARGET test_orcamodel_hofx_ssh_parallel
+ecbuild_add_test( TARGET test_orcamodel_hofx_ssh_parallel_checkerboard
                   OMP 1
                   MPI 2
+                  ARGS testinput/hofx_nc_ssh_checkerboard.yaml
+                  COMMAND orcamodel_hofx.x )
+
+ecbuild_add_test( TARGET test_orcamodel_hofx_ssh_parallel_serial
+                  OMP 1
+                  MPI 2
+                  ARGS testinput/hofx_nc_ssh.yaml
+                  COMMAND orcamodel_hofx.x )
+
+ecbuild_add_test( TARGET test_orcamodel_hofx_ssh
                   ARGS testinput/hofx_nc_ssh.yaml
                   COMMAND orcamodel_hofx.x )