notebooks/time_generalization.py.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "48db886d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import json\n",
    "from functools import partial \n",
    "\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns \n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "\n",
    "\n",
    "SNAKEMAKE = snakemake\n",
    "inputs = SNAKEMAKE.input\n",
    "outputs = SNAKEMAKE.output\n",
    "config = SNAKEMAKE.config\n",
    "plt_cfg = config[\"plotting\"]\n",
    "partition = SNAKEMAKE.wildcards.partition\n",
    "experiment = SNAKEMAKE.wildcards.experiment\n",
    "\n",
    "\n",
    "plt.rcParams.update(plt_cfg[\"rcparams\"])\n",
    "\n",
    "\n",
    "TASKS = [\n",
    "    \"air_temperature\",\n",
    "    \"dew_point_temperature\",\n",
    "    \"surface_air_pressure\",\n",
    "    \"relative_humidity\",\n",
    "    \"water_vapor_mixing_ratio\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d2262234",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dir = Path(outputs[0])\n",
    "out_dir.mkdir(parents=True, exist_ok=True)\n",
    "print(f\"Partition: {partition} \\n\")\n",
    "print(\"Experiment configuration: \\n\")\n",
    "exp_config = config[\"experiments\"][experiment]\n",
    "print(json.dumps(exp_config, indent=4))\n",
    "approaches = exp_config[\"approaches\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a474cf19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_predictions(ds, reftimes=None):\n",
    "    params = {}\n",
    "    path = Path(ds.encoding[\"source\"])\n",
    "    params[\"approach\"] = path.parents[2].name\n",
    "    params[\"split\"] = int(path.parents[4].name.split(\"~\")[1])\n",
    "    params[\"seed\"] = int(path.parents[3].name.split(\"~\")[1])\n",
    "    dims = list(params.keys())\n",
    "    ds = ds.assign_coords(params).expand_dims(dims)\n",
    "    if reftimes is not None:\n",
    "        ds = ds.reindex(forecast_reference_time=reftimes).load()\n",
    "    return ds\n",
    "\n",
    "def ds_to_df(ds, name):\n",
    "    return (\n",
    "        ds\n",
    "        .reset_coords(drop=True)\n",
    "        .to_array(\"variable\")\n",
    "        .to_dataframe(name)\n",
    "        .reset_index()\n",
    "    )\n",
    "\n",
    "def unstack(ds):\n",
    "    dims = [\"forecast_reference_time\", \"t\", \"station\"]\n",
    "    samples = pd.MultiIndex.from_arrays([ds[dim].values for dim in dims], names=dims)\n",
    "    ds = ds.reset_coords(drop=True).assign_coords(s=samples).unstack(\"s\")\n",
    "    return ds\n",
    "\n",
    "def remove_source_prefix(ds):\n",
    "    for var in ds.data_vars:\n",
    "        _, name = var.split(\":\")\n",
    "        ds = ds.rename({var: name})\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3a7240a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "if partition in [\"train\", \"val\"]:\n",
    "    obs = remove_source_prefix(unstack(xr.load_dataset(inputs[\"y\"])))\n",
    "    reftimes = obs.forecast_reference_time\n",
    "    pred = xr.open_mfdataset(inputs[\"predictions\"], preprocess=partial(preprocess_predictions, reftimes=reftimes), parallel=True)\n",
    "    pred = pred.dropna(\"forecast_reference_time\", \"all\")\n",
    "elif partition == \"test\":\n",
    "    pred = xr.open_mfdataset(inputs[\"predictions\"], preprocess=preprocess_predictions, parallel=True)\n",
    "    obs = remove_source_prefix(unstack(xr.load_dataset(inputs[\"y\"])))\n",
    "\n",
    "obs = obs.reindex_like(pred).load().chunk(\"auto\")\n",
    "pred = pred.load().chunk(\"auto\")\n",
    "obs, pred = xr.broadcast(obs, pred)\n",
    "err = pred - obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b83c2e11",
   "metadata": {},
   "outputs": [],
   "source": [
    "mae = abs(err).mean([\"forecast_reference_time\",\"t\",\"station\"])\n",
    "reduce_dims = [\"var\"]\n",
    "nmae = (mae / obs.std()).to_array(\"var\").mean(reduce_dims).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b63322b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "q50, q90, q99, q999 = obs.air_temperature.quantile([0.5, 0.9, 0.99, 0.999]).compute().values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2b8a4a03",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_nmae = []\n",
    "res_mae = []\n",
    "res_bias = []\n",
    "\n",
    "for q in (q50, q90, q99, q999):\n",
    "    bias = err.where(obs.air_temperature > q).mean([\"forecast_reference_time\",\"t\",\"station\"])\n",
    "    mae = abs(err).where(obs.air_temperature > q).mean([\"forecast_reference_time\",\"t\",\"station\"])\n",
    "    reduce_dims = [\"var\"]\n",
    "    nmae = (mae / obs.std()).to_array(\"var\").mean(reduce_dims).compute()\n",
    "    res_nmae.append(nmae)\n",
    "    res_mae.append(mae)\n",
    "    res_bias.append(bias)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "92eee1e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "nmae_ = xr.concat(res_nmae, dim=\"quantile\").assign_coords(quantile=[0.5, 0.9, 0.99, 0.999])\n",
    "mae_ = xr.concat(res_mae, dim=\"quantile\").assign_coords(quantile=[0.5, 0.9, 0.99, 0.999])\n",
    "bias_ = xr.concat(res_bias, dim=\"quantile\").assign_coords(quantile=[0.5, 0.9, 0.99, 0.999])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f74b5456",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, figsize=(6,5))\n",
    "\n",
    "\n",
    "df = (nmae_.to_dataframe(name=\"Aggregated NMAE\").reset_index()\n",
    "      .replace(plt_cfg[\"approach_names\"])\n",
    "      .rename(columns={\"approach\":\"Approach\"})\n",
    "     )\n",
    "\n",
    "\n",
    "sns.boxplot(\n",
    "    data=df,\n",
    "    x=\"quantile\",\n",
    "    y=\"Aggregated NMAE\",\n",
    "    hue=\"Approach\",\n",
    "    hue_order=list(plt_cfg[\"approach_names\"].values()),\n",
    "    palette=plt_cfg[\"approach_colors\"],\n",
    "    showfliers=False,\n",
    "    ax=ax\n",
    ")\n",
    "\n",
    "ax = plt.gca()\n",
    "ax.set_xlabel(\"Quantile of $T$\")\n",
    "lgd = ax.legend(\n",
    "        bbox_to_anchor=(0.13, 0.885, 0.78, 0.1),\n",
    "        loc=\"lower left\",\n",
    "        ncol=2,\n",
    "        mode=\"expand\",\n",
    "        borderaxespad=0.,\n",
    "        frameon=False,\n",
    "        fontsize=11,\n",
    "        bbox_transform=plt.gcf().transFigure\n",
    "    )\n",
    "\n",
    "plt.savefig(out_dir / \"NMAE_vs_quantile.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27bdf93a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d037f00d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c76a0b2c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}