{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# 12. Seismic regularization\nThe problem of *seismic data regularization* (or interpolation) is a very\nsimple one to write, yet ill-posed and very hard to solve.\n\nThe forward modelling operator is a simple :py:class:`pylops.Restriction`\noperator which is applied along the spatial direction(s).\n\n\\begin{align}\\mathbf{y} = \\mathbf{R} \\mathbf{x}\\end{align}\n\nHere $\\mathbf{y} = [\\mathbf{y}_{R_1}^T, \\mathbf{y}_{R_2}^T,\\ldots,\n\\mathbf{y}_{R_N^T}]^T$ where each vector $\\mathbf{y}_{R_i}$\ncontains all time samples recorded in the seismic data at the specific\nreceiver $R_i$. Similarly, $\\mathbf{x} = [\\mathbf{x}_{r_1}^T,\n\\mathbf{x}_{r_2}^T,\\ldots, \\mathbf{x}_{r_M}^T]$, contains all traces at the\nregularly and finely sampled receiver locations $r_i$.\n\nBy inverting such an equation we can create a regularized data with\ndensely and regularly spatial direction(s).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\nfrom scipy.signal import convolve\n\nimport pylops\nfrom pylops.utils.seismicevents import linear2d, makeaxis\nfrom pylops.utils.wavelets import ricker\n\nnp.random.seed(0)\nplt.close(\"all\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's start by creating a very simple 2d data composed of 3 linear events\ninput parameters\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "par = {\"ox\": 0, \"dx\": 2, \"nx\": 70, \"ot\": 0, \"dt\": 0.004, \"nt\": 80, \"f0\": 20}\n\nv = 1500\nt0_m = [0.1, 0.2, 0.28]\ntheta_m = [0, 30, -80]\nphi_m = [0]\namp_m = [1.0, -2, 0.5]\n\n# axis\ntaxis, t2, xaxis, y = makeaxis(par)\n\n# wavelet\nwav = ricker(taxis[:41], f0=par[\"f0\"])[0]\n\n# model\n_, x = linear2d(xaxis, taxis, v, t0_m, theta_m, amp_m, wav)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can now define the spatial locations along which the data has been\nsampled. In this specific example we will assume that we have access only to\n40% of the 'original' locations.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "perc_subsampling = 0.6\nnxsub = int(np.round(par[\"nx\"] * perc_subsampling))\n\niava = np.sort(np.random.permutation(np.arange(par[\"nx\"]))[:nxsub])\n\n# restriction operator\nRop = pylops.Restriction((par[\"nx\"], par[\"nt\"]), iava, axis=0, dtype=\"float64\")\n\n# data\ny = Rop * x.ravel()\ny = y.reshape(nxsub, par[\"nt\"])\n\n# mask\nymask = Rop.mask(x.ravel())\n\n# inverse\nxinv = Rop / y.ravel()\nxinv = xinv.reshape(par[\"nx\"], par[\"nt\"])\n\nfig, axs = plt.subplots(1, 2, sharey=True, figsize=(5, 4))\naxs[0].imshow(\n    x.T, cmap=\"gray\", vmin=-2, vmax=2, extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0])\n)\naxs[0].set_title(\"Model\")\naxs[0].axis(\"tight\")\naxs[1].imshow(\n    ymask.T,\n    cmap=\"gray\",\n    vmin=-2,\n    vmax=2,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[1].set_title(\"Masked model\")\naxs[1].axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As we can see, inverting the restriction operator is not possible without\nadding any prior information into the inverse problem. In the following we\nwill consider two possible routes:\n\n* regularized inversion with second derivative along the spatial axis\n\n  .. math::\n       J = \\|\\mathbf{y} - \\mathbf{R} \\mathbf{x}\\|_2 +\n       \\epsilon_\\nabla ^2 \\|\\nabla \\mathbf{x}\\|_2\n\n* sparsity-promoting inversion with :py:class:`pylops.FFT2` operator used\n  as sparsyfing transform\n\n  .. math::\n       J = \\|\\mathbf{y} - \\mathbf{R} \\mathbf{F}^H \\mathbf{x}\\|_2 +\n       \\epsilon \\|\\mathbf{F}^H \\mathbf{x}\\|_1\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# smooth inversion\nD2op = pylops.SecondDerivative((par[\"nx\"], par[\"nt\"]), axis=0, dtype=\"float64\")\n\nxsmooth, _, _ = pylops.waveeqprocessing.SeismicInterpolation(\n    y,\n    par[\"nx\"],\n    iava,\n    kind=\"spatial\",\n    **dict(epsRs=[np.sqrt(0.1)], damp=np.sqrt(1e-4), iter_lim=50, show=0)\n)\n\n# sparse inversion with FFT2\nnfft = 2**8\nFFTop = pylops.signalprocessing.FFT2D(\n    dims=[par[\"nx\"], par[\"nt\"]], nffts=[nfft, nfft], sampling=[par[\"dx\"], par[\"dt\"]]\n)\nX = FFTop * x.ravel()\nX = np.reshape(X, (nfft, nfft))\n\nxl1, Xl1, cost = pylops.waveeqprocessing.SeismicInterpolation(\n    y,\n    par[\"nx\"],\n    iava,\n    kind=\"fk\",\n    nffts=(nfft, nfft),\n    sampling=(par[\"dx\"], par[\"dt\"]),\n    **dict(niter=50, eps=1e-1)\n)\n\nfig, axs = plt.subplots(1, 4, sharey=True, figsize=(13, 4))\naxs[0].imshow(\n    x.T, cmap=\"gray\", vmin=-2, vmax=2, extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0])\n)\naxs[0].set_title(\"Model\")\naxs[0].axis(\"tight\")\naxs[1].imshow(\n    ymask.T,\n    cmap=\"gray\",\n    vmin=-2,\n    vmax=2,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[1].set_title(\"Masked model\")\naxs[1].axis(\"tight\")\naxs[2].imshow(\n    xsmooth.T,\n    cmap=\"gray\",\n    vmin=-2,\n    vmax=2,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[2].set_title(\"Smoothed model\")\naxs[2].axis(\"tight\")\naxs[3].imshow(\n    xl1.T,\n    cmap=\"gray\",\n    vmin=-2,\n    vmax=2,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[3].set_title(\"L1 model\")\naxs[3].axis(\"tight\")\n\nfig, axs = plt.subplots(1, 3, figsize=(10, 2))\naxs[0].imshow(\n    np.fft.fftshift(np.abs(X[:, : nfft // 2 - 1]), axes=0).T,\n    extent=(\n        np.fft.fftshift(FFTop.f1)[0],\n        np.fft.fftshift(FFTop.f1)[-1],\n        FFTop.f2[nfft // 2 - 1],\n        FFTop.f2[0],\n    ),\n)\naxs[0].set_title(\"Model in f-k domain\")\naxs[0].axis(\"tight\")\naxs[0].set_xlim(-0.1, 0.1)\naxs[0].set_ylim(50, 0)\naxs[1].imshow(\n    np.fft.fftshift(np.abs(Xl1[:, : nfft // 2 - 1]), axes=0).T,\n    extent=(\n        np.fft.fftshift(FFTop.f1)[0],\n        np.fft.fftshift(FFTop.f1)[-1],\n        FFTop.f2[nfft // 2 - 1],\n        FFTop.f2[0],\n    ),\n)\naxs[1].set_title(\"Reconstructed model in f-k domain\")\naxs[1].axis(\"tight\")\naxs[1].set_xlim(-0.1, 0.1)\naxs[1].set_ylim(50, 0)\naxs[2].plot(cost, \"k\", lw=3)\naxs[2].set_title(\"FISTA convergence\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We see how adding prior information to the inversion can help improving the\nestimate of the regularized seismic data. Nevertheless, in both cases the\nreconstructed data is not perfect. A better sparsyfing transform could in\nfact be chosen here to be the linear\n:py:class:`pylops.signalprocessing.Radon2D` transform in spite of the\n:py:class:`pylops.FFT2` transform.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "npx = 40\npxmax = 1e-3\npx = np.linspace(-pxmax, pxmax, npx)\nRadop = pylops.signalprocessing.Radon2D(taxis, xaxis, px, engine=\"numba\")\n\nRRop = Rop * Radop\n\n# adjoint\nXadj_fromx = Radop.H * x.ravel()\nXadj_fromx = Xadj_fromx.reshape(npx, par[\"nt\"])\n\nXadj = RRop.H * y.ravel()\nXadj = Xadj.reshape(npx, par[\"nt\"])\n\n# L1 inverse\nxl1, Xl1, cost = pylops.waveeqprocessing.SeismicInterpolation(\n    y,\n    par[\"nx\"],\n    iava,\n    kind=\"radon-linear\",\n    spataxis=xaxis,\n    taxis=taxis,\n    paxis=px,\n    centeredh=True,\n    **dict(niter=50, eps=1e-1)\n)\n\nfig, axs = plt.subplots(2, 3, sharey=True, figsize=(12, 7))\naxs[0][0].imshow(\n    x.T, cmap=\"gray\", vmin=-2, vmax=2, extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0])\n)\naxs[0][0].set_title(\"Data\", fontsize=12)\naxs[0][0].axis(\"tight\")\naxs[0][1].imshow(\n    ymask.T,\n    cmap=\"gray\",\n    vmin=-2,\n    vmax=2,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[0][1].set_title(\"Masked data\", fontsize=12)\naxs[0][1].axis(\"tight\")\naxs[0][2].imshow(\n    xl1.T,\n    cmap=\"gray\",\n    vmin=-2,\n    vmax=2,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[0][2].set_title(\"Reconstructed data\", fontsize=12)\naxs[0][2].axis(\"tight\")\naxs[1][0].imshow(\n    Xadj_fromx.T,\n    cmap=\"gray\",\n    vmin=-70,\n    vmax=70,\n    extent=(px[0], px[-1], taxis[-1], taxis[0]),\n)\naxs[1][0].set_title(\"Adj. Radon on data\", fontsize=12)\naxs[1][0].axis(\"tight\")\naxs[1][1].imshow(\n    Xadj.T, cmap=\"gray\", vmin=-50, vmax=50, extent=(px[0], px[-1], taxis[-1], taxis[0])\n)\naxs[1][1].set_title(\"Adj. Radon on subsampled data\", fontsize=12)\naxs[1][1].axis(\"tight\")\naxs[1][2].imshow(\n    Xl1.T, cmap=\"gray\", vmin=-0.2, vmax=0.2, extent=(px[0], px[-1], taxis[-1], taxis[0])\n)\naxs[1][2].set_title(\"Inverse Radon on subsampled data\", fontsize=12)\naxs[1][2].axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, let's take now a more realistic dataset. We will use once again the\nlinear :py:class:`pylops.signalprocessing.Radon2D` transform but we will\ntake advantnge of the :py:class:`pylops.signalprocessing.Sliding2D` operator\nto perform such a transform locally instead of globally to the entire\ndataset.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "inputfile = \"../testdata/marchenko/input.npz\"\ninputdata = np.load(inputfile)\n\nx = inputdata[\"R\"][50, :, ::2]\nx = x / np.abs(x).max()\ntaxis, xaxis = inputdata[\"t\"][::2], inputdata[\"r\"][0]\n\npar = {}\npar[\"nx\"], par[\"nt\"] = x.shape\npar[\"dx\"] = inputdata[\"r\"][0, 1] - inputdata[\"r\"][0, 0]\npar[\"dt\"] = inputdata[\"t\"][1] - inputdata[\"t\"][0]\n\n# add wavelet\nwav = inputdata[\"wav\"][::2]\nwav_c = np.argmax(wav)\nx = np.apply_along_axis(convolve, 1, x, wav, mode=\"full\")\nx = x[:, wav_c:][:, : par[\"nt\"]]\n\n# gain\ngain = np.tile((taxis**2)[:, np.newaxis], (1, par[\"nx\"])).T\nx = x * gain\n\n# subsampling locations\nperc_subsampling = 0.5\nNsub = int(np.round(par[\"nx\"] * perc_subsampling))\niava = np.sort(np.random.permutation(np.arange(par[\"nx\"]))[:Nsub])\n\n# restriction operator\nRop = pylops.Restriction((par[\"nx\"], par[\"nt\"]), iava, axis=0, dtype=\"float64\")\n\ny = Rop * x.ravel()\nxadj = Rop.H * y.ravel()\n\ny = y.reshape(Nsub, par[\"nt\"])\nxadj = xadj.reshape(par[\"nx\"], par[\"nt\"])\n\n# apply mask\nymask = Rop.mask(x.ravel())\n\n# sliding windows with radon transform\ndx = par[\"dx\"]\nnwins = 4\nnwin = 27\nnover = 3\nnpx = 31\npxmax = 5e-4\npx = np.linspace(-pxmax, pxmax, npx)\ndimsd = x.shape\ndims = (nwins * npx, dimsd[1])\n\nOp = pylops.signalprocessing.Radon2D(\n    taxis,\n    np.linspace(-par[\"dx\"] * nwin // 2, par[\"dx\"] * nwin // 2, nwin),\n    px,\n    centeredh=True,\n    kind=\"linear\",\n    engine=\"numba\",\n)\nSlidop = pylops.signalprocessing.Sliding2D(\n    Op, dims, dimsd, nwin, nover, tapertype=\"cosine\"\n)\n\n# adjoint\nRSop = Rop * Slidop\n\nXadj_fromx = Slidop.H * x.ravel()\nXadj_fromx = Xadj_fromx.reshape(npx * nwins, par[\"nt\"])\n\nXadj = RSop.H * y.ravel()\nXadj = Xadj.reshape(npx * nwins, par[\"nt\"])\n\n# inverse\nxl1, Xl1, _ = pylops.waveeqprocessing.SeismicInterpolation(\n    y,\n    par[\"nx\"],\n    iava,\n    kind=\"sliding\",\n    spataxis=xaxis,\n    taxis=taxis,\n    paxis=px,\n    nwins=nwins,\n    nwin=nwin,\n    nover=nover,\n    **dict(niter=50, eps=1e-2)\n)\n\nfig, axs = plt.subplots(2, 3, sharey=True, figsize=(12, 14))\naxs[0][0].imshow(\n    x.T,\n    cmap=\"gray\",\n    vmin=-0.1,\n    vmax=0.1,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[0][0].set_title(\"Data\")\naxs[0][0].axis(\"tight\")\naxs[0][1].imshow(\n    ymask.T,\n    cmap=\"gray\",\n    vmin=-0.1,\n    vmax=0.1,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[0][1].set_title(\"Masked data\")\naxs[0][1].axis(\"tight\")\naxs[0][2].imshow(\n    xl1.T,\n    cmap=\"gray\",\n    vmin=-0.1,\n    vmax=0.1,\n    extent=(xaxis[0], xaxis[-1], taxis[-1], taxis[0]),\n)\naxs[0][2].set_title(\"Reconstructed data\")\naxs[0][2].axis(\"tight\")\naxs[1][0].imshow(\n    Xadj_fromx.T,\n    cmap=\"gray\",\n    vmin=-1,\n    vmax=1,\n    extent=(px[0], px[-1], taxis[-1], taxis[0]),\n)\naxs[1][0].set_title(\"Adjoint Radon on data\")\naxs[1][0].axis(\"tight\")\naxs[1][1].imshow(\n    Xadj.T,\n    cmap=\"gray\",\n    vmin=-0.6,\n    vmax=0.6,\n    extent=(px[0], px[-1], taxis[-1], taxis[0]),\n)\naxs[1][1].set_title(\"Adjoint Radon on subsampled data\")\naxs[1][1].axis(\"tight\")\naxs[1][2].imshow(\n    Xl1.T,\n    cmap=\"gray\",\n    vmin=-0.03,\n    vmax=0.03,\n    extent=(px[0], px[-1], taxis[-1], taxis[0]),\n)\naxs[1][2].set_title(\"Inverse Radon on subsampled data\")\naxs[1][2].axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As expected the linear :py:class:`pylops.signalprocessing.Radon2D` is\nable to locally explain events in the input data and leads to a satisfactory\nrecovery. Note that increasing the number of iterations and sliding windows\ncan further refine the result, especially the accuracy of weak events, as\nshown in this companion\n[notebook](https://github.com/mrava87/pylops_notebooks/blob/master/developement/SeismicInterpolation.ipynb).\n\n"
      ]
    }
  ],
  "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
}