{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ARMA(1,1) Noise Model for Pastas\n",
    "*R.A. Collenteur, University of Graz, May 2020*\n",
    "\n",
    "In this notebook an Autoregressive-Moving-Average (ARMA(1,1)) noise model is developed for Pastas models. This new noise model is tested against synthetic data generated with Numpy or Statsmodels' ARMA model. This noise model is tested on head time series with a regular time step.\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "<b>Warning</b>\n",
    "\n",
    "It should be noted that the time step may be non-equidistant in this formulation, but this model is not yet tested for irregular time steps.\n",
    "   \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.special import gammainc, gammaincinv\n",
    "\n",
    "import pastas as ps\n",
    "\n",
    "ps.set_log_level(\"ERROR\")\n",
    "ps.show_versions(numba=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Develop the ARMA(1,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} - \\upsilon(t_{i-1}) \\text{e}^{-\\Delta t_i / \\beta}$$\n",
    "\n",
    "where $\\upsilon$ is the noise, $\\Delta t_i$ the time step between the residuals ($r$), and respectively $\\alpha$ [days] and $\\beta$ [days] the parameters of the AR and MA parts of the model. The model is named `ArmaModel` and can be found in `noisemodel.py`. It is added to a Pastas model as follows: `ml.add_noisemodel(ps.ArmaModel())`\n",
    "\n",
    "## 2. Generate synthetic head time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in some data\n",
    "rain = pd.read_csv('../examples/data/rain_260.csv', index_col=0, parse_dates=[0]).squeeze() / 1000\n",
    "\n",
    "# Set the True parameters\n",
    "Atrue = 800\n",
    "ntrue = 1.1\n",
    "atrue = 200\n",
    "dtrue = 20\n",
    "\n",
    "# Generate the head\n",
    "step = ps.Gamma().block([Atrue, ntrue, atrue])\n",
    "h = dtrue * np.ones(len(rain) + step.size)\n",
    "for i in range(len(rain)):\n",
    "    h[i:i + step.size] += rain[i] * step\n",
    "head = pd.DataFrame(index=rain.index, data=h[:len(rain)],)\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)');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Generate ARMA(1,1) noise and add it to the synthetic heads\n",
    "In the following code-block, noise is generated using an ARMA(1,1) process using Numpy. An alternative procedure is available from Statsmodels (commented out now). More information about the ARMA model can be found on the [statsmodels website](https://www.statsmodels.org/dev/generated/statsmodels.tsa.arima_process.ArmaProcess.html). The noise is added to the head series generated in the previous code-block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reproduction of random numbers\n",
    "np.random.seed(1234)\n",
    "alpha= 0.95\n",
    "beta = 0.1\n",
    "\n",
    "# generate samples using Statsmodels\n",
    "# import statsmodels.api as stats\n",
    "# ar = np.array([1, -alpha])\n",
    "# ma = np.r_[1, beta]\n",
    "# arma = stats.tsa.ArmaProcess(ar, ma)\n",
    "# noise = arma.generate_sample(head[0].index.size)*np.std(head.values) * 0.1\n",
    "\n",
    "# generate samples using Numpy\n",
    "random_seed = np.random.RandomState(1234)\n",
    "\n",
    "noise = random_seed.normal(0,1,len(head)) * np.std(head.values) * 0.1\n",
    "a = np.zeros_like(head[0])\n",
    "\n",
    "for i in range(1, noise.size):\n",
    "    a[i] = noise[i] + noise[i - 1] * beta + a[i - 1] * alpha\n",
    "\n",
    "head_noise = head[0] + a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(a)\n",
    "plt.plot(noise)"
   ]
  },
  {
   "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.ArmaModel())\n",
    "\n",
    "ml.solve(tmin=\"1991\", tmax='2015-06-29', noise=True, report=True)\n",
    "axes = ml.plots.results(figsize=(10,5));\n",
    "axes[-2].plot(ps.Gamma().step([Atrue, ntrue, atrue]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Did we find back the original ARMA parameters?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.exp(-1./ml.parameters.loc[\"noise_alpha\", \"optimal\"]).round(2), \"vs\", alpha)\n",
    "print(np.exp(-1./np.abs(ml.parameters.loc[\"noise_beta\", \"optimal\"])).round(2), \"vs.\", beta)"
   ]
  },
  {
   "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));"
   ]
  },
  {
   "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, compared to its AR(1) counterpart. \n",
    "\n",
    "## 7. What happens if we use an AR(1) 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",
    "\n",
    "ml.solve(tmin=\"1991\", tmax='2015-06-29', noise=True, 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./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));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Significant autocorrelation is still present at lag 1 and the parameter of the AR(1) is overestimated, trying to correct for the lack of an MA(1) part. This is to be expected, as the MA(1) process generates a strong autocorrelation at this first time lag. The negative autocorrelation in the first few time steps is a result of the overestimation of the AR(1) parameter.\n",
    "\n",
    "A possible effect of failing to remove the autocorrelation at lag 1 may be that the parameter standard errors are under- or overestimated. Although that does not seem the case for this synthetic, real life examples may suffer from this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Test the ArmaModel for irregular time steps\n",
    "In this final step the ArmaModel is tested for irregular timesteps, using the indices from a real groundwater level time series. It is clear from the example below that the ArmaModel does not yet work for irregular time steps, as (unlike the AR(1) model in Pastas) no weights are applied yet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = pd.read_csv(\"../examples/data/test_index.csv\", parse_dates=True, \n",
    "                    index_col=0).index.round(\"D\").drop_duplicates()\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, name='recharge', settings='prec')\n",
    "ml.add_stressmodel(sm)\n",
    "ml.add_noisemodel(ps.ArmaModel())\n",
    "\n",
    "ml.solve(tmin=\"1991\", tmax='2015-06-29', noise=True, 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./ml.parameters.loc[\"noise_alpha\", \"optimal\"]).round(2), \"vs\", alpha)\n",
    "print(np.exp(-1./ml.parameters.loc[\"noise_beta\", \"optimal\"]).round(2), \"vs.\", beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml.plots.diagnostics(figsize=(10,4));"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}