EAMxx: correctly handle mask fields in VerticalRemapper

Author: bartgol

components/eamxx/src/share/grid/remap/vertical_remapper.cpp

@@ -160,8 +160,7 @@ registration_ends_impl ()
     const auto& src = m_src_fields[i];
           auto& tgt = m_tgt_fields[i];
 
-    // Clone src layout, since we may strip dims later for mask creation
-    auto src_layout = src.get_header().get_identifier().get_layout().clone();
+    const auto& src_layout = src.get_header().get_identifier().get_layout().clone();
 
     if (src_layout.has_tag(LEV) or src_layout.has_tag(ILEV)) {
       // Determine if this field can be handled with packs, and whether it's at midpoints
@@ -185,47 +184,27 @@ registration_ends_impl ()
         // NOTE: for now we assume that masking is determined only by the COL,LEV location in space
         //       and that fields with multiple components will have the same masking for each component
         //       at a specific COL,LEV
-        src_layout.strip_dims({CMP});
+
+        auto tgt_layout = create_tgt_layout(src_layout);
+        tgt_layout.strip_dims({CMP});
 
         // I this mask has already been created, retrieve it, otherwise create it
-        const auto mask_name = m_tgt_grid->name() + "_" + ekat::join(src_layout.names(),"_") + "_mask";
-        Field tgt_mask;
-        if (m_field2type.count(mask_name)==0) {
+        const auto mask_name = m_tgt_grid->name() + "_" + ekat::join(tgt_layout.names(),"_") + "_mask";
+        auto& mask = m_masks[mask_name];
+        if (not mask.is_allocated()) {
           auto nondim = ekat::units::Units::nondimensional();
           // Create this src/tgt mask fields, and assign them to these src/tgt fields extra data
 
-          FieldIdentifier src_mask_fid (mask_name, src_layout, nondim, m_src_grid->name() );
-          FieldIdentifier tgt_mask_fid = create_tgt_fid(src_mask_fid);
-
-          Field src_mask (src_mask_fid);
-          src_mask.allocate_view();
-
-          tgt_mask  = Field (tgt_mask_fid);
-          tgt_mask.allocate_view();
-
-          // Initialize the src mask values to 1.0
-          src_mask.deep_copy(1.0);
-
-          m_src_masks.push_back(src_mask);
-          m_tgt_masks.push_back(tgt_mask);
-
-          auto& mt = m_field2type[src_mask_fid.name()];
-          mt.packed = false;
-          mt.midpoints = src_layout.has_tag(LEV);
-        } else {
-          for (size_t i=0; i<m_tgt_masks.size(); ++i) {
-            if (m_tgt_masks[i].name()==mask_name) {
-              tgt_mask = m_tgt_masks[i];
-              break;
-            }
-          }
+          FieldIdentifier mask_fid (mask_name, tgt_layout, nondim, m_tgt_grid->name() );
+          mask  = Field (mask_fid);
+          mask.allocate_view();
         }
 
         EKAT_REQUIRE_MSG(not tgt.get_header().has_extra_data("mask_field"),
             "[VerticalRemapper::registration_ends_impl] Error! Target field already has mask data assigned.\n"
             " - tgt field name: " + tgt.name() + "\n");
 
-        tgt.get_header().set_extra_data("mask_field",tgt_mask);
+        tgt.get_header().set_extra_data("mask_field",mask);
 
         // Since we do mask (at top and/or bot), the tgt field MAY be contain fill_value entries
         tgt.get_header().set_may_be_filled(true);
@@ -383,6 +362,8 @@ create_layout (const FieldLayout& from_layout,
 
 void VerticalRemapper::remap_fwd_impl ()
 {
+  using namespace ShortFieldTagsNames;
+
   // 1. Setup any interp object that was created (if nullptr, no fields need it)
   if (m_lin_interp_mid_packed) {
     setup_lin_interp(*m_lin_interp_mid_packed,m_src_pmid,m_tgt_pmid);
@@ -397,9 +378,12 @@ void VerticalRemapper::remap_fwd_impl ()
     setup_lin_interp(*m_lin_interp_int_scalar,m_src_pint,m_tgt_pint);
   }
 
-  using namespace ShortFieldTagsNames;
+  // 2. Init all masks fields (if any) to 1 (signaling no masked entries)
+  for (auto& [name, mask] : m_masks) {
+    mask.deep_copy(1);
+  }
 
-  // 2. Interpolate the fields
+  // 3. Interpolate the fields
   for (int i=0; i<m_num_fields; ++i) {
     const auto& f_src    = m_src_fields[i];
           auto& f_tgt    = m_tgt_fields[i];
@@ -435,29 +419,6 @@ void VerticalRemapper::remap_fwd_impl ()
     }
   }
 
-  // 3. Interpolate the mask fields
-  for (unsigned i=0; i<m_tgt_masks.size(); ++i) {
-          auto& f_src = m_src_masks[i];
-          auto& f_tgt = m_tgt_masks[i];
-    const auto& type = m_field2type.at(f_src.name());
-
-    // Dispatch interpolation to the proper lin interp object
-    if (type.midpoints) {
-      if (type.packed) {
-        apply_vertical_interpolation(*m_lin_interp_mid_packed,f_src,f_tgt,m_src_pmid,m_tgt_pmid);
-      } else {
-        apply_vertical_interpolation(*m_lin_interp_mid_scalar,f_src,f_tgt,m_src_pmid,m_tgt_pmid);
-      }
-      extrapolate(f_src,f_tgt,m_src_pmid,m_tgt_pmid);
-    } else {
-      if (type.packed) {
-        apply_vertical_interpolation(*m_lin_interp_int_packed,f_src,f_tgt,m_src_pint,m_tgt_pint);
-      } else {
-        apply_vertical_interpolation(*m_lin_interp_int_scalar,f_src,f_tgt,m_src_pint,m_tgt_pint);
-      }
-      extrapolate(f_src,f_tgt,m_src_pint,m_tgt_pint);
-    }
-  }
 }
 
 template<int Packsize>
@@ -641,6 +602,12 @@ extrapolate (const Field& f_src,
   auto etop = m_etype_top;
   auto ebot = m_etype_bot;
   auto mid = nlevs_tgt / 2;
+  auto do_mask = etop==Mask or ebot==Mask;
+  decltype(f_tgt.get_view<Real**>()) mask_v;
+  if (do_mask) {
+    mask_v = f_tgt.get_header().get_extra_data<Field>("mask_field").get_view<Real**>();
+  }
+
   switch(f_src.rank()) {
     case 2:
     {
@@ -674,6 +641,7 @@ extrapolate (const Field& f_src,
                 y_tgt[ilev] = y_src[nlevs_src-1];
               } else {
                 y_tgt[ilev] = fill_val;
+                mask_v(icol,ilev) = 0;
               }
             }
           } else {
@@ -683,6 +651,7 @@ extrapolate (const Field& f_src,
                 y_tgt[ilev] = y_src[0];
               } else {
                 y_tgt[ilev] = fill_val;
+                mask_v(icol,ilev) = 0;
               }
             }
           }
@@ -725,6 +694,7 @@ extrapolate (const Field& f_src,
                 y_tgt[ilev] = y_src[nlevs_src-1];
               } else {
                 y_tgt[ilev] = fill_val;
+                mask_v(icol,ilev) = 0;
               }
             }
           } else {
@@ -734,6 +704,7 @@ extrapolate (const Field& f_src,
                 y_tgt[ilev] = y_src[0];
               } else {
                 y_tgt[ilev] = fill_val;
+                mask_v(icol,ilev) = 0;
               }
             }
           }

components/eamxx/src/share/grid/remap/vertical_remapper.hpp

@@ -112,9 +112,8 @@ class VerticalRemapper : public AbstractRemapper
 
   ekat::Comm            m_comm;
 
-  // Source and target fields
-  std::vector<Field>    m_src_masks;
-  std::vector<Field>    m_tgt_masks;
+  // Tgt grid masks (in case extrap type at top or bot is Mask)
+  std::map<std::string,Field>    m_masks;
 
   // Vertical profile fields, both for source and target
   Field                 m_src_pmid;

components/eamxx/src/share/io/tests/io_remap_test.cpp

@@ -675,7 +675,6 @@ ekat::ParameterList set_output_params(const std::string& name, const std::string
 
   params.set<std::string>("filename_prefix",name);
   params.set<std::string>("averaging_type","instant");
-  params.set<int>("max_snapshots_per_file",1);
   params.set<std::string>("floating_point_precision","real");
   auto& oc = params.sublist("output_control");
   oc.set<int>("frequency",1);