Read all ORCA field data onto a single processor (#71)

* Add nemo_io::ReadServer and const correctness
* Refactor and remove unnecessary methods
* add Doxygen docs coverage
* cpplinting pass

src/orca-jedi/CMakeLists.txt

@@ -27,6 +27,9 @@ nemo_io/NemoFieldReader.h
 nemo_io/NemoFieldReader.cc
 nemo_io/NemoFieldWriter.h
 nemo_io/NemoFieldWriter.cc
+nemo_io/OrcaIndex.h
+nemo_io/ReadServer.h
+nemo_io/ReadServer.cc
 utilities/OrcaModelTraits.h
 variablechanges/VariableChangeParameters.h
 variablechanges/VariableChange.h

src/orca-jedi/nemo_io/NemoFieldReader.h

@@ -18,9 +18,6 @@
 
 #include "atlas/runtime/Exception.h"
 #include "atlas/util/Point.h"
-#include "atlas/field.h"
-#include "atlas/mesh.h"
-#include "atlas/array/ArrayView.h"
 
 namespace orcamodel {
 
@@ -29,29 +26,17 @@ class NemoFieldReader : private util::ObjectCounter<NemoFieldReader> {
  public:
   static const std::string classname() {return "orcamodel::NemoFieldReader";}
 
-  explicit NemoFieldReader(eckit::PathName& filename);
+  explicit NemoFieldReader(const eckit::PathName& filename);
 
-  std::vector<atlas::PointXY> read_locs();
-  size_t read_dim_size(const std::string& name);
   void read_datetimes();
-  size_t get_nearest_datetime_index(const util::DateTime& datetime);
-  template<class T> T read_fillvalue(const std::string& name);
+  std::vector<atlas::PointXY> read_locs() const;
+  size_t read_dim_size(const std::string& name) const;
+  size_t get_nearest_datetime_index(const util::DateTime& datetime) const;
+  template<typename T> T read_fillvalue(const std::string& name) const;
   std::vector<double> read_var_slice(const std::string& varname,
-      const size_t t_indx, const size_t z_indx);
-  void read_surf_var(const std::string& varname, const size_t t_indx,
-      atlas::array::ArrayView<double, 2>& field_view);
-  void read_volume_var(const std::string& varname, const size_t t_indx,
-      atlas::array::ArrayView<double, 2>& field_view);
-  void read_vertical_var(const std::string& varname,
-      atlas::array::ArrayView<double, 2>& field_view);
-  void read_vertical_var(const std::string& varname, const atlas::Mesh& mesh,
-      atlas::array::ArrayView<double, 2>& field_view);
-
-  void read_surf_var(const std::string& varname, const atlas::Mesh& mesh,
-      const size_t t_indx, atlas::array::ArrayView<double, 2>& field_view);
-  void read_volume_var(const std::string& varname,
-     const atlas::Mesh& mesh, const size_t t_indx,
-     atlas::array::ArrayView<double, 2>& field_view);
+      const size_t t_indx, const size_t z_indx) const;
+  std::vector<double> read_vertical_var(const std::string& varname,
+      const size_t nlevels) const;
 
  private:
   NemoFieldReader() : ncFile() {}

src/orca-jedi/nemo_io/OrcaIndex.h

@@ -0,0 +1,58 @@
+/*
+ * (C) British Crown Copyright 2024 Met Office
+ */
+
+#pragma once
+
+#include <algorithm>
+#include <sstream>
+#include <limits>
+#include <map>
+
+#include "eckit/exception/Exceptions.h"
+
+#include "oops/util/Logger.h"
+
+#include "atlas-orca/grid/OrcaGrid.h"
+
+namespace orcamodel {
+
+/// \brief Indexer into a 1D array of an ORCA model field, match each point to the corresponding
+///        point in the atlas-orca ij space.
+struct OrcaIndex {
+    int32_t ix_glb_max;
+    int32_t iy_glb_max;
+    int32_t glbarray_offset;
+    int32_t glbarray_jstride;
+    int32_t nx_halo_WE;
+    int32_t ny_halo_NS;
+
+    explicit OrcaIndex(const atlas::OrcaGrid& orcaGrid) {
+        iy_glb_max = orcaGrid.ny() + orcaGrid.haloNorth() - 1;
+        ix_glb_max = orcaGrid.nx() + orcaGrid.haloEast() - 1;
+
+        nx_halo_WE = orcaGrid.nx() + orcaGrid.haloEast() + orcaGrid.haloWest();
+        ny_halo_NS = orcaGrid.ny() + orcaGrid.haloNorth()
+          + orcaGrid.haloSouth();
+
+        // vector of local indices: necessary for remote indices of ghost nodes
+        int iy_glb_min = -orcaGrid.haloSouth();
+        int ix_glb_min = -orcaGrid.haloWest();
+        glbarray_offset  = -(nx_halo_WE * iy_glb_min) - ix_glb_min;
+        glbarray_jstride = nx_halo_WE;
+    }
+
+    /// \brief Index of a 1D array corresponding to point i, j
+    /// \param i
+    /// \param j
+    /// \return index of a matching 1D array
+    int64_t operator()(const int i, const int j) const {
+        ATLAS_ASSERT(i <= ix_glb_max,
+            std::to_string(i) + " > " + std::to_string(ix_glb_max));
+        ATLAS_ASSERT(j <= iy_glb_max,
+            std::to_string(j) + " > " + std::to_string(iy_glb_max));
+        return glbarray_offset + j * glbarray_jstride + i;
+    }
+};
+
+}  // namespace orcamodel

src/orca-jedi/nemo_io/ReadServer.cc

@@ -0,0 +1,176 @@
+/*
+ * (C) British Crown Copyright 2024 Met Office
+ */
+
+
+
+#include "orca-jedi/nemo_io/ReadServer.h"
+#include <string>
+#include <vector>
+#include <memory>
+
+#include "eckit/exception/Exceptions.h"
+
+namespace orcamodel {
+
+ReadServer::ReadServer(const eckit::PathName& file_path, const atlas::Mesh& mesh) :
+  mesh_(mesh),
+  index_glbarray_(atlas::OrcaGrid(mesh.grid())) {
+  if (myrank == mpiroot) {
+    reader_ = std::make_unique<NemoFieldReader>(file_path);
+  }
+}
+
+/// \brief Read in a 2D horizontal slice of variable data on the root processor only
+/// \param var_name
+/// \param t_index Index of the time slice.
+/// \param z_index Index of the vertical slice.
+/// \param buffer Vector to store the data.
+void ReadServer::read_var_on_root(const std::string& var_name,
+              const size_t t_index,
+              const size_t z_index,
+              std::vector<double>& buffer) const {
+  size_t size = index_glbarray_.nx_halo_WE * index_glbarray_.ny_halo_NS;
+  if (myrank == mpiroot) {
+    buffer = reader_->read_var_slice(var_name, t_index, z_index);
+  } else {
+    buffer.resize(size);
+  }
+}
+
+/// \brief Read 1D vertical variable data on the root processor only
+/// \param var_name NetCDF name of the vertical variable.
+/// \param n_levels Number of levels to read from the file
+/// \param buffer Vector to store the data.
+void ReadServer::read_vertical_var_on_root(const std::string& var_name,
+              const size_t n_levels,
+              std::vector<double>& buffer) const {
+  size_t size = index_glbarray_.nx_halo_WE * index_glbarray_.ny_halo_NS;
+  if (myrank == mpiroot) {
+    buffer = reader_->read_vertical_var(var_name, n_levels);
+  } else {
+    buffer.resize(size);
+  }
+}
+
+/// \brief Move data from a buffer into an atlas arrayView.
+/// \param buffer Vector of data to read
+/// \param z_index Index of the vertical slice.
+/// \param field_view View into the atlas field to store the data.
+void ReadServer::fill_field(const std::vector<double>& buffer,
+              const size_t z_index,
+      atlas::array::ArrayView<double, 2>& field_view) const {
+    auto ghost = atlas::array::make_view<int32_t, 1>(this->mesh_.nodes().ghost());
+    auto ij = atlas::array::make_view<int32_t, 2>(this->mesh_.nodes().field("ij"));
+    const size_t numNodes = field_view.shape(0);
+    // "ReadServer buffer size does not equal the number of horizontal nodes in the field_view"
+    assert(numNodes <= buffer.size());
+    for (size_t inode = 0; inode < numNodes; ++inode) {
+      if (ghost(inode)) continue;
+      const int64_t ibuf = index_glbarray_(ij(inode, 0), ij(inode, 1));
+      field_view(inode, z_index) = buffer[ibuf];
+    }
+}
+
+/// \brief Move vertical data from a buffer into an atlas arrayView.
+/// \param buffer Vector of data to read
+/// \param field_view View into the atlas field to store the data.
+void ReadServer::fill_vertical_field(const std::vector<double>& buffer,
+      atlas::array::ArrayView<double, 2>& field_view) const {
+    auto ghost = atlas::array::make_view<int32_t, 1>(this->mesh_.nodes().ghost());
+    const size_t num_nodes = field_view.shape(0);
+    const size_t num_levels = field_view.shape(1);
+    // "ReadServer buffer size does not equal the number of levels in the field_view"
+    assert(num_levels <= buffer.size());
+
+    // even for 1D depths, store the data in an atlas 3D field - inefficient but flexible
+    for (size_t inode = 0; inode < num_nodes; ++inode) {
+      for (size_t k = 0; k < num_levels; ++k) {
+        if (ghost(inode)) continue;
+        field_view(inode, k) = buffer[k];
+      }
+    }
+}
+
+/// \brief Read a NetCDF variable into an atlas field.
+/// \param var_name The netCDF name of the variable to read.
+/// \param t_index The time index for the data to read.
+/// \param field_view View into the atlas field to store the data.
+void ReadServer::read_var(const std::string& var_name,
+    const size_t t_index,
+    atlas::array::ArrayView<double, 2>& field_view) {
+
+  size_t n_levels = field_view.shape(1);
+  size_t size = index_glbarray_.nx_halo_WE * index_glbarray_.ny_halo_NS;
+
+  std::vector<double> buffer;
+  // For each level
+  for (size_t iLev = 0; iLev < n_levels; iLev++) {
+    // read the data for that level onto the root processor
+    this->read_var_on_root(var_name, t_index, iLev, buffer);
+
+    // mpi distribute that data out to all processors
+    atlas::mpi::comm().broadcast(&buffer.front(), size, mpiroot);
+
+    // each processor fills out its field_view from the buffer
+    this->fill_field(buffer, iLev, field_view);
+  }
+}
+
+/// \brief Read a vertical variable into an atlas field.
+/// \param var_name The netCDF name of the variable to read.
+/// \param field_view View into the atlas field to store the data.
+void ReadServer::read_vertical_var(const std::string& var_name,
+    atlas::array::ArrayView<double, 2>& field_view) {
+
+  size_t n_levels = field_view.shape(1);
+  size_t size = index_glbarray_.nx_halo_WE * index_glbarray_.ny_halo_NS;
+
+  std::vector<double> buffer;
+
+  // read the data onto the root processor
+  this->read_vertical_var_on_root(var_name, n_levels, buffer);
+
+  // mpi distribute that data out to all processors
+  atlas::mpi::comm().broadcast(&buffer.front(), size, mpiroot);
+
+  // each processor fills out its field_view from the buffer
+  this->fill_vertical_field(buffer, field_view);
+}
+
+/// \brief Find the nearest datetime index to a datetime on the MPI root only.
+/// \param datetime Search for the index of the time slice in the file nearest this datetime.
+/// \return The index of the nearest time slice in the file.
+size_t ReadServer::get_nearest_datetime_index(const util::DateTime& datetime) const {
+  size_t t_index;
+
+  if (myrank == mpiroot) {
+    t_index = reader_->get_nearest_datetime_index(datetime);
+  }
+
+  // mpi distribute that data out to all processors
+  atlas::mpi::comm().broadcast(t_index, mpiroot);
+
+  return t_index;
+}
+
+/// \brief Read the _FillValue for a variable, defaulting to the minimum value
+/// for the datatype. Read on the MPI root process only.
+/// \param name Name of the netCDF variable containing the _FillValue attribute to retrieve.
+/// \return The fill value for this netCDF variable.
+template<class T> T ReadServer::read_fillvalue(const std::string& name) const {
+  T fillvalue;
+
+  if (myrank == mpiroot) {
+    fillvalue = reader_->read_fillvalue<T>(name);
+  }
+
+  // mpi distribute that data out to all processors
+  atlas::mpi::comm().broadcast(fillvalue, mpiroot);
+
+  return fillvalue;
+}
+template int ReadServer::read_fillvalue<int>(const std::string& name) const;
+template float ReadServer::read_fillvalue<float>(const std::string& name) const;
+template double ReadServer::read_fillvalue<double>(const std::string& name) const;
+}  // namespace orcamodel

src/orca-jedi/nemo_io/ReadServer.h

@@ -0,0 +1,59 @@
+/*
+ * (C) British Crown Copyright 2024 Met Office
+ */
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include <memory>
+
+#include "oops/util/DateTime.h"
+#include "atlas/parallel/mpi/mpi.h"
+#include "atlas/array.h"
+#include "atlas/mesh.h"
+#include "atlas/grid.h"
+#include "atlas-orca/grid/OrcaGrid.h"
+
+#include "orca-jedi/nemo_io/OrcaIndex.h"
+#include "orca-jedi/nemo_io/NemoFieldReader.h"
+
+#include "eckit/exception/Exceptions.h"
+
+namespace orcamodel {
+class ReadServer {
+ public:
+  explicit ReadServer(const eckit::PathName& file_path,
+      const atlas::Mesh& mesh);
+  ReadServer(ReadServer &&) = default;
+  ReadServer(const ReadServer &) = default;
+  ReadServer &operator=(ReadServer &&) = default;
+  ReadServer &operator=(const ReadServer &) = default;
+  void read_var(const std::string& var_name,
+      const size_t t_index,
+      atlas::array::ArrayView<double, 2>& field_view);
+  void read_vertical_var(const std::string& var_name,
+      atlas::array::ArrayView<double, 2>& field_view);
+  size_t get_nearest_datetime_index(const util::DateTime& datetime) const;
+  template<class T> T read_fillvalue(const std::string& nemo_var_name) const;
+
+ private:
+  void read_var_on_root(const std::string& var_name,
+      const size_t t_index,
+      const size_t z_index,
+      std::vector<double>& buffer) const;
+  void read_vertical_var_on_root(const std::string& var_name,
+      const size_t n_levels,
+      std::vector<double>& buffer) const;
+  void fill_field(const std::vector<double>& buffer,
+      const size_t z_index,
+      atlas::array::ArrayView<double, 2>& field_view) const;
+  void fill_vertical_field(const std::vector<double>& buffer,
+      atlas::array::ArrayView<double, 2>& field_view) const;
+  const size_t mpiroot = 0;
+  const size_t myrank = atlas::mpi::rank();
+  const atlas::Mesh& mesh_;
+  const OrcaIndex index_glbarray_;
+  std::unique_ptr<NemoFieldReader> reader_;
+};
+}  // namespace orcamodel

src/orca-jedi/state/StateIOUtils.cc

@@ -15,7 +15,7 @@
 
 #include "orca-jedi/geometry/Geometry.h"
 #include "orca-jedi/state/State.h"
-#include "orca-jedi/nemo_io/NemoFieldReader.h"
+#include "orca-jedi/nemo_io/ReadServer.h"
 #include "orca-jedi/nemo_io/NemoFieldWriter.h"
 
 #define NEMO_FILL_TOL 1e-6
@@ -43,12 +43,12 @@ void readFieldsFromFile(
     auto nemo_field_path = eckit::PathName(nemo_file_name);
     oops::Log::debug() << "orcamodel::readFieldsFromFile:: nemo_field_path "
                        << nemo_field_path << std::endl;
-    NemoFieldReader nemo_file(nemo_field_path);
+    ReadServer nemo_reader(nemo_field_path, geom.mesh());
 
     // Read fields from Nemo field file
     // field names in the atlas fieldset are assumed to match their names in
     // the field file
-    const size_t time_indx = nemo_file.get_nearest_datetime_index(valid_date);
+    const size_t time_indx = nemo_reader.get_nearest_datetime_index(valid_date);
     oops::Log::debug() << "orcamodel::readFieldsFromFile:: time_indx "
                        << time_indx << std::endl;
 
@@ -71,15 +71,12 @@ void readFieldsFromFile(
                          << std::endl;
       if (geom.variable_in_variable_type(fieldName, variable_type)) {
         auto field_view = atlas::array::make_view<double, 2>(field);
-        if (varCoordTypeMap[fieldName] == "surface") {
-          nemo_file.read_surf_var(nemoName, geom.mesh(), time_indx, field_view);
-        } else if (varCoordTypeMap[fieldName] == "vertical") {
-          nemo_file.read_vertical_var(nemoName, geom.mesh(), field_view);
+        if (varCoordTypeMap[fieldName] == "vertical") {
+          nemo_reader.read_vertical_var(nemoName, field_view);
         } else {
-          nemo_file.read_volume_var(nemoName, geom.mesh(), time_indx,
-              field_view);
+          nemo_reader.read_var(nemoName, time_indx, field_view);
         }
-        auto missing_value = nemo_file.read_fillvalue<double>(nemoName);
+        auto missing_value = nemo_reader.read_fillvalue<double>(nemoName);
         field.metadata().set("missing_value", missing_value);
         field.metadata().set("missing_value_type", "approximately-equals");
         field.metadata().set("missing_value_epsilon", NEMO_FILL_TOL);