{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seasonal-Trend decomposition using LOESS (STL)\n",
    "\n",
    "This note book illustrates the use of `STL` to decompose a time series into three components: trend, season(al) and residual. STL uses LOESS (locally estimated scatterplot smoothing) to extract smooths estimates of the three components.  The key inputs into `STL` are:\n",
    "\n",
    "* `season` - The length of the seasonal smoother. Must be odd.\n",
    "* `trend` - The length of the trend smoother, usually around 150% of `season`.  Must be odd and larger than `season`.\n",
    "* `low_pass` - The length of the low-pass estimation window, usually the smallest odd number larger than the periodicity of the data.\n",
    "\n",
    "First we import the required packages, prepare the graphics environment, and prepare the data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from pandas.plotting import register_matplotlib_converters\n",
    "\n",
    "register_matplotlib_converters()\n",
    "sns.set_style(\"darkgrid\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rc(\"figure\", figsize=(16, 12))\n",
    "plt.rc(\"font\", size=13)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Atmospheric CO2\n",
    "\n",
    "The example in Cleveland, Cleveland, McRae, and Terpenning (1990) uses CO2 data, which is in the list below.  This monthly data (January 1959 to December 1987) has a clear trend and seasonality across the sample. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co2 = [\n",
    "    315.58,\n",
    "    316.39,\n",
    "    316.79,\n",
    "    317.82,\n",
    "    318.39,\n",
    "    318.22,\n",
    "    316.68,\n",
    "    315.01,\n",
    "    314.02,\n",
    "    313.55,\n",
    "    315.02,\n",
    "    315.75,\n",
    "    316.52,\n",
    "    317.10,\n",
    "    317.79,\n",
    "    319.22,\n",
    "    320.08,\n",
    "    319.70,\n",
    "    318.27,\n",
    "    315.99,\n",
    "    314.24,\n",
    "    314.05,\n",
    "    315.05,\n",
    "    316.23,\n",
    "    316.92,\n",
    "    317.76,\n",
    "    318.54,\n",
    "    319.49,\n",
    "    320.64,\n",
    "    319.85,\n",
    "    318.70,\n",
    "    316.96,\n",
    "    315.17,\n",
    "    315.47,\n",
    "    316.19,\n",
    "    317.17,\n",
    "    318.12,\n",
    "    318.72,\n",
    "    319.79,\n",
    "    320.68,\n",
    "    321.28,\n",
    "    320.89,\n",
    "    319.79,\n",
    "    317.56,\n",
    "    316.46,\n",
    "    315.59,\n",
    "    316.85,\n",
    "    317.87,\n",
    "    318.87,\n",
    "    319.25,\n",
    "    320.13,\n",
    "    321.49,\n",
    "    322.34,\n",
    "    321.62,\n",
    "    319.85,\n",
    "    317.87,\n",
    "    316.36,\n",
    "    316.24,\n",
    "    317.13,\n",
    "    318.46,\n",
    "    319.57,\n",
    "    320.23,\n",
    "    320.89,\n",
    "    321.54,\n",
    "    322.20,\n",
    "    321.90,\n",
    "    320.42,\n",
    "    318.60,\n",
    "    316.73,\n",
    "    317.15,\n",
    "    317.94,\n",
    "    318.91,\n",
    "    319.73,\n",
    "    320.78,\n",
    "    321.23,\n",
    "    322.49,\n",
    "    322.59,\n",
    "    322.35,\n",
    "    321.61,\n",
    "    319.24,\n",
    "    318.23,\n",
    "    317.76,\n",
    "    319.36,\n",
    "    319.50,\n",
    "    320.35,\n",
    "    321.40,\n",
    "    322.22,\n",
    "    323.45,\n",
    "    323.80,\n",
    "    323.50,\n",
    "    322.16,\n",
    "    320.09,\n",
    "    318.26,\n",
    "    317.66,\n",
    "    319.47,\n",
    "    320.70,\n",
    "    322.06,\n",
    "    322.23,\n",
    "    322.78,\n",
    "    324.10,\n",
    "    324.63,\n",
    "    323.79,\n",
    "    322.34,\n",
    "    320.73,\n",
    "    319.00,\n",
    "    318.99,\n",
    "    320.41,\n",
    "    321.68,\n",
    "    322.30,\n",
    "    322.89,\n",
    "    323.59,\n",
    "    324.65,\n",
    "    325.30,\n",
    "    325.15,\n",
    "    323.88,\n",
    "    321.80,\n",
    "    319.99,\n",
    "    319.86,\n",
    "    320.88,\n",
    "    322.36,\n",
    "    323.59,\n",
    "    324.23,\n",
    "    325.34,\n",
    "    326.33,\n",
    "    327.03,\n",
    "    326.24,\n",
    "    325.39,\n",
    "    323.16,\n",
    "    321.87,\n",
    "    321.31,\n",
    "    322.34,\n",
    "    323.74,\n",
    "    324.61,\n",
    "    325.58,\n",
    "    326.55,\n",
    "    327.81,\n",
    "    327.82,\n",
    "    327.53,\n",
    "    326.29,\n",
    "    324.66,\n",
    "    323.12,\n",
    "    323.09,\n",
    "    324.01,\n",
    "    325.10,\n",
    "    326.12,\n",
    "    326.62,\n",
    "    327.16,\n",
    "    327.94,\n",
    "    329.15,\n",
    "    328.79,\n",
    "    327.53,\n",
    "    325.65,\n",
    "    323.60,\n",
    "    323.78,\n",
    "    325.13,\n",
    "    326.26,\n",
    "    326.93,\n",
    "    327.84,\n",
    "    327.96,\n",
    "    329.93,\n",
    "    330.25,\n",
    "    329.24,\n",
    "    328.13,\n",
    "    326.42,\n",
    "    324.97,\n",
    "    325.29,\n",
    "    326.56,\n",
    "    327.73,\n",
    "    328.73,\n",
    "    329.70,\n",
    "    330.46,\n",
    "    331.70,\n",
    "    332.66,\n",
    "    332.22,\n",
    "    331.02,\n",
    "    329.39,\n",
    "    327.58,\n",
    "    327.27,\n",
    "    328.30,\n",
    "    328.81,\n",
    "    329.44,\n",
    "    330.89,\n",
    "    331.62,\n",
    "    332.85,\n",
    "    333.29,\n",
    "    332.44,\n",
    "    331.35,\n",
    "    329.58,\n",
    "    327.58,\n",
    "    327.55,\n",
    "    328.56,\n",
    "    329.73,\n",
    "    330.45,\n",
    "    330.98,\n",
    "    331.63,\n",
    "    332.88,\n",
    "    333.63,\n",
    "    333.53,\n",
    "    331.90,\n",
    "    330.08,\n",
    "    328.59,\n",
    "    328.31,\n",
    "    329.44,\n",
    "    330.64,\n",
    "    331.62,\n",
    "    332.45,\n",
    "    333.36,\n",
    "    334.46,\n",
    "    334.84,\n",
    "    334.29,\n",
    "    333.04,\n",
    "    330.88,\n",
    "    329.23,\n",
    "    328.83,\n",
    "    330.18,\n",
    "    331.50,\n",
    "    332.80,\n",
    "    333.22,\n",
    "    334.54,\n",
    "    335.82,\n",
    "    336.45,\n",
    "    335.97,\n",
    "    334.65,\n",
    "    332.40,\n",
    "    331.28,\n",
    "    330.73,\n",
    "    332.05,\n",
    "    333.54,\n",
    "    334.65,\n",
    "    335.06,\n",
    "    336.32,\n",
    "    337.39,\n",
    "    337.66,\n",
    "    337.56,\n",
    "    336.24,\n",
    "    334.39,\n",
    "    332.43,\n",
    "    332.22,\n",
    "    333.61,\n",
    "    334.78,\n",
    "    335.88,\n",
    "    336.43,\n",
    "    337.61,\n",
    "    338.53,\n",
    "    339.06,\n",
    "    338.92,\n",
    "    337.39,\n",
    "    335.72,\n",
    "    333.64,\n",
    "    333.65,\n",
    "    335.07,\n",
    "    336.53,\n",
    "    337.82,\n",
    "    338.19,\n",
    "    339.89,\n",
    "    340.56,\n",
    "    341.22,\n",
    "    340.92,\n",
    "    339.26,\n",
    "    337.27,\n",
    "    335.66,\n",
    "    335.54,\n",
    "    336.71,\n",
    "    337.79,\n",
    "    338.79,\n",
    "    340.06,\n",
    "    340.93,\n",
    "    342.02,\n",
    "    342.65,\n",
    "    341.80,\n",
    "    340.01,\n",
    "    337.94,\n",
    "    336.17,\n",
    "    336.28,\n",
    "    337.76,\n",
    "    339.05,\n",
    "    340.18,\n",
    "    341.04,\n",
    "    342.16,\n",
    "    343.01,\n",
    "    343.64,\n",
    "    342.91,\n",
    "    341.72,\n",
    "    339.52,\n",
    "    337.75,\n",
    "    337.68,\n",
    "    339.14,\n",
    "    340.37,\n",
    "    341.32,\n",
    "    342.45,\n",
    "    343.05,\n",
    "    344.91,\n",
    "    345.77,\n",
    "    345.30,\n",
    "    343.98,\n",
    "    342.41,\n",
    "    339.89,\n",
    "    340.03,\n",
    "    341.19,\n",
    "    342.87,\n",
    "    343.74,\n",
    "    344.55,\n",
    "    345.28,\n",
    "    347.00,\n",
    "    347.37,\n",
    "    346.74,\n",
    "    345.36,\n",
    "    343.19,\n",
    "    340.97,\n",
    "    341.20,\n",
    "    342.76,\n",
    "    343.96,\n",
    "    344.82,\n",
    "    345.82,\n",
    "    347.24,\n",
    "    348.09,\n",
    "    348.66,\n",
    "    347.90,\n",
    "    346.27,\n",
    "    344.21,\n",
    "    342.88,\n",
    "    342.58,\n",
    "    343.99,\n",
    "    345.31,\n",
    "    345.98,\n",
    "    346.72,\n",
    "    347.63,\n",
    "    349.24,\n",
    "    349.83,\n",
    "    349.10,\n",
    "    347.52,\n",
    "    345.43,\n",
    "    344.48,\n",
    "    343.89,\n",
    "    345.29,\n",
    "    346.54,\n",
    "    347.66,\n",
    "    348.07,\n",
    "    349.12,\n",
    "    350.55,\n",
    "    351.34,\n",
    "    350.80,\n",
    "    349.10,\n",
    "    347.54,\n",
    "    346.20,\n",
    "    346.20,\n",
    "    347.44,\n",
    "    348.67,\n",
    "]\n",
    "co2 = pd.Series(\n",
    "    co2, index=pd.date_range(\"1-1-1959\", periods=len(co2), freq=\"M\"), name=\"CO2\"\n",
    ")\n",
    "co2.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The decomposition requires 1 input, the data series. If the data series does not have a frequency, then you must also specify `period`. The default value for `seasonal` is 7, and so should also be changed in most applications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from statsmodels.tsa.seasonal import STL\n",
    "\n",
    "stl = STL(co2, seasonal=13)\n",
    "res = stl.fit()\n",
    "fig = res.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Robust Fitting\n",
    "Setting `robust` uses a data-dependent weighting function that re-weights data when estimating the LOESS (and so is using LOWESS). Using robust estimation allows the model to tolerate larger errors that are visible on the bottom plot.\n",
    "\n",
    "Here we use a series the measures the production of electrical equipment in the EU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from statsmodels.datasets import elec_equip as ds\n",
    "\n",
    "elec_equip = ds.load().data.iloc[:, 0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we estimate the model with and without robust weighting.  The difference is minor and is most pronounced during the financial crisis of 2008. The non-robust estimate places equal weights on all observations and so produces smaller errors, on average.  The weights vary between 0 and 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_stl_plot(fig, res, legend):\n",
    "    \"\"\"Add 3 plots from a second STL fit\"\"\"\n",
    "    axs = fig.get_axes()\n",
    "    comps = [\"trend\", \"seasonal\", \"resid\"]\n",
    "    for ax, comp in zip(axs[1:], comps):\n",
    "        series = getattr(res, comp)\n",
    "        if comp == \"resid\":\n",
    "            ax.plot(series, marker=\"o\", linestyle=\"none\")\n",
    "        else:\n",
    "            ax.plot(series)\n",
    "            if comp == \"trend\":\n",
    "                ax.legend(legend, frameon=False)\n",
    "\n",
    "stl = STL(elec_equip, period=12, robust=True)\n",
    "res_robust = stl.fit()\n",
    "fig = res_robust.plot()\n",
    "res_non_robust = STL(elec_equip, period=12, robust=False).fit()\n",
    "add_stl_plot(fig, res_non_robust, [\"Robust\", \"Non-robust\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16, 5))\n",
    "lines = plt.plot(res_robust.weights, marker=\"o\", linestyle=\"none\")\n",
    "ax = plt.gca()\n",
    "xlim = ax.set_xlim(elec_equip.index[0], elec_equip.index[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LOESS degree\n",
    "The default configuration estimates the LOESS model with both a constant and a trend.  This can be changed to only include a constant by setting `COMPONENT_deg` to 0. Here the degree makes little difference except in the trend around the financial crisis of 2008."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stl = STL(\n",
    "    elec_equip, period=12, seasonal_deg=0, trend_deg=0, low_pass_deg=0, robust=True\n",
    ")\n",
    "res_deg_0 = stl.fit()\n",
    "fig = res_robust.plot()\n",
    "add_stl_plot(fig, res_deg_0, [\"Degree 1\", \"Degree 0\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance\n",
    "Three options can be used to reduce the computational cost of the STL decomposition:\n",
    "\n",
    "* `seasonal_jump`\n",
    "* `trend_jump`\n",
    "* `low_pass_jump`\n",
    "\n",
    "When these are non-zero, the LOESS for component `COMPONENT` is only estimated ever `COMPONENT_jump` observations, and linear interpolation is used between points. These values should not normally be more than 10-20% of the size of `seasonal`, `trend` or `low_pass`, respectively. \n",
    "\n",
    "The example below shows how these can reduce the computational cost by a factor of 15 using simulated data with both a low-frequency cosinusoidal trend and a sinusoidal seasonal pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "rs = np.random.RandomState(0xA4FD94BC)\n",
    "tau = 2000\n",
    "t = np.arange(tau)\n",
    "period = int(0.05 * tau)\n",
    "seasonal = period + ((period % 2) == 0)  # Ensure odd\n",
    "e = 0.25 * rs.standard_normal(tau)\n",
    "y = np.cos(t / tau * 2 * np.pi) + 0.25 * np.sin(t / period * 2 * np.pi) + e\n",
    "plt.plot(y)\n",
    "plt.title(\"Simulated Data\")\n",
    "xlim = plt.gca().set_xlim(0, tau)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, the base line model is estimated with all jumps equal to 1. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod = STL(y, period=period, seasonal=seasonal)\n",
    "%timeit mod.fit()\n",
    "res = mod.fit()\n",
    "fig = res.plot(observed=False, resid=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The jumps are all set to 15% of their window length. Limited linear interpolation makes little difference to the fit of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "low_pass_jump = seasonal_jump = int(0.15 * (period + 1))\n",
    "trend_jump = int(0.15 * 1.5 * (period + 1))\n",
    "mod = STL(\n",
    "    y,\n",
    "    period=period,\n",
    "    seasonal=seasonal,\n",
    "    seasonal_jump=seasonal_jump,\n",
    "    trend_jump=trend_jump,\n",
    "    low_pass_jump=low_pass_jump,\n",
    ")\n",
    "%timeit mod.fit()\n",
    "res = mod.fit()\n",
    "fig = res.plot(observed=False, resid=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Forecasting with STL\n",
    "\n",
    "``STLForecast`` simplifies the process of using STL to remove seasonalities and then using a standard time-series model to forecast the trend and cyclical components. \n",
    "\n",
    "Here we use STL to handle the seasonality and then an ARIMA(1,1,0) to model the deseasonalized data. The seasonal component is forecast from the find full cycle where \n",
    "\n",
    "$$E[S_{T+h}|\\mathcal{F}_T]=\\hat{S}_{T-k}$$\n",
    "\n",
    "where $k= m - h + m \\lfloor \\frac{h-1}{m} \\rfloor$. The forecast automatically adds the seasonal component forecast to the ARIMA forecast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from statsmodels.tsa.arima.model import ARIMA\n",
    "from statsmodels.tsa.forecasting.stl import STLForecast\n",
    "\n",
    "elec_equip.index.freq = elec_equip.index.inferred_freq\n",
    "stlf = STLForecast(elec_equip, ARIMA, model_kwargs=dict(order=(1, 1, 0), trend=\"t\"))\n",
    "stlf_res = stlf.fit()\n",
    "\n",
    "forecast = stlf_res.forecast(24)\n",
    "plt.plot(elec_equip)\n",
    "plt.plot(forecast)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``summary`` contains information about both the time-series model and the STL decomposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(stlf_res.summary())"
   ]
  }
 ],
 "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
