Merge branch 'copilot/support-latlon-input-class' (PR #8472)

Fix FieldReader lat/lon decomposition for point-grid reads

[BFB]

Author: bartgol

components/eamxx/src/share/field/field_reader.cpp

@@ -143,58 +143,92 @@ void FieldReader::set_dim_decomp(const Field& gids, const ekat::Comm& comm)
 
 void FieldReader::setup_internals ()
 {
-  if (m_dim_decomp_name!="" and m_reader_state & (NEW_FILE | NEW_DECOMP)) {
-    if (std::is_same_v<std::int64_t,scorpio::offset_t>) {
-      scorpio::set_dim_decomp(m_filename, m_dim_decomp_name, m_dim_decomp_offsets);
-    } else {
-      std::vector<scorpio::offset_t> offsets(m_dim_decomp_offsets.begin(),
-                                             m_dim_decomp_offsets.end());
-      scorpio::set_dim_decomp(m_filename, m_dim_decomp_name, offsets);
-    }
+  using namespace ShortFieldTagsNames;
+
+  auto name_lat = m_tag_rename.count("lat")>0 ? m_tag_rename["lat"] : "lat";
+  auto name_lon = m_tag_rename.count("lon")>0 ? m_tag_rename["lon"] : "lon";
+  auto name_col = m_tag_rename.count("col")>0 ? m_tag_rename["col"] : "ncol";
+
+  // 1. Check if lat/lon dims are in the file
+  bool latlon_in_file = scorpio::has_dim(m_filename,name_lat) and scorpio::has_dim(m_filename,name_lon);
+
+  if (m_reader_state & NEW_FILE) {
+    EKAT_REQUIRE_MSG (not (latlon_in_file and scorpio::has_dim(m_filename,name_col)),
+        "Error! We cannot read from files that mix lat/lon vars and ncol vars.\n"
+        " - file name: " + m_filename + "\n");
   }
 
+  // 2. Check fields are in the file, with correct dim names
+  bool latlon_decomp = false;
   if (m_reader_state & (NEW_FIELDS | NEW_FILE)) {
-    // First, check fields are in the file, with correct layout
     for (const auto & f : m_fields) {
       // Check that the variable is in the file.
       EKAT_REQUIRE_MSG (scorpio::has_var(m_filename,f.name()),
           "Error! Input file does not store a required variable.\n"
           " - filename: " + m_filename + "\n"
           " - varname : " + f.name() + "\n");
 
-      const auto& layout = f.get_header().get_identifier().get_layout();
-      auto io_dims = layout.names();
-      for (int i=0; i<layout.rank(); ++i) {
-        if (io_dims[i]=="dim")
-          io_dims[i] += std::to_string(layout.dim(i));
-        if (m_tag_rename.count(io_dims[i])>0) {
-          io_dims[i] = m_tag_rename.at(io_dims[i]);
-        }
+      const auto& fl = f.get_header().get_identifier().get_layout();
+      const auto& var = scorpio::get_var(m_filename,f.name());
+
+      auto f_dims = fl.names();
+      auto f_extents = fl.dims();
+      auto var_dims = var.dim_names();
+
+      if (fl.has_tag(COL) and latlon_in_file) {
+        auto idx = fl.dim_idx(COL);
+        f_extents.erase(f_extents.begin()+idx);
+        ekat::erase(f_dims,name_col);
+        ekat::erase(var_dims,name_col);
+        ekat::erase(var_dims,name_lat);
+        ekat::erase(var_dims,name_lon);
+        latlon_decomp |= m_dim_decomp_name==fl.name(idx);
       }
 
-      const auto& var = scorpio::get_var(m_filename,f.name());
-      EKAT_REQUIRE_MSG (var.dim_names()==io_dims,
+      EKAT_REQUIRE_MSG (var_dims==f_dims,
           "Error! Layout mismatch for input file variable.\n"
           " - filename: " + m_filename + "\n"
           " - varname : " + f.name() + "\n"
-          " - expected dim names: " + ekat::join(io_dims,",") + "\n"
-          " - dims from file    : " + ekat::join(var.dim_names(),",") + "\n");
+          " - field dim names: " + ekat::join(fl.names(),",") + "\n"
+          " - file var dim names: " + ekat::join(var.dim_names(),",") + "\n");
 
-      for (int i=0; i<layout.rank(); ++i) {
-        if (io_dims[i]!=m_dim_decomp_name) {
-          const int file_len  = scorpio::get_dimlen(m_filename,io_dims[i]);
-          EKAT_REQUIRE_MSG (layout.dim(i)==file_len,
+      for (auto [f_dim,var_dim,f_ext] : ekat::zip(f_dims,var_dims,f_extents)) {
+        if (f_dim!=m_dim_decomp_name) {
+          const int file_len  = scorpio::get_dimlen(m_filename,f_dim);
+          EKAT_REQUIRE_MSG (f_ext==file_len,
               "Error! Dimension mismatch for input file variable.\n"
             " - filename : " + m_filename + "\n"
             " - varname  : " + f.name() + "\n"
-            " - var dims : " + ekat::join(io_dims,",") + "\n"
-            " - dim name : " + io_dims[i] + "\n"
-            " - expected extent : " + std::to_string(layout.dim(i)) + "\n"
+            " - var dims : " + ekat::join(var_dims,",") + "\n"
+            " - dim name : " + var_dim + "\n"
+            " - expected extent : " + std::to_string(f_ext) + "\n"
             " - extent from file: " + std::to_string(file_len) + "\n");
         }
       }
     }
+  }
 
+  // 3. Set the decomp
+  if (m_dim_decomp_name!="" and m_reader_state & (NEW_FILE | NEW_FIELDS | NEW_DECOMP)) {
+    if (std::is_same_v<std::int64_t,scorpio::offset_t>) {
+      if (latlon_decomp) {
+        scorpio::set_dims_decomp(m_filename,{name_lat,name_lon},m_dim_decomp_offsets);
+      } else {
+        scorpio::set_dim_decomp(m_filename, m_dim_decomp_name, m_dim_decomp_offsets);
+      }
+    } else {
+      std::vector<scorpio::offset_t> offsets(m_dim_decomp_offsets.begin(),
+                                             m_dim_decomp_offsets.end());
+      if (latlon_decomp) {
+        scorpio::set_dims_decomp(m_filename,{name_lat,name_lon},offsets);
+      } else {
+        scorpio::set_dim_decomp(m_filename, m_dim_decomp_name, offsets);
+      }
+    }
+  }
+
+  // 4. Create IO fields
+  if (m_reader_state & (NEW_FIELDS | NEW_FILE)) {
     // Now create io_fields
     // IO fields MUST be contiguous, so if padded or with parent, create a new one.
     // Otherwise, we can alias input field, possibly with renamed tags

components/eamxx/src/share/field/tests/field_reader_tests.cpp

@@ -140,7 +140,7 @@ TEST_CASE ("read_fields_no_decomp")
 //  Test: decomposed read (distributed column dimension)
 // ============================================================ //
 
-TEST_CASE ("read_fields_with_decomp", "[field_utils][io]")
+TEST_CASE ("read_fields_with_decomp")
 {
   using namespace ShortFieldTagsNames;
   using namespace ekat::units;
@@ -236,6 +236,101 @@ TEST_CASE ("read_fields_with_decomp", "[field_utils][io]")
   scorpio::finalize_subsystem();
 }
 
+// ============================================================ //
+//  Test: decomposed read from lat/lon file into point-grid fields
+// ============================================================ //
+
+TEST_CASE ("read_fields_with_latlon_decomp", "[field_utils][io]")
+{
+  using namespace ShortFieldTagsNames;
+
+  ekat::Comm comm(MPI_COMM_WORLD);
+  scorpio::init_subsystem(comm);
+
+  const std::string filename =
+      "field_io_test_latlon_decomp_np" + std::to_string(comm.size()) + ".nc";
+
+  const int lcols_per_rank = 3;
+  const int nlat           = comm.size();
+  const int nlon           = lcols_per_rank;
+  const int nlevs          = 4;
+
+  FieldLayout lay2d({COL, LEV}, {lcols_per_rank, nlevs});
+  FieldLayout lay1d({COL},      {lcols_per_rank});
+
+  auto my_gids = make_field("gids",lay1d,DataType::IntType);
+  auto my_gids_beg = my_gids.get_internal_view_data<int,Host>();
+  auto my_gids_end = my_gids_beg + lcols_per_rank;
+  std::iota(my_gids_beg,my_gids_end,1 + comm.rank() * lcols_per_rank);
+  my_gids.sync_to_dev();
+
+  // ---- Write phase ---- //
+  {
+    scorpio::register_file(filename, scorpio::Write);
+    scorpio::define_dim(filename, "lat", nlat);
+    scorpio::define_dim(filename, "lon", nlon);
+    scorpio::define_dim(filename, "lev", nlevs);
+
+    std::vector<scorpio::offset_t> offsets(lcols_per_rank);
+    for (int i = 0; i < lcols_per_rank; ++i) {
+      offsets[i] = my_gids_beg[i] - 1;
+    }
+    scorpio::set_dims_decomp(filename, {"lat","lon"}, offsets);
+
+    scorpio::define_var(filename, "f2d", {"lat", "lon", "lev"}, "real", false);
+    scorpio::define_var(filename, "f1d", {"lat", "lon"},        "int",  false);
+    scorpio::enddef(filename);
+
+    Field f2d = make_field("f2d", lay2d);
+    Field f1d = make_field("f1d", lay1d, DataType::IntType);
+    f2d.sync_to_host();
+    f1d.sync_to_host();
+
+    auto v2d = f2d.get_view<Real**, Host>();
+    auto v1d = f1d.get_view<int*,  Host>();
+    for (int li = 0; li < lcols_per_rank; ++li) {
+      const int gcol = my_gids_beg[li];
+      v1d(li) = gcol;
+      for (int j = 0; j < nlevs; ++j) {
+        v2d(li, j) = Real(gcol * nlevs + j);
+      }
+    }
+    f2d.sync_to_dev();
+    f1d.sync_to_dev();
+
+    write_field_to_file<Real>(filename, f2d);
+    write_field_to_file<int>(filename, f1d);
+
+    scorpio::release_file(filename);
+  }
+
+  // ---- Read phase ---- //
+  {
+    Field f2d = make_field("f2d", lay2d);
+    Field f1d = make_field("f1d", lay1d, DataType::IntType);
+    f2d.deep_copy(-1);
+    f1d.deep_copy(-1);
+
+    read_fields(filename, {f2d, f1d}, my_gids, comm);
+
+    f2d.sync_to_host();
+    f1d.sync_to_host();
+
+    auto v2d = f2d.get_view<const Real**, Host>();
+    auto v1d = f1d.get_view<const int*,  Host>();
+
+    for (int li = 0; li < lcols_per_rank; ++li) {
+      const int gcol = my_gids_beg[li];
+      REQUIRE (v1d(li) == gcol);
+      for (int j = 0; j < nlevs; ++j) {
+        REQUIRE (v2d(li, j) == Real(gcol * nlevs + j));
+      }
+    }
+  }
+
+  scorpio::finalize_subsystem();
+}
+
 // ============================================================ //
 //  Test: reading a specific time slice
 // ============================================================ //