{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Total Variation (TV) Regularization\n\nThis set of examples shows how to add Total Variation (TV) regularization to an\ninverse problem in order to enforce blockiness in the reconstructed model.\n\nTo do so we will use the generalizated Split Bregman iterations by means of\n:func:`pylops.optimization.sparsity.SplitBregman` solver.\n\nThe first example is concerned with denoising of a piece-wise step function\nwhich has been contaminated by noise. The forward model is:\n\n\\begin{align}\\mathbf{y} = \\mathbf{x} + \\mathbf{n}\\end{align}\n\nmeaning that we have an identity operator ($\\mathbf{I}$) and inverting\nfor $\\mathbf{x}$ from $\\mathbf{y}$ is impossible without adding\nprior information. We will enforce blockiness in the solution by adding a\nregularization term that enforces sparsity in the first derivative of\nthe solution:\n\n\\begin{align}J = \\mu/2  ||\\mathbf{y} - \\mathbf{I} \\mathbf{x}||_2 +\n        || \\nabla \\mathbf{x}||_1\\end{align}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n\n# sphinx_gallery_thumbnail_number = 5\nimport numpy as np\n\nimport pylops\n\nplt.close(\"all\")\nnp.random.seed(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's start by creating the model and data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nx = 101\nx = np.zeros(nx)\nx[: nx // 2] = 10\nx[nx // 2 : 3 * nx // 4] = -5\n\nIop = pylops.Identity(nx)\n\nn = np.random.normal(0, 1, nx)\ny = Iop * (x + n)\n\nplt.figure(figsize=(10, 5))\nplt.plot(x, \"k\", lw=3, label=\"x\")\nplt.plot(y, \".k\", label=\"y=x+n\")\nplt.legend()\nplt.title(\"Model and data\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To start we will try to use a simple L2 regularization that enforces\nsmoothness in the solution. We can see how denoising is succesfully achieved\nbut the solution is much smoother than we wish for.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "D2op = pylops.SecondDerivative(nx, edge=True)\nlamda = 1e2\n\nxinv = pylops.optimization.leastsquares.regularized_inversion(\n    Iop, y, [D2op], epsRs=[np.sqrt(lamda / 2)], **dict(iter_lim=30)\n)[0]\n\nplt.figure(figsize=(10, 5))\nplt.plot(x, \"k\", lw=3, label=\"x\")\nplt.plot(y, \".k\", label=\"y=x+n\")\nplt.plot(xinv, \"r\", lw=5, label=\"xinv\")\nplt.legend()\nplt.title(\"L2 inversion\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we impose blockiness in the solution using the Split Bregman solver\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Dop = pylops.FirstDerivative(nx, edge=True, kind=\"backward\")\nmu = 0.01\nlamda = 0.3\nniter_out = 50\nniter_in = 3\n\nxinv = pylops.optimization.sparsity.splitbregman(\n    Iop,\n    y,\n    [Dop],\n    niter_outer=niter_out,\n    niter_inner=niter_in,\n    mu=mu,\n    epsRL1s=[lamda],\n    tol=1e-4,\n    tau=1.0,\n    **dict(iter_lim=30, damp=1e-10)\n)[0]\n\nplt.figure(figsize=(10, 5))\nplt.plot(x, \"k\", lw=3, label=\"x\")\nplt.plot(y, \".k\", label=\"y=x+n\")\nplt.plot(xinv, \"r\", lw=5, label=\"xinv\")\nplt.legend()\nplt.title(\"TV inversion\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, we repeat the same exercise on a 2-dimensional image. In this case\nwe mock a medical imaging problem: the data is created by appling a 2D\nFourier Transform to the input model and by randomly sampling 60% of\nits values.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "x = np.load(\"../testdata/optimization/shepp_logan_phantom.npy\")\nx = x / x.max()\nny, nx = x.shape\n\nperc_subsampling = 0.6\nnxsub = int(np.round(ny * nx * perc_subsampling))\niava = np.sort(np.random.permutation(np.arange(ny * nx))[:nxsub])\nRop = pylops.Restriction(ny * nx, iava, axis=0, dtype=np.complex128)\nFop = pylops.signalprocessing.FFT2D(dims=(ny, nx))\n\nn = np.random.normal(0, 0.0, (ny, nx))\ny = Rop * Fop * (x.ravel() + n.ravel())\nyfft = Fop * (x.ravel() + n.ravel())\nyfft = np.fft.fftshift(yfft.reshape(ny, nx))\n\nymask = Rop.mask(Fop * (x.ravel()) + n.ravel())\nymask = ymask.reshape(ny, nx)\nymask.data[:] = np.fft.fftshift(ymask.data)\nymask.mask[:] = np.fft.fftshift(ymask.mask)\n\nfig, axs = plt.subplots(1, 3, figsize=(14, 5))\naxs[0].imshow(x, vmin=0, vmax=1, cmap=\"gray\")\naxs[0].set_title(\"Model\")\naxs[0].axis(\"tight\")\naxs[1].imshow(np.abs(yfft), vmin=0, vmax=1, cmap=\"rainbow\")\naxs[1].set_title(\"Full data\")\naxs[1].axis(\"tight\")\naxs[2].imshow(np.abs(ymask), vmin=0, vmax=1, cmap=\"rainbow\")\naxs[2].set_title(\"Sampled data\")\naxs[2].axis(\"tight\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's attempt now to reconstruct the model using the Split Bregman\nwith anisotropic TV regularization (aka sum of L1 norms of the\nfirst derivatives over x and y):\n\n\\begin{align}J = \\mu/2 ||\\mathbf{y} - \\mathbf{R} \\mathbf{F} \\mathbf{x}||_2\n        + || \\nabla_x \\mathbf{x}||_1 + || \\nabla_y \\mathbf{x}||_1\\end{align}\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Dop = [\n    pylops.FirstDerivative(\n        (ny, nx), axis=0, edge=False, kind=\"backward\", dtype=np.complex128\n    ),\n    pylops.FirstDerivative(\n        (ny, nx), axis=1, edge=False, kind=\"backward\", dtype=np.complex128\n    ),\n]\n\n# TV\nmu = 1.5\nlamda = [0.1, 0.1]\nniter_out = 20\nniter_in = 10\n\nxinv = pylops.optimization.sparsity.splitbregman(\n    Rop * Fop,\n    y.ravel(),\n    Dop,\n    niter_outer=niter_out,\n    niter_inner=niter_in,\n    mu=mu,\n    epsRL1s=lamda,\n    tol=1e-4,\n    tau=1.0,\n    show=False,\n    **dict(iter_lim=5, damp=1e-4)\n)[0]\nxinv = np.real(xinv.reshape(ny, nx))\n\nfig, axs = plt.subplots(1, 2, figsize=(9, 5))\naxs[0].imshow(x, vmin=0, vmax=1, cmap=\"gray\")\naxs[0].set_title(\"Model\")\naxs[0].axis(\"tight\")\naxs[1].imshow(xinv, vmin=0, vmax=1, cmap=\"gray\")\naxs[1].set_title(\"TV Inversion\")\naxs[1].axis(\"tight\")\nplt.tight_layout()\n\nfig, axs = plt.subplots(2, 1, figsize=(10, 5))\naxs[0].plot(x[ny // 2], \"k\", lw=5, label=\"x\")\naxs[0].plot(xinv[ny // 2], \"r\", lw=3, label=\"xinv TV\")\naxs[0].set_title(\"Horizontal section\")\naxs[0].legend()\naxs[1].plot(x[:, nx // 2], \"k\", lw=5, label=\"x\")\naxs[1].plot(xinv[:, nx // 2], \"r\", lw=3, label=\"xinv TV\")\naxs[1].set_title(\"Vertical section\")\naxs[1].legend()\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Note that more optimized variations of the Split Bregman algorithm have been\nproposed in the literature for this specific problem, both improving the\noverall quality of the inversion and the speed of convergence.\n\nIn PyLops we however prefer to implement the generalized Split Bergman\nalgorithm as this can used for any sort of problem where we wish to\nadd any number of L1 and/or L2 regularization terms to the cost function\nto minimize.\n\n"
      ]
    }
  ],
  "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
}