{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "03aa22e7",
   "metadata": {},
   "source": [
    "# Stochastic Response\n",
    "Richard M. Murray, 6 Feb 2022 (updated 9 Feb 2023)\n",
    "\n",
    "This notebook illustrates the implementation of random processes and stochastic response.  We focus on a system of the form\n",
    "$$\n",
    "  \\dot X = A X + F V \\qquad X \\in {\\mathbb R}^n\n",
    "$$\n",
    "\n",
    "where $V$ is a white noise process and the system is a first order linear system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "902af902",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt\n",
    "import control as ct\n",
    "from math import sqrt, exp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77d58303",
   "metadata": {},
   "source": [
    "## First order linear system\n",
    "\n",
    "We start by looking at the stochastic response for a first order linear system\n",
    "\n",
    "$$\n",
    "\\begin{gathered}\n",
    "  \\dot X = -a X + V, \\qquad Y = C X \\\\\n",
    "  \\mathbb{E}(V) = 0, \\quad \\mathbb{E}(V^\\mathsf{T}(t_1) V(t_2)) = 0.1\\, \\delta(t_1 - t_2)\n",
    "\\end{gathered}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60192a8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First order system\n",
    "a = 1\n",
    "c = 1\n",
    "sys = ct.tf(c, [1, a])\n",
    "\n",
    "# Create the time vector that we want to use\n",
    "Tf = 5\n",
    "T = np.linspace(0, Tf, 1000)\n",
    "dt = T[1] - T[0]\n",
    "\n",
    "# Create the basis for a white noise signal\n",
    "# Note: use sqrt(Q/dt) for desired covariance\n",
    "Q = np.array([[0.1]])\n",
    "# V = np.random.normal(0, sqrt(Q[0,0]/dt), T.shape)\n",
    "V = ct.white_noise(T, Q)\n",
    "\n",
    "plt.plot(T, V[0])\n",
    "plt.xlabel('Time [s]')\n",
    "plt.ylabel('$V$');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4629e2c",
   "metadata": {},
   "source": [
    "Note that the magnitude of the signal seems to be much larger than $Q$.  This is because we have a Guassian process $\\implies$ the signal consists of a sequence of \"impulse-like\" functions that have magnitude that increases with the time step $dt$ as $1/\\sqrt{dt}$ (this gives covariance $\\mathbb{E}(V(t_1) V^T(t_2)) = Q \\delta(t_2 - t_1)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23319dc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate the sample properties and make sure they match\n",
    "print(\"mean(V) [0.0] = \", np.mean(V))\n",
    "print(\"cov(V) * dt [%0.3g] = \" % Q, np.round(np.cov(V), decimals=3) * dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bdaaccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Response of the first order system\n",
    "# Scale white noise by sqrt(dt) to account for impulse\n",
    "T, Y = ct.forced_response(sys, T, V)\n",
    "plt.plot(T, Y)\n",
    "plt.xlabel('Time [s]')\n",
    "plt.ylabel('$Y$');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ead0232e",
   "metadata": {},
   "source": [
    "This is a first order system, and so we can use the calculation from the course\n",
    "notes to compute the analytical correlation function and compare this to the \n",
    "sampled data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d31ce324",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare static properties to what we expect analytically\n",
    "def r(tau):\n",
    "    return c**2 * Q / (2 * a) * exp(-a * abs(tau))\n",
    "    \n",
    "print(\"* mean(Y) [%0.3g] = %0.3g\" % (0, np.mean(Y)))\n",
    "print(\"* cov(Y) [%0.3g] = %0.3g\" % (r(0), np.cov(Y)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cf5a4b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Correlation function for the input\n",
    "# Scale by dt to take time step into account\n",
    "# r_V = sp.signal.correlate(V, V) * dt / Tf\n",
    "# tau = sp.signal.correlation_lags(len(V), len(V)) * dt\n",
    "tau, r_V = ct.correlation(T, V)\n",
    "\n",
    "plt.plot(tau, r_V, 'r-')\n",
    "plt.xlabel(r'$\\tau$')\n",
    "plt.ylabel(r'$r_V(\\tau)$');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62af90a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Correlation function for the output\n",
    "# r_Y = sp.signal.correlate(Y, Y) * dt / Tf\n",
    "# tau = sp.signal.correlation_lags(len(Y), len(Y)) * dt\n",
    "tau, r_Y = ct.correlation(T, Y)\n",
    "plt.plot(tau, r_Y)\n",
    "plt.xlabel(r'$\\tau$')\n",
    "plt.ylabel(r'$r_Y(\\tau)$')\n",
    "\n",
    "# Compare to the analytical answer\n",
    "plt.plot(tau, [r(t)[0, 0] for t in tau], 'k--');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a2785e9",
   "metadata": {},
   "source": [
    "The analytical curve may or may not line up that well with the correlation function based on the sample.  Try running the code again from the top to see how things change based on the specific random sequence chosen at the start.\n",
    "\n",
    "Note: the _right_ way to compute the correlation function would be to run a lot of different samples of white noise filtered through the system dynamics and compute $R(t_1, t_2)$ across those samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd5dfc75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# As a crude approximation, compute the average correlation\n",
    "r_avg = np.zeros_like(r_Y)\n",
    "for i in range(100):\n",
    "    V = ct.white_noise(T, Q)\n",
    "    _, Y = ct.forced_response(sys, T, V)\n",
    "    tau, r_Y = ct.correlation(T, Y)\n",
    "    r_avg = r_avg + r_Y\n",
    "r_avg = r_avg / i\n",
    "plt.plot(tau, r_avg)\n",
    "plt.xlabel(r'$\\tau$')\n",
    "plt.ylabel(r'$r_Y(\\tau)$')\n",
    "\n",
    "# Compare to the analytical answer\n",
    "plt.plot(tau, [r(t)[0, 0] for t in tau], 'k--');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f07ec584",
   "metadata": {},
   "source": [
    "## Dryden gust model\n",
    "\n",
    "Friedland, _Control Systems Design_, Example 10B\n",
    "\n",
    "Based on experimental data, the power spectral density for the vertical component of random wind velocity in turbulent air can be modeled as\n",
    "$$\n",
    "S(\\omega) = \\sigma_\\text{z}^2 T \\frac{1 + 3 (\\omega T)^2}{[1 + (\\omega T)^2]^2},\n",
    "$$\n",
    "where $\\sigma_\\text{z}$ and $T$ are parameters that depend on the wind characteristics.\n",
    "\n",
    "This power spectral density can be modeled using white noise by running it through a linear system with transfer fucntion\n",
    "$$\n",
    "H(s) = \\frac{1 + \\sqrt{3} T}{(1 + T s)^2}.\n",
    "$$\n",
    "A state space realization for this transfer function is given by\n",
    "$$\n",
    "\\begin{aligned}\n",
    "  \\dot X &= \\begin{bmatrix} 0 & 1 \\\\ -\\frac{1}{T^2} & -\\frac{2}{T} \\end{bmatrix} X \n",
    "    + \\begin{bmatrix} 0 \\\\ 1 \\end{bmatrix} V \\\\\n",
    "  Y &= \\begin{bmatrix} \\frac{1}{T^2} & \\frac{\\sqrt{3}}{T} \\end{bmatrix}\n",
    "  \\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d09fc03a",
   "metadata": {},
   "source": [
    "To create a disturbance signal with the characteristics of the Dryden gust model, we create a linear system with the given parameters and computing the input/output response to white noise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8df16a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_z = 1\n",
    "T = 1\n",
    "filter = ct.ss([[0, 1], [-1/T**2, -2/T]], [[0], [1]], [[1/T**2, sqrt(3)/T]], 0)\n",
    "\n",
    "timepts = np.linspace(0, 10, 1000)\n",
    "V = ct.white_noise(timepts, sigma_z**2)\n",
    "resp = ct.input_output_response(filter, timepts, V)\n",
    "\n",
    "plt.plot(resp.time, resp.outputs);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d6604ee",
   "metadata": {},
   "source": [
    "We can compute the correlation function and power spectral density to confirm that we match the desired characteristics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "febc8b80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the correlation function\n",
    "tau, R = ct.correlation(resp.time, resp.outputs)\n",
    "\n",
    "# Analytical expression for the correlation function (see Friedland)\n",
    "def dryden_corrfcn(tau, sigma_z=1, T=1):\n",
    "    return sigma_z**2 * np.exp(-np.abs(tau)/T) * (1- np.abs(tau)/(2*T))\n",
    "\n",
    "# Plot the correlation function\n",
    "fig, axs = plt.subplots(1, 2)\n",
    "axs[0].plot(tau, R)\n",
    "axs[0].plot(tau, dryden_corrfcn(tau))\n",
    "axs[0].set_xlabel(\"$\\tau$\")\n",
    "axs[0].set_ylabel(\"$r(\\tau)$\")\n",
    "axs[0].set_title(\"Correlation function\")\n",
    "\n",
    "# Compute the power spectral density\n",
    "dt = timepts[1] - timepts[0]\n",
    "S = sp.fft.rfft(R) * dt * 2            # rfft returns omega >= 0 => muliple mag by 2\n",
    "omega = sp.fft.rfftfreq(R.size, dt)\n",
    "\n",
    "# Analytical expression for the correlation function (see Friedland)\n",
    "def dryden_psd(omega, sigma_z=1., T=1.):\n",
    "    return sigma_z**2 * T * (1 + 3 * (omega * T)**2) / (1 + (omega * T)**2)**2\n",
    "\n",
    "# Plot the power spectral density\n",
    "axs[1].loglog(omega[1:], np.abs(S[1:]))\n",
    "axs[1].loglog(omega[1:], dryden_psd(omega[1:]))\n",
    "axs[0].set_xlabel(\"$\\omega$ [rad/sec]$\")\n",
    "axs[0].set_ylabel(\"$S(\\omega)$\")\n",
    "axs[0].set_title(\"Power spectral density\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1516ff6a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}