From c078e7226468e5939c36796da83affddb4bf9cb8 Mon Sep 17 00:00:00 2001 From: alperaltuntas Date: Wed, 19 Aug 2026 09:32:51 -0600 Subject: [PATCH 1/2] Name diag_table streams after their output files The label of each entry in the Files section of `diag_table.yaml` is now the stream name, i.e., the part of the output file name that follows the case name and `mom6`: `sigma2_hist` becomes `h.rho2`, `hist_z_space` becomes `h.z`, `visc_and_diff_daily_avg` becomes `h.visc`, and so on. Users see these names on disk and in the archive, so these are also the names they will refer to once we add the ability to modify the diag_table via `user_nl_mom`. The date template of the file name is no longer written by hand: it is derived from how often a new file is started, since a name that lacks that time resolution gets overwritten by the next file of the same stream. The `suffix` entry is therefore gone. Entries that need a configuration dependent name provide an optional guarded `name` entry instead (only the MARBL entries do, since they rename their files in spinup runs.) Settings shared by all files moved to a new `FileDefaults` section, and `new_file_freq` now defaults to one file per unit of the output frequency. Together, these bring `diag_table.yaml` from 769 to 404 lines without losing expressiveness: the base name of every stream was already the same in all 32 combinations of `TEST`, `OCN_DIAG_MODE`, `OCN_GRID`, and `OCN_DIAG_SECTIONS`. `FType_diag_table` is restructured around the stream model that now lives in `diag_table_streams`, which has no CIME dependency so that it can be tested on its own. Generating the table in terms of streams rather than template entries also fixes a few latent problems: - Two entries that resolve to the same file are now merged into one stream. MARBL's `low` and `low_native_z` do resolve to the same file when `MARBL_HIST_VERT_GRID` is native or both, and used to emit a duplicate file entry in Section-1 along with duplicate `geolat` and `geolon` fields. (Its `medium` entries have the same problem, worked around in `MARBL_diags_to_diag_table` by combining them.) - A fields or lists block whose guards are all false reduces to None, which is now skipped rather than passed to len() or max(). Only the ordering of the blocks kept this from being reachable. - The check for a field listed more than once in the same file only ever looked within a single fields block. - Whether `new_file_freq_units` gets written was decided by testing for `time_axis_units`. - Stream settings are now validated, and a malformed template gets an error naming the entry it came from rather than a TypeError or a KeyError from deeper down. Also drops an unused import in `FType_input_nml`. status: no answer changes. MOM6 reads the table no differently, and diagnostic contents are unchanged in all the configurations compared, with these exceptions: `h.visc` in tests and, with `MARBL_HIST_VERT_GRID` native or both, `h.bgc.native` in tests both started a new file every day while their names only resolved to the month, so they used to overwrite themselves and now get a daily date template. The section files say "mean" where they used to say ".true."; FMS maps both to the same time average (see `init_output_field` in `diag_util.F90`). Column padding is now driven by what is actually written to the table. testing: rendered the diag_table for 96 combinations of `TEST`, `OCN_DIAG_MODE`, `OCN_GRID`, `OCN_DIAG_SECTIONS`, and `MARBL_HIST_VERT_GRID`, and compared each against the current output. The scripts in `tests/` pass. aux_mom has not been run yet. Co-Authored-By: Claude Opus 5 --- .../MARBL_diags_to_diag_table.py | 117 ++-- cime_config/MOM_RPS/FType_diag_table.py | 325 ++++++----- cime_config/MOM_RPS/FType_input_nml.py | 1 - cime_config/MOM_RPS/diag_table_streams.py | 487 ++++++++++++++++ param_templates/diag_table.yaml | 543 +++--------------- param_templates/json/diag_table.json | 538 ++--------------- 6 files changed, 835 insertions(+), 1176 deletions(-) create mode 100644 cime_config/MOM_RPS/diag_table_streams.py diff --git a/cime_config/MARBL_scripts/MARBL_diags_to_diag_table.py b/cime_config/MARBL_scripts/MARBL_diags_to_diag_table.py index 25d2c720..eb80d2b6 100755 --- a/cime_config/MARBL_scripts/MARBL_diags_to_diag_table.py +++ b/cime_config/MARBL_scripts/MARBL_diags_to_diag_table.py @@ -65,10 +65,10 @@ def __init__(self, vert_grid): # "medium" frequency should be treated like "mom6.h.native" stream -- annual in spinup runs, monthly otherwise # i. 2D vars new_file_freq_units = "days" if self._nstep_output else None - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_annual%4yr", - "$TEST == True": "h.bgc.native%4yr-%2mo-%2dy", - "else": "h.bgc.native%4yr-%2mo", + name_dict = { + '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_annual", + "$TEST == True": "h.bgc.native", + "else": "h.bgc.native", } output_freq_units_dict = { '$OCN_DIAG_MODE == "spinup"': "years", @@ -77,32 +77,32 @@ def __init__(self, vert_grid): "else": "months", } self._diag_table_dict["medium"] = self._dict_template( - suffix_dict, output_freq_units_dict, new_file_freq_units=new_file_freq_units + name_dict, output_freq_units_dict, new_file_freq_units=new_file_freq_units ) # ii. 3D vars on interpolated grid if vert_grid in ["interpolated", "both"]: - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.z_annual%4yr", - "$TEST == True": "h.bgc.z%4yr-%2mo-%2dy", - f"{self._nstep_output} == True": "h.bgc.z_nstep%4yr-%2mo-%2dy", - "else": "h.bgc.z%4yr-%2mo", + name_dict = { + '$OCN_DIAG_MODE == "spinup"': "h.bgc.z_annual", + "$TEST == True": "h.bgc.z", + f"{self._nstep_output} == True": "h.bgc.z_nstep", + "else": "h.bgc.z", } self._diag_table_dict["medium_z"] = self._dict_template( - suffix_dict, + name_dict, output_freq_units_dict, new_file_freq_units=new_file_freq_units, module="ocean_model_z", ) # iii. 3D vars on native grid if vert_grid in ["native", "both"]: - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_annual%4yr", - "$TEST == True": "h.bgc.native%4yr-%2mo", - f"{self._nstep_output} == True": "h.bgc.native_nstep%4yr-%2mo-%2dy", - "else": "h.bgc.native%4yr-%2mo", + name_dict = { + '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_annual", + "$TEST == True": "h.bgc.native", + f"{self._nstep_output} == True": "h.bgc.native_nstep", + "else": "h.bgc.native", } self._diag_table_dict["medium_native_z"] = self._dict_template( - suffix_dict, + name_dict, output_freq_units_dict, new_file_freq_units=new_file_freq_units, module="ocean_model", @@ -111,9 +111,9 @@ def __init__(self, vert_grid): # "high" frequency should be treated like "mom6.h.sfc" stream -- 5-day averages in spinup, daily otherwise # unlike "sfc", this stream will write one file per month instead of per year (except in spinup) # i. 2D vars - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.daily5%4yr", - "else": "h.bgc.daily%4yr-%2mo", + name_dict = { + '$OCN_DIAG_MODE == "spinup"': "h.bgc.daily5", + "else": "h.bgc.daily", } output_freq_dict = {'$OCN_DIAG_MODE == "spinup"': 5, "else": 1} new_file_freq_units_dict = { @@ -121,16 +121,16 @@ def __init__(self, vert_grid): "else": "months", } self._diag_table_dict["high"] = self._dict_template( - suffix_dict, "days", new_file_freq_units_dict, output_freq_dict + name_dict, "days", new_file_freq_units_dict, output_freq_dict ) # ii. 3D vars on interpolated grid if vert_grid in ["interpolated", "both"]: - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.z_daily5%4yr", - "else": "h.bgc.z_daily%4yr-%2mo", + name_dict = { + '$OCN_DIAG_MODE == "spinup"': "h.bgc.z_daily5", + "else": "h.bgc.z_daily", } self._diag_table_dict["high_z"] = self._dict_template( - suffix_dict, + name_dict, "days", new_file_freq_units_dict, output_freq_dict, @@ -138,12 +138,8 @@ def __init__(self, vert_grid): ) # iii. 3D vars on native grid if vert_grid in ["native", "both"]: - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_daily5%4yr", - "else": "h.bgc.native_daily5%4yr-%2mo", - } self._diag_table_dict["high_native_z"] = self._dict_template( - suffix_dict, + "h.bgc.native_daily5", "days", new_file_freq_units_dict, output_freq_dict, @@ -152,28 +148,28 @@ def __init__(self, vert_grid): # "low" frequency should be treated as annual averages # i. 2D vars - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_annual2%4yr", - "else": "h.bgc.native_annual%4yr", + name_dict = { + '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_annual2", + "else": "h.bgc.native_annual", } - self._diag_table_dict["low"] = self._dict_template(suffix_dict, "years") + self._diag_table_dict["low"] = self._dict_template(name_dict, "years") # ii. 3D vars on interpolated grid if vert_grid in ["interpolated", "both"]: - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.z_annual2%4yr", - "else": "h.bgc.z_annual%4yr", + name_dict = { + '$OCN_DIAG_MODE == "spinup"': "h.bgc.z_annual2", + "else": "h.bgc.z_annual", } self._diag_table_dict["low_z"] = self._dict_template( - suffix_dict, "years", module="ocean_model_z" + name_dict, "years", module="ocean_model_z" ) # iii. 3D vars on native grid if vert_grid in ["native", "both"]: - suffix_dict = { - '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_annual2%4yr", - "else": "h.bgc.native_annual%4yr", + name_dict = { + '$OCN_DIAG_MODE == "spinup"': "h.bgc.native_annual2", + "else": "h.bgc.native_annual", } self._diag_table_dict["low_native_z"] = self._dict_template( - suffix_dict, "years", module="ocean_model" + name_dict, "years", module="ocean_model" ) def update(self, varname, frequency, is2D, lMARBL_output_all, vert_grid): @@ -274,11 +270,10 @@ def dump_to_json(self, filename): def _dict_template( self, - suffix, + name, output_freq_units, new_file_freq_units=None, - output_freq=1, - new_file_freq=1, + output_freq=None, module="ocean_model", ): """ @@ -286,34 +281,30 @@ def _dict_template( Variables will be added to output file by appending to template["fields"]['$OCN_DIAG_MODE != "none"']["lists"][0] Parameters: - * suffix: string used to identify output file; could also be a dictionary - where keys are logical evaluations + * name: name of the output stream, which is also the segment of the file + name that follows the case name and the component name; could + also be a dictionary where keys are logical evaluations * output_freq_units: units used to determine how often to output; similar - to suffix, this can also be a dictionary + to name, this can also be a dictionary * new_file_freq_units: units used to determine how often to generate new stream files; if None, will use output_freq_units (default: None) - * output_freq: how frequently to output (default: 1) - * new_file_freq: how frequently to create new files (default: 1) + * output_freq: how frequently to output; if None, the diag_table default + of one output per output_freq_units is used (default: None) * module: string that determines vertical grid; "ocean_model_z" maps to Z space, "ocean_model" stays on native grid, "ocean_model_rho2" is sigma2 """ template = dict() - template["suffix"] = suffix - template["output_freq"] = output_freq - template["new_file_freq"] = new_file_freq + template["name"] = name + if output_freq is not None: + template["output_freq"] = output_freq + # Note that this cannot be left to the diag_table default, which orders + # its guards differently: a run that is both a test and a spinup run + # outputs daily here, but annually by default. template["output_freq_units"] = output_freq_units if new_file_freq_units: template["new_file_freq_units"] = new_file_freq_units - else: - template["new_file_freq_units"] = output_freq_units - template["time_axis_units"] = "days" - template["reduction_method"] = "mean" - template["regional_section"] = "none" + template["packing"] = "= 1 if $TEST or $MARBL_DIAG_MODE == 'test_suite' else 2" template["fields"] = { - '$OCN_DIAG_MODE != "none"': { - "module": module, - "packing": "= 1 if $TEST or $MARBL_DIAG_MODE == 'test_suite' else 2", - "lists": [[]], - } + '$OCN_DIAG_MODE != "none"': {"module": module, "lists": [[]]} } return template diff --git a/cime_config/MOM_RPS/FType_diag_table.py b/cime_config/MOM_RPS/FType_diag_table.py index b0c8b8f6..5b0bab58 100644 --- a/cime_config/MOM_RPS/FType_diag_table.py +++ b/cime_config/MOM_RPS/FType_diag_table.py @@ -1,12 +1,26 @@ import os +from collections import OrderedDict + from CIME.ParamGen.paramgen import ParamGen +from diag_table_streams import ( + CASENAME, + DiagTableError, + Field, + FieldGroup, + int_setting, + REQUIRED_SETTINGS, + Stream, + STREAM_SETTINGS, + write_diag_table, +) + class FType_diag_table(ParamGen): """Encapsulates data and read/write methods for MOM6 diag_table input file.""" - @classmethod - def resolve(cls, unresolved_diag_table_path, resolved_diag_table_path, casename): + @staticmethod + def resolve(unresolved_diag_table_path, resolved_diag_table_path, casename): """Resolve the casename in an unresolved diag_table. Parameters @@ -25,178 +39,171 @@ def resolve(cls, unresolved_diag_table_path, resolved_diag_table_path, casename) with open(resolved_diag_table_path, "w") as resolved_diag_table: with open(unresolved_diag_table_path, "r") as diag_table_unresolved: for line in diag_table_unresolved: - resolved_diag_table.write(line.replace("${CASE}", casename)) + resolved_diag_table.write(line.replace(CASENAME, casename)) def write(self, output_path, case, MOM_input_final): - def get_all_fields(fields_block): - """Given a fields block, returns a list of all fields.""" - all_fields = [] - if fields_block is not None: - all_fields = [] - all_lists_blocks = [ - fields_block[lists_label] - for lists_label in fields_block - if lists_label.startswith("lists") - ] - for lists_block in all_lists_blocks: - if lists_block is not None: - all_fields.extend(sum(lists_block, [])) - return all_fields - - def is_empty_file(file_block): - """Returns true if the fields list of file is empty.""" - all_fields_blocks = [ - file_block[fields_label] - for fields_label in file_block - if fields_label.startswith("fields") - ] - for fields_block in all_fields_blocks: - if fields_block is None: - continue - all_lists_blocks = [ - fields_block[lists_label] - for lists_label in fields_block - if lists_label.startswith("lists") - ] - for lists_block in all_lists_blocks: - if len(lists_block) > 0: - return False - return True + """Writes out the diag_table of a case. + + Parameters + ---------- + output_path : str + The path of the diag_table to be created. The case name is left + unresolved in it, to be substituted later by the resolve method. + case : CIME.case.Case + The case whose diag_table is to be written. + MOM_input_final : FType_MOM_params + The MOM6 parameters of the case, i.e., MOM_input updated with + MOM_override. Consulted for expandable variables that are MOM6 + parameters rather than case variables. + """ def expand_func(varname): val = case.get_value(varname) if val is None: - val = MOM_input_final.data["Global"][varname]["value"] + val = ( + MOM_input_final.data.get("Global", {}).get(varname, {}).get("value") + ) if val is None: - raise RuntimeError("Cannot determine the value of variable: " + varname) + raise DiagTableError( + "Cannot determine the value of the variable {} appearing in " + "the diag_table template: it is neither a case variable nor a " + "MOM6 parameter of this case.".format(varname) + ) return val - # From the general template (diag_table.yaml), reduce a custom diag_table for this case + # From the general template (diag_table.yaml), reduce a custom diag_table + # for this case, and turn its file entries into streams. self.reduce(expand_func) - - with open(os.path.join(output_path), "w") as diag_table: - - # Print header: - casename = "${CASE}" - diag_table.write( - '"MOM6 diagnostic fields table for CESM case: ' + casename + '"\n' + write_diag_table(list(self._streams().values()), output_path) + + def _streams(self): + """Returns the streams of this case, keyed and ordered by stream name.""" + assert self.reduced, "May only collect streams from a reduced diag_table." + defaults = self.data.get("FileDefaults") or {} + _check_defaults(defaults) + if not self.data.get("Files"): + raise DiagTableError( + "The diag_table template has no Files section, and so describes " + "no output at all." ) - diag_table.write("1 1 1 0 0 0\n") # TODO - filename = lambda suffix: '"' + casename + ".mom6." + suffix + '"' - - # max filename length: - mfl = ( - max( - [ - len(filename(self._data["Files"][file_block_name]["suffix"])) - for file_block_name in self._data["Files"] - ] - ) - + 4 - ) # quotation marks and tabbing - - # Section 1: File section - diag_table.write("### Section-1: File List\n") - diag_table.write("#========================\n") + streams = OrderedDict() + for label, entry in self.data["Files"].items(): + stream = _stream_from_entry(label, entry, defaults) + if stream.name in streams: + streams[stream.name].merge(stream) + else: + streams[stream.name] = stream + return streams + + +def _check_defaults(defaults): + """Checks that the FileDefaults section only provides stream settings.""" + unknown = [key for key in defaults if key not in STREAM_SETTINGS] + if unknown: + raise DiagTableError( + "Unknown setting(s) {} in the FileDefaults section of the diag_table " + "template. Only the settings of a file may be given a default: " + "{}.".format(", ".join(unknown), ", ".join(STREAM_SETTINGS)) + ) - for file_block_name in self._data["Files"]: - file_block = self._data["Files"][file_block_name] - fname = filename(file_block["suffix"]) - # if the fields list(s) is empty, skip to the next file: - if is_empty_file(file_block): - continue +def _value_of(setting, entry, defaults, fallback=None): + """Returns the value of a setting: the entry's, else the default, else fallback.""" + for source in (entry, defaults): + value = source.get(setting) + if value is not None: + return value + return fallback + + +def _stream_from_entry(label, entry, defaults): + """Builds a Stream from one entry of the Files section of the template. + + Parameters + ---------- + label : str + The label of the entry, which is also the name of the stream unless the + entry provides an explicit, possibly configuration dependent, name. + entry : dict + The reduced entry, i.e., its settings and its fields blocks. + defaults : dict + The reduced FileDefaults section, providing the value of any setting that + the entry itself does not specify. + """ + unknown = [ + key + for key in entry + if key != "name" and key not in STREAM_SETTINGS and not key.startswith("fields") + ] + if unknown: + raise DiagTableError( + "Unknown setting(s) {} in diag_table template entry {}. Valid " + "settings are: {}, name, and fields blocks.".format( + ", ".join(unknown), label, ", ".join(STREAM_SETTINGS) + ) + ) + # A stream is named after its entry, unless the entry names itself. Only the + # MARBL entries do, in renaming their files to reflect the frequency of a + # spinup run, which a label cannot express because guards live in values. + stream = Stream(entry.get("name") or label, labels=[label]) + for setting in REQUIRED_SETTINGS: + stream.settings[setting] = _value_of(setting, entry, defaults) + + # A file that is written only once is never rolled over, and so has neither a + # new_file_freq nor a date template in its name. Otherwise a new file is + # started once per unit of the output frequency, unless said otherwise. + if int_setting(stream.settings["output_freq"], "output_freq", stream.source) > 0: + stream.settings["new_file_freq"] = _value_of( + "new_file_freq", entry, defaults, 1 + ) + stream.settings["new_file_freq_units"] = _value_of( + "new_file_freq_units", + entry, + defaults, + stream.settings["output_freq_units"], + ) - file_descr_str = ( - "{fname:" - + str(mfl) - + "s} {output_freq:3s} {output_freq_units:9s} 1, " - '{time_axis_units:9s} "time"' - ).format( - fname=fname + ",", - output_freq=str(file_block["output_freq"]) + ",", - output_freq_units='"' + file_block["output_freq_units"] + '",', - time_axis_units='"' + file_block["time_axis_units"] + '",', + for fields_label in [key for key in entry if key.startswith("fields")]: + fields_block = entry[fields_label] + if fields_block is None: # the guards of this block are all false + continue + try: + unknown = [ + key + for key in fields_block + if key != "module" and not key.startswith("lists") + ] + if unknown: + raise DiagTableError( + "the {} block has unknown key(s) {}. A fields block may only " + "have a module and lists of fields.".format( + fields_label, ", ".join(unknown) + ) ) - - if "new_file_freq" in file_block: - file_descr_str += ", " + str(file_block["new_file_freq"]) + ", " - if "time_axis_units" in file_block: - file_descr_str += ( - '"' + str(file_block["new_file_freq_units"]) + '"' - ) - diag_table.write(file_descr_str + "\n") - - diag_table.write("\n") - - ## Field section (per file): - diag_table.write("### Section-2: Fields List\n") - diag_table.write("#=========================\n") - for file_block_name in self._data["Files"]: - file_block = self._data["Files"][file_block_name] - fname = filename(file_block["suffix"]) - - # if the fields list(s) is empty, skip to the next file: - if is_empty_file(file_block): + if fields_block.get("module") is None: + raise DiagTableError( + "the {} block has no module. Add one, naming the MOM6 " + "diagnostics module that its fields come from.".format(fields_label) + ) + group = FieldGroup(fields_block["module"]) + for lists_label in [key for key in fields_block if key.startswith("lists")]: + lists_block = fields_block[lists_label] + if lists_block is None: # the guards of this block are all false continue - - # write the header for the fields list of this file block - diag_table.write("# {fname}\n".format(fname=fname)) - - # keep a record of all fields in this file to make sure no duplicate field exists - all_fields = [] - - # all of the fields blocks, i.e., blocks starting with "fields" prefix - all_fields_blocks = [ - file_block[fields_label] - for fields_label in file_block - if fields_label.startswith("fields") - ] - - # Loop over fields blocks - for field_block in all_fields_blocks: - module = field_block["module"] - packing = field_block["packing"] - field_list_1d = get_all_fields(field_block) - - # seperate field_name, alias, and reduction method - # (the latter two are optional) - field_list_1d_seperated = [] - for field in field_list_1d: - field_split = field.split(":") - field_name = field_split[0] - alias = field_name - reduction = file_block["reduction_method"] - assert 1 <= len(field_split) <= 3, ( - "Invalid field format: " + field - ) - if len(field_split) >= 2: - alias = field_split[1] - if len(field_split) >= 3: - reduction = field_split[2] - - field_list_1d_seperated.append((field_name, alias, reduction)) - - # check if there are any duplicate fields in the same file: - field_set = set() - for field_name, alias, reduction in field_list_1d_seperated: - if alias in field_set: - raise ValueError( - 'Field "' - + alias - + '" is listed more than once' - + " in file: " - + file_block["suffix"] + for field_list in lists_block: + if not isinstance(field_list, list): + raise DiagTableError( + "the {} block holds {!r} where a list of fields was " + "expected. Each element of a lists block is itself a " + "list, so that field lists can be combined.".format( + lists_label, field_list ) - field_set.add(alias) - - mfnl = max([len(field) for field in field_list_1d]) + 3 - mfnl = min(16, mfnl) # limit to 16 - w = lambda s: f'"{s}",' # wrap string in quotes and add comma - for field_name, alias, reduction in field_list_1d_seperated: - diag_table.write( - f'{w(module)} {w(field_name):{mfnl}}{w(alias):{mfnl}}{fname}, "all", ' - f'{w(reduction)} {w(file_block["regional_section"])} {packing}\n' ) - - diag_table.write("\n") + group.fields.extend(Field(spec) for spec in field_list) + stream.groups.append(group) + except DiagTableError as error: + raise DiagTableError( + "Cannot write {}: {}".format(stream.source, error) + ) from error + + return stream diff --git a/cime_config/MOM_RPS/FType_input_nml.py b/cime_config/MOM_RPS/FType_input_nml.py index ac457ec7..dd6e130f 100644 --- a/cime_config/MOM_RPS/FType_input_nml.py +++ b/cime_config/MOM_RPS/FType_input_nml.py @@ -1,4 +1,3 @@ -import os from CIME.ParamGen.paramgen import ParamGen diff --git a/cime_config/MOM_RPS/diag_table_streams.py b/cime_config/MOM_RPS/diag_table_streams.py new file mode 100644 index 00000000..a1352909 --- /dev/null +++ b/cime_config/MOM_RPS/diag_table_streams.py @@ -0,0 +1,487 @@ +"""Stream model and writer for the MOM6 diag_table. + +A *stream* is one MOM6 output file: one entry in the file list of the diag_table +(its Section-1) together with the fields written to that file (its Section-3). +Streams are built from param_templates/diag_table.yaml by FType_diag_table. + +The name of a stream is also the file name segment that follows the case name +and the component name, i.e., a stream named "h.native" is written to +${CASE}.mom6.h.native..nc. The part is not specified by hand: it is +derived from how often a new file is started (new_file_freq_units). +""" + +import re +from collections import OrderedDict + +# File name date templates +DATE_TEMPLATES = OrderedDict( + [ + ("years", "%4yr"), + ("months", "%4yr-%2mo"), + ("days", "%4yr-%2mo-%2dy"), + ("hours", "%4yr-%2mo-%2dy-%2hr"), + ("minutes", "%4yr-%2mo-%2dy-%2hr%2mi"), + ("seconds", "%4yr-%2mo-%2dy-%2hr%2mi%2sc"), + ] +) + +# Units accepted for output_freq and new_file_freq. +FREQ_UNITS = tuple(DATE_TEMPLATES) + +# MOM6 diagnostics modules, one per vertical coordinate of the output. +MODULES = ("ocean_model", "ocean_model_z", "ocean_model_rho2") + +# The case name is not known when the diag_table is written, so the file names +# carry this placeholder, which FType_diag_table.resolve substitutes later. +CASENAME = "${CASE}" + +# Reduction methods accepted by FMS. Within a group, the alternative spellings +# are equivalent (see init_output_field in FMS diag_util.F90). +REDUCTION_METHODS = ( + ".true.", + "mean", + "average", + "avg", + ".false.", + "none", + "point", + "rms", + "min", + "minimum", + "max", + "maximum", + "sum", + "cumsum", +) + +# Reduction methods that take a trailing sample count, e.g., diurnal8. +_COUNTED_REDUCTIONS = ("diurnal", "pow") + +# Number of bytes per output value: 1 is double precision, 2 single precision. +PACKING_VALUES = (1, 2) + +# Settings of a stream, i.e., of an entry in the file list of the diag_table. +# new_file_freq and new_file_freq_units are absent for streams that are written +# only once (output_freq < 0). +STREAM_SETTINGS = ( + "output_freq", + "output_freq_units", + "new_file_freq", + "new_file_freq_units", + "time_axis_units", + "reduction_method", + "regional_section", + "packing", +) + +# Settings that every stream must have a value for. The two of STREAM_SETTINGS +# that are missing here, new_file_freq and new_file_freq_units, are absent for a +# stream that is written only once, and so is never rolled over. +REQUIRED_SETTINGS = ( + "output_freq", + "output_freq_units", + "time_axis_units", + "reduction_method", + "regional_section", + "packing", +) + +_FIELD_NAME = re.compile(r"^\w+$") +_REGIONAL_SECTION = re.compile(r"^\s*(-?[\d.]+\s+){5}-?[\d.]+\s*$") + + +class DiagTableError(Exception): + """Raised when a diag_table cannot be generated as specified.""" + + +def is_valid_reduction(reduction): + """Returns True if reduction is a reduction method accepted by FMS. + + Example + ------- + >>> is_valid_reduction("mean") and is_valid_reduction("diurnal8") + True + >>> is_valid_reduction("median") or is_valid_reduction("pow") + False + """ + if reduction in REDUCTION_METHODS: + return True + return any( + re.match(r"^" + prefix + r"\d+$", str(reduction)) + for prefix in _COUNTED_REDUCTIONS + ) + + +def int_setting(value, setting, source): + """Returns the value of a whole number setting, or reports what is wrong. + + int() is not called on a setting directly, so that a missing or non-numeric + value is reported in terms of the stream it belongs to rather than escaping + as a TypeError or a ValueError. + + Parameters + ---------- + value: + The value to interpret, which may be missing, i.e., None. + setting: str + Name of the setting, for the error message. + source: str + Where the setting came from, for the error message. + + Example + ------- + >>> int_setting(2, "packing", 'stream "h.native"') + 2 + """ + if value is None: + raise DiagTableError( + "Cannot write {}: it has no value for {}.".format(source, setting) + ) + try: + return int(value) + except (TypeError, ValueError): + raise DiagTableError( + 'Cannot write {}: its {} is "{}", which is not a whole ' + "number.".format(source, setting, value) + ) from None + + +def is_valid_regional_section(section): + """Returns True if section is "none" (global) or six coordinate bounds. + + Example + ------- + >>> is_valid_regional_section("none") + True + >>> is_valid_regional_section("-5.75 19.0 78.93 78.93 -1 -1") + True + >>> is_valid_regional_section("-5.75 19.0") + False + """ + return str(section) == "none" or bool(_REGIONAL_SECTION.match(str(section))) + + +class Field: + """A single diagnostic field written to a stream. + + A field is given as "name[:output_name[:reduction]]", where output_name + defaults to name, and reduction defaults to the reduction method of the + stream that the field belongs to. + + Example + ------- + >>> field = Field("KPP_OBLdepth:oml_max:max") + >>> field.name, field.output_name, field.reduction + ('KPP_OBLdepth', 'oml_max', 'max') + >>> Field("tos").output_name, Field("tos").reduction + ('tos', None) + """ + + def __init__(self, spec): + self.spec = str(spec).strip() + parts = [part.strip() for part in self.spec.split(":")] + if not 1 <= len(parts) <= 3 or not all(parts): + raise DiagTableError( + 'Invalid field: "{}". Expected "name", "name:output_name", or ' + '"name:output_name:reduction".'.format(self.spec) + ) + for part in parts[:2]: + if not _FIELD_NAME.match(part): + raise DiagTableError( + 'Invalid field name "{}" in "{}". Field and output names may ' + "contain letters, digits and underscores only.".format( + part, self.spec + ) + ) + self.name = parts[0] + self.output_name = parts[1] if len(parts) > 1 else self.name + self.reduction = parts[2] if len(parts) > 2 else None + if self.reduction is not None and not is_valid_reduction(self.reduction): + raise DiagTableError( + 'Invalid reduction method "{}" in field "{}". Valid methods are: ' + "{}, diurnal, pow.".format( + self.reduction, self.spec, ", ".join(REDUCTION_METHODS) + ) + ) + + def __repr__(self): + return "Field({!r})".format(self.spec) + + +class FieldGroup: + """The fields that a stream receives from one MOM6 diagnostics module.""" + + def __init__(self, module, fields=None): + if module not in MODULES: + raise DiagTableError( + 'Invalid diagnostics module "{}". Valid modules are: {}.'.format( + module, ", ".join(MODULES) + ) + ) + self.module = module + self.fields = list(fields) if fields else [] + + def __repr__(self): + return "FieldGroup({!r}, {} fields)".format(self.module, len(self.fields)) + + +class Stream: + """One MOM6 output file, i.e., one entry in the diag_table file list. + + Attributes + ---------- + name: str + Name of the stream, e.g., "h.native". Also the file name segment that + follows the case name and the component name. + settings: dict + Stream settings, keyed by the names in STREAM_SETTINGS. + groups: list of FieldGroup + The fields written to this stream, grouped by diagnostics module. + labels: list of str + Labels of the diag_table template entries that this stream was built + from. Used in error messages only. + """ + + def __init__(self, name, settings=None, groups=None, labels=None): + self.name = name + self.settings = dict(settings) if settings else {} + self.groups = list(groups) if groups else [] + self.labels = list(labels) if labels else [] + + @property + def source(self): + """Describes this stream and its origin, for error messages.""" + return 'stream "{}" (from template entry {})'.format( + self.name, " and ".join(self.labels) or "(unknown)" + ) + + @property + def date_template(self): + """The date part of the file name, derived from new_file_freq_units.""" + if self.settings.get("new_file_freq") is None: + return "" + units = self.settings.get("new_file_freq_units") + if units not in DATE_TEMPLATES: + raise DiagTableError( + 'Cannot write {}: its new_file_freq_units is "{}", which is not ' + "one of: {}.".format(self.source, units, ", ".join(FREQ_UNITS)) + ) + return DATE_TEMPLATES[units] + + @property + def suffix(self): + """The full file name segment, i.e., the name plus the date template.""" + return self.name + self.date_template + + @property + def is_empty(self): + """Returns True if no field is written to this stream.""" + return not any(group.fields for group in self.groups) + + def output_names(self): + """Returns the output names of all fields written to this stream.""" + return [field.output_name for group in self.groups for field in group.fields] + + def group_for(self, module): + """Returns the field group of a module, creating it if necessary.""" + for group in self.groups: + if group.module == module: + return group + group = FieldGroup(module) + self.groups.append(group) + return group + + def merge(self, other): + """Merges the fields of another stream of the same name into this one. + + Two template entries may resolve to the same stream name, in which case + they describe the same output file and must agree on its settings. Fields + that both entries request are written only once. + """ + for setting in STREAM_SETTINGS: + mine, theirs = self.settings.get(setting), other.settings.get(setting) + if mine != theirs: + raise DiagTableError( + 'Template entries {} and {} both describe the stream "{}", ' + "but they disagree on {}: {!r} vs {!r}.".format( + " and ".join(self.labels), + " and ".join(other.labels), + self.name, + setting, + mine, + theirs, + ) + ) + for group in other.groups: + target = self.group_for(group.module) + existing = {field.output_name: field for field in target.fields} + for field in group.fields: + duplicate = existing.get(field.output_name) + if duplicate is None: + target.fields.append(field) + elif duplicate.spec != field.spec: + raise DiagTableError( + 'Template entries {} and {} both write "{}" to stream ' + '"{}", but they request it differently: "{}" vs ' + '"{}".'.format( + " and ".join(self.labels), + " and ".join(other.labels), + field.output_name, + self.name, + duplicate.spec, + field.spec, + ) + ) + self.labels.extend(other.labels) + + def validate(self): + """Checks that this stream can be written to the diag_table.""" + settings = self.settings + for setting in REQUIRED_SETTINGS: + if settings.get(setting) is None: + raise DiagTableError( + "Cannot write {}: it has no value for {}.".format( + self.source, setting + ) + ) + int_setting(settings["output_freq"], "output_freq", self.source) + if settings["output_freq_units"] not in FREQ_UNITS: + raise DiagTableError( + 'Cannot write {}: its output_freq_units is "{}", which is not one ' + "of: {}.".format( + self.source, settings["output_freq_units"], ", ".join(FREQ_UNITS) + ) + ) + if not is_valid_reduction(settings["reduction_method"]): + raise DiagTableError( + 'Cannot write {}: its reduction_method is "{}", which is not one ' + "of: {}, diurnal, pow.".format( + self.source, + settings["reduction_method"], + ", ".join(REDUCTION_METHODS), + ) + ) + if not is_valid_regional_section(settings["regional_section"]): + raise DiagTableError( + 'Cannot write {}: its regional_section is "{}", which is neither ' + '"none" nor six space separated bounds.'.format( + self.source, settings["regional_section"] + ) + ) + packing = int_setting(settings["packing"], "packing", self.source) + if packing not in PACKING_VALUES: + raise DiagTableError( + "Cannot write {}: its packing is {}, which is not one of: " + "{}.".format( + self.source, + settings["packing"], + ", ".join(str(value) for value in PACKING_VALUES), + ) + ) + self.date_template # raises if new_file_freq_units is invalid + + # A given output name may be written to a file only once. + seen = set() + for name in self.output_names(): + if name in seen: + raise DiagTableError( + 'Cannot write {}: the field "{}" is written to it more than ' + "once.".format(self.source, name) + ) + seen.add(name) + + def __repr__(self): + return "Stream({!r}, {} fields)".format(self.name, len(self.output_names())) + + +def write_diag_table(streams, output_path): + """Writes a diag_table for the given streams. + + Parameters + ---------- + streams: list of Stream + The streams of the case, in the order their files are to be listed. + Streams with no fields are not written out. + output_path: str + Path of the diag_table to write. + """ + + def filename(stream): + return '"{}.mom6.{}"'.format(CASENAME, stream.suffix) + + def quoted(value): + """Renders one column: the value in quotes, and the separating comma.""" + return '"{}",'.format(value) + + # The streams are gone over twice, once to validate them and once to write + # them out, so a one-shot iterable would yield a table with no files in it. + assert isinstance( + streams, (list, tuple) + ), "write_diag_table needs a list of streams, not a one-shot iterable." + + for stream in streams: + stream.validate() + written = [stream for stream in streams if not stream.is_empty] + + # Width of the file name column, including the quotes and the comma. Only + # the files that are actually written have a say in it. + name_width = max((len(filename(stream)) for stream in written), default=0) + 4 + + with open(output_path, "w") as diag_table: + diag_table.write( + '"MOM6 diagnostic fields table for CESM case: {}"\n'.format(CASENAME) + ) + diag_table.write("1 1 1 0 0 0\n") # TODO + diag_table.write("### Section-1: File List\n") + diag_table.write("#========================\n") + for stream in written: + settings = stream.settings + entry = '{fname:{width}s} {output_freq:3s} {output_freq_units:9s} 1, {time_axis_units:9s} "time"'.format( + width=name_width, + fname=filename(stream) + ",", + output_freq="{},".format(settings["output_freq"]), + output_freq_units='"{}",'.format(settings["output_freq_units"]), + time_axis_units='"{}",'.format(settings["time_axis_units"]), + ) + if settings.get("new_file_freq") is not None: + entry += ', {}, "{}"'.format( + settings["new_file_freq"], settings["new_file_freq_units"] + ) + diag_table.write(entry + "\n") + diag_table.write("\n") + + diag_table.write("### Section-2: Fields List\n") + diag_table.write("#=========================\n") + for stream in written: + settings = stream.settings + fname = filename(stream) + diag_table.write("# {}\n".format(fname)) + for group in stream.groups: + if not group.fields: + continue + # Width of the two field name columns + field_width = min( + 16, + max( + len(quoted(name)) + for field in group.fields + for name in (field.name, field.output_name) + ), + ) + for field in group.fields: + diag_table.write( + "{module} {name:{width}}{output_name:{width}}{fname}, " + '"all", {reduction} {regional_section} {packing}\n'.format( + module=quoted(group.module), + name=quoted(field.name), + output_name=quoted(field.output_name), + width=field_width, + fname=fname, + reduction=quoted( + field.reduction or settings["reduction_method"] + ), + regional_section=quoted(settings["regional_section"]), + packing=settings["packing"], + ) + ) + diag_table.write("\n") diff --git a/param_templates/diag_table.yaml b/param_templates/diag_table.yaml index 252aa2c2..e8b16a13 100644 --- a/param_templates/diag_table.yaml +++ b/param_templates/diag_table.yaml @@ -1,6 +1,6 @@ ############################################################################### # Section 1: (Reusable) List of Fields -# Field lists to be added in Files defined in Section 2. +# Field lists to be added to the files defined in Section 3. ############################################################################### --- FieldLists: @@ -82,62 +82,49 @@ FieldLists: - &cmip7_sfc ["tossq", "rsdo", "T_adx_2d", "T_ady_2d"] ############################################################################### -# Section 2: File lists: -# List of files to be added in diag_table +# Section 2: Defaults for all files, unless overridden in the file entry. +############################################################################### + +FileDefaults: + output_freq: 1 + output_freq_units: + $TEST == True: "days" + $OCN_DIAG_MODE == "spinup": "years" + else: "months" + time_axis_units: "days" + reduction_method: "mean" # time average + regional_section: "none" # global + packing: = 1 if $TEST else 2 # 1: double precision, 2: single precision + +############################################################################### +# Section 3: File list +# The files, i.e., the output streams, to be added to the diag_table. The +# label of each entry is the name of the stream: it ends up in the file name +# as ${CASE}.mom6.