diff --git a/components/eamxx/scripts/compare-nc-files b/components/eamxx/scripts/compare-nc-files index f8857b530ce5..ad6bf60221af 100755 --- a/components/eamxx/scripts/compare-nc-files +++ b/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")) diff --git a/components/eamxx/scripts/compare_nc_files.py b/components/eamxx/scripts/compare_nc_files.py index 130bc39d359f..13da8250a7fc 100644 --- a/components/eamxx/scripts/compare_nc_files.py +++ b/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,7 +11,7 @@ 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() @@ -20,6 +19,7 @@ def __init__(self,src_file,tgt_file=None,compare=None): "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): ########################################################################### diff --git a/components/eamxx/src/share/field/field_utils.cpp b/components/eamxx/src/share/field/field_utils.cpp index 2d0c390e747c..3af1c56d7538 100644 --- a/components/eamxx/src/share/field/field_utils.cpp +++ b/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" diff --git a/components/eamxx/src/share/io/scorpio_output.cpp b/components/eamxx/src/share/io/scorpio_output.cpp index 4a9af1e9c973..a92f542c3779 100644 --- a/components/eamxx/src/share/io/scorpio_output.cpp +++ b/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("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()); + 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()); + } else { + scorpio::write_var(filename,count.name(),count.get_internal_view_data()); + } auto func_finish = std::chrono::steady_clock::now(); auto duration_loc = std::chrono::duration_cast(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()); + 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()); + } else { + scorpio::write_var(filename,field_name,f_out.get_internal_view_data()); + } auto func_finish = std::chrono::steady_clock::now(); auto duration_loc = std::chrono::duration_cast(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); diff --git a/components/eamxx/src/share/io/scorpio_output.hpp b/components/eamxx/src/share/io/scorpio_output.hpp index 243c032a45a6..84446b546c9d 100644 --- a/components/eamxx/src/share/io/scorpio_output.hpp +++ b/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> m_field_mgrs; + std::map m_helper_fields; std::shared_ptr m_io_grid; std::shared_ptr m_horiz_remapper; diff --git a/components/eamxx/src/share/scorpio_interface/eamxx_scorpio_interface.cpp b/components/eamxx/src/share/scorpio_interface/eamxx_scorpio_interface.cpp index 1a2b36aa0e71..7dd23ab2ce13 100644 --- a/components/eamxx/src/share/scorpio_interface/eamxx_scorpio_interface.cpp +++ b/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; ioffsets==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 decomp_dim; + for (size_t i=0; ioffsets) { + 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_dim2_..._dimk - std::shared_ptr 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(); 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 gdimlen = {decomp->dim->length}; + // Note, if the offsets for the dimension are null then this is by definition + // a non-decomposed dimension. + std::vector gdimlen = {}; int non_decomp_dim_prod = 1; - for (int idim=1; idimlength); - 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; idofoffsets.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; idofoffsets.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; iioffsets[ii * dim_loc_len + idof] = + ii * decomp_dim->length + dof_offset; + } + } } // Create PIO decomp diff --git a/components/eamxx/tests/single-process/shoc/CMakeLists.txt b/components/eamxx/tests/single-process/shoc/CMakeLists.txt index e2ea7e933523..c1a2c5214f4a 100644 --- a/components/eamxx/tests/single-process/shoc/CMakeLists.txt +++ b/components/eamxx/tests/single-process/shoc/CMakeLists.txt @@ -26,8 +26,14 @@ GetInputFile(cam/topo/${EAMxx_tests_TOPO_FILE}) configure_file(${CMAKE_CURRENT_SOURCE_DIR}/input.yaml ${CMAKE_CURRENT_BINARY_DIR}/input.yaml) +set (TRANSPOSED_FLAG False) +set (FPREFIX standalone) configure_file(${CMAKE_CURRENT_SOURCE_DIR}/output.yaml ${CMAKE_CURRENT_BINARY_DIR}/output.yaml) +set (TRANSPOSED_FLAG True) +set (FPREFIX transposed) +configure_file(${CMAKE_CURRENT_SOURCE_DIR}/output.yaml + ${CMAKE_CURRENT_BINARY_DIR}/output_transposed.yaml) # Compare output files produced by npX tests, to ensure they are bfb include (CompareNCFiles) @@ -55,6 +61,20 @@ foreach (NRANKS RANGE ${TEST_RANK_START} ${TEST_RANK_END}) FIXTURES_REQUIRED ${FIXTURES_BASE_NAME}_np${NRANKS}_omp1) endforeach() +# Check transpose test +foreach (NRANKS RANGE ${TEST_RANK_START} ${TEST_RANK_END}) + set (script ${SCREAM_BASE_DIR}/scripts/compare-nc-files) + set (fname_base shoc_standalone_output.INSTANT.nsteps_x1.np${NRANKS}.${RUN_T0}.nc) + set (fname_comp shoc_transposed_output.INSTANT.nsteps_x1.np${NRANKS}.${RUN_T0}.nc) + set (tname shoc_transpose_check_np${NRANKS}) + add_test (NAME ${tname} + COMMAND ${script} -s ${fname_base} -t ${fname_comp} -a + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}) + set_tests_properties (${tname} PROPERTIES + LABELS "shoc;physics" + FIXTURES_REQUIRED ${FIXTURES_BASE_NAME}_np${NRANKS}_omp1) +endforeach() + # Check that the content of sliced vars U/V and corresponding surf_mom_flux # components in all output NC files matches the content of corresponding components # of the vector quantities diff --git a/components/eamxx/tests/single-process/shoc/input.yaml b/components/eamxx/tests/single-process/shoc/input.yaml index 80af0bfd16aa..2f01c2f987df 100644 --- a/components/eamxx/tests/single-process/shoc/input.yaml +++ b/components/eamxx/tests/single-process/shoc/input.yaml @@ -46,5 +46,5 @@ initial_conditions: # The parameters for I/O control scorpio: - output_yaml_files: ["output.yaml"] + output_yaml_files: ["output.yaml", "output_transposed.yaml"] ... diff --git a/components/eamxx/tests/single-process/shoc/output.yaml b/components/eamxx/tests/single-process/shoc/output.yaml index 3a86ba3f9f25..3b6be797caeb 100644 --- a/components/eamxx/tests/single-process/shoc/output.yaml +++ b/components/eamxx/tests/single-process/shoc/output.yaml @@ -1,7 +1,8 @@ %YAML 1.1 --- -filename_prefix: shoc_standalone_output +filename_prefix: shoc_${FPREFIX}_output averaging_type: instant +transpose: ${TRANSPOSED_FLAG} fields: physics: field_names: