Add a planner that aligns CDEPS stream years with the run dates

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"]

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 ``<file>`` 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 ``<file>`` 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

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_<comp>_streams``.
+
+Narrowing the file list matters as much as the alignment: the ``<datafiles>``
+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 ``<stream>:<key> = <value>`` line per setting to
+``user_nl_<comp>_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}