{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1c0ca885c60d2d8",
   "metadata": {},
   "source": [
    "# Fit data with a frequency response function\n",
    "\n",
    "Objects from ``thztools``: ``apply_frf``, ``fit``, ``NoiseModel``, ``timebase``, ``wave``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import thztools as thz\n",
    "\n",
    "from matplotlib.figure import figaspect\n",
    "from scipy import stats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4191a224c6ed5cf",
   "metadata": {},
   "source": [
    "## Set the simulation parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d19a67bd5ab2a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 256  # Number of samples\n",
    "dt = 0.05  # Sampling time [ps]\n",
    "\n",
    "sigma_alpha = 1e-4  # Additive noise amplitude [signal units]\n",
    "sigma_beta = 1e-2  # Multiplicative noise amplitude [dimensionless]\n",
    "sigma_tau = 1e-3  # Time base noise amplitude [ps]\n",
    "\n",
    "p0 = (0.5, 1.0)  # Frequency response function parameters (amplitude, delay)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77b4188a8914dd9c",
   "metadata": {},
   "source": [
    "## Simulate an input waveform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "930396cc3130007",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = thz.timebase(n, dt=dt)\n",
    "mu = thz.wave(n, dt=dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89464fda05406570",
   "metadata": {},
   "source": [
    "## Define the frequency response function\n",
    "The frequency response function ``frfun`` rescales the input waveform by ``a`` and delays it by ``eta``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9edcfefcdbac9713",
   "metadata": {},
   "outputs": [],
   "source": [
    "def frfun(omega, a, eta):\n",
    "    return a * np.exp(1j * omega * eta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b4d1233071318d6",
   "metadata": {},
   "source": [
    "## Apply the frequency response function\n",
    "Generate the time-domain output waveform `psi` by applying `frfun` to `mu` with the `apply_frf` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d62ac6199ccf0b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "psi = thz.apply_frf(frfun, mu, dt=dt, args=p0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a48b312c1323fa2b",
   "metadata": {},
   "source": [
    "## Add noise\n",
    "\n",
    "Create an instance `noise_model` of the `NoiseModel` class and use the `noise_sim` method to add simulated noise to each waveform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6150d23e036de3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_model = thz.NoiseModel(\n",
    "    sigma_alpha=sigma_alpha, \n",
    "    sigma_beta=sigma_beta, \n",
    "    sigma_tau=sigma_tau, \n",
    "    dt=dt\n",
    ")\n",
    "\n",
    "x = mu + noise_model.noise_sim(mu, seed=0)\n",
    "y = psi + noise_model.noise_sim(psi)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e469611eab6ad0d0",
   "metadata": {},
   "source": [
    "## Fit the frequency response function to the noisy input and output waveforms\n",
    "\n",
    "Fit the parameters of the frequency response function to the noisy input and output waveforms. The fitted parameters are consistent with the true parameters, and the goodness of fit statistic `result.resnorm` is consistent with a $\\chi^2$-distribution with $\\nu = n_\\mu - n_p - n_a - n_b = 253$ statistical degrees of freedom, where $n_a = 1$ and $n_b = 0$ because we exclude zero frequency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1511b665444f238",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = thz.fit(frfun, x, y, p0, noise_parms=(sigma_alpha, sigma_beta, sigma_tau), dt=dt)\n",
    "print(f\"Amplitude parameter: {result.p_opt[0]:.4f} ± {np.sqrt(result.p_cov[0, 0]):.4f}\")\n",
    "print(f\"Delay parameter: {result.p_opt[1]:.4f} ± {np.sqrt(result.p_cov[1, 1]):.4f}\")\n",
    "print(f\"Goodness of fit: {result.resnorm:.1f}\")\n",
    "print(f\"Degrees of freedom: {result.dof:d}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7346066d03f966a4",
   "metadata": {},
   "source": [
    "## Show the fit residuals\n",
    "\n",
    "As expected, the total-least-squares (TLS) residuals, `result.r_tls`, are Gaussian-distributed, and show no evidence of correlation in the time domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc5e60de716f4eba",
   "metadata": {},
   "outputs": [],
   "source": [
    "norm_res = result.r_tls\n",
    "osm, osr = stats.probplot(\n",
    "    norm_res, fit=False\n",
    ")\n",
    "\n",
    "_, ax = plt.subplots()\n",
    "\n",
    "ax.plot(osr, osm, '.', ms=2)\n",
    "ax.plot([-3, 3], [-3, 3], '--', c='gray')\n",
    "ax.grid()\n",
    "\n",
    "ax.set_xlim(-5, 5)\n",
    "ax.set_ylim(stats.norm.ppf([0.0005, 0.9995]))\n",
    "\n",
    "ax.set_xticks(np.arange(-4, 4.5, 2))\n",
    "ax.set_yticks(stats.norm.ppf([0.005, 0.1, 0.5, 0.9, 0.995]))\n",
    "\n",
    "ax.set_yticklabels(['0.005', '0.1', '0.5', '0.9', '0.995'])\n",
    "\n",
    "ax.set_xlabel('Normed residual')\n",
    "ax.set_ylabel('Probability')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68c88d75441d80f",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots()\n",
    "\n",
    "markerline, stemlines, baseline = ax.stem(\n",
    "    t, norm_res, linefmt='-', markerfmt='.'\n",
    ")\n",
    "markerline.set_markersize(2)\n",
    "stemlines.set_linewidth(0.5)\n",
    "baseline.set_linewidth(1)\n",
    "\n",
    "ax.set_xlim(0, 10)\n",
    "ax.set_ylim(-3.5, 3.5)\n",
    "\n",
    "ax.set_xlabel('Time (ps)')\n",
    "ax.set_ylabel('Normed residual')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}