{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "rect = lambda x: (np.abs(x) <= 1/2) * 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## логика"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def swap_halves(a):\n",
    "    n = len(a)\n",
    "    assert n % 2 == 0\n",
    "    \n",
    "    m = n // 2\n",
    "    \n",
    "    b = a\n",
    "    b = np.r_[\n",
    "        b[m:],\n",
    "        b[:m],\n",
    "    ]\n",
    "    \n",
    "    return b\n",
    "\n",
    "def left_right_pad(a, m):\n",
    "    n = len(a)\n",
    "    \n",
    "    assert n % 2 == 0\n",
    "    assert m % 2 == 0\n",
    "    assert m >= n\n",
    "    \n",
    "    l = (m - n) // 2\n",
    "    r = l + n\n",
    "    \n",
    "    b = np.zeros(m, dtype=a.dtype)\n",
    "    b[l:r] = a\n",
    "    \n",
    "    return b\n",
    "\n",
    "def left_right_unpad(b, n):\n",
    "    m = len(b)\n",
    "    \n",
    "    assert n % 2 == 0\n",
    "    assert m % 2 == 0\n",
    "    assert m >= n\n",
    "    \n",
    "    l = (m - n) // 2\n",
    "    r = l + n\n",
    "    \n",
    "    return b[l:r]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def finite_fourier_1(xs1, ys1, m):\n",
    "    n = len(xs1)\n",
    "    assert n % 2 == 0\n",
    "    assert m % 2 == 0\n",
    "\n",
    "    a1 = xs1[-1]\n",
    "    h_x1 = xs1[1] - xs1[0]\n",
    "\n",
    "    fft_arg = ys1\n",
    "    fft_arg = left_right_pad(fft_arg, m)\n",
    "    fft_arg = swap_halves(fft_arg)\n",
    "    \n",
    "    fft_res = np.fft.fft(fft_arg)\n",
    "    ys2 = fft_res * h_x1\n",
    "    ys2 = swap_halves(ys2)\n",
    "    ys2 = left_right_unpad(ys2, n)\n",
    "\n",
    "\n",
    "    a2 = n ** 2 / (4 * a1 * m)\n",
    "    xs2 = np.linspace(-a2, a2, n)\n",
    "    \n",
    "    return xs2, ys2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def finite_fourier_2(xs1, ys1, xs2):\n",
    "    n = len(xs1)\n",
    "    h_x1 = xs1[1] - xs1[0]\n",
    "    \n",
    "    # first dimension - x\n",
    "    x_2d = np.broadcast_to(xs1[:, np.newaxis], (n, n))\n",
    "    \n",
    "    # second dimension - u\n",
    "    u_2d = np.broadcast_to(xs2[np.newaxis, :], (n, n))\n",
    "    \n",
    "    # exp values\n",
    "    A = np.exp((-2 * np.pi * 1j) * x_2d * u_2d)\n",
    "    \n",
    "    # scale rows by f(x)\n",
    "    A = A * np.broadcast_to(ys1[:, np.newaxis], (n, n))\n",
    "\n",
    "    int_weights = np.ones(n)\n",
    "    int_weights[0] = 1/2\n",
    "    int_weights[-1] = 1/2\n",
    "    int_weights *= h_x1\n",
    "    \n",
    "    # scale rows by int_weights\n",
    "    A = A * np.broadcast_to(int_weights[:, np.newaxis], (n, n))\n",
    "\n",
    "    ys2 = np.sum(A, axis=0)\n",
    "    \n",
    "    return ys2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## примеры"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = lambda x: np.exp(-x**2)\n",
    "F = lambda x: np.sqrt(np.pi) * np.exp(-x**2 * np.pi**2)\n",
    "\n",
    "a = 5\n",
    "b_given = 1.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = np.sinc\n",
    "F = rect\n",
    "\n",
    "a = 5\n",
    "b_given = 1.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = lambda x: np.exp(2 * np.pi * 1j * x) + np.exp(-5 * np.pi * 1j * x)\n",
    "F = lambda x: (4 * x + 3) * np.sin(8 * np.pi * x) / np.pi / (x - 1) / (2 * x + 5)\n",
    "\n",
    "a = 1\n",
    "b_given = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.6\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126332 characters]\n",
      "text/plain": [
       "<Figure size 1920x1080 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n = 1 << 10\n",
    "m = 1 << 15\n",
    "\n",
    "xs1 = np.linspace(-a, a, n)\n",
    "ys1 = f(xs1)\n",
    "ys1 = np.array(ys1, dtype=complex)\n",
    "\n",
    "xs2, ys2 = finite_fourier_1(xs1, ys1, m)\n",
    "print(xs2[-1])\n",
    "\n",
    "ys3 = F(xs2)\n",
    "ys3 = np.array(ys3, dtype=complex)\n",
    "\n",
    "subplot(2, 2, 1)\n",
    "plot(xs1, np.abs(ys1))\n",
    "xlabel('x')\n",
    "title('$\\\\left|f\\\\right|$')\n",
    "\n",
    "subplot(2, 2, 3)\n",
    "plot(xs1, np.angle(ys1))\n",
    "ylim(-np.pi * 1.1, np.pi * 1.1)\n",
    "xlabel('x')\n",
    "title('$arg f$')\n",
    "\n",
    "subplot(2, 2, 2)\n",
    "plot(xs2, np.abs(ys2))\n",
    "plot(xs2, np.abs(ys3))\n",
    "xlabel('u')\n",
    "title('$\\\\left|F\\\\right|$')\n",
    "\n",
    "subplot(2, 2, 4)\n",
    "plot(xs2, np.angle(ys2))\n",
    "plot(xs2, np.angle(ys3))\n",
    "ylim(-np.pi * 1.1, np.pi * 1.1)\n",
    "xlabel('u')\n",
    "title('$arg F$')\n",
    "\n",
    "gcf().suptitle('FINITE FFT LIBRARY')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 131228 characters]\n",
      "text/plain": [
       "<Figure size 1920x1080 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xs1 = np.linspace(-a, a, n)\n",
    "ys1 = vectorize(f)(xs1)\n",
    "ys1 = np.array(ys1, dtype=complex)\n",
    "\n",
    "xs2 = np.linspace(-b_given, b_given, n)\n",
    "ys2 = finite_fourier_2(xs1, ys1, xs2)\n",
    "\n",
    "ys3 = F(xs2)\n",
    "ys3 = np.array(ys3, dtype=complex)\n",
    "\n",
    "subplot(2, 2, 1)\n",
    "plot(xs1, np.abs(ys1))\n",
    "xlabel('x')\n",
    "title('$\\\\left|f\\\\right|$')\n",
    "\n",
    "subplot(2, 2, 3)\n",
    "plot(xs1, np.angle(ys1))\n",
    "ylim(-np.pi * 1.1, np.pi * 1.1)\n",
    "xlabel('x')\n",
    "title('$arg f$')\n",
    "\n",
    "subplot(2, 2, 2)\n",
    "plot(xs2, np.abs(ys2))\n",
    "plot(xs2, np.abs(ys3))\n",
    "xlabel('u')\n",
    "title('$\\\\left|F\\\\right|$')\n",
    "\n",
    "subplot(2, 2, 4)\n",
    "plot(xs2, np.angle(ys2))\n",
    "plot(xs2, np.angle(ys3))\n",
    "ylim(-np.pi * 1.1, np.pi * 1.1)\n",
    "xlabel('u')\n",
    "title('$arg F$')\n",
    "\n",
    "gcf().suptitle('INTEGRAL FFT')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## логика"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def swap_halves_2d(a):\n",
    "    n = len(a)\n",
    "    assert n % 2 == 0\n",
    "    \n",
    "    m = n // 2\n",
    "    \n",
    "    b = a\n",
    "    b = np.r_[\n",
    "        b[m:, :],\n",
    "        b[:m, :],\n",
    "    ]\n",
    "    b = np.c_[\n",
    "        b[:, m:],\n",
    "        b[:, :m],\n",
    "    ]\n",
    "    \n",
    "    return b\n",
    "\n",
    "def left_right_pad_2d(a, m):\n",
    "    n = len(a)\n",
    "    \n",
    "    assert n % 2 == 0\n",
    "    assert m % 2 == 0\n",
    "    assert m >= n\n",
    "    \n",
    "    l = (m - n) // 2\n",
    "    r = l + n\n",
    "    \n",
    "    b = np.zeros((m, m), dtype=a.dtype)\n",
    "    b[l:r, l:r] = a\n",
    "    \n",
    "    return b\n",
    "\n",
    "def left_right_unpad_2d(b, n):\n",
    "    m = len(b)\n",
    "    \n",
    "    assert n % 2 == 0\n",
    "    assert m % 2 == 0\n",
    "    assert m >= n\n",
    "    \n",
    "    l = (m - n) // 2\n",
    "    r = l + n\n",
    "    \n",
    "    return b[l:r, l:r]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def finite_fourier_1_2d(xs1, ys1, m):\n",
    "    n = len(xs1)\n",
    "    assert n % 2 == 0\n",
    "    assert m % 2 == 0\n",
    "\n",
    "    a1 = xs1[-1]\n",
    "    h_x1 = xs1[1] - xs1[0]\n",
    "\n",
    "    fft_arg = ys1\n",
    "    fft_arg = left_right_pad_2d(fft_arg, m)\n",
    "    fft_arg = swap_halves_2d(fft_arg)\n",
    "    \n",
    "    fft_res = np.fft.fft2(fft_arg)\n",
    "    ys2 = fft_res * h_x1 * h_x1\n",
    "    ys2 = swap_halves_2d(ys2)\n",
    "    ys2 = left_right_unpad_2d(ys2, n)\n",
    "\n",
    "\n",
    "    a2 = n ** 2 / (4 * a1 * m)\n",
    "    xs2 = np.linspace(-a2, a2, n)\n",
    "    \n",
    "    return xs2, ys2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def finite_fourier_2_2d(xs1, ys1, xs2):\n",
    "    n = len(xs1)\n",
    "    h_x1 = xs1[1] - xs1[0]\n",
    "    \n",
    "    shape = (n, n, n, n)\n",
    "    \n",
    "    # first dimension - x\n",
    "    x_4d = np.broadcast_to(xs1[:, np.newaxis, np.newaxis, np.newaxis], shape)\n",
    "    # second dimension - y\n",
    "    y_4d = np.broadcast_to(xs1[np.newaxis, :, np.newaxis, np.newaxis], shape)\n",
    "    \n",
    "    # third dimension - u\n",
    "    u_4d = np.broadcast_to(xs2[np.newaxis, np.newaxis, :, np.newaxis], shape)\n",
    "    # forth dimension - v\n",
    "    v_4d = np.broadcast_to(xs2[np.newaxis, np.newaxis, np.newaxis, :], shape)\n",
    "    \n",
    "    # exp values\n",
    "    A = np.exp((-2 * np.pi * 1j) * (x_4d * u_4d + y_4d * v_4d))\n",
    "    \n",
    "    # scale d1 and d2 by f(x, y)\n",
    "    A = A * np.broadcast_to(ys1[:, :, np.newaxis, np.newaxis], shape)\n",
    "\n",
    "    int_weights = np.ones(n)\n",
    "    int_weights[0] = 1/2\n",
    "    int_weights[-1] = 1/2\n",
    "    int_weights *= h_x1\n",
    "    \n",
    "    # scale d1 by int_weights\n",
    "    A = A * np.broadcast_to(int_weights[:, np.newaxis, np.newaxis, np.newaxis], shape)\n",
    "    # scale d2 by int_weights\n",
    "    A = A * np.broadcast_to(int_weights[np.newaxis, :, np.newaxis, np.newaxis], shape)\n",
    "\n",
    "    A = np.sum(A, axis=0)\n",
    "    A = np.sum(A, axis=0)\n",
    "    \n",
    "    return A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## примеры"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_2d = lambda a: np.exp(-(a[:, :, 0]**2 + a[:, :, 1]**2))\n",
    "F_2d = lambda a: np.pi * np.exp(-(a[:, :, 0]**2 + a[:, :, 1]**2) * np.pi**2)\n",
    "\n",
    "y1_max = 1\n",
    "y2_max = np.pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_2d = lambda a: np.sinc(a[:, :, 0]) * np.sinc(a[:, :, 1])\n",
    "F_2d = lambda a: rect(a[:, :, 0]) * rect(a[:, :, 1])\n",
    "\n",
    "y1_max = 1\n",
    "y2_max = 1\n",
    "\n",
    "a = 5\n",
    "b = 0.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_1d = lambda x: np.exp(2 * np.pi * 1j * x) + np.exp(-5 * np.pi * 1j * x)\n",
    "F_1d = lambda x: (4 * x + 3) * np.sin(8 * np.pi * x) / np.pi / (x - 1) / (2 * x + 5)\n",
    "\n",
    "\n",
    "f_2d = lambda a: f_1d(a[:, :, 0]) * f_1d(a[:, :, 1])\n",
    "F_2d = lambda a: F_1d(a[:, :, 0]) * F_1d(a[:, :, 1])\n",
    "\n",
    "y1_max = 2\n",
    "y2_max = 1\n",
    "\n",
    "a = 1\n",
    "b = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 107960 characters]\n",
      "text/plain": [
       "<Figure size 1920x1080 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n = 1 << 6\n",
    "m = 1 << 8\n",
    "\n",
    "xs1 = np.linspace(-a, a, n)\n",
    "X1 = np.reshape(list(itertools.product(xs1, xs1)), (n, n, 2))\n",
    "ys1 = f_2d(X1)\n",
    "\n",
    "xs2, ys2 = finite_fourier_1_2d(xs1, ys1, m)\n",
    "print(xs2[-1])\n",
    "X2 = np.reshape(list(itertools.product(xs2, xs2)), (n, n, 2))\n",
    "\n",
    "ys3 = F_2d(X2)\n",
    "ys3 = np.array(ys3, dtype=complex)\n",
    "\n",
    "subplot(2, 3, 1)\n",
    "imshow(\n",
    "    np.abs(ys1), \n",
    "    extent=[xs1[0], xs1[-1], xs1[0], xs1[-1]],\n",
    "    vmin=0,\n",
    "    vmax=y1_max,\n",
    ")\n",
    "colorbar()\n",
    "title('$\\\\left|f\\\\right|$')\n",
    "\n",
    "subplot(2, 3, 4)\n",
    "imshow(\n",
    "    np.angle(ys1),\n",
    "    extent=[xs1[0], xs1[-1], xs1[0], xs1[-1]],\n",
    "    vmin=-np.pi, \n",
    "    vmax=np.pi,\n",
    ")\n",
    "colorbar()\n",
    "title('$arg f$')\n",
    "\n",
    "subplot(2, 3, 2)\n",
    "imshow(\n",
    "    np.abs(ys2), \n",
    "    extent=[xs2[0], xs2[-1], xs2[0], xs2[-1]],\n",
    "    vmin=0,\n",
    "    vmax=y2_max,\n",
    ")\n",
    "colorbar()\n",
    "title('$\\\\left|\\hat F\\\\right|$')\n",
    "\n",
    "subplot(2, 3, 5)\n",
    "imshow(\n",
    "    np.angle(ys2), \n",
    "    extent=[xs2[0], xs2[-1], xs2[0], xs2[-1]],\n",
    "    vmin=-np.pi, \n",
    "    vmax=np.pi,\n",
    ")\n",
    "colorbar()\n",
    "title('$arg \\hat F$')\n",
    "\n",
    "subplot(2, 3, 3)\n",
    "imshow(\n",
    "    np.abs(ys3), \n",
    "    extent=[xs2[0], xs2[-1], xs2[0], xs2[-1]],\n",
    "    vmin=0,\n",
    "    vmax=y2_max,\n",
    ")\n",
    "colorbar()\n",
    "title('$\\\\left|F\\\\right|$')\n",
    "\n",
    "subplot(2, 3, 6)\n",
    "imshow(\n",
    "    np.angle(ys3), \n",
    "    extent=[xs2[0], xs2[-1], xs2[0], xs2[-1]],\n",
    "    vmin=-np.pi, \n",
    "    vmax=np.pi,\n",
    ")\n",
    "colorbar()\n",
    "title('$arg F$')\n",
    "\n",
    "gcf().suptitle('FINITE FFT LIBRARY')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 102168 characters]\n",
      "text/plain": [
       "<Figure size 1920x1080 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n = 1 << 6\n",
    "\n",
    "xs1 = np.linspace(-a, a, n)\n",
    "X1 = np.reshape(list(itertools.product(xs1, xs1)), (n, n, 2))\n",
    "ys1 = f_2d(X1)\n",
    "\n",
    "xs2 = np.linspace(-b, b, n)\n",
    "ys2 = finite_fourier_2_2d(xs1, ys1, xs2)\n",
    "\n",
    "ys3 = F_2d(X2)\n",
    "ys3 = np.array(ys3, dtype=complex)\n",
    "\n",
    "subplot(2, 3, 1)\n",
    "imshow(\n",
    "    np.abs(ys1), \n",
    "    extent=[xs1[0], xs1[-1], xs1[0], xs1[-1]],\n",
    "    vmin=0,\n",
    "    vmax=y1_max,\n",
    ")\n",
    "colorbar()\n",
    "title('$\\\\left|f\\\\right|$')\n",
    "\n",
    "subplot(2, 3, 4)\n",
    "imshow(\n",
    "    np.angle(ys1),\n",
    "    extent=[xs1[0], xs1[-1], xs1[0], xs1[-1]],\n",
    "    vmin=-np.pi, \n",
    "    vmax=np.pi,\n",
    ")\n",
    "colorbar()\n",
    "title('$arg f$')\n",
    "\n",
    "subplot(2, 3, 2)\n",
    "imshow(\n",
    "    np.abs(ys2), \n",
    "    extent=[xs2[0], xs2[-1], xs2[0], xs2[-1]],\n",
    "    vmin=0,\n",
    "    vmax=y2_max,\n",
    ")\n",
    "colorbar()\n",
    "title('$\\\\left|\\hat F\\\\right|$')\n",
    "\n",
    "subplot(2, 3, 5)\n",
    "imshow(\n",
    "    np.angle(ys2), \n",
    "    extent=[xs2[0], xs2[-1], xs2[0], xs2[-1]],\n",
    "    vmin=-np.pi, \n",
    "    vmax=np.pi,\n",
    ")\n",
    "colorbar()\n",
    "title('$arg \\hat F$')\n",
    "\n",
    "subplot(2, 3, 3)\n",
    "imshow(\n",
    "    np.abs(ys3), \n",
    "    extent=[xs2[0], xs2[-1], xs2[0], xs2[-1]],\n",
    "    vmin=0,\n",
    "    vmax=y2_max,\n",
    ")\n",
    "colorbar()\n",
    "title('$\\\\left|F\\\\right|$')\n",
    "\n",
    "subplot(2, 3, 6)\n",
    "imshow(\n",
    "    np.angle(ys3), \n",
    "    extent=[xs2[0], xs2[-1], xs2[0], xs2[-1]],\n",
    "    vmin=-np.pi, \n",
    "    vmax=np.pi,\n",
    ")\n",
    "colorbar()\n",
    "title('$arg F$')\n",
    "\n",
    "gcf().suptitle('INTEGRAL FFT')\n",
    "None"
   ]
  }
 ],
 "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
