{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Non-stationary Convolution\nThis example shows how to use the :py:class:`pylops.signalprocessing.NonStationaryConvolve1D`\nand :py:class:`pylops.signalprocessing.NonStationaryConvolve2D` operators to perform non-stationary\nconvolution between two signals.\n\nSimilar to their stationary counterparts, these operators can be used in the forward model of\nseveral common application in signal processing that require filtering of an input signal for a\ntime- or space-varying instrument response.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\nfrom scipy.signal.windows import gaussian\n\nimport pylops\nfrom pylops.utils.wavelets import ricker\n\nplt.close(\"all\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We will start by creating a zero signal of length `nt` and we will\nplace a comb of unitary spikes. We also create a non-stationary filter defined by\n5 equally spaced [Ricker wavelets](http://subsurfwiki.org/wiki/Ricker_wavelet)\nwith dominant frequencies of $f = 10, 15, 20, 25$ and $30$ Hz.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nt = 601\ndt = 0.004\nt = np.arange(nt) * dt\ntw = ricker(t[:51], f0=5)[1]\n\nfs = [10, 15, 20, 25, 30]\nwavs = np.stack([ricker(t[:51], f0=f)[0] for f in fs])\n\nCop = pylops.signalprocessing.NonStationaryConvolve1D(\n    dims=nt, hs=wavs, ih=(101, 201, 301, 401, 501)\n)\n\nx = np.zeros(nt)\nx[64 : nt - 64 : 64] = 1.0\n\ny = Cop @ x\n\nplt.figure(figsize=(10, 3))\nplt.plot(t, x, \"k\")\nplt.plot(t, y, \"k\")\nplt.xlabel(\"Time [sec]\")\nplt.title(\"Input and output\")\nplt.xlim(0, t[-1])\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's now visualize the filters in time and frequency domain\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(10, 3))\nplt.pcolormesh(t, tw, Cop.hsinterp.T, cmap=\"gray\")\nplt.xlabel(\"Time [sec]\")\nplt.ylabel(\"Wavelet Time [sec]\")\nplt.title(\"Wavelets\")\nplt.xlim(0, t[-1])\nplt.tight_layout()\n\n# Spectra\nf = np.fft.rfftfreq(nt, dt)\nSh = np.abs(np.fft.rfft(Cop.hsinterp.T, n=nt, axis=0))\n\nplt.figure(figsize=(10, 3))\nplt.pcolormesh(t, f, Sh, cmap=\"jet\", vmax=5e0)\nplt.ylabel(\"Frequency [Hz]\")\nplt.xlabel(\"Time [sec]\")\nplt.title(\"Wavelet spectrogram\")\nplt.ylim(0, 70)\nplt.xlim(0, t[-1])\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, we repeat the same exercise with a 2-dimensional non-stationary\nfilter\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nx, nz = 601, 501\n\nwav1a, _, _ = ricker(t[:17], f0=12)\nwav1b, _, _ = ricker(t[:17], f0=30)\nwav2a = gaussian(35, 2.0)\nwav2b = gaussian(35, 4.0)\n\nwav11 = np.outer(wav1a, wav2a[np.newaxis]).T\nwav12 = np.outer(wav1b, wav2a[np.newaxis]).T\nwav21 = np.outer(wav1b, wav2b[np.newaxis]).T\nwav22 = np.outer(wav1a, wav2b[np.newaxis]).T\nwavsize = wav11.shape\n\nhs = np.zeros((2, 2, *wavsize))\nhs[0, 0] = wav11\nhs[0, 1] = wav12\nhs[1, 0] = wav21\nhs[1, 1] = wav22\n\nfig, axs = plt.subplots(2, 2, figsize=(10, 5))\naxs[0, 0].imshow(wav11, cmap=\"gray\")\naxs[0, 0].axis(\"tight\")\naxs[0, 0].set_title(r\"$H_{1,1}$\")\naxs[0, 1].imshow(wav12, cmap=\"gray\")\naxs[0, 1].axis(\"tight\")\naxs[0, 1].set_title(r\"$H_{1,2}$\")\naxs[1, 0].imshow(wav21, cmap=\"gray\")\naxs[1, 0].axis(\"tight\")\naxs[1, 0].set_title(r\"$H_{2,1}$\")\naxs[1, 1].imshow(wav22, cmap=\"gray\")\naxs[1, 1].axis(\"tight\")\naxs[1, 1].set_title(r\"$H_{2,2}$\")\nplt.tight_layout()\n\nCop = pylops.signalprocessing.NonStationaryConvolve2D(\n    hs=hs, ihx=(201, 401), ihz=(201, 401), dims=(nx, nz), engine=\"numba\"\n)\n\nx = np.zeros((nx, nz))\nline1 = (np.arange(nx) * np.tan(np.deg2rad(25))).astype(int) + (nz - 1) // 4\nline2 = (np.arange(nx) * np.tan(np.deg2rad(-25))).astype(int) + (3 * (nz - 1)) // 4\nx[np.arange(nx), np.clip(line1, 0, nz - 1)] = 1.0\nx[np.arange(nx), np.clip(line2, 0, nz - 1)] = -1.0\n\ny = Cop @ x\n\nfig, axs = plt.subplots(1, 2, figsize=(10, 4))\naxs[0].imshow(x.T, cmap=\"gray\", vmin=-1.0, vmax=1.0)\naxs[0].axis(\"tight\")\naxs[0].set_title(\"Input\")\naxs[1].imshow(y.T, cmap=\"gray\", vmin=-3.0, vmax=3.0)\naxs[1].axis(\"tight\")\naxs[1].set_title(\"Output\")\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
}