{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# State space modeling: Local Linear Trends"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook describes how to extend the statsmodels statespace classes to create and estimate a custom model. Here we develop a local linear trend model.\n",
    "\n",
    "The Local Linear Trend model has the form (see Durbin and Koopman 2012, Chapter 3.2 for all notation and details):\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "y_t & = \\mu_t + \\varepsilon_t \\qquad & \\varepsilon_t \\sim\n",
    "    N(0, \\sigma_\\varepsilon^2) \\\\\n",
    "\\mu_{t+1} & = \\mu_t + \\nu_t + \\xi_t & \\xi_t \\sim N(0, \\sigma_\\xi^2) \\\\\n",
    "\\nu_{t+1} & = \\nu_t + \\zeta_t & \\zeta_t \\sim N(0, \\sigma_\\zeta^2)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "It is easy to see that this can be cast into state space form as:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "y_t & = \\begin{pmatrix} 1 & 0 \\end{pmatrix} \\begin{pmatrix} \\mu_t \\\\ \\nu_t \\end{pmatrix} + \\varepsilon_t \\\\\n",
    "\\begin{pmatrix} \\mu_{t+1} \\\\ \\nu_{t+1} \\end{pmatrix} & = \\begin{bmatrix} 1 & 1 \\\\ 0 & 1 \\end{bmatrix} \\begin{pmatrix} \\mu_t \\\\ \\nu_t \\end{pmatrix} + \\begin{pmatrix} \\xi_t \\\\ \\zeta_t \\end{pmatrix}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Notice that much of the state space representation is composed of known values; in fact the only parts in which parameters to be estimated appear are in the variance / covariance matrices:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "H_t & = \\begin{bmatrix} \\sigma_\\varepsilon^2 \\end{bmatrix} \\\\\n",
    "Q_t & = \\begin{bmatrix} \\sigma_\\xi^2 & 0 \\\\ 0 & \\sigma_\\zeta^2 \\end{bmatrix}\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.stats import norm\n",
    "import statsmodels.api as sm\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To take advantage of the existing infrastructure, including Kalman filtering and maximum likelihood estimation, we create a new class which extends from `statsmodels.tsa.statespace.MLEModel`. There are a number of things that must be specified:\n",
    "\n",
    "1. **k_states**, **k_posdef**: These two parameters must be provided to the base classes in initialization. The inform the statespace model about the size of, respectively, the state vector, above $\\begin{pmatrix} \\mu_t & \\nu_t \\end{pmatrix}'$, and   the state error vector, above $\\begin{pmatrix} \\xi_t & \\zeta_t \\end{pmatrix}'$. Note that the dimension of the endogenous vector does not have to be specified, since it can be inferred from the `endog` array.\n",
    "2. **update**: The method `update`, with argument `params`, must be specified (it is used when `fit()` is called to calculate the MLE). It takes the parameters and fills them into the appropriate state space matrices. For example, below, the `params` vector contains variance parameters $\\begin{pmatrix} \\sigma_\\varepsilon^2 & \\sigma_\\xi^2 & \\sigma_\\zeta^2\\end{pmatrix}$, and the `update` method must place them in the observation and state covariance matrices. More generally, the parameter vector might be mapped into many different places in all of the statespace matrices.\n",
    "3. **statespace matrices**: by default, all state space matrices (`obs_intercept, design, obs_cov, state_intercept, transition, selection, state_cov`) are set to zeros. Values that are fixed (like the ones in the design and transition matrices here) can be set in initialization, whereas values that vary with the parameters should be set in the `update` method. Note that it is easy to forget to set the selection matrix, which is often just the identity matrix (as it is here), but not setting it will lead to a very different model (one where there is not a stochastic component to the transition equation).\n",
    "4. **start params**: start parameters must be set, even if it is just a vector of zeros, although often good start parameters can be found from the data. Maximum likelihood estimation by gradient methods (as employed here) can be sensitive to the starting parameters, so it is important to select good ones if possible. Here it does not matter too much (although as variances, they should't be set zero).\n",
    "5. **initialization**: in addition to defined state space matrices, all state space models must be initialized with the mean and variance for the initial distribution of the state vector. If the distribution is known, `initialize_known(initial_state, initial_state_cov)` can be called, or if the model is stationary (e.g. an ARMA model), `initialize_stationary` can be used. Otherwise, `initialize_approximate_diffuse` is a reasonable generic initialization (exact diffuse initialization is not yet available). Since the local linear trend model is not stationary (it is composed of random walks) and since the distribution is not generally known, we use `initialize_approximate_diffuse` below.\n",
    "\n",
    "The above are the minimum necessary for a successful model. There are also a number of things that do not have to be set, but which may be helpful or important for some applications:\n",
    "\n",
    "1. **transform / untransform**: when `fit` is called, the optimizer in the background will use gradient methods to select the parameters that maximize the likelihood function. By default it uses unbounded optimization, which means that it may select any parameter value. In many cases, that is not the desired behavior; variances, for example, cannot be negative. To get around this, the `transform` method takes the unconstrained vector of parameters provided by the optimizer and returns a constrained vector of parameters used in likelihood evaluation. `untransform` provides the reverse operation.\n",
    "2. **param_names**: this internal method can be used to set names for the estimated parameters so that e.g. the summary provides meaningful names. If not present, parameters are named `param0`, `param1`, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Univariate Local Linear Trend Model\n",
    "\"\"\"\n",
    "class LocalLinearTrend(sm.tsa.statespace.MLEModel):\n",
    "    def __init__(self, endog):\n",
    "        # Model order\n",
    "        k_states = k_posdef = 2\n",
    "\n",
    "        # Initialize the statespace\n",
    "        super(LocalLinearTrend, self).__init__(\n",
    "            endog, k_states=k_states, k_posdef=k_posdef,\n",
    "            initialization='approximate_diffuse',\n",
    "            loglikelihood_burn=k_states\n",
    "        )\n",
    "\n",
    "        # Initialize the matrices\n",
    "        self.ssm['design'] = np.array([1, 0])\n",
    "        self.ssm['transition'] = np.array([[1, 1],\n",
    "                                       [0, 1]])\n",
    "        self.ssm['selection'] = np.eye(k_states)\n",
    "\n",
    "        # Cache some indices\n",
    "        self._state_cov_idx = ('state_cov',) + np.diag_indices(k_posdef)\n",
    "\n",
    "    @property\n",
    "    def param_names(self):\n",
    "        return ['sigma2.measurement', 'sigma2.level', 'sigma2.trend']\n",
    "\n",
    "    @property\n",
    "    def start_params(self):\n",
    "        return [np.std(self.endog)]*3\n",
    "\n",
    "    def transform_params(self, unconstrained):\n",
    "        return unconstrained**2\n",
    "\n",
    "    def untransform_params(self, constrained):\n",
    "        return constrained**0.5\n",
    "\n",
    "    def update(self, params, *args, **kwargs):\n",
    "        params = super(LocalLinearTrend, self).update(params, *args, **kwargs)\n",
    "        \n",
    "        # Observation covariance\n",
    "        self.ssm['obs_cov',0,0] = params[0]\n",
    "\n",
    "        # State covariance\n",
    "        self.ssm[self._state_cov_idx] = params[1:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using this simple model, we can estimate the parameters from a local linear trend model. The following example is from Commandeur and Koopman (2007), section 3.4., modeling motor vehicle fatalities in Finland."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "from io import BytesIO\n",
    "from zipfile import ZipFile\n",
    "    \n",
    "# Download the dataset\n",
    "ck = requests.get('http://staff.feweb.vu.nl/koopman/projects/ckbook/OxCodeAll.zip').content\n",
    "zipped = ZipFile(BytesIO(ck))\n",
    "df = pd.read_table(\n",
    "    BytesIO(zipped.read('OxCodeIntroStateSpaceBook/Chapter_2/NorwayFinland.txt')),\n",
    "    skiprows=1, header=None, sep='\\s+', engine='python',\n",
    "    names=['date','nf', 'ff']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we defined the local linear trend model as extending from `MLEModel`, the `fit()` method is immediately available, just as in other statsmodels maximum likelihood classes. Similarly, the returned results class supports many of the same post-estimation results, like the `summary` method.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Load Dataset\n",
    "df.index = pd.date_range(start='%d-01-01' % df.date[0], end='%d-01-01' % df.iloc[-1, 0], freq='YS')\n",
    "\n",
    "# Log transform\n",
    "df['lff'] = np.log(df['ff'])\n",
    "\n",
    "# Setup the model\n",
    "mod = LocalLinearTrend(df['lff'])\n",
    "\n",
    "# Fit it using MLE (recall that we are fitting the three variance parameters)\n",
    "res = mod.fit(disp=False)\n",
    "print(res.summary())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can do post-estimation prediction and forecasting. Notice that the end period can be specified as a date."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Perform prediction and forecasting\n",
    "predict = res.get_prediction()\n",
    "forecast = res.get_forecast('2014')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10,4))\n",
    "\n",
    "# Plot the results\n",
    "df['lff'].plot(ax=ax, style='k.', label='Observations')\n",
    "predict.predicted_mean.plot(ax=ax, label='One-step-ahead Prediction')\n",
    "predict_ci = predict.conf_int(alpha=0.05)\n",
    "predict_index = np.arange(len(predict_ci))\n",
    "ax.fill_between(predict_index[2:], predict_ci.iloc[2:, 0], predict_ci.iloc[2:, 1], alpha=0.1)\n",
    "\n",
    "forecast.predicted_mean.plot(ax=ax, style='r', label='Forecast')\n",
    "forecast_ci = forecast.conf_int()\n",
    "forecast_index = np.arange(len(predict_ci), len(predict_ci) + len(forecast_ci))\n",
    "ax.fill_between(forecast_index, forecast_ci.iloc[:, 0], forecast_ci.iloc[:, 1], alpha=0.1)\n",
    "\n",
    "# Cleanup the image\n",
    "ax.set_ylim((4, 8));\n",
    "legend = ax.legend(loc='lower left');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "    Commandeur, Jacques J. F., and Siem Jan Koopman. 2007.\n",
    "    An Introduction to State Space Time Series Analysis.\n",
    "    Oxford ; New York: Oxford University Press.\n",
    "\n",
    "    Durbin, James, and Siem Jan Koopman. 2012.\n",
    "    Time Series Analysis by State Space Methods: Second Edition.\n",
    "    Oxford University Press."
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
