Annual max summer heat index

!pip install heatindex

import matplotlib.pyplot as plt
import numpy as np
import os
import shutil
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, LocalCluster

# Auto-select each catalog: POSIX where /gdex/data is mounted (NCAR HPC, CIRRUS
# Binder), otherwise the OSDF catalog over the network.
_era5_posix  = "/gdex/data/d633000/catalogs/d633000-posix.json"
_jra3q_posix = "/gdex/data/d640000/catalogs/d640000-posix.json"
era5_catalog  = _era5_posix  if os.path.exists(_era5_posix)  else "https://data.gdex.ucar.edu/d633000/catalogs/d633000-osdf.json"
jra3q_catalog = _jra3q_posix if os.path.exists(_jra3q_posix) else "https://data.gdex.ucar.edu/d640000/catalogs/d640000-osdf.json"
print(era5_catalog)
print(jra3q_catalog)

# Set up your sratch folder path
username       = os.environ["USER"]
glade_scratch  = "/glade/derecho/scratch/" + username
print(glade_scratch)

def get_pbs_cluster(n_workers=5):
    """PBSCluster on NCAR HPC (Casper/Derecho)."""
    scratch = f"/glade/derecho/scratch/{os.environ.get('USER', '')}"
    cluster = PBSCluster(
        job_name        = "dask-wk25",
        cores           = 1,
        memory          = "8GiB",
        processes       = 1,
        local_directory = f"{scratch}/dask/spill/",
        log_directory   = f"{scratch}/dask/logs/",
        resource_spec   = "select=1:ncpus=1:mem=8GB",
        queue           = "casper",
        walltime        = "5:00:00",
        interface       = "ext",
    )
    cluster.scale(n_workers)
    return cluster


def get_local_cluster():
    """LocalCluster for a laptop or a CIRRUS Binder pod (no PBS scheduler)."""
    return LocalCluster()

# Auto-detect the environment: a `qsub` on PATH means a PBS scheduler is present
# (Casper/Derecho) -> PBSCluster; otherwise fall back to LocalCluster (CIRRUS
# Binder or a laptop). No editing needed to move between environments.
on_hpc  = bool(shutil.which("qsub"))
cluster = get_pbs_cluster() if on_hpc else get_local_cluster()
client  = Client(cluster)
if on_hpc:
    client.wait_for_workers(n_workers=5)

print(f"Detected {'NCAR HPC -> PBSCluster' if on_hpc else 'no PBS -> LocalCluster'}")
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         = 2025

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

# ── Output ────────────────────────────────────────────────────────────────────
OUTPUT_ZARR = glade_scratch + "/max_jja_ehi_boulder_chicago.zarr"

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_t_ds  = xr.open_zarr(era5_t_path,  chunks={})[[ERA5_T_VAR]]
era5_td_ds = xr.open_zarr(era5_td_path, chunks={})[[ERA5_TD_VAR]]
era5_ds = xr.merge([era5_t_ds, era5_td_ds])
#
era5_ds

jra3q_ds   = xr.open_dataset(jra3q_t_path, engine='kerchunk', chunks={})[[JRA3Q_T_VAR, JRA3Q_RH_VAR]]
jra3q_ds

# Subset to analysis period and JJA
era5_ds  = era5_ds.sel(time=slice(str(YEAR_START_ERA5),  str(YEAR_END)))
jra3q_ds = jra3q_ds.sel(time=slice(str(YEAR_START_JRA3Q), str(YEAR_END)))

era5_ds  = era5_ds.sel(time=era5_ds.time.dt.month.isin(JJA_MONTHS))
jra3q_ds = jra3q_ds.sel(time=jra3q_ds.time.dt.month.isin(JJA_MONTHS))

#
# Extract nearest grid cell for each site
era5_points  = {}
jra3q_points = {}

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

    era5_pt  = era5_ds.sel(latitude=lat, longitude=lon, method="nearest")
    jra3q_pt = jra3q_ds.sel(lat=lat,     lon=lon,       method="nearest")

    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):
    """
    Compute the extended heat index (Lu and Romps, 2022) using the
    Python implementation of Romps (2024).
    
    Parameters
    ----------
    T_K : xr.DataArray  — 2-m temperature in Kelvin
    rh  : xr.DataArray  — relative humidity in [0, 1]

    Returns
    -------
    xr.DataArray — heat index in Kelvin
    """
    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


# Compute heat index and annual maximum JJA value for each site
results = {}

for site in SITES:
    results[site] = {}
    for label, points in [("ERA5", era5_points), ("JRA-3Q", jra3q_points)]:
        ds   = points[site]
        T_K  = ds[ERA5_T_VAR]  if label == "ERA5" else ds[JRA3Q_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 {label} / {site} ...")
        results[site][label] = (ann_max - 273.15).compute()
        print(f"  Done. Years: {int(results[site][label].year[0])}–"
              f"{int(results[site][label].year[-1])}")

out_vars = {}

for site in SITES:
    slug = site.lower()
    for label, da in results[site].items():
        key = f"{label.lower().replace('-', '')}_{slug}_hi_max_jja"
        da  = da.drop_vars(["latitude", "longitude", "lat", "lon"], errors="ignore")
        da  = da.rename(key)
        out_vars[key] = da

ds_out = xr.Dataset(out_vars)
ds_out.to_zarr(OUTPUT_ZARR, mode="w")
print(f"Results written to: {OUTPUT_ZARR}")
print(xr.open_zarr(OUTPUT_ZARR))

%%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–2025)")
ax.legend(framealpha=0.9)
ax.grid(linestyle=":", alpha=0.5)
plt.tight_layout()
plt.show()

# Close the cluster
cluster.close()