Feature/bump covariance (#123)

* Add 3DVar and error covariance toolbox (dirac test) main programs. * Add initial version of 3dvar and dirac ctest. * Add saber to citest. * Another attempt to fix CI test. Add state tofieldset test. * Fixes to yaml files. Dirac test working. 3DVar needs adjoint interpolator adding. * Add interpolator increment apply and adjoint methods and add increment to state method. Update yaml 3DVar yaml file so 3DVar runs through (but produces nans currently). * Initial attempt to implement an increment interpolator unit test. * Fixes to new interpolator unit tests. * interpolator adjoint deal with missing values. * Add some additional code to respect fillvalues. * Update ci container (#117) * Update ci container * allow manual trigger of ci * temporary change to force ci to run * revert temporary change as ci now gets past cmake config step * Fix coding norms. * Tidying up. * Fix a merge error. * Apply suggestions from code review Co-authored-by: Toby Searle <[email protected]> * Update copyrights. * Add reference data for 3dvar and dirac application tests. Turn off creation of json schema again for 3dvar and errorcovariancetoolbox applications. * Update 3dvar 3dvar 3dvar minimizer to DRPCG as recommended. * Rename ice to sic to match other tests etc. * Add BUMP NICA grid calculation. Change to use BUMP background error correlations in Dirac and 3DVar. * Add BUMP NICA grid calculation. Change to use BUMP background error correlations in Dirac and 3DVar. * Add BUMP NICA grid calculation. Change to use BUMP background error correlations in Dirac and 3DVar. * Add BUMP NICA grid calculation. Change to use BUMP background error correlations in Dirac and 3DVar. * Remove old 3dvar yaml file. * Changes to allow increments to deal with fillvalues produced by BUMP. * Temporary fix to new atlas error about non-linear interpolation. Add some input data to the repo so the tests work. * Make masking out duplicate points for bump more generic. * Change bump settings to make test for setting up grid quicker. Attempt to revert change to not use missing-if-all-missing in atlas interpolator. * Filter out all observations with missing values in the innovations. Temporary fix for "Adjoint interpolation only works for interpolation schemes that are linear" error in 3dvar. * Updates to fix atlas adjoint interpolation error message. * Add bump nicas test reference data. Add units tests. Tidy up code/comments. * Add support for missing values to more increment methods. * Bug fix to missing value detection. * Add identity 3dvar test (old test). Add some halo exchanges when going from increments to fieldset and vice versa. * Change fill value. * Update fill value * Apply suggestions from code review Co-authored-by: Toby Searle <[email protected]> * Changes in response to review. In particular removing the unneeded custom routine to determine the i and j points. * Add a bit more info to comment about extra field area. * Add a description to the set_gmask (set geometry mask) optional parameter. --------- Co-authored-by: Toby Searle <[email protected]>
MetOffice · Nov 20, 2024 · 0ec97d1 · 0ec97d1
1 parent 343dd1b
commit 0ec97d1
Show file tree

Hide file tree

Showing 24 changed files with 679 additions and 53 deletions.
diff --git a/README.md b/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).

diff --git a/src/orca-jedi/geometry/Geometry.cc b/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
diff --git a/src/orca-jedi/geometry/Geometry.h b/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
diff --git a/src/orca-jedi/geometry/GeometryParameters.h b/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