Add interpolator increment apply and adjoint methods and add increment to

frld · frld · commit fa7c893be16a · 2024-07-30T17:28:32.000+01:00
state method. Update yaml 3DVar yaml file so 3DVar runs through (but produces nans currently).
diff --git a/src/orca-jedi/interpolator/Interpolator.cc b/src/orca-jedi/interpolator/Interpolator.cc
@@ -27,6 +27,7 @@
 
 #include "orca-jedi/geometry/Geometry.h"
 #include "orca-jedi/state/State.h"
+#include "orca-jedi/increment/Increment.h"
 
 #include "orca-jedi/interpolator/Interpolator.h"
 
@@ -177,6 +178,149 @@ template void Interpolator::executeInterpolation<float>(
     const std::vector<bool> & mask,
     std::vector<double>::iterator& iter) const;
 
+void Interpolator::apply(const oops::Variables& vars, const Increment& inc,
+           const std::vector<bool> & mask,
+           std::vector<double>& result) const {
+
+  // DJL question can the atlas templates help?
+
+  // input is inc output is result
+
+  const size_t nvars = vars.size();
+
+  for (size_t j=0; j < nvars; ++j) {
+    if (!inc.variables().has(vars[j])) {
+      std::stringstream err_stream;
+      err_stream << "orcamodel::Interpolator::apply varname \" "
+                 << "\" " << vars[j]
+                 << " not found in the model increment." << std::endl;
+      err_stream << "    add the variable to the increment variables and "
+                 << "add a mapping from the geometry to that variable."
+                 << std::endl;
+      throw eckit::BadParameter(err_stream.str(), Here());
+    }
+  }
+
+  const std::vector<size_t> varSizes =
+    inc.geometry()->variableSizes(vars);
+  size_t nvals = 0;
+  for (size_t jvar=0; jvar < nvars; ++jvar) nvals += nlocs_ * varSizes[jvar];
+  result.resize(nvals);
+
+  std::size_t out_idx = 0;
+  for (size_t jvar=0; jvar < nvars; ++jvar) {
+    auto gv_varname = vars[jvar].name();
+    atlas::Field tgt_field = atlasObsFuncSpace_.createField<double>(
+        atlas::option::name(gv_varname) |
+        atlas::option::levels(varSizes[jvar]));
+    interpolator_.execute(inc.incrementFields()[gv_varname], tgt_field);
+    auto field_view = atlas::array::make_view<double, 2>(tgt_field);
+    atlas::field::MissingValue mv(inc.incrementFields()[gv_varname]);
+    bool has_mv = static_cast<bool>(mv);
+    for (std::size_t klev=0; klev < varSizes[jvar]; ++klev) {
+      for (std::size_t iloc=0; iloc < nlocs_; iloc++) {
+        if (has_mv && mv(field_view(iloc, klev))) {
+          result[out_idx] = util::missingValue<double>();
+        } else {
+          result[out_idx] = field_view(iloc, klev);
+        }
+        ++out_idx;
+      }
+    }
+  }
+}
+
+void Interpolator::applyAD(const oops::Variables& vars, Increment& inc,
+             const std::vector<bool> & mask,
+             const std::vector<double> & resultin) const
+{
+
+  // input is resultin output is inc
+
+  oops::Log::trace() << "orcamodel::Interpolator::applyAD start "
+                     << std::endl;
+
+  oops::Log::debug() << "DJL ** Interpolator::applyAD this needs checking **" << std::endl;
+
+// ** Not sure what I'm doing yet - Just trying to do the opposite of applyAD
+
+  const size_t nvars = vars.size();
+
+  for (size_t j=0; j < nvars; ++j) {
+    if (!inc.variables().has(vars[j])) {
+      std::stringstream err_stream;
+      err_stream << "orcamodel::Interpolator::apply varname \" "
+                 << "\" " << vars[j]
+                 << " not found in the model increment." << std::endl;
+      err_stream << "    add the variable to the increment variables and "
+                 << "add a mapping from the geometry to that variable."
+                 << std::endl;
+      throw eckit::BadParameter(err_stream.str(), Here());
+    }
+  }
+
+  const std::vector<size_t> varSizes =
+    inc.geometry()->variableSizes(vars);
+  size_t nvals = 0;
+
+//    boost::uuids::uuid uuid = boost::uuids::random_generator()();    
+//    std::shared_ptr<const Geometry> geom = inc.geometry();
+//    writeGenFieldsToFile("applyADpre"+ boost::uuids::to_string(uuid) +".nc", *geom, inc.validTime(), inc.incrementFields());
+
+  for (size_t jvar=0; jvar < nvars; ++jvar) nvals += nlocs_ * varSizes[jvar];
+//    result.resize(nvals);
+
+  std::size_t out_idx = 0;
+  for (size_t jvar=0; jvar < nvars; ++jvar) {
+    oops::Log::debug() << "DJL ** jvar " << jvar << " " << nvars
+                       << "varSizes " << varSizes[jvar]
+                       << std::endl;
+    auto gv_varname = vars[jvar].name();
+//      atlas::Field tgt_field = atlasObsFuncSpace_.createField<double>(
+//          atlas::option::name(gv_varname) |
+//          atlas::option::levels(varSizes[jvar]));
+//      atlas::Field incField = inc.incrementFields()[jvar];
+
+    auto tgt_field = atlasObsFuncSpace_.createField<double>(
+      atlas::option::name(gv_varname) |
+      atlas::option::levels(varSizes[jvar]));
+
+    auto field_view = atlas::array::make_view<double, 2>(tgt_field);
+//    field_view.assign(0.0);
+
+    for (std::size_t klev=0; klev < varSizes[jvar]; ++klev) {
+      for (std::size_t iloc=0; iloc < nlocs_; iloc++) {
+//          if (has_mv && mv(field_view(iloc, klev))) {
+////            result[out_idx] = util::missingValue(result[out_idx]);
+//          } else {
+        oops::Log::debug() << "DJL iloc " << iloc << " klev " << klev << " out_idx "
+                          << out_idx << " resultin[out_idx] " << resultin[out_idx] << std::endl;
+        field_view(iloc, klev) = resultin[out_idx];
+//          }
+        ++out_idx;
+      }
+    }
+
+
+//      atlas::field::MissingValue mv(inc.incrementFields()[gv_varname]);
+//      bool has_mv = static_cast<bool>(mv);
+    interpolator_.execute_adjoint(inc.incrementFields()[gv_varname], tgt_field);
+
+  }    // jvar
+
+  // DJL write to file
+
+//  boost::uuids::uuid uuid = boost::uuids::random_generator()();    
+//  std::shared_ptr<const Geometry> geom = inc.geometry();
+//    std::ostringstream out;
+//    out << std::setfill('0') << std::setw(6) << uuid;
+//  writeGenFieldsToFile("applyAD"+ boost::uuids::to_string(uuid) +".nc", *geom, inc.validTime(), inc.incrementFields());
+//    fileCounter_++;
+
+  oops::Log::trace() << "orcamodel::Interpolator::applyAD done "
+                     << std::endl;
+}
+
 void Interpolator::print(std::ostream & os) const {
   os << "orcamodel::Interpolator: " << std::endl;
   os << "  Obs function space " << atlasObsFuncSpace_ << std::endl;
diff --git a/src/orca-jedi/interpolator/Interpolator.h b/src/orca-jedi/interpolator/Interpolator.h
@@ -58,16 +58,10 @@ class Interpolator : public util::Printable,
              std::vector<double>& result) const;
   void apply(const oops::Variables& vars, const Increment& inc,
              const std::vector<bool> & mask,
-             std::vector<double>& result) const {
-    throw eckit::NotImplemented("Increment interpolation not implemented",
-                                Here());
-  }
-  void applyAD(const oops::Variables& vars, const Increment& inc,
+             std::vector<double>& result) const;
+  void applyAD(const oops::Variables& vars, Increment& inc,
                const std::vector<bool> & mask,
-               const std::vector<double> &) const {
-    throw eckit::NotImplemented("Adjoint interpolation not implemented",
-                                Here());
-  }
+               const std::vector<double> &) const;
 
  private:
   template<class T> void executeInterpolation(
diff --git a/src/orca-jedi/interpolator/InterpolatorParameters.h b/src/orca-jedi/interpolator/InterpolatorParameters.h
@@ -22,6 +22,7 @@ class OrcaAtlasInterpolatorParameters : public oops::Parameters {
   oops::OptionalParameter<std::string> non_linear{"non_linear", this};
   oops::OptionalParameter<double>
     max_fraction_elems_to_try{"max_fraction_elems_to_try", this};
+  oops::OptionalParameter<bool> adjoint{"adjoint", this};
 };
 
 class OrcaInterpolatorParameters : public oops::Parameters {
diff --git a/src/orca-jedi/state/State.cc b/src/orca-jedi/state/State.cc
@@ -165,10 +165,35 @@ State & State::operator=(const State & rhs) {
 // Interactions with Increments
 
 State & State::operator+=(const Increment & dx) {
-  std::string err_message =
-      "orcamodel::State::State::operator+= not implemented";
-  throw eckit::NotImplemented(err_message, Here());
+
   oops::Log::trace() << "State(ORCA)::add increment starting" << std::endl;
+
+  ASSERT(this->validTime() == dx.validTime());
+
+// DJL assumes state and increments have the same fields in the field set
+// DJL add something to enforce/check this
+
+  auto ghost = atlas::array::make_view<int32_t, 1>(
+      geom_->mesh().nodes().ghost());
+  for (int i = 0; i< stateFields_.size();i++)
+  {
+    atlas::Field field = stateFields_[i];
+    atlas::Field fieldi = dx.incrementFields()[i];
+
+    std::string fieldName = field.name();
+    std::string fieldNamei = fieldi.name();
+    oops::Log::debug() << "orcamodel::Increment::add:: state field name = " << fieldName
+                       << " increment field name = " << fieldNamei
+                       << std::endl;
+    auto field_view = atlas::array::make_view<double, 2>(field);
+    auto field_viewi = atlas::array::make_view<double, 2>(fieldi);
+    for (atlas::idx_t j = 0; j < field_view.shape(0); ++j) {
+      for (atlas::idx_t k = 0; k < field_view.shape(1); ++k) {
+        if (!ghost(j)) field_view(j, k) += field_viewi(j, k);
+      }
+    }
+  }
+ 
   oops::Log::trace() << "State(ORCA)::add increment done" << std::endl;
   return *this;
 }
diff --git a/src/tests/testinput/3dvar_ice_obs.yaml b/src/tests/testinput/3dvar_ice_obs.yaml
@@ -44,15 +44,15 @@ cost function:
         obsdataout:
           engine:
             type: H5File
-            obsfile: testoutput/test_3dvar.nc
+            obsfile: testoutput/test_3dvar_obsdataout.nc
         simulated variables: [ice_area_fraction]
       get values:
         time interpolation: linear
         atlas-interpolator:
           type: unstructured-bilinear-lonlat
           non_linear: missing-if-all-missing
           max_fraction_elems_to_try: 0.0
-#          adjoint: true
+          adjoint: true
       obs operator:
         name: Composite
         components:
@@ -91,20 +91,17 @@ variational:
       write increment: true
       increment:
         state component:
-          datadir: Data
-          output path: testoutput/increments_ice
+          output path: testoutput/increments_ice.nc
           date: 2021-06-30T00:00:00Z
-          exp: 3dvar.iter1
-          type: in
+#          type: in
 final:
   diagnostics:
     departures: oman
 output:
   date: 2021-06-30T00:00:00Z
   state variables: [ ice_area_fraction ]
-  nemo field file: testoutput/3dvar_ice
-#  exp: 3dvar
-#  steps: ["2021-06-29T12:00:00Z"]               # state write not implemented 
-
+  nemo field file: testoutput/3dvar_ice2.nc
+  output nemo field file: testoutput/3dvar_ice_analysis.nc
+#  steps: ["2021-06-29T12:00:00Z"]
 #test:
 #  reference filename: testoutput/test_3dvar_ice_obs.ref
