{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# 18. Deblending\nThe cocktail party problem arises when sounds from different sources mix before reaching our ears\n(or any recording device), requiring the brain (or any hardware in the recording device) to estimate\nindividual sources from the received mixture. In seismic acquisition, an analog problem is present\nwhen multiple sources are fired simultaneously. This family of acquisition methods is usually referred to as\nsimultaneous shooting and the problem of separating the blended shot gathers into their individual\ncomponents is called deblending. Whilst various firing strategies can be adopted, in this example\nwe consider the continuous blending problem where a single source is fired sequentially at an interval\nshorter than the amount of time required for waves to travel into the Earth and come back.\n\nSimply stated the forward problem can be written as:\n\n\\begin{align}\\mathbf{d}^b = \\boldsymbol\\Phi \\mathbf{d}\\end{align}\n\nHere $\\mathbf{d} = [\\mathbf{d}_1^T, \\mathbf{d}_2^T,\\ldots,\n\\mathbf{d}_N^T]^T$ is a stack of $N$ individual shot gathers,\n$\\boldsymbol\\Phi=[\\boldsymbol\\Phi_1, \\boldsymbol\\Phi_2,\\ldots,\n\\boldsymbol\\Phi_N]$ is the blending operator, $\\mathbf{d}^b$ is the\nso-called supergather than contains all shots superimposed to each other.\n\nIn order to successfully invert this severely under-determined problem, two key\ningredients must be introduced:\n\n- the firing time of each source (i.e., shifts of the blending operator) must be\n  chosen to be dithered around a nominal regular, periodic firing interval.\n  In our case, we consider shots of duration $T=4\\,\\text{s}$, regular firing time of $T_s=2\\,\\text{s}$\n  and a dithering code as follows $\\Delta t = U(-1,1)$;\n- prior information about the data to reconstruct, either in the form of regularization\n  or preconditioning must be introduced. In our case we will use a patch-FK transform\n  as preconditioner and solve the problem imposing sparsity in the transformed domain.\n\nIn other words, we aim to solve the following problem:\n\n\\begin{align}J = \\|\\mathbf{d}^b - \\boldsymbol\\Phi \\mathbf{S}^H \\mathbf{x}\\|_2 + \\epsilon \\|\\mathbf{x}\\|_1\\end{align}\n\nfor which we will use the :py:class:`pylops.optimization.sparsity.fista` solver.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\nfrom scipy.sparse.linalg import lobpcg as sp_lobpcg\n\nimport pylops\n\nnp.random.seed(10)\nplt.close(\"all\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can now load and display a small portion of the MobilAVO dataset composed\nof 60 shots and a single receiver. This data is unblended.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data = np.load(\"../testdata/deblending/mobil.npy\")\nns, nt = data.shape\n\ndt = 0.004\nt = np.arange(nt) * dt\n\nfig, ax = plt.subplots(1, 1, figsize=(12, 8))\nax.imshow(\n    data.T,\n    cmap=\"gray\",\n    vmin=-50,\n    vmax=50,\n    extent=(0, ns, t[-1], 0),\n    interpolation=\"none\",\n)\nax.set_title(\"CRG\")\nax.set_xlabel(\"#Src\")\nax.set_ylabel(\"t [s]\")\nax.axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We are now ready to define the blending operator, blend our data, and apply\nthe adjoint of the blending operator to it. This is usually referred as\npseudo-deblending: as we will see brings back each source to its own nominal\nfiring time, but since sources partially overlap in time, it will also generate\nsome burst like noise in the data. Deblending can hopefully fix this.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "overlap = 0.5\nignition_times = 2.0 * np.random.rand(ns) - 1.0\nignition_times = np.arange(0, overlap * nt * ns, overlap * nt) * dt + ignition_times\nignition_times[0] = 0.0\nBop = pylops.waveeqprocessing.BlendingContinuous(\n    nt, 1, ns, dt, ignition_times, dtype=\"complex128\"\n)\n\ndata_blended = Bop * data[:, np.newaxis]\ndata_pseudo = Bop.H * data_blended\ndata_pseudo = data_pseudo.reshape(ns, nt)\n\nfig, ax = plt.subplots(1, 1, figsize=(12, 8))\nax.imshow(\n    data_pseudo.T.real,\n    cmap=\"gray\",\n    vmin=-50,\n    vmax=50,\n    extent=(0, ns, t[-1], 0),\n    interpolation=\"none\",\n)\nax.set_title(\"Pseudo-deblended CRG\")\nax.set_xlabel(\"#Src\")\nax.set_ylabel(\"t [s]\")\nax.axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We are finally ready to solve our deblending inverse problem\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Patched FK\ndimsd = data.shape\nnwin = (20, 80)\nnover = (10, 40)\nnop = (128, 128)\nnop1 = (128, 65)\nnwins = (5, 24)\ndims = (nwins[0] * nop1[0], nwins[1] * nop1[1])\n\nFop = pylops.signalprocessing.FFT2D(nwin, nffts=nop, real=True)\nSop = pylops.signalprocessing.Patch2D(\n    Fop.H, dims, dimsd, nwin, nover, nop1, tapertype=\"hanning\"\n)\n# Overall operator\nOp = Bop * Sop\n\n# Compute max eigenvalue (we do this explicitly to be able to run this fast)\nOp1 = pylops.LinearOperator(Op.H * Op, explicit=False)\nX = np.random.rand(Op1.shape[0], 1).astype(Op1.dtype)\nmaxeig = sp_lobpcg(Op1, X=X, maxiter=5, tol=1e-10)[0][0]\nalpha = 1.0 / maxeig\n\n# Deblend\nniter = 60\ndecay = (np.exp(-0.05 * np.arange(niter)) + 0.2) / 1.2\n\np_inv = pylops.optimization.sparsity.fista(\n    Op,\n    data_blended.ravel(),\n    niter=niter,\n    eps=5e0,\n    alpha=alpha,\n    decay=decay,\n    show=True,\n)[0]\ndata_inv = Sop * p_inv\ndata_inv = data_inv.reshape(ns, nt)\n\nfig, axs = plt.subplots(1, 4, sharey=False, figsize=(12, 8))\naxs[0].imshow(\n    data.T.real,\n    cmap=\"gray\",\n    extent=(0, ns, t[-1], 0),\n    vmin=-50,\n    vmax=50,\n    interpolation=\"none\",\n)\naxs[0].set_title(\"CRG\")\naxs[0].set_xlabel(\"#Src\")\naxs[0].set_ylabel(\"t [s]\")\naxs[0].axis(\"tight\")\naxs[1].imshow(\n    data_pseudo.T.real,\n    cmap=\"gray\",\n    extent=(0, ns, t[-1], 0),\n    vmin=-50,\n    vmax=50,\n    interpolation=\"none\",\n)\naxs[1].set_title(\"Pseudo-deblended CRG\")\naxs[1].set_xlabel(\"#Src\")\naxs[1].axis(\"tight\")\naxs[2].imshow(\n    data_inv.T.real,\n    cmap=\"gray\",\n    extent=(0, ns, t[-1], 0),\n    vmin=-50,\n    vmax=50,\n    interpolation=\"none\",\n)\naxs[2].set_xlabel(\"#Src\")\naxs[2].set_title(\"Deblended CRG\")\naxs[2].axis(\"tight\")\naxs[3].imshow(\n    data.T.real - data_inv.T.real,\n    cmap=\"gray\",\n    extent=(0, ns, t[-1], 0),\n    vmin=-50,\n    vmax=50,\n    interpolation=\"none\",\n)\naxs[3].set_xlabel(\"#Src\")\naxs[3].set_title(\"Blending error\")\naxs[3].axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, let's look a bit more at what really happened under the hood. We\ndisplay a number of patches and their associated FK spectrum\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Sop1 = pylops.signalprocessing.Patch2D(\n    Fop.H, dims, dimsd, nwin, nover, nop1, tapertype=None\n)\n\n# Original\np = Sop1.H * data.ravel()\npreshape = p.reshape(nwins[0], nwins[1], nop1[0], nop1[1])\n\nix = 16\nfig, axs = plt.subplots(2, 4, figsize=(12, 5))\nfig.suptitle(\"Data patches\")\nfor i in range(4):\n    axs[0][i].imshow(np.fft.fftshift(np.abs(preshape[i, ix]).T, axes=1))\n    axs[0][i].axis(\"tight\")\n    axs[1][i].imshow(\n        np.real((Fop.H * preshape[i, ix].ravel()).reshape(nwin)).T,\n        cmap=\"gray\",\n        vmin=-30,\n        vmax=30,\n        interpolation=\"none\",\n    )\n    axs[1][i].axis(\"tight\")\nplt.tight_layout()\n\n# Pseudo-deblended\np_pseudo = Sop1.H * data_pseudo.ravel()\np_pseudoreshape = p_pseudo.reshape(nwins[0], nwins[1], nop1[0], nop1[1])\n\nix = 16\nfig, axs = plt.subplots(2, 4, figsize=(12, 5))\nfig.suptitle(\"Pseudo-deblended patches\")\nfor i in range(4):\n    axs[0][i].imshow(np.fft.fftshift(np.abs(p_pseudoreshape[i, ix]).T, axes=1))\n    axs[0][i].axis(\"tight\")\n    axs[1][i].imshow(\n        np.real((Fop.H * p_pseudoreshape[i, ix].ravel()).reshape(nwin)).T,\n        cmap=\"gray\",\n        vmin=-30,\n        vmax=30,\n        interpolation=\"none\",\n    )\n    axs[1][i].axis(\"tight\")\nplt.tight_layout()\n\n# Deblended\np_inv = Sop1.H * data_inv.ravel()\np_invreshape = p_inv.reshape(nwins[0], nwins[1], nop1[0], nop1[1])\n\nix = 16\nfig, axs = plt.subplots(2, 4, figsize=(12, 5))\nfig.suptitle(\"Deblended patches\")\nfor i in range(4):\n    axs[0][i].imshow(np.fft.fftshift(np.abs(p_invreshape[i, ix]).T, axes=1))\n    axs[0][i].axis(\"tight\")\n    axs[1][i].imshow(\n        np.real((Fop.H * p_invreshape[i, ix].ravel()).reshape(nwin)).T,\n        cmap=\"gray\",\n        vmin=-30,\n        vmax=30,\n        interpolation=\"none\",\n    )\n    axs[1][i].axis(\"tight\")\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
}