notebooks/loss_alpha.py.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68165074",
   "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",
    "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",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9974ee9",
   "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",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff574f8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_predictions(ds, reftimes=None):\n",
    "    path = Path(ds.encoding[\"source\"])\n",
    "    params = {}\n",
    "    for kv in path.parent.name.split(\"-\"):\n",
    "        k, v = kv.split(\"~\")\n",
    "        params[k] = v\n",
    "    params[\"approach\"] = path.parents[2].name\n",
    "    params[\"split\"] = int(path.parents[3].name.split(\"~\")[1])\n",
    "    params[\"seed\"] = int(path.parents[4].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": null,
   "id": "feab6517",
   "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({\"forecast_reference_time\": 200})\n",
    "pred = pred.squeeze().load().chunk({\"forecast_reference_time\": 200})\n",
    "obs = obs.squeeze()\n",
    "obs, pred = xr.broadcast(obs, pred)\n",
    "err = pred - obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "480de138",
   "metadata": {},
   "outputs": [],
   "source": [
    "mae = abs(err).mean([\"forecast_reference_time\",\"t\"])\n",
    "mae = (\n",
    "    ds_to_df(mae, \"MAE\")\n",
    "    .replace(plt_cfg[\"approach_names\"])\n",
    "    .replace(plt_cfg[\"var_short_names\"])\n",
    ")\n",
    "\n",
    "\n",
    "fig, axs = plt.subplots(1, len(TASKS), figsize=(5,5), sharey=False)\n",
    "for i, var in enumerate(TASKS):\n",
    "    var = plt_cfg[\"var_short_names\"][var]\n",
    "    sns.boxplot(\n",
    "        data=mae.query(f\"variable=='{var}'\"),\n",
    "        x=\"variable\",\n",
    "        y=\"MAE\",\n",
    "        palette=plt_cfg[\"approach_colors\"],\n",
    "        ax=axs[i]\n",
    "    )\n",
    "    \n",
    "    axs[i].set_ylim(0., None)\n",
    "    \n",
    "for ax in axs:\n",
    "    ax.set_ylabel(\"\")\n",
    "    ax.set_xlabel(\"\")\n",
    "\n",
    "plt.subplots_adjust(wspace=1.5, hspace=0)\n",
    "axs[0].set_ylabel(\"Mean absolute error\")\n",
    "\n",
    "\n",
    "plt.savefig(out_dir / \"mae_boxplots.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d32dd085",
   "metadata": {},
   "outputs": [],
   "source": [
    "mae = np.abs(err).mean([\"forecast_reference_time\",\"t\"])\n",
    "mae[\"aggreated\"] = mae.to_array(\"v\").mean(\"v\")\n",
    "mae = ds_to_df(mae, \"mae\")\n",
    "mae = mae.replace(plt_cfg[\"approach_names\"])\n",
    "mae = mae.replace(plt_cfg[\"var_short_names\"])\n",
    "cp = sns.catplot(data=mae, x=\"loss.alpha\", y=\"mae\", col=\"variable\", palette=\"Blues\", kind=\"box\", sharey=False, col_wrap=3, aspect=0.5)\n",
    "for ax in cp.axes:\n",
    "    ax.set_title(ax.title.get_text().split(\"=\")[1][1:])\n",
    "    ax.set_ylim(0., None)\n",
    "\n",
    "plt.savefig(out_dir / \"alpha_vs_mae.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a759078d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rh_from_t_td(t, t_d):\n",
    "    \n",
    "    rh = xr.where(\n",
    "        t >= 0,\n",
    "        100 * np.exp((17.368 * t_d)/(238.83+t_d) - (17.368 * t)/(238.83+t)),\n",
    "        100 * np.exp((17.856 * t_d)/(245.52+t_d) - (17.856 * t)/(245.52+t))\n",
    "    )\n",
    "    \n",
    "    return rh\n",
    "\n",
    "def e_from_t_rh(t, rh):\n",
    "    e = xr.where(\n",
    "        t >= 0,\n",
    "        rh / 100 * 6.107 * np.exp((17.368 * t) / (238.83 + t)),\n",
    "        rh / 100 * 6.108 * np.exp((17.856 * t) / (245.52 + t)),\n",
    "    )\n",
    "    return e\n",
    "\n",
    "\n",
    "def td_from_e_t(e, t):\n",
    "\n",
    "    td = xr.where(\n",
    "        t >= 0.0,\n",
    "        -238.83 * (np.log(e / 6.107)) / (np.log(e / 6.107) - 17.368),\n",
    "        -245.52 * (np.log(e / 6.108)) / (np.log(e / 6.108) - 17.856),\n",
    "    )\n",
    "\n",
    "    return td\n",
    "\n",
    "def r_from_e_p(e, p):\n",
    "    return 622.0 * (e / (p - e))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82d3f74b",
   "metadata": {},
   "outputs": [],
   "source": [
    "rh = pred[\"relative_humidity\"]\n",
    "t = pred[\"air_temperature\"]\n",
    "t_d = pred[\"dew_point_temperature\"]\n",
    "r = pred[\"water_vapor_mixing_ratio\"]\n",
    "p = pred[\"surface_air_pressure\"]\n",
    "e = e_from_t_rh(t, rh)\n",
    "\n",
    "\n",
    "rh_derived = rh_from_t_td(t, t_d)\n",
    "rh_residual = (rh_derived - rh) ** 2\n",
    "\n",
    "r_derived = r_from_e_p(e, p)\n",
    "r_residual = (r_derived - r) ** 2\n",
    "\n",
    "physical_penalty = abs(rh_residual) / obs.relative_humidity.var() + abs(r_residual) / obs.water_vapor_mixing_ratio.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad380498",
   "metadata": {},
   "outputs": [],
   "source": [
    "dims = [\"forecast_reference_time\",\"t\",\"station\"]\n",
    "mae = np.abs(err).mean(dim=dims)\n",
    "mae = (mae / obs.std()).to_array().mean(\"variable\")\n",
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "ax1.plot(pred[\"loss.alpha\"].values, physical_penalty.mean(dim=dims), color=\"red\")\n",
    "ax1.tick_params(axis='x', rotation=90)\n",
    "ax1.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "ax1.tick_params(axis='y', labelcolor=\"red\")\n",
    "ax1.set_ylabel(r\"$\\mathcal{P}$\")\n",
    "ax2 = ax1.twinx()\n",
    "ax2.plot(pred[\"loss.alpha\"].values, mae.values, color=\"blue\")\n",
    "ax2.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "ax2.tick_params(axis='y', labelcolor=\"blue\")\n",
    "ax2.set_ylabel(r\"NMAE\")\n",
    "ax1.set_xlabel(r\"Physical penalty weight $\\alpha$\")\n",
    "ax1.set_ylim(0., 0.04)\n",
    "ax2.set_ylim(0.18, 0.22)\n",
    "plt.tight_layout()\n",
    "plt.savefig(out_dir / \"p_mae_vs_alpha.png\")"
   ]
  }
 ],
 "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.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "83068b732181b14f66262f02228ff7bb18ccf2d7f61d29cd0d6c11f3d6147031"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}