{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nMP, OMP, ISTA and FISTA\n=======================\n\nThis example shows how to use the :py:class:`pylops.optimization.sparsity.OMP`,\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 uses a L0 norm and\nmathematically translates to solving the following constrained problem:\n\n\\begin{align}||\\mathbf{x}||_0 \\quad  subj. to \\quad ||\\mathbf{Op}\\mathbf{x}-\\mathbf{b}||_2 <= \\sigma,\\end{align}\n\nwhile ISTA and FISTA solve an uncostrained problem with a L1 regularization term:\n\n\\begin{align}J = ||\\mathbf{d} - \\mathbf{Op} \\mathbf{x}||_2 + \\epsilon ||\\mathbf{x}||_1\\end{align}\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport matplotlib.pyplot as plt\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# ISTA\neps = 1e-2\nmaxit = 500\nx_mp = pylops.optimization.sparsity.OMP(Aop, y, maxit, niter_inner=0,\n                                        sigma=1e-4)[0]\nx_omp = pylops.optimization.sparsity.OMP(Aop, y, maxit, sigma=1e-4)[0]\nx_ista = pylops.optimization.sparsity.ISTA(Aop, y, maxit, eps=eps,\n                                           tol=1e-3, returninfo=True)[0]\n\nfig, ax = plt.subplots(1, 1, figsize=(8, 3))\nm, s, b = ax.stem(x, linefmt='k', basefmt='k',\n        markerfmt='ko', label='True')\nplt.setp(m, markersize = 15)\nm, s, b = ax.stem(x_mp, linefmt='--c', basefmt='--c',\n        markerfmt='co', label='MP')\nplt.setp(m, markersize = 10)\nm, s, b = ax.stem(x_omp, linefmt='--g', basefmt='--g',\n        markerfmt='go', label='OMP')\nplt.setp(m, markersize = 7)\nm, s, b = ax.stem(x_ista, linefmt='--r',\n        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.RegularizedInversion` 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] = -.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,\n                                         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 = \\\n    pylops.optimization.sparsity.OMP(Cop, y, niter_outer=200, sigma=1e-8)\n\nxista, niteri, costi = \\\n    pylops.optimization.sparsity.ISTA(Cop, y, niter=400, eps=5e-1,\n                                      tol=1e-8, returninfo=True)\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)\n\nax.set_title('Noise-free deconvolution', fontsize=14, fontweight='bold')\nax.legend()\nplt.tight_layout()\n\n# noisy\nxls = \\\n    pylops.optimization.leastsquares.RegularizedInversion(Cop, [], yn,\n                                                          returninfo=False,\n                                                          **dict(damp=1e-1,\n                                                                 atol=1e-3,\n                                                                 iter_lim=100,\n                                                                 show=0))\n\nxista, niteri, costi = \\\n    pylops.optimization.sparsity.ISTA(Cop, yn, niter=100, eps=5e-1,\n                                      tol=1e-5, returninfo=True)\n\nxfista, niterf, costf = \\\n    pylops.optimization.sparsity.FISTA(Cop, yn, niter=100, eps=5e-1,\n                                       tol=1e-5, returninfo=True)\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.6.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}