Feature/bump covariance

README.md

@@ -23,6 +23,7 @@ _orca-jedi_ includes executables to calculate model values at observation locati
   * [JCSDA/ufo](https://github.com/JCSDA/ufo)
   * [ecmwf/atlas](https://github.com/ecmwf/atlas)
   * [ecmwf/atlas-orca](https://github.com/ecmwf/atlas-orca)
+  * [saber](https://github.com/JCSDA/saber)
 
 ### Installing
 
@@ -68,6 +69,8 @@ OOPS_DEBUG=true
 OOPS_TRACE=true
 ```
 
+For help with SABER specifically please refer to the [SABER](https://jointcenterforsatellitedataassimilation-jedi-docs.readthedocs-hosted.com/en/latest/inside/jedi-components/saber/index.html) section in the JEDI documentation. 
+
 ## Authors
 
 The current lead maintainer is [@twsearle](https://github.com/twsearle) along with a large amount of help from Met Office contributors (see the "Contributors" page on github).

src/orca-jedi/geometry/Geometry.cc

@@ -2,15 +2,21 @@
  * (C) British Crown Copyright 2024 Met Office
  */
 
+#include <tuple>
+
 #include "orca-jedi/geometry/Geometry.h"
+#include "orca-jedi/utilities/Types.h"
 
 #include "atlas/field/Field.h"
 #include "atlas/field/FieldSet.h"
+#include "atlas/field/MissingValue.h"
 #include "atlas/functionspace/StructuredColumns.h"
 #include "atlas/mesh.h"
 #include "atlas/meshgenerator.h"
 #include "atlas/parallel/mpi/mpi.h"
 
+#include "atlas-orca/grid/OrcaGrid.h"
+
 #include "eckit/mpi/Comm.h"
 #include "eckit/config/Configuration.h"
 #include "eckit/exception/Exceptions.h"
@@ -109,6 +115,11 @@ Geometry::Geometry(const eckit::Configuration & config,
     funcSpace_ = atlas::functionspace::NodeColumns(
         mesh_, atlas::option::halo(halo));
     log_status();
+
+    if (params_.extraFieldsInit.value().value_or(false)) {
+      // Fill extra geometry fields for BUMP / SABER
+      create_extrafields();
+    }
 }
 
 // -----------------------------------------------------------------------------
@@ -124,6 +135,96 @@ const std::string Geometry::nemo_var_name(const std::string std_name) const {
   throw eckit::BadValue(err_stream.str(), Here());
 }
 
+// -----------------------------------------------------------------------------
+/// \brief Create extrafields used by BUMP in SABER.
+///        These are area, vunit, hmask, gmask.
+void Geometry::create_extrafields() {
+  extraFields_ = atlas::FieldSet();
+
+  atlas::OrcaGrid orcaGrid = mesh_.grid();
+  int nx = orcaGrid.nx() + orcaGrid.haloWest() + orcaGrid.haloEast();
+  int ny = orcaGrid.ny() + orcaGrid.haloNorth();
+  oops::Log::debug() << "orcagrid nx " << nx << " ny " << ny
+                     << std::endl;
+
+  // Create vertical unit field - the value used is not tuned.
+  atlas::Field vunit = funcSpace_.createField<double>(
+    atlas::option::name("vunit") | atlas::option::levels(n_levels_));
+  auto field_view = atlas::array::make_view<double, 2>(vunit);
+  for (atlas::idx_t j = 0; j < field_view.shape(0); ++j) {
+    for (atlas::idx_t k = 0; k < field_view.shape(1); ++k) {
+       field_view(j, k) = 1.;
+    }
+  }
+  oops::Log::debug() << "orcamodel::Geometry: adding vunit to extraFields."
+                     << std::endl;
+  extraFields_->add(vunit);
+
+  // Create owned or halo mask.
+  atlas::Field hmask = funcSpace_.createField<int32_t>(
+    atlas::option::name("owned") | atlas::option::levels(n_levels_));
+  auto ghost = atlas::array::make_view<int32_t, 1>(mesh_.nodes().ghost());
+  auto ij = atlas::array::make_view<int32_t, 2>(mesh_.nodes().field("ij"));
+
+  auto field_view1 = atlas::array::make_view<int32_t, 2>(hmask);
+  for (atlas::idx_t j = 0; j < field_view1.shape(0); ++j) {
+    for (atlas::idx_t k = 0; k < field_view1.shape(1); ++k) {
+      int x = ij(j, 0) + 1;
+      int y = ij(j, 1) + 1;
+      // 0 mask, 1 ocean
+      // setting some edge points to the mask value to prevent BUMP giving duplicate points error.
+      if (ghost(j) || x >= nx - 1 || y >= ny - 1) {
+        field_view1(j, k) = 0;
+      } else {
+        field_view1(j, k) = 1;
+      }
+    }
+  }
+  oops::Log::debug()
+      << "orcamodel::Geometry: adding owned to extraFields (set to all ocean except halo)."
+      << std::endl;
+  extraFields_->add(hmask);
+
+  // Create geometry mask or gmask.
+  atlas::Field gmask = funcSpace_.createField<int32_t>(
+    atlas::option::name("gmask") | atlas::option::levels(n_levels_));
+
+  auto field_view2 = atlas::array::make_view<int32_t, 2>(gmask);
+  for (atlas::idx_t j = 0; j < field_view2.shape(0); ++j) {
+    for (atlas::idx_t k = 0; k < field_view2.shape(1); ++k) {
+      int x = ij(j, 0) + 1;
+      int y = ij(j, 1) + 1;
+      // 0 mask, 1 ocean
+      // Setting some edge points to the mask value to prevent BUMP giving duplicate points error.
+      if (ghost(j) || x >= nx - 1 || y >= ny - 1) {
+        field_view2(j, k) = 0;
+      } else {
+        field_view2(j, k) = 1;
+      }
+    }
+  }
+  oops::Log::debug() << "orcamodel::Geometry: adding gmask (set to all ocean except halo)."
+                     << std::endl;
+  extraFields_->add(gmask);
+
+  // Create grid cell area field /m^2 - the value used is not tuned.
+  // Curerently ~= area of 2 degree square grid cell at the equator.
+  atlas::Field area = funcSpace_.createField<double>(
+    atlas::option::name("area") | atlas::option::levels(n_levels_));
+  auto field_view3 = atlas::array::make_view<double, 2>(area);
+
+  for (atlas::idx_t j = 0; j < field_view3.shape(0); ++j) {
+    for (atlas::idx_t k = 0; k < field_view3.shape(1); ++k) {
+      field_view3(j, k) = 4e10;
+    }
+  }
+  log_status();
+
+  oops::Log::debug() << "orcamodel::Geometry: adding area to extraFields."
+                     << std::endl;
+  extraFields_->add(area);
+}
+
 // -----------------------------------------------------------------------------
 /// \brief Give the number of levels for each provided level - surface variables
 ///        have 1 level, volumetric variables have "number levels" levels.
@@ -273,4 +374,40 @@ void Geometry::log_status() const {
       << " Gb" << std::endl;
 }
 
+/// \brief Set gmask extra variable in geometry based on input field missing values.
+/// \param[in]    atlas::Field field.
+void Geometry::set_gmask(atlas::Field & field) const {
+  oops::Log::debug() << "orcamodel::Geometry setting gmask from field "
+                     << field.name() << " missing values" << std::endl;
+
+  atlas::Field gmask = extraFields_.field("gmask");
+
+  atlas::field::MissingValue mv(field);
+  bool has_mv = static_cast<bool>(mv);
+  oops::Log::debug() << "has_mv " << has_mv << std::endl;
+
+  const auto setGmaskField = [&](auto typeVal) {
+    using T = decltype(typeVal);
+    auto field_viewin = atlas::array::make_view<T, 2>(field);
+    auto field_viewgm = atlas::array::make_view<int32_t, 2>(gmask);
+    if (has_mv) {
+      for (atlas::idx_t j = 0; j < field_viewgm.shape(0); ++j) {
+        for (atlas::idx_t k = 0; k < field_viewgm.shape(1); ++k) {
+          // Only change values that are currently unmasked ( 0 mask, 1 ocean ).
+          if (field_viewgm(j, k) == 1) {
+            if (mv(field_viewin(j, k))) {
+              field_viewgm(j, k) = 0;
+            }
+          }
+        }
+      }
+    }
+  };
+  ApplyForFieldType(setGmaskField,
+                    fieldPrecision(field.name()),
+                    std::string("orcamodel::Geometry::set_gmask ")
+                    + field.name() + "' field type not recognised");
+  log_status();
+}
+
 }  // namespace orcamodel

src/orca-jedi/geometry/Geometry.h

@@ -9,6 +9,7 @@
 #include <string>
 #include <vector>
 #include <memory>
+#include <tuple>
 
 #include "atlas/field/Field.h"
 #include "atlas/field/FieldSet.h"
@@ -54,10 +55,13 @@ class Geometry : public util::Printable {
       const;
   const eckit::mpi::Comm & getComm() const {return comm_;}
   const oops::Variables & variables() const;
+  void create_extrafields();
   void latlon(std::vector<double> & lats, std::vector<double> & lons,
               const bool halo) const;
   const atlas::FunctionSpace & functionSpace() const {return funcSpace_;}
-  const atlas::FieldSet & fields() const {return nofields_;}
+  const atlas::FieldSet & extraFields() const {return extraFields_;}
+  const atlas::FieldSet & fields() const {return extraFields_;}
+  atlas::FieldSet & extraFields() {return extraFields_;}
 
   const atlas::Grid & grid() const {return grid_;}
   const atlas::Mesh & mesh() const {return mesh_;}
@@ -70,6 +74,7 @@ class Geometry : public util::Printable {
   FieldDType fieldPrecision(std::string variable_name) const;
   std::shared_ptr<eckit::Timer> timer() const {return eckit_timer_;}
   void log_status() const;
+  void set_gmask(atlas::Field &) const;
 
  private:
   void print(std::ostream &) const;
@@ -81,9 +86,10 @@ class Geometry : public util::Printable {
   atlas::grid::Partitioner partitioner_;
   atlas::Mesh mesh_;
   atlas::functionspace::NodeColumns funcSpace_;
-  atlas::FieldSet nofields_;
   std::shared_ptr<eckit::Timer> eckit_timer_;
+  atlas::FieldSet extraFields_;
 };
+
 // -----------------------------------------------------------------------------
 
 }  // namespace orcamodel

src/orca-jedi/geometry/GeometryParameters.h

@@ -13,6 +13,7 @@
 #include "oops/base/ParameterTraitsVariables.h"
 #include "oops/util/parameters/Parameter.h"
 #include "oops/util/parameters/RequiredParameter.h"
+#include "oops/util/parameters/OptionalParameter.h"
 #include "orca-jedi/geometry/GeometryParameterTraitsFieldDType.h"
 #include "orca-jedi/utilities/Types.h"
 
@@ -57,6 +58,7 @@ class OrcaGeometryParameters : public oops::Parameters {
         " The default will not distribute the data ('serial').",
       "serial",
       this};
+  oops::OptionalParameter<bool> extraFieldsInit{"initialise extra fields", this};
 };
 
 }  //  namespace orcamodel