{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# PhaseShift operator\nThis example shows how to use the :class:`pylops.waveeqprocessing.PhaseShift`\noperator to perform frequency-wavenumber shift of an input multi-dimensional\nsignal. Such a procedure is applied in a variety of disciplines including\ngeophysics, medical imaging and non-destructive testing.\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 first create a synthetic dataset composed of a number of hyperbolas\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "par = {\n    \"ox\": -300,\n    \"dx\": 20,\n    \"nx\": 31,\n    \"oy\": -200,\n    \"dy\": 20,\n    \"ny\": 21,\n    \"ot\": 0,\n    \"dt\": 0.004,\n    \"nt\": 201,\n    \"f0\": 20,\n    \"nfmax\": 210,\n}\n\n# Create axis\nt, t2, x, y = pylops.utils.seismicevents.makeaxis(par)\n\n# Create wavelet\nwav = pylops.utils.wavelets.ricker(np.arange(41) * par[\"dt\"], f0=par[\"f0\"])[0]\n\nvrms = [900, 1300, 1800]\nt0 = [0.2, 0.3, 0.6]\namp = [1.0, 0.6, -2.0]\n\n_, m = pylops.utils.seismicevents.hyperbolic2d(x, t, t0, vrms, amp, wav)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can now apply a taper at the edges and also pad the input to avoid\nartifacts during the phase shift\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pad = 11\ntaper = pylops.utils.tapers.taper2d(par[\"nt\"], par[\"nx\"], 5)\nmpad = np.pad(m * taper, ((pad, pad), (0, 0)), mode=\"constant\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We perform now forward propagation in a constant velocity $v=2000$ for\na depth of $z_{prop} = 100 m$. We should expect the hyperbolas to move\nforward in time and become flatter.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "vel = 1500.0\nzprop = 100\nfreq = np.fft.rfftfreq(par[\"nt\"], par[\"dt\"])\nkx = np.fft.fftshift(np.fft.fftfreq(par[\"nx\"] + 2 * pad, par[\"dx\"]))\nPop = pylops.waveeqprocessing.PhaseShift(vel, zprop, par[\"nt\"], freq, kx)\n\nmdown = Pop * mpad.T.ravel()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We now take the forward propagated wavefield and apply backward propagation,\nwhich is in this case simply the adjoint of our operator.\nWe should expect the hyperbolas to move backward in time and show the same\ntraveltime as the original dataset. Of course, as we are only performing the\nadjoint operation we should expect some small differences between this\nwavefield and the input dataset.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mup = Pop.H * mdown.ravel()\n\nmdown = np.real(mdown.reshape(par[\"nt\"], par[\"nx\"] + 2 * pad)[:, pad:-pad])\nmup = np.real(mup.reshape(par[\"nt\"], par[\"nx\"] + 2 * pad)[:, pad:-pad])\n\nfig, axs = plt.subplots(1, 3, figsize=(10, 6), sharey=True)\nfig.suptitle(\"2D Phase shift\", fontsize=12, fontweight=\"bold\")\naxs[0].imshow(\n    m.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[0].set_title(\"Original data\")\naxs[1].imshow(\n    mdown,\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[1].set_xlabel(r\"$x(m)$\")\naxs[1].set_title(\"Forward propagation\")\naxs[2].imshow(\n    mup,\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[2].set_xlabel(r\"$x(m)$\")\naxs[2].set_title(\"Backward propagation\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally we perform the same for a 3-dimensional signal\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "_, m = pylops.utils.seismicevents.hyperbolic3d(x, y, t, t0, vrms, vrms, amp, wav)\n\npad = 11\ntaper = pylops.utils.tapers.taper3d(par[\"nt\"], (par[\"ny\"], par[\"nx\"]), (3, 3))\nmpad = np.pad(m * taper, ((pad, pad), (pad, pad), (0, 0)), mode=\"constant\")\n\nkx = np.fft.fftshift(np.fft.fftfreq(par[\"nx\"] + 2 * pad, par[\"dx\"]))\nky = np.fft.fftshift(np.fft.fftfreq(par[\"ny\"] + 2 * pad, par[\"dy\"]))\nPop = pylops.waveeqprocessing.PhaseShift(vel, zprop, par[\"nt\"], freq, kx, ky)\n\nmdown = Pop * mpad.transpose(2, 1, 0).ravel()\n\nmup = Pop.H * mdown.ravel()\n\nmdown = np.real(\n    mdown.reshape(par[\"nt\"], par[\"nx\"] + 2 * pad, par[\"ny\"] + 2 * pad)[\n        :, pad:-pad, pad:-pad\n    ]\n)\nmup = np.real(\n    mup.reshape(par[\"nt\"], par[\"nx\"] + 2 * pad, par[\"ny\"] + 2 * pad)[\n        :, pad:-pad, pad:-pad\n    ]\n)\n\nfig, axs = plt.subplots(2, 3, figsize=(10, 12), sharey=True)\nfig.suptitle(\"3D Phase shift\", fontsize=12, fontweight=\"bold\")\naxs[0][0].imshow(\n    m[:, par[\"nx\"] // 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][0].set_xlabel(r\"$y(m)$\")\naxs[0][0].set_ylabel(r\"$t(s)$\")\naxs[0][0].set_title(\"Original data\")\naxs[0][1].imshow(\n    mdown[:, par[\"nx\"] // 2],\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][1].set_xlabel(r\"$y(m)$\")\naxs[0][1].set_title(\"Forward propagation\")\naxs[0][2].imshow(\n    mup[:, par[\"nx\"] // 2],\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][2].set_xlabel(r\"$y(m)$\")\naxs[0][2].set_title(\"Backward propagation\")\naxs[1][0].imshow(\n    m[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[1][0].set_xlabel(r\"$x(m)$\")\naxs[1][0].set_ylabel(r\"$t(s)$\")\naxs[1][0].set_title(\"Original data\")\naxs[1][1].imshow(\n    mdown[:, :, par[\"ny\"] // 2],\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[1][1].set_xlabel(r\"$x(m)$\")\naxs[1][1].set_title(\"Forward propagation\")\naxs[1][2].imshow(\n    mup[:, :, par[\"ny\"] // 2],\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[1][2].set_xlabel(r\"$x(m)$\")\naxs[1][2].set_title(\"Backward propagation\")\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
}