notebooks/physical_consistency.py.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b05518fd",
   "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",
    "import datashader as dsh\n",
    "import datashader.transfer_functions as dtf\n",
    "from datashader.mpl_ext import dsshow\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",
    "TASKS = [\n",
    "    \"air_temperature\",\n",
    "    \"dew_point_temperature\",\n",
    "    \"surface_air_pressure\",\n",
    "    \"relative_humidity\",\n",
    "    \"water_vapor_mixing_ratio\",\n",
    "]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "844be57f",
   "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\"]\n",
    "try:\n",
    "    approaches.remove(\"offline_constrained\")\n",
    "except:\n",
    "    pass\n",
    "if \"architecture_constrained\" in approaches:\n",
    "    approaches.remove(\"architecture_constrained\")\n",
    "    approaches.append(\"architecture_constrained\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dd6e658",
   "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": null,
   "id": "39f19f75",
   "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": null,
   "id": "365c50b6",
   "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",
    "def e_from_td(t_d, t):\n",
    "    e = xr.where(\n",
    "        t >= 0,\n",
    "        6.107 * np.exp((17.368 * t_d) / (238.83 + t_d)),\n",
    "        6.108 * np.exp((17.856 * t_d) / (245.52 + t_d)),\n",
    "    )\n",
    "    return e   \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))\n",
    "\n",
    "def r_from_p_td(p, t_d, t):\n",
    "\n",
    "    e = e_from_td(t_d, t)\n",
    "\n",
    "    return r_from_e_p(e, p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "396d6bea",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_categories(ds, var, bins, labels):\n",
    "    assert len(bins) + 1 == len(labels)\n",
    "    da = ds[var]\n",
    "    bins = (-np.inf, *bins, np.inf)\n",
    "    \n",
    "    cat = np.empty_like(da, dtype=\"U100\")\n",
    "    for i in range(len(bins)-1):\n",
    "        print(bins[i], bins[i+1])\n",
    "        sel = (bins[i] <= da) & (da < bins[i+1])\n",
    "        cat[sel] = labels[i]\n",
    "    \n",
    "    cat[cat == \"\"] = labels[2]\n",
    "    cat = xr.DataArray(cat, coords=da.coords)    \n",
    "    return cat\n",
    "\n",
    "def shade_plot(df, ax, **kwargs):\n",
    "\n",
    "    agg = kwargs.pop(\"aggregator\", dsh.count())\n",
    "    norm = kwargs.pop(\"norm\", \"eq_hist\")\n",
    "    \n",
    "    artist = dsshow(\n",
    "        df,\n",
    "        dsh.Point(x,\"y\"),\n",
    "        aggregator=agg,\n",
    "        norm=norm,\n",
    "        ax=ax,\n",
    "        aspect=\"auto\",\n",
    "        **kwargs\n",
    "    )\n",
    "    return artist\n",
    "\n",
    "def RH_deviation(ds):\n",
    "    ds = ds.reset_coords(drop=True)\n",
    "    t = ds[\"air_temperature\"]\n",
    "    t_d = ds[\"dew_point_temperature\"]\n",
    "    rh = ds[\"relative_humidity\"]\n",
    "    rh_derived = rh_from_t_td(t, t_d) \n",
    "    y = rh - rh_derived\n",
    "    return y\n",
    "\n",
    "def r_deviation(ds):\n",
    "    ds = ds.reset_coords(drop=True)\n",
    "    t = ds[\"air_temperature\"]\n",
    "    t_d = ds[\"dew_point_temperature\"]\n",
    "    p = ds[\"surface_air_pressure\"]\n",
    "    r = ds[\"water_vapor_mixing_ratio\"]\n",
    "    e = e_from_td(t_d, t)\n",
    "    r_derived = r_from_e_p(e, p)\n",
    "    y =  r - r_derived\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46039aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = \"water_vapor_mixing_ratio\"\n",
    "\n",
    "fig, axs = plt.subplots(1, len(approaches), figsize=(6.5, 5), sharey=True, sharex=True)\n",
    "for i, approach in enumerate(approaches):\n",
    "    ds = pred.sel(approach=approach).reset_coords(drop=True)\n",
    "    y = r_deviation(ds)\n",
    "    df = xr.Dataset({x:ds[x], \"y\":y}).to_dataframe()\n",
    "    artist = shade_plot(df, axs[i], norm=\"eq_hist\")\n",
    "    if approach == \"unconstrained\":\n",
    "        xmin, xmax, ymin, ymax = artist.get_extent()\n",
    "    axs[i].set_title(plt_cfg[\"approach_names\"][approach], fontsize=10)\n",
    "    axs[i].set_xlabel(r\"$r$ [g $kg^{-1}$]\")#, fontsize=13)\n",
    "axs[0].set_ylabel(r\"$r - g(P, T_d)$\")#, fontsize=13)\n",
    "axs[0].set_ylim(ymin - 2, ymax + 2)\n",
    "axs[0].set_xlim(xmin - 2, xmax + 5)\n",
    "plt.tight_layout()\n",
    "plt.savefig(out_dir / \"mixing_ratio_deviations.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fde675c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = \"relative_humidity\"\n",
    "\n",
    "fig, axs = plt.subplots(1, len(approaches), figsize=(6.5, 5), sharey=True, sharex=True)\n",
    "for i, approach in enumerate(approaches):\n",
    "    ds = pred.sel(approach=approach).reset_coords(drop=True)\n",
    "    y = RH_deviation(ds)\n",
    "    df = xr.Dataset({x:ds[x], \"y\":y}).to_dataframe()\n",
    "    artist = shade_plot(df, axs[i])\n",
    "    if approach == \"unconstrained\":\n",
    "        xmin, xmax, ymin, ymax = artist.get_extent()\n",
    "    axs[i].set_title(plt_cfg[\"approach_names\"][approach], fontsize=10)    \n",
    "    axs[i].set_xlabel(r\"$RH$ [%]\")#, fontsize=13)\n",
    "axs[0].set_ylabel(r\"$RH - f(T, T_d)$\")#, fontsize=13)\n",
    "axs[0].set_ylim(-32, 20)\n",
    "axs[0].set_xlim(xmin - 2, xmax + 2)\n",
    "plt.tight_layout()\n",
    "plt.savefig(out_dir / \"relative_humidity_deviations.png\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.5 ('dlpp')",
   "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.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "83068b732181b14f66262f02228ff7bb18ccf2d7f61d29cd0d6c11f3d6147031"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}