{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Wavelet estimation\nThis example shows how to use the :py:class:`pylops.avo.prestack.PrestackWaveletModelling` to\nestimate a wavelet from pre-stack seismic data. This problem can be written in mathematical\nform as:\n\n\\begin{align}\\mathbf{d}=  \\mathbf{G} \\mathbf{w}\\end{align}\n\nwhere $\\mathbf{G}$ is an operator that convolves an angle-variant reflectivity series\nwith the wavelet $\\mathbf{w}$ that we aim to retrieve.\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 filtfilt\n\nimport pylops\nfrom pylops.utils.wavelets import ricker\n\nplt.close(\"all\")\nnp.random.seed(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's start by creating the input elastic property profiles and wavelet\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nt0 = 501\ndt0 = 0.004\nntheta = 21\n\nt0 = np.arange(nt0) * dt0\nthetamin, thetamax = 0, 40\ntheta = np.linspace(thetamin, thetamax, ntheta)\n\n# Elastic property profiles\nvp = (\n    2000\n    + 5 * np.arange(nt0)\n    + 2 * filtfilt(np.ones(5) / 5.0, 1, np.random.normal(0, 160, nt0))\n)\nvs = 600 + vp / 2 + 3 * filtfilt(np.ones(5) / 5.0, 1, np.random.normal(0, 100, nt0))\nrho = 1000 + vp + filtfilt(np.ones(5) / 5.0, 1, np.random.normal(0, 120, nt0))\nvp[201:] += 1500\nvs[201:] += 500\nrho[201:] += 100\n\n# Wavelet\nntwav = 41\nwavoff = 10\nwav, twav, wavc = ricker(t0[: ntwav // 2 + 1], 20)\nwav_phase = np.hstack((wav[wavoff:], np.zeros(wavoff)))\n\n# vs/vp profile\nvsvp = vs / vp\n\n# Model\nm = np.stack((np.log(vp), np.log(vs), np.log(rho)), axis=1)\n\nfig, axs = plt.subplots(1, 3, figsize=(9, 7), sharey=True)\naxs[0].plot(vp, t0, \"k\", lw=3)\naxs[0].set(xlabel=\"[m/s]\", ylabel=r\"$t$ [s]\", ylim=[t0[0], t0[-1]], title=\"Vp\")\naxs[0].grid()\naxs[1].plot(vp / vs, t0, \"k\", lw=3)\naxs[1].set(title=\"Vp/Vs\")\naxs[1].grid()\naxs[2].plot(rho, t0, \"k\", lw=3)\naxs[2].set(xlabel=\"[kg/m\u00b3]\", title=\"Rho\")\naxs[2].invert_yaxis()\naxs[2].grid()\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We create now the operators to model a synthetic pre-stack seismic gather\nwith a zero-phase as well as a mixed phase wavelet.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Create operators\nWavesop = pylops.avo.prestack.PrestackWaveletModelling(\n    m, theta, nwav=ntwav, wavc=wavc, vsvp=vsvp, linearization=\"akirich\"\n)\nWavesop_phase = pylops.avo.prestack.PrestackWaveletModelling(\n    m, theta, nwav=ntwav, wavc=wavc, vsvp=vsvp, linearization=\"akirich\"\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's apply those operators to the elastic model and create some synthetic data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "d = (Wavesop * wav).reshape(ntheta, nt0).T\nd_phase = (Wavesop_phase * wav_phase).reshape(ntheta, nt0).T\n\n# add noise\ndn = d + np.random.normal(0, 3e-2, d.shape)\n\nfig, axs = plt.subplots(1, 3, figsize=(13, 7), sharey=True)\naxs[0].imshow(\n    d, cmap=\"gray\", extent=(theta[0], theta[-1], t0[-1], t0[0]), vmin=-0.1, vmax=0.1\n)\naxs[0].axis(\"tight\")\naxs[0].set(xlabel=r\"$\\theta$ [\u00b0]\", ylabel=r\"$t$ [s]\")\naxs[0].set_title(\"Data with zero-phase wavelet\", fontsize=10)\naxs[1].imshow(\n    d_phase,\n    cmap=\"gray\",\n    extent=(theta[0], theta[-1], t0[-1], t0[0]),\n    vmin=-0.1,\n    vmax=0.1,\n)\naxs[1].axis(\"tight\")\naxs[1].set_title(\"Data with non-zero-phase wavelet\", fontsize=10)\naxs[1].set_xlabel(r\"$\\theta$ [\u00b0]\")\naxs[2].imshow(\n    dn, cmap=\"gray\", extent=(theta[0], theta[-1], t0[-1], t0[0]), vmin=-0.1, vmax=0.1\n)\naxs[2].axis(\"tight\")\naxs[2].set_title(\"Noisy Data with zero-phase wavelet\", fontsize=10)\naxs[2].set_xlabel(r\"$\\theta$ [\u00b0]\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can invert the data. First we will consider noise-free data, subsequently\nwe will add some noise and add a regularization terms in the inversion\nprocess to obtain a well-behaved wavelet also under noise conditions.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wav_est = Wavesop / d.T.ravel()\nwav_phase_est = Wavesop_phase / d_phase.T.ravel()\nwavn_est = Wavesop / dn.T.ravel()\n\n# Create regularization operator\nD2op = pylops.SecondDerivative(ntwav, dtype=\"float64\")\n\n# Invert for wavelet\n(\n    wavn_reg_est,\n    istop,\n    itn,\n    r1norm,\n    r2norm,\n) = pylops.optimization.leastsquares.regularized_inversion(\n    Wavesop,\n    dn.T.ravel(),\n    [D2op],\n    epsRs=[np.sqrt(0.1)],\n    **dict(damp=np.sqrt(1e-4), iter_lim=200, show=0)\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As expected, the regularization helps to retrieve a smooth wavelet\neven under noisy conditions.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# sphinx_gallery_thumbnail_number = 3\nfig, axs = plt.subplots(2, 1, sharex=True, figsize=(8, 6))\naxs[0].plot(wav, \"k\", lw=6, label=\"True\")\naxs[0].plot(wav_est, \"--r\", lw=4, label=\"Estimated (noise-free)\")\naxs[0].plot(wavn_est, \"--g\", lw=4, label=\"Estimated (noisy)\")\naxs[0].plot(wavn_reg_est, \"--m\", lw=4, label=\"Estimated (noisy regularized)\")\naxs[0].set_title(\"Zero-phase wavelet\")\naxs[0].grid()\naxs[0].legend(loc=\"upper right\")\naxs[0].axis(\"tight\")\naxs[1].plot(wav_phase, \"k\", lw=6, label=\"True\")\naxs[1].plot(wav_phase_est, \"--r\", lw=4, label=\"Estimated\")\naxs[1].set_title(\"Wavelet with phase\")\naxs[1].grid()\naxs[1].legend(loc=\"upper right\")\naxs[1].axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally we repeat the same exercise, but this time we use a *preconditioner*.\nInitially, our preconditioner is a :py:class:`pylops.Symmetrize` operator\nto ensure that our estimated wavelet is zero-phase. After we chain\nthe :py:class:`pylops.Symmetrize` and the :py:class:`pylops.Smoothing1D`\noperators to also guarantee a smooth wavelet.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Create symmetrize operator\nSop = pylops.Symmetrize((ntwav + 1) // 2)\n\n# Create smoothing operator\nSmop = pylops.Smoothing1D(5, dims=((ntwav + 1) // 2,), dtype=\"float64\")\n\n# Invert for wavelet\nwavn_prec_est = pylops.optimization.leastsquares.preconditioned_inversion(\n    Wavesop, dn.T.ravel(), Sop, **dict(damp=np.sqrt(1e-4), iter_lim=200, show=0)\n)[0]\n\nwavn_smooth_est = pylops.optimization.leastsquares.preconditioned_inversion(\n    Wavesop, dn.T.ravel(), Sop * Smop, **dict(damp=np.sqrt(1e-4), iter_lim=200, show=0)\n)[0]\n\nfig, ax = plt.subplots(1, 1, sharex=True, figsize=(8, 3))\nax.plot(wav, \"k\", lw=6, label=\"True\")\nax.plot(wav_est, \"--r\", lw=4, label=\"Estimated (noise-free)\")\nax.plot(wavn_prec_est, \"--g\", lw=4, label=\"Estimated (noisy symmetric)\")\nax.plot(wavn_smooth_est, \"--m\", lw=4, label=\"Estimated (noisy smoothed)\")\nax.set_title(\"Zero-phase wavelet\")\nax.grid()\nax.legend(loc=\"upper right\")\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
}