Merge branch 'bartgol/eamxx/data-interpolation-support-p3d-from-ps' (PR #2919)

Support using a source pressure profile reconstructed from ps and hyam/hybm.

Author: bartgol

components/eamxx/src/share/grid/abstract_grid.cpp

@@ -294,7 +294,7 @@ AbstractGrid::get_geometry_data (const std::string& name) const {
 }
 
 Field
-AbstractGrid::create_geometry_data (const FieldIdentifier& fid)
+AbstractGrid::create_geometry_data (const FieldIdentifier& fid, const int pack_size)
 {
   const auto& name = fid.name();
 
@@ -307,6 +307,7 @@ AbstractGrid::create_geometry_data (const FieldIdentifier& fid)
 
   // Create field and the read only copy as well
   auto& f = m_geo_fields[name] = Field(fid);
+  f.get_header().get_alloc_properties().request_allocation(pack_size);
   f.allocate_view();
   return f;
 }

components/eamxx/src/share/grid/abstract_grid.hpp

@@ -136,11 +136,12 @@ class AbstractGrid : public ekat::enable_shared_from_this<AbstractGrid>
   Field get_geometry_data (const std::string& name) const;
 
   // Create geometry data, throws if already existing. Returns writable field
-  Field create_geometry_data (const FieldIdentifier& fid);
+  Field create_geometry_data (const FieldIdentifier& fid, const int pack_size = 1);
   Field create_geometry_data (const std::string& name, const FieldLayout& layout,
                               const ekat::units::Units& units = ekat::units::Units::invalid(),
-                              const DataType data_type = DataType::RealType) {
-    return create_geometry_data(FieldIdentifier(name,layout,units,this->name(),data_type));
+                              const DataType data_type = DataType::RealType,
+                              const int pack_size = 1) {
+    return create_geometry_data(FieldIdentifier(name,layout,units,this->name(),data_type),pack_size);
   }
 
   // Sets pre-existing field as geometry data.

components/eamxx/src/share/io/scream_scorpio_interface.cpp

@@ -1204,19 +1204,6 @@ void read_var (const std::string &filename, const std::string &varname, T* buf,
         " - time len: " + std::to_string(f.time_dim->length));
     err = PIOc_setframe(f.ncid,var.ncid,frame);
     check_scorpio_noerr (err,f.name,"variable",varname,"read_var","setframe");
-  } else if (time_index>=0) {
-    // This is a bit of a hacky usage. We want to read a time index of a var,
-    // but the time dim is NOT unlimited in the input file. Apparently, SCORPIO
-    // supports this, and some E3SM input files are indeed like this, so we
-    // need to support it here too, hacky as it may be.
-    frame = time_index;
-
-    EKAT_REQUIRE_MSG (frame<var.dims[0]->length,
-        "Error! First dim index out of bounds.\n"
-        " - filename     : " + filename + "\n"
-        " - varname      : " + varname + "\n"
-        " - first dim idx: " + std::to_string(time_index) + "\n"
-        " - first dim len: " + std::to_string(var.dims[0]->length) + "\n");
   }
 
   std::string pioc_func;

components/eamxx/src/share/tests/data_interpolation_setup.cpp

@@ -56,9 +56,12 @@ TEST_CASE ("data_interpolation_setup")
       scorpio::define_var(fname+suffix,"s3d_i",{"ncol",ilev},        "real", true);
       scorpio::define_var(fname+suffix,"v3d_i",{"ncol","dim2",ilev}, "real", true);
 
-      // We keep p1d and p3d NOT time-dep
+      // We keep p1d, p2d, p3d, hyam, and hybm as NOT time-dep
       scorpio::define_var(fname+suffix,"p1d",  {"lev"},"real", false);
+      scorpio::define_var(fname+suffix,"p2d",  {"ncol"},"real", false);
       scorpio::define_var(fname+suffix,"p3d",  {"ncol","lev"},"real", false);
+      scorpio::define_var(fname+suffix,"hyam", {"lev"},"real", false);
+      scorpio::define_var(fname+suffix,"hybm", {"lev"},"real", false);
 
       scorpio::enddef(fname+suffix);
     }
@@ -81,14 +84,22 @@ TEST_CASE ("data_interpolation_setup")
     //   - two to be used for linear-hystory interp
     util::TimeStamp time = get_first_slice_time ();
 
-    // We keep p1d and p3d NOT time-dep, so we write outside the loop
-    auto p1d = base_fields.back();
+    // We keep pressures fields NOT time-dep, so we write outside the loop. Also write hyam/hybm here
+    auto p1d = base_fields[6];
     auto p3d = base_fields[2].alias("p3d");
+    auto p2d = base_fields[0].alias("p2d");
+    auto hybm = p1d.alias("hybm");
+    auto hyam = hybm.clone("hyam");
     p1d.sync_to_host();
+    p2d.sync_to_host();
     p3d.sync_to_host();
+    hyam.deep_copy(0); hyam.sync_to_host();
     for (const std::string& fname : files) {
       scorpio::write_var(fname+suffix,p1d.name(),p1d.get_internal_view_data<Real,Host>());
+      scorpio::write_var(fname+suffix,p2d.name(),p2d.get_internal_view_data<Real,Host>());
       scorpio::write_var(fname+suffix,p3d.name(),p3d.get_internal_view_data<Real,Host>());
+      scorpio::write_var(fname+suffix,hybm.name(),hybm.get_internal_view_data<Real,Host>());
+      scorpio::write_var(fname+suffix,hyam.name(),hyam.get_internal_view_data<Real,Host>());
     }
 
     int nfields = fields.size() - 1; // Don't handle p1d, since it's done above

components/eamxx/src/share/tests/data_interpolation_tests.cpp

@@ -14,7 +14,12 @@ namespace scream {
 // evaluation may be different (in finite prec) from the one in the class.
 constexpr auto tol = std::numeric_limits<Real>::epsilon()*10;
 
+// NOTE: due to how we build p3d, it is true that p3d is the tensor product of
+//       p1d and p2d. Since hybm=p1d and hyam=0, the P2D case is the same as
+//       the P3D one, once the 3d pressure has been reconstructed
+constexpr auto NOP = DataInterpolation::None;
 constexpr auto P1D = DataInterpolation::Static1D;
+constexpr auto P2D = DataInterpolation::Dynamic3DRef;
 constexpr auto P3D = DataInterpolation::Dynamic3D;
 
 using strvec_t = std::vector<std::string>;
@@ -73,7 +78,9 @@ void run_tests (const std::shared_ptr<const AbstractGrid>& grid,
     f.deep_copy(1);
   }
 
-  std::string data_pname = vr_type==P1D ? "p1d" : "p3d";  // if vr_type==None, it's not used anyways
+  std::string data_pname =  // if vr_type==None, it's not used anyways
+    vr_type==P1D ? "p1d" :
+                   (vr_type==P3D ? "p3d" : "p2d");
   auto model_pmid = base[2].clone("pmid");
   auto model_pint = base[4].clone("pint");
   if (vr_type==P1D) {
@@ -95,7 +102,7 @@ void run_tests (const std::shared_ptr<const AbstractGrid>& grid,
   }
 
   std::vector<Field> expected_vcoarse, base_vcoarse, ones_vcoarse;
-  if (vr_type==P1D or vr_type==P3D) {
+  if (vr_type!=NOP) {
     // If we do remap, there is some P0 extrapolation,
     // for which we need to know the data at the top/bot
     // NOTE: the data has NO ilev coord in this case
@@ -158,6 +165,7 @@ void run_tests (const std::shared_ptr<const AbstractGrid>& grid,
       const auto& layout = fields[i].get_header().get_identifier().get_layout();
       bool midpoints  = layout.has_tag(LEV);
       bool interfaces = layout.has_tag(ILEV);
+      bool has_levels = midpoints or interfaces;
 
       // Compute expected field, possibly clipping the top/bot value, if vr_type!=None
       // Since we run with 2x the number of levels as in the data, we get ONE level
@@ -167,7 +175,7 @@ void run_tests (const std::shared_ptr<const AbstractGrid>& grid,
       // of midpoints and interfaces on the data grid.
       expected[i].update(base[i],1,0);
       expected[i].update(ones[i],delta,1.0);
-      if ( (vr_type==P1D or vr_type==P3D) and (midpoints or interfaces)) {
+      if ( vr_type!=NOP and has_levels) {
         auto vlen = midpoints ? nlevs : nlevs+1;
         auto scalar = not layout.is_vector_layout();
 
@@ -202,6 +210,10 @@ void run_tests (const std::shared_ptr<const AbstractGrid>& grid,
       diff[i].update(fields[i],1,1);
       diff[i].update(expected[i],-1,1);
       diff[i].scale(1.0 / frobenius_norm<Real>(expected[i]));
+      if (frobenius_norm<Real>(diff[i])>=tol) {
+        print_field_hyperslab(expected[i]);
+        print_field_hyperslab(fields[i]);
+      }
       REQUIRE (frobenius_norm<Real>(diff[i])<tol);
     }
   }
@@ -265,6 +277,11 @@ TEST_CASE ("interpolation")
         run_tests (vfine_grid,files_no_ilev,t_beg,timeline,P1D);
         root_print(comm,"  timeline=PERIODIC, horiz_remap=NO,  vert_remap=p1d ......... PASS\n");
       }
+      SECTION ("p2d-vert") {
+        root_print(comm,"  timeline=PERIODIC, horiz_remap=NO,  vert_remap=p2d .........\n");
+        run_tests (vfine_grid,files_no_ilev,t_beg,timeline,P2D);
+        root_print(comm,"  timeline=PERIODIC, horiz_remap=NO,  vert_remap=p2d ......... PASS\n");
+      }
       SECTION ("p3d-vert") {
         root_print(comm,"  timeline=PERIODIC, horiz_remap=NO,  vert_remap=p3d .........\n");
         run_tests (vfine_grid,files_no_ilev,t_beg,timeline,P3D);
@@ -283,6 +300,11 @@ TEST_CASE ("interpolation")
         run_tests (hvfine_grid,files_no_ilev,t_beg,timeline,P1D);
         root_print(comm,"  timeline=PERIODIC, horiz_remap=YES, vert_remap=p1d ......... PASS\n");
       }
+      SECTION ("p2d-vert") {
+        root_print(comm,"  timeline=PERIODIC, horiz_remap=YES, vert_remap=p2d .........\n");
+        run_tests (hvfine_grid,files_no_ilev,t_beg,timeline,P2D);
+        root_print(comm,"  timeline=PERIODIC, horiz_remap=YES, vert_remap=p2d ......... PASS\n");
+      }
       SECTION ("p3d-vert") {
         root_print(comm,"  timeline=PERIODIC, horiz_remap=YES, vert_remap=p3d .........\n");
         run_tests (hvfine_grid,files_no_ilev,t_beg,timeline,P3D);
@@ -309,6 +331,11 @@ TEST_CASE ("interpolation")
         run_tests (vfine_grid,files_no_ilev,t_beg,timeline,P1D);
         root_print(comm,"  timeline=LINEAR,   horiz_remap=NO,  vert_remap=p1d ......... PASS\n");
       }
+      SECTION ("p2d-vert") {
+        root_print(comm,"  timeline=LINEAR,   horiz_remap=NO,  vert_remap=p2d .........\n");
+        run_tests (vfine_grid,files_no_ilev,t_beg,timeline,P2D);
+        root_print(comm,"  timeline=LINEAR,   horiz_remap=NO,  vert_remap=p2d ......... PASS\n");
+      }
       SECTION ("p3d-vert") {
         root_print(comm,"  timeline=LINEAR,   horiz_remap=NO,  vert_remap=p3d .........\n");
         run_tests (vfine_grid,files_no_ilev,t_beg,timeline,P3D);
@@ -327,6 +354,11 @@ TEST_CASE ("interpolation")
         run_tests (hvfine_grid,files_no_ilev,t_beg,timeline,P1D);
         root_print(comm,"  timeline=LINEAR,   horiz_remap=YES, vert_remap=p1d ......... PASS\n");
       }
+      SECTION ("p2d-vert") {
+        root_print(comm,"  timeline=LINEAR,   horiz_remap=YES, vert_remap=p2d .........\n");
+        run_tests (hvfine_grid,files_no_ilev,t_beg,timeline,P2D);
+        root_print(comm,"  timeline=LINEAR,   horiz_remap=YES, vert_remap=p2d ......... PASS\n");
+      }
       SECTION ("p3d-vert") {
         root_print(comm,"  timeline=LINEAR,   horiz_remap=YES, vert_remap=p3d .........\n");
         run_tests (hvfine_grid,files_no_ilev,t_beg,timeline,P3D);

components/eamxx/src/share/tests/data_interpolation_tests.hpp

@@ -58,7 +58,6 @@ create_fields (const std::shared_ptr<const AbstractGrid>& grid,
   int nlevs = grid->get_num_vertical_levels();
 
   // Create fields
-
   auto layout_s2d   = grid->get_2d_scalar_layout();
   auto layout_v2d   = grid->get_2d_vector_layout(ncmps);
   auto layout_s3d_m = grid->get_3d_scalar_layout(true);

components/eamxx/src/share/util/eamxx_data_interpolation.cpp

@@ -6,6 +6,7 @@
 #include "share/io/scream_scorpio_interface.hpp"
 #include "share/io/scream_io_utils.hpp"
 #include "share/util/scream_universal_constants.hpp"
+#include "physics/share/physics_constants.hpp"
 
 #include <filesystem>
 #include <fstream>
@@ -59,17 +60,55 @@ void DataInterpolation::run (const util::TimeStamp& ts)
     out.deep_copy(beg);
     out.update(end,alpha,1-alpha);
   }
-  // For Dynamic3D profile we also need to compute the source pressure profile
+
+  // For Dynamic3D/Dynamic3D profile we also need to compute the source pressure profile
   // NOTE: this can't be done in the loop above, since src_p is not a "remapped"
   //       field in the vertical remapper (also, we need to use ad different ptr)
   if (m_vr_type==Dynamic3D) {
     // The pressure field is THE LAST registered in the horiz remappers
     const auto p_beg = m_horiz_remapper_beg->get_tgt_field(m_nfields);
     const auto p_end = m_horiz_remapper_end->get_tgt_field(m_nfields);
 
-    auto p = m_vremap->get_source_pressure(true); // mid or int doesn't matter
+    auto p = m_helper_pressure_fields["src_p"];
     p.deep_copy(p_beg);
     p.update(p_end,alpha,1-alpha);
+  } else if (m_vr_type==Dynamic3DRef) {
+    // The surface pressure field is THE LAST registered in the horiz remappers
+    const auto ps_beg = m_horiz_remapper_beg->get_tgt_field(m_nfields);
+    const auto ps_end = m_horiz_remapper_end->get_tgt_field(m_nfields);
+
+    auto p  = m_helper_pressure_fields["src_p"];
+    auto ps = m_helper_pressure_fields["p_data"];
+    ps.deep_copy(ps_beg);
+    ps.update(ps_end,alpha,1-alpha);
+
+    // Reconstruct reference p from ps, hyam, and hybm
+    using KT = KokkosTypes<DefaultDevice>;
+    using ExeSpace = typename KT::ExeSpace;
+    using MemberType = typename KT::MemberType;
+    using ESU = ekat::ExeSpaceUtils<ExeSpace>;
+    using C = scream::physics::Constants<Real>;
+    using PT = ekat::Pack<Real,SCREAM_PACK_SIZE>;
+
+    auto ps_v = ps.get_view<const Real*>();
+    auto p_v  = p.get_view<PT**>();
+    auto hyam = m_vert_remapper->get_src_grid()->get_geometry_data("hyam").get_view<const PT*>();
+    auto hybm = m_vert_remapper->get_src_grid()->get_geometry_data("hybm").get_view<const PT*>();
+
+    constexpr auto P0 = C::P0;
+
+    const int ncols = ps_v.extent(0);
+    const int num_vert_packs = p_v.extent(1);
+    const auto policy = ESU::get_default_team_policy(ncols, num_vert_packs);
+
+    Kokkos::parallel_for("spa_compute_p_src_loop", policy,
+      KOKKOS_LAMBDA (const MemberType& team) {
+      const int icol = team.league_rank();
+      Kokkos::parallel_for(Kokkos::TeamVectorRange(team,num_vert_packs),
+                           [&](const int k) {
+        p_v(icol,k) = ps_v(icol) * hybm(k)  + P0 * hyam(k);
+      });
+    });
   }
 
   m_vert_remapper->remap_fwd();
@@ -94,7 +133,7 @@ update_end_fields ()
     fields.push_back(m_horiz_remapper_end->get_src_field(i));
   }
 
-  if (m_vr_type==Dynamic3D) {
+  if (m_vr_type==Dynamic3D or m_vr_type==Dynamic3DRef) {
     // We also need to read the src pressure profile
     fields.push_back(m_horiz_remapper_end->get_src_field(m_nfields));
   }
@@ -276,6 +315,7 @@ setup_remappers (const std::string& hremap_filename,
   int nlevs_data = get_input_files_dimlen ("lev");
   grid_after_hremap->reset_num_vertical_lev(nlevs_data);
 
+  std::shared_ptr<VerticalRemapper> vremap;
   if (hremap_filename!="") {
     m_horiz_remapper_beg = std::make_shared<RefiningRemapperP2P>(grid_after_hremap,hremap_filename);
     m_horiz_remapper_end = std::make_shared<RefiningRemapperP2P>(grid_after_hremap,hremap_filename);
@@ -306,10 +346,49 @@ setup_remappers (const std::string& hremap_filename,
       }
     };
 
-    m_vert_remapper = m_vremap = std::make_shared<VerticalRemapper>(grid_after_hremap,m_model_grid);
-    m_vremap->set_extrapolation_type(s2et(extrap_type_top),VerticalRemapper::Top);
-    m_vremap->set_extrapolation_type(s2et(extrap_type_bot),VerticalRemapper::Bot);
-    m_vremap->set_mask_value(mask_value);
+    auto vremap = std::make_shared<VerticalRemapper>(grid_after_hremap,m_model_grid);
+    
+    vremap->set_extrapolation_type(s2et(extrap_type_top),VerticalRemapper::Top);
+    vremap->set_extrapolation_type(s2et(extrap_type_bot),VerticalRemapper::Bot);
+    vremap->set_mask_value(mask_value);
+
+    // Setup vertical pressure profiles (which can add 1 extra field to hremap)
+    // This MUST be done before registering in vremap, since register_field_from_tgt
+    // REQUIRES to have source pressure profiles set BEFORE.
+
+    auto p_layout = vr_type==Static1D ? grid_after_hremap->get_vertical_layout(true)
+                                      : grid_after_hremap->get_3d_scalar_layout(true);
+    auto& src_p = m_helper_pressure_fields ["src_p"];
+    src_p = Field (FieldIdentifier("src_p",p_layout,ekat::units::Pa,grid_after_hremap->name()));
+    src_p.get_header().get_alloc_properties().request_allocation(SCREAM_PACK_SIZE);
+    src_p.allocate_view();
+    if (vr_type==Dynamic3D) {
+      // We load a full 3d profile, so p_data IS src_p
+      m_helper_pressure_fields ["p_data"] = src_p.alias(data_pname);
+    } else if (vr_type==Dynamic3DRef) {
+      // We load the surface pressure, and reconstruct src_p via p=ps*hybm(k) + p0*hyam(k)
+      auto& ps = m_helper_pressure_fields ["p_data"];
+      ps = Field(FieldIdentifier(data_pname,grid_after_hremap->get_2d_scalar_layout(),ekat::units::Pa,grid_after_hremap->name()));
+      ps.allocate_view();
+
+      // We need to reconstruct the 3d pressure from ps, hybm, and hyam.
+      // We read and store hyam/hybm in the vremap src grid
+      auto layout = grid_after_hremap->get_vertical_layout(true);
+      auto nondim = ekat::units::Units::nondimensional();
+      DataType real_t = DataType::RealType;
+      auto hyam = grid_after_hremap->create_geometry_data("hyam",layout,nondim,real_t,SCREAM_PACK_SIZE);
+      auto hybm = grid_after_hremap->create_geometry_data("hybm",layout,nondim,real_t,SCREAM_PACK_SIZE);
+      AtmosphereInput hvcoord_reader (m_time_database.files.front(),grid_after_hremap,{hyam,hybm},true);
+      hvcoord_reader.read_variables();
+    } else if (vr_type==Static1D) {
+      // Can load p now, since it's static
+      AtmosphereInput src_p_reader (m_time_database.files.front(),grid_after_hremap,{src_p.alias(data_pname)},true);
+      src_p_reader.read_variables();
+    }
+    vremap->set_source_pressure (m_helper_pressure_fields["src_p"],VerticalRemapper::Both);
+    vremap->set_target_pressure (model_pmid,model_pint);
+
+    m_vert_remapper = vremap;
   } else {
     // If no vert remap is requested, model_grid and grid_after_hremap MUST have same nlevs
     int model_nlevs = m_model_grid->get_num_vertical_levels();
@@ -319,36 +398,6 @@ setup_remappers (const std::string& hremap_filename,
         " - input data num vert levels: " + std::to_string(nlevs_data) + "\n");
     m_vert_remapper = std::make_shared<IDR>(grid_after_hremap,SAT);
   }
-
-  // Setup vertical pressure profiles (which can add 1 extra field to hremap)
-  // This MUST be done before registering in vremap, since register_field_from_tgt
-  // REQUIRES to have source pressure profiles set BEFORE.
-  Field data_p;
-  if (vr_type==Dynamic3D) {
-    // We also need to load and remap the pressure from the input files
-    auto hr_tgt_grid = m_horiz_remapper_beg->get_tgt_grid();
-    auto p_layout = hr_tgt_grid->get_3d_scalar_layout(true);
-    data_p = Field (FieldIdentifier(data_pname,p_layout,ekat::units::Pa,hr_tgt_grid->name()));
-    data_p.allocate_view();
-
-    m_vremap->set_source_pressure (data_p,VerticalRemapper::Both);
-    m_vremap->set_target_pressure (model_pmid,model_pint);
-  } else if (vr_type==Static1D) {
-    auto hr_tgt_grid = m_horiz_remapper_beg->get_tgt_grid();
-    auto p_layout = hr_tgt_grid->get_vertical_layout(true);
-    data_p = Field (FieldIdentifier(data_pname,p_layout,ekat::units::Pa,hr_tgt_grid->name()));
-    data_p.allocate_view();
-
-    // Use raw scorpio to read this var, since it's not decomposed. Use any file, since it's static
-    auto filename = m_time_database.files.front();
-    scorpio::register_file(filename,scorpio::Read);
-    scorpio::read_var(filename,data_pname,data_p.get_internal_view_data<Real,Host>());
-    scorpio::release_file(filename);
-    data_p.sync_to_dev();
-
-    m_vremap->set_source_pressure (data_p,VerticalRemapper::Both);
-    m_vremap->set_target_pressure (model_pmid,model_pint);
-  }
   m_vr_type = vr_type;
 
   // Register fields in the remappers. Vertical first, since we only have model-grid fields
@@ -365,9 +414,10 @@ setup_remappers (const std::string& hremap_filename,
     m_horiz_remapper_beg->register_field_from_tgt(f.clone());
     m_horiz_remapper_end->register_field_from_tgt(f.clone());
   }
-  if (vr_type==Dynamic3D) {
-    m_horiz_remapper_beg->register_field_from_tgt(data_p.clone());
-    m_horiz_remapper_end->register_field_from_tgt(data_p.clone());
+  if (vr_type==Dynamic3D or vr_type==Dynamic3DRef) {
+    const auto& data_p = m_helper_pressure_fields["p_data"];
+    m_horiz_remapper_beg->register_field_from_tgt(data_p.clone(data_pname));
+    m_horiz_remapper_end->register_field_from_tgt(data_p.clone(data_pname));
   }
   m_horiz_remapper_beg->registration_ends();
   m_horiz_remapper_end->registration_ends();