{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "eaa846c5",
            "metadata": {},
            "source": [
                "# Threshold Drought\n",
                "\n",
                "## Load packages"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "247f31ac",
            "metadata": {},
            "outputs": [],
            "source": [
                "import pandas as pd\n",
                "from scipy import stats as sps\n",
                "\n",
                "import spei as si"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "7b1fa2a9",
            "metadata": {},
            "source": [
                "## Load data"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "34ae712e",
            "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\")\n",
                "head = df[\"Head [m] B32C0572_DEBILT\"].rename(\"head\").dropna()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5830c471",
            "metadata": {},
            "source": [
                "## Calculate precipitation surplus"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "f750f7ae",
            "metadata": {},
            "outputs": [],
            "source": [
                "surplusd = prec - evap\n",
                "surplus = surplusd.resample(\"MS\").sum()\n",
                "surplus.plot()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "f013e9f1",
            "metadata": {},
            "source": [
                "## Fit distribution"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "6ae83c3e",
            "metadata": {},
            "outputs": [],
            "source": [
                "dist = sps.fisk\n",
                "sispei = si.SI(\n",
                "    series=surplus,\n",
                "    dist=dist,\n",
                "    timescale=0,\n",
                ")\n",
                "sispei.fit_distribution()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "41ab4ace",
            "metadata": {},
            "source": [
                "## Get threshold\n",
                "\n",
                "Choose arbitrary threshold based on quantile of the distribution. Can be any threshold the user wants as well. Only then the threshold time series has to be created manually."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "62530ad6",
            "metadata": {},
            "outputs": [],
            "source": [
                "speithr = sispei.ppf(0.3)  # 30% quantile threshold"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "efed2416",
            "metadata": {},
            "source": [
                "## Plot threshold"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "c7492e76",
            "metadata": {},
            "outputs": [],
            "source": [
                "ax = si.plot.threshold(\n",
                "    series=sispei.series,\n",
                "    threshold=speithr,\n",
                "    fill_color=\"red\",\n",
                ")\n",
                "_ = ax.set_xlim(pd.Timestamp(\"2010\"), pd.Timestamp(\"2020\"))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "b54b1d3f",
            "metadata": {},
            "source": [
                "## Repeat for head time series"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "1a7f3565",
            "metadata": {},
            "outputs": [],
            "source": [
                "timescale = 6\n",
                "sisgi = si.SI(\n",
                "    head,\n",
                "    dist=sps.norm,\n",
                "    timescale=timescale,\n",
                "    fit_freq=\"MS\",\n",
                "    normal_scores_transform=True,\n",
                "    agg_func=\"mean\",\n",
                ")\n",
                "sgithr = sisgi.ppf(0.4)  # choose arbitrary threshold"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "6cdf67c7",
            "metadata": {},
            "outputs": [],
            "source": [
                "ax = si.plot.threshold(\n",
                "    series=head.iloc[timescale - 1 :],\n",
                "    threshold=sgithr,\n",
                "    fill_color=\"red\",\n",
                ")\n",
                "_ = ax.set_xlim(pd.Timestamp(\"2010\"), pd.Timestamp(\"2020\"))"
            ]
        }
    ],
    "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.12.3"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}