diff --git a/cdds/bin/update_grid_names b/cdds/bin/update_grid_names new file mode 100755 index 000000000..051997119 --- /dev/null +++ b/cdds/bin/update_grid_names @@ -0,0 +1,9 @@ +#!/usr/bin/env python3 +# (C) British Crown Copyright 2026, Met Office. +# Please see LICENSE.md for license details. +import sys + +from cdds.utils.grid_labels.command_line import main + +if __name__ == '__main__': + sys.exit(main()) diff --git a/cdds/cdds/utils/grid_labels/__init__.py b/cdds/cdds/utils/grid_labels/__init__.py new file mode 100644 index 000000000..066f7c978 --- /dev/null +++ b/cdds/cdds/utils/grid_labels/__init__.py @@ -0,0 +1,2 @@ +# (C) British Crown Copyright 2026, Met Office. +# Please see LICENSE.md for license details. diff --git a/cdds/cdds/utils/grid_labels/command_line.py b/cdds/cdds/utils/grid_labels/command_line.py new file mode 100644 index 000000000..3511d3e25 --- /dev/null +++ b/cdds/cdds/utils/grid_labels/command_line.py @@ -0,0 +1,39 @@ +# (C) British Crown Copyright 2026, Met Office. +# Please see LICENSE.md for license details. +import argparse + +from cdds.utils.grid_labels.main import map_variables_to_grid_names, write_grid_names_config +from cdds.utils.grid_labels.mappings import get_mip_convert_mappings +from cdds.utils.grid_labels.parsers import parse_mappings_json +from cdds.utils.grid_labels.stashmaster import stash_records + + +def main(): + """ + Main entry point for the update_grid_names command line tool. + """ + arguments = parse_args() + + plugins = ["ukcm2", "ukesm1p3"] + + records = stash_records(arguments.stashmaster) + + for plugin in plugins: + mappings = get_mip_convert_mappings(plugin) + ocean_grids = parse_mappings_json(arguments.mappings, plugin, "ocean") + seaice_grids = parse_mappings_json(arguments.mappings, plugin, "seaice") + grid_names = map_variables_to_grid_names(mappings, ocean_grids, seaice_grids, records, plugin) + write_grid_names_config(grid_names, f"grids_{plugin}.cfg") + + +def parse_args(): + """ + Parse command line arguments for the update_grid_names tool. + """ + + parser = argparse.ArgumentParser(description="Update grid names based on model parameters.") + + parser.add_argument('stashmaster', type=str, help='Path to a copy of STASHmaster',) + parser.add_argument('mappings', type=str, help='Path to the mappings.json from the CDDS-CMIP7-mappings repo') + + return parser.parse_args() diff --git a/cdds/cdds/utils/grid_labels/main.py b/cdds/cdds/utils/grid_labels/main.py new file mode 100644 index 000000000..d08fca0d4 --- /dev/null +++ b/cdds/cdds/utils/grid_labels/main.py @@ -0,0 +1,129 @@ +# (C) British Crown Copyright 2026, Met Office. +# Please see LICENSE.md for license details. +from collections import defaultdict + +from cdds.utils.grid_labels.mappings import Mapping +from cdds.utils.grid_labels.parsers import parse_icemod_grids, parse_ocean_grids +from cdds.utils.grid_labels.stashmaster import stash_records + +grid_name_to_grid_id = { + "latlon-native": {1, 2, 3, 4, 5, 26, 21, 17, 22}, + "latlon-uvgrid": {11, 12, 13, 14, 15}, + "latlon-ugrid": {18, 27}, + "latlon-vgrid": {19}, +} + +substream_to_grid_name = { + "grid-T": "tripolar-native", + "grid-U": "tripolar-ugrid", + "grid-V": "tripolar-vgrid", + "grid-W": "tripolar-native", + "diaptr": "tripolar-native", + "scalar": "tripolar-native", + # for ukesmp13 + "ptrc-T": "tripolar-native", + "diad-T": "tripolar-native", +} + + +grid_type = { + "latlon-native": "atmos", + "latlon-uvgrid": "atmos", + "latlon-ugrid": "atmos", + "latlon-vgrid": "atmos", + "tripolar-ugrid": "ocean", + "tripolar-vgrid": "ocean", + "tripolar-native": "ocean", + "seaice-native": "ocean", +} + + +ancils = { + "areacello_ti-u-hxy-u": "tripolar-native", + "basin_ti-u-hxy-u": "tripolar-native", + "deptho_ti-u-hxy-sea": "tripolar-native", + "dxto_ti-u-hxy-u": "tripolar-native", + "dxuo_ti-u-hxy-u": "tripolar-ugrid", + "dxvo_ti-u-hxy-u": "tripolar-vgrid", + "dyto_ti-u-hxy-u": "tripolar-native", + "dyuo_ti-u-hxy-u": "tripolar-ugrid", + "dyvo_ti-u-hxy-u": "tripolar-vgrid", + "sftof_ti-u-hxy-u": "tripolar-native", + "hfgeou_ti-u-hxy-sea": "tripolar-native", + "hfsnthermds_tavg-ol-hxy-sea": "tripolar-native", + "rsdo_tavg-ol-hxy-sea": "tripolar-native", +} + +seaice = { + # all seaice variables are assumed to be tripolar-native unless specified here + "sidmasstranx_tavg-u-hxy-u": "tripolar-ugrid", + "sidmasstrany_tavg-u-hxy-u": "tripolar-vgrid", + # these are seaice variable that will have their coordinates replaced by processor + "sistrxdtop_tavg-u-hxy-si": "tripolar-ugrid", + "sistrydtop_tavg-u-hxy-si": "tripolar-vgrid", + "sistrxubot_tavg-u-hxy-si": "tripolar-ugrid", + "sistryubot_tavg-u-hxy-si": "tripolar-vgrid", + "siu_tavg-u-hxy-si": "tripolar-ugrid", + "siv_tavg-u-hxy-si": "tripolar-vgrid", + "siforceintstrx_tavg-u-hxy-si": "tripolar-ugrid", + "siforceintstry_tavg-u-hxy-si": "tripolar-vgrid", + "siforcetiltx_tavg-u-hxy-si": "tripolar-ugrid", + "siforcetilty_tavg-u-hxy-si": "tripolar-vgrid", +} + + +def grid_ids_to_grid_name(ids: set[int]) -> str | None: + for label, label_ids in grid_name_to_grid_id.items(): + if ids.issubset(label_ids): + return label + return None + + +def stash_to_grid_name(mapping, records): + grid_ids = {int(records[code].Grid) for code in mapping.stash} + if grid_name := grid_ids_to_grid_name(grid_ids): + return grid_name + else: + raise ValueError( + f"Failed to find grid name for mapping: {mapping.name}, MIP Table: {mapping.mip_table}, Stash codes: {mapping.stash}, Grid ids: {grid_ids}" + ) + + +def map_variables_to_grid_names(mappings: dict[str, Mapping], ocean_grids, seaice_grids, records, plugin): + grid_names = defaultdict(dict) + + for variable, mapping in mappings.items(): + grid_name = None + + if mapping.stash: + grid_name = stash_to_grid_name(mapping, records) + elif variable in ocean_grids: + grid_name = substream_to_grid_name[ocean_grids[variable]] + elif variable in ancils: + grid_name = ancils[variable] + # ukcm2 seaice + elif plugin == "ukcm2": + if variable in seaice_grids and variable not in seaice: + grid_name = "tripolar-native" + if variable in seaice: + grid_name = seaice[variable] + # ukesm1p3 seaice + elif plugin == "ukesm1p3": + if variable in seaice_grids: + grid_name = "seaice-native" + + if not grid_name: + print(f"{plugin} Failed to identify a grid name for variable: {variable}, MIP Table: {mapping.mip_table}") + else: + grid_names[mapping.mip_table][variable] = (grid_type[grid_name], grid_name) + + return grid_names + + +def write_grid_names_config(grid_names: dict[str, dict[str, tuple[str, str]]], output_file: str) -> None: + with open(output_file, "w") as fh: + for mip_table, variables in grid_names.items(): + fh.write(f"[{mip_table}]\n") + for variable, (grid_type, grid_name) in variables.items(): + fh.write(f"{variable} = {grid_type} {grid_name}\n") + fh.write("\n") diff --git a/cdds/cdds/utils/grid_labels/mappings.py b/cdds/cdds/utils/grid_labels/mappings.py new file mode 100644 index 000000000..c48e654cf --- /dev/null +++ b/cdds/cdds/utils/grid_labels/mappings.py @@ -0,0 +1,44 @@ +# (C) British Crown Copyright 2026, Met Office. +# Please see LICENSE.md for license details. +import glob +import os +from typing import Any +from configparser import ConfigParser, ExtendedInterpolation +from dataclasses import dataclass + +from cdds.utils.grid_labels.stashmaster import extract_stash_codes + +from mip_convert import plugins + + +@dataclass +class Mapping: + name: str + expression: str + stash: list[str] + mip_table: str + + +def get_mip_convert_mappings(plugin: str) -> dict[str, Mapping]: + base_plugin_path = plugins.__file__ + glob_string = os.path.join(os.path.dirname(base_plugin_path), plugin, "data", '*mappings.cfg') + cfg_files = glob.glob(glob_string) + + mappings = {} + + for cfg_file in cfg_files: + mappings_config_object = ConfigParser(interpolation=ExtendedInterpolation()) + mappings_config_object.read(cfg_file) + for mapping_name, values in mappings_config_object.items(): + if mapping_name in ['DEFAULT', 'COMMON']: + continue + if expression := values.get('expression'): + stash_codes = extract_stash_codes(expression) + mappings[mapping_name] = Mapping( + name=mapping_name, + expression=expression, + stash=stash_codes, + mip_table=cfg_file.split("_")[-2]) + else: + raise ValueError(f"Missing expression for mapping: {mapping_name} in file: {cfg_file}") + return mappings diff --git a/cdds/cdds/utils/grid_labels/parsers.py b/cdds/cdds/utils/grid_labels/parsers.py new file mode 100644 index 000000000..ecec87d3a --- /dev/null +++ b/cdds/cdds/utils/grid_labels/parsers.py @@ -0,0 +1,57 @@ +# (C) British Crown Copyright 2026, Met Office. +# Please see LICENSE.md for license details. +import json +import re + + +def parse_ocean_grids(ocean_grids_file: str) -> dict[str, str]: + with open(ocean_grids_file, "r") as fh: + data = fh.readlines() + + ocean_grids = {} + for line in data: + variable, grid = line.strip().split(":") + ocean_grids[variable] = grid + + return ocean_grids + + +def parse_icemod_grids(icemod_grids_file: str) -> dict[str, str]: + with open(icemod_grids_file, "r") as fh: + data = fh.readlines() + + icemod_grids = {} + + regex = r"float (.*)\(.*(grid_\w)" + + for line in data: + match = re.search(regex, line) + if match: + variable = match.group(1) + grid = match.group(2) + icemod_grids[variable] = grid.replace("_", "-") + return icemod_grids + + +def parse_mappings_json(mappings_file: str, plugin: str, realm: str) -> dict[str, str]: + with open(mappings_file, "r") as fh: + mappings = json.load(fh) + + model_alias = {"ukcm2": "UKCM2", "ukesm1p3": "UKESM1-3"} + + mappings = [mapping for mapping in mappings if mapping["XIOS entries"]] + + grids = {} + + if realm == "ocean": + regex = re.compile(r"^o\w{2}/([\w-]*)") + elif realm == "seaice": + regex = re.compile(r"^(i\w{2})") + + for mapping in mappings: + if model_alias[plugin] in mapping["XIOS entries"]: + match = regex.match(mapping["XIOS entries"][model_alias[plugin]]) + if match: + grids[mapping["Data Request information"]["Branded variable name"]] = match.group(1) + + return grids diff --git a/cdds/cdds/utils/grid_labels/stashmaster.py b/cdds/cdds/utils/grid_labels/stashmaster.py new file mode 100644 index 000000000..ebfcda5d7 --- /dev/null +++ b/cdds/cdds/utils/grid_labels/stashmaster.py @@ -0,0 +1,117 @@ +# (C) British Crown Copyright 2026, Met Office. +# Please see LICENSE.md for license details. +import re +from configparser import ConfigParser +from dataclasses import dataclass + + +@dataclass +class StashMasterRecord: + Model: str + Sectn: str + Item: str + Name: str + Space: str + Point: str + Time: str + Grid: str + LevelT: str + LevelF: str + LevelL: str + PseudT: str + PseudF: str + PseudL: str + LevCom: str + Option_Codes: str + Version_Mask: str + Halo: str + DataT: str + DumpP: str + PC: str + Rotate: str + PPF: str + USER: str + LBVC: str + BLEV: str + TLEV: str + RBLEVV: str + CFLL: str + CFFF: str + + +def extract_stash_codes(expression): + regex = r"m\d*s\d*i\d*" + return re.findall(regex, expression) + + +def parse_stashmaster(stashmaster: str) -> list[str]: + """ + #|Model |Sectn | Item |Name | + #|Space |Point | Time | Grid |LevelT|LevelF|LevelL|PseudT|PseudF|PseudL|LevCom| + #| Option Codes | Version Mask | Halo | + #|DataT |DumpP | PC1 PC2 PC3 PC4 PC5 PC6 PC7 PC8 PC9 PCA | + #|Rotate| PPF | USER | LBVC | BLEV | TLEV |RBLEVV| CFLL | CFFF | + """ + + regex = r"^1\|(.*)\n^2\|(.*)\n^3\|(.*)\n^4\|(.*)\n^5\|(.*)" + + with open(stashmaster, "r") as fh: + data = fh.read() + + result = re.findall(regex, data, re.MULTILINE) + processed_result = [] + + for x in result: + x = "".join(x).split("|") + x = [x.strip() for x in x] + processed_result.append(x) + + return processed_result + + +def parse_stashmaster_meta(section): + section_name = f"{section}=" + + config = ConfigParser() + config.read("STASHmaster-meta.conf") + grids = {} + for section in config.sections(): + if section_name in section: + print(section) + print(config.get(section, "help")) + grids[section.split(section_name)[1]] = config.get(section, "help") + return grids + + +def parse_stashmaster_meta_stash(section: str): + config = ConfigParser() + config.read("STASHmaster-meta.conf") + grids = {} + + for section in config.sections(): + if "grid=" in section: + print(section) + print(config.get(section, "help")) + grids[section.split("grid=")[1]] = config.get(section, "help") + return grids + + +def to_formatted_stash_code(record: StashMasterRecord) -> str: + stash = "".join([ + f"m{int(record.Model):02d}", + f"s{int(record.Sectn):02d}", + f"i{int(record.Item):03d}", + ]) + return stash + + +def stash_records(stashmaster_file) -> dict[str, StashMasterRecord]: + stash_records = {} + for x in parse_stashmaster(stashmaster_file): + if len(x) == 31: + stash_record = StashMasterRecord(*x[:-1]) + stash_code = to_formatted_stash_code(stash_record) + stash_records[stash_code] = stash_record + else: + print(f"Unexpected number of fields: {len(x)} in record: {x[3]}") + return stash_records