{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Automatic differentiation\n",
    "\n",
    "Being based on `pytorch`, `torchtt` can handle automatic differentiation with respect to the TT cores. \n",
    "\n",
    "\n",
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as tn\n",
    "try: \n",
    "    import torchtt as tntt\n",
    "except:\n",
    "    print('Installing torchTT...')\n",
    "    %pip install git+https://github.com/ion-g-ion/torchTT\n",
    "    import torchtt as tntt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, a function to differentiate is created and some tensors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = [2,3,4,5]\n",
    "A = tntt.randn([(n,n) for n in N],[1]+[2]*(len(N)-1)+[1])\n",
    "y = tntt.randn(N,A.R)\n",
    "x = tntt.ones(N)\n",
    "\n",
    "def f(x,A,y):\n",
    "    z = tntt.dot(A @ (x-y),(x-y))\n",
    "    return z.norm()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to compute the derivative of a scalar with respect to all cores of a TT object, the AD graph recording has to be started:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tntt.grad.watch(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `torchtt.grad.grad()` method, the gradient is computed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "val = f(x,A,y)\n",
    "grad_cores = tntt.grad.grad(val, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The variable `grad_cores` is a list of tensors representing the derivatives of `f()` with resect to the individual core entries.\n",
    "For checking, we compute the derivative of teh function with respect to one element of the core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1.2106e-09, dtype=torch.float64, grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "h = 1e-7\n",
    "x1 = x.clone()\n",
    "x1.cores[1][0,0,0] += h\n",
    "x2 = x.clone()\n",
    "x2.cores[1][0,0,0] -= h\n",
    "derivative = (f(x1,A,y)-f(x2,A,y))/(2*h)\n",
    "print(tn.abs(derivative-grad_cores[1][0,0,0])/tn.abs(derivative))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The functions `torchtt.grad.grad()` and `torchtt.grad.watch()` can take an additional list of modes `core_indices` as argument which decides which cores are watched and differentiaated with respect to."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}