From cb4768ac920ead3224b34c6f9a3c44c2642bbbbf Mon Sep 17 00:00:00 2001 From: manishvenu Date: Fri, 23 Oct 2026 15:08:55 -0600 Subject: [PATCH] Add plan_stream_years to align DATM/DROF stream years with a run Multi-year data streams default to cycling their full available range. plan_stream_years reads the stream definitions for a component's mode from the CESM tree and returns, per stream, the year_first/year_last/year_align (and optionally the narrowed data file list) that cover a run's years with padding, skipping single-year or non-covering streams. It writes nothing; callers apply the plan to user_nl__streams. --- tests/2_integration/test_stream_years.py | 67 +++++++++++ visualCaseGen/cdeps_streams.py | 136 +++++++++++++++++++++++ visualCaseGen/stream_years.py | 96 ++++++++++++++++ 3 files changed, 299 insertions(+) create mode 100644 tests/2_integration/test_stream_years.py create mode 100644 visualCaseGen/cdeps_streams.py create mode 100644 visualCaseGen/stream_years.py diff --git a/tests/2_integration/test_stream_years.py b/tests/2_integration/test_stream_years.py new file mode 100644 index 0000000..a4d7ed3 --- /dev/null +++ b/tests/2_integration/test_stream_years.py @@ -0,0 +1,67 @@ +"""Tests for planning CDEPS stream years (visualCaseGen.stream_years). + +These read the real CDEPS stream definitions from the CESM tree that the other +integration tests use; no inputdata is needed. The exact file paths asserted +below double as a guard on the private CDEPS helpers that cdeps_streams calls, +so a CESM update that changes them fails here rather than silently producing +paths that do not exist. +""" + +import pytest + +from visualCaseGen.cime_interface import CIME_interface +from visualCaseGen.stream_years import plan_stream_years + + +class FakeCase: + """Stands in for a CIME Case: plan_stream_years only calls get_value.""" + + def get_value(self, name): + return {"DIN_LOC_ROOT": "/inputdata"}.get(name) + + +def test_multi_year_streams_are_aligned_around_the_run(): + """A 2021 run keeps 2019-2022 (a lead-in year, plus one year of padding on + each side), aligns the data year with the model year, and narrows each + stream's files to those years, each year carrying its own datestamp.""" + plan = plan_stream_years( + CIME_interface().srcroot, + "datm", + "CORE_IAF_JRA_1p5_2023", + FakeCase(), + 2021, + 2021, + ) + + assert plan["status"] == "configured" + assert len(plan["streams"]) == 8 + for settings in plan["streams"].values(): + assert settings["year_first"] == settings["year_align"] == 2019 + assert settings["year_last"] == 2022 + assert plan["streams"]["CORE_IAF_JRA_1p5_2023.GCGCS.PREC"]["datafiles"] == [ + f"/inputdata/ocn/jra55/v1.5_noleap/JRA.v1.5.prec.TL319.{year}.{stamp}.nc" + for year, stamp in [ + (2019, 210504), + (2020, 210504), + (2021, 220505), + (2022, 230718), + ] + ] + + +@pytest.mark.parametrize( + "mode, year, reason", + [ + ("CORE2_NYF", 2000, "single-year climatology"), + ("CORE_IAF_JRA", 2020, "fall outside the 1958-2016"), + ], +) +def test_unalignable_streams_are_skipped(mode, year, reason): + """A repeating single year has nothing to align, and a run outside the + stream's years cannot be aligned with real data. Both are reported instead + of planned, so the case keeps its CDEPS defaults.""" + plan = plan_stream_years( + CIME_interface().srcroot, "datm", mode, FakeCase(), year, year + ) + assert plan["status"] == "skipped" + assert reason in plan["reason"] diff --git a/visualCaseGen/cdeps_streams.py b/visualCaseGen/cdeps_streams.py new file mode 100644 index 0000000..d017b7c --- /dev/null +++ b/visualCaseGen/cdeps_streams.py @@ -0,0 +1,136 @@ +"""Read CDEPS stream definitions through CIME's own parsers. + +These are thin wrappers over ``StreamCDEPS`` and ``NamelistDefinition`` rather +than XML parsing of our own, so that the stream list and the expanded file +paths stay identical to what the CDEPS ``buildnml`` would generate. That +matters for streams whose files carry more than one datestamp: JRA v1.5 splits +1958-2023 across four ```` entries stamped 210504/220505/230718/240531, +and a hand-rolled expansion would silently emit paths that do not exist. + +``coverage`` and ``datafiles`` call private ``StreamCDEPS`` helpers +(``_get_stream_first_and_last_dates``, ``_resolve_values``, ``_sub_paths``). +Every such call is kept in this module, so a CDEPS change that breaks them +needs fixing in one place. +""" + +import os +import sys + + +def _cdeps_root(cesmroot): + return os.path.join(str(cesmroot), "components", "cdeps") + + +def _component_config(cesmroot, component): + return os.path.join(_cdeps_root(cesmroot), component, "cime_config") + + +def _add_to_path(directory): + if directory not in sys.path: + sys.path.append(directory) + + +def _ensure_cime(cesmroot): + """Put CIME on sys.path. + + CIME_interface normally does this already, but doing it here too lets + this module be used (and tested) without first constructing a + CIME_interface. + """ + _add_to_path(os.path.join(str(cesmroot), "cime")) + + +def _streams(cesmroot, component): + """A StreamCDEPS object for ``component`` ("datm" or "drof"). + + Rebuilt on every call rather than memoized here: CIME's ``GenericXML`` + keeps its own cache of parsed trees, so repeated construction costs + well under a millisecond. + """ + _ensure_cime(cesmroot) + shared_config = os.path.join(_cdeps_root(cesmroot), "cime_config") + # stream_cdeps.py is a loose module in the CDEPS config dir, not a package. + _add_to_path(shared_config) + from stream_cdeps import StreamCDEPS + + return StreamCDEPS( + os.path.join( + _component_config(cesmroot, component), + f"stream_definition_{component}.xml", + ), + os.path.join(shared_config, "stream_definition_v2.0.xsd"), + ) + + +def stream_names(cesmroot, component, mode): + """The stream names active for a DATM_MODE/DROF_MODE, in buildnml's order. + + The mode-to-streams mapping is matched by regex on the mode name, so this + cannot be replaced by prefixing: CORE_IAF_JRA_1p5_2023 yields streams named + CORE_IAF_JRA_1p5_2023.GCGCS.PREC, and the DROF mode IAF_JRA yields + rof.iaf_jra. + """ + _ensure_cime(cesmroot) + from CIME.XML.namelist_definition import NamelistDefinition + + nmldef = NamelistDefinition( + os.path.join( + _component_config(cesmroot, component), + f"namelist_definition_{component}.xml", + ) + ) + nmldef.set_nodes() + return nmldef.get_default_value( + "streamslist", attribute={f"{component}_mode": mode} + ) + + +def _stream_entry(cesmroot, component, stream_name): + streams = _streams(cesmroot, component) + for node in streams.get_children("stream_entry"): + if streams.get(node, "name") == stream_name: + return streams, node + raise KeyError(f"No stream_entry named {stream_name} in {component}") + + +def coverage(cesmroot, component, stream_name, cime_case): + """The (year_first, year_last) the stream definition itself declares.""" + streams, node = _stream_entry(cesmroot, component, stream_name) + return streams._get_stream_first_and_last_dates(node, cime_case) + + +def datafiles(cesmroot, component, stream_name, year_first, year_last, cime_case): + """Resolved file paths covering [year_first, year_last]. + + Mirrors the ``stream_datafiles`` branch of ``StreamCDEPS.create_stream_xml``: + each ```` contributes only the years its own first_year/last_year + attributes cover, so multi-datestamp streams resolve correctly. + """ + streams, node = _stream_entry(cesmroot, component, stream_name) + paths = [] + for child in streams.get_children("stream_datafiles", root=node): + for entry in streams.get_children(root=child): + template = streams._resolve_values(cime_case, entry.xml_element.text) + attrib = entry.xml_element.attrib + if "first_year" not in attrib or "last_year" not in attrib: + paths.append(template.strip()) + continue + first = max( + int(streams._resolve_values(cime_case, attrib["first_year"])), + year_first, + ) + last = min( + int(streams._resolve_values(cime_case, attrib["last_year"])), + year_last, + ) + if first > last: + continue + expanded = streams._sub_paths( + stream_name, + template, + first, + last, + int(attrib.get("filename_advance_days", 0)), + ) + paths += [line for line in expanded.strip().split("\n") if line] + return paths diff --git a/visualCaseGen/stream_years.py b/visualCaseGen/stream_years.py new file mode 100644 index 0000000..0429c68 --- /dev/null +++ b/visualCaseGen/stream_years.py @@ -0,0 +1,96 @@ +"""Plan the stream year settings that align CDEPS data streams with a run. + +CDEPS hardcodes ``year_align=1`` and the full year span for every JRA stream in +``stream_definition_{datm,drof}.xml``, and none of the JRA, NYF or ERA5 stream +entries reads the ``DATM_YR_ALIGN``/``_START``/``_END`` xml variables (only +the land-forcing and CPLHIST streams do). So a JRA case silently reads an +arbitrary data year, and ``xmlchange`` cannot fix it. The only lever CDEPS +offers is ``user_nl__streams``. + +Narrowing the file list matters as much as the alignment: the ```` +list is expanded from the stream definition's own year tags *before* user mods +are merged, so retagging the years alone still leaves all 59 JRA files +enumerated, and CDEPS reads the time axis of each one in turn until it reaches +the year it needs. + +``plan_stream_years`` only computes the settings; it writes nothing. To apply a +"configured" plan, write one ``: = `` line per setting to +``user_nl__streams``, joining ``datafiles`` with commas. +""" + +from visualCaseGen import cdeps_streams + + +def plan_stream_years( + cesmroot, component, mode, case, start_year, end_year, *, padding=1, narrow=True +): + """Work out the stream year settings for one component, or why there are none. + + Parameters + ---------- + cesmroot : str or Path + The CESM root directory (the one holding cime/ and components/cdeps/). + component : str + The CDEPS data component, e.g., "datm" or "drof". + mode : str + The component's stream set, i.e., the case's DATM_MODE or DROF_MODE. + case : CIME Case or similar + Only its get_value method is used, to expand $-variables such as + $DIN_LOC_ROOT in the stream definitions. + start_year, end_year : int + The first and last model years of the run. + padding : int, optional + Extra years of forcing to keep on either side of the run. + narrow : bool, optional + If True, also restrict each stream's datafiles to the years kept. + + Returns + ------- + dict + ``{"status": "configured", "mode": mode, "streams": {name: settings}}``, + where settings holds year_first, year_last, year_align and, if narrow, + datafiles (a list of paths). Otherwise ``{"status": "skipped", + "reason": str}``: a component that cannot be aligned is reported + rather than raised, so the caller can leave it at its CDEPS defaults. + """ + names = cdeps_streams.stream_names(cesmroot, component, mode) + if not names: + return {"status": "skipped", "reason": f"no streams defined for {mode}"} + + streams = {} + for name in names: + avail_first, avail_last = cdeps_streams.coverage( + cesmroot, component, name, case + ) + # Single-year streams are climatologies (CORE2_NYF is 1/1/1, the + # CORE_RYF*_JRA repeat years are e.g. 1961/1961). There is no + # alignment to get right and no files to narrow. + if avail_first == avail_last: + return { + "status": "skipped", + "reason": f"{mode} is a single-year climatology", + } + if start_year < avail_first or end_year > avail_last: + return { + "status": "skipped", + "reason": ( + f"run years {start_year}-{end_year} fall outside the " + f"{avail_first}-{avail_last} covered by {mode}" + ), + } + + # year_align == year_first makes the data year equal the model year + # for every year in range: dYear = yrFirst + modulo(mYear - yrFirst, n). + # The extra leading year is needed to interpolate at the run start. + year_first = max(avail_first, start_year - 1 - padding) + year_last = min(avail_last, end_year + padding) + streams[name] = { + "year_first": year_first, + "year_last": year_last, + "year_align": year_first, + } + if narrow: + streams[name]["datafiles"] = cdeps_streams.datafiles( + cesmroot, component, name, year_first, year_last, case + ) + return {"status": "configured", "mode": mode, "streams": streams}