{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy as sc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1 - Simple FFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19592 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 15024 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# using sc.fft\n",
    "\n",
    "def f(x):\n",
    "    return np.exp(-np.pi * x**2)\n",
    "\n",
    "n_segments = 2**13\n",
    "# FFT works best with 2**n segments\n",
    "x = np.linspace(-10, 10, n_segments)\n",
    "# plt.xlim([-5, 5])\n",
    "fx = f(x)\n",
    "\n",
    "plt.plot(x, fx)\n",
    "plt.show()\n",
    "fk = sc.fft.fft(\n",
    "    fx,\n",
    "    # norm = 'ortho'\n",
    ")\n",
    "k = sc.fft.fftfreq(n_segments, 2/n_segments)\n",
    "plt.scatter(k, np.abs(fk))\n",
    "# plt.xlim([-5, 5])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2 - signal analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43104 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1024 1024\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 14596 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def signal(x, A=1, B=1):\n",
    "    return A * np.sin(2 * np.pi * 2 * x) + B * np.sin(2 * np.pi * 4 * x)\n",
    "\n",
    "n_segments = 2**10\n",
    "\n",
    "domain = [-1, 1]\n",
    "x = np.linspace(domain[0], domain[1], n_segments)\n",
    "\n",
    "fx = signal(x)\n",
    "plt.plot(x, fx)\n",
    "plt.show()\n",
    "\n",
    "## decompose signal\n",
    "fk = sc.fft.fft(fx)\n",
    "# fk = np.fft.fftshift(fk)\n",
    "k = sc.fft.fftfreq(n_segments, 2/n_segments)\n",
    "print(k.size, fk.size)\n",
    "plt.xlim([-10, 10])\n",
    "plt.scatter(k, np.abs(fk))\n",
    "# plt.plot(np.abs(fk))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 3 - Gibbs phenomenon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16208 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16632 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 22220 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def step_function(x, L=1):\n",
    "    def heaviside(x):\n",
    "        ret = np.zeros_like(x)\n",
    "        ret[x >= 0] = 1\n",
    "        return ret\n",
    "    return 2 * heaviside(x / L) - 1\n",
    "\n",
    "n_segments = 2**8\n",
    "length = 5\n",
    "x = np.linspace(-length, length, n_segments)\n",
    "fx = step_function(x, length)\n",
    "plt.plot(x, fx)\n",
    "plt.show()\n",
    "\n",
    "# fourier space\n",
    "fk = sc.fft.fft(fx)\n",
    "k = sc.fft.fftfreq(n_segments, 2/n_segments)\n",
    "plt.scatter(k, np.abs(fk))\n",
    "plt.show()\n",
    "\n",
    "# reconstruct signal\n",
    "fx_reconstructed = np.zeros_like(fk)\n",
    "for i in range(fx_reconstructed.size):\n",
    "    fx_reconstructed[i] = np.sum(\n",
    "        fk * np.exp(1j * k * x[i])\n",
    "    ) * 2 / n_segments\n",
    "\n",
    "fx_reconstructed_easy = sc.fft.ifft(fk)\n",
    "plt.plot(x, fx_reconstructed_easy)\n",
    "plt.plot(x, np.abs(fx_reconstructed))\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 4 - derivatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24528 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 41712 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 38448 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 23520 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def f1(x):\n",
    "    return np.exp(-x**2)\n",
    "\n",
    "def f2(x):\n",
    "    return np.sin(x)\n",
    "\n",
    "def f3(x):\n",
    "    return 1/x\n",
    "\n",
    "def f4(x):\n",
    "    return np.cos(x) * np.exp(-x**2)\n",
    "\n",
    "\n",
    "def fft_derivative(f:callable):\n",
    "    def derivative(x):\n",
    "        fk = sc.fft.fft(f(x))\n",
    "        k = sc.fft.fftfreq(x.size, 2/x.size)\n",
    "        df_dk = 1j * k * fk\n",
    "        df_dx = sc.fft.ifft(df_dk)\n",
    "        return df_dx\n",
    "    return derivative\n",
    "\n",
    "n_segments = 2**7\n",
    "x = np.linspace(-10, 10, n_segments)\n",
    "\n",
    "\n",
    "f1x = f1(x)\n",
    "df1_dx = fft_derivative(f1)(x)\n",
    "plot = plt.plot(x, f1x)\n",
    "plot = plt.plot(x, df1_dx)\n",
    "plt.show()\n",
    "\n",
    "f2x = f2(x)\n",
    "df2_dx = fft_derivative(f2)(x)\n",
    "plot = plt.plot(x, f2x)\n",
    "plot = plt.plot(x, df2_dx)\n",
    "plt.show()\n",
    "\n",
    "f3x = f3(x)\n",
    "df3_dx = fft_derivative(f3)(x)\n",
    "plot = plt.plot(x, f3x)\n",
    "plot = plt.plot(x, df3_dx)\n",
    "plt.show()\n",
    "\n",
    "f4x = f4(x)\n",
    "df4_dx = fft_derivative(f4)(x)\n",
    "plot = plt.plot(x, f4x)\n",
    "plot = plt.plot(x, df4_dx)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "projects-X-9bmgL6",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
