JJA heat-index calculation

!pip install heatindex

import matplotlib.pyplot as plt
import numpy as np
import os
import xarray as xr
import intake
import intake_esm
import pandas as pd
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import dask 
from dask_jobqueue import PBSCluster
from dask.distributed import Client

import zarr
zarr.config.set({"async.concurrency": 4})   # cap simultaneous range requests; drop to 2 if it still trips

# Set up your scratch folder path - Ignore these lines if you are not on NCAR's cluster
username       = os.environ["USER"]
glade_scratch  = "/glade/derecho/scratch/" + username
print(glade_scratch)

# OSDF based data access 
era5_catalog  = 'https://data.gdex.ucar.edu/d633000/catalogs/d633000-osdf.json'
jra3q_catalog = 'https://data.gdex.ucar.edu/d640000/catalogs/d640000-osdf.json' 
# era5_catalog  = 'https://data.gdex.ucar.edu/d633000/catalogs/d633000-https.json'
# jra3q_catalog = 'https://data.gdex.ucar.edu/d640000/catalogs/d640000-https.json'

USE_PBS_SCHEDULER = True # Use NCAR's HPC cluster 
USE_DASK_GATEWAY  = False 

# Create a PBS cluster object
def get_pbs_cluster():
    """ Create cluster through dask_jobqueue.   
    """
    from dask_jobqueue import PBSCluster
    cluster = PBSCluster(
        job_name = 'dask-osdf-24',
        cores = 1,
        memory = '8GiB',
        processes = 1,
        local_directory = glade_scratch + '/dask/spill',
        log_directory = glade_scratch + '/dask/logs/',
        resource_spec = 'select=1:ncpus=1:mem=8GB',
        queue = 'casper',
        walltime = '3:00:00',
        #interface = 'ib0'
        interface = 'ext'
    )
    return cluster

def get_gateway_cluster():
    """ Create cluster through dask_gateway
    """
    from dask_gateway import Gateway

    gateway = Gateway()
    cluster = gateway.new_cluster()
    cluster.adapt(minimum=2, maximum=4)
    return cluster

def get_local_cluster():
    """ Create cluster using the Jupyter server's resources
    """
    from distributed import LocalCluster, performance_report
    cluster = LocalCluster()    

    cluster.scale(5)
    return cluster

# Obtain dask cluster in one of three ways
if USE_PBS_SCHEDULER:
    cluster = get_pbs_cluster()
elif USE_DASK_GATEWAY:
    cluster = get_gateway_cluster()
else:
    cluster = get_local_cluster()

# Connect to cluster
from distributed import Client
client = Client(cluster)

# Scale the cluster and display cluster dashboard URL
n_workers =3
cluster.scale(n_workers)
client.wait_for_workers(n_workers = n_workers)
cluster

%%time
era5_cat = intake.open_esm_datastore(era5_catalog)
jra3q_cat = intake.open_esm_datastore(jra3q_catalog)
era5_cat

# ── Time range ────────────────────────────────────────────────────────────────
YEAR_START_ERA5  = 1990
YEAR_START_JRA3Q = 1990
YEAR_END         = 2005

JJA_MONTHS = [6, 7, 8]

# ── Site coordinates (degrees north, degrees east 0–360) ──────────────────────
SITES = {
    "Boulder_CO" : {"lat":  40.01, "lon": 254.73},
    "Chicago_IL" : {"lat":  41.88, "lon": 272.37},
}

# ── Reproducibility ───────────────────────────────────────────────────────────
RANDOM_SEED = 42

jra3q_cat.df[['variable', 'short_name', 'long_name']].drop_duplicates().loc[
    jra3q_cat.df['long_name'].str.contains('2m temp|humid', case=False, na=False)]

# ERA5: 2-m temperature and 2-m dew-point (GDEX convention)
ERA5_T_VAR  = "VAR_2T"
ERA5_TD_VAR = "VAR_2D"

era5_search = era5_cat.search(variable=[ERA5_T_VAR, ERA5_TD_VAR])
print(f"ERA5 matched: {len(era5_search.df):,} entries")

# JRA-3Q: 2-m temperature and 2-m relative humidity
# Variable names confirmed from the discovery step in the previous cell:
#   tmp2m-hgt-an-ll125  →  2-metre temperature, height level, analysis
#   rh2m-hgt-an-ll125   →  2-metre relative humidity, height level, analysis
JRA3Q_T_VAR  = "tmp2m-hgt-an-ll125"
JRA3Q_RH_VAR = "rh2m-hgt-an-ll125"

jra3q_search = jra3q_cat.search(variable=[JRA3Q_T_VAR, JRA3Q_RH_VAR])
print(f"JRA-3Q matched: {len(jra3q_search.df):,} entries")

#Use paths to the load the dataset
era5_t_path  = era5_search.df.loc[era5_search.df['variable'] == ERA5_T_VAR,  'path'].item()
era5_td_path = era5_search.df.loc[era5_search.df['variable'] == ERA5_TD_VAR, 'path'].item()

jra3q_t_path  = jra3q_search.df.loc[jra3q_search.df['variable'] == JRA3Q_T_VAR,  'path'].item()
jra3q_rh_path = jra3q_search.df.loc[jra3q_search.df['variable'] == JRA3Q_RH_VAR, 'path'].item()

print(era5_t_path)
print(era5_td_path)
print(jra3q_t_path)
print(jra3q_rh_path)

%%time
era5_2t = xr.open_zarr("osdf:///ncar/gdex/d633000/e5.oper.an.sfc.zarr/e5.oper.an.sfc.2t.zarr")
era5_2d = xr.open_zarr("osdf:///ncar/gdex/d633000/e5.oper.an.sfc.zarr/e5.oper.an.sfc.2d.zarr")
era5_dsets = xr.merge([era5_2t[[ERA5_T_VAR]], era5_2d[[ERA5_TD_VAR]]])   # VAR_2T, VAR_2D only
era5_dsets

%%time
# Read JRA-3Q over OSDF / PelicanFS, but DIRECT FROM ORIGIN (direct_reads=True),
# bypassing the OSDF cache layer. A cache node in the rotation serves corrupted
# bytes for some d640000 chunks — deterministic zlib "incorrect header check"
# from Casper — while the origin copy is clean (verified: the same chunks decode
# over direct reads and over plain HTTPS). direct_reads keeps the osdf:// /
# PelicanFS streaming path while routing around the faulty cache.
jra3q_dsets = xr.open_dataset("https://data.gdex.ucar.edu/d640000/kerchunk/anl_surf125-remote-osdf.parq",
    engine="kerchunk",storage_options={"remote_protocol": "osdf", "remote_options": {"direct_reads": True}},)[[JRA3Q_T_VAR, JRA3Q_RH_VAR]]
jra3q_dsets

import dask

def load_jra_point(ds, lat, lon, retries=8):
    """Read a JRA-3Q point time series as TRUE in-memory NumPy.

    The kerchunk reference exposes the archive as a per-timestep-chunked
    (dask) store. On this stack `.load()`/`.compute()` return *lazy* dask
    arrays that re-read the store later — and with a distributed Client active
    those reads run on the workers, where the notebook-process
    async.concurrency throttle does NOT apply, so the OSDF cache truncates
    responses under the fan-out. To read it reliably we:
      * apply the point + JJA selection, then
      * force the read IN-PROCESS via the synchronous scheduler (so the zarr
        async.concurrency cap actually limits simultaneous range requests), and
      * pull `.values` (true NumPy) rather than `.load()` (lazy dask),
    retrying the transient OSDF truncation. Returns a small, fully in-memory
    Dataset that needs no further store access.
    """
    sub = (ds.sel(time=slice(str(YEAR_START_JRA3Q), str(YEAR_END)))
              .sel(lat=lat, lon=lon, method="nearest"))
    sub = sub.sel(time=sub.time.dt.month.isin(JJA_MONTHS))     # JJA only -> fewer reads
    for attempt in range(1, retries + 1):
        try:
            with dask.config.set(scheduler="synchronous"):
                return xr.Dataset(
                    {v: ("time", np.asarray(sub[v].values))
                     for v in (JRA3Q_T_VAR, JRA3Q_RH_VAR)},
                    coords={"time": sub["time"].values,
                            "lat": float(sub["lat"]), "lon": float(sub["lon"])},
                )
        except Exception as exc:
            if attempt == retries:
                raise
            print(f"    JRA-3Q read attempt {attempt}/{retries} failed "
                  f"({type(exc).__name__}); retrying ...")


# Subset to analysis period and JJA
era5_points  = {}
jra3q_points = {}

for site, coords in SITES.items():
    lat, lon = coords["lat"], coords["lon"]

    # ERA5: Zarr-backed and reliable; keep it lazy (JJA filter + point), compute later
    era5_pt = era5_dsets.sel(time=slice(str(YEAR_START_ERA5), str(YEAR_END)))
    era5_pt = era5_pt.sel(time=era5_pt.time.dt.month.isin(JJA_MONTHS))
    era5_pt = era5_pt.sel(latitude=lat, longitude=lon, method="nearest")

    # JRA-3Q: read the JJA point series into NumPy once (in-process, throttled, retried)
    print(f"Loading JRA-3Q / {site} into memory ...")
    jra3q_pt = load_jra_point(jra3q_dsets, lat, lon)

    print(f"ERA5   / {site}: requested ({lat:.2f}°N, {lon:.2f}°E) → "
          f"nearest ({float(era5_pt.latitude):.3f}°N, {float(era5_pt.longitude):.3f}°E)")
    print(f"JRA-3Q / {site}: requested ({lat:.2f}°N, {lon:.2f}°E) → "
          f"nearest ({float(jra3q_pt.lat):.3f}°N, {float(jra3q_pt.lon):.3f}°E)")

    era5_points[site]  = era5_pt
    jra3q_points[site] = jra3q_pt

# Randomly subsample ERA5 to one observation per 6-hour window
# to match the JRA-3Q temporal cadence prior to computing the annual maximum

def subsample_6h_random(ds, seed=RANDOM_SEED):
    """
    Retain one randomly selected observation per 6-hour window.
    The selection mask is built from the time coordinate only,
    leaving data variables Dask-lazy.
    """
    times        = pd.DatetimeIndex(ds["time"].values)
    window_start = times.floor("6h")
    unique_wins  = np.unique(window_start)

    rng           = np.random.default_rng(seed)
    offsets_h     = rng.integers(0, 6, size=len(unique_wins)).astype(int)
    win_to_offset = dict(zip(unique_wins, offsets_h))

    pos_h = ((times - window_start).total_seconds() / 3600).astype(int)
    keep  = np.fromiter(
        (pos_h[i] == win_to_offset[window_start[i]] for i in range(len(times))),
        dtype=bool,
        count=len(times),
    )

    print(f"  {ds.sizes['time']:,} → {keep.sum():,} time steps retained "
          f"({keep.sum() / len(times) * 100:.1f}%, expected ≈16.7%)")
    return ds.isel(time=keep)


for site in SITES:
    era5_points[site] = subsample_6h_random(era5_points[site])
    print(f"ERA5 / {site}: subsampling complete")

# Derive relative humidity from ERA5 2-m temperature and dew-point temperature.
# Following Romps (2024, supplementary section 1.5), relative humidity is computed
# as the ratio of the saturation vapor pressure evaluated at the dew-point temperature
# to that evaluated at the air temperature. The saturation vapor pressure is
# approximated using the August-Roche-Magnus formula (Alduchov and Eskridge, 1996).
#
# Alduchov, O.A. and Eskridge, R.E., 1996. Improved Magnus form approximation of
# saturation vapor pressure. Journal of Applied Meteorology and Climatology, 35(4),
# pp.601-609. DOI: 10.1175/1520-0450(1996)035<0601:IMFAOS>2.0.CO;2
#
# JRA-3Q provides relative humidity directly in %; convert to [0, 1].

def rh_from_T_Td(T_K, Td_K):
    a, b = 17.625, 243.04
    T_C  = T_K  - 273.15
    Td_C = Td_K - 273.15
    rh   = np.exp(a * Td_C / (b + Td_C)) / np.exp(a * T_C  / (b + T_C))
    return rh.clip(0.0, 1.0)


for site in SITES:
    era5_points[site] = era5_points[site].assign(
        rh=rh_from_T_Td(era5_points[site][ERA5_T_VAR],era5_points[site][ERA5_TD_VAR]))
    
    
    jra3q_points[site] = jra3q_points[site].assign(rh=(jra3q_points[site][JRA3Q_RH_VAR] / 100.0).clip(0.0, 1.0))

print("Relative humidity prepared for all sites.")

%%time
from heatindex import heatindex

def compute_heatindex_xr(T_K, rh):
    hi_K = xr.apply_ufunc(
        heatindex, T_K, rh,
        dask="parallelized",
        output_dtypes=[float],
    )
    hi_K.attrs = {"units": "K", "long_name": "Heat Index (Lu and Romps, 2022)"}
    return hi_K

results = {}

# ── ERA5 first — stable Zarr-backed Dask computation ─────────────────────────
for site in SITES:
    results[site] = {}
    ds   = era5_points[site]
    T_K  = ds[ERA5_T_VAR]
    rh   = ds["rh"]
    hi_K = compute_heatindex_xr(T_K, rh)
    ann_max = hi_K.groupby("time.year").max(dim="time")
    print(f"Computing ERA5 / {site} ...")
    results[site]["ERA5"] = (ann_max - 273.15).compute()
    print(f"  Done. Years: {int(results[site]['ERA5'].year[0])}–"
          f"{int(results[site]['ERA5'].year[-1])}")

print("ERA5 complete. Starting JRA-3Q ...")

# ── JRA-3Q — load point data into memory and process locally ─────────────────
for site in SITES:
    ds  = jra3q_points[site]   # already in memory as NumPy (load_jra_point)
    T_K = ds[JRA3Q_T_VAR]
    rh  = ds["rh"]

    hi_K    = xr.apply_ufunc(heatindex, T_K, rh, output_dtypes=[float])
    ann_max = hi_K.groupby("time.year").max(dim="time")

    print(f"Computing JRA-3Q / {site} ...")
    results[site]["JRA-3Q"] = (ann_max - 273.15)
    print(f"  Done. Years: {int(results[site]['JRA-3Q'].year[0])}–"
          f"{int(results[site]['JRA-3Q'].year[-1])}")

results

%%time
fig, ax = plt.subplots(figsize=(10, 5))

colors = {"Boulder_CO": "tab:blue", "Chicago_IL": "tab:red"}
ls     = {"ERA5": "-", "JRA-3Q": "--"}

for site in SITES:
    for label, da in results[site].items():
        ax.plot(
            da.year,
            da.values,
            color     = colors[site],
            linestyle = ls[label],
            linewidth = 1.5,
            label     = f"{site.replace('_', ', ')} — {label}",
        )

ax.set_xlabel("Year")
ax.set_ylabel("Annual maximum JJA heat index (°C)")
ax.set_title("Annual Maximum JJA Heat Index\nBoulder, CO and Chicago, IL (1990–2005)")
ax.legend(framealpha=0.9)
ax.grid(linestyle=":", alpha=0.5)
plt.tight_layout()
plt.show()

# Close the cluster
cluster.close()