Finding the data: CESM2 LENS on AWS

import importlib.metadata
import os
import re
import sys
from datetime import datetime

import intake
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import xarray as xr
from lensml import (
    FRONT_RANGE,
    NINO34,
    PARAMS_RECORD,
    SEASON,
    credit_own,
    derive_member_winters,
    load_cache,
    show,
    show_fig,
)

now = datetime.now().astimezone().isoformat(timespec="seconds")
print(f"run started {now}   kernel python {sys.version.split()[0]}")
PACKAGES = [
    "numpy",
    "pandas",
    "scipy",
    "scikit-learn",
    "matplotlib",
    "xarray",
    "zarr",
    "s3fs",
    "cftime",
    "intake",
    "intake-esm",
]
for pkg in PACKAGES:
    try:
        print(f"{pkg:13s} {importlib.metadata.version(pkg)}")
    except importlib.metadata.PackageNotFoundError:
        print(f"{pkg:13s} NOT INSTALLED: re-create the environment")

def span(box):
    """A box as degrees with hemispheres, from its slices on the 0-360 grid."""
    y0, y1 = box["lat"].start, box["lat"].stop
    x0, x1 = box["lon"].start, box["lon"].stop

    def lat(y):
        return f"{abs(y):g}°{'S' if y < 0 else 'N'}"

    def lon(x):
        return f"{x:g}°E" if x <= 180 else f"{360 - x:g}°W"

    return f"{lat(y0)}–{lat(y1)}, {lon(x0)}–{lon(x1)}"


fig, ax = plt.subplots(figsize=(9, 3.8))
ax.set(xlim=(165, 275), ylim=(-24, 56), ylabel="latitude")
ax.set_title(
    "The question's two boxes: the autumn (SON) Niño 3.4 state as the input, "
    f"{SEASON} Front Range precipitation as the target",
    fontsize=9,
)
xt = list(range(170, 276, 15))
ax.set_xticks(xt)
ax.set_xticklabels(
    [f"{x}°E" if x <= 180 else f"{360 - x}°W" for x in xt], fontsize=8
)
yt = list(range(-20, 61, 20))
ax.set_yticks(yt)
ax.set_yticklabels(
    [f"{abs(y)}°{'S' if y < 0 else 'N'}" for y in yt], fontsize=8
)
ax.axhline(0, lw=0.8, c="#6A7077", ls=":")
ax.grid(True, lw=0.5)
boxes = [
    (
        "NINO34",
        NINO34,
        "input: SON mean of TS (the standard Niño 3.4 box)",
        (215, -13),
        "top",
    ),
    (
        "FRONT_RANGE",
        FRONT_RANGE,
        f"target: {SEASON} mean of PRECC + PRECL (Front Range box around Denver: "
        "west of the Divide to the plains)\nselection bounds; the 2 × 3 grid cells averaged reach about "
        "38.6–40.5°N, 106.9–103.1°W",
        (236, 50),
        "bottom",
    ),
]
for name, box, note, at, va in boxes:
    x0, x1 = box["lon"].start, box["lon"].stop
    y0, y1 = box["lat"].start, box["lat"].stop
    ax.add_patch(
        plt.Rectangle(
            (x0, y0),
            x1 - x0,
            y1 - y0,
            facecolor="#C9A23A",
            edgecolor="#8F7022",
            lw=1.2,
            zorder=3,
        )
    )
    ax.annotate(
        f"{name}  {span(box)}\n{note}",
        xy=((x0 + x1) / 2, y0 if va == "top" else y1),
        xytext=at,
        fontsize=7.5,
        ha="center",
        va=va,
        arrowprops=dict(arrowstyle="-", color="#6A7077", lw=0.8),
    )
ax.annotate(
    "",
    xy=(FRONT_RANGE["lon"].start, FRONT_RANGE["lat"].start),
    xytext=(NINO34["lon"].stop, NINO34["lat"].stop),
    arrowprops=dict(
        arrowstyle="->",
        color="#8F7022",
        lw=1.6,
        connectionstyle="arc3,rad=-0.25",
    ),
)
ax.text(
    243,
    22,
    f"SON of year y − 1\n→ {SEASON} winter y",
    fontsize=8,
    ha="right",
    va="center",
    color="#8F7022",
)
credit_own(fig, "a diagram of the parameter boxes; no CESM2 LENS data drawn")
plt.tight_layout()
show_fig(
    fig,
    (
        "Schematic map on a longitude–latitude frame with no coastlines: "
        "a wide gold rectangle on the equator marks the Niño 3.4 box "
        "(5°S–5°N, 170°–120°W) over which the autumn input is averaged; a"
        " small gold rectangle near 39.5°N, 105°W marks the Front Range "
        "box over which the winter precipitation target is averaged; a "
        "curved arrow from the first to the second is labelled with the "
        "season step from the autumn of one year to the following winter."
    ),
)

CATALOG_URL = "https://ncar-cesm2-lens.s3-us-west-2.amazonaws.com/catalogs/aws-cesm2-le.json"
try:
    cat = intake.open_esm_datastore(CATALOG_URL)
    subset = cat.search(
        component="atm",
        frequency="monthly",
        experiment="historical",
        variable=["PRECC", "PRECL", "TS"],
    )
    print(
        f"catalogue last_updated: {cat.esmcat.last_updated}   "
        f"rows in catalogue: {len(cat.df)}   rows selected: {len(subset.df)}"
    )
    cols = [
        "variable",
        "forcing_variant",
        "units",
        "start_time",
        "end_time",
        "path",
    ]
    with pd.option_context(
        "display.max_colwidth", None
    ):  # the six paths in full
        display(subset.df[cols])
except Exception as err:  # no network, a DNS failure, S3 unreachable
    if os.environ.get("CI"):  # CI has network: a failure there is real
        raise
    cat = subset = None
    print(
        f"OFFLINE: the catalogue at {CATALOG_URL} could not be read "
        f"({type(err).__name__}). The Metadata Document's catalogue and store "
        "rows are *not printed this run*; every later cell runs from the cache."
    )

try:
    probe = xr.open_zarr(
        "s3://ncar-cesm2-lens/atm/monthly/cesm2LE-historical-cmip6-TS.zarr",
        consolidated=True,
        storage_options={"anon": True},
    )
except Exception as err:  # offline: the structure rows are not printed
    if os.environ.get("CI"):  # CI has network: a failure there is real
        raise
    probe = None
    print(
        f"OFFLINE: the TS store could not be opened ({type(err).__name__}); "
        "the dims, calendar and chunking rows are *not printed this run*."
    )
if probe is not None:
    print("dims      :", dict(probe.sizes))
    print("coords    :", list(probe.coords))
    print(
        "members   :",
        probe.member_id.values[:3],
        "…",
        "calendar:",
        type(probe.time.values[0]).__name__,
    )
    ch = tuple(int(c) for c in probe["TS"].encoding["chunks"])
    per_chunk = (
        int(np.prod(ch)) * probe["TS"].dtype.itemsize
    )  # bytes, uncompressed
    n_time_chunks = -(
        -probe.sizes["time"] // ch[1]
    )  # chunks along time per member
    per_member_var = n_time_chunks * per_chunk
    n_members = (
        2 * probe.member_id.size
    )  # two halves, each the size of this store
    print(
        f"chunking  : {ch} {probe['TS'].dtype} → {per_chunk / 1e6:.0f} MB per chunk "
        f"uncompressed; one member-variable ({probe.sizes['time']} months) = "
        f"{n_time_chunks} chunks × {per_chunk / 1e6:.0f} MB ≈ {per_member_var / 1e6:.0f} "
        f"MB read for any box subset; the derivation ({n_members} members × 3 "
        f"variables) ≈ {n_members * 3 * per_member_var / 1e9:.0f} GB"
    )

show(derive_member_winters, load_cache)

REBUILD = False  # True re-derives data/member_winters.csv from S3
N_MEMBERS_PER_HALF = (
    50  # used only when REBUILD = True; 50 = the full ensemble
)

table, cache_header, access_date = load_cache(REBUILD, N_MEMBERS_PER_HALF)
print(cache_header)
print(
    "cache header matches lensml.py:", PARAMS_RECORD, f"· access={access_date}"
)
n_header = 2 * int(
    re.search(r"n_members_per_half=(\d+)", cache_header).group(1)
)
n_store = (
    2 * probe.member_id.size if probe is not None else None
)  # unknown offline
if n_store:
    kind = (
        "A SAMPLE"
        if table.member_id.nunique() < n_store
        else "the full ensemble"
    )
    of_store = f"of {n_store} in the two stores ({kind})"
else:
    of_store = (
        f"(offline: the header says it was derived from {n_header} members)"
    )
print(
    f"{len(table)} member-winters · {table.member_id.nunique()} members {of_store} "
    f"· winters {table.winter.min()}–{table.winter.max()}"
)
table.head()