{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linear solvers in the TT format\n",
    "\n",
    "This tutorial addresses solving multilinear systems $\\mathsf{Ax}=\\mathsf{b}$ in the TT format.\n",
    "\n",
    "Imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch as tn\n",
    "import datetime\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": [
    "### Small example\n",
    "\n",
    "A random tensor operator $\\mathsf{A}$ is created in the TT format. We create a random right-hand side $\\mathsf{b} = \\mathsf{Ax}$, where $\\mathsf{x}$ is a random tensor in the TT format. This way the solution of $\\mathsf{Ax}=\\mathsf{b}$ is known and we can compare it as a reference.\n",
    "This works only for small random tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = tntt.random([(4,4),(5,5),(6,6)],[1,2,3,1]) \n",
    "x = tntt.random([4,5,6],[1,2,3,1])\n",
    "b = A @ x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve the multilinear system $\\mathsf{Ax}=\\mathsf{b}$ using the method torchtt.solvers.amen_solve().\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = tntt.solvers.amen_solve(A,b, x0 = b, eps = 1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The relative residual norm and the relative error of the solution are reported:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(xs)\n",
    "print('Relative residual error ',(A@xs-b).norm()/b.norm())\n",
    "print('Relative error of the solution  ',(xs-x).norm()/x.norm())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finite differences\n",
    "\n",
    "We now solve the problem $\\Delta u = 1$ in $[0,1]^d$ with $ u = 0 $ on the entire boundary using finite differences.\n",
    "First, set the size of the problem (n is the mode size and d is the number of dimensions):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = tn.float64 \n",
    "n =  256\n",
    "d = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the finite differences matrix corresponding to the problem. The operator is constructed directly in the TT format as it follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L1d = -2*tn.eye(n, dtype = dtype)+tn.diag(tn.ones(n-1,dtype = dtype),-1)+tn.diag(tn.ones(n-1,dtype = dtype),1)\n",
    "L1d[0,1] = 0\n",
    "L1d[-1,-2] = 0\n",
    "L1d /= (n-1)\n",
    "L1d = tntt.TT(L1d, [(n,n)])\n",
    "\n",
    "L_tt = tntt.zeros([(n,n)]*d)\n",
    "for i in range(1,d-1):\n",
    "    L_tt = L_tt+tntt.eye([n]*i)**L1d**tntt.eye([n]*(d-1-i))\n",
    "L_tt = L_tt + L1d**tntt.eye([n]*(d-1)) +  tntt.eye([n]*(d-1))**L1d\n",
    "L_tt = L_tt.round(1e-14)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The right hand site of the finite difference system is also computed in the TT format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b1d = tn.ones(n, dtype=dtype)\n",
    "b1d[0] = 0\n",
    "b1d[-1] = 0\n",
    "b1d = tntt.TT(b1d)\n",
    "b_tt = b1d\n",
    "for i in range(d-1):\n",
    "    b_tt = b_tt**b1d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve the system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time = datetime.datetime.now()\n",
    "x = tntt.solvers.amen_solve(L_tt, b_tt ,x0 = b_tt, nswp = 20, eps = 1e-8, verbose = True)\n",
    "time = datetime.datetime.now() - time\n",
    "print('Relative residual: ',(L_tt@x-b_tt).norm()/b_tt.norm())\n",
    "print('Solver time: ',time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Display the structure of the TT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try one more time on the GPU (if available)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if tn.cuda.is_available():\n",
    "    cuda_dev = 'cuda:0'\n",
    "    time = datetime.datetime.now()\n",
    "    x = tntt.solvers.amen_solve(L_tt.to(cuda_dev), b_tt.to(cuda_dev) ,x0 = b_tt.to(cuda_dev), nswp = 20, eps = 1e-8, verbose = True, preconditioner=None)\n",
    "    time = datetime.datetime.now() - time\n",
    "    x = x.cpu()\n",
    "    print('Relative residual: ',(L_tt@x-b_tt).norm()/b_tt.norm())\n",
    "    print('Solver time: ',time)\n",
    "else:\n",
    "    print('GPU not available...')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "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.10.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
