Use VADER for insitu/potential TL/AD

src/soca/LinearVariableChange/LinearVariableChange.cc

@@ -5,15 +5,18 @@
  * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
  */
 
+#include <map>
 #include <ostream>
 #include <string>
 
 #include "soca/LinearVariableChange/LinearVariableChange.h"
+#include "soca/VariableChange/Model2GeoVaLs/Model2GeoVaLs.h"
 
 #include "soca/Geometry/Geometry.h"
 #include "soca/Increment/Increment.h"
 #include "soca/State/State.h"
 
+#include "oops/util/FieldSetHelpers.h"
 #include "oops/util/Logger.h"
 
 using oops::Log;
@@ -22,10 +25,24 @@ namespace soca {
 
 // -----------------------------------------------------------------------------
 
+std::map<std::string, std::vector<std::string>> SocaLinVaderCookbook {
+  {"sea_water_temperature", {"SeaWaterTemperature_A", "SeaWaterTemperature_B"}}
+};
+
+// -----------------------------------------------------------------------------
+
 LinearVariableChange::LinearVariableChange(const Geometry & geom,
                                            const eckit::Configuration & config)
-  : geom_(geom), params_(), linVarChas_() {
+  : geom_(geom), params_(), linVarChas_(), vader_(), default_(false) {
     params_.deserialize(config);
+
+  // setup vader
+  eckit::LocalConfiguration vaderConfig, vaderCookbookConfig;
+  for (auto kv : SocaLinVaderCookbook) vaderCookbookConfig.set(kv.first, kv.second);
+  vaderConfig.set(vader::configCookbookKey, vaderCookbookConfig);
+  vader_.reset(new vader::Vader(params_.vader, vaderConfig));
+
+  default_ = (params_.linearVariableChangesWrapper.value() == boost::none);
 }
 
 // -----------------------------------------------------------------------------
@@ -36,9 +53,37 @@ LinearVariableChange::~LinearVariableChange() {}
 
 void LinearVariableChange::changeVarTraj(const State & xfg,
                                          const oops::Variables & vars) {
-  Log::trace() << "LinearVariableChange::setTrajectory starting" << std::endl;
+  Log::trace() << "LinearVariableChange::changeVarTraj starting" << std::endl;
+
+  // The following is TEMPORARY.
+  // ----------------------------------------------------------------------------
+  // We need to get some variables BEFORE we run VADER.
+  State xfg_vadertraj(geom_, xfg);
+  if (default_ && vars.has("sea_water_temperature")) {
+    Log::debug() << "LinearVariableChange::changeVarTraj Pre-VADER variable changes. " << std::endl;
+    oops::Variables preVaderVars(std::vector<std::string>{
+      "latitude",
+      "longitude",
+      "sea_water_depth",
+      "sea_area_fraction",
+      });
+    preVaderVars += xfg.variables();
+    xfg_vadertraj.updateFields(preVaderVars);
+    Model2GeoVaLs m2gv(geom_, params_.toConfiguration());
+    m2gv.changeVar(xfg, xfg_vadertraj);
+    Log::debug() << "LinearVariableChange::changeVarTraj variables after var change: "
+        << xfg_vadertraj.variables() << std::endl;
+  }
 
-  // TODO(travis): do something with vars?
+  // call Vader
+  // ----------------------------------------------------------------------------
+  oops::Variables varsFilled = xfg_vadertraj.variables();
+  oops::Variables varsVader = vars;
+  varsVader -= varsFilled;  // pass only the needed variables
+  atlas::FieldSet xfs;
+  xfg_vadertraj.toFieldSet(xfs);
+  xfs.haloExchange();
+  vader_->changeVarTraj(xfs, varsVader);
 
   // TODO(travis): this is not ideal. We are saving the first trajectory and
   // assuming it is the background. This should all get ripped out when the
@@ -72,11 +117,21 @@ void LinearVariableChange::changeVarTraj(const State & xfg,
   Log::trace() << "LinearVariableChange::setTrajectory done" << std::endl;
 }
 
+// -------------------------------------------------------------------------------------------------
+
+void LinearVariableChange::initVaderTLAD(oops::Variables & ingredientVars) const {
+  oops::Log::trace() << "LinearVariableChange::initVaderTLAD starting" << std::endl;
+  oops::Variables originalIngredientVars = ingredientVars;
+  varsVaderPopulates_ = vader_->initTLAD(ingredientVars);
+  varsVaderPopulates_ -= originalIngredientVars;
+  oops::Log::trace() << "LinearVariableChange::initVaderTLAD done" << std::endl;
+}
+
 // -----------------------------------------------------------------------------
 
 void LinearVariableChange::changeVarTL(Increment & dx,
                                        const oops::Variables & vars) const {
-  Log::trace() << "LinearVariableChange::multiply starting" << std::endl;
+  Log::trace() << "LinearVariableChange::changeVarTL starting" << std::endl;
 
   // If all variables already in incoming state just remove the no longer
   // needed fields
@@ -87,6 +142,26 @@ void LinearVariableChange::changeVarTL(Increment & dx,
   //   return;
   // }
 
+  // Make sure vader is fully initialized
+  if (vader_->needsTLADInit()) {
+    oops::Variables ingredientVars = dx.variables();
+    initVaderTLAD(ingredientVars);
+    oops::Log::info() << " changevarTL: vader will populate variables: "
+                      << varsVaderPopulates_ << std::endl;
+  }
+  // If Vader is doing anything, call Vader
+  if (varsVaderPopulates_.size() > 0) {
+    atlas::FieldSet dxfs;
+    dx.toFieldSet(dxfs);
+    vader_->changeVarTL(dxfs);
+    // Set intermediate state for the Increment containing original fields plus the ones
+    // Vader has done
+    oops::Variables varsVader = dx.variables();
+    varsVader += varsVaderPopulates_;
+    dx.updateFields(varsVader);
+    dx.fromFieldSet(dxfs);
+  }
+
   // Create output state
   Increment dxout(dx.geometry(), vars, dx.validTime());
 
@@ -98,12 +173,6 @@ void LinearVariableChange::changeVarTL(Increment & dx,
      dx = dxout;
   }
 
-  // Allocate any extra fields and remove fields no longer needed
-  // dx.updateFields(vars);
-
-  // Copy data from temporary state
-  // dx = dxout;
-
   Log::trace() << "LinearVariableChange::multiply done" << std::endl;
 }
 
@@ -132,17 +201,60 @@ void LinearVariableChange::changeVarInverseTL(Increment & dx,
 
 void LinearVariableChange::changeVarAD(Increment & dx,
                                        const oops::Variables & vars) const {
-  Log::trace() << "LinearVariableChange::multiplyAD starting" << std::endl;
+  Log::trace() << "LinearVariableChange::changeVarAD starting" << std::endl;
+
+  // Make sure vader is fully initialized
+  if (vader_->needsTLADInit()) {
+    oops::Variables ingredientVars(vars);
+    initVaderTLAD(ingredientVars);
+    oops::Log::info() << " changevarAD: vader will populate variables: "
+                      << varsVaderPopulates_ << std::endl;
+  }
+
+  // Create dx_for_vader as a copy of dx, with just the variables created by Vader
+  // (in the forward direction). Remove the variables created by vader from dx.
+  // This way we ensure the model code will not be able to do the adjoint for these vars
+  Increment dx_for_vader(dx, true);  // true => full copy
+  oops::Variables varsVaderDidntPopulate = dx.variables();
+  varsVaderDidntPopulate -= varsVaderPopulates_;
+  dx_for_vader.updateFields(varsVaderPopulates_);
+  dx.updateFields(varsVaderDidntPopulate);
+
+  // Create empty output state
   Increment dxout(dx.geometry(), vars, dx.validTime());
 
-  // Call variable change(s)
+  // Call model's adjoint variable change.
   for (ircst_ it = linVarChas_.rbegin(); it != linVarChas_.rend(); ++it) {
     dxout.zero();
     it->multiplyAD(dx, dxout);
     dx.updateFields(vars);
     dx = dxout;
   }
 
+  // dxout needs to temporarily have the variables that Vader populated put into it before
+  // being passed into vader_.changeVarAD, so Vader can do its adjoints.
+  if (varsVaderPopulates_.size() > 0) {
+    atlas::FieldSet dxfs;
+    dx.toFieldSet(dxfs);
+    oops::Variables varsVaderWillAdjoint = varsVaderPopulates_;
+    atlas::FieldSet dxvaderfs;
+    dx_for_vader.toFieldSet(dxvaderfs);
+    for (const auto field : dxvaderfs) {
+      dxfs.add(field);
+    }
+
+    oops::Variables varsAdjointed = vader_->changeVarAD(dxfs);
+    varsVaderWillAdjoint -= varsAdjointed;
+    // After changeVarAD, vader should have removed everything from varsVaderWillAdjoint,
+    // indicating it did all the adjoints we expected it to.
+    ASSERT(varsVaderWillAdjoint.size() == 0);
+    // remove the fields that were adjointed by vader from dxout_fs
+    util::removeFieldsFromFieldSet(dxfs, varsAdjointed.variables());
+    // Copy dxout into dx for return
+    dx.updateFields(vars);
+    dx.fromFieldSet(dxfs);
+  }
+
   Log::trace() << "LinearVariableChange::multiplyAD done" << std::endl;
 }
 

src/soca/LinearVariableChange/LinearVariableChange.h

@@ -22,6 +22,8 @@
 #include "oops/util/parameters/RequiredParameter.h"
 #include "oops/util/Printable.h"
 
+#include "vader/vader.h"
+
 #include "soca/LinearVariableChange/Base/LinearVariableChangeBase.h"
 
 namespace soca {
@@ -38,6 +40,7 @@ class LinearVariableChangeParameters : public oops::Parameters {
   oops::OptionalParameter<oops::Variables> outputVariables{"output variables", this};
   oops::OptionalParameter<std::vector<LinearVariableChangeParametersWrapper>>
          linearVariableChangesWrapper{"linear variable changes", this};
+  oops::Parameter<vader::VaderParameters> vader{"vader", {}, this};
 };
 
 // -----------------------------------------------------------------------------
@@ -65,10 +68,15 @@ class LinearVariableChange : public util::Printable {
 
  private:
   void print(std::ostream &) const override;
+  void initVaderTLAD(oops::Variables & ingredientVars) const;
   Parameters_ params_;
   const Geometry & geom_;
   std::unique_ptr<State> bkg_;
   LinVarChaVec_ linVarChas_;
+
+  std::unique_ptr<vader::Vader> vader_;
+  mutable oops::Variables varsVaderPopulates_;
+  bool default_;
 };
 
 // -----------------------------------------------------------------------------

test/testinput/3dvar_nicas.yml

@@ -176,7 +176,13 @@ cost function:
             obsfile: data_static/obs/prof.nc
         simulated variables: [waterTemperature]
       obs operator:
-        name: InsituTemperature
+        name: VertInterp
+        observation alias file: testinput/obsop_name_map.yml
+        variables:
+          - waterTemperature
+        vertical coordinate: sea_water_depth
+        observation vertical coordinate: depth
+        interpolation method: linear
       obs error:
         covariance model: diagonal
       obs filters:

test/testref/3dvar_nicas.test

@@ -1,25 +1,25 @@
-CostJo   : Nonlinear Jo(SeaSurfaceTemp) = 6.8893641637637177e+02, nobs = 113, Jo/n = 6.0967824458086000e+00, err = 3.6655656239312961e-01
+CostJo   : Nonlinear Jo(SeaSurfaceTemp) = 6.8893641637637154e+02, nobs = 113, Jo/n = 6.0967824458085973e+00, err = 3.6655656239312961e-01
 CostJo   : Nonlinear Jo(SeaSurfaceSalinity) = 9.2074193097584697e+01, nobs = 49, Jo/n = 1.8790651652568306e+00, err = 1.0000000000000000e+00
-CostJo   : Nonlinear Jo(ADT) = 2.2281591911422575e+02, nobs = 89, Jo/n = 2.5035496529688288e+00, err = 1.0000000149011612e-01
-CostJo   : Nonlinear Jo(InsituTemperature) = 3.4206375767190184e+02, nobs = 203, Jo/n = 1.6850431412408957e+00, err = 9.0034813852418905e-01
-CostJo   : Nonlinear Jo(InsituSalinity) = 3.0723498032711791e+02, nobs = 218, Jo/n = 1.4093347721427427e+00, err = 6.0819699545777528e-01
+CostJo   : Nonlinear Jo(ADT) = 2.2281591911422569e+02, nobs = 89, Jo/n = 2.5035496529688279e+00, err = 1.0000000149011612e-01
+CostJo   : Nonlinear Jo(InsituTemperature) = 3.4207282026717877e+02, nobs = 203, Jo/n = 1.6850877845673831e+00, err = 9.0034813852418905e-01
+CostJo   : Nonlinear Jo(InsituSalinity) = 3.0723498032711768e+02, nobs = 218, Jo/n = 1.4093347721427416e+00, err = 6.0819699545777528e-01
 CostJb   : Nonlinear Jb = 0.0000000000000000e+00
-CostFunction: Nonlinear J = 1.6531252665872018e+03
-RPCGMinimizer: reduction in residual norm = 1.2006776473098589e-01
+CostFunction: Nonlinear J = 1.6531343291824785e+03
+RPCGMinimizer: reduction in residual norm = 1.2414458192527290e-01
 CostFunction::addIncrement: Analysis: 
   Valid time: 2018-04-15T00:00:00Z
        sea_water_cell_thickness   min=0.0009999999999977   max=1345.6400000000003274   mean=112.4337347359288941
-             sea_water_salinity   min=0.0000000000000000   max=38.6959149218921183   mean=29.0069721693647047
-sea_water_potential_temperature   min=-1.8883899372702533   max=31.0125853407566474   mean=5.2973741711447007
+             sea_water_salinity   min=0.0000000000000000   max=38.6959149218921183   mean=29.0074916345918368
+sea_water_potential_temperature   min=-1.8883899372702533   max=31.7246904174557400   mean=5.3128599069475166
     eastward_sea_water_velocity   min=-0.6975788196638589   max=0.4840625031242498   mean=0.0010009071341291
    northward_sea_water_velocity   min=-0.7661101215480253   max=0.9402890903466523   mean=0.0021201094019194
- sea_surface_height_above_geoid   min=-2.0936283531745894   max=0.8190071369678245   mean=-0.2049739891582961
+ sea_surface_height_above_geoid   min=-2.1508977733523640   max=0.8678867820476882   mean=-0.2004871633702467
     ocean_mixed_layer_thickness   min=0.0005000000000000   max=4593.1523423819935488   mean=165.9225478094618893
                 sea_water_depth   min=0.0005000000000000   max=5587.8978052886013757   mean=1053.5232951172040430
-CostJo   : Nonlinear Jo(SeaSurfaceTemp) = 275.5343698338789409, nobs = 113, Jo/n = 2.4383572551670705, err = 0.3665565623931296
-CostJo   : Nonlinear Jo(SeaSurfaceSalinity) = 92.7854784143652296, nobs = 49, Jo/n = 1.8935811921299026, err = 1.0000000000000000
-CostJo   : Nonlinear Jo(ADT) = 178.4631772821379343, nobs = 89, Jo/n = 2.0052042391251454, err = 0.1000000014901161
-CostJo   : Nonlinear Jo(InsituTemperature) = 280.3382013630258029, nobs = 203, Jo/n = 1.3809763613942159, err = 0.9003481385241890
-CostJo   : Nonlinear Jo(InsituSalinity) = 119.7825517174736092, nobs = 218, Jo/n = 0.5494612464104294, err = 0.6081969954577753
-CostJb   : Nonlinear Jb = 29.8302471366160731
-CostFunction: Nonlinear J = 976.7340257474976397
+CostJo   : Nonlinear Jo(SeaSurfaceTemp) = 218.3295992979803088, nobs = 113, Jo/n = 1.9321203477697373, err = 0.3665565623931296
+CostJo   : Nonlinear Jo(SeaSurfaceSalinity) = 91.0601724565712800, nobs = 49, Jo/n = 1.8583708664606384, err = 1.0000000000000000
+CostJo   : Nonlinear Jo(ADT) = 169.9471585510565319, nobs = 89, Jo/n = 1.9095186354051297, err = 0.1000000014901161
+CostJo   : Nonlinear Jo(InsituTemperature) = 420.0581423685579239, nobs = 203, Jo/n = 2.0692519328500389, err = 0.9003481385241890
+CostJo   : Nonlinear Jo(InsituSalinity) = 121.1026570349107629, nobs = 218, Jo/n = 0.5555167753894990, err = 0.6081969954577753
+CostJb   : Nonlinear Jb = 39.6288485214363888
+CostFunction: Nonlinear J = 1060.1265782305131324