{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1eafa11e-18d8-4bc5-be0c-22a4ec3cc945",
   "metadata": {},
   "source": [
    "## IMPORTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b531dcb-9364-4115-a61d-bac9167c4f33",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import uxarray as ux\n",
    "import xarray as xr\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c26acb7c-dfe1-403e-8265-6bf31feadf3d",
   "metadata": {},
   "source": [
    "## DASK CONFIG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01c78d4-ffde-481e-93ed-032252bb52db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "\n",
    "# client = Client(n_workers=1, threads_per_worker=6, processes=False, memory_limit=\"100GB\")\n",
    "client = Client(n_workers=6, threads_per_worker=2, memory_limit=\"12GB\")\n",
    "client.dashboard_link"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeed2cb5-71f4-4f0c-b71a-624fe0144053",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Will measure peak notbook (kernel only) memory usage and displaye at the end\n",
    "from distributed.diagnostics.memory_sampler import MemorySampler\n",
    "ms = MemorySampler()\n",
    "cm = ms.sample(\"full notebook\")\n",
    "cm.__enter__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e82865e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Will measure peak combined (kernel + dask workers) memory usage and displaye at the end\n",
    "import os\n",
    "import threading\n",
    "import psutil\n",
    "\n",
    "_peak_rss_bytes = 0\n",
    "_stop_rss_watch = threading.Event()\n",
    "\n",
    "\n",
    "def _watch_rss(interval=5):\n",
    "    global _peak_rss_bytes\n",
    "    root = psutil.Process(os.getpid())\n",
    "    while not _stop_rss_watch.is_set():\n",
    "        try:\n",
    "            procs = [root] + root.children(recursive=True)\n",
    "            total = sum(p.memory_info().rss for p in procs if p.is_running())\n",
    "            _peak_rss_bytes = max(_peak_rss_bytes, total)\n",
    "        except psutil.Error:\n",
    "            pass\n",
    "        _stop_rss_watch.wait(interval)\n",
    "\n",
    "\n",
    "threading.Thread(target=_watch_rss, daemon=True).start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "609ed9dc-ac91-4d9b-a664-f3344259a7a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The distributed.shuffle messages can be too much - suppress them!\n",
    "\n",
    "import logging\n",
    "logging.getLogger(\"distributed.shuffle\").setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dbc332e-bcd0-4c8f-b283-4a0be5f2371a",
   "metadata": {},
   "outputs": [],
   "source": [
    "TIME_CHUNK = 40\n",
    "FACE_CHUNK = 2000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44eefc7b-4f0f-42d0-b668-fee999e6b7e4",
   "metadata": {},
   "source": [
    "## RUN CONFIG (Native Grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc2f4d16-fc4c-489f-b9c0-0ddca59fe06b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Determine if notebook is to be run on local device or HPC\n",
    "RUN_ON = 'hpc'  # Notebook to be run on either of \"local\" | \"hpc\"\n",
    "\n",
    "if RUN_ON == 'local':\n",
    "    KERCHUNK_PATH = \"https://data.gdex.ucar.edu/d651007/kerchunk/\"\n",
    "    KERCHUNK_ACCESS = \"-remote-osdf\"\n",
    "    KERCHUNK_REMOTE_PROTOCOL = \"osdf\"\n",
    "    GRID_FILE = \"ne120np4_pentagons_100310.nc\"  # Local\n",
    "elif RUN_ON == 'hpc':\n",
    "    KERCHUNK_PATH = \"/gdex/data/d651007/kerchunk/\"\n",
    "    KERCHUNK_ACCESS = \"\"  # Kerchunk dataset access, either of \"\" | \"-remote-https\" | \"-remote-osdf\"\n",
    "    KERCHUNK_REMOTE_PROTOCOL = \"file\"\n",
    "    # GRID_FILE = \"/glade/p/cesmdata/cseg/inputdata/share/scripgrids/ne120np4_pentagons_100310.nc\"  # Glade - this is older?\n",
    "    GRID_FILE = \"/glade/campaign/cesm/cesmdata/cseg/inputdata/share/scripgrids/ne120np4_pentagons_100310.nc\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a969d7b-32f3-4773-8c42-5561891f697d",
   "metadata": {},
   "outputs": [],
   "source": [
    "MEMBERS = range(2, 11)   # member 1 has no 6-hourly stream\n",
    "\n",
    "MODEL = \"atm\"\n",
    "FREQ = \"hour_6\"  # or \"day_1\"\n",
    "VAR = \"PRECT\"\n",
    "VARS = ['TS', 'PSL', 'PRECT']\n",
    "\n",
    "EXTREME_MEMBER = 10  # to be used for one-member cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f19c432-3752-411a-85f1-93a14a48e5c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "T_START = \"1979-01-01\"\n",
    "T_END = \"1988-12-31\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51a4ac08-5588-45f0-bd33-d435a0030abc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unit conversions\n",
    "\n",
    "SEC_PER_DAY = 24 * 3600  # 6-hourly output\n",
    "M_TO_MM = 1000.0\n",
    "\n",
    "# there are four 6-hour timesteps per day, and exactly 365 days per year in this data;\n",
    "# it avoids dealing with leap years, and has cftime.DatetimeNoLeap time coordinates.\n",
    "STEPS_PER_DAY = 4\n",
    "DAYS_PER_YEAR = 365\n",
    "STEPS_PER_YEAR = STEPS_PER_DAY * DAYS_PER_YEAR"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cffab96-bec4-44dc-bed3-bba420a5c99b",
   "metadata": {},
   "source": [
    "## UXARRAY (Explore Grid and a single ensemble member)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c794f8c2-61d9-44c8-ba20-3603656f696c",
   "metadata": {},
   "source": [
    "### Grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d701b80c-d353-41a1-b2c0-26226f8fcec7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "uxgrid = ux.open_grid(GRID_FILE)\n",
    "uxgrid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1609448-eda2-4c7b-ae84-ddf3a84d4622",
   "metadata": {},
   "source": [
    "### Single ensemble member"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e88c640c-ec6e-4ae2-a3d0-b5219b2ad943",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A single ensemble file for sampling\n",
    "data_filename = \"b.e13.BHISTC5.ne120_t12.cesm-ihesp-hires1.0.46-1920-2005.\" + f\"{EXTREME_MEMBER:03d}\" + \".\" + MODEL + \".\" + FREQ + KERCHUNK_ACCESS + \".parq\"\n",
    "data_file = KERCHUNK_PATH + data_filename"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7897b7a4-0d70-46d6-910d-2eef0c6566af",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# One ensemble member\n",
    "uxds = ux.open_dataset(\n",
    "    GRID_FILE, data_file, \n",
    "    engine=\"kerchunk\", \n",
    "    decode_timedelta=False, \n",
    "    storage_options={\"remote_protocol\": KERCHUNK_REMOTE_PROTOCOL},\n",
    "    chunks={}\n",
    ")\n",
    "uxds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e920789-7745-4823-a815-8c0db7ec1c80",
   "metadata": {},
   "source": [
    "#### Naive plot of a member's PRECT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d650959b-2252-4966-a744-879bd9ea9d98",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "uxds[VAR].isel(time=0).plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a761482c-dabd-454c-8992-03fcb0adacc8",
   "metadata": {},
   "source": [
    "## YEARLY GLOBAL STATS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6306cff2-2e1b-40bf-81a6-fb8ba919ce78",
   "metadata": {},
   "outputs": [],
   "source": [
    "uxds_yearly = uxds[VARS].isel(time=slice(None, None, STEPS_PER_YEAR)).chunk()\n",
    "print(f'Size: {uxds_yearly.nbytes/1e9:.2f} GB')\n",
    "uxds_yearly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f673ecbc-7d09-48ad-85d3-f5b59e010a08",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "uxds_yearly = uxds_yearly.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d021603-3ce2-4244-a34e-f102612812e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# naive average: just take mean of all faces across the globe\n",
    "mean_yearly = uxds_yearly.mean('n_face')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ca3c260-b146-42cd-bf58-07141eff9a46",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# face-area-weighted mean: wmean = ds.weighted_mean()\n",
    "# (not yet implemented on UxDataset, so can't run that yet.\n",
    "# workaround: convert to UxDataset, take mean, then convert back.)\n",
    "weighted_mean_yearly = uxds_yearly.to_array(dim='variable').weighted_mean().to_dataset(dim='variable')\n",
    "\n",
    "# Also, similar workaround as above: convert to xarray.DataArray\n",
    "# weighted_mean_yearly = weighted_mean_yearly.to_xarray()\n",
    "# wmean_jan1 = weighted_mean_yearly  # alias so we can still refer back to this in later sections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9739146c-9a09-46b2-951a-0a2b23992c19",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# visualize the results (pure xarray; nothing in this cell uses uxarray anymore)\n",
    "fig, axes = plt.subplots(nrows=len(VARS), ncols=1, figsize=(5, 2.5*len(VARS)))\n",
    "for ax, var in zip(axes, VARS):\n",
    "    weighted_mean_yearly[var].to_xarray().plot(ax=ax, label=\"weighted mean\", marker='o', ms=2)\n",
    "    mean_yearly[var].to_xarray().plot(ax=ax, label=\"mean\", ls='--', marker='o', ms=2)\n",
    "# labeling & formatting\n",
    "axes[0].legend()\n",
    "for ax in axes[:-1]:\n",
    "    ax.set(xlabel=\"\", xticklabels=[])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d04d8a62-4f63-4542-a46c-4389f9b2897f",
   "metadata": {},
   "source": [
    "## WHOLE ENSEMBLE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "864bf9c3-ce82-4eb2-8cc6-fdb49d756fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function to open a single ensemble member\n",
    "def open_member(member):\n",
    "    hits = sorted(Path(KERCHUNK_PATH).rglob(f\"*.{member:03d}.{MODEL}.{FREQ}{KERCHUNK_ACCESS}.parq\"))\n",
    "\n",
    "    if not hits:\n",
    "        return None\n",
    "    assert len(hits) == 1, f\"expected exactly one file for member {member}, found {len(hits)}: {hits}\"\n",
    "\n",
    "    return ux.open_dataset(\n",
    "        GRID_FILE, \n",
    "        str(hits[0]), \n",
    "        engine=\"kerchunk\", \n",
    "        decode_timedelta=False, \n",
    "        storage_options={\"remote_protocol\": KERCHUNK_REMOTE_PROTOCOL},\n",
    "        chunks={\"time\": TIME_CHUNK}\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c7e5b84-398b-4c99-92b6-21c0de8bdd57",
   "metadata": {},
   "source": [
    "### Read the ensemble members"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "196c3b7c-d466-4b1a-8a2d-76920e981dc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Opened lazily only, no eager persist here. Native storage isn't spatially\n",
    "# chunked (chunks=[1, 777602], one full-globe row per timestep), so every chunk\n",
    "# read pays a ~33x read/decompress tax before the CONUS mask discards ~97% of it.\n",
    "members = {m: da for m in MEMBERS if (da := open_member(m)) is not None}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac051d45-7803-40e4-bd8a-b88f5d1b9e88",
   "metadata": {},
   "source": [
    "#### CONUS - region to subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd731e8e-6a2e-40d2-804f-2bd49d480f70",
   "metadata": {},
   "outputs": [],
   "source": [
    "lat, lon = uxgrid.face_lat.values, uxgrid.face_lon.values\n",
    "# lon180 = ((lon + 180) % 360) - 180\n",
    "\n",
    "conus = (lat > 25) & (lat < 50) & (lon > -125) & (lon < -66)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1e9c446-62eb-4c8b-89f5-961aa40c8140",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "members = {key: ds[VAR].sel(time=slice(T_START, T_END)).isel(n_face=conus) for key, ds in members.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f56de63-0448-47c2-9dfe-d8495f203b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Persist just this one member now, staged at the point it's first needed (it's\n",
    "# reused again below for extreme_metrics) — instead of persisting all 9 members\n",
    "# upfront, which concentrated the ~33x full-globe read/decompress tax into one\n",
    "# burst across the whole cluster.\n",
    "members[EXTREME_MEMBER] = members[EXTREME_MEMBER].persist()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fa5cfcb-e6f1-421c-934d-8fff5b865d3c",
   "metadata": {},
   "source": [
    "## PRECIPITATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09c52632",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"Source units:\", members[EXTREME_MEMBER].attrs.get(\"units\", \"unknown\"))\n",
    "\n",
    "# CESM/CAM PRECT is a total precipitation rate in m/s -> convert to mm/day\n",
    "precip_mm = members[EXTREME_MEMBER] * M_TO_MM * SEC_PER_DAY\n",
    "precip_mm.attrs[\"units\"] = \"mm/day\"\n",
    "precip_mm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddfc7c20-e537-4679-92a1-e430ea22f6e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "precip_mm.mean(dim=\"time\").plot(\n",
    "    cmap=\"YlGnBu\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=f\"6-hourly PRECT (mm/day) - 10-yr avg — Ensemble member {EXTREME_MEMBER} (USCLIVAR ne120 native grid)\" ,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe5c9989",
   "metadata": {},
   "source": [
    "## EXTREME PRECIPITATION METRICS\n",
    "\n",
    "Using the 6-hourly `PRECT` of ensemble members, compute per-gridcell extreme precipitation metrics over CONUS and map their spatial distribution on the native ne120 mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea8dbe18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function to calculate a couple of preferred percentiles, max, and the total 95th-percentile\n",
    "\n",
    "def extreme_precip_metrics(uxda, high_q=0.95, extreme_q=0.99, face_chunk=FACE_CHUNK):\n",
    "    \"\"\"Per-gridcell extreme precipitation metrics along the \"time\" dimension.\"\"\"\n",
    "    \n",
    "    uxda = uxda.chunk({\"time\": -1, \"n_face\": face_chunk})\n",
    "\n",
    "    q_high = uxda.quantile(high_q, dim=\"time\", skipna=True).drop_vars(\"quantile\")\n",
    "    q_extreme = uxda.quantile(extreme_q, dim=\"time\", skipna=True).drop_vars(\"quantile\")\n",
    "    uxda_max = uxda.max(dim=\"time\", skipna=True)\n",
    "\n",
    "    above_high_total = uxda.where(uxda >= q_high).sum(dim=\"time\", skipna=True)\n",
    "    total = uxda.sum(dim=\"time\", skipna=True)\n",
    "    r95ptot = (above_high_total / total) * 100\n",
    "\n",
    "    freq_extreme = (uxda >= q_extreme).sum(dim=\"time\") / uxda.sizes[\"time\"] * 100\n",
    "\n",
    "    metrics = xr.Dataset(\n",
    "        {\n",
    "            f\"p{int(high_q * 100)}\": q_high,\n",
    "            f\"p{int(extreme_q * 100)}\": q_extreme,\n",
    "            \"max\": uxda_max,\n",
    "            \"r95ptot\": r95ptot,\n",
    "            f\"freq_p{int(extreme_q * 100)}\": freq_extreme,\n",
    "        }\n",
    "    )\n",
    "    return ux.UxDataset(metrics, uxgrid=uxda.uxgrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2938d360",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "extreme_metrics = extreme_precip_metrics(precip_mm).compute()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77b9dd01",
   "metadata": {},
   "source": [
    "### 95th and 99th percentile, max, and r95ptot over 10 years"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "264ecf3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "extreme_metrics[\"p95\"].plot(\n",
    "    cmap=\"YlGnBu\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=f\"6-hourly PRECT (mm/day) - 95th percentile (10 year) — ensemble member {EXTREME_MEMBER} (USCLIVAR ne120 native grid)\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a3b3d2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "extreme_metrics[\"p99\"].plot(\n",
    "    cmap=\"YlGnBu\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=f\"6-hourly PRECT (mm/day) - 99th percentile (10 year) — ensemble member {EXTREME_MEMBER} (USCLIVAR ne120 native grid)\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43dabb00",
   "metadata": {},
   "outputs": [],
   "source": [
    "extreme_metrics[\"max\"].plot(\n",
    "    cmap=\"magma\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=f\"6-hourly PRECT (mm/day) - 10-year maximum — ensemble member {EXTREME_MEMBER} (USCLIVAR ne120 native grid)\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2399f428",
   "metadata": {},
   "outputs": [],
   "source": [
    "extreme_metrics[\"r95ptot\"].plot(\n",
    "    cmap=\"YlGnBu\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=f\"6-hourly PRECT (mm/day) - % of total precip from ≥ p95 — ensemble member {EXTREME_MEMBER} (USCLIVAR ne120 grid)\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e18987f",
   "metadata": {},
   "source": [
    "### Ensemble mean of the 99th percentile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d751283-d999-45c1-bbe8-81f3b54021f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Build the p99 metric lazily for every available member, then compute the mean once \n",
    "# so dask can schedule the whole ensemble mean as a single parallel graph.\n",
    "\n",
    "member_p99 = {}\n",
    "for m, uxda in members.items():\n",
    "    if m == EXTREME_MEMBER:\n",
    "        member_p99[m] = extreme_metrics[\"p99\"]\n",
    "        continue\n",
    "    precip = (uxda * M_TO_MM * SEC_PER_DAY).chunk({\"time\": -1, \"n_face\": FACE_CHUNK})\n",
    "    member_p99[m] = precip.quantile(0.99, dim=\"time\", skipna=True).drop_vars(\"quantile\")\n",
    "\n",
    "ensemble_p99 = xr.concat(list(member_p99.values()), dim=\"member\").assign_coords(\n",
    "    member=list(member_p99.keys())\n",
    ")\n",
    "\n",
    "ensemble_p99_mean = ensemble_p99.mean(dim=\"member\").compute()\n",
    "\n",
    "ensemble_p99_mean_ux = ux.UxDataArray(\n",
    "    ensemble_p99_mean, uxgrid=members[EXTREME_MEMBER].uxgrid, name=\"p99_ensemble_mean\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c290f1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble_p99_mean_ux.plot(\n",
    "    cmap=\"YlGnBu\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=\"6-hourly PRECT (mm/day) - 99th percentile (10 year) — Ensemble mean (USCLIVAR ne120 native grid)\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e2d3b4d",
   "metadata": {},
   "source": [
    "## NATIVE GRID vs CONSERVATIVELY REMAPPED GRID\n",
    "\n",
    "To see the effect of conservative regridding, we will now start look ating the conservatively regridded dataset for the workshop as well, and we will take the difference between the results from the native grid and this regridded (lat/lon) grid — this isolates how much conservative aggregation smooths out extremes relative to the native resolution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eb759a1-dd85-41c9-b7cb-6f344e09edb2",
   "metadata": {},
   "source": [
    "### USCLIVAR's regridded data\n",
    "Read in the regridded CESM-HR global 6-hourly variables used to identify and track extreme weather events (`data_file_regridded_large`) below. The [workshop's website](https://usclivar.org/meetings/high-resolution-modeling-data-access) points to Globus Guest Collections for the datasets, but we will access them through Glade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64c207e5-aa26-40f8-b7ca-71f7e89121be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# High-frequency (6-hr) outputs for extreme weather event identification and tracking:\n",
    "data_file_regridded_large = '/glade/campaign/univ/utam0017/Extreme_Weather_Risk_dataset'\n",
    "\n",
    "# # There was also a smaller dataset intended for the practicum (not used for this notebook):\n",
    "# data_file_regridded_small = '/glade/campaign/univ/utam0017/High-Resolution_Modeling_Workshop_MESACLIP_dataset'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69646d40-9341-4e82-8437-7f32141d94d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "_REGRID_MEMBER_TEMPLATE = \"b.e13.HF-TNST.rcp85.ne120_t12.1920-2100.{member:03d}\"\n",
    "_REGRID_FILENAME_RE = re.compile(r\"\\.(\\d{10})-(\\d{10})\\.nc$\")\n",
    "\n",
    "\n",
    "def open_regridded_member(member, lat_slice=None, lon_slice=None):\n",
    "    \"\"\"Read one ensemble member's externally-regridded (0.25 deg lat/lon) field.\"\"\"\n",
    "    \n",
    "    member_dir = Path(data_file_regridded_large) / _REGRID_MEMBER_TEMPLATE.format(member=member) / VAR\n",
    "    hits = sorted(member_dir.glob(f\"*.{VAR}.latlon_0.25x0.25_0E.*.nc\"),\n",
    "                  key=lambda p: _REGRID_FILENAME_RE.search(p.name).group(1))\n",
    "    if not hits:\n",
    "        return None\n",
    "\n",
    "    t_start_key, t_end_key = T_START.replace(\"-\", \"\"), T_END.replace(\"-\", \"\")\n",
    "    hits = [\n",
    "        p for p in hits\n",
    "        if not (_REGRID_FILENAME_RE.search(p.name).group(2)[:8] < t_start_key\n",
    "                or _REGRID_FILENAME_RE.search(p.name).group(1)[:8] > t_end_key)\n",
    "    ]\n",
    "\n",
    "    ds = xr.open_mfdataset(\n",
    "        hits,\n",
    "        combine=\"nested\",\n",
    "        concat_dim=\"time\",\n",
    "        decode_times=xr.coders.CFDatetimeCoder(use_cftime=True),\n",
    "        decode_timedelta=False,\n",
    "        # Native chunks are [1, 720, 1440] — one full-globe row per timestep, same\n",
    "        # \"not spatially chunked\" situation as the kerchunk source above.\n",
    "        chunks={\"time\": 40},\n",
    "        data_vars=\"minimal\", coords=\"minimal\", compat=\"override\",\n",
    "    )\n",
    "    \n",
    "    da = ds[VAR].sel(time=slice(T_START, T_END))\n",
    "    # lat/lon here are a genuine regular grid, so a contiguous slice() is basic\n",
    "    # indexing dask can push down cheaply — unlike the CONUS boolean mask on the\n",
    "    # native ne120 mesh, this doesn't force reading the full globe per timestep.\n",
    "    if lat_slice is not None:\n",
    "        da = da.sel(lat=lat_slice)\n",
    "    if lon_slice is not None:\n",
    "        da = da.sel(lon=lon_slice)\n",
    "    return da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d1f3ecf-afe4-4f0f-b706-4bd8da818a56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CONUS bbox in this dataset's 0-360 longitude convention (native mesh used -125/-66)\n",
    "precip_regridded_raw = open_regridded_member(\n",
    "    EXTREME_MEMBER, lat_slice=slice(25, 50), lon_slice=slice(235, 294),\n",
    ")\n",
    "precip_regridded_mm = precip_regridded_raw * M_TO_MM * SEC_PER_DAY\n",
    "precip_regridded_mm.attrs[\"units\"] = \"mm/day\"\n",
    "precip_regridded_mm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25568979-1d95-468b-81db-8a484e7bb0a0",
   "metadata": {},
   "source": [
    "### 99th-percentile (6-hourly PRECT, 10 years) on Regridded Data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae1d75c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "p99_regridded = (\n",
    "    precip_regridded_mm\n",
    "    .chunk({\"time\": -1, \"lat\": 50, \"lon\": 50})\n",
    "    .quantile(0.99, dim=\"time\", skipna=True)\n",
    "    .drop_vars(\"quantile\")\n",
    "    .compute()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ea853bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "p99_regridded.hvplot(\n",
    "    cmap=\"YlGnBu\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=f\"6-hourly PRECT (mm/day) - 99th percentile (10 year) — ensemble member {EXTREME_MEMBER} (USCLIVAR regridded 0.25° grid)\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "726bb97b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "p99_native = extreme_metrics[\"p99\"]\n",
    "p99_native.plot(\n",
    "    cmap=\"YlGnBu\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=f\"6-hourly PRECT (mm/day) - 99th percentile (10 year) — ensemble member {EXTREME_MEMBER} (USCLIVAR ne120 native grid)\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b369673d",
   "metadata": {},
   "source": [
    "### Difference\n",
    "\n",
    "Let's look at the difference of results between the native-mesh `p99_native` read directly by UXarray and the regridded 0.25° grid `p99_regridded`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2c03e3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "regridded_lon_centers = ((p99_regridded[\"lon\"].values + 180) % 360) - 180  # this product's lon is 0-360\n",
    "regridded_lat_centers = p99_regridded[\"lat\"].values\n",
    "uxgrid_from_regridded = ux.Grid.from_structured(lon=regridded_lon_centers, lat=regridded_lat_centers)\n",
    "\n",
    "p99_native_remapped_to_regridded_grid = p99_native.remap.nearest_neighbor(uxgrid_from_regridded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef01e9d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "p99_native_remapped_to_regridded_grid.plot(\n",
    "    cmap=\"YlGnBu\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=\"6-hourly PRECT (mm/day) - 99th percentile (10 year) — native ne120 mesh remapped onto the 0.25° grid\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f9296bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "p99_diff_external = ux.UxDataArray(\n",
    "    xr.DataArray(\n",
    "        p99_native_remapped_to_regridded_grid.values - p99_regridded.values.ravel(),\n",
    "        dims=[\"n_face\"], name=\"p99_diff_external\",\n",
    "    ),\n",
    "    uxgrid=uxgrid_from_regridded,\n",
    ")\n",
    "p99_diff_external.plot(\n",
    "    cmap=\"RdBu_r\",\n",
    "    features=[\"borders\", \"coastline\"],\n",
    "    title=(\n",
    "        \"6-hourly PRECT (mm/day) - 99th percentile (10 year) — ensemble member 10\\n\"\n",
    "        \"USCLIVAR native grid minus USCLIVAR 0.25° regridded grid\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a82cf76f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"MESACLIP native mesh max p99:                                              \", float(p99_native.max()))\n",
    "print(\"MESACLIP native mesh remapped onto MESACLIP regridded 0.25° grid, max p99: \", float(p99_native_remapped_to_regridded_grid.max()))\n",
    "print(\"MESACLIP regridded 0.25° grid max p99:                                     \", float(p99_regridded.max()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d76ad3c-48da-4831-9651-e71d440dc6c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "cm.__exit__(None, None, None)\n",
    "ms.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bec380f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "_stop_rss_watch.set()\n",
    "print(f\"Peak combined RSS (kernel + dask workers): {_peak_rss_bytes / 1e9:.2f} GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a804536-65f8-4fa0-9b75-99abbec41586",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be75719b-0207-48c9-8eea-bb94d0d90503",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:uxarray-usclivar]",
   "language": "python",
   "name": "conda-env-uxarray-usclivar-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.14.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}