{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Spread How-to\nThis example focuses on the :py:class:`pylops.basicoperators.Spread` operator,\nwhich is a highly versatile operator in PyLops to perform spreading/stacking\noperations in a vectorized manner (or efficiently via Numba-jitted ``for`` loops).\n\nThe :py:class:`pylops.basicoperators.Spread` is powerful in its generality, but\nit may not be obvious for at first how to structure your code to leverage it properly.\nWhile it is highly recommended for advanced users to inspect the\n:py:class:`pylops.signalprocessing.Radon2D` and\n:py:class:`pylops.signalprocessing.Radon3D` operators since\nthey are built using the :py:class:`pylops.basicoperators.Spread` class,\nhere we provide a simple example on how to get started.\n\nIn this example we will recreate a simplified version of the famous linear\n[Radon operator](https://en.wikipedia.org/wiki/Radon_transform), which stacks\ndata along straight lines with a given intercept and slope.\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 define the time and space axes as well as some auxiliary input\nparameters that we will use to create a Ricker wavelet\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "par = {\n    \"ox\": -200,\n    \"dx\": 2,\n    \"nx\": 201,\n    \"ot\": 0,\n    \"dt\": 0.004,\n    \"nt\": 501,\n    \"f0\": 20,\n    \"nfmax\": 210,\n}\n\n# Create axis\nt, _, x, _ = pylops.utils.seismicevents.makeaxis(par)\n\n# Create centered Ricker wavelet\nt_wav = np.arange(41) * par[\"dt\"]\nwav, _, _ = pylops.utils.wavelets.ricker(t_wav, f0=par[\"f0\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We will create a 2d data with a number of crossing linear events, to which we will\nlater apply our Radon transforms. We use the convenience function\n:py:func:`pylops.utils.seismicevents.linear2d`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "v = 1500  # m/s\nt0 = [0.2, 0.7, 1.6]  # seconds\ntheta = [40, 0, -60]  # degrees\namp = [1.0, 0.6, -2.0]\n\nmlin, mlinwav = pylops.utils.seismicevents.linear2d(x, t, v, t0, theta, amp, wav)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's now define the slowness axis and use :py:class:`pylops.signalprocessing.Radon2D`\nto implement our benchmark linear Radon. Refer to the documentation of the\noperator for a more detailed mathematical description of linear Radon.\nNote that ``pxmax`` is in s/m, which explains the small value. Its highest value\ncorresponds to the lowest value of velocity in the transform. In this case we choose that\nto be 1000 m/s.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "npx, pxmax = 41, 1e-3\npx = np.linspace(-pxmax, pxmax, npx)\n\nRLop = pylops.signalprocessing.Radon2D(\n    t, x, px, centeredh=False, kind=\"linear\", interp=False, engine=\"numpy\"\n)\n\n# Compute adjoint = Radon transform\nmlinwavR = RLop.H * mlinwav"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now, let's try to reimplement this operator from scratch using :py:class:`pylops.basicoperators.Spread`.\nUsing the on-the-fly approach, and we need to create a function which takes\nindices of the model domain, here $(p_x, t_0)$\nwhere $p_x$ is the slope and $t_0$ is the intercept of the\nparametric curve $t(x) = t_0 + p_x x$ we wish to spread the model over\nin the data domain. The function must return an array of size ``nx``, containing\nthe indices corresponding to $t(x)$.\n\nThe on-the-fly approach is useful when storing the indices in RAM may exhaust\nresources, especially when computing the indices is fast. When there is\nenough memory to store the full table of indices\n(an array of size $n_x \\times n_t \\times n_{p_x}$) the\n:py:class:`pylops.basicoperators.Spread` operator can be used with tables instead.\nWe will see an example of this later.\n\nReturning to our on-the-fly example, we need to create a function which only depends on\n``ipx`` and ``it0``, so we create a closure around it with all our other auxiliary\nvariables.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def create_radon_fh(xaxis, taxis, pxaxis):\n    ot = taxis[0]\n    dt = taxis[1] - taxis[0]\n    nt = len(taxis)\n\n    def fh(ipx, it0):\n        tx = t[it0] + xaxis * pxaxis[ipx]\n        it0_frac = (tx - ot) / dt\n        itx = np.rint(it0_frac)\n        # Indices outside time axis set to nan\n        itx[np.isin(itx, range(nt), invert=True)] = np.nan\n        return itx\n\n    return fh\n\n\nfRad = create_radon_fh(x, t, px)\nROTFOp = pylops.Spread((npx, par[\"nt\"]), (par[\"nx\"], par[\"nt\"]), fh=fRad)\n\nmlinwavROTF = ROTFOp.H * mlinwav"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compare the results between the native Radon transform and the one using our\non-the-fly :py:class:`pylops.basicoperators.Spread`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(1, 3, figsize=(9, 5), sharey=True)\naxs[0].imshow(\n    mlinwav.T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-1,\n    vmax=1,\n    cmap=\"gray\",\n    extent=(x.min(), x.max(), t.max(), t.min()),\n)\naxs[0].set_title(\"Linear events\", fontsize=12, fontweight=\"bold\")\naxs[0].set_xlabel(r\"$x$ [m]\")\naxs[0].set_ylabel(r\"$t$ [s]\")\naxs[1].imshow(\n    mlinwavR.T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-10,\n    vmax=10,\n    cmap=\"gray\",\n    extent=(px.min(), px.max(), t.max(), t.min()),\n)\naxs[1].set_title(\"Native Linear Radon\", fontsize=12, fontweight=\"bold\")\naxs[1].set_xlabel(r\"$p_x$ [s/m]\")\naxs[1].ticklabel_format(style=\"sci\", axis=\"x\", scilimits=(0, 0))\n\naxs[2].imshow(\n    mlinwavROTF.T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-10,\n    vmax=10,\n    cmap=\"gray\",\n    extent=(px.min(), px.max(), t.max(), t.min()),\n)\naxs[2].set_title(\"On-the-fly Linear Radon\", fontsize=12, fontweight=\"bold\")\naxs[2].set_xlabel(r\"$p_x$ [s/m]\")\naxs[2].ticklabel_format(style=\"sci\", axis=\"x\", scilimits=(0, 0))\nfig.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, we will re-implement the example above using pre-computed tables.\nThis is useful when ``fh`` is expensive to compute, or requires manual edition\nprior to usage.\n\nUsing a table instead of a function is simple, we just need to apply ``fh`` to\nall our points and store the results.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def create_table(npx, nt, nx):\n    table = np.full((npx, nt, nx), fill_value=np.nan)\n    for ipx in range(npx):\n        for it0 in range(nt):\n            table[ipx, it0, :] = fRad(ipx, it0)\n    return table\n\n\ntable = create_table(npx, par[\"nt\"], par[\"nx\"])\nRPCOp = pylops.Spread((npx, par[\"nt\"]), (par[\"nx\"], par[\"nt\"]), table=table)\n\nmlinwavRPC = RPCOp.H * mlinwav"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compare the results between the pre-computed or on-the-fly Radon transforms\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, axs = plt.subplots(1, 3, figsize=(9, 5), sharey=True)\naxs[0].imshow(\n    mlinwav.T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-1,\n    vmax=1,\n    cmap=\"gray\",\n    extent=(x.min(), x.max(), t.max(), t.min()),\n)\naxs[0].set_title(\"Linear events\", fontsize=12, fontweight=\"bold\")\naxs[0].set_xlabel(r\"$x$ [m]\")\naxs[0].set_ylabel(r\"$t$ [s]\")\naxs[1].imshow(\n    mlinwavRPC.T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-10,\n    vmax=10,\n    cmap=\"gray\",\n    extent=(px.min(), px.max(), t.max(), t.min()),\n)\naxs[1].set_title(\"Pre-computed Linear Radon\", fontsize=12, fontweight=\"bold\")\naxs[1].set_xlabel(r\"$p_x$ [s/m]\")\naxs[1].ticklabel_format(style=\"sci\", axis=\"x\", scilimits=(0, 0))\n\naxs[2].imshow(\n    mlinwavROTF.T,\n    aspect=\"auto\",\n    interpolation=\"nearest\",\n    vmin=-10,\n    vmax=10,\n    cmap=\"gray\",\n    extent=(px.min(), px.max(), t.max(), t.min()),\n)\naxs[2].set_title(\"On-the-fly Linear Radon\", fontsize=12, fontweight=\"bold\")\naxs[2].set_xlabel(r\"$p_x$ [s/m]\")\naxs[2].ticklabel_format(style=\"sci\", axis=\"x\", scilimits=(0, 0))\nfig.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
}