{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# 01. Automatic Differentiation\nThis tutorial focuses on one of the two main benefits of re-implementing\nsome of PyLops linear operators within the PyTorch framework, namely the\npossibility to perform Automatic Differentiation (AD) on chains of operators\nwhich can be:\n\n- native PyTorch mathematical operations (e.g., :func:`torch.log`,\n  :func:`torch.sin`, :func:`torch.tan`, :func:`torch.pow`, ...)\n- neural network operators in :mod:`torch.nn`\n- PyLops and/or PyLops-gpu linear operators\n\nThis opens up many opportunities, such as easily including linear regularization\nterms to nonlinear cost functions or using linear preconditioners with nonlinear\nmodelling operators.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport torch\nimport matplotlib.pyplot as plt\nfrom torch.autograd import gradcheck\n\nimport pylops_gpu\nfrom pylops_gpu.utils.backend import device\n\ndev = device()\nplt.close('all')\nnp.random.seed(10)\ntorch.manual_seed(10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In this example we consider a simple multidimensional functional:\n\n\\begin{align}\\mathbf{y} = \\mathbf{A} sin(\\mathbf{x})\\end{align}\n\nand we use AD to compute the gradient with respect to the input vector\nevaluated at $\\mathbf{x}=\\mathbf{x}_0$ :\n$\\mathbf{g} = d\\mathbf{y} / d\\mathbf{x} |_{\\mathbf{x}=\\mathbf{x}_0}$.\n\nLet's start by defining the Jacobian:\n\n  .. math::\n       \\textbf{J} = \\begin{bmatrix}\n       dy_1 / dx_1 & ... & dy_1 / dx_M \\\\\n       ... & ... & ... \\\\\n       dy_N / dx_1 & ... & dy_N / dx_M\n       \\end{bmatrix} = \\begin{bmatrix}\n       a_{11} cos(x_1) & ... & a_{1M} cos(x_M) \\\\\n       ... & ... & ... \\\\\n       a_{N1} cos(x_1) & ... & a_{NM} cos(x_M)\n       \\end{bmatrix} = \\textbf{A} cos(\\mathbf{x})\n\nSince both input and output are multidimensional,\nPyTorch ``backward`` actually computes the product between the transposed\nJacobian and a vector $\\mathbf{v}$:\n$\\mathbf{g}=\\mathbf{J^T} \\mathbf{v}$.\n\nTo validate the correctness of the AD result, we can in this simple case\nalso compute the Jacobian analytically and apply it to the same vector\n$\\mathbf{v}$ that we have provided to PyTorch ``backward``.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nx, ny = 10, 6\nx0 = torch.arange(nx, dtype=torch.double, requires_grad=True)\n\n# Forward\nA = torch.normal(0., 1., (ny, nx), dtype=torch.double)\nAop = pylops_gpu.TorchOperator(pylops_gpu.MatrixMult(A))\ny = Aop.apply(torch.sin(x0))\n\n# AD\nv = torch.ones(ny, dtype=torch.double)\ny.backward(v, retain_graph=True)\nadgrad = x0.grad\n\n# Analytical\nJ = (A * torch.cos(x0))\nanagrad = torch.matmul(J.T, v)\n\nprint('Input: ', x0)\nprint('AD gradient: ', adgrad)\nprint('Analytical gradient: ', anagrad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Similarly we can use the :func:`torch.autograd.gradcheck` directly from\nPyTorch. Note that doubles must be used for this to succeed with very small\n`eps` and `atol`\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "input = (torch.arange(nx, dtype=torch.double, requires_grad=True),\n         Aop.matvec, Aop.rmatvec, Aop.pylops, Aop.device)\ntest = gradcheck(Aop.Top, input, eps=1e-6, atol=1e-4)\nprint(test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Note that while matrix-vector multiplication could have been performed using\nthe native PyTorch operator :func:`torch.matmul`, in this case we have shown\nthat we are also able to use a PyLops-gpu operator wrapped in\n:class:`pylops_gpu.TorchOperator`. As already mentioned, this gives us the\nability to use much more complex linear operators provided by PyLops within\na chain of mixed linear and nonlinear AD-enabled operators.\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.6.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}