{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Total Variation (TV) Regularization\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_gpu.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\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# sphinx_gallery_thumbnail_number = 3\nimport torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport pylops_gpu\n\nfrom pylops_gpu.utils.backend import device\n\ndev = device()\nprint('PyLops-gpu working on %s...' % dev)\nplt.close('all')\n\ntorch.manual_seed(0)\nnp.random.seed(1)\ndtype = torch.float32"
      ]
    },
    {
      "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 = int(101)\nx = torch.zeros(nx, dtype=dtype).to(dev)\nx[:nx//2] = 10\nx[nx//2:3*nx//4] = -5\nIop = pylops_gpu.Identity(nx, device=dev, dtype=dtype)\nnoise = torch.from_numpy(np.random.normal(0, 1, nx).astype(np.float32)).to(dev)\n\ny = Iop * (x + noise)\n\nplt.figure(figsize=(10, 5))\nplt.plot(x.cpu(), 'k', lw=3, label='x')\nplt.plot(y.cpu(), '.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_gpu.SecondDerivative(nx, device=dev, dtype=dtype)\nlamda = 1e2\n\nxinv = xinv = \\\n    pylops_gpu.optimization.leastsquares.NormalEquationsInversion(Op=Iop,\n                                                                  Regs=[D2op],\n                                                                  epsRs=[np.sqrt(lamda/2)],\n                                                                  data=y,\n                                                                  device=dev,\n                                                                  **dict(niter=30))\n\nplt.figure(figsize=(10, 5))\nplt.plot(x.cpu(), 'k', lw=3, label='x')\nplt.plot(y.cpu(), '.k', label='y=x+n')\nplt.plot(xinv.cpu(), '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_gpu.FirstDerivative(nx, device=dev, dtype=dtype)\nmu = 0.01\nlamda = 0.3\nniter_out = 50\nniter_in = 3\n\nxinv, niter = \\\n    pylops_gpu.optimization.sparsity.SplitBregman(Iop, [Dop], y, niter_out,\n                                                  niter_in, mu=mu, epsRL1s=[lamda],\n                                                  tol=1e-4, tau=1.,\n                                                  **dict(niter=30, epsI=1e-10))\n\nplt.figure(figsize=(10, 5))\nplt.plot(x.cpu(), 'k', lw=3, label='x')\nplt.plot(y.cpu(), '.k', label='y=x+n')\nplt.plot(xinv.cpu(), 'r', lw=5, label='xinv')\nplt.legend()\nplt.title('TV inversion')\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.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}