{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nonlinear Thresholding\n",
    "When noise and signal share a common frequency band, spectral filtering would lead to a loss in signal. For these cases, we need different techniques to reduce the disturibing noise. We start with a reocorde signal $x(t)$, contains some signal $s(t)$ and additive noise $n(t)$:\n",
    "$$\n",
    "x(t) = s(t) + n(t)\n",
    "$$\n",
    "We can rewrite the equation above in time-frequency domain as\n",
    "$$\n",
    "X(t, f) = S(t, f) + N(t, f) ~.\n",
    "$$\n",
    "Assuming that some time-frequnecy coefficients can be associated by noise, we set all coefficients, which a below this threshold, to zero:\n",
    "$$\n",
    "\\tilde{X}(t, f) = \\left\\{\n",
    "    \\begin{array}{@{}ll@{}}\n",
    "         X(t, f) & \\mathrm{if}~ |X(t, f)| \\geq \\beta(f) \\\\\n",
    "          0          & \\mathrm{otherwise}\n",
    "    \\end{array}\\right.\n",
    "    ~,\n",
    "$$\n",
    "where $\\beta(f)$ denotes the threshold function. The threshold function is defined as \n",
    "$$\n",
    "\\beta(f) = \\mathrm{ECDF}_f^{-1} (P = 0.99) ~,\n",
    "$$\n",
    "where $\\mathrm{ECDF}_f^{-1}$ denotes the inverse cumulative distribution function or quantile function, e.g. before the first arrival of earthquake waves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load all required packages\n",
    "import os\n",
    "import numpy as np\n",
    "from scipy.signal import stft, istft\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load our earthquake data and plot\n",
    "d = np.load(os.path.join(os.path.expanduser('~'),'work', 'data', 'events', 'bug2019mgoh_Z.npz'))\n",
    "#d = np.load(os.path.join(os.path.expanduser('~'),'work', 'data', 'events', 'bug2019gbbo_Z.npz'))\n",
    "#d = np.load(os.path.join(os.path.expanduser('~'),'work', 'data', 'events', 'bug2019ibsd_Z.npz'))\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.plot(d['data']);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute spectrogram of data\n",
    "f, t, X = stft(d['data'], fs=1/0.01, nfft=198, nperseg=99)   # Be careful with choice of nfft and nperseg, \n",
    "                                                             # because istft does not work for all pairs\n",
    "print(X.shape)\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.pcolormesh(t, f, np.abs(X))\n",
    "plt.xlabel(\"Frequency (Hz)\")\n",
    "plt.ylabel(\"Time (s)\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Functions for thresholding\n",
    "def threshold(X, quantile=0.99):\n",
    "    \"\"\"\n",
    "    X: time-frequency coefficients\n",
    "    q: quantile [0-1]\n",
    "    \"\"\"\n",
    "    # Loop over frequencies to build threshold function\n",
    "    beta = np.zeros(X.shape[0])\n",
    "    for i in range(X.shape[0]):\n",
    "        beta[i] = np.quantile(np.abs(X[i, :]), q=quantile)\n",
    "    \n",
    "    return beta\n",
    "\n",
    "def modify_spectrogram(X, beta):\n",
    "    # Loop over all items in X and apply threshold\n",
    "    # Task: modify X!\n",
    "    \n",
    "    return X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute STFT of data before first arrival of P\n",
    "_, _, X_thres = stft(d['data'][:2500], fs=1/0.01, nfft=198, nperseg=99)\n",
    "# Estimate threshold fucntion\n",
    "beta = threshold(X_thres)\n",
    "\n",
    "# Modifiy original spectrogram\n",
    "X_mod = modify_spectrogram(X, beta)\n",
    "\n",
    "# Plot modified spectrogram\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.pcolormesh(t, f, np.abs(X_mod))\n",
    "plt.xlabel(\"Frequency (Hz)\")\n",
    "plt.ylabel(\"Time (s)\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inverse STFT of modified spectrogram\n",
    "t, x_mod = istft(X_mod, fs=1/0.01, nfft=198, nperseg=99)\n",
    "\n",
    "# Plot corrected seismogram\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.plot(t, x_mod);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}