EAMxx: Transposed Output

components/eamxx/scripts/compare-nc-files

@@ -49,6 +49,10 @@ OR
     parser.add_argument("-c","--compare",nargs='+', default=[],
                         help="Compare variables from src file against variables from tgt file")
 
+    # Allow variables to be transposed
+    parser.add_argument("-a","--allow_transpose",action="store_true", default=False,
+                        help="Allow comparison if variables are transposed in dimensions")
+
     return parser.parse_args(args[1:])
 
 ###############################################################################
@@ -60,7 +64,8 @@ def _main_func(description):
     print (f" **** Comparing nc files **** \n"
            f"  src file: {pncf._src_file}\n"
            f"  tgt file: {pncf._tgt_file}\n"
-           f"  comparisons: {cmp_str}")
+           f"  comparisons: {cmp_str}\n"
+           f"  allow transpose: {pncf._allow_transpose}")
     success = pncf.run()
 
     print("Comparisons result: {}".format("SUCCESS" if success else "FAIL"))

components/eamxx/scripts/compare_nc_files.py

@@ -1,8 +1,7 @@
-from utils import expect, ensure_netcdf4
+from utils import expect, _ensure_pylib_impl
+_ensure_pylib_impl("xarray")
 
-ensure_netcdf4()
-
-from netCDF4 import Dataset
+import xarray as xr
 import numpy as np
 
 import pathlib
@@ -12,14 +11,15 @@ class CompareNcFiles(object):
 ###############################################################################
 
     ###########################################################################
-    def __init__(self,src_file,tgt_file=None,compare=None):
+    def __init__(self,src_file,tgt_file=None,compare=None,allow_transpose=False):
     ###########################################################################
 
         self._src_file = pathlib.Path(src_file).resolve().absolute()
         expect (self._src_file.exists(),
                 "Error! File '{}' does not exist.".format(self._src_file))
 
         self._compare  = compare
+        self._allow_transpose = allow_transpose
 
         if tgt_file is None:
             self._tgt_file = self._src_file
@@ -62,10 +62,27 @@ def get_name_and_dims(self,name_dims):
     def compare_variables(self):
     ###########################################################################
 
-        ds_src = Dataset(self._src_file,'r')
-        ds_tgt = Dataset(self._tgt_file,'r')
+        ds_src = xr.open_dataset(self._src_file)
+        ds_tgt = xr.open_dataset(self._tgt_file)
 
         success = True
+        if self._compare is None or self._compare == []:
+            self._compare = []
+            # If compare is an empty list, compare all variables
+            print(f"Specific comparison variables not provided,\n"
+                  f"will compare ALL variables in \n"
+                  f"{self._src_file}\n"
+                  f"with\n"
+                  f"{self._tgt_file}\n")
+            for var in ds_src.variables:
+                if var not in ds_tgt.variables:
+                    print (f" Comparison failed! Variable not found.\n"
+                           f"   - var name: {var}\n"
+                           f"   - file name: {self._tgt_file}")
+                    success = False
+                    continue
+                self._compare.append(var+"="+var)
+
         for expr in self._compare:
             # Split the expression, to get the output var name
             tokens = expr.split('=')
@@ -89,11 +106,11 @@ def compare_variables(self):
                        f"   - file name: {self._tgt_file}")
                 success = False
                 continue
-            lvar = ds_src.variables[lname];
-            rvar = ds_tgt.variables[rname];
+            lvar = ds_src[lname];
+            rvar = ds_tgt[rname];
 
-            lvar_rank = len(lvar.dimensions)
-            rvar_rank = len(rvar.dimensions)
+            lvar_rank = len(lvar.dims)
+            rvar_rank = len(rvar.dims)
 
             expect (len(ldims)==0 or len(ldims)==lvar_rank,
                     f"Invalid slice specification for {lname}.\n"
@@ -104,53 +121,33 @@ def compare_variables(self):
                     f"  input request: ({','.join(rdims)})\n"
                     f"  variable rank: {rvar_rank}")
 
-
-            lslices = [[idim,slice] for idim,slice in enumerate(ldims) if slice!=":"]
-            rslices = [[idim,slice] for idim,slice in enumerate(rdims) if slice!=":"]
-
-            lrank = lvar_rank - len(lslices)
-            rrank = rvar_rank - len(rslices)
-
-            if lrank!=rrank:
-                print (f" Comparison failed. Rank mismatch.\n"
-                       f"  - input comparison: {expr}\n"
-                       f"  - upon slicing, rank({lname}) = {lrank}\n"
-                       f"  - upon slicing, rank({rname}) = {rrank}")
-                success = False
-                continue
-
-            lvals = self.slice_variable(lvar,lvar[:],lslices)
-            rvals = self.slice_variable(rvar,rvar[:],rslices)
-
-            if not np.array_equal(lvals,rvals):
-                #  print (f"lvals: {lvals}")
-                #  print (f"rvals: {rvals}")
-                item = np.argwhere(lvals!=rvals)[0]
-                rval = self.slice_variable(rvar,rvals,
-                                           [[idim,slice] for idim,slice in enumerate(item)])
-                lval = self.slice_variable(lvar,lvals,
-                                           [[idim,slice] for idim,slice in enumerate(item)])
-                loc = ",".join([str(i+1) for i in item])
-                print (f" Comparison failed. Values differ.\n"
-                       f"  - input comparison: {expr}\n"
-                       f'  - upon slicing, {lname}({loc}) = {lval}\n'
-                       f'  - upon slicing, {rname}({loc}) = {rval}')
+            lslices = {lvar.dims[idim]:int(slice)-1 for idim,slice in enumerate(ldims) if slice!=":"}
+            rslices = {rvar.dims[idim]:int(slice)-1 for idim,slice in enumerate(rdims) if slice!=":"}
+            lvar_sliced = lvar.sel(lslices)
+            rvar_sliced = rvar.sel(rslices)
+            expect (set(lvar_sliced.dims) == set(rvar_sliced.dims),
+                    f"Error, even when sliced these two elements do not share the same dimensionsn\n"
+                    f"   - left var name   : {lname}\n"
+                    f"   - right var name  : {rname}\n"
+                    f"   - left dimensions : {lvar_sliced.dims}\n"
+                    f"   - right dimensions: {rvar_sliced.dims}\n")
+
+            if self._allow_transpose:
+              rvar_sliced = rvar_sliced.transpose(*lvar_sliced.dims)
+
+            equal = (lvar_sliced.data==rvar_sliced.data).all()
+            if not equal:
+                rse = np.sqrt((lvar_sliced.data-rvar_sliced.data)**2)
+                nonmatch_count = np.count_nonzero(rse)
+                print (f" Comparison failed. Values differ at {nonmatch_count} out of {rse.size} locations.\n"
+                      f"  - input comparison: {expr}\n"
+                      f'  - max L2 error, {rse.max()}\n'
+                      f'  - max L2 location, [{",".join(map(str,(np.array(np.unravel_index(rse.argmax(),rse.shape))+1).tolist()))}]\n'
+                      f'  - dimensions, {lvar_sliced.dims}')
                 success = False
 
         return success
 
-    ###########################################################################
-    def slice_variable(self,var,vals,slices):
-    ###########################################################################
-
-        if len(slices)==0:
-            return vals
-
-        idim, slice_idx = slices.pop(-1)
-        vals = vals.take(int(slice_idx)-1,axis=int(idim))
-
-        return self.slice_variable(var,vals,slices)
-
     ###########################################################################
     def run(self):
     ###########################################################################

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

@@ -359,12 +359,16 @@ void transpose (const Field& src, Field& tgt)
       " - src field layout: " + src_id.get_layout().to_string() + "\n"
       " - tgt field layout: " + tgt_id.get_layout().to_string() + "\n");
 
-  EKAT_REQUIRE_MSG (src_id.get_units()==tgt_id.get_units(),
-      "Error! Input transpose field units are incompatible with src field.\n"
-      " - src field name: " + src.name() + "\n"
-      " - tgt field name: " + tgt.name() + "\n"
-      " - src field units: " + src_id.get_units().get_si_string() + "\n"
-      " - tgt field units: " + tgt_id.get_units().get_si_string() + "\n");
+  // Skip the next check if the units for either field are listed as INVALID
+  using namespace ekat::units;
+  if (src_id.get_units() != Units::invalid() and tgt_id.get_units() != Units::invalid()) {
+    EKAT_REQUIRE_MSG (src_id.get_units()==tgt_id.get_units(),
+        "Error! Input transpose field units are incompatible with src field.\n"
+        " - src field name: " + src.name() + "\n"
+        " - tgt field name: " + tgt.name() + "\n"
+        " - src field units: " + src_id.get_units().get_si_string() + "\n"
+        " - tgt field units: " + tgt_id.get_units().get_si_string() + "\n");
+  }
 
   EKAT_REQUIRE_MSG (src_id.data_type()==tgt_id.data_type(),
       "Error! Input transpose field data type is incompatible with src field.\n"

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

@@ -100,6 +100,11 @@ AtmosphereOutput::AtmosphereOutput(const ekat::Comm &comm, const ekat::Parameter
   // the name of the yaml file where the options are read from.
   m_stream_name = params.name();
 
+  // Is this output set to be transposed?
+  if (params.isParameter("transpose")) {
+    m_transpose = params.get<bool>("transpose");
+  }
+
   auto gm = field_mgr->get_grids_manager();
 
   // Figure out what kind of averaging is requested
@@ -337,7 +342,23 @@ void AtmosphereOutput::init()
 
     // Store the field layout, so that calls to setup_output_file are easier
     const auto& layout = fid.get_layout();
-    m_vars_dims[fname] = get_var_dimnames(layout);
+    m_vars_dims[fname] = get_var_dimnames(m_transpose ? layout.transpose() : layout);
+
+    // Initialize a helper_field for each unique layout.  This can be used for operations
+    // such as writing transposed output.
+    if (m_transpose) {
+      const auto helper_layout = layout.transpose();
+      const std::string helper_name   = "transposed_"+helper_layout.to_string();
+      if (m_helper_fields.find(helper_name) == m_helper_fields.end()) {
+        // We can add a new helper field for this layout
+        using namespace ekat::units;
+        FieldIdentifier fid_helper(helper_name,helper_layout,Units::invalid(),fid.get_grid_name());
+        Field helper(fid_helper);
+        helper.get_header().get_alloc_properties().request_allocation();
+        helper.allocate_view();
+        m_helper_fields[helper_name] = helper;
+      }
+    }
 
     // Now check that all the dims of this field are already set to be registered.
     const auto& tags = layout.tags();
@@ -495,7 +516,15 @@ run (const std::string& filename,
         count.sync_to_host();
 
         auto func_start = std::chrono::steady_clock::now();
-        scorpio::write_var(filename,count.name(),count.get_internal_view_data<int,Host>());
+        if (m_transpose) {
+          const auto& fl = count.get_header().get_identifier().get_layout().to_string();
+          const std::string helper_name   = "transposed_"+fl;
+          auto& temp = m_helper_fields.at(helper_name);
+          transpose(count,temp);
+          scorpio::write_var(filename,count.name(),temp.get_internal_view_data<int,Host>());
+        } else {
+          scorpio::write_var(filename,count.name(),count.get_internal_view_data<int,Host>());
+        }
         auto func_finish = std::chrono::steady_clock::now();
         auto duration_loc = std::chrono::duration_cast<std::chrono::milliseconds>(func_finish - func_start);
         duration_write += duration_loc.count();
@@ -573,7 +602,15 @@ run (const std::string& filename,
 
       // Write to file
       auto func_start = std::chrono::steady_clock::now();
-      scorpio::write_var(filename,field_name,f_out.get_internal_view_data<Real,Host>());
+      if (m_transpose) {
+        const auto& fl = f_out.get_header().get_identifier().get_layout().to_string();
+        const std::string helper_name   = "transposed_"+fl;
+        auto& temp = m_helper_fields.at(helper_name);
+        transpose(f_out,temp);
+        scorpio::write_var(filename,field_name,temp.get_internal_view_data<Real,Host>());
+      } else {
+        scorpio::write_var(filename,field_name,f_out.get_internal_view_data<Real,Host>());
+      }
       auto func_finish = std::chrono::steady_clock::now();
       auto duration_loc = std::chrono::duration_cast<std::chrono::milliseconds>(func_finish - func_start);
       duration_write += duration_loc.count();
@@ -651,7 +688,7 @@ void AtmosphereOutput::set_avg_cnt_tracking(const std::string& name, const Field
   }
 
   // We have not created this avg count field yet.
-  m_vars_dims[avg_cnt_name] = get_var_dimnames(layout);
+  m_vars_dims[avg_cnt_name] = get_var_dimnames(m_transpose ? layout.transpose() : layout);
 
   auto nondim = ekat::units::Units::nondimensional();
   FieldIdentifier count_id (avg_cnt_name,layout,nondim,m_io_grid->name(),DataType::IntType);

components/eamxx/src/share/io/scorpio_output.hpp

@@ -172,6 +172,7 @@ class AtmosphereOutput
 
   // --- Internal variables --- //
   ekat::Comm m_comm;
+  bool m_transpose = false;
 
   // We store separate shared pointers for field mgrs at different stages of IO:
   // More specifically, the order of operations is as follows:
@@ -201,6 +202,7 @@ class AtmosphereOutput
     Scorpio // Output fields to pass to scorpio (may differ from the above in case of packing)
   };
   std::map<Phase, std::shared_ptr<fm_type>> m_field_mgrs;
+  std::map<std::string, Field> m_helper_fields;
 
   std::shared_ptr<const grid_type> m_io_grid;
   std::shared_ptr<remapper_type> m_horiz_remapper;

components/eamxx/src/share/scorpio_interface/eamxx_scorpio_interface.cpp

@@ -749,28 +749,34 @@ void reset_time_dim_len (const std::string& filename, const int new_length)
 void set_var_decomp (PIOVar& var,
                      const std::string& filename)
 {
-  for (size_t i=1; i<var.dims.size(); ++i) {
-    EKAT_REQUIRE_MSG (var.dims[i]->offsets==nullptr,
-        "Error! We currently only allow decomposition on slowest-striding dimension.\n"
-        "       Generalizing is not complicated, but it was not a priority.\n"
+  std::shared_ptr<const PIODim> decomp_dim;
+  for (size_t i=0; i<var.dims.size(); ++i) {
+    if (var.dims[i]->offsets) {
+      EKAT_REQUIRE_MSG (decomp_dim==nullptr,
+        "Error! Call set_var_decomp, more than one dimension set as decomposition dimension.\n"
         " - filename: " + filename + "\n"
-        " - varname : " + var.name + "\n"
-        " - var dims: " + ekat::join(var.dims,get_entity_name,",") + "\n"
-        " - bad dim : " + var.dims[i]->name + "\n");
+        " - varname : " + var.name + "\n");
+      decomp_dim = var.dims[i];
+    }
   }
-  EKAT_REQUIRE_MSG (var.dims[0]->offsets!=nullptr,
-      "Error! Calling set_var_decomp, but the var first dimension does not appear to be decomposed.\n"
-      " - filename: " + filename + "\n"
-      " - varname : " + var.name  + "\n"
-      " - var dims: " + ekat::join(var.dims,get_entity_name,",") + "\n");
+  EKAT_REQUIRE_MSG(decomp_dim!=nullptr,
+    "Error! Calling set_var_decomp, none of the dimensions have been set as the decomposition dimension.\n"
+    " - filename: " + filename + "\n"
+    " - varname : " + var.name + "\n");
   EKAT_REQUIRE_MSG (var.decomp==nullptr,
       "Error! You should have invalidated var.decomp before attempting to reset it.\n"
       " - filename  : " + filename + "\n"
       " - varname   : " + var.name  + "\n"
       " - var decomp: " + var.decomp->name  + "\n");
+  EKAT_REQUIRE_MSG (decomp_dim==var.dims[0] or decomp_dim==var.dims.back(),
+      "Error! Variable decomposition only supports decompostion over the first or last dimension.\n"
+      " - filename  : " + filename + "\n"
+      " - varname   : " + var.name  + "\n"
+      " - var decomp: " + var.decomp->name  + "\n");
+  bool decomp_first_dim = var.dims[0]->offsets ? true : false;
+
 
   // Create decomp name: dtype-dim1<len1>_dim2<len2>_..._dimk<lenN>
-  std::shared_ptr<const PIODim> decomp_dim;
   std::string decomp_tag = var.dtype + "-";
   for (auto d : var.dims) {
     decomp_tag += d->name + "<" + std::to_string(d->length) + ">_";
@@ -800,28 +806,42 @@ void set_var_decomp (PIOVar& var,
     // We haven't create this decomp yet. Go ahead and create one
     decomp = std::make_shared<PIODecomp>();
     decomp->name = decomp_tag;
-    decomp->dim = var.dims[0];
+    decomp->dim = decomp_dim;
 
     int ndims = var.dims.size();
 
     // Get ALL dims global lengths, and compute prod of *non-decomposed* dims
-    std::vector<int> gdimlen = {decomp->dim->length};
+    // Note, if the offsets for the dimension are null then this is by definition
+    // a non-decomposed dimension.
+    std::vector<int> gdimlen = {};
     int non_decomp_dim_prod = 1;
-    for (int idim=1; idim<ndims; ++idim) {
+    for (int idim=0; idim<ndims; ++idim) {
       auto d = var.dims[idim];
       gdimlen.push_back(d->length);
-      non_decomp_dim_prod *= d->length;
+      if (d->offsets==nullptr) {
+        non_decomp_dim_prod *= d->length;
+      }
     }
 
     // Create offsets list
     const auto& dim_offsets = *decomp->dim->offsets;
     int dim_loc_len = dim_offsets.size();
     decomp->offsets.resize (non_decomp_dim_prod*dim_loc_len);
-    for (int idof=0; idof<dim_loc_len; ++idof) {
-      auto dof_offset = dim_offsets[idof];
-      auto beg = decomp->offsets.begin()+ idof*non_decomp_dim_prod;
-      auto end = beg + non_decomp_dim_prod;
-      std::iota (beg,end,non_decomp_dim_prod*dof_offset);
+    if (decomp_first_dim) {
+      for (int idof=0; idof<dim_loc_len; ++idof) {
+        auto dof_offset = dim_offsets[idof];
+        auto beg = decomp->offsets.begin()+ idof*non_decomp_dim_prod;
+        auto end = beg + non_decomp_dim_prod;
+        std::iota (beg,end,non_decomp_dim_prod*dof_offset);
+      }
+    } else {
+      for (int ii = 0; ii<non_decomp_dim_prod; ++ii) {
+        for (int idof=0; idof<dim_loc_len; ++idof) {
+          auto dof_offset = dim_offsets[idof];
+          decomp->offsets[ii * dim_loc_len + idof] =
+              ii * decomp_dim->length + dof_offset;
+        }
+      }
     }
 
     // Create PIO decomp