examples/somata_getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ff8aec6e",
   "metadata": {},
   "source": [
    "# Getting started with SOMATA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d05246d7",
   "metadata": {},
   "source": [
    "This notebook contains three self-guided exercises to get started with the SOMATA package.\n",
    "\n",
    "Anyone new to SOMATA is recommended to work through this notebook before analyzing data or building new algorithms with SOMATA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e3bb55f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Author: Mingjian He <mh1@stanford.edu>\n",
    "import numpy as np\n",
    "from pprint import pprint\n",
    "from scipy import signal\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Global setting of matplotlib figures\n",
    "plt.rcParams.update({\n",
    "    \"text.usetex\": True,\n",
    "    \"font.family\": \"Helvetica\",\n",
    "    \"figure.constrained_layout.use\": True,\n",
    "    \"savefig.dpi\": 300\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c3b6abf",
   "metadata": {},
   "source": [
    "## Exercise 1: getting familiar with the syntaxes\n",
    "Learning objective: understand common syntaxes to manipulate SOMATA basic models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad83a0a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the four different basic models in SOMATA\n",
    "from somata.basic_models import StateSpaceModel as Ssm\n",
    "from somata.basic_models import GeneralSSModel as Gen\n",
    "from somata.basic_models import OscillatorModel as Osc\n",
    "from somata.basic_models import AutoRegModel as Arn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92f6e661",
   "metadata": {},
   "source": [
    "### 1.1 `print()` and `append()`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8109b12",
   "metadata": {},
   "source": [
    "These constructor methods can be called without any arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dec66dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm()  # create an empty state-space model instance without any parameters\n",
    "s1  # this returns the __repr__() method output for a Python object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58475bdc",
   "metadata": {},
   "source": [
    "Invoking `print()` gives helpful summary information about the state-space model object.\n",
    "\n",
    "This is one of the most heavily used methods in SOMATA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f003df2b-d172-4e27-b149-f0932cbf3254",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(s1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecc20468",
   "metadata": {},
   "source": [
    "Appending one model to another model augments the state-space parameters in block-diagonal form."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b117b112-b4b2-4e10-b046-b437df35375b",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm(F=1)\n",
    "s1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7f48d4c",
   "metadata": {},
   "source": [
    "The first model calling `append()` gets modified in place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4e40880",
   "metadata": {},
   "outputs": [],
   "source": [
    "s2 = Ssm(F=0.9)\n",
    "s1.append(s2)\n",
    "print(s1)\n",
    "print('s1.F', s1.F)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bae2b541",
   "metadata": {},
   "source": [
    "### 1.2 The concept of `components`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbc8ae0a",
   "metadata": {},
   "source": [
    "An alternative way to think about this new state-space model `s1` is that it is a state-space model with two components.\n",
    "\n",
    "The first component has a univariate state in the hidden state space, so does the second component.\n",
    "\n",
    "We can try to recreate such state-space model using the `components` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f3ec8fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "g1 = Gen(F=1)\n",
    "g2 = Gen(F=0.9)\n",
    "s1 = Ssm(components=[g1, g2])\n",
    "s1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2feb942",
   "metadata": {},
   "source": [
    "Notice that the `__repr__()` method now returns `Ssm(2)<....>`, which indicates that this state-space model now has two components (`ncomp=2`).\n",
    "\n",
    "Utilizing `components` to represent state-space models is advantageous in many ways. To list a few:\n",
    "1. It gives a much clearer understanding of the structure of the state-space model in a quick glance.\n",
    "2. It can be used for updating model parameters in parallel during the M-step of EM algorithm due to the convenient block-diagonal structure.\n",
    "3. It provides a breadboard for different combinations of state-space model components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4400afe6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(s1)\n",
    "print('s1.F', s1.F)\n",
    "print('s1.G', s1.G)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1aafa735",
   "metadata": {},
   "source": [
    "The four characters in `<>` are the last four digits of the memory address ID of the model instance, just in case it becomes difficult to keep track of multiple models.\n",
    "\n",
    "Notice how `g1` and `g2` have become the component models within `s1`, indicated by the same memory addresses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be315d37",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(g1)\n",
    "print(s1.components[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d291e0bb",
   "metadata": {},
   "source": [
    "However, when `s1` undergoes EM, the component models `g1` and `g2` will not get updated automatically.\n",
    "\n",
    "This is so that computations are not slowed down by unnecessary `setattr()` calls. If you would like to map parameters in `s1` back to its components, use `fill_components()`.\n",
    "\n",
    "If you want to break the memory address link, simply cascade `copy()` to the model constructor call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a965ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "s2 = Ssm(components=[g1, g2]).copy()\n",
    "print(s2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b295918e",
   "metadata": {},
   "source": [
    "Since `StateSpaceModel` is the parent state-space model class object for the other basic models in SOMATA, itself cannot be a component.\n",
    "\n",
    "This constraint is to avoid ambiguous handling of component state-space models. All other basic models can be components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cfa486b",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm(F=1)\n",
    "s2 = Ssm(F=0.9)\n",
    "# s3 = Ssm(components=[s1, s2])  # this won't work by design"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22850f18",
   "metadata": {},
   "source": [
    "### 1.3 `GeneralSSModel` class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b461c518",
   "metadata": {},
   "source": [
    "`GeneralSSModel` prints the same attributes as `StateSpaceModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3293b89",
   "metadata": {},
   "outputs": [],
   "source": [
    "g1 = Gen(F=1)\n",
    "g1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c010d3f",
   "metadata": {},
   "source": [
    "Notice that the `components` attribute is automatically populated with a \"space-holder\" component model, which did not happen for `StateSpaceModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7c36435",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(g1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a35e80ba",
   "metadata": {},
   "source": [
    "This component model does not have any parameter unlike in the constructor example above using the `components=` argument.\n",
    "\n",
    "**It is important to emphasize that `components` models should not be used directly throughout SOMATA!**\n",
    "\n",
    "Regardless of carrying parameters or not, they act as space holders to organize the overarching state-space model rather than being used for computations themselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c140e32",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(g1.components[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f149788",
   "metadata": {},
   "source": [
    "### 1.4 `OscillatorModel` class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b785ee6",
   "metadata": {},
   "source": [
    "`OscillatorModel` has special attributes that are printed to help quickly understand the oscillator model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81a21317",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1 = Osc()\n",
    "o1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2374438",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(o1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cfc9cd1",
   "metadata": {},
   "source": [
    "This also means that we can create an `OscillatorModel` using more concise parameter arguments.\n",
    "\n",
    "The default $\\sigma^2=3$ gets filled in when `sigma2` isn't provided, but you would likely want to set it explicitly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a2f349c",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1 = Osc(a=0.95, w=np.pi/10)\n",
    "print(o1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eb33c2a",
   "metadata": {},
   "source": [
    "Notice how the `freq` display changes to the physical `Hz` unit when sampling rate (`Fs`) is also provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04e176d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1 = Osc(a=0.95, w=np.pi/10, Fs=250)\n",
    "print(o1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35715b54",
   "metadata": {},
   "source": [
    "We can directly provide rotating frequency in `Hz` along with `Fs` to create an `OscillatorModel`, which is very useful in practice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24493711",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1 = Osc(a=0.99, freq=10, Fs=100)\n",
    "print(o1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ab3ab9a",
   "metadata": {},
   "source": [
    "### 1.5 `AutoRegModel` class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9b7ca10",
   "metadata": {},
   "source": [
    "`AutoRegModel` has a more intelligent `__repr__()` since the AR order is a more relevant information when there is only one component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f55d550b",
   "metadata": {},
   "outputs": [],
   "source": [
    "a1 = Arn()\n",
    "a1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6de135d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "a2 = Arn(coeff=[0.9, -0.5, 0.3])\n",
    "a2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e797e3cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(a2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c640dbbe",
   "metadata": {},
   "source": [
    "You can also create a multi-component AR model by passing a list of lists of parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b70367e",
   "metadata": {},
   "outputs": [],
   "source": [
    "a3 = Arn(coeff=[[0.9, -0.5], [0.3, 0.1]])\n",
    "a3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeda8bc4",
   "metadata": {},
   "source": [
    "Like `OscillatorModel`, the default $\\sigma^2=3$ gets filled in, but you would likely want to set it explicitly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18ff335f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(a3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab4b61f5",
   "metadata": {},
   "source": [
    "Note that `AutoRegModel` does not support multivariate AR models because the display gets too clunky.\n",
    "\n",
    "Use `StateSpaceModel` to hold high-dimensional state-space models instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a256cd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mva1 = Arn(F=[[0.9, -0.5], [0.3, 0.1]], Q=[[0.9, 0.2], [0.2, 0.5]])  # this won't work by design\n",
    "mva1 = Ssm(F=[[0.9, -0.5], [0.3, 0.1]], Q=[[0.9, 0.2], [0.2, 0.5]])\n",
    "print(mva1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5163e1e5",
   "metadata": {},
   "source": [
    "### 1.6 State-space models with heterogeneous `components`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d1a340e",
   "metadata": {},
   "source": [
    "Now let's try something fancy. What if we want to build a state-space model with both an AR model and an oscillator model as its components?\n",
    "\n",
    "The first way is to pass the AR and oscillator models as `components` into a new state-space model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b045799d",
   "metadata": {},
   "outputs": [],
   "source": [
    "a1 = Arn(coeff=[0.9, -0.5, 0.3])\n",
    "o1 = Osc(a=0.99, freq=10, Fs=100)\n",
    "s1 = Ssm(components=[a1, o1])\n",
    "print(s1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d56790a",
   "metadata": {},
   "source": [
    "The second way is to create the two models and then use the `concat_()` method instead of `append()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9bf0691",
   "metadata": {},
   "outputs": [],
   "source": [
    "a1 = Arn(coeff=[0.9, -0.5, 0.3])\n",
    "o1 = Osc(a=0.99, freq=10, Fs=100)\n",
    "s1 = a1.concat_(o1)  # this returns a new model instance instead of modifying a1 in place like in append()\n",
    "print(s1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7076059",
   "metadata": {},
   "source": [
    "Notice that unlike in the first way, the memory addresses of the component models in `s1` are different from those of `a1` and `o1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bc18c0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(a1)\n",
    "print(o1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "931d5fa6",
   "metadata": {},
   "source": [
    "### 1.7 Automatic initialization of `components`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5188edd1",
   "metadata": {},
   "source": [
    "The `components` argument can also be used to automatically construct a specific model type as components from parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe567696",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm(components='Arn', F=0.9, Q=1)\n",
    "print(s1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0fbdeb2",
   "metadata": {},
   "source": [
    "We can use this syntax to automatically parse block-diagonal parameters into multiple components. \n",
    "\n",
    "For example, two AR1 models with filled parameters are created below as `components`. We will later see that this is convenient for `OscillatorModel` as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "081ff4a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm(components='Arn', F=[[0.9, 0], [0, 0.8]], Q=[[1, 0], [0, 1]])\n",
    "print(s1)\n",
    "print(s1.components[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1c470b9",
   "metadata": {},
   "source": [
    "### 1.8 Model stacking and arrays of models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f545526d",
   "metadata": {},
   "source": [
    "Unlike `append()` or `concat_()`, adding one model to another model will \"stack\" in the third dimension of mismatched parameters (`F` in this case) and increase `nmodel` by 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b504909",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm(F=1, Q=0.5)\n",
    "s2 = Ssm(F=0.9, Q=0.5)\n",
    "s3 = s1 + s2\n",
    "s3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parameters that are matched between state-space models (`Q` in this case) are kept intact to be memory-efficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2070b5f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(s3)\n",
    "print('s3.F', s3.F)\n",
    "print('s3.Q', s3.Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d2622e6",
   "metadata": {},
   "source": [
    "If we multiply two models with different parameters, models with combinations of their parameters are stacked together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "914a84e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm(F=1, Q=0.5, R=0.3)\n",
    "s2 = Ssm(F=0.9, Q=0.5, R=0.5)\n",
    "s3 = s1 * s2\n",
    "print(s3)\n",
    "print('s3.F', s3.F)\n",
    "print('s3.Q', s3.Q)\n",
    "print('s3.R', s3.R)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82715c9f",
   "metadata": {},
   "source": [
    "Raising a model to any positive exponential power will form combinations across its own parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1035e13",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm(F=1, Q=0.1, R=0.3)\n",
    "s2 = Ssm(F=0.9, Q=0.2, R=0.5)\n",
    "s3 = s1 + s2\n",
    "print(s3)\n",
    "s4 = s3 ** 2\n",
    "print(s4)\n",
    "print('s4.F', s4.F)\n",
    "print('s4.Q', s4.Q)\n",
    "print('s4.R', s4.R)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65f9a782",
   "metadata": {},
   "source": [
    "`s3` and `s4` are essentially condensed models representing multiple alternative state-space models, which differ in some parameters.\n",
    "\n",
    "We can turn them into arrays of state-space models for easy looping with the `stack_to_array()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a24e76f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(s4.stack_to_array())\n",
    "print(type(s4.stack_to_array()))\n",
    "\n",
    "for model_no, model in enumerate(s4.stack_to_array()):\n",
    "    print(f'Model {model_no}: F =', model.F, 'Q =', model.Q, 'R =', model.R)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d8610f7",
   "metadata": {},
   "source": [
    "### 1.9 Now that you have learned how to work with SOMATA basic models, let's give it a try\n",
    "\n",
    "**Create a seasonal adjustment model with the following components:**\n",
    "\n",
    "1. A random walk of variance $0.8$, i.e., $x_t = x_{t-1} + w_t$ and $w_t \\sim \\mathcal{N}(0, 0.8)$.\n",
    "2. An oscillation with damping factor of 0.96, central frequency of $20$ Hz, and mean-zero process noise of variance $1.2$.\n",
    "3. An autoregressive process with coefficients $[0.7, -0.35, 0.5]$ and process noise of variance $0.3$.\n",
    "4. An unnamed dynamic with transition matrix $\\begin{bmatrix} 0.3 & 0.4\\\\ 0.5 & 0.6 \\end{bmatrix}$, state noise covariance $\\begin{bmatrix} 0.6 & 0.23\\\\ 0.23 & 0.4 \\end{bmatrix}$, and observation matrix $\\begin{bmatrix} 0.112 & 3.58 \\end{bmatrix}$.\n",
    "\n",
    "The seasonal adjustment model is used to track a stream of univariate observations collected at $250$ Hz sampling rate, and you estimate the observation noise variance to be at $1.0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b9533e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "m1 = Arn(F=1, Q=0.8)\n",
    "m2 = Osc(a=0.96, freq=20, sigma2=1.2, Fs=250)\n",
    "m3 = Arn(coeff=[0.7, -0.35, 0.5], sigma2=0.3)\n",
    "m4 = Gen(F=[[0.3, 0.4], [0.5, 0.6]], Q=[[0.6, 0.23], [0.23, 0.4]], G=[0.112, 3.58])\n",
    "\n",
    "season_model = Ssm(components=[m1, m2, m3, m4], R=1.0)\n",
    "print(season_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84817e37",
   "metadata": {},
   "source": [
    "## Exercise 2: running EM algorithm with state-space models\n",
    "Learning objective: understand how to do state inference and parameter estimation on SOMATA basic models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80d09a1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import two of the basic models in SOMATA\n",
    "from somata.basic_models import StateSpaceModel as Ssm\n",
    "from somata.basic_models import OscillatorModel as Osc\n",
    "\n",
    "# Import a data loading helper function from the tests directory\n",
    "from tests.test_load_data import _load_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34684412",
   "metadata": {},
   "source": [
    "### 2.1 Load the data and visualize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e676c78d",
   "metadata": {},
   "source": [
    "Load a 10-s example time series data and create a state-space model.\n",
    "\n",
    "Notice that `components=None` in `s1`. We will come back to this later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e16f94a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load example state-space model parameters from a .mat file\n",
    "ssm_params = _load_data(\"kalman_inputs.mat\")\n",
    "_ = ssm_params.pop('R_weights')  # remove extra variable\n",
    "ssm_params['Fs'] = 100  # add sampling frequency to the dictionary\n",
    "\n",
    "s1 = Ssm(F=ssm_params['F'], Q=ssm_params['Q'], mu0=ssm_params['mu0'], S0=ssm_params['S0'],\n",
    "         G=ssm_params['G'], R=ssm_params['R'], y=ssm_params['y'], Fs=ssm_params['Fs'])\n",
    "print(s1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2f71580",
   "metadata": {},
   "source": [
    "Plot the raw time trace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5785800",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(s1.y.T)\n",
    "plt.xlabel('Time (samples)')\n",
    "plt.ylabel('Arbitrary units')\n",
    "_ = plt.title('Example time series data')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a5fe8e0",
   "metadata": {},
   "source": [
    "The parameters in `ssm_params` are already set up as two independent oscillators. Let's create a more informative state-space model object using `OscillatorModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b626a8fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1 = Osc(F=ssm_params['F'], Q=ssm_params['Q'], mu0=ssm_params['mu0'], S0=ssm_params['S0'],\n",
    "         G=ssm_params['G'], R=ssm_params['R'], y=ssm_params['y'], Fs=ssm_params['Fs'])  # automatically parses oscillators\n",
    "o1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "863b81a3",
   "metadata": {},
   "source": [
    "This state equation contains two oscillators in the hidden state space: one at $1.48$ Hz, and one at $12.66$ Hz. These parameters have been updated after 50 iterations of EM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edbce0bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(o1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b6ac99d",
   "metadata": {},
   "source": [
    "### 2.2 State inference with Kalman filtering and smoothing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7565886c",
   "metadata": {},
   "source": [
    "Performing Kalman filtering and smoothing with SOMATA basic models is very simple."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "958c821d",
   "metadata": {},
   "outputs": [],
   "source": [
    "kalman_results = o1.kalman_filt_smooth(return_dict=True)\n",
    "_ = [print(x) for x in kalman_results.keys()]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "299334ad",
   "metadata": {},
   "source": [
    "When we are doing EM, we don't need all these outputs so we can get a more concise dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aad2cfa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "em_kalman_results = o1.kalman_filt_smooth(EM=True)\n",
    "_ = [print(x) for x in em_kalman_results.keys()]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "168ca5c7",
   "metadata": {},
   "source": [
    "Let's look at the Kalman filtered estimates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2c38d2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(4, 1, figsize=(10, 10))\n",
    "ax[0].plot(kalman_results['x_t_t'][0, :], label='Oscillator 1 - Real')\n",
    "ax[1].plot(kalman_results['x_t_t'][1, :], label='Oscillator 1 - Imaginary')\n",
    "ax[2].plot(kalman_results['x_t_t'][2, :], label='Oscillator 2 - Real')\n",
    "ax[3].plot(kalman_results['x_t_t'][3, :], label='Oscillator 2 - Imaginary')\n",
    "for axx in ax:\n",
    "    axx.legend(loc='upper right')\n",
    "    axx.set_ylabel('Arbitrary units')\n",
    "ax[-1].set_xlabel('Time (samples)')\n",
    "_ = ax[0].set_title('Kalman filtering estimates')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7209471d",
   "metadata": {},
   "source": [
    "If we have multiple models stacked together or an array of models, we can run Kalman filtering and smoothing in parallel with `par_kalman()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6541685",
   "metadata": {},
   "outputs": [],
   "source": [
    "o2 = o1.copy()\n",
    "o2.freq = np.array([2, 13])  # manually change the oscillator rotation center frequencies\n",
    "o2.fill_ssm_param()  # this propagates the changes to the state-space model parameters (F, Q)\n",
    "print(o2)\n",
    "print('o1.F', o1.F)\n",
    "print('o2.F', o2.F)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a05e429f",
   "metadata": {},
   "source": [
    "We can create a stacked model then run Kalman filtering and smoothing across all underlying models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15cb42da",
   "metadata": {},
   "outputs": [],
   "source": [
    "o3 = o1 + o2  # this creates a stacked model\n",
    "print(o3)\n",
    "\n",
    "# We can directly call the par_kalman() method on the stacked model\n",
    "par_kalman_results = Ssm.par_kalman(o3, return_dict=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd053413",
   "metadata": {},
   "source": [
    "Each of the `_all` keys points to a list of corresponding results for the array of models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5699da8",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = [print(x) for x in par_kalman_results.keys()]\n",
    "par_kalman_results['x_t_n_all']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adad66fa",
   "metadata": {},
   "source": [
    "Equivalently, we can create an array of models and call the `par_kalman()` method.\n",
    "\n",
    "This syntax is quite flexible, since the array can contain arbitrary SOMATA state-space models, as long as Kalman filtering and smoothing can be performed.\n",
    "\n",
    "One default check is that these models should all have the same observed data `y` to discourage erroneous model comparisons.\n",
    "\n",
    "If one really intends to run with different data, pass in `skip_check_observed=True` but be mindful of estimation differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1121ab86",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_array = o3.stack_to_array()  # equivalent to model_array = [o1, o2]\n",
    "par_kalman_results = Ssm.par_kalman(model_array, return_dict=True, skip_check_observed=False)\n",
    "_ = [print(x) for x in par_kalman_results.keys()]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2703ab7c",
   "metadata": {},
   "source": [
    "### 2.3 Run EM algorithm on state-space models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ea03e89",
   "metadata": {},
   "source": [
    "It is very simple to run the vanilla EM algorithm with maximum likelihood estimation (MLE).\n",
    "\n",
    "However, there are a few things to take note of:\n",
    "1. The `**` syntax simply passes the dictionary key-value pairs as keyword arguments into the `m_estimate()` method.\n",
    "2. The `EM=True` flag is necessary since otherwise there are extra arguments created from the full dictionary of Kalman results.\n",
    "3. The model calling `m_estimate()` is modified in place, including all parameters but not the component models.\n",
    "4. Both `kalman_filt_smooth()` and `m_estimate()` can take in different observed data using the `y=` argument, which is useful for handling data segments.\n",
    "4. While a plain model, like `s1` above, can run Kalman filtering and smoothing, a model must have `components` in order to invoke the `m_estimate()` method.\n",
    "5. The `m_estimate()` method looks into `components` to retrieve component-specific update methods and to update matrix blocks in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9ccbd0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1.m_estimate(**o1.kalman_filt_smooth(EM=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b84b6144",
   "metadata": {},
   "source": [
    "Therefore, in order to follow the oscillator specific update equations, one needs to use the `OscillatorModel` constructor from the beginning.\n",
    "\n",
    "Or one can pass in `components='Osc'` when using the `StateSpaceModel` constructor. This ensures that the component models are of `OscillatorModel` type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e53709a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# s1.m_estimate(**s1.kalman_filt_smooth(EM=True))  # this won't work by design\n",
    "\n",
    "s1 = Ssm(F=ssm_params['F'], Q=ssm_params['Q'], mu0=ssm_params['mu0'], S0=ssm_params['S0'],\n",
    "         G=ssm_params['G'], R=ssm_params['R'], y=ssm_params['y'], Fs=ssm_params['Fs'], components='Osc')\n",
    "print(s1)\n",
    "s1.m_estimate(**s1.kalman_filt_smooth(EM=True))  # this works because component models are OscillatorModel instances\n",
    "print('s1.Q', s1.Q)  # this is diagonal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d86f90e4",
   "metadata": {},
   "source": [
    "A different set of update rules will be employed if a different component model class is assumed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd512051",
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Ssm(F=ssm_params['F'], Q=ssm_params['Q'], mu0=ssm_params['mu0'], S0=ssm_params['S0'],\n",
    "         G=ssm_params['G'], R=ssm_params['R'], y=ssm_params['y'], Fs=ssm_params['Fs'], components='Gen')\n",
    "print(s1)\n",
    "s1.m_estimate(**s1.kalman_filt_smooth(EM=True))  # this uses the full matrix update rules for the GeneralSSModel instances\n",
    "print('s1.Q', s1.Q)  # no longer diagonal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36032d86",
   "metadata": {},
   "source": [
    "One can easily run multiple EM iterations in a for loop. Let's start with some arbitrary guesses of parameters and see how EM increases the log-likelihood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71207fcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1 = Osc(a=[0.99, 0.99], freq=[1, 10], y=ssm_params['y'], Fs=ssm_params['Fs'])\n",
    "print('Before EM: o1.freq', o1.freq)\n",
    "logL = []\n",
    "for _ in range(50):  # run 50 iterations of EM\n",
    "    kalman_results = o1.kalman_filt_smooth(EM=True)\n",
    "    logL.append(kalman_results['logL'].sum())  # sum over time points\n",
    "    o1.m_estimate(**kalman_results)\n",
    "print('After EM: o1.freq', o1.freq)\n",
    "\n",
    "plt.plot(logL)\n",
    "plt.xlabel('EM iteration')\n",
    "plt.ylabel('Log-likelihood')\n",
    "_ = plt.title('Log-likelihood over EM iterations')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db6fd346",
   "metadata": {},
   "source": [
    "Priors can be used as long as they are implemented in the corresponding `_m_update_...()` methods of the component model class.\n",
    "\n",
    "Details of these priors can be found in the implementations of model parameter update methods called within `m_estimate()` during the EM algorithm.\n",
    "\n",
    "One can also control which parameters get updated using the `update_param=` or `keep_param=` arguments. Examples are omitted here, but you are encouraged to explore their usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70b602a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "priors = o1.initialize_priors()  # a list of priors one dictionary per component\n",
    "pprint(priors)\n",
    "o1.m_estimate(**o1.kalman_filt_smooth(EM=True), priors=priors)  # maximum a posteriori (MAP) estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cce74b4",
   "metadata": {},
   "source": [
    "### 2.4 Comparison with Hilbert transform after bandpass filtering"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b38b00b",
   "metadata": {},
   "source": [
    "A common approach neuroscientists use to extract the amplitude and phase of neural oscillations is by bandpass filtering followed by Hilbert transform.\n",
    "\n",
    "Let's compare the results of this approach to state-space oscillator modeling of the ~$12$ Hz oscillation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbf65210",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First select the order of the lowest order filter to achieve the desired filter properties\n",
    "N, Wn = signal.buttord(wp=[10, 14], ws=[9, 15], gpass=1, gstop=50, fs=ssm_params['Fs'])\n",
    "\n",
    "# Then we design the bandpass filter\n",
    "sos = signal.butter(N, Wn, btype=\"bandpass\", output='sos', fs=ssm_params['Fs'])\n",
    "\n",
    "# Visualize the filter frequency response - it's a good habit to always check your filter design\n",
    "w, h = signal.sosfreqz(sos, fs=ssm_params['Fs'])\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(w, 20 * np.log10(np.maximum(np.abs(h), 1e-3)))\n",
    "ax.set_title('Butterworth filter frequency response')\n",
    "ax.set_xlabel('Frequency [Hz]')\n",
    "ax.set_ylabel('Amplitude [dB]')\n",
    "ax.margins(0, 0.1)\n",
    "ax.grid(which='both', axis='both')\n",
    "_ = ax.set_xlim([0, ssm_params['Fs']/2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67db84e1",
   "metadata": {},
   "source": [
    "Apply the filter we designed with the above frequency response and visualize the filtered time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "351fa6c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This preserves the phase by filtering forward and backward\n",
    "y_filt = signal.sosfiltfilt(sos, np.squeeze(ssm_params['y']))\n",
    "\n",
    "# Visualize the filtered time series\n",
    "plt.plot(y_filt)\n",
    "plt.xlabel('Time (samples)')\n",
    "plt.ylabel('Arbitrary units')\n",
    "_ = plt.title('Filtered data in 10-14 Hz')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed213e27",
   "metadata": {},
   "source": [
    "Next, we apply Hilbert transform to obtain the complex analytic signal. Also recompute the smoothing estimates from state-space modeling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a79977d",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_hilb = signal.hilbert(y_filt)\n",
    "\n",
    "# Kalman smoothing using saved parameters that have been learned after 50 EM iterations\n",
    "kalman_results = Osc(**ssm_params).kalman_filt_smooth(return_dict=True)\n",
    "y_ssm = kalman_results['x_t_n'][2, :] + 1j * kalman_results['x_t_n'][3, :]  # form a complex signal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc869ffa",
   "metadata": {},
   "source": [
    "Compare the analytic signal to the state-space modeling smoothing estimates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "051c2a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 1, figsize=(10, 5))\n",
    "ax[0].plot(y_ssm.real, label='Oscillator 2 - Real')\n",
    "ax[0].plot(y_hilb.real, label='Analytic signal - Real')\n",
    "ax[1].plot(y_ssm.imag, label='Oscillator 2 - Imaginary')\n",
    "ax[1].plot(y_hilb.imag, label='Analytic signal - Imaginary')\n",
    "\n",
    "for axx in ax:\n",
    "    axx.legend(loc='lower left')\n",
    "    axx.set_ylabel('Arbitrary units')\n",
    "ax[-1].set_xlabel('Time (samples)')\n",
    "_ = ax[0].set_title('Kalman smoothing estimates and Hilbert analytic signal')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1147f5f5",
   "metadata": {},
   "source": [
    "We can compare the extracted oscillation instantaneous amplitudes and phases as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe1f97bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 1, figsize=(10, 5))\n",
    "ax[0].plot(np.abs(y_ssm), label='Oscillator 2')\n",
    "ax[0].plot(np.abs(y_hilb), label='Analytic signal')\n",
    "ax[1].plot(np.unwrap(np.angle(y_ssm)), label='Oscillator 2')  # unwrap phase to see the difference more clearly\n",
    "ax[1].plot(np.unwrap(np.angle(y_hilb)), label='Analytic signal')\n",
    "\n",
    "for axx in ax:\n",
    "    axx.legend(loc='upper left')\n",
    "ax[0].set_ylabel('Amplitude (a.u.)')\n",
    "ax[1].set_ylabel('Phase (rad)')\n",
    "ax[-1].set_xlabel('Time (samples)')\n",
    "_ = ax[0].set_title('Amplitude and phase comparisons')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28de0e3e",
   "metadata": {},
   "source": [
    "The take-home message is that the best \"Hilbert transform after bandpass filtering\" can do is to (almost) par the state-space modeling estimates.\n",
    "\n",
    "However, the analytic approach with Hilbert transform has a few common pitfalls:\n",
    "- The oscillation needs to be strong and to dominate the filtered frequency range.\n",
    "- The filter needs to be carefully designed and tuned to the underlying oscillation.\n",
    "- Even then, the instantaneous amplitude of filtered signal is often smoother than reality.\n",
    "\n",
    "On the other hand, state-space modeling provides reliable estimates after EM learning, which adjusts to the observed data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79f1bdac",
   "metadata": {},
   "source": [
    "## Exercise 3: searching for oscillations in data\n",
    "Learning objective: learn to use and interpret the oscillator search methods in SOMATA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fbf9a74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the oscillator basic model in SOMATA\n",
    "from somata.basic_models import OscillatorModel as Osc\n",
    "\n",
    "# Import the two oscillator search methods implemented in SOMATA\n",
    "from somata.oscillator_search import IterativeOscillatorModel as IterOsc\n",
    "from somata.oscillator_search import DecomposedOscillatorModel as DecOsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4703fc02",
   "metadata": {},
   "source": [
    "### 3.1 Simulate oscillatory data with state-space models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9aff7c44",
   "metadata": {},
   "source": [
    "First create an oscillator model instance with two oscillators and simulate $10$-s data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02c084ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "o1 = Osc(a=[0.996, 0.95], freq=[0.1, 10], sigma2=[0.4, 0.2], R=1.2, Fs=100)\n",
    "print(o1)\n",
    "\n",
    "np.random.seed(1)  # for reproducibility\n",
    "_, y = o1.simulate(duration=10)\n",
    "\n",
    "plt.plot(y.T)\n",
    "plt.xlabel('Time (samples)')\n",
    "plt.ylabel('Arbitrary units')\n",
    "_ = plt.title('Simulated time series data')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e851fda5",
   "metadata": {},
   "source": [
    "### 3.2 Iterative oscillator search algorithm (iOsc+)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b59b5f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a iterative oscillator model object\n",
    "io1 = IterOsc(y, o1.Fs, noise_start=None, osc_range=7)\n",
    "# noise_start determines the frequency above which is used to estimate the observation noise; default: (Nyquist - 20 Hz)\n",
    "# osc_range is the maximum number of total oscillators; default: 7"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f49f534",
   "metadata": {},
   "source": [
    "Oscillator search model objects have convenient visualization methods implemented."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbaf7642",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = io1.plot_mtm()  # plot multitaper spectrogram and mean spectrum\n",
    "_ = io1.plot_trace()  # plot raw time trace"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d952bf3",
   "metadata": {},
   "source": [
    "Now run iterations to search for oscillations present in the simulated data while visualizing every iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3978f9f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "io1.iterate(freq_res=1, plot_fit=True, verbose=True)  # this is the iOsc+ algorithm\n",
    "# freq_res is the minimal resolution in Hz from existing frequencies when adding a new oscillator\n",
    "# plot_fit=True plots innovation spectrum and AR fitting during each iteration\n",
    "# verbose=True prints parameters throughout the method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddf4ee0e",
   "metadata": {},
   "source": [
    "Examine the final oscillator model selected by the iOsc+ algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c19e239",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(io1.get_knee_osc())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c135e20",
   "metadata": {},
   "source": [
    "Plot log-likelihood and the selected model (may not be the highest likelihood)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a982aec",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = io1.plot_log_likelihoods()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "421485f2",
   "metadata": {},
   "source": [
    "Plot the fitted oscillators in the frequency domain with both theoretical and empirical spectra.\n",
    "\n",
    "Note that the empirical spectrum is based on the smoothing estimates of the hidden states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab134e04",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = io1.plot_fit_spectra()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2d6ac49",
   "metadata": {},
   "source": [
    "Plot the time traces of estimated $\\mathbf{x}_t$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bcd7da8",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = io1.plot_fit_traces()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ad168bc",
   "metadata": {},
   "source": [
    "### 3.3 Decomposed oscillator search algorithm (dOsc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8674859",
   "metadata": {},
   "source": [
    "dOsc has a very similar constructor method as iOsc+. The `plot_mtm()` and `plot_trace()` methods work the same way, so we skip them here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56ca12d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "do1 = DecOsc(y, o1.Fs, noise_start=None, osc_range=7)\n",
    "# noise_start determines the frequency above which is used to estimate the observation noise; default: (Nyquist - 20 Hz)\n",
    "# osc_range is the maximum number of total oscillators; default: 7"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61dc7157",
   "metadata": {},
   "source": [
    "The method call to run the iterations to search for oscillations is also called `iterate()`, albeit with slightly different arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ddcbb8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "do1.iterate(plot_fit=True)  # this is the dOsc algorithm\n",
    "# plot_fit=True plots fitted theoretical spectra during each iteration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c19fb79",
   "metadata": {},
   "source": [
    "Also examine the final oscillator model selected by the dOsc algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82e92fef",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(do1.get_knee_osc())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8febd6f",
   "metadata": {},
   "source": [
    "The same sanity check methods are available for the dOsc algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d6c63ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = do1.plot_log_likelihoods()\n",
    "_ = do1.plot_fit_spectra()\n",
    "_ = do1.plot_fit_traces()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd1462ba",
   "metadata": {},
   "source": [
    "### 3.4 Diagnostic residual plots and statistical tests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6aaa63f",
   "metadata": {},
   "source": [
    "We can plot the model fitting residuals (one-step prediction error), i.e., $y_t - \\hat{y}_{t|t-1}$, and run diagnostic statistical tests to check for auto-correlations and normality.\n",
    "\n",
    "The same `diagnose_residual()` method applies to both oscillator search methods. Below we demonstrate it for `DecomposedOscillatorModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85a60f79",
   "metadata": {},
   "outputs": [],
   "source": [
    "do1.diagnose_residual()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a820ba7",
   "metadata": {},
   "source": [
    "## You are now ready to use SOMATA for state-space modeling!"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}