{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AR(1) noise model with irregular synthetic time series\n",
    "*R.A. Collenteur, University of Graz, May 2020*\n",
    "\n",
    "In this notebook the classical Autoregressive AR(1) noise model is tested for Pastas models. This noise model is tested against synthetic data generated with NumPy or Statsmodels' ARMA model. This noise model is tested on head time series with regular and irregular time steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import pastas as ps\n",
    "\n",
    "ps.set_log_level(\"ERROR\")\n",
    "ps.show_versions()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Develop the AR(1) Noise Model for Pastas\n",
    "\n",
    "The following formula is used to calculate the noise according to the ARMA(1,1) process:\n",
    "\n",
    "$$\\upsilon(t_i) = r(t_i) - r(t_{i-1}) \\text{e}^{-\\Delta t_i / \\alpha}$$\n",
    "\n",
    "where $\\upsilon$ is the noise, $\\Delta t_i$ the time step between the residuals ($r$), and $\\alpha$ [days] the AR parameter of the model. The model is named `NoiseModel` and can be found in `noisemodel.py`. It can be added to a Pastas model as follows: `ml.add_noisemodel(ps.ArNoiseModel())`\n",
    "\n",
    "## 2. Generate synthetic head time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in some data\n",
    "rain = (\n",
    "    pd.read_csv(\"../examples/data/rain_260.csv\", index_col=0, parse_dates=[0]).squeeze()\n",
    "    / 1000\n",
    ")\n",
    "\n",
    "# Set the True parameters\n",
    "Atrue = 800\n",
    "ntrue = 1.4\n",
    "atrue = 200\n",
    "dtrue = 20\n",
    "\n",
    "# Generate the head\n",
    "step = ps.Gamma().block([Atrue, ntrue, atrue], cutoff=0.9999)\n",
    "h = dtrue * np.ones(len(rain) + step.size)\n",
    "for i in range(len(rain)):\n",
    "    h[i : i + step.size] += rain.iloc[i] * step\n",
    "head = pd.DataFrame(\n",
    "    index=rain.index,\n",
    "    data=h[: len(rain)],\n",
    ")\n",
    "head = head[\"1990\":\"2015\"]\n",
    "\n",
    "# Plot the head without noise\n",
    "plt.figure(figsize=(10, 2))\n",
    "plt.plot(head, \"k.\", label=\"head\")\n",
    "plt.legend(loc=0)\n",
    "plt.ylabel(\"Head (m)\")\n",
    "plt.xlabel(\"Time (years)\");"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Generate AR(1) noise and add it to the synthetic heads\n",
    "In the following code-block, noise is generated using an AR(1) process using Numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reproduction of random numbers\n",
    "np.random.seed(1234)\n",
    "alpha = 0.8\n",
    "\n",
    "# generate samples using NumPy\n",
    "random_seed = np.random.RandomState(1234)\n",
    "noise = random_seed.normal(0, 1, len(head)) * np.std(head.values) * 0.2\n",
    "a = np.zeros_like(head[0])\n",
    "\n",
    "for i in range(1, noise.size):\n",
    "    a[i] = noise[i] + a[i - 1] * alpha\n",
    "\n",
    "head_noise = head[0] + a"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Create and solve a Pastas Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = ps.Model(head_noise)\n",
    "sm = ps.StressModel(rain, ps.Gamma(), name=\"recharge\", settings=\"prec\")\n",
    "ml.add_stressmodel(sm)\n",
    "ml.add_noisemodel(ps.ArNoiseModel())\n",
    "\n",
    "ml.solve(tmin=\"1991\", tmax=\"2015-06-29\", report=False)\n",
    "\n",
    "# Plot the results\n",
    "axes = ml.plots.results(figsize=(10, 5))\n",
    "axes[-2].plot(ps.Gamma().step([Atrue, ntrue, atrue], cutoff=0.999));"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Did we find back the original AR parameter?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.exp(-1.0 / ml.parameters.loc[\"noise_alpha\", \"optimal\"]).round(2), \"vs\", alpha)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The estimated parameters for the noise model are almost the same as the true parameters, showing that the model works for regular time steps.\n",
    "\n",
    "## 6. So is the autocorrelation removed correctly?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.plots.diagnostics(figsize=(10, 4));"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That seems okay. It is important to understand that this noisemodel will only help in removing autocorrelations at the first time lag, but not at larger time lags."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Test the NoiseModel for irregular time steps\n",
    "In this final step the Ar(1) noisemodel is tested for irregular timesteps, using the indices from a real groundwater level time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = (\n",
    "    pd.read_csv(\"../examples/data/test_index.csv\", parse_dates=True, index_col=0)\n",
    "    .index.round(\"D\")\n",
    "    .drop_duplicates()\n",
    ")\n",
    "head_irregular = head_noise.reindex(index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml = ps.Model(head_irregular)\n",
    "sm = ps.StressModel(rain, ps.Gamma(cutoff=0.99), name=\"recharge\", settings=\"prec\")\n",
    "ml.add_stressmodel(sm)\n",
    "ml.add_noisemodel(ps.ArNoiseModel())\n",
    "\n",
    "# ml.set_parameter(\"noise_alpha\", 9.5, vary=False)\n",
    "# ml.set_parameter(\"recharge_A\", 800, vary=False)\n",
    "# ml.set_parameter(\"recharge_n\", 1.2, vary=False)\n",
    "# ml.set_parameter(\"recharge_a\", 200, vary=False)\n",
    "\n",
    "ml.solve(tmin=\"1991\", tmax=\"2015-06-29\", report=False)\n",
    "axes = ml.plots.results(figsize=(10, 5))\n",
    "axes[-2].plot(ps.Gamma().step([Atrue, ntrue, atrue]))\n",
    "\n",
    "print(np.exp(-1.0 / ml.parameters.loc[\"noise_alpha\", \"optimal\"]).round(2), \"vs\", alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.plots.diagnostics(figsize=(10, 4), acf_options={\"bin_width\": 0.5});"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This autocorrelation plot looks good too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = ml.noise()\n",
    "print(\"\", v.loc[:\"2010-01-01\"].std())\n",
    "print(\"\", v.loc[\"2010-01-01\":].std())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pastas",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}