{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Plotting and manipulating FFTs for filtering\n",
    "\n",
    "Plot the power of the FFT of a signal and inverse FFT back to reconstruct\n",
    "a signal.\n",
    "\n",
    "This example demonstrate :func:`scipy.fft.fft`,\n",
    ":func:`scipy.fft.fftfreq` and :func:`scipy.fft.ifft`. It\n",
    "implements a basic filter that is very suboptimal, and should not be\n",
    "used.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate the signal\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x10bbc9dc0>]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 45416 characters]",
      "text/plain": [
       "<Figure size 600x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Seed the random number generator\n",
    "rng = np.random.default_rng(27446968)\n",
    "\n",
    "time_step = 0.02\n",
    "period = 5.0\n",
    "\n",
    "time_vec = np.arange(0, 20, time_step)\n",
    "sig = np.sin(2 * np.pi / period * time_vec) + 0.5 * rng.normal(size=time_vec.size)\n",
    "\n",
    "plt.figure(figsize=(6, 5))\n",
    "plt.plot(time_vec, sig, label=\"Original signal\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute and plot the power\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 33712 characters]",
      "text/plain": [
       "<Figure size 600x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# The FFT of the signal\n",
    "sig_fft = sp.fft.fft(sig)\n",
    "\n",
    "# And the power (sig_fft is of complex dtype)\n",
    "power = np.abs(sig_fft) ** 2\n",
    "\n",
    "# The corresponding frequencies\n",
    "sample_freq = sp.fft.fftfreq(sig.size, d=time_step)\n",
    "\n",
    "# Plot the FFT power\n",
    "plt.figure(figsize=(6, 5))\n",
    "plt.plot(sample_freq, power)\n",
    "plt.xlabel(\"Frequency [Hz]\")\n",
    "plt.ylabel(\"plower\")\n",
    "\n",
    "# Find the peak frequency: we can focus on only the positive frequencies\n",
    "pos_mask = np.where(sample_freq > 0)\n",
    "freqs = sample_freq[pos_mask]\n",
    "peak_freq = freqs[power[pos_mask].argmax()]\n",
    "\n",
    "# Check that it does indeed correspond to the frequency that we generate\n",
    "# the signal with\n",
    "np.allclose(peak_freq, 1.0 / period)\n",
    "\n",
    "# An inner plot to show the peak frequency\n",
    "axes = plt.axes((0.55, 0.3, 0.3, 0.5))\n",
    "plt.title(\"Peak frequency\")\n",
    "plt.plot(freqs[:8], power[pos_mask][:8])\n",
    "plt.setp(axes, yticks=[])\n",
    "\n",
    "# scipy.signal.find_peaks_cwt can also be used for more advanced\n",
    "# peak detection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remove all the high frequencies\n",
    "\n",
    " We now remove all the high frequencies and transform back from\n",
    " frequencies to signal.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/envs/pyscixx/lib/python3.12/site-packages/matplotlib/cbook.py:1719: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return math.isfinite(val)\n",
      "/opt/miniconda3/envs/pyscixx/lib/python3.12/site-packages/matplotlib/cbook.py:1355: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return np.asarray(x, float)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x11cf1ae70>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 70592 characters]",
      "text/plain": [
       "<Figure size 600x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "high_freq_fft = sig_fft.copy()\n",
    "high_freq_fft[np.abs(sample_freq) > peak_freq] = 0\n",
    "filtered_sig = sp.fft.ifft(high_freq_fft)\n",
    "\n",
    "plt.figure(figsize=(6, 5))\n",
    "plt.plot(time_vec, sig, label=\"Original signal\")\n",
    "plt.plot(time_vec, filtered_sig, linewidth=3, label=\"Filtered signal\")\n",
    "plt.xlabel(\"Time [s]\")\n",
    "plt.ylabel(\"Amplitude\")\n",
    "\n",
    "plt.legend(loc=\"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note** This is actually a bad way of creating a filter: such brutal\n",
    "cut-off in frequency space does not control distortion on the signal.\n",
    "\n",
    "Filters should be created using the SciPy filter design code\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
