{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Analyzing Multi-Year Droughts\n",
                "\n",
                "After van Mourik, J. and Ruijsch, D. and van der Wiel, K. and Hazeleger, W. and Wanders, N. (2025) - [Regional drivers and characteristics of multi-year droughts](https://doi.org/10.1016/j.wace.2025.100748)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Load packages"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import matplotlib as mpl\n",
                "import matplotlib.pyplot as plt\n",
                "import pandas as pd\n",
                "\n",
                "import spei as si  # si for standardized index\n",
                "\n",
                "print(si.show_versions())"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Load meteo data"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "df = pd.read_csv(\"data/DEBILT.csv\", index_col=0, parse_dates=True)\n",
                "prec = df[\"Prec [m/d] 260_DEBILT\"].multiply(1e3).rename(\"prec\")\n",
                "evap = df[\"Evap [m/d] 260_DEBILT\"].multiply(1e3).rename(\"evap\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Calculate SPEI over different time scales\n",
                "\n",
                "1 month to 24 months"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "pe = (\n",
                "    (prec - evap).dropna().resample(\"ME\").sum()\n",
                ")  # calculate  monthly precipitation excess\n",
                "spei1 = si.spei(pe, timescale=1).rename(\"1\")\n",
                "spei3 = si.spei(pe, timescale=3).rename(\"3\")\n",
                "spei6 = si.spei(pe, timescale=6).rename(\"6\")\n",
                "spei9 = si.spei(pe, timescale=9).rename(\"9\")\n",
                "spei12 = si.spei(pe, timescale=12).rename(\"12\")\n",
                "spei24 = si.spei(pe, timescale=24).rename(\"24\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Visualization"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Plot as time series"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "f, ax = plt.subplots(6, figsize=(7.0, 10.0), sharex=True)\n",
                "si.plot.si(spei1, ax=ax[0], add_category=False)\n",
                "si.plot.si(spei3, ax=ax[1], add_category=False)\n",
                "si.plot.si(spei6, ax=ax[2], add_category=False)\n",
                "si.plot.si(spei9, ax=ax[3], add_category=False)\n",
                "si.plot.si(spei12, ax=ax[4], add_category=False)\n",
                "si.plot.si(spei24, ax=ax[5], add_category=False)\n",
                "ax[0].set_ylabel(\"SPEI-1\")\n",
                "ax[1].set_ylabel(\"SPEI-3\")\n",
                "ax[2].set_ylabel(\"SPEI-6\")\n",
                "ax[3].set_ylabel(\"SPEI-9\")\n",
                "ax[4].set_ylabel(\"SPEI-12\")\n",
                "ax[5].set_ylabel(\"SPEI-24\")\n",
                "ax[5].xaxis.set_major_locator(mpl.dates.YearLocator())\n",
                "ax[5].xaxis.set_minor_locator(mpl.dates.MonthLocator())\n",
                "ax[5].set_xlim(pd.Timestamp(\"2005\"), pd.Timestamp(\"2015\"))"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Plot as heatmap\n",
                "\n",
                "Only droughts"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "ax = si.plot.heatmap([spei1, spei3, spei6, spei9, spei12, spei24])\n",
                "ax.xaxis.set_major_locator(mpl.dates.YearLocator())\n",
                "ax.xaxis.set_minor_locator(mpl.dates.MonthLocator())\n",
                "ax.get_figure().axes[-1].set_yticks(range(-3, 0))\n",
                "ax.set_xlim(pd.Timestamp(\"2010\"), pd.Timestamp(\"2015\"))"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Or with wet periods"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "f, ax = plt.subplots(figsize=(8.0, 3.0))\n",
                "ax = si.plot.heatmap(\n",
                "    [spei1, spei3, spei6, spei9, spei12, spei24],\n",
                "    cmap=\"vik_r\",\n",
                "    vmax=3,\n",
                "    add_category=True,\n",
                "    ax=ax,\n",
                ")\n",
                "ax.xaxis.set_major_locator(mpl.dates.YearLocator())\n",
                "ax.xaxis.set_minor_locator(mpl.dates.MonthLocator())\n",
                "ax.set_xlim(pd.Timestamp(\"2010\"), pd.Timestamp(\"2015\"))"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "SPEI",
            "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.13.1"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}