{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# MP, OMP, ISTA and FISTA\n\nThis example shows how to use the :py:class:`pylops.optimization.sparsity.omp`,\n:py:class:`pylops.optimization.sparsity.irls`,\n:py:class:`pylops.optimization.sparsity.ista`, and\n:py:class:`pylops.optimization.sparsity.fista` solvers.\n\nThese solvers can be used when the model to retrieve is supposed to have\na sparse representation in a certain domain. MP and OMP use a L0 norm and\nmathematically translates to solving the following constrained problem:\n\n\\begin{align}\\quad \\|\\mathbf{Op}\\mathbf{x}-  \\mathbf{b}\\|_2 <= \\sigma,\\end{align}\n\nwhile IRLS, ISTA and FISTA solve an uncostrained problem with a L1\nregularization term:\n\n\\begin{align}J = \\|\\mathbf{d} - \\mathbf{Op} \\mathbf{x}\\|_2 + \\epsilon \\|\\mathbf{x}\\|_1\\end{align}\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\")\nnp.random.seed(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's start with a simple example, where we create a dense mixing matrix\nand a sparse signal and we use OMP and ISTA to recover such a signal.\nNote that the mixing matrix leads to an underdetermined system of equations\n($N < M$) so being able to add some extra prior information regarding\nthe sparsity of our desired model is essential to be able to invert\nsuch a system.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "N, M = 15, 20\nA = np.random.randn(N, M)\nA = A / np.linalg.norm(A, axis=0)\nAop = pylops.MatrixMult(A)\n\nx = np.random.rand(M)\nx[x < 0.9] = 0\ny = Aop * x\n\n# MP/OMP\neps = 1e-2\nmaxit = 500\nx_mp = pylops.optimization.sparsity.omp(\n    Aop, y, niter_outer=maxit, niter_inner=0, sigma=1e-4\n)[0]\nx_omp = pylops.optimization.sparsity.omp(Aop, y, niter_outer=maxit, sigma=1e-4)[0]\n\n# IRLS\nx_irls = pylops.optimization.sparsity.irls(\n    Aop, y, nouter=50, epsI=1e-5, kind=\"model\", **dict(iter_lim=10)\n)[0]\n\n# ISTA\nx_ista = pylops.optimization.sparsity.ista(\n    Aop,\n    y,\n    niter=maxit,\n    eps=eps,\n    tol=1e-3,\n)[0]\n\nfig, ax = plt.subplots(1, 1, figsize=(8, 3))\nm, s, b = ax.stem(x, linefmt=\"k\", basefmt=\"k\", markerfmt=\"ko\", label=\"True\")\nplt.setp(m, markersize=15)\nm, s, b = ax.stem(x_mp, linefmt=\"--c\", basefmt=\"--c\", markerfmt=\"co\", label=\"MP\")\nplt.setp(m, markersize=10)\nm, s, b = ax.stem(x_omp, linefmt=\"--g\", basefmt=\"--g\", markerfmt=\"go\", label=\"OMP\")\nplt.setp(m, markersize=7)\nm, s, b = ax.stem(x_irls, linefmt=\"--m\", basefmt=\"--m\", markerfmt=\"mo\", label=\"IRLS\")\nplt.setp(m, markersize=7)\nm, s, b = ax.stem(x_ista, linefmt=\"--r\", basefmt=\"--r\", markerfmt=\"ro\", label=\"ISTA\")\nplt.setp(m, markersize=3)\nax.set_title(\"Model\", size=15, fontweight=\"bold\")\nax.legend()\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We now consider a more interesting problem problem, *wavelet deconvolution*\nfrom a signal that we assume being composed by a train of spikes convolved\nwith a certain wavelet. We will see how solving such a problem with a\nleast-squares solver such as\n:py:class:`pylops.optimization.leastsquares.regularized_inversion` does not\nproduce the expected results (especially in the presence of noisy data),\nconversely using the :py:class:`pylops.optimization.sparsity.ista` and\n:py:class:`pylops.optimization.sparsity.fista` solvers allows us\nto succesfully retrieve the input signal even in the presence of noise.\n:py:class:`pylops.optimization.sparsity.fista` shows faster convergence which\nis particularly useful for this problem.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nt = 61\ndt = 0.004\nt = np.arange(nt) * dt\nx = np.zeros(nt)\nx[10] = -0.4\nx[int(nt / 2)] = 1\nx[nt - 20] = 0.5\n\nh, th, hcenter = pylops.utils.wavelets.ricker(t[:101], f0=20)\n\nCop = pylops.signalprocessing.Convolve1D(nt, h=h, offset=hcenter, dtype=\"float32\")\ny = Cop * x\nyn = y + np.random.normal(0, 0.1, y.shape)\n\n# noise free\nxls = Cop / y\n\nxomp, nitero, costo = pylops.optimization.sparsity.omp(\n    Cop, y, niter_outer=200, sigma=1e-8\n)\n\nxista, niteri, costi = pylops.optimization.sparsity.ista(\n    Cop,\n    y,\n    niter=400,\n    eps=5e-1,\n    tol=1e-8,\n)\n\nfig, ax = plt.subplots(1, 1, figsize=(8, 3))\nax.plot(t, x, \"k\", lw=8, label=r\"$x$\")\nax.plot(t, y, \"r\", lw=4, label=r\"$y=Ax$\")\nax.plot(t, xls, \"--g\", lw=4, label=r\"$x_{LS}$\")\nax.plot(t, xomp, \"--b\", lw=4, label=r\"$x_{OMP} (niter=%d)$\" % nitero)\nax.plot(t, xista, \"--m\", lw=4, label=r\"$x_{ISTA} (niter=%d)$\" % niteri)\nax.set_title(\"Noise-free deconvolution\", fontsize=14, fontweight=\"bold\")\nax.legend()\nplt.tight_layout()\n\n# noisy\nxls = pylops.optimization.leastsquares.regularized_inversion(\n    Cop, yn, [], **dict(damp=1e-1, atol=1e-3, iter_lim=100, show=0)\n)[0]\n\nxista, niteri, costi = pylops.optimization.sparsity.ista(\n    Cop,\n    yn,\n    niter=100,\n    eps=5e-1,\n    tol=1e-5,\n)\n\nxfista, niterf, costf = pylops.optimization.sparsity.fista(\n    Cop,\n    yn,\n    niter=100,\n    eps=5e-1,\n    tol=1e-5,\n)\n\nfig, ax = plt.subplots(1, 1, figsize=(8, 3))\nax.plot(t, x, \"k\", lw=8, label=r\"$x$\")\nax.plot(t, y, \"r\", lw=4, label=r\"$y=Ax$\")\nax.plot(t, yn, \"--b\", lw=4, label=r\"$y_n$\")\nax.plot(t, xls, \"--g\", lw=4, label=r\"$x_{LS}$\")\nax.plot(t, xista, \"--m\", lw=4, label=r\"$x_{ISTA} (niter=%d)$\" % niteri)\nax.plot(t, xfista, \"--y\", lw=4, label=r\"$x_{FISTA} (niter=%d)$\" % niterf)\nax.set_title(\"Noisy deconvolution\", fontsize=14, fontweight=\"bold\")\nax.legend()\nplt.tight_layout()\n\nfig, ax = plt.subplots(1, 1, figsize=(8, 3))\nax.semilogy(costi, \"m\", lw=2, label=r\"$x_{ISTA} (niter=%d)$\" % niteri)\nax.semilogy(costf, \"y\", lw=2, label=r\"$x_{FISTA} (niter=%d)$\" % niterf)\nax.set_title(\"Cost function\", size=15, fontweight=\"bold\")\nax.set_xlabel(\"Iteration\")\nax.legend()\nax.grid(True, which=\"both\")\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
}