notebooks/performance.py.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc934e7c",
   "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",
    "from scipy.stats import norm \n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "from itertools import combinations\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",
    "\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50f5b246",
   "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\"][\"default\"]\n",
    "print(json.dumps(exp_config, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26149a8e",
   "metadata": {},
   "source": [
    "# Data import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a4b0702",
   "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[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",
    "\n",
    "def process_nwp(ds, parameters):\n",
    "    for var in ds.data_vars:\n",
    "        prefix, name = var.split(\":\")\n",
    "        name = name.removesuffix(\"_ensavg\")\n",
    "        ds = ds.rename({var:name})\n",
    "    ds = ds[parameters]\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\n",
    "\n",
    "\n",
    "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",
    "\n",
    "\n",
    "reftimes = obs.forecast_reference_time.values\n",
    "stations = obs.station.values\n",
    "features = xr.open_zarr(\"/scratch/fzanetta/pcpp-workflow/data/01_raw/features.zarr\").sel(station=stations, forecast_reference_time=reftimes)\n",
    "nwp = process_nwp(features, TASKS)\n",
    "nwp = nwp.compute().chunk(\"auto\").persist()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba220961",
   "metadata": {},
   "source": [
    "# Plotting and tables functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aea1319",
   "metadata": {},
   "outputs": [],
   "source": [
    "def boxplots(ds, cfg, figsize=(6,6), ylim_mul = (1., 1.), yticks_inside=False, approaches=None, return_legend=False):\n",
    "    \n",
    "    df = (\n",
    "        ds_to_df(ds, \"metric\")\n",
    "        .replace(cfg[\"approach_names\"])\n",
    "        .replace(cfg[\"var_short_names\"])\n",
    "        .rename(columns={\"approach\":\"Approach\"})\n",
    "    )\n",
    "    \n",
    "    if approaches is None:\n",
    "        names = cfg[\"approach_names\"].values()\n",
    "    else:\n",
    "        names = [v for k, v in cfg[\"approach_names\"].items() if k in approaches]\n",
    "        \n",
    "    fig, axs = plt.subplots(1, len(TASKS), figsize=figsize, sharey=False)\n",
    "    for i, var in enumerate(TASKS):\n",
    "        var = cfg[\"var_short_names\"][var]\n",
    "        sns.boxplot(\n",
    "            data=df.query(f\"variable=='{var}'\"),\n",
    "            x=\"variable\",\n",
    "            y=\"metric\",\n",
    "            hue=\"Approach\",\n",
    "            hue_order=names,\n",
    "            palette=cfg[\"approach_colors\"],\n",
    "            ax=axs[i],\n",
    "            showfliers=False\n",
    "        )\n",
    "\n",
    "    for i, ax in enumerate(axs):\n",
    "        ax.legend().remove()\n",
    "        ax.set_ylabel(\"\")\n",
    "        ax.set_xlabel(\"\")\n",
    "        ymin, ymax = ax.get_ylim()\n",
    "        new_ylims = (ymin * ylim_mul[0], ymax * ylim_mul[1])\n",
    "        if TASKS[i] == \"relative_humidity\":\n",
    "            new_ylims = (new_ylims[0] * 0.9, new_ylims[1] * 1.1)\n",
    "        ax.set_ylim(*new_ylims)\n",
    "        if yticks_inside:\n",
    "            ymin_ = round(ymin * 0.4 + new_ylims[0] * 0.6, 2)\n",
    "            ymax_ = round(ymax * 0.4 + new_ylims[1] * 0.6, 2)\n",
    "            ax.tick_params(axis=\"y\",direction=\"in\", pad=-40)\n",
    "            ax.set_yticks([ymin_, ymax_], labels=[ymin_, ymax_])\n",
    "            \n",
    "\n",
    "#             ax.text(0., ymin_, str(round(ymin_,2)), ha=\"center\")\n",
    "#             ax.text(0., ymax_, str(round(ymax_,2)), ha=\"center\")\n",
    "        \n",
    "\n",
    "    lgd = axs[0].legend(\n",
    "        bbox_to_anchor=(0.075, 0.87, 0.88, 0.1),\n",
    "        loc=\"lower left\",\n",
    "        ncol=2,\n",
    "        mode=\"expand\",\n",
    "        borderaxespad=0.,\n",
    "        frameon=False,\n",
    "        # fontsize=9,\n",
    "        bbox_transform=plt.gcf().transFigure\n",
    "    )\n",
    "    \n",
    "    plt.subplots_adjust(wspace=1.5, hspace=0)\n",
    "    if return_legend:\n",
    "        return fig, axs, lgd\n",
    "    else:\n",
    "        return fig, axs\n",
    "\n",
    "def latex_table(results, cfg, scores=[\"mae\", \"msss\"], labels=[\"MAE\", \"MSSS\"]):\n",
    "\n",
    "    out = \"\"\n",
    "\n",
    "    # preamble\n",
    "    out += \"\\\\begin{table*}\\n\"\n",
    "    out += \"    \\caption{This is a table. Add something here.}\\n\"\n",
    "    out += \"    \\\\renewcommand{\\\\arraystretch}{1.1} \\n\"\n",
    "    out += \"    \\\\begin{tabular*}{\\hsize}{@{\\extracolsep\\\\fill}l lllll lllll@{}}\\n\"\n",
    "\n",
    "    # body\n",
    "    var_symbols = list(cfg[\"var_short_names\"].values())\n",
    "    out += \"        \\\\topline\\n\"\n",
    "    out += \"        &\" + \" & \".join([\"\\multicolumn{2}{c}\" + f\"{'{'}{v}{'}'}\" for v in var_symbols]) + \"\\\\\\ \\n\"\n",
    "    out += \"        \\cmidrule(lr){2-3} \\cmidrule(lr){4-5} \\cmidrule(lr){6-7} \\cmidrule(lr){8-9} \\cmidrule(lr){10-11} \\n\"\n",
    "    out += f\"        & {labels[0]} & {labels[1]} & {labels[0]} & {labels[1]} & {labels[0]} & {labels[1]} & {labels[0]} & {labels[1]} & {labels[0]} & {labels[1]} \\\\\\ \\n\"\n",
    "    for approach, name in cfg[\"approach_names\"].items():\n",
    "        row = \"        \"\n",
    "        row += f\"{name} & \" + \" & \".join(\n",
    "            [f\"{results['mae'][var][approach]:.3f} & {results['msss'][var][approach]:.3f}\"\n",
    "             for var in results[\"mae\"].keys()]\n",
    "        )\n",
    "        row += \" \\\\\\ \\n\"\n",
    "        out += row\n",
    "    out += \"        \\\\\\ \\n\"\n",
    "    \n",
    "    # postamble\n",
    "    out += \"   \\end{tabular*}\\n\"\n",
    "    out += \"\\end{table*}\"\n",
    "    return out\n",
    "\n",
    "def latex_table_counts(results, cfg):\n",
    "\n",
    "    out = \"\"\n",
    "\n",
    "    # preamble\n",
    "    out += \"\\\\begin{table*}\\n\"\n",
    "    out += \"    \\caption{This is a table. Add something here.}\\n\"\n",
    "    out += \"    \\\\renewcommand{\\\\arraystretch}{1.1} \\n\"\n",
    "    out += \"    \\\\begin{tabular*}{\\hsize}{@{\\extracolsep\\\\fill}l lllll@{}}\\n\"\n",
    "\n",
    "    # body\n",
    "    var_symbols = list(cfg[\"var_short_names\"].values())\n",
    "    out += \"        \\\\topline\\n\"\n",
    "    out += \"        & \" + \" & \".join([f\"{v}\" for v in var_symbols]) + \"\\\\\\ \\n\"\n",
    "    out += \"        \\\\midline\\n\"\n",
    "    for approach, name in cfg[\"approach_names\"].items():\n",
    "        row = \"        \"\n",
    "        row += f\"{name} & \" + \" & \".join(\n",
    "            [f\"{results[var][approach]}\"\n",
    "             for var in results.keys()]\n",
    "        )\n",
    "        row += \" \\\\\\ \\n\"\n",
    "        out += row\n",
    "    out += \"        \\\\\\ \\n\"\n",
    "    \n",
    "    # postamble\n",
    "    out += \"   \\end{tabular*}\\n\"\n",
    "    out += \"\\end{table*}\"\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec4d5e38",
   "metadata": {},
   "source": [
    "# Metrics "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "176c8857",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mean_absolute_error(pred, obs, reduce_dims):\n",
    "    return abs(pred - obs).mean(reduce_dims)\n",
    "\n",
    "def mean_squared_error(pred, obs, reduce_dims):\n",
    "    return ((pred - obs) ** 2).mean(reduce_dims)\n",
    "\n",
    "def mean_squared_skill_score(pred, ref, obs, reduce_dims):\n",
    "    mse_pred = ((pred - obs) ** 2).mean(reduce_dims)\n",
    "    mse_ref  = ((ref - obs) ** 2).mean(reduce_dims)\n",
    "    return 1 - mse_pred / mse_ref\n",
    "\n",
    "def r_squared(pred, obs, reduce_dims):\n",
    "    res = pred - obs\n",
    "    ss_tot = ((obs.mean(reduce_dims) - obs) ** 2).sum(reduce_dims)\n",
    "    ss_res = (res ** 2).sum(reduce_dims)\n",
    "    return 1 - ss_res / ss_tot\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "773aae18",
   "metadata": {},
   "source": [
    "# Boxplots\n",
    "\n",
    "These plots show the performance metrics averaged over the forecast reference time, station and leadtime. The distributions represent the variability due to cross-validation split and random seed.\n",
    "\n",
    "In the case of the MSSS and the R$^2$, we first compute the metrics for each station individually and take the average afterwards. This way we remove the variability related to a station's climatological mean. If we would not do that, we would have large values (close to 1), but it would not be a fair evaluation.\n",
    "\n",
    "### Mean Absolute Error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9277451d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mae = mean_absolute_error(pred, obs, [\"forecast_reference_time\", \"station\",\"t\"])\n",
    "fig, axs = boxplots(mae, plt_cfg, figsize=(6,5), ylim_mul=(0.97, 1.03), yticks_inside=True)\n",
    "axs[0].set_ylabel(\"Mean absolute error\")\n",
    "plt.tight_layout()\n",
    "plt.subplots_adjust(wspace=0.1)\n",
    "plt.savefig(out_dir / \"mae_boxplots.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11b5cccd",
   "metadata": {},
   "source": [
    "### Mean Squared Skill Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99816450",
   "metadata": {},
   "outputs": [],
   "source": [
    "msss = mean_squared_skill_score(pred, nwp, obs, [\"forecast_reference_time\", \"t\"]).mean(\"station\")\n",
    "fig, axs = boxplots(msss, plt_cfg, figsize=(6,5), ylim_mul=(0.9, 1.1), yticks_inside=True)\n",
    "axs[0].set_ylabel(\"Mean Squared Skill Score\")\n",
    "plt.tight_layout()\n",
    "plt.subplots_adjust(wspace=0.1)\n",
    "plt.savefig(out_dir / \"msss_boxplots.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3988b38",
   "metadata": {},
   "source": [
    "### Coefficient of determination"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "451c2e88",
   "metadata": {},
   "outputs": [],
   "source": [
    "r2 = r_squared(pred, obs, [\"forecast_reference_time\", \"t\"]).mean(\"station\")\n",
    "fig, axs = boxplots(r2, plt_cfg, figsize=(6,5), ylim_mul=(0.99, 1.015), yticks_inside=True)\n",
    "axs[0].set_ylabel(\"Coefficient of determination\")\n",
    "plt.tight_layout()\n",
    "plt.subplots_adjust(wspace=0.1)\n",
    "plt.savefig(out_dir / \"r_squared_boxplots.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8734a15c",
   "metadata": {},
   "source": [
    "# Scatter plots\n",
    "\n",
    "Here we plot the NWP baseline performance metric on the x-axis and the postprocessed predictions on the y-axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7933b77",
   "metadata": {},
   "outputs": [],
   "source": [
    "def station_metrics_scatterplot(err_pp, err_nwp):\n",
    "    \n",
    "    fig, axs = plt.subplots(2,3, figsize=(15,10), layout=\"constrained\")\n",
    "    for var, ax in zip(TASKS, axs.ravel()[:-1]):\n",
    "    #     var = \"relative_humidity\"\n",
    "    #     fig, ax = plt.subplots()\n",
    "        for i, approach in enumerate(exp_config[\"approaches\"]):\n",
    "            ax.scatter(\n",
    "                err_nwp[var].values,\n",
    "                err_pp[var].sel(approach=approach).values,\n",
    "                c=plt_cfg[\"approach_colors\"][i],\n",
    "                s=12,\n",
    "                alpha=0.7,\n",
    "                label=plt_cfg[\"approach_names\"][approach]\n",
    "            )\n",
    "        xlim = ax.get_xlim()\n",
    "        ylim = ax.get_ylim()\n",
    "\n",
    "        lim = (min(xlim[0], ylim[0]), max(xlim[1], ylim[1]))\n",
    "        ax.axline((np.mean(lim),np.mean(lim)), slope=1, linestyle=\"--\", c=\"k\", linewidth=1)\n",
    "        ax.set(\n",
    "            ylim=lim, xlim=lim, \n",
    "            title=plt_cfg[\"var_long_names\"][var],\n",
    "            aspect=1\n",
    "        )\n",
    "    handles, labels = ax.get_legend_handles_labels()\n",
    "    fig.legend(handles, labels, bbox_to_anchor=(0.95, 0.3))\n",
    "    # leg = fig.legend(loc=\"lower right\")\n",
    "    axs[1,2].remove()\n",
    "    \n",
    "    return fig, axs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e74af08",
   "metadata": {},
   "source": [
    "### Mean Absolute Error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "945c7d39",
   "metadata": {},
   "outputs": [],
   "source": [
    "mae_nwp = mean_absolute_error(nwp, obs, [\"forecast_reference_time\",\"seed\",\"split\",\"t\"]).mean([\"approach\"])\n",
    "mae_pp = mean_absolute_error(pred, obs, [\"forecast_reference_time\",\"seed\",\"split\",\"t\"])\n",
    "\n",
    "fig, axs = station_metrics_scatterplot(mae_pp, mae_nwp)\n",
    "for var, ax in zip(TASKS, axs.ravel()):\n",
    "    ax.set(\n",
    "        ylabel=f\"MAE [{plt_cfg['var_units'][var]}] of postprocessing\", \n",
    "        xlabel=f\"MAE [{plt_cfg['var_units'][var]}] of NWP\",\n",
    "    )\n",
    "    \n",
    "plt.savefig(out_dir / \"mae_scatterplot.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9dcbf70",
   "metadata": {},
   "source": [
    "### Mean squared error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc2f1b3",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "mae_nwp = mean_squared_error(nwp, obs, [\"forecast_reference_time\",\"seed\",\"split\",\"t\"]).mean([\"approach\"])\n",
    "mae_pp = mean_squared_error(pred, obs, [\"forecast_reference_time\",\"seed\",\"split\",\"t\"])\n",
    "\n",
    "fig, axs = station_metrics_scatterplot(mae_pp, mae_nwp)\n",
    "for var, ax in zip(TASKS, axs.ravel()):\n",
    "    ax.set(\n",
    "        ylabel=f\"MSE [{plt_cfg['var_units'][var]}]$^2$ of postprocessing\", \n",
    "        xlabel=f\"MSE [{plt_cfg['var_units'][var]}]$^2$ of NWP\",\n",
    "    )\n",
    "    \n",
    "plt.savefig(out_dir / \"mse_scatterplot.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29fe12ed",
   "metadata": {},
   "source": [
    "# Overall results & tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63092118",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mae\n",
    "mae = mean_absolute_error(pred, obs, [\"forecast_reference_time\",\"seed\", \"split\",\"t\"]).mean(\"station\")\n",
    "mae.to_dataframe().to_json(out_dir / \"mae_results.json\", indent=4)\n",
    "\n",
    "    \n",
    "# msss\n",
    "msss = mean_squared_skill_score(pred, nwp, obs, [\"forecast_reference_time\",\"seed\", \"split\", \"t\"]).mean(\"station\")\n",
    "msss.to_dataframe().to_json(out_dir / \"msss_results.json\", indent=4)\n",
    "\n",
    "\n",
    "# r squared\n",
    "r2 = r_squared(pred, obs, [\"forecast_reference_time\",\"seed\", \"split\", \"t\"]).mean(\"station\")\n",
    "r2.to_dataframe().to_json(out_dir / \"r_squared_results.json\", indent=4)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76259d70",
   "metadata": {},
   "source": [
    "### Main results table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12766314",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = latex_table({\"mae\":mae.to_dataframe(), \"msss\":msss.to_dataframe()}, plt_cfg)\n",
    "with open(out_dir / \"latex_table_results.txt\", \"w\") as f:\n",
    "    f.write(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6093480f",
   "metadata": {},
   "source": [
    "## Significance testing\n",
    "We perform tests of equal predictive performance using the Diebold-Mariano test, and apply the Benjamin-Hochberg procedure to account for the fals discovery rate in multiple testing. We apply tests between each pair of models, for every combination of leadtime and station. The loss vectors are first averaged over the cross-validation split and the random seeds. \n",
    "We assume independence between forecast errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbd5a912",
   "metadata": {},
   "outputs": [],
   "source": [
    "def diebold_mariano_test(l1, l2):\n",
    "    \"\"\"\n",
    "    Diebold-Mariano (DM) test, assuming independence between\n",
    "    loss values (as in Schulz and Lerch (2022)). \n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    l1: np.ndarray\n",
    "        Loss of the first forecast.\n",
    "    l2: np.ndarray\n",
    "        Loss of the second forecast.\n",
    "\n",
    "    Returns:\n",
    "    -------\n",
    "    dm_stat: the DM test statistic \n",
    "    p_value: the p-value\n",
    "    \n",
    "    \"\"\"\n",
    "    D = l1 - l2\n",
    "    n = D.shape[0]\n",
    "    D_mean = np.mean(D)\n",
    "    D_var = np.var(D)\n",
    "    se = np.sqrt(D_var / n)\n",
    "    dm_stat = D_mean / se\n",
    "    p_value = 2 * (1 - norm.cdf(np.abs(dm_stat)))\n",
    "    return dm_stat, p_value\n",
    "\n",
    "\n",
    "def benjamin_hochberg_correction(p_values, q=0.05):\n",
    "    \"\"\"\n",
    "    Returns significant p-values after applying Benjamin-Hochberg correction.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    p_values: np.ndarray\n",
    "        The array of p-values.\n",
    "    q: float, optional\n",
    "        The significance level for the correction.\n",
    "    \"\"\"\n",
    "    sorted_p_values = np.sort(p_values)\n",
    "    m = len(p_values)\n",
    "    adjusted_p_values = sorted_p_values * m / np.arange(1, m+1)\n",
    "    k = np.max(np.where(adjusted_p_values <= np.arange(1, m+1) / m * q))\n",
    "    significant_p_values = p_values <= sorted_p_values[k]\n",
    "    return significant_p_values\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "360e7eee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predictive_performance_test(errors: xr.DataArray):\n",
    "    \"\"\"\n",
    "    For a given target variable, perform Diebold-Mariano tests between each pair\n",
    "    of settings, for every combination of leadime and stations.\n",
    "\n",
    "    Parameters:\n",
    "    -----------\n",
    "    errors: xr.DataArray\n",
    "        The errors for the given target variable.\n",
    "    \n",
    "    Returns:\n",
    "    --------\n",
    "    results: pd.DataFrame\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    results = np.zeros((4,4)).astype(int) * np.nan \n",
    "    # loop for pairwise model comparisons\n",
    "    for a1, a2 in list(combinations(exp_config[\"approaches\"], 2)):\n",
    "        res = []\n",
    "        n = 0\n",
    "        # loop over stations and leadtimes\n",
    "        for station in errors.station.values:\n",
    "            for leadtime in errors.t.values:\n",
    "                l1 = errors.sel(approach=a1, station=station, t=leadtime).dropna(\"forecast_reference_time\")\n",
    "                l2 = errors.sel(approach=a2, station=station, t=leadtime).dropna(\"forecast_reference_time\")\n",
    "\n",
    "                # diebold-mariano test\n",
    "                dm, p = diebold_mariano_test(l1.values, l2.values)\n",
    "                res.append((station, leadtime, dm, p))\n",
    "                n += 1\n",
    "\n",
    "        # aggregate results\n",
    "        res = pd.DataFrame(res, columns=[\"station\",\"leadtime\", \"DM\",\"p\"]).set_index(\"station\")\n",
    "\n",
    "        # table index\n",
    "        a1_idx = exp_config[\"approaches\"].index(a1)\n",
    "        a2_idx = exp_config[\"approaches\"].index(a2)\n",
    "\n",
    "        # forecasts of model 1 significantly better than 2\n",
    "        results[a1_idx, a2_idx] = int(((res[\"DM\"] < 0.) & benjamin_hochberg_correction(res[\"p\"])).sum()) / n * 100\n",
    "\n",
    "        # forecasts of model 2 significantly better than 1\n",
    "        results[a2_idx, a1_idx] = int(((res[\"DM\"] > 0.) & benjamin_hochberg_correction(res[\"p\"])).sum()) / n * 100\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efb9fa02",
   "metadata": {},
   "outputs": [],
   "source": [
    "errors = abs(pred - obs).mean([\"split\",\"seed\"]).compute()\n",
    "\n",
    "dfs = []\n",
    "for var in TASKS:\n",
    "    results = predictive_performance_test(errors[var])\n",
    "    df = pd.DataFrame(results, columns=exp_config[\"approaches\"], index=exp_config[\"approaches\"])\n",
    "    df[\"Winning average\"] = df.mean(axis=1)\n",
    "    df.loc[\"Losing average\"] = (*df.mean(axis=0).values[:-1], np.nan)\n",
    "    dfs.append(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bbda244",
   "metadata": {},
   "outputs": [],
   "source": [
    "def latex_table_tests(results, cfg):\n",
    "\n",
    "    out = \"\"\n",
    "\n",
    "    # preamble\n",
    "    out += \"\\\\begin{table*}\\n\"\n",
    "    out += \"    \\caption{This is a table. Add something here.}\\n\"\n",
    "    out += \"    \\\\renewcommand{\\\\arraystretch}{1.} \\n\"\n",
    "    out += \"    \\\\begin{tabular*}{\\hsize}{@{\\extracolsep\\\\fill}l llll l@{}}\\n\"\n",
    "\n",
    "    # body\n",
    "    colnames = list(cfg[\"approach_names\"].values()) + [\"Winning average\"]\n",
    "    rownames = list(cfg[\"approach_names\"].values()) + [\"Losing average\"]\n",
    "    out += \"        \\\\topline\\n\"\n",
    "    out += \"        & \" + \" & \".join([f\"{v}\" for v in colnames]) + \"\\\\\\ \\n\"\n",
    "    out += \"        \\\\midline\\n\"\n",
    "    for name in results.index.values:\n",
    "        f = lambda v: \"\" if np.isnan(v) else f\"{v:.2f}\"\n",
    "        row = \"        \"\n",
    "        if name == \"Losing average\":\n",
    "            row += \"\\\\midline\\n        \"\n",
    "        row += f\"{cfg['approach_names'].get(name, name)} & \" + \" & \".join([f(results.loc[name, col]) for col in results.columns])\n",
    "        row += \" \\\\\\ \\n\"\n",
    "        out += row\n",
    "    out += \"        \\\\\\ \\n\"\n",
    "    \n",
    "    # postamble\n",
    "    out += \"   \\end{tabular*}\\n\"\n",
    "    out += \"\\end{table*}\"\n",
    "    return out\n",
    "\n",
    "for var, df in zip(TASKS, dfs):\n",
    "    out = latex_table_tests(df, plt_cfg)\n",
    "    with open(out_dir / f\"signficance_table_{var}.txt\", \"w\") as f:\n",
    "        f.write(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f90ce0ea",
   "metadata": {},
   "source": [
    "## Extra: suspect scores\n",
    "One might reasonably believe that the skill scores (MSSS) presented above are suspiciously high, expecially for pressure. The reason for these large values is that the errors of the NWP model are indeed very large: this is typical in situations of complex topography (such as Switzerland) where quantities related to elevation often have very large biases due to the difference between the model elevation and the true elevation.\n",
    "\n",
    "In the code below we show that the mean bias of the NWP forecast at certain stations varies between -80 and +87 hPa."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24bd8cab",
   "metadata": {},
   "outputs": [],
   "source": [
    "nwp_err = nwp - obs\n",
    "nwp_err[\"surface_air_pressure\"].mean([\"seed\",\"split\",\"forecast_reference_time\",\"t\",\"approach\"]).compute().astype(int)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90ff406d",
   "metadata": {},
   "source": [
    "This bias is easily corrected by the postprocessing models, and gets close to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "165deec8",
   "metadata": {},
   "outputs": [],
   "source": [
    "err = pred - obs\n",
    "err[\"surface_air_pressure\"].mean([\"seed\",\"split\",\"forecast_reference_time\",\"t\", \"approach\"]).compute().astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fec9af6",
   "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.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "83068b732181b14f66262f02228ff7bb18ccf2d7f61d29cd0d6c11f3d6147031"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}