{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Patching\nThis example shows how to use the :py:class:`pylops.signalprocessing.Patch2D`\nand :py:class:`pylops.signalprocessing.Patch3D` operators to perform repeated\ntransforms over small patches of a 2-dimensional or 3-dimensional\narray. The transforms that we apply in this example are the\n:py:class:`pylops.signalprocessing.FFT2D` and\n:py:class:`pylops.signalprocessing.FFT3D` but this operator has been\ndesigned to allow a variety of transforms as long as they operate with signals\nthat are 2- or 3-dimensional in nature, respectively.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\n\nimport pylops\n\nplt.close(\"all\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's start by creating an 2-dimensional array of size $n_x \\times n_t$\ncomposed of 3 parabolic events\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "par = {\"ox\": -140, \"dx\": 2, \"nx\": 140, \"ot\": 0, \"dt\": 0.004, \"nt\": 200, \"f0\": 20}\n\nv = 1500\nt0 = [0.2, 0.4, 0.5]\npx = [0, 0, 0]\npxx = [1e-5, 5e-6, 1e-20]\namp = [1.0, -2, 0.5]\n\n# Create axis\nt, t2, x, y = pylops.utils.seismicevents.makeaxis(par)\n\n# Create wavelet\nwav = pylops.utils.wavelets.ricker(t[:41], f0=par[\"f0\"])[0]\n\n# Generate model\n_, data = pylops.utils.seismicevents.parabolic2d(x, t, t0, px, pxx, amp, wav)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We want to divide this 2-dimensional data into small overlapping\npatches in the spatial direction and apply the adjoint of the\n:py:class:`pylops.signalprocessing.FFT2D` operator to each patch. This is\ndone by simply using the adjoint of the\n:py:class:`pylops.signalprocessing.Patch2D` operator. Note that for non-\northogonal operators, this must be replaced by an inverse.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nwin = (20, 34)  # window size in data domain\nnop = (\n    128,\n    128 // 2 + 1,\n)  # window size in model domain; we use real FFT, second axis is half\nnover = (10, 4)  # overlap between windows\ndimsd = data.shape\n\n# Sliding window transform without taper\nOp = pylops.signalprocessing.FFT2D(nwin, nffts=(128, 128), real=True)\n\nnwins, dims, mwin_inends, dwin_inends = pylops.signalprocessing.patch2d_design(\n    dimsd, nwin, nover, (128, 65)\n)\nPatch = pylops.signalprocessing.Patch2D(\n    Op.H, dims, dimsd, nwin, nover, nop, tapertype=None\n)\nfftdata = Patch.H * data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We now create a similar operator but we also add a taper to the overlapping\nparts of the patches. We then apply the forward to restore the original\nsignal.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Patch = pylops.signalprocessing.Patch2D(\n    Op.H, dims, dimsd, nwin, nover, nop, tapertype=\"hanning\"\n)\n\nreconstructed_data = Patch * fftdata"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally we re-arrange the transformed patches so that we can also display\nthem\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fftdatareshaped = np.zeros((nop[0] * nwins[0], nop[1] * nwins[1]), dtype=fftdata.dtype)\n\niwin = 1\nfor ix in range(nwins[0]):\n    for it in range(nwins[1]):\n        fftdatareshaped[\n            ix * nop[0] : (ix + 1) * nop[0], it * nop[1] : (it + 1) * nop[1]\n        ] = np.fft.fftshift(fftdata[ix, it])\n        iwin += 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's finally visualize all the intermediate results as well as our final\ndata reconstruction after inverting the\n:py:class:`pylops.signalprocessing.Sliding2D` operator.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(1, 3, figsize=(12, 5))\nim = axs[0].imshow(data.T, cmap=\"gray\")\naxs[0].set_title(\"Original data\")\nplt.colorbar(im, ax=axs[0])\naxs[0].axis(\"tight\")\nim = axs[1].imshow(reconstructed_data.real.T, cmap=\"gray\")\naxs[1].set_title(\"Reconstruction from adjoint\")\nplt.colorbar(im, ax=axs[1])\naxs[1].axis(\"tight\")\naxs[2].imshow(np.abs(fftdatareshaped).T, cmap=\"jet\")\naxs[2].set_title(\"FFT data\")\naxs[2].axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We repeat now the same exercise in 3d\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "par = {\n    \"oy\": -60,\n    \"dy\": 2,\n    \"ny\": 60,\n    \"ox\": -50,\n    \"dx\": 2,\n    \"nx\": 50,\n    \"ot\": 0,\n    \"dt\": 0.004,\n    \"nt\": 100,\n    \"f0\": 20,\n}\n\nv = 1500\nt0 = [0.05, 0.2, 0.3]\nvrms = [500, 700, 1700]\namp = [1.0, -2, 0.5]\n\n# Create axis\nt, t2, x, y = pylops.utils.seismicevents.makeaxis(par)\n\n# Create wavelet\nwav = pylops.utils.wavelets.ricker(t[:41], f0=par[\"f0\"])[0]\n\n# Generate model\n_, data = pylops.utils.seismicevents.hyperbolic3d(x, y, t, t0, vrms, vrms, amp, wav)\n\n\nfig, axs = plt.subplots(1, 3, figsize=(12, 5))\nfig.suptitle(\"Original data\", fontsize=12, fontweight=\"bold\", y=0.95)\naxs[0].imshow(\n    data[par[\"ny\"] // 2].T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-2,\n    vmax=2,\n    cmap=\"gray\",\n    extent=(x.min(), x.max(), t.max(), t.min()),\n)\naxs[0].set_xlabel(r\"$x(m)$\")\naxs[0].set_ylabel(r\"$t(s)$\")\naxs[1].imshow(\n    data[:, par[\"nx\"] // 2].T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-2,\n    vmax=2,\n    cmap=\"gray\",\n    extent=(y.min(), y.max(), t.max(), t.min()),\n)\naxs[1].set_xlabel(r\"$y(m)$\")\naxs[1].set_ylabel(r\"$t(s)$\")\naxs[2].imshow(\n    data[:, :, par[\"nt\"] // 2],\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-2,\n    vmax=2,\n    cmap=\"gray\",\n    extent=(x.min(), x.max(), y.max(), x.min()),\n)\naxs[2].set_xlabel(r\"$x(m)$\")\naxs[2].set_ylabel(r\"$y(m)$\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's create now the :py:class:`pylops.signalprocessing.Patch3D` operator\napplying the adjoint of the :py:class:`pylops.signalprocessing.FFT3D`\noperator to each patch.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nwin = (20, 20, 34)  # window size in data domain\nnop = (\n    128,\n    128,\n    128 // 2 + 1,\n)  # window size in model domain; we use real FFT, third axis is half\nnover = (10, 10, 4)  # overlap between windows\ndimsd = data.shape\n\n# Sliding window transform without taper\nOp = pylops.signalprocessing.FFTND(nwin, nffts=(128, 128, 128), real=True)\n\nnwins, dims, mwin_inends, dwin_inends = pylops.signalprocessing.patch3d_design(\n    dimsd, nwin, nover, (128, 128, 65)\n)\nPatch = pylops.signalprocessing.Patch3D(\n    Op.H, dims, dimsd, nwin, nover, nop, tapertype=None\n)\nfftdata = Patch.H * data\n\nPatch = pylops.signalprocessing.Patch3D(\n    Op.H, dims, dimsd, nwin, nover, nop, tapertype=\"hanning\"\n)\nreconstructed_data = np.real(Patch * fftdata)\n\nfig, axs = plt.subplots(1, 3, figsize=(12, 5))\nfig.suptitle(\"Reconstructed data\", fontsize=12, fontweight=\"bold\", y=0.95)\naxs[0].imshow(\n    reconstructed_data[par[\"ny\"] // 2].T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-2,\n    vmax=2,\n    cmap=\"gray\",\n    extent=(x.min(), x.max(), t.max(), t.min()),\n)\naxs[0].set_xlabel(r\"$x(m)$\")\naxs[0].set_ylabel(r\"$t(s)$\")\naxs[1].imshow(\n    reconstructed_data[:, par[\"nx\"] // 2].T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-2,\n    vmax=2,\n    cmap=\"gray\",\n    extent=(y.min(), y.max(), t.max(), t.min()),\n)\naxs[1].set_xlabel(r\"$y(m)$\")\naxs[1].set_ylabel(r\"$t(s)$\")\naxs[2].imshow(\n    reconstructed_data[:, :, par[\"nt\"] // 2],\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-2,\n    vmax=2,\n    cmap=\"gray\",\n    extent=(x.min(), x.max(), y.max(), x.min()),\n)\naxs[2].set_xlabel(r\"$x(m)$\")\naxs[2].set_ylabel(r\"$y(m)$\")\nplt.tight_layout()"
      ]
    }
  ],
  "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.9.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}