diff --git a/data/cesm_regrid_maps.yaml b/data/cesm_grids.yaml similarity index 50% rename from data/cesm_regrid_maps.yaml rename to data/cesm_grids.yaml index 330b89a..dd3a48e 100644 --- a/data/cesm_regrid_maps.yaml +++ b/data/cesm_grids.yaml @@ -1,5 +1,5 @@ # ESMF regrid weight files for CESM, keyed by resolution. -# See noresm_regrid_maps.yaml for the schema. +# See noresm_grids.yaml for the schema. inputdata_dir: /glade/campaign/cesm/cesmdata/inputdata/ @@ -11,35 +11,39 @@ resolutions: conservative: cpl/gridmaps/tx2_3v2/map_t232_TO_1x1d_aave.251023.nc bilinear: cpl/gridmaps/tx2_3v2/map_t232_TO_1x1d_blin.251023.nc - # 'regular' means the variable already carries lat/lon, so nothing is - # regridded (see the resolution == "regular" branch in regrid.py) and no + # 'custom' means the variable already carries lat/lon, so nothing is + # regridded (see the resolution == "custom" branch in regrid.py) and no # weights are ever applied. The entry exists only because _regrid_fx_once # builds a regridder before checking whether it needs one, so the file must # exist to be opened and discarded. These are the maps the old implicit # fallback used; they are named here to keep that behaviour visible. - regular: + custom: conservative: cpl/gridmaps/tx2_3v2/map_t232_TO_1x1d_aave.251023.nc bilinear: cpl/gridmaps/tx2_3v2/map_t232_TO_1x1d_blin.251023.nc +# Per-realm grids; see noresm_grids.yaml for the fields and a template. +# +# No ne16 entry: 'resolutions' above defines no ne16 weights for CESM, so an +# ne16 CESM run could never have regridded anything. +# +# The ocean, sea-ice and land-ice grid_labels are carried over unchanged and +# still say g106 (regular 1 x 1), which cannot be right for a tripolar or a +# projected grid. They are left alone because this change is about where the +# input grid is stated, not about correcting CESM's labels. grid_names_per_realm: - ne16: - atmos: g123 - aerosol: g123 - atmosChem: g123 - land: g123 - seaIce: g202 - landIce: g194 ne30: - atmos: g106 - aerosol: g106 - atmosChem: g106 - land: g106 - seaIce: g106 - landIce: g106 + atmos: {input_grid: ne30, grid_label: g106} + aerosol: {input_grid: ne30, grid_label: g106} + atmosChem: {input_grid: ne30, grid_label: g106} + land: {input_grid: ne30, grid_label: g106} + ocean: {input_grid: tx2_3v2, grid_label: g106} + ocnBgchem: {input_grid: tx2_3v2, grid_label: g106} + seaIce: {input_grid: tx2_3v2, grid_label: g106} + landIce: {input_grid: custom, grid_label: g106} # Static dataset (global attribute) metadata for the CMOR dataset JSON builder; -# see noresm_regrid_maps.yaml for the schema. CESM writes CESM3 at every +# see noresm_grids.yaml for the schema. CESM writes CESM3 at every # resolution for now (placeholder). dataset: base_source_id: CESM3 diff --git a/data/intensive_vars.yaml b/data/intensive_vars.yaml index 0f48c52..2d0089f 100644 --- a/data/intensive_vars.yaml +++ b/data/intensive_vars.yaml @@ -5,7 +5,7 @@ # Fluxes must conserve their integral over a cell and use the conservative map, # which is the default for anything not listed here. # -# Unlike the include-pattern and regrid-map tables, this file is NOT per model: +# Unlike the include-pattern and grid tables, this file is NOT per model: # whether a quantity is intensive is a property of the variable, not of the # model that produced it. Keeping one copy stops CESM and NorESM output from # diverging in how the same variable was regridded. diff --git a/data/noresm_grids.yaml b/data/noresm_grids.yaml new file mode 100644 index 0000000..30dd7cb --- /dev/null +++ b/data/noresm_grids.yaml @@ -0,0 +1,101 @@ +# ESMF regrid weight files for NorESM, keyed by resolution. +# +# Paths are relative to 'inputdata_dir'. Each resolution supplies a +# conservative map and, optionally, a bilinear one; regrid.py uses the +# bilinear map for the variables listed in intensive_vars.yaml and the +# conservative map for everything else, including all fx fields. +# +# Both maps for a resolution must target the SAME destination grid, or a +# single run emits variables on two different grids. + +inputdata_dir: /nird/datalake/NS9560K/diagnostics/land_xesmf_diag_data/ + +resolutions: + ne30: + conservative: map_ne30pg3_to_1x1d_aave.nc + bilinear: map_ne30pg3_to_1x1d_bilin.nc + ne16: + conservative: map_ne16pg3_to_2x2_aave_c20260904.nc + bilinear: map_ne16pg3_to_2x2_blin_c20260904.nc + tnx1v4: + conservative: map_tnx1v4_to_1x1_aave_c260531.nc + bilinear: map_tnx1v4_to_1x1_blin_c260531.nc + + # 'custom' means the variable already carries lat/lon, so nothing is + # regridded (see the resolution == "custom" branch in regrid.py) and no + # weights are ever applied. The entry exists only because _regrid_fx_once + # builds a regridder before checking whether it needs one, so the file must + # exist to be opened and discarded. These are the maps the old implicit + # fallback used; they are named here to keep that behaviour visible. + custom: + conservative: map_tnx1v4_to_1x1_aave_c260531.nc + bilinear: map_tnx1v4_to_1x1_blin_c260531.nc + +# Per-realm grids, for each resolution the case may have been run at. +# +# input_grid the grid the realm's history files are on, naming an entry in +# 'resolutions' above. 'custom' means nothing is regridded: +# CISM land-ice output is georeferenced from its projected x/y +# coordinates instead. +# grid_label the grid the output lands on, as a CMIP7 grid code. +# +# Only the atmosphere and land change with the resolution. The ocean, sea-ice +# and land-ice rows are deliberately identical between ne16 and ne30, because +# those components run the same grids in both configurations. +# +grid_names_per_realm: + NorESM3-LM: + atmos: {input_grid: ne16, grid_label: g123} + aerosol: {input_grid: ne16, grid_label: g123} + atmosChem: {input_grid: ne16, grid_label: g123} + land: {input_grid: ne16, grid_label: g123} + ocean: {input_grid: tnx1v4, grid_label: g202} + ocnBgchem: {input_grid: tnx1v4, grid_label: g202} + seaIce: {input_grid: tnx1v4, grid_label: g202} + landIce: {input_grid: custom, grid_label: g194} + NorESM3-MM: + atmos: {input_grid: ne30, grid_label: g106} + aerosol: {input_grid: ne30, grid_label: g106} + atmosChem: {input_grid: ne30, grid_label: g106} + land: {input_grid: ne30, grid_label: g106} + ocean: {input_grid: tnx1v4, grid_label: g202} + ocnBgchem: {input_grid: tnx1v4, grid_label: g202} + seaIce: {input_grid: tnx1v4, grid_label: g202} + landIce: {input_grid: custom, grid_label: g194} + + # To add a resolution, uncomment and fill in every field for the grids that + # case ran. Name it 'custom' to run a case this table does not otherwise + # describe; a 'custom' block is used like any other, and without one nothing + # is regridded. Any grid may be named, including new ones, so long as each + # has an entry in 'resolutions' above. Nothing is inherited from the blocks above: a + # resolution running a different ocean or ice-sheet grid says so here. Each + # input_grid needs an entry in 'resolutions', and each grid_label must be + # registered in the controlled vocabulary, or CMOR refuses to write. + # + # : + # atmos: {input_grid: , grid_label: } + # aerosol: {input_grid: , grid_label: } + # atmosChem: {input_grid: , grid_label: } + # land: {input_grid: , grid_label: } + # ocean: {input_grid: , grid_label: } + # ocnBgchem: {input_grid: , grid_label: } + # seaIce: {input_grid: , grid_label: } + # landIce: {input_grid: , grid_label: } + +# Static dataset (global attribute) metadata for the CMOR dataset JSON builder. +# base_source_id is the key looked up in cmor-cvs.json for institution_id and +# source text; it must be a CV source_id entry. source_ids/nominal_resolution +# are the branded, per-resolution values actually written to the output; the +# default_* keys catch any resolution not listed (e.g. non-CMIP test runs). +dataset: + base_source_id: NorESM3 + source_type: AOGCM + calendar: noleap + source_ids: + NorESM3-LM: NorESM3-LM + NorESM3-MM: NorESM3-MM + default_source_id: NorESM3 + nominal_resolution: + NorESM3-LM: 250 km + NorESM3-MM: 100 km + default_nominal_resolution: 100 km diff --git a/data/noresm_regrid_maps.yaml b/data/noresm_regrid_maps.yaml deleted file mode 100644 index 1240ba1..0000000 --- a/data/noresm_regrid_maps.yaml +++ /dev/null @@ -1,66 +0,0 @@ -# ESMF regrid weight files for NorESM, keyed by resolution. -# -# Paths are relative to 'inputdata_dir'. Each resolution supplies a -# conservative map and, optionally, a bilinear one; regrid.py uses the -# bilinear map for the variables listed in intensive_vars.yaml and the -# conservative map for everything else, including all fx fields. -# -# Both maps for a resolution must target the SAME destination grid, or a -# single run emits variables on two different grids. - -inputdata_dir: /nird/datalake/NS9560K/diagnostics/land_xesmf_diag_data/ - -resolutions: - ne30: - conservative: map_ne30pg3_to_1x1d_aave.nc - bilinear: map_ne30pg3_to_1x1d_bilin.nc - ne16: - conservative: map_ne16pg3_to_2x2_aave_c20260904.nc - bilinear: map_ne16pg3_to_2x2_blin_c20260904.nc - tnx1v4: - conservative: map_tnx1v4_to_1x1_aave_c260531.nc - bilinear: map_tnx1v4_to_1x1_blin_c260531.nc - - # 'regular' means the variable already carries lat/lon, so nothing is - # regridded (see the resolution == "regular" branch in regrid.py) and no - # weights are ever applied. The entry exists only because _regrid_fx_once - # builds a regridder before checking whether it needs one, so the file must - # exist to be opened and discarded. These are the maps the old implicit - # fallback used; they are named here to keep that behaviour visible. - regular: - conservative: map_tnx1v4_to_1x1_aave_c260531.nc - bilinear: map_tnx1v4_to_1x1_blin_c260531.nc - -grid_names_per_realm: - ne16: - atmos: g123 - aerosol: g123 - atmosChem: g123 - land: g123 - seaIce: g202 - landIce: g194 - ne30: - atmos: g106 - aerosol: g106 - atmosChem: g106 - land: g106 - seaIce: g106 - landIce: g106 - -# Static dataset (global attribute) metadata for the CMOR dataset JSON builder. -# base_source_id is the key looked up in cmor-cvs.json for institution_id and -# source text; it must be a CV source_id entry. source_ids/nominal_resolution -# are the branded, per-resolution values actually written to the output; the -# default_* keys catch any resolution not listed (e.g. non-CMIP test runs). -dataset: - base_source_id: NorESM3 - source_type: AOGCM - calendar: noleap - source_ids: - ne16: NorESM3-LM - ne30: NorESM3-MM - default_source_id: NorESM3 - nominal_resolution: - ne16: 250 km - ne30: 100 km - default_nominal_resolution: 100 km diff --git a/scripts/cmor_driver.py b/scripts/cmor_driver.py index 7f78214..ac3c7f0 100755 --- a/scripts/cmor_driver.py +++ b/scripts/cmor_driver.py @@ -53,7 +53,7 @@ patterns_for_variable, ) from cmip7_prep.mapping_compat import Mapping -from cmip7_prep.grids import ATM_RESOLUTIONS, resolution_for +from cmip7_prep.grids import MODEL_RESOLUTIONS, CICE_GRID_VARS, resolution_for from cmip7_prep.regrid import zonal_mean_on_pressure_grid, regrid_to_latlon_ds from cmip7_prep.pipeline import ( realize_regrid_prepare, @@ -146,6 +146,17 @@ def parse_args(): "gen_timeseries.py or another time series utility" ), ) + required.add_argument( + "--model-res", + type=str, + choices=list(MODEL_RESOLUTIONS), + required=True, + help=( + "Model resolution, as named in data/_grids.yaml. It selects " + "the grid each realm's input is on and the grid label its output " + "carries, so every realm needs it." + ), + ) required.add_argument( "--experiment", type=str, @@ -155,18 +166,6 @@ def parse_args(): case = parser.add_argument_group("other properties of the case") - case.add_argument( - "--atmos-res", - type=str, - choices=list(ATM_RESOLUTIONS), - default=None, - help=( - "Grid the atmosphere and land were run on. Required when --realm " - "is atmos, atmosChem, aerosol or land; ignored otherwise, since " - "ocean and sea ice are on the model's own tripolar grid and land " - "ice is written on its native projected grid." - ), - ) case.add_argument( "--ocn-static-file", type=str, @@ -369,9 +368,15 @@ def _prepare_seaice_native(mapping, ds_native, varname, frequency): means carry Northern and Southern Hemisphere variants that become separate published datasets. - TLAT/TLON ride along as coordinates, but the *_bounds variables are data - variables and would be dropped by the realize_all projection, so they are - reassigned for the 2-D variants that need them. + realize_all hands back one variable at a time, and a new dataset is built + around it. The lat/lon arrays come along, because xarray carries + coordinates with a variable. The bounds arrays are ordinary variables + rather than coordinates, so they stay behind in the original dataset and + have to be copied over here. Only variants that are still maps need them; + the hemispheric sums are single numbers. + + All of CICE's lat/lon and bounds arrays are copied, not just the T ones, so + that for example siu and siv keep ULAT/ULON. Returns (cmor_items, status). """ @@ -381,9 +386,19 @@ def _prepare_seaice_native(mapping, ds_native, varname, frequency): if "time_bounds" in ds_native and "time_bounds" not in ds_v: ds_v = ds_v.assign(time_bounds=ds_native["time_bounds"]) if "nj" in da.dims and "ni" in da.dims: - for gname in ("TLAT", "TLON", "latt_bounds", "lont_bounds"): + for gname in CICE_GRID_VARS: if gname in ds_native and gname not in ds_v: ds_v = ds_v.assign({gname: ds_native[gname]}) + # realize_all need not preserve attributes, and without this the + # writer cannot tell which grid point the variable is on. + native = ds_native.get(varname) + # xarray moves this attribute into encoding when it decodes + # coordinates, so both places are checked. + coords_attr = getattr(native, "attrs", {}).get("coordinates") or getattr( + native, "encoding", {} + ).get("coordinates") + if coords_attr and "coordinates" not in ds_v[varname].attrs: + ds_v[varname].attrs["coordinates"] = coords_attr cmor_items.append((ds_v, variant_cfg)) return ( cmor_items, @@ -487,6 +502,7 @@ def process_one_var( tables_root, outdir, resolution, + case_resolution, model, realm="atmos", frequency="mon", @@ -495,7 +511,12 @@ def process_one_var( ice_sheet=None, experiment=None, ) -> list[tuple[str, str]]: - """Compute+write one CMIP variable. Returns a list of (varname, 'ok' or error message) tuples.""" + """Compute+write one CMIP variable. Returns a list of (varname, 'ok' or error message) tuples. + + ``resolution`` is the grid this realm regrids from; ``case_resolution`` is + the resolution the case was run at, which identifies it for the metadata. + They are the same string only for the atmosphere and land realms. + """ varname = cmip_var.branded_variable_name.name # At this point you have a cmip_var (metadata from database query for the target variable) @@ -640,9 +661,15 @@ def process_one_var( # JSON files. frequency, variant indices, experiment metadata # and the per-realm grid_label are already baked in, so only the # per-variable region is set below. + # case_resolution, not resolution: source_id, + # nominal_resolution and grid_label are keyed by the resolution + # the case was run at (ne16 is NorESM3-LM at 250 km, ne30 is + # NorESM3-MM at 100 km), while resolution is the grid this realm + # regrids from. Passing the latter (tnx1v4 for sea ice) matches + # nothing, and the dataset is labelled with the fallbacks. dataset_cfg = build_dataset_cfg( model=model, - resolution=resolution, + resolution=case_resolution, experiment=experiment, frequency=frequency, ripf=ripf_index, @@ -721,7 +748,7 @@ def main(): OUTDIR = args.outdir # The grid to regrid from follows from the model and realm; see grids.py. try: - resolution = resolution_for(args.model, args.realm, args.atmos_res) + resolution = resolution_for(args.model, args.realm, args.model_res) except ValueError as exc: logger.error("%s", exc) sys.exit(2) @@ -729,7 +756,7 @@ def main(): "input grid for realm %s is %s (atmosphere/land ran at %s)", args.realm, resolution, - args.atmos_res, + args.model_res, ) model = args.model frequency = args.frequency @@ -992,6 +1019,7 @@ def main(): tables_root, OUTDIR, resolution, + args.model_res, model, realm=realm, frequency=frequency, @@ -1012,6 +1040,7 @@ def main(): tables_root, OUTDIR, resolution, + args.model_res, model, realm=realm, frequency=frequency, diff --git a/scripts/run_reference_case.py b/scripts/run_reference_case.py index eb27e90..1f69575 100644 --- a/scripts/run_reference_case.py +++ b/scripts/run_reference_case.py @@ -35,7 +35,7 @@ # pylint: disable=wrong-import-position from cmip7_prep.cv_lookup import validate as validate_cv -from cmip7_prep.grids import ATM_RESOLUTIONS +from cmip7_prep.grids import MODEL_RESOLUTIONS from cmip7_prep.reference_run import STAGES, Plan, Step, build_plan logger = logging.getLogger("run_reference_case") @@ -73,6 +73,16 @@ def parse_arguments(): "The wrong one produces plausible but wrong output." ), ) + required.add_argument( + "--model-res", + choices=list(MODEL_RESOLUTIONS), + required=True, + help=( + "Model resolution, as named in data/_grids.yaml. It selects " + "the grid each realm's input is on and the grid label its output " + "carries, so every realm needs it." + ), + ) required.add_argument( "--experiment", required=True, @@ -83,16 +93,6 @@ def parse_arguments(): ) selection = parser.add_argument_group("selecting what to run") - selection.add_argument( - "--atmos-res", - choices=list(ATM_RESOLUTIONS), - default=None, - help=( - "Grid the atmosphere and land were run on. Required unless the " - "realms being processed are all ocean, sea ice or land ice, which " - "derive their own grid." - ), - ) selection.add_argument( "--realms", nargs="+", @@ -351,7 +351,7 @@ def main(): frequencies=args.frequencies, years=args.years, stages=args.stages, - atmos_res=args.atmos_res, + model_res=args.model_res, experiment=args.experiment, workers=args.workers, ice_sheet=args.ice_sheet, diff --git a/src/cmip7_prep/cmor_utils.py b/src/cmip7_prep/cmor_utils.py index 29847ed..e50d4f1 100644 --- a/src/cmip7_prep/cmor_utils.py +++ b/src/cmip7_prep/cmor_utils.py @@ -16,7 +16,7 @@ import xarray as xr # Local imports avoid an import cycle with regrid_maps. -from .regrid_maps import load_regrid_maps, get_grid_names +from .regrid_maps import load_regrid_maps, get_realm_row _FILL_DEFAULT = 1.0e20 @@ -134,17 +134,17 @@ def build_dataset_cfg( Institutional, source, license and experiment metadata are read from the controlled vocabulary (cmor-cvs.json); the branded source_id and nominal resolution come from the per-model ``dataset`` block in - ``data/_regrid_maps.yaml``. Replaces the packaged cmor_dataset*.json. + ``data/_grids.yaml``. Replaces the packaged cmor_dataset*.json. """ cv = _load_controlled_vocabulary(str(tables_root)) meta = load_regrid_maps(model).get("dataset") or {} base_source_id = meta.get("base_source_id") - # grid_label is the per-realm grid code from the regrid-map table; its + # grid_label is the per-realm grid code from the model grid table; its # human-readable description comes from the CV keyed by that same code. try: - grid_label = get_grid_names(model, resolution, realm) - except ValueError: + grid_label = get_realm_row(model, resolution, realm)["grid_label"] + except (ValueError, KeyError): logger.warning( "No grid name for model=%s resolution=%s realm=%s; defaulting " "grid_label to 'gr'", diff --git a/src/cmip7_prep/cmor_writer.py b/src/cmip7_prep/cmor_writer.py index f275473..501eb4e 100644 --- a/src/cmip7_prep/cmor_writer.py +++ b/src/cmip7_prep/cmor_writer.py @@ -22,6 +22,7 @@ import numpy as np import xarray as xr +from .grids import CICE_BOUNDS_BY_COORD from .cmor_utils import ( get_cmor_attr, set_cmor_attr, @@ -121,6 +122,80 @@ def period_in_axis_units(dataset, units: str) -> float | None: return float(count) +def is_latitude(da) -> bool: + """Return whether a variable is a latitude coordinate.""" + attrs = getattr(da, "attrs", {}) + return attrs.get("standard_name") == "latitude" or str( + attrs.get("units", "") + ).startswith("degrees_north") + + +def is_longitude(da) -> bool: + """Return whether a variable is a longitude coordinate.""" + attrs = getattr(da, "attrs", {}) + return attrs.get("standard_name") == "longitude" or str( + attrs.get("units", "") + ).startswith("degrees_east") + + +def horizontal_coords_of(dataset, var_da, fallback_lat, fallback_lon): + """Return the latitude and longitude a variable is actually defined on. + + This matters on a staggered grid: CICE writes velocities on the B-grid + velocity point, so siu and siv are on ULAT/ULON while the thermodynamic + fields are on TLAT/TLON. Taking the centre for everything put the + velocities in the wrong place. + + The variable's own coordinates are used, wherever they are to be found -- + attached to it, or named in a ``coordinates`` attribute. Which of the two + depends on how the file was opened. They are identified by their CF + attributes rather than by name, so the names need not be known here, and + the given fallback is used when neither yields a pair. + """ + # The 'coordinates' attribute is what distinguishes the grid points. The + # coordinates xarray attaches cannot: TLAT and ULAT have the same dimensions, + # so every variable on (nj, ni) carries both. xarray moves the attribute + # into encoding when it decodes coordinates, so look in both places. + named = str( + getattr(var_da, "attrs", {}).get("coordinates") + or getattr(var_da, "encoding", {}).get("coordinates") + or "" + ).split() + candidates = [dataset[name] for name in named if name in dataset] + + # Only if the variable names none: anything attached will at least be a + # coordinate of the right shape. + if not candidates: + candidates = list(getattr(var_da, "coords", {}).values()) + + lat = next((da for da in candidates if is_latitude(da)), None) + lon = next((da for da in candidates if is_longitude(da)), None) + if lat is None or lon is None: + return fallback_lat, fallback_lon + return lat, lon + + +def vertex_bounds_of(dataset, coord, coord_name: str): + """Return the vertex bounds of a coordinate, or None if the file has none. + + The coordinate's own ``bounds`` attribute is used when present. Otherwise a + name built from the coordinate is tried -- CICE writes ULAT's bounds as + ``latu_bounds`` -- and nothing else: bounds from another grid point would + pair cell corners with the wrong cell centres. + """ + named = getattr(coord, "attrs", {}).get("bounds") + if isinstance(named, str) and named in dataset: + return dataset[named] + for candidate in ( + CICE_BOUNDS_BY_COORD.get(coord_name.upper()), + f"{coord_name}_bnds", + f"{coord_name.lower()}_bounds", + ): + if candidate and candidate in dataset: + return dataset[candidate] + return None + + class CmorSession( AbstractContextManager ): # pylint: disable=too-many-instance-attributes @@ -452,10 +527,16 @@ def _define_mom6_grid(self, ds, var_name, var_da, var_dims): def _define_cice_grid(self, ds, var_name, var_da): """Register a native CICE (nj, ni) tripole grid via cmor.grid(). - The CICE grid is logically-rectangular: it carries TLAT/TLON cell centers - and vertex bounds (latt_bounds/lont_bounds, shape (nj, ni, nvertices)) - inline. Defines i_index/j_index axes from the CMIP7 grids table and returns - the resulting grid id, which stands in for both the nj and ni dimensions. + The CICE grid is logically-rectangular and staggered: each point carries + its own latitude, longitude and vertex bounds inline -- TLAT/TLON with + latt_bounds/lont_bounds at the cell centre, ULAT/ULON with + latu_bounds/lonu_bounds at the velocity point, all of shape + (nj, ni, nvertices) for the bounds. The variable's own 'coordinates' + attribute says which point it is on, and that pair is what the grid is + defined from. + + Defines i_index/j_index axes from the CMIP7 grids table and returns the + resulting grid id, which stands in for both the nj and ni dimensions. """ logger.debug( "[CMOR axis debug] Defining native CICE (nj, ni) grid for %s.", @@ -472,30 +553,35 @@ def _coord(dsi, *names): return dsi[nm] return None - tlat = _coord(ds, "TLAT", "tlat", "lat", "latitude") - tlon = _coord(ds, "TLON", "tlon", "lon", "longitude") - if tlat is None or tlon is None: + centre_lat = _coord(ds, "TLAT", "tlat", "lat", "latitude") + centre_lon = _coord(ds, "TLON", "tlon", "lon", "longitude") + + # Which point of the staggered grid this variable is on comes from its + # own 'coordinates' attribute: the velocities are on the B-grid velocity + # point (ULAT/ULON), the thermodynamic fields on the centre (TLAT/TLON). + # Taking the centre for everything put siu and siv in the wrong place. + grid_lat, grid_lon = horizontal_coords_of(ds, var_da, centre_lat, centre_lon) + if grid_lat is None or grid_lon is None: raise KeyError( - "CICE native grid requires TLAT/TLON cell-center coordinates " - f"in the dataset for variable '{var_name}'." + "CICE native grid requires latitude/longitude coordinates in " + f"the dataset for variable '{var_name}'." ) + logger.debug( + "[CMOR axis debug] %s is on %s/%s", + var_name, + getattr(grid_lat, "name", "?"), + getattr(grid_lon, "name", "?"), + ) - lat_bnds_da = None - lon_bnds_da = None - bname = tlat.attrs.get("bounds") - if isinstance(bname, str) and bname in ds: - lat_bnds_da = ds[bname] - bname = tlon.attrs.get("bounds") - if isinstance(bname, str) and bname in ds: - lon_bnds_da = ds[bname] - if lat_bnds_da is None: - lat_bnds_da = _coord(ds, "latt_bounds", "lat_bounds", "TLAT_bnds") - if lon_bnds_da is None: - lon_bnds_da = _coord(ds, "lont_bounds", "lon_bounds", "TLON_bnds") + lat_bnds_da = vertex_bounds_of(ds, grid_lat, str(getattr(grid_lat, "name", ""))) + lon_bnds_da = vertex_bounds_of(ds, grid_lon, str(getattr(grid_lon, "name", ""))) if lat_bnds_da is None or lon_bnds_da is None: raise KeyError( "CICE native grid requires latitude/longitude vertex bounds " - f"(e.g. latt_bounds/lont_bounds) for variable '{var_name}'." + f"for {getattr(grid_lat, 'name', '?')}/{getattr(grid_lon, 'name', '?')} " + f"(e.g. latt_bounds/lont_bounds) for variable '{var_name}'. " + "Bounds from another grid point are not substituted: they would " + "pair cell corners with the wrong cell centres." ) # The horizontal grid is static. During the multi-file merge, xarray @@ -506,8 +592,8 @@ def _coord(dsi, *names): # rank does not match number of axes passed via axis_ids". Collapse any # non-horizontal dimension (e.g. time) by taking the first index, since # the grid geometry does not vary in time. - tlat = _horizontal_only(tlat, ("nj", "ni")) - tlon = _horizontal_only(tlon, ("nj", "ni")) + grid_lat = _horizontal_only(grid_lat, ("nj", "ni")) + grid_lon = _horizontal_only(grid_lon, ("nj", "ni")) # vertex bounds keep their trailing vertices dim in addition to nj, ni lat_bnds_da = _horizontal_only( lat_bnds_da, ("nj", "ni") + lat_bnds_da.dims[-1:] @@ -516,8 +602,8 @@ def _coord(dsi, *names): lon_bnds_da, ("nj", "ni") + lon_bnds_da.dims[-1:] ) - lat_vals = np.asarray(tlat.values, dtype="f8") - lon_vals = np.mod(np.asarray(tlon.values, dtype="f8"), 360.0) + lat_vals = np.asarray(grid_lat.values, dtype="f8") + lon_vals = np.mod(np.asarray(grid_lon.values, dtype="f8"), 360.0) lat_vert = np.asarray(lat_bnds_da.values, dtype="f8") lon_vert = np.mod(np.asarray(lon_bnds_da.values, dtype="f8"), 360.0) diff --git a/src/cmip7_prep/grids.py b/src/cmip7_prep/grids.py index c07476c..963c9fd 100644 --- a/src/cmip7_prep/grids.py +++ b/src/cmip7_prep/grids.py @@ -1,84 +1,76 @@ """The grid each realm's data is on. Model components run on different grids, so the grid to regrid from depends on -which realm is being processed: +which realm is being processed. Which grid that is is stated per realm, per +resolution, in ``data/_grids.yaml`` under ``grid_names_per_realm`` +as each row's ``input_grid``, rather than being hardcoded here, so that adding +a model or changing a component's grid is a data change. - atmos, atmosChem, aerosol, land the atmosphere/land grid the case was - run at: ne30, ne16, or 'regular' when - the output already carries lat/lon - ocean, ocnBgchem, seaIce the model's tripolar grid: tnx1v4 for - NorESM, tx2_3v2 for CESM - landIce its own projected grid; CISM output is - georeferenced from the x/y coordinates - in the files, not regridded - -:func:`resolution_for` returns it, so a caller gives only the atmosphere/land -resolution -- and only when processing a realm that uses it. +:func:`resolution_for` reads it. ``'custom'`` is the way to run a case the +table does not otherwise describe: give it a ``custom`` block naming whatever +grids it ran, adding entries to ``resolutions`` for any that are new. Those +grids may be regular lat/lon, in which case nothing is regridded, but they need +not be -- that is one possibility, not what ``custom`` means. With no block at +all, nothing is regridded, since that is all that can be assumed. """ from __future__ import annotations -# Grids the atmosphere and land may be run on. 'regular' means the output -# already carries lat/lon and is not regridded. -ATM_RESOLUTIONS = ("ne30", "ne16", "regular") - -# Realms whose input grid is the atmosphere/land grid the case was run on. -ATM_GRID_REALMS = frozenset({"atmos", "atmosChem", "aerosol", "land"}) - -# The ocean/sea-ice grid each model writes on. These realms never share the -# atmosphere's grid, so passing one an atmosphere resolution would regrid -# through the wrong weights. -OCEAN_GRID = {"noresm": "tnx1v4", "cesm": "tx2_3v2"} +from .regrid_maps import get_realm_row, load_regrid_maps -OCEAN_GRID_REALMS = frozenset({"ocean", "ocnBgchem", "seaIce"}) +# Resolutions the case may have been run at, as accepted on the command line. +MODEL_RESOLUTIONS = ("ne30", "NorESM3-LM", "NorESM3-MM", "custom") -# Realms written on their native grid, with no ESMF regridding. CISM land-ice -# output carries projected x/y coordinates, which cmor_writer georeferences with -# the ice sheet's own map projection (see cism_grid.project_xy_to_latlon), so no -# weight files are involved. The value below is a pass-through: it names the -# entry that builds a regridder and discards it, which is what the unregridded -# path expects. -NATIVE_GRID_REALMS = frozenset({"landIce"}) -NATIVE_GRID = "regular" +# The extension point: a case the table does not otherwise describe, whose +# grids are declared by filling in the commented-out block at the end of +# 'grid_names_per_realm'. Any grid may be named there, new ones included. +# With no block, nothing is regridded. +CUSTOM = "custom" -def needs_atmos_res(realm: str) -> bool: - """Return whether a realm's grid depends on the atmosphere resolution. +def resolution_for(model: str, realm: str, model_res: str | None = None) -> str: + """Return the input grid name for one realm. - Only the realms sharing the atmosphere and land grid do. Sea ice, ocean - and land ice have their own, so asking for an atmosphere resolution to - process them would be asking for something that cannot matter. + ``model_res`` is the resolution the case was run at, which the table is + keyed by along with the realm. """ - return realm in ATM_GRID_REALMS + if model_res is None: + raise ValueError( + f"The resolution the case was run at must be given to find " + f"realm {realm!r}'s input grid" + ) + if model_res not in MODEL_RESOLUTIONS: + raise ValueError( + f"Unknown atmosphere resolution {model_res!r}; " + f"choose from {list(MODEL_RESOLUTIONS)}" + ) + declared = load_regrid_maps(model).get("grid_names_per_realm") or {} + if model_res == CUSTOM and model_res not in declared: + # A custom case whose grids have not been declared: nothing is + # regridded, which is all that can be assumed about input the table + # says nothing about. + return CUSTOM + return get_realm_row(model, model_res, realm)["input_grid"] -def resolution_for(model: str, realm: str, atmos_res: str | None = None) -> str: - """Return the input grid name for one realm. +# CICE staggers its grid: the thermodynamic fields sit at the cell centre (T), +# the velocities at the B-grid velocity point (U), and CICE6 adds the N and E +# points of the C grid. Each point has its own latitude, longitude and vertex +# bounds, and a variable says which it is on through its 'coordinates' +# attribute. +CICE_BOUNDS_BY_COORD = { + "TLAT": "latt_bounds", + "TLON": "lont_bounds", + "ULAT": "latu_bounds", + "ULON": "lonu_bounds", + "NLAT": "latn_bounds", + "NLON": "lonn_bounds", + "ELAT": "late_bounds", + "ELON": "lone_bounds", +} - ``atmos_res`` is the grid the atmosphere and land were run on. It is - required only for the realms on that grid, and ignored for the rest, so a - sea-ice run need not supply one. - """ - if realm in ATM_GRID_REALMS: - if atmos_res is None: - raise ValueError( - f"Realm {realm!r} is on the atmosphere/land grid, so its " - "resolution must be given" - ) - if atmos_res not in ATM_RESOLUTIONS: - raise ValueError( - f"Unknown atmosphere resolution {atmos_res!r}; " - f"choose from {list(ATM_RESOLUTIONS)}" - ) - return atmos_res - if realm in OCEAN_GRID_REALMS: - try: - return OCEAN_GRID[model] - except KeyError: - raise ValueError( - f"No ocean grid known for model={model!r}; " - f"known models: {sorted(OCEAN_GRID)}" - ) from None - if realm in NATIVE_GRID_REALMS: - return NATIVE_GRID - raise ValueError(f"No input grid known for realm {realm!r}") +# Every coordinate and bounds variable of that grid. Realizing a variable +# builds a new dataset around it, which leaves the bounds behind, so these are +# copied over -- all of them, so a velocity variable keeps ULAT/ULON and not +# just the centre (issue #115). +CICE_GRID_VARS = tuple(CICE_BOUNDS_BY_COORD) + tuple(CICE_BOUNDS_BY_COORD.values()) diff --git a/src/cmip7_prep/reference_run.py b/src/cmip7_prep/reference_run.py index 781fd26..512470e 100644 --- a/src/cmip7_prep/reference_run.py +++ b/src/cmip7_prep/reference_run.py @@ -23,7 +23,7 @@ from pathlib import Path from typing import Sequence -from .grids import ATM_RESOLUTIONS, needs_atmos_res +from .grids import MODEL_RESOLUTIONS from .include_patterns import all_include_patterns, load_include_patterns # Which component directory holds each realm's history files. Mirrors @@ -153,7 +153,7 @@ def build_plan( frequencies: Sequence[str] | None = None, years: str | None = None, stages: Sequence[str] = STAGES, - atmos_res: str | None = None, + model_res: str | None = None, experiment: str, workers: int = 4, ice_sheet: str | None = None, @@ -171,8 +171,9 @@ def build_plan( history directory is absent is recorded in ``Plan.skipped`` rather than failing the run, since an archived case need not hold every component. - ``atmos_res`` is the grid the atmosphere and land were run on; every - other realm's grid is derived from it or from the model. + ``model_res`` is the resolution the case was run at, named by its + atmosphere grid. Every realm needs it, because it identifies the case; + the grid a realm regrids from is derived from it. ``workers`` applies to time series generation only; CMORization always runs with one worker (see CMOR_WORKERS). @@ -190,16 +191,18 @@ def build_plan( if unknown: raise ValueError(f"Unknown stage(s) {unknown}; choose from {list(STAGES)}") wanted_realms = list(realms) if realms else realms_for(model) - on_atmos_grid = [realm for realm in wanted_realms if needs_atmos_res(realm)] - if on_atmos_grid and atmos_res is None: + # Needed whatever the realms, because it identifies the case rather than + # only the grid the atmosphere realms regrid from: source_id, + # nominal_resolution and every realm's grid label are keyed by it. + if model_res is None: raise ValueError( - "An atmosphere resolution is needed for " - f"{on_atmos_grid}; the other realms derive their own grid" + "The resolution the case was run at must be given; it identifies " + "the case, not only the grid the atmosphere realms regrid from" ) - if atmos_res is not None and atmos_res not in ATM_RESOLUTIONS: + if model_res not in MODEL_RESOLUTIONS: raise ValueError( - f"Unknown atmosphere resolution {atmos_res!r}; " - f"choose from {list(ATM_RESOLUTIONS)}" + f"Unknown atmosphere resolution {model_res!r}; " + f"choose from {list(MODEL_RESOLUTIONS)}" ) case_dir = Path(case_dir) @@ -272,7 +275,7 @@ def build_plan( ts_dir=ts_dir, cmor_root=cmor_root, model=model, - atmos_res=(atmos_res if needs_atmos_res(realm) else None), + model_res=model_res, experiment=experiment, sheet=sheet, variant_label=variant_label, @@ -357,7 +360,7 @@ def _cmor_step( ts_dir, cmor_root, model, - atmos_res, + model_res, experiment, sheet, variant_label=None, @@ -382,8 +385,8 @@ def _cmor_step( "--workers", str(CMOR_WORKERS), ] - if atmos_res: - command += ["--atmos-res", atmos_res] + if model_res: + command += ["--model-res", model_res] if sheet: command += ["--ice-sheet", sheet] if variant_label: diff --git a/src/cmip7_prep/regrid.py b/src/cmip7_prep/regrid.py index 71259cb..9175761 100644 --- a/src/cmip7_prep/regrid.py +++ b/src/cmip7_prep/regrid.py @@ -380,14 +380,14 @@ def regrid_to_latlon( # In future: handle if the lat, lon coords are already present, but still on the wrong grid # or other changes to the grid should be made. - if resolution == "regular": + if resolution == "custom": if "lat" in var_da.dims and "lon" in var_da.dims: logger.info( "Variable already has 'lat' and 'lon' dims; skipping regridding." ) return var_da raise KeyError( - "regular grid is requested but variable does not have 'lat','lon' dims" + "custom grid is requested but variable does not have 'lat','lon' dims" ) if "ncol" not in var_da.dims and "lndgrid" not in var_da.dims: diff --git a/src/cmip7_prep/regrid_maps.py b/src/cmip7_prep/regrid_maps.py index b99f6e3..67bb7f3 100644 --- a/src/cmip7_prep/regrid_maps.py +++ b/src/cmip7_prep/regrid_maps.py @@ -1,6 +1,6 @@ """ESMF regrid weight files, loaded from the packaged YAML tables. -Each supported model has a ``data/_regrid_maps.yaml`` table mapping a +Each supported model has a ``data/_grids.yaml`` table mapping a resolution to its weight files:: inputdata_dir: /nird/datalake/NS9560K/diagnostics/land_xesmf_diag_data/ @@ -39,10 +39,10 @@ @lru_cache(maxsize=None) def _load_cached(model: str) -> dict: - """Read and cache one model's regrid-map table.""" - path = DATA_DIR / f"{model}_regrid_maps.yaml" + """Read and cache one model's grid table.""" + path = DATA_DIR / f"{model}_grids.yaml" if not path.is_file(): - raise ValueError(f"No regrid-map table for model={model!r}; expected {path}") + raise ValueError(f"No grid table for model={model!r}; expected {path}") with open(path, encoding="utf-8") as handle: return yaml.safe_load(handle) or {} @@ -79,10 +79,12 @@ def get_map_paths(model: str, resolution: str) -> dict[str, Path]: return {method: root / name for method, name in entry.items()} -def get_grid_names(model: str, resolution: str, realm: str) -> dict[str, str]: - """Return the grid names for one model and resolution. +def get_realm_row(model: str, resolution: str, realm: str) -> dict[str, str]: + """Return one realm's row: the grid its input is on and its output label. - The result has keys like 'atmos', 'aerosol', 'land', 'seaIce', 'landIce'. + 'input_grid' names an entry in 'resolutions'; 'grid_label' is the CMIP7 + code for the grid the output lands on. grids.py reads the first and + cmor_utils.py the second. """ table = _load_cached(model) grid_names = table.get("grid_names_per_realm") or {} @@ -98,7 +100,7 @@ def get_grid_names(model: str, resolution: str, realm: str) -> dict[str, str]: f"available: {sorted(realm_grids)}" ) - return grid_names[resolution][realm] or {} + return dict(realm_grids[realm] or {}) def get_grid_desc(model: str, resolution: str, realm: str) -> str: diff --git a/tests/test_cice_staggered_coords.py b/tests/test_cice_staggered_coords.py new file mode 100644 index 0000000..ce03260 --- /dev/null +++ b/tests/test_cice_staggered_coords.py @@ -0,0 +1,195 @@ +"""Tests for putting CICE variables on the grid point they belong to. + +CICE staggers its grid: the thermodynamic fields sit at the cell centre +(TLAT/TLON) and the velocities at the B-grid velocity point (ULAT/ULON). Each +variable names its own coordinates, so they are read from the variable rather +than assumed, which is what issue #115 reports going wrong -- siu and siv were +written with centre coordinates. +""" + +import numpy as np +import xarray as xr + +from cmip7_prep.grids import CICE_BOUNDS_BY_COORD, CICE_GRID_VARS +from cmip7_prep.cmor_writer import ( + horizontal_coords_of, + is_latitude, + is_longitude, + vertex_bounds_of, +) + +NJ, NI, NVERT = 2, 3, 4 + + +def _cice_dataset(with_u_bounds=True): + """Return a dataset shaped like CICE output, with T and U grid points.""" + shape = (NJ, NI) + ds = xr.Dataset( + { + "siconc": ( + ("nj", "ni"), + np.zeros(shape), + {"coordinates": "TLON TLAT time"}, + ), + "siu": ( + ("nj", "ni"), + np.zeros(shape), + {"coordinates": "ULON ULAT time"}, + ), + "no_coords": (("nj", "ni"), np.zeros(shape), {}), + }, + coords={ + "TLAT": ( + ("nj", "ni"), + np.full(shape, 10.0), + {"units": "degrees_north", "bounds": "latt_bounds"}, + ), + "TLON": ( + ("nj", "ni"), + np.full(shape, 20.0), + {"units": "degrees_east", "bounds": "lont_bounds"}, + ), + # The velocity point deliberately carries no 'bounds' attribute, so + # the name map is what has to find them. + "ULAT": (("nj", "ni"), np.full(shape, 11.0), {"units": "degrees_north"}), + "ULON": (("nj", "ni"), np.full(shape, 21.0), {"units": "degrees_east"}), + }, + ) + ds["latt_bounds"] = (("nj", "ni", "nvertices"), np.zeros((NJ, NI, NVERT))) + ds["lont_bounds"] = (("nj", "ni", "nvertices"), np.zeros((NJ, NI, NVERT))) + if with_u_bounds: + ds["latu_bounds"] = (("nj", "ni", "nvertices"), np.ones((NJ, NI, NVERT))) + ds["lonu_bounds"] = (("nj", "ni", "nvertices"), np.ones((NJ, NI, NVERT))) + return ds + + +class TestIdentifyingCoordinates: + """Latitude and longitude are recognised by their CF attributes.""" + + def test_latitude_by_units(self): + """degrees_north marks a latitude whatever the variable is called.""" + assert is_latitude(xr.DataArray([0.0], attrs={"units": "degrees_north"})) + + def test_latitude_by_standard_name(self): + """A standard_name of latitude is enough on its own.""" + assert is_latitude(xr.DataArray([0.0], attrs={"standard_name": "latitude"})) + + def test_longitude_by_units(self): + """degrees_east marks a longitude.""" + assert is_longitude(xr.DataArray([0.0], attrs={"units": "degrees_east"})) + + def test_a_latitude_is_not_a_longitude(self): + """The two are told apart rather than both matching.""" + lat = xr.DataArray([0.0], attrs={"units": "degrees_north"}) + assert not is_longitude(lat) + + +class TestChoosingTheGridPoint: + """Each variable gets the coordinates its own attribute names.""" + + def test_velocity_gets_the_velocity_point(self): + """siu is on ULAT/ULON, which is what issue #115 reports.""" + ds = _cice_dataset() + lat, lon = horizontal_coords_of(ds, ds["siu"], ds["TLAT"], ds["TLON"]) + assert lat.name == "ULAT" + assert lon.name == "ULON" + + def test_thermodynamic_field_gets_the_centre(self): + """siconc stays on TLAT/TLON.""" + ds = _cice_dataset() + lat, lon = horizontal_coords_of(ds, ds["siconc"], ds["TLAT"], ds["TLON"]) + assert lat.name == "TLAT" + assert lon.name == "TLON" + + def test_the_two_points_are_different_values(self): + """The choice changes the numbers, not only the name. + + If it did not, the bug in the issue would have had no effect. + """ + ds = _cice_dataset() + u_lat, _ = horizontal_coords_of(ds, ds["siu"], ds["TLAT"], ds["TLON"]) + t_lat, _ = horizontal_coords_of(ds, ds["siconc"], ds["TLAT"], ds["TLON"]) + assert float(u_lat[0, 0]) != float(t_lat[0, 0]) + + def test_the_attribute_decides_when_both_pairs_are_attached(self): + """Attached coordinates cannot distinguish the grid points. + + TLAT and ULAT have the same dimensions, so xarray gives every variable + on (nj, ni) all four of them. Only the variable's own 'coordinates' + attribute says which pair it is on, which is why that is read first. + """ + ds = _cice_dataset() + attached = {str(name) for name in ds["siu"].coords} + assert {"TLAT", "ULAT"} <= attached, "both pairs should be attached" + lat, lon = horizontal_coords_of(ds, ds["siu"], ds["TLAT"], ds["TLON"]) + assert (lat.name, lon.name) == ("ULAT", "ULON") + + def test_the_attribute_is_read_from_encoding_too(self): + """Decoding coordinates moves the attribute out of attrs. + + xarray puts it in encoding instead, so a dataset read from a file keeps + the information there rather than in attrs. + """ + ds = _cice_dataset() + del ds["siu"].attrs["coordinates"] + ds["siu"].encoding["coordinates"] = "ULON ULAT time" + lat, lon = horizontal_coords_of(ds, ds["siu"], ds["TLAT"], ds["TLON"]) + assert (lat.name, lon.name) == ("ULAT", "ULON") + + def test_a_variable_without_the_attribute_falls_back(self): + """Output that names no coordinates uses the centre, as before.""" + ds = _cice_dataset() + lat, lon = horizontal_coords_of(ds, ds["no_coords"], ds["TLAT"], ds["TLON"]) + assert lat.name == "TLAT" + assert lon.name == "TLON" + + def test_an_attribute_naming_absent_variables_falls_back(self): + """Coordinates named but not present do not leave us empty-handed.""" + ds = _cice_dataset() + ds["siu"].attrs["coordinates"] = "NOWHERE_LON NOWHERE_LAT" + lat, _ = horizontal_coords_of(ds, ds["siu"], ds["TLAT"], ds["TLON"]) + assert lat.name == "TLAT" + + +class TestFindingVertexBounds: + """Bounds come from the same grid point as the coordinate.""" + + def test_the_bounds_attribute_is_followed(self): + """TLAT says where its bounds are, so that is used.""" + ds = _cice_dataset() + bounds = vertex_bounds_of(ds, ds["TLAT"], "TLAT") + assert bounds.name == "latt_bounds" + + def test_the_name_map_covers_a_missing_attribute(self): + """ULAT carries no bounds attribute, so the map finds latu_bounds.""" + ds = _cice_dataset() + assert vertex_bounds_of(ds, ds["ULAT"], "ULAT").name == "latu_bounds" + assert vertex_bounds_of(ds, ds["ULON"], "ULON").name == "lonu_bounds" + + def test_the_velocity_bounds_are_not_the_centre_bounds(self): + """Substituting centre bounds would misplace every cell corner.""" + ds = _cice_dataset() + u_bounds = vertex_bounds_of(ds, ds["ULAT"], "ULAT") + t_bounds = vertex_bounds_of(ds, ds["TLAT"], "TLAT") + assert u_bounds.name != t_bounds.name + assert float(u_bounds[0, 0, 0]) != float(t_bounds[0, 0, 0]) + + def test_absent_bounds_return_nothing(self): + """With no bounds for this point, None is returned rather than another + point's bounds, leaving the caller to report it.""" + ds = _cice_dataset(with_u_bounds=False) + assert vertex_bounds_of(ds, ds["ULAT"], "ULAT") is None + + def test_every_cice_grid_point_is_mapped(self): + """The four grid points CICE writes all have an entry.""" + for point in ("T", "U", "N", "E"): + assert f"{point}LAT" in CICE_BOUNDS_BY_COORD + assert f"{point}LON" in CICE_BOUNDS_BY_COORD + + def test_the_carried_list_covers_coordinates_and_bounds(self): + """What the driver copies over is every name in the map, both halves.""" + assert set(CICE_GRID_VARS) == set(CICE_BOUNDS_BY_COORD) | set( + CICE_BOUNDS_BY_COORD.values() + ) + assert "ULAT" in CICE_GRID_VARS + assert "latu_bounds" in CICE_GRID_VARS diff --git a/tests/test_grids.py b/tests/test_grids.py index c1134c3..0260c6a 100644 --- a/tests/test_grids.py +++ b/tests/test_grids.py @@ -1,61 +1,156 @@ -"""Tests for deriving each realm's input grid.""" +"""Tests for reading each realm's input grid from the model tables. + +Nothing here names a grid or a resolution. Both are data, stated per model in +``data/_grids.yaml``, and restating them in a test would only +assert that this file and that one were edited together. So the tests check +the properties the table must have, and that the lookup returns what the table +says, whatever the table says. +""" import pytest -from cmip7_prep.grids import ( - ATM_RESOLUTIONS, - needs_atmos_res, - NATIVE_GRID, - OCEAN_GRID, - resolution_for, -) - - -class TestResolutionPerRealm: - """Tests for choosing each realm's input grid. - - The grid is not a free choice: sea ice is on the model's tripolar grid - whatever the atmosphere was run on, so passing one realm another's grid - would regrid through the wrong weights and yield plausible wrong output. - """ - - @pytest.mark.parametrize("realm", ["atmos", "atmosChem", "aerosol", "land"]) - @pytest.mark.parametrize("atm", ATM_RESOLUTIONS) - def test_atmosphere_realms_take_the_given_grid(self, realm, atm): - """Atmosphere and land use the resolution the case was run at.""" - assert resolution_for("noresm", realm, atm) == atm - - @pytest.mark.parametrize("realm", ["seaIce", "ocean", "ocnBgchem"]) - @pytest.mark.parametrize("model", ["noresm", "cesm"]) - def test_ocean_realms_take_the_model_grid(self, realm, model): - """Ocean and sea ice ignore the atmosphere resolution entirely.""" - assert resolution_for(model, realm, "ne16") == OCEAN_GRID[model] - - @pytest.mark.parametrize("realm", ["seaIce", "ocean", "ocnBgchem", "landIce"]) - def test_realms_with_their_own_grid_need_no_resolution(self, realm): - """A sea-ice or land-ice run need not supply an atmosphere resolution.""" - assert resolution_for("noresm", realm) - assert not needs_atmos_res(realm) - - @pytest.mark.parametrize("realm", ["atmos", "atmosChem", "aerosol", "land"]) - def test_atmosphere_realms_say_so_when_it_is_missing(self, realm): - """A realm on the atmosphere grid refuses to guess its resolution.""" - assert needs_atmos_res(realm) - with pytest.raises(ValueError, match="resolution must be given"): - resolution_for("noresm", realm) - - def test_unknown_atmosphere_resolution_is_rejected(self): - """A resolution with no weight files fails before any run starts.""" +from cmip7_prep.grids import MODEL_RESOLUTIONS, CUSTOM, resolution_for +from cmip7_prep.regrid_maps import load_regrid_maps + +MODELS = ("noresm", "cesm") + + +def _rows(model): + """Yield (resolution, realm, row) for every row in a model's table.""" + for resolution, realms in load_regrid_maps(model)["grid_names_per_realm"].items(): + for realm, row in realms.items(): + yield resolution, realm, row + + +class TestTheLookupReturnsWhatTheTableSays: + """resolution_for is a reader of the table, and decides nothing itself.""" + + @pytest.mark.parametrize("model", MODELS) + def test_every_row_is_returned_verbatim(self, model): + """Each realm gets its own row's input_grid, for each resolution.""" + for resolution, realm, row in _rows(model): + assert ( + resolution_for(model, realm, resolution) == row["input_grid"] + ), f"{model}/{resolution}/{realm}" + + @pytest.mark.parametrize("model", MODELS) + def test_the_resolution_selects_the_row(self, model): + """Two resolutions that differ for a realm give different answers. + + This is what the original bug got wrong: a realm was given another + realm's grid, which regrids through the wrong weights and produces + output that looks fine. + """ + table = load_regrid_maps(model)["grid_names_per_realm"] + resolutions = sorted(table) + differing = [ + (realm, [table[res][realm]["input_grid"] for res in resolutions]) + for realm in table[resolutions[0]] + if len({table[res][realm]["input_grid"] for res in resolutions}) > 1 + ] + if not differing: + pytest.skip(f"{model} has only one resolution in its table") + for realm, expected in differing: + got = [resolution_for(model, realm, res) for res in resolutions] + assert got == expected, realm + + +class TestEveryRowIsUsable: + """Each row must name things the rest of the table can act on.""" + + @pytest.mark.parametrize("model", MODELS) + def test_every_input_grid_has_weights(self, model): + """A grid with no 'resolutions' entry would fail at run time.""" + defined = set(load_regrid_maps(model)["resolutions"]) + for resolution, realm, row in _rows(model): + assert row["input_grid"] in defined, f"{model}/{resolution}/{realm}" + + @pytest.mark.parametrize("model", MODELS) + def test_every_row_has_both_fields_and_nothing_else(self, model): + """Neither field may be left out, and no third field is read.""" + for resolution, realm, row in _rows(model): + assert set(row) == { + "input_grid", + "grid_label", + }, f"{model}/{resolution}/{realm}" + + @pytest.mark.parametrize("model", MODELS) + def test_every_resolution_defines_the_same_realms(self, model): + """A realm present at one resolution and missing at another would fail + only for the run that asked for it.""" + table = load_regrid_maps(model)["grid_names_per_realm"] + realms = [set(rows) for rows in table.values()] + assert all(group == realms[0] for group in realms) + + +class TestWhatIsRefused: + """A missing or unknown value fails, rather than being guessed.""" + + def test_a_missing_resolution_is_rejected(self): + """The resolution is never guessed, for any realm.""" + with pytest.raises(ValueError, match="must be given"): + resolution_for("noresm", next(iter(_rows("noresm")))[1]) + + def test_an_unknown_resolution_is_rejected(self): + """A resolution the command line does not offer fails up front.""" + unknown = "".join(MODEL_RESOLUTIONS) + "x" with pytest.raises(ValueError, match="Unknown atmosphere resolution"): - resolution_for("noresm", "atmos", "ne120") + resolution_for("noresm", "atmos", unknown) + + def test_a_resolution_the_model_lacks_is_rejected(self): + """An offered resolution with no rows for this model fails. + + Not every model runs every resolution, and the message names the ones + it has rather than falling back to another. + """ + for model in MODELS: + has_rows = set(load_regrid_maps(model)["grid_names_per_realm"]) + missing = [ + res + for res in MODEL_RESOLUTIONS + if res not in has_rows and res != CUSTOM + ] + for res in missing: + with pytest.raises(ValueError, match="No grid names defined"): + resolution_for(model, "atmos", res) - @pytest.mark.parametrize("model", ["noresm", "cesm"]) - @pytest.mark.parametrize("atm", ATM_RESOLUTIONS) - def test_landice_is_never_regridded(self, model, atm): - """land ice takes the pass-through grid whatever else was asked for. + def test_an_unknown_realm_is_rejected(self): + """A misspelled realm is answered with the table's own keys.""" + resolution = next(iter(_rows("noresm")))[0] + with pytest.raises(ValueError, match="No grid names defined"): + resolution_for("noresm", "atmosphere", resolution) - CISM output is written on its native projected grid, georeferenced from - its x/y coordinates, so no weight files apply and the atmosphere's - resolution is irrelevant to it. + +class TestACustomCase: + """A case that is not one of the standard configurations.""" + + @pytest.mark.parametrize("model", MODELS) + def test_undeclared_grids_mean_nothing_is_regridded(self, model): + """With no block of its own, every realm is left alone. + + Nothing can be assumed about the grids of a case the table says + nothing about, so the one safe answer is not to regrid. """ - assert resolution_for(model, "landIce", atm) == NATIVE_GRID + declared = load_regrid_maps(model)["grid_names_per_realm"] + if CUSTOM in declared: + pytest.skip(f"{model} declares a {CUSTOM} block") + for _, realm, _ in _rows(model): + assert resolution_for(model, realm, CUSTOM) == CUSTOM + + @pytest.mark.parametrize("model", MODELS) + def test_declared_grids_are_used(self, model): + """Filling in the block makes the case behave like any other. + + This is what the commented-out template at the end of the table is + for, so a custom case on, say, regular lat/lon is not forced through + the no-regrid path. + """ + declared = load_regrid_maps(model)["grid_names_per_realm"] + if CUSTOM not in declared: + pytest.skip(f"{model} declares no {CUSTOM} block") + for realm, row in declared[CUSTOM].items(): + assert resolution_for(model, realm, CUSTOM) == row["input_grid"] + + def test_it_is_offered_on_the_command_line(self): + """It is a --model-res choice, so the lookup has to accept it.""" + assert CUSTOM in MODEL_RESOLUTIONS diff --git a/tests/test_reference_run.py b/tests/test_reference_run.py index 0a0365b..a7cefae 100644 --- a/tests/test_reference_run.py +++ b/tests/test_reference_run.py @@ -62,7 +62,7 @@ def _plan(case_dir, outdir, **kwargs): wrong output, so a caller must say which it means. """ kwargs.setdefault("model", "noresm") - kwargs.setdefault("atmos_res", "ne16") + kwargs.setdefault("model_res", "NorESM3-LM") kwargs.setdefault("experiment", "piControl") return build_plan(case_dir, outdir, **kwargs) @@ -313,11 +313,11 @@ def test_resolution_and_experiment_reach_cmor(self, case_dir, tmp_path): case_dir, tmp_path / "out", realms=["atmos"], - atmos_res="ne30", + model_res="NorESM3-MM", experiment="historical", ) command = _command_of(plan, "cmor-atmos-mon") - assert command[command.index("--atmos-res") + 1] == "ne30" + assert command[command.index("--model-res") + 1] == "NorESM3-MM" assert command[command.index("--experiment") + 1] == "historical" def test_variant_label_reaches_both_later_stages(self, case_dir, tmp_path): @@ -434,13 +434,13 @@ def test_case_properties_are_required(omitted): """build_plan refuses to guess a property of the case. Model and experiment change what the output means while leaving it looking - valid, so neither may be defaulted. The atmosphere resolution is not here - because it is only needed by some realms; see - TestAtmosResolutionIsConditional. + valid, so neither may be defaulted. The resolution is not here only + because it is accepted as a keyword and refused later; see + TestAtmosResolutionIsAlwaysRequired. """ supplied = { "model": "noresm", - "atmos_res": "ne16", + "model_res": "NorESM3-LM", "experiment": "piControl", } del supplied[omitted] @@ -448,48 +448,49 @@ def test_case_properties_are_required(omitted): build_plan("case", "out", **supplied) -class TestAtmosResolutionIsConditional: - """The atmosphere resolution is asked for only when a realm uses it.""" +class TestAtmosResolutionIsAlwaysRequired: + """Every realm needs the resolution the case was run at. - def test_not_needed_for_realms_with_their_own_grid(self, case_dir, tmp_path): - """Sea ice and land ice plan fully without one.""" - plan = build_plan( - case_dir, - tmp_path / "out", - model="noresm", - experiment="piControl", - realms=["seaIce", "landIce"], - ) - assert plan.steps - for step in plan.for_stage("cmor"): - assert "--atmos-res" not in step.command - - @pytest.mark.parametrize("realm", ["atmos", "land"]) - def test_needed_for_realms_on_the_atmosphere_grid(self, case_dir, tmp_path, realm): - """A realm on that grid refuses to plan without one, and names itself.""" - with pytest.raises(ValueError, match=f"needed for \\['{realm}'\\]"): + It names the atmosphere grid, but it identifies the case rather than only + the grid a realm regrids from: source_id, nominal_resolution and every + realm's grid label are keyed by it. Leaving it out of a sea-ice run + labelled the output with fallbacks instead of the case's own identity. + """ + + def test_planning_without_one_refuses(self, case_dir, tmp_path): + """No realm plans without it, whatever grid its input is on.""" + with pytest.raises(ValueError, match="resolution the case was run at"): build_plan( case_dir, tmp_path / "out", model="noresm", experiment="piControl", - realms=[realm], + realms=["seaIce", "landIce"], ) - def test_the_whole_matrix_names_every_realm_that_needs_it(self, case_dir, tmp_path): - """Planning every realm without one lists the four that require it.""" - with pytest.raises(ValueError, match="atmos.*atmosChem.*aerosol.*land"): + def test_an_unknown_resolution_is_rejected(self, case_dir, tmp_path): + """A resolution outside the known set is named back. + + cmor_driver.py would reject it too, but only after the timeseries + stage has already run. + """ + with pytest.raises(ValueError, match="ne99"): build_plan( - case_dir, tmp_path / "out", model="noresm", experiment="piControl" + case_dir, + tmp_path / "out", + model="noresm", + experiment="piControl", + model_res="ne99", ) - def test_only_atmosphere_realms_carry_the_flag(self, case_dir, tmp_path): - """With one supplied, only the realms that use it are given it.""" - plan = _plan(case_dir, tmp_path / "out", realms=["atmos", "seaIce"]) - atm = _command_of(plan, "cmor-atmos-mon") - sea = _command_of(plan, "cmor-seaIce-mon") - assert atm[atm.index("--atmos-res") + 1] == "ne16" - assert "--atmos-res" not in sea + @pytest.mark.parametrize("realm", realms_for("noresm")) + def test_every_realm_is_given_it(self, case_dir, tmp_path, realm): + """Each realm's CMOR steps carry it, not just the atmosphere's.""" + plan = _plan(case_dir, tmp_path / "out", realms=[realm]) + steps = plan.for_stage("cmor") + assert steps + for step in steps: + assert step.command[step.command.index("--model-res") + 1] == "NorESM3-LM" # ------------------------------------------- forwarding the grid decision @@ -516,5 +517,4 @@ def test_landice_runs_every_stage_without_extra_input(self, case_dir, tmp_path): assert {step.stage for step in plan.steps} == set(STAGES) assert not plan.skipped command = _command_of(plan, "cmor-landIce-yr") - assert "--atmos-res" not in command assert command[command.index("--ice-sheet") + 1] == "gris" diff --git a/tests/test_regrid_latlon.py b/tests/test_regrid_latlon.py index 6c9b515..07e6113 100644 --- a/tests/test_regrid_latlon.py +++ b/tests/test_regrid_latlon.py @@ -254,14 +254,14 @@ def test_tables_cover_driver_resolution_choices(): The driver's choices and the YAML keys are edited in different files, so they can drift apart; a missing key turns a valid run into a ValueError. """ - driver_choices = {"ne16", "ne30", "tx2_3v2", "tnx1v4", "regular"} + driver_choices = {"ne16", "ne30", "tx2_3v2", "tnx1v4", "custom"} defined = set() for model in ("cesm", "noresm"): defined |= set(load_regrid_maps(model)["resolutions"]) assert driver_choices <= defined -def test_regular_resolution_has_maps(): - """'regular' skips regridding, but the fx path still asks for a map.""" +def test_custom_resolution_has_maps(): + """'custom' skips regridding, but the fx path still asks for a map.""" for model in ("cesm", "noresm"): - assert "conservative" in get_map_paths(model, "regular") + assert "conservative" in get_map_paths(model, "custom")