EAMxx: make horiz interp remappers (coarsening and refining) handle fields without COL tag

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

@@ -34,26 +34,7 @@ CoarseningRemapper (const grid_ptr_type& src_grid,
         continue;
       }
       const auto& src_data = src_grid->get_geometry_data(name);
-      const auto& src_data_fid = src_data.get_header().get_identifier();
-      const auto& layout = src_data_fid.get_layout();
-      if (layout.tags().empty()) {
-        // This is a scalar field, so won't be coarsened.
-        // Simply copy it in the tgt grid, but we still need to assign the new grid name.
-        FieldIdentifier tgt_data_fid(src_data_fid.name(),src_data_fid.get_layout(),src_data_fid.get_units(),m_tgt_grid->name());
-        auto tgt_data = m_coarse_grid->create_geometry_data(tgt_data_fid);
-        tgt_data.deep_copy(src_data);
-      } else if (layout.tags()[0]!=COL) {
-        // Not a field to be coarsened (perhaps a vertical coordinate field).
-        // Simply copy it in the tgt grid, but we still need to assign the new grid name.
-        FieldIdentifier tgt_data_fid(src_data_fid.name(),src_data_fid.get_layout(),src_data_fid.get_units(),m_tgt_grid->name());
-        auto tgt_data = m_coarse_grid->create_geometry_data(tgt_data_fid);
-        tgt_data.deep_copy(src_data);
-      } else {
-        // This field needs to be remapped
-        auto tgt_data_fid = create_tgt_fid(src_data_fid);
-        auto tgt_data = m_coarse_grid->create_geometry_data(tgt_data_fid);
-        register_field(src_data,tgt_data);
-      }
+      register_field_from_src(src_data);
     }
     registration_ends();
     if (get_num_fields()>0) {
@@ -202,10 +183,15 @@ void CoarseningRemapper::remap_fwd_impl ()
     return (ap.get_last_extent() % SCREAM_PACK_SIZE) == 0;
   };
 
-  // Loop over each field
+  // First, perform the local mat-vec. Recall that in these y=Ax products,
+  // x is the src field, and y is the overlapped tgt field.
   for (int i=0; i<m_num_fields; ++i) {
-    // First, perform the local mat-vec. Recall that in these y=Ax products,
-    // x is the src field, and y is the overlapped tgt field.
+    if (m_needs_remap[i]==0) {
+      // No need to do a mat-vec here. Just deep copy and move on
+      m_tgt_fields[i].deep_copy(m_src_fields[i]);
+      continue;
+    }
+
     const auto& f_src = m_src_fields[i];
     const auto& f_ov  = m_ov_fields[i];
 
@@ -611,16 +597,15 @@ void CoarseningRemapper::pack_and_send ()
   const auto pid_lid_start = m_send_pid_lids_start;
   const auto lids_pids = m_send_lids_pids;
   const auto buf = m_send_buffer;
-  constexpr auto COL = FieldTag::Column;
 
   for (int ifield=0; ifield<m_num_fields; ++ifield) {
+    if (m_needs_remap[ifield]==0)
+      continue; // Not a field to coarsen
+
     const auto& f  = m_ov_fields[ifield];
     const auto& fl = f.get_header().get_identifier().get_layout();
     const auto f_pid_offsets = ekat::subview(m_send_f_pid_offsets,ifield);
 
-    if (not fl.has_tag(COL))
-      continue; // Not a field to coarsen
-
     switch (fl.rank()) {
       case 1:
       {
@@ -752,6 +737,9 @@ void CoarseningRemapper::recv_and_unpack ()
   const auto recv_lids_end = m_recv_lids_end;
   const auto recv_lids_pidpos = m_recv_lids_pidpos;
   for (int ifield=0; ifield<m_num_fields; ++ifield) {
+    if (m_needs_remap[ifield]==0)
+      continue; // Not a field to coarsen
+
           auto& f  = m_tgt_fields[ifield];
     const auto& fl = f.get_header().get_identifier().get_layout();
     const auto f_pid_offsets = ekat::subview(m_recv_f_pid_offsets,ifield);
@@ -966,6 +954,9 @@ void CoarseningRemapper::setup_mpi_data_structures ()
   std::vector<int> field_col_size (m_num_fields);
   int sum_fields_col_sizes = 0;
   for (int i=0; i<m_num_fields; ++i) {
+    if (m_needs_remap[i]==0)
+      continue;
+
     const auto& f  = m_src_fields[i];
     const auto& fl = f.get_header().get_identifier().get_layout();
     field_col_size[i] = fl.clone().strip_dim(COL).size();
@@ -1012,6 +1003,9 @@ void CoarseningRemapper::setup_mpi_data_structures ()
   for (int pid=0,pos=0; pid<m_comm.size(); ++pid) {
     send_pid_offsets[pid] = pos;
     for (int i=0; i<m_num_fields; ++i) {
+      if (m_needs_remap[i]==0)
+        continue;
+
       send_f_pid_offsets_h(i,pid) = pos;
       pos += field_col_size[i]*pid2lids_send[pid].size();
     }
@@ -1110,6 +1104,9 @@ void CoarseningRemapper::setup_mpi_data_structures ()
     recv_pid_offsets[pid] = pos;
     const int num_recv_gids = recv_pid_start[pid+1] - recv_pid_start[pid];
     for (int i=0; i<m_num_fields; ++i) {
+      if (m_needs_remap[i]==0)
+        continue;
+
       recv_f_pid_offsets_h(i,pid) = pos;
       pos += field_col_size[i]*num_recv_gids;
     }

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

@@ -90,6 +90,9 @@ HorizInterpRemapperBase::
 
 void HorizInterpRemapperBase::registration_ends_impl ()
 {
+  using namespace ShortFieldTagsNames;
+
+  m_needs_remap.resize(m_num_fields,1);
   for (int i=0; i<m_num_fields; ++i) {
     const auto& src_dt = m_src_fields[i].get_header().get_identifier().data_type();
     const auto& tgt_dt = m_tgt_fields[i].get_header().get_identifier().data_type();
@@ -99,6 +102,17 @@ void HorizInterpRemapperBase::registration_ends_impl ()
         "  - tgt field name: " + m_tgt_fields[i].name() + "\n"
         "  - src data type : " + e2str(src_dt) + "\n"
         "  - tgt data type : " + e2str(tgt_dt) + "\n");
+
+    const auto& src_fl = m_src_fields[i].get_header().get_identifier().get_layout();
+    if (not src_fl.has_tag(COL)) {
+      // This field will be skipped in several of the remap steps
+      m_needs_remap[i] = 0;
+    } else {
+      EKAT_REQUIRE_MSG (src_fl.tag(0)==COL,
+          "[HorizInterpRemapperBase::registration_ends_impl] Error! If present, the COL dimension MUST be the first one.\n"
+          " - field name: " + m_src_fields[i].name() + "\n"
+          " - field layout: " + src_fl.to_string() + "\n");
+    }
   }
 
   create_ov_fields ();
@@ -107,13 +121,22 @@ void HorizInterpRemapperBase::registration_ends_impl ()
 
 void HorizInterpRemapperBase::create_ov_fields ()
 {
+  using namespace ShortFieldTagsNames;
+
   m_ov_fields.reserve(m_num_fields);
   const auto num_ov_gids = m_ov_coarse_grid->get_num_local_dofs();
   const auto ov_gn = m_ov_coarse_grid->name();
   const auto dt = DataType::RealType;
   for (int i=0; i<m_num_fields; ++i) {
     const auto& f = m_type==InterpType::Refine ? m_tgt_fields[i] : m_src_fields[i];
     const auto& fid = f.get_header().get_identifier();
+    if (m_needs_remap[i]==0) {
+      // This field won't be remapped. We can simply emplace an empty field (which won't be used),
+      // to make sure m_ov_fields[i] always returns the ov field for the i-th field
+      m_ov_fields.emplace_back();
+      continue;
+    }
+
     const auto layout = fid.get_layout().clone().reset_dim(0,num_ov_gids);
     FieldIdentifier ov_fid (fid.name(),layout,fid.get_units(),ov_gn,dt);
 
@@ -255,6 +278,7 @@ void HorizInterpRemapperBase::clean_up ()
   m_src_fields.clear();
   m_tgt_fields.clear();
   m_ov_fields.clear();
+  m_needs_remap.clear();
 
   // Reset the state of the base class
   m_state = RepoState::Clean;

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

@@ -76,6 +76,10 @@ class HorizInterpRemapperBase : public AbstractRemapper
 
   ekat::Comm      m_comm;
 
+  // Store whether each field needs remap. Only fields with COL dim do.
+  // NOTE: use int and NOT bool, as vector<bool> is evil
+  std::vector<int>    m_needs_remap;
+
   static std::map<std::string,HorizRemapperData> s_remapper_data;
 };
 

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

@@ -50,6 +50,12 @@ void RefiningRemapperP2P::remap_fwd_impl ()
   // Loop over each field, perform mat-vec
   constexpr auto COL = ShortFieldTagsNames::COL;
   for (int i=0; i<m_num_fields; ++i) {
+    if (m_needs_remap[i]==0) {
+      // No need to do a mat-vec here. Just deep copy and move on
+      m_tgt_fields[i].deep_copy(m_src_fields[i]);
+      continue;
+    }
+
     auto& f_tgt = m_tgt_fields[i];
 
     // It's ok to register fields that do not have the COL tag
@@ -87,6 +93,9 @@ void RefiningRemapperP2P::setup_mpi_data_structures ()
   // Get cumulative col size of each field (to be used to compute offsets)
   m_fields_col_sizes_scan_sum.resize(m_num_fields+1,0);
   for (int i=0; i<m_num_fields; ++i) {
+    if (m_needs_remap[i]==0)
+      continue;
+
     const auto& f = m_src_fields[i];
     const auto& fl = f.get_header().get_identifier().get_layout();
 
@@ -166,8 +175,6 @@ void RefiningRemapperP2P::pack_and_send ()
   using TeamMember  = typename KT::MemberType;
   using ESU         = ekat::ExeSpaceUtils<typename KT::ExeSpace>;
 
-  constexpr auto COL = ShortFieldTagsNames::COL;
-
   auto export_pids = m_imp_exp->export_pids();
   auto export_lids = m_imp_exp->export_lids();
   auto ncols_send  = m_imp_exp->num_exports_per_pid();
@@ -176,12 +183,12 @@ void RefiningRemapperP2P::pack_and_send ()
   const int num_exports = export_pids.size();
   const int total_col_size = m_fields_col_sizes_scan_sum.back();
   for (int ifield=0; ifield<m_num_fields; ++ifield) {
-    const auto& f = m_src_fields[ifield];
-    const auto& fl = f.get_header().get_identifier().get_layout();
-    if (not fl.has_tag(COL)) {
+    if (m_needs_remap[ifield]==0)
       // No need to process this field. We'll deep copy src->tgt later
       continue;
-    }
+
+    const auto& f = m_src_fields[ifield];
+    const auto& fl = f.get_header().get_identifier().get_layout();
     const auto f_col_sizes_scan_sum = m_fields_col_sizes_scan_sum[ifield];
     switch (fl.rank()) {
       case 1:
@@ -316,8 +323,6 @@ void RefiningRemapperP2P::recv_and_unpack ()
   using TeamMember  = typename KT::MemberType;
   using ESU         = ekat::ExeSpaceUtils<typename KT::ExeSpace>;
 
-  constexpr auto COL = ShortFieldTagsNames::COL;
-
   auto import_pids = m_imp_exp->import_pids();
   auto import_lids = m_imp_exp->import_lids();
   auto ncols_recv  = m_imp_exp->num_imports_per_pid();
@@ -326,12 +331,12 @@ void RefiningRemapperP2P::recv_and_unpack ()
   const int num_imports = import_pids.size();
   const int total_col_size = m_fields_col_sizes_scan_sum.back();
   for (int ifield=0; ifield<m_num_fields; ++ifield) {
-          auto& f  = m_ov_fields[ifield];
-    const auto& fl = f.get_header().get_identifier().get_layout();
-    if (not fl.has_tag(COL)) {
+    if (m_needs_remap[ifield]==0)
       // No need to process this field. We'll deep copy src->tgt later
       continue;
-    }
+
+          auto& f  = m_ov_fields[ifield];
+    const auto& fl = f.get_header().get_identifier().get_layout();
     const auto f_col_sizes_scan_sum = m_fields_col_sizes_scan_sum[ifield];
     switch (fl.rank()) {
       case 1:

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

@@ -113,6 +113,8 @@ Field create_field (const std::string& name, const LayoutType lt, const Abstract
   const auto& gn = grid.name();
   Field f;
   switch (lt) {
+    case LayoutType::Scalar1D:
+      f = Field(FieldIdentifier(name,grid.get_vertical_layout(midpoints),u,gn)); break;
     case LayoutType::Scalar2D:
       f = Field(FieldIdentifier(name,grid.get_2d_scalar_layout(),u,gn));  break;
     case LayoutType::Vector2D:
@@ -161,6 +163,10 @@ Field all_gather_field_impl (const Field& f, const ekat::Comm& comm) {
   constexpr auto COL = ShortFieldTagsNames::COL;
   const auto& fid = f.get_header().get_identifier();
   const auto& fl  = fid.get_layout();
+  if (not fl.has_tag(COL)) {
+    // Not partitioned
+    return f;
+  }
   int col_size = fl.clone().strip_dim(COL).size();
   auto tags = fl.tags();
   auto dims = fl.dims();
@@ -310,6 +316,7 @@ TEST_CASE("coarsening_remap")
   // Here we will simplify and just remap a simple 2D horizontal field.
   auto tgt_grid = remap->get_coarse_grid();
 
+  auto src_s1d   = create_field("s1d",  LayoutType::Scalar1D, *src_grid, true, engine);
   auto src_s2d   = create_field("s2d",  LayoutType::Scalar2D, *src_grid, false, engine);
   auto src_v2d   = create_field("v2d",  LayoutType::Vector2D, *src_grid, false, engine);
   auto src_t2d   = create_field("t2d",  LayoutType::Tensor2D, *src_grid, false, engine);
@@ -320,6 +327,7 @@ TEST_CASE("coarsening_remap")
   auto src_t3d_m = create_field("t3d_m",LayoutType::Tensor3D, *src_grid, true,  engine);
   auto src_t3d_i = create_field("t3d_i",LayoutType::Tensor3D, *src_grid, false, engine);
 
+  auto tgt_s1d   = create_field("s1d",  LayoutType::Scalar1D, *tgt_grid, true);
   auto tgt_s2d   = create_field("s2d",  LayoutType::Scalar2D, *tgt_grid, false);
   auto tgt_v2d   = create_field("v2d",  LayoutType::Vector2D, *tgt_grid, false);
   auto tgt_t2d   = create_field("t2d",  LayoutType::Tensor2D, *tgt_grid, false);
@@ -330,8 +338,8 @@ TEST_CASE("coarsening_remap")
   auto tgt_t3d_m = create_field("t3d_m",LayoutType::Tensor3D, *tgt_grid, true );
   auto tgt_t3d_i = create_field("t3d_i",LayoutType::Tensor3D, *tgt_grid, false);
 
-  std::vector<Field> src_f = {src_s2d,src_v2d,src_t2d,src_s3d_m,src_s3d_i,src_v3d_m,src_v3d_i,src_t3d_m,src_t3d_i};
-  std::vector<Field> tgt_f = {tgt_s2d,tgt_v2d,tgt_t2d,tgt_s3d_m,tgt_s3d_i,tgt_v3d_m,tgt_v3d_i,tgt_t3d_m,tgt_t3d_i};
+  std::vector<Field> src_f = {src_s1d,src_s2d,src_v2d,src_t2d,src_s3d_m,src_s3d_i,src_v3d_m,src_v3d_i,src_t3d_m,src_t3d_i};
+  std::vector<Field> tgt_f = {tgt_s1d,tgt_s2d,tgt_v2d,tgt_t2d,tgt_s3d_m,tgt_s3d_i,tgt_v3d_m,tgt_v3d_i,tgt_t3d_m,tgt_t3d_i};
 
   // -------------------------------------- //
   //     Register fields in the remapper    //
@@ -373,6 +381,11 @@ TEST_CASE("coarsening_remap")
       root_print (msg + "\n",comm);
       bool ok = true;
       switch (l.type()) {
+        case LayoutType::Scalar1D:
+        {
+          CHECK ( views_are_equal(gsrc,gtgt) );
+          ok &= catch_capture.lastAssertionPassed();
+        } break;
         case LayoutType::Scalar2D:
         {
           const auto v_src = gsrc.get_view<const Real*,Host>();

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

@@ -25,6 +25,8 @@ Field create_field (const std::string& name, const LayoutType lt, const Abstract
   const auto  ndims = 2;
   Field f;
   switch (lt) {
+    case LayoutType::Scalar1D:
+      f = Field(FieldIdentifier(name,grid.get_vertical_layout(true),u,gn));  break;
     case LayoutType::Scalar2D:
       f = Field(FieldIdentifier(name,grid.get_2d_scalar_layout(),u,gn));  break;
     case LayoutType::Vector2D:
@@ -64,6 +66,10 @@ Field all_gather_field (const Field& f, const ekat::Comm& comm) {
   constexpr auto COL = ShortFieldTagsNames::COL;
   const auto& fid = f.get_header().get_identifier();
   const auto& fl  = fid.get_layout();
+  if (not fl.has_tag(COL)) {
+    // Not partitioned
+    return f;
+  }
   int col_size = fl.clone().strip_dim(COL).size();
   auto tags = fl.tags();
   auto dims = fl.dims();
@@ -205,17 +211,20 @@ TEST_CASE ("refining_remapper") {
   auto src_grid = r->get_src_grid();
 
   auto bundle_src = create_field("bundle3d_src",LayoutType::Vector3D,*src_grid,engine);
+  auto s1d_src   = create_field("s1d_src",LayoutType::Scalar1D,*src_grid,engine);
   auto s2d_src   = create_field("s2d_src",LayoutType::Scalar2D,*src_grid,engine);
   auto v2d_src   = create_field("v2d_src",LayoutType::Vector2D,*src_grid,engine);
   auto s3d_src   = create_field("s3d_src",LayoutType::Scalar3D,*src_grid,engine);
   auto v3d_src   = create_field("v3d_src",LayoutType::Vector3D,*src_grid,engine);
 
   auto bundle_tgt = create_field("bundle3d_tgt",LayoutType::Vector3D,*tgt_grid);
+  auto s1d_tgt   = create_field("s1d_tgt",LayoutType::Scalar1D,*tgt_grid);
   auto s2d_tgt   = create_field("s2d_tgt",LayoutType::Scalar2D,*tgt_grid);
   auto v2d_tgt   = create_field("v2d_tgt",LayoutType::Vector2D,*tgt_grid);
   auto s3d_tgt   = create_field("s3d_tgt",LayoutType::Scalar3D,*tgt_grid);
   auto v3d_tgt   = create_field("v3d_tgt",LayoutType::Vector3D,*tgt_grid);
 
+  r->register_field(s1d_src,s1d_tgt);
   r->register_field(s2d_src,s2d_tgt);
   r->register_field(v2d_src,v2d_tgt);
   r->register_field(s3d_src,s3d_tgt);
@@ -229,19 +238,37 @@ TEST_CASE ("refining_remapper") {
   r->remap_fwd();
 
   // Gather global copies (to make checks easier) and check src/tgt fields
+  auto gs1d_src = all_gather_field(s1d_src,comm);
   auto gs2d_src = all_gather_field(s2d_src,comm);
   auto gv2d_src = all_gather_field(v2d_src,comm);
   auto gs3d_src = all_gather_field(s3d_src,comm);
   auto gv3d_src = all_gather_field(v3d_src,comm);
   auto gbundle_src = all_gather_field(bundle_src,comm);
 
+  auto gs1d_tgt = all_gather_field(s1d_tgt,comm);
   auto gs2d_tgt = all_gather_field(s2d_tgt,comm);
   auto gv2d_tgt = all_gather_field(v2d_tgt,comm);
   auto gs3d_tgt = all_gather_field(s3d_tgt,comm);
   auto gv3d_tgt = all_gather_field(v3d_tgt,comm);
   auto gbundle_tgt = all_gather_field(bundle_tgt,comm);
 
   Real avg;
+  // Scalar 1D
+  {
+    if (comm.am_i_root()) {
+      printf(" -> Checking 1d scalars .........\n");
+    }
+    bool ok = true;
+    gs1d_src.sync_to_host();
+    gs1d_tgt.sync_to_host();
+
+    CHECK (views_are_equal(gs1d_src,gs1d_tgt));
+    ok &= catch_capture.lastAssertionPassed();
+    if (comm.am_i_root()) {
+      printf(" -> Checking 1d scalars ......... %s\n",ok ? "PASS" : "FAIL");
+    }
+  }
+
   // Scalar 2D
   {
     if (comm.am_i_root()) {