diff --git a/.github/scripts/pr_mod_file_tests.py b/.github/scripts/pr_mod_file_tests.py index 6b1ba1375..150febff8 100755 --- a/.github/scripts/pr_mod_file_tests.py +++ b/.github/scripts/pr_mod_file_tests.py @@ -140,15 +140,18 @@ def _main_prog(): print("Generating list of modified files...") - #This should eventually be passed in via a command-line - #argument, and include everything inside the "lib" directory -JN: - testable_files = {"lib/adf_base.py", - "lib/adf_config.py", - "lib/adf_file_utils.py", - "lib/adf_info.py", - "lib/adf_obs.py", - "lib/adf_web.py", - "lib/adf_diag.py"} + # This should eventually be passed in via a command-line + # argument, and include everything inside the "lib" directory -JN: + testable_files = { + "lib/adf_base.py", + "lib/adf_config.py", + "lib/adf_file_utils.py", + "lib/adf_info.py", + "lib/adf_obs.py", + "lib/adf_units.py", + "lib/adf_web.py", + "lib/adf_diag.py", + } #+++++++++++++++++++++++ #Read in input arguments diff --git a/lib/adf_dataset.py b/lib/adf_dataset.py index 6caedee6f..227e3d0c5 100644 --- a/lib/adf_dataset.py +++ b/lib/adf_dataset.py @@ -5,6 +5,7 @@ import xarray as xr import adf_utils as utils +from adf_units import units_equivalent warnings.formatwarning = utils.my_formatwarning # "reference data" @@ -294,6 +295,16 @@ def load_climo_ds(self, case, variablename): ds = self.load_dataset(fils) if ds is None: return None + if (scale_factor != 1 or add_offset != 0) and self.already_converted( + ds[variablename].attrs, variablename + ): + self.adf.debug_log( + f"\t INFO: '{variablename}' climo file is already" + " in converted units; not converting again." + ) + scale_factor = 1 + add_offset = 0 + # End if # xarray arithmetic drops attrs, so carry them across by hand -- otherwise # the regridded files lose 'units' and the plotting scripts KeyError on it. attrs = ds[variablename].attrs.copy() @@ -368,10 +379,10 @@ def load_reference_climo_ds(self, case, variablename, apply_scaling=True): if ds is None: return None vname = self.ref_var_nam[variablename] # name of variable in the reference data - # Check if already transformed (via attribute or units) - new_unit = self.adf.variable_defaults.get(variablename, {}).get('new_unit') - unit_match = new_unit is not None and ds[vname].attrs.get('units') == new_unit - if ds[vname].attrs.get('transformed', False) or unit_match: + # Check if already transformed. The units comparison was a literal + # string match, which almost never fired: an observation file says + # "W/m2" where the defaults say "Wm$^{-2}$". + if self.already_converted(ds[vname].attrs, variablename): apply_scaling = False if not apply_scaling: add_offset = 0 @@ -381,6 +392,13 @@ def load_reference_climo_ds(self, case, variablename, apply_scaling=True): ds[vname] = ds[vname] * scale_factor + add_offset ds[vname].attrs = attrs if scale_factor != 1 or add_offset != 0: + # Rename the units as load_climo_ds does for the test cases. Without + # this the reference kept the units it arrived with while holding + # converted values, so the file said 'm/s' over data in mm/day. + new_unit = self.adf.variable_defaults.get(variablename, {}).get("new_unit") + if new_unit: + ds[vname].attrs["units"] = new_unit + # End if # int, not bool: netCDF4 cannot store a Python bool as an attribute ds[vname].attrs['transformed'] = 1 return ds @@ -394,7 +412,13 @@ def load_reference_climo_da(self, case, variablename, apply_scaling=True): scale_factor = 1 else: add_offset, scale_factor = self.get_value_converters(case, variablename) - return self.load_da(fils, vname, add_offset=add_offset, scale_factor=scale_factor) + return self.load_da( + fils, + vname, + field=variablename, + add_offset=add_offset, + scale_factor=scale_factor, + ) def get_reference_climo_file(self, var): """Return a list of files to be used as reference (aka baseline) for variable var.""" @@ -510,7 +534,15 @@ def load_reference_regrid_da(self, case, field, apply_scaling=None): file_field = self.ref_var_nam[field] if self.adf.compare_obs else field add_offset, scale_factor = self._regrid_converters(fils, file_field, case, field, apply_scaling) - return self.load_da(fils, file_field, add_offset=add_offset, scale_factor=scale_factor) + # field, not file_field: an observation file names the variable its own + # way, and the variable defaults are keyed by the ADF name + return self.load_da( + fils, + file_field, + field=field, + add_offset=add_offset, + scale_factor=scale_factor, + ) def _regrid_converters(self, fils, file_field, case, field, apply_scaling): """Return the (add_offset, scale_factor) to use for a regridded file. @@ -532,7 +564,7 @@ def _regrid_converters(self, fils, file_field, case, field, apply_scaling): if apply_scaling or (scale_factor == 1 and add_offset == 0): return add_offset, scale_factor ds = self.load_dataset(fils) - if ds is not None and ds[file_field].attrs.get('transformed', 0): + if ds is not None and self.already_converted(ds[file_field].attrs, field): return 0, 1 return add_offset, scale_factor @@ -569,10 +601,101 @@ def load_dataset(self, fils, use_time_bounds=False): # End if return ds - def load_da(self, fils, variablename, use_time_bounds=False, **kwargs): + def already_converted(self, attrs, field): + """Report whether the conversion for `field` has already been applied. + + Parameters + ---------- + attrs : dict + attributes of the variable as it was read from the file + field : str + ADF name of the variable, i.e. the key into the variable defaults + + Returns + ------- + bool + ``True`` when the file already holds converted values. + + Notes + ----- + Two things say so. A file the ADF wrote carries ``transformed``, + stamped by whichever load applied the conversion. A file it did not + write -- an observation file, or a regridded file from an older ADF -- + carries no stamp, and then the only evidence is that its units are + already the units the defaults are converting to. That comparison is + made with :func:`units_equivalent`, because the two strings come from + different hands: CAM writes ``W/m2`` where the defaults say + ``Wm$^{-2}$``, and comparing them literally answers the wrong + question. + """ + if attrs.get("transformed", False): + return True + # End if + new_unit = self.adf.variable_defaults.get(field, {}).get("new_unit") + return units_equivalent(attrs.get("units"), new_unit) + + def apply_conversion(self, data, field, case=None): + """Apply the variable-defaults conversion to `data`, once. + + For a script that opens a file itself rather than through one of the + loaders here -- the TEM files and the vector plots do -- so that it + makes the same decision they do. + + Parameters + ---------- + data : xarray.DataArray + the variable as it was read from the file + field : str + ADF name of the variable, i.e. the key into the variable defaults + case : str, optional + case the data belongs to, which decides whether the observation + converters are used. Defaults to `field`'s test-case converters. + + Returns + ------- + xarray.DataArray + The converted data, carrying its attributes, with ``units`` set + to the defaults' ``new_unit`` and ``transformed`` stamped on when + a conversion was applied. Returned unchanged when the file + already holds converted values. + """ + if data is None: + return None + # End if + add_offset, scale_factor = self.get_value_converters( + case if case is not None else field, field + ) + if scale_factor == 1 and add_offset == 0: + return data + # End if + if self.already_converted(data.attrs, field): + dmsg = f"\t INFO: '{field}' is already in converted units in the" + dmsg += " file, so the conversion is not applied again." + self.adf.debug_log(dmsg) + return data + # End if + attrs = data.attrs.copy() + converted = data * scale_factor + add_offset + converted.attrs = attrs + new_unit = self.adf.variable_defaults.get(field, {}).get("new_unit") + if new_unit: + converted.attrs["units"] = new_unit + # End if + # int, not bool: netCDF4 cannot store a Python bool as an attribute + converted.attrs["transformed"] = 1 + return converted + + def load_da(self, fils, variablename, use_time_bounds=False, field=None, **kwargs): """Return xarray DataArray from file(s) w/ optional scale factor, offset, new units. `use_time_bounds` is passed to `load_dataset`; see there. + + `field` is the ADF name of the variable when it differs from its name + in the file, which is the case for observations. The variable + defaults are keyed by the ADF name. + + A conversion that has already been applied to the file is not applied + again; see :meth:`already_converted`. """ ds = self.load_dataset(fils, use_time_bounds=use_time_bounds) if ds is None: @@ -581,12 +704,25 @@ def load_da(self, fils, variablename, use_time_bounds=False, **kwargs): da = ds[variablename].squeeze() scale_factor = kwargs.get('scale_factor', 1) add_offset = kwargs.get('add_offset', 0) + + if (scale_factor != 1 or add_offset != 0) and self.already_converted( + da.attrs, field if field is not None else variablename + ): + dmsg = f"\t INFO: '{variablename}' is already in converted units" + dmsg += " in the file, so the conversion is not applied again." + self.adf.debug_log(dmsg) + scale_factor = 1 + add_offset = 0 + # End if + attrs = da.attrs.copy() da = da * scale_factor + add_offset da.attrs = attrs if scale_factor != 1 or add_offset != 0: - new_unit = self.adf.variable_defaults.get(variablename, {}).get("new_unit") + new_unit = self.adf.variable_defaults.get( + field if field is not None else variablename, {} + ).get("new_unit") if new_unit: da.attrs['units'] = new_unit # Stamp on any conversion, not only one that renames the units -- diff --git a/lib/adf_units.py b/lib/adf_units.py new file mode 100644 index 000000000..46308d474 --- /dev/null +++ b/lib/adf_units.py @@ -0,0 +1,236 @@ +""" +Unit string handling. + +Kept apart from adf_utils so it can be unit tested without the scientific +stack: the ADF unit test workflow installs only PyYAML and pytest, so anything +that imports xarray/geocat at module level cannot be exercised in CI. This +module imports nothing but re. + +Functions +--------- +normalize_units(units) + Reduce a unit string to a canonical form for comparison. +units_equivalent(first, second) + Report whether two unit strings mean the same thing. + +Notes +----- +The ADF compares units to decide whether the conversion named in the variable +defaults has already been applied to a file. The two strings being compared +come from different places -- one written by CAM, the other by whoever edited +`adf_variable_defaults.yaml` -- so they agree on the physics far more often +than they agree character for character. `W/m2`, `W m-2`, `Wm^-2` and the +LaTeX `W m$^{-2}$` are one unit written four ways, and the shipped defaults +alone contain both `Wm$^{-2}$` and `W m$^{-2}$`. Comparing the raw strings +answers "were these typed the same way", which is not the question being +asked, and getting it wrong scales the data twice. + +Rendering a unit for somewhere with no LaTeX renderer, such as a table cell, +is `adf_utils.plain_text_units`; this module only compares. +""" + +import re + +# Unit names that mean the same thing. Keys and values are compared after +# normalization, so only real synonyms belong here, not spelling variants: +_ALIASES = { + "fraction": "1", + "frac": "1", + "unitless": "1", + "dimensionless": "1", + "none": "1", + "-": "1", + "[-]": "1", + "percent": "%", + "pct": "%", + "degrees_kelvin": "k", + "deg_k": "k", + "kelvin": "k", + "degrees_celsius": "degc", + "celsius": "degc", + "deg_c": "degc", + "degrees_east": "degrees_east", + "meters": "m", + "meter": "m", + "metre": "m", + "seconds": "s", + "second": "s", + "sec": "s", + "days": "d", + "day": "d", + "grams": "g", + "gram": "g", + "micron": "um", + "microns": "um", +} + +# Factors written without a space between them. Case matters and is still +# intact here, which is the point: "Nm" is newton-metre while "nm" is +# nanometre, and lower-casing first makes them the same string. Only forms +# that actually occur are listed -- a general splitter guessing where to cut a +# run of letters reads "Sv" as siemens-volt and "nm" as newton-metre, the +# second of which reports two different units as equal. Note that "ms" is +# read as metre-second: no CAM field is reported in milliseconds, and +# "ms$^{-1}$" for wind speed is in the shipped variable defaults today. +_RUN_TOGETHER = { + "Wm": "W m", + "Nm": "N m", + "Jm": "J m", + "Km": "K m", + "kgm": "kg m", + "gm": "g m", + "ms": "m s", + "Pas": "Pa s", +} + +# LaTeX and unicode fragments that carry no meaning for a comparison: +_LATEX = ( + (r"\mathrm", ""), + (r"\text", ""), + (r"\mu", "u"), + (r"\,", " "), + (r"\;", " "), + (r"\ ", " "), + ("µ", "u"), # micro sign + ("μ", "u"), # greek small letter mu + ("·", " "), # middle dot, used as a multiplication sign + ("**", "^"), +) + +# Superscript digits, which appear in units copied out of documents: +_SUPERSCRIPTS = str.maketrans( + { + "⁰": "0", + "¹": "1", + "²": "2", + "³": "3", + "⁴": "4", + "⁵": "5", + "⁶": "6", + "⁷": "7", + "⁸": "8", + "⁹": "9", + "⁻": "-", + } +) + + +def _strip_markup(units): + """Return `units` with LaTeX, superscripts and stray braces removed.""" + text = units.translate(_SUPERSCRIPTS) + for old, new in _LATEX: + text = text.replace(old, new) + # A LaTeX exponent is written ^{-2}; the braces are noise once the + # superscript digits are gone: + text = re.sub(r"\^\s*\{([^}]*)\}", r"^\1", text) + text = text.replace("$", "").replace("{", "").replace("}", "") + # Split the run-together factors before anything lower-cases the string, + # because which factors they are depends on their case: + for joined, apart in _RUN_TOGETHER.items(): + text = re.sub(rf"(? 'has_lat', 'has_lev', with T/F values - - #Notify user of level dimension: - if validate_lat_lev['has_lev']: - print(f"\t INFO: {var} has lev dimension.") - has_lev = True - else: - has_lev = False - - # # Seasonal Averages - # Note: xarray can do seasonal averaging, but depends on having time accessor, - # which these prototype climo files don't. - # - - #Create new dictionaries: - mseasons = {} - oseasons = {} - - #Loop over season dictionary: - for s in seasons: - plot_name = plot_loc / f"{var}_{s}_Meridional_Mean.{plot_type}" - - # Check redo_plot. If set to True: remove old plot, if it already exists: - if (not redo_plot) and plot_name.is_file(): - #Add already-existing plot to website (if enabled): - adfobj.debug_log(f"'{plot_name}' exists and clobber is false.") - adfobj.add_website_data(plot_name, var, case_name, season=s, - plot_type="Meridional") - #Continue to next iteration: - continue - elif (redo_plot) and plot_name.is_file(): - plot_name.unlink() - - mseasons[s] = utils.seasonal_mean(mdata, season=s, is_climo=True) - oseasons[s] = utils.seasonal_mean(odata, season=s, is_climo=True) - - - #Create new plot: - pf.plot_meridional_mean_and_save(plot_name, case_nickname, base_nickname, - [syear_cases[case_idx],eyear_cases[case_idx]], - [syear_baseline,eyear_baseline], - mseasons[s], oseasons[s], has_lev, latbounds=slice(-5,5), obs=obs, **vres) - - #Add plot to website (if enabled): - adfobj.add_website_data(plot_name, var, case_name, season=s, - plot_type="Meridional") - - #End for (seasons loop) - #End for (case names loop) - #End for (obs/baseline loop) - #End for (variables loop) + mseasons[s] = utils.seasonal_mean(mdata, season=s, is_climo=True) + oseasons[s] = utils.seasonal_mean(odata, season=s, is_climo=True) + + # Create new plot: + pf.plot_meridional_mean_and_save( + plot_name, + case_nickname, + adfobj.data.ref_nickname, + [syear_cases[case_idx], eyear_cases[case_idx]], + [syear_baseline, eyear_baseline], + mseasons[s], + oseasons[s], + has_lev, + latbounds=slice(-5, 5), + obs=adfobj.compare_obs, + **vres, + ) + + # Add plot to website (if enabled): + adfobj.add_website_data( + plot_name, var, case_name, season=s, plot_type="Meridional" + ) + + # End for (seasons loop) + # End for (case names loop) + # End for (variables loop) #Notify user that script has ended: print(" ...Meridional mean plots have been generated successfully.") @@ -254,5 +297,41 @@ def meridional_mean(adfobj): ######### +def _all_plots_present(plot_locations, case_names, var, seasons, plot_type, merid_skip): + """ + Report whether every plot this variable would produce is already there. + + Parameters + ---------- + plot_locations : list + output plot directory for each case + case_names : list + names of the test cases + var : str + ADF name of the variable being plotted + seasons : dict + the seasons being plotted, keyed by name + plot_type : str + file extension the plots are written with, e.g. "png" + merid_skip : set + plots found by the pre-flight pass above + + Returns + ------- + bool + ``True`` when nothing is left to draw for this variable, so the + reference and test data never have to be opened. + """ + for case_idx in range(len(case_names)): + plot_loc = Path(plot_locations[case_idx]) + for s in seasons: + if (plot_loc / f"{var}_{s}_Meridional_Mean.{plot_type}") not in merid_skip: + return False + # End if + # End for + # End for + return True + + ############## -#END OF SCRIPT \ No newline at end of file +# END OF FILE diff --git a/scripts/plotting/polar_map.py b/scripts/plotting/polar_map.py index 7893fee8e..5d634b49a 100644 --- a/scripts/plotting/polar_map.py +++ b/scripts/plotting/polar_map.py @@ -4,6 +4,7 @@ # ADF library import plotting_functions as pf +import plotting_utils as plot_utils import adf_utils as utils def get_hemisphere(hemi_type): @@ -138,10 +139,32 @@ def polar_map(adfobj): vres = res.get(var, {}) web_category = vres.get("category", None) + # A complete set of plots can be recognised from the file names alone, + # because a 2-D variable and a 3-D one are named differently. Doing + # that first means a re-run never opens the regridded files just to + # find out which of the two it is looking at. + settled = _existing_plot_set( + plot_locations, case_names, var, seasons, pres_levs, plot_type + ) + if settled and not redo_plot: + for path, web_name, case_name, season, hemi_type in settled: + adfobj.add_website_data( + path, + web_name, + case_name, + category=web_category, + season=season, + plot_type=hemi_type, + ) + # End for + print(f"\t Skipping {var} - all plots already exist") + continue + # End if + # Get all plot info and check existence plot_info = [] all_plots_exist = True - + for case_idx, case_name in enumerate(case_names): plot_loc = Path(plot_locations[case_idx]) @@ -266,5 +289,73 @@ def polar_map(adfobj): ############## #END OF `polar_map` function + +def _existing_plot_set(plot_locations, case_names, var, seasons, pres_levs, plot_type): + """ + Return the complete set of polar plots for `var`, if one is on disk. + + A 2-D variable is plotted as ``{var}_{season}_{hemisphere}_Mean`` and a + 3-D one as ``{var}_{pressure}hpa_{season}_{hemisphere}_Mean``, so which + of the two a variable is can be read off the file names. That is what + lets a re-run decide there is nothing to do without opening the regridded + file to look for a ``lev`` dimension. + + Parameters + ---------- + plot_locations : list + output plot directory for each case + case_names : list + names of the test cases + var : str + ADF name of the variable being plotted + seasons : list or dict + the seasons being plotted + pres_levs : list + the configured pressure levels, empty when none are set + plot_type : str + file extension the plots are written with, e.g. "png" + + Returns + ------- + list + ``(path, website name, case, season, hemisphere)`` for every plot, + when a complete set exists; an empty list when it does not. + """ + flat = [] + levelled = [] + for case_idx, case_name in enumerate(case_names): + plot_loc = Path(plot_locations[case_idx]) + for season in seasons: + for hemi_type in ["NHPolar", "SHPolar"]: + flat.append( + ( + plot_loc / f"{var}_{season}_{hemi_type}_Mean.{plot_type}", + var, + case_name, + season, + hemi_type, + ) + ) + for pres in pres_levs: + levelled.append( + ( + plot_loc + / ( + f"{var}_{pres}hpa_{season}_{hemi_type}" + f"_Mean.{plot_type}" + ), + f"{var}_{pres}hpa", + case_name, + season, + hemi_type, + ) + ) + # End for + # End for + # End for + # End for + return plot_utils.first_complete_plot_set([levelled, flat]) + + ############## -# END OF FILE \ No newline at end of file +# END OF FILE diff --git a/scripts/plotting/tem.py b/scripts/plotting/tem.py index 13f9c49d4..a72c153ec 100644 --- a/scripts/plotting/tem.py +++ b/scripts/plotting/tem.py @@ -408,19 +408,17 @@ def tem(adf): mdata = ds[var].squeeze() odata = ds_base[var].squeeze() - #Apply the unit conversion from the variable defaults. TEM - #files carry little metadata, so the new unit is taken from the - #defaults too. Observations have their own scaling, which is - #assumed to bring them to the same units, so they keep the unit - #string they arrived with. - mdata = mdata * vres.get("scale_factor", 1) + vres.get("add_offset", 0) - mdata.attrs['units'] = vres.get("new_unit", mdata.attrs.get('units', 'none')) + # Apply the unit conversion from the variable defaults, through + # the ADF's data layer so that a TEM file already holding + # converted values is not scaled a second time. Observations + # have their own scaling, which is assumed to bring them to the + # same units, so they keep the unit string they arrived with. + mdata = adf.data.apply_conversion(mdata, var) if obs: odata = (odata * vres.get("obs_scale_factor", 1) + vres.get("obs_add_offset", 0)) else: - odata = odata * vres.get("scale_factor", 1) + vres.get("add_offset", 0) - odata.attrs['units'] = vres.get("new_unit", odata.attrs.get('units', 'none')) + odata = adf.data.apply_conversion(odata, var) #End if #Month-length weighted seasonal (or annual) mean. The weighted @@ -571,4 +569,4 @@ def _mesh(da): print(" ...TEM plots have been generated successfully.") # Helper functions -################## \ No newline at end of file +################## diff --git a/scripts/plotting/zonal_mean.py b/scripts/plotting/zonal_mean.py index 2f8630fb3..97a19928f 100644 --- a/scripts/plotting/zonal_mean.py +++ b/scripts/plotting/zonal_mean.py @@ -100,7 +100,7 @@ def zonal_mean(adfobj): for var in var_list: for s in seasons: #Check zonal log-p: - plot_name_log = plot_loc / f"{var}_{s}_Zonal_logp_Mean.{plot_type}" + plot_name_log = plot_loc / f"{var}_logp_{s}_Zonal_Mean.{plot_type}" # Check redo_plot. If set to True: remove old plot, if it already exists: if (not redo_plot) and plot_name_log.is_file(): @@ -252,10 +252,6 @@ def zonal_mean(adfobj): # because we can let any pressure-level interpolation happen there # This could be re-visited for efficiency or improved code structure. - #Seasonal Averages - mseasons[s] = utils.seasonal_mean(mdata, season=s, is_climo=True) - oseasons[s] = utils.seasonal_mean(odata, season=s, is_climo=True) - #Set the file name plot_name = plot_loc / f"{var}_{s}_Zonal_Mean.{plot_type}" plot_name_log = None @@ -265,6 +261,19 @@ def zonal_mean(adfobj): plot_name_log = plot_loc / f"{var}_logp_{s}_Zonal_Mean.{plot_type}" #End if + # Nothing to draw for this season: both plots were found above + # and redo_plot is false. The seasonal averages below are the + # expensive part of a re-run, so work that out before doing them. + if (plot_name in zonal_skip) and ( + (plot_name_log is None) or (plot_name_log in logp_zonal_skip) + ): + continue + # End if + + # Seasonal Averages + mseasons[s] = utils.seasonal_mean(mdata, season=s, is_climo=True) + oseasons[s] = utils.seasonal_mean(odata, season=s, is_climo=True) + #Create plots if plot_name not in zonal_skip: @@ -299,4 +308,4 @@ def zonal_mean(adfobj): ############## -#END OF SCRIPT \ No newline at end of file +#END OF SCRIPT