{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TT decomposition in torchtt\n",
    "\n",
    "This notebook is a tutorial on how to use the basic functionalities of the `torchtt` package. \n",
    "\n",
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": [
    "### Decomposition of a full tensor in TT format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now create a 4d `torch.tensor` which we will use later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tens_full = tn.reshape(tn.arange(32*16*8*10, dtype = tn.float64),[32,16,8,10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The TT approximation of a given tensor is $\\mathsf{x}_{i_1i_2...i_d} \\approx \\sum\\limits_{r_1,...,r_{d-1}=1}^{R_1,...,R_{d-1}} \\mathsf{g}^{(1)}_{1i_1r_1}\\cdots\\mathsf{g}^{(d)}_{r_{d-1}i_d1} $. Using the constructor `torchtt.TT()` a full tensor can be decomposed in the TT format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tens_tt = tntt.TT(tens_full)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The newly instantiated object contains the cores as a list, the mode sizes and the rank."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('TT cores', tens_tt.cores)\n",
    "print('Mode size ', tens_tt.N)\n",
    "print('TT rank ', tens_tt.R)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the TT decomposition is not exact in most of the cases, an approximation is made. If the argument `eps` is provided to the `torchtt.TT()` function the decomposition can be performed upt to the given relative accuracy.\n",
    "Moreover the maximum rank can also be provided as the argument `rmax`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tens_full2 = tens_full+1e-5*tn.randn(tens_full.shape, dtype=tens_full.dtype)\n",
    "tens_tt2 = tntt.TT(tens_full2, eps = 1e-4)\n",
    "print(tens_tt2.R)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The original tensor can be recovered using the `torchtt.TT.full()` method (also check if it equals the original full tensor):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tens_full_rec = tens_tt.full()\n",
    "print(tn.linalg.norm(tens_full-tens_full_rec)/tn.linalg.norm(tens_full))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `print()` function, information about the newly created `torchtt.TT` instance can be displayed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tens_tt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tensor operators\n",
    "\n",
    "As a generalization of the matrix vector algebra, one can define tensor operators that act on tensors. If the tensor si $d$-dimensional, the tensor operator will be $2d$-dimensional.\n",
    "The goal is to perform a product $\\mathsf{Ax}\\in\\mathbb{M_1\\times \\cdots \\times M_d}$ between a tensor $\\mathsf{x}\\in\\mathbb{R}^{N_1\\times \\cdots \\times N_d}$ and the operator $\\mathsf{A}\\in \\mathbb{R}^{(M_1\\times \\cdots \\times M_d)\\times(N_1\\times \\cdots \\times N_d)}$. For the operators the following TT matrix format is used $\\mathsf{A}_{i_1...i_d,j_1...j_d}\\approx \\sum\\limits_{r_1,...,r_{d-1}=1}^{R_1,...,R_{d-1}} \\mathsf{g}^{(1)}_{1i_1j_1r_1}\\cdots\\mathsf{g}^{(d)}_{r_{d-1}i_dj_d1}$.\n",
    "\n",
    "If a tensor operator needs to be decomposed from full, the additional argument `shape` of the `torchtt.TT()` constructor has to be used to provide the shape.\n",
    "If the tensor operator has the shape $(M_1\\times \\cdots \\times M_d)\\times(N_1\\times \\cdots \\times N_d)$ the argument must be passed as `[(M1,N1),(M2,N2),(M3,N2),...]`.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_full = tn.reshape(tn.arange(8*4*6*3*7*9, dtype = tn.float64),[8,4,6,3,7,9])\n",
    "# create an instance of torchtt.TT\n",
    "A_ttm = tntt.TT(A_full, eps = 1e-12, shape = [(8,3),(4,7),(6,9)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Slicing\n",
    "\n",
    "Slicing operation can be performed on a tensor in TT format. If all the dimensions are indexed with an integer and the multiindices are valid, a torch.tensor with the corresponding value is returned. Slices can be also used, however the returned object in this case is again a torchtt.TT instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tens_tt[1,2,3,4])\n",
    "print(tens_tt[1,1:4,2,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TT rank rounding\n",
    "\n",
    "In some cases the TT rank becomes too large and a reduction is desired. The goal is to perform a reduction of the rank while maintaining an accuracy.\n",
    "The problem statement of the rounding operation is: given a tensor $\\mathsf{x}$ in the TT format with the TT rank $\\mathbf{R}$ and an $\\epsilon>0$, find a tensor $\\tilde{\\mathsf{x}}$ with TT rank $\\tilde{\\mathbf{R}}\\leq \\mathbf{R}$ such that $ ||\\mathsf{x}-\\tilde{\\mathsf{x}}||_F\\leq \\epsilon || \\mathsf{x} ||_F$.\n",
    "This is implemented using the member method of a TT object `torchtt.TT.round()`. The argument `epsilon` is passed to the function as well as the optional argument `rmax` which also restricts the rank of the rounding.\n",
    "\n",
    "We will create a tensor of TT rank $(1,6,6,6,1)$ in the TT format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t1 = tntt.randn([10,20,30,40],[1,2,2,2,1])\n",
    "t2 = tntt.randn([10,20,30,40],[1,2,2,2,1])\n",
    "t3 = tntt.randn([10,20,30,40],[1,2,2,2,1])\n",
    "t1, t2, t3 = t1/t1.norm(), t2/t2.norm(), t3/t3.norm()\n",
    "tt = t1+1e-3*t2+1e-6*t3\n",
    "t_full = tt.full()\n",
    "print(tt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rounding the tensor to a relative `epsilon` of 1e-5 yields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt1 = tt.round(1e-5)\n",
    "print(tt1)\n",
    "print('Error ',tn.linalg.norm(tt1.full()-tt.full())/tn.linalg.norm(tt.full()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is equivalent to removing the `t3` from `tt` and the error will be less than 1e-6.\n",
    "If a truncation with `epsilon=1e-2` is done, the resulting tensor will have the rank `[1,2,2,2,1]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt1 = tt.round(1e-2)\n",
    "print(tt1)\n",
    "print('Error ',tn.linalg.norm(tt1.full()-tt.full())/tn.linalg.norm(tt.full()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The maximum rank of a truncation can also be provided as argument."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Special tensors\n",
    "\n",
    "Some tensors can be directly constructed in the TT format: the one tensor, the zeros tensor, the identity tensor operator adn random tensors with a given rank."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The one tensor can be created directly in the TT format using `torchtt.ones()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tntt.ones([2,3,4]).full())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The zero tensor ca be created in the TT format using `torchtt.zeros()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tntt.zeros([2,3,4]).full())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The identity tensor operator is created using `torchtt.eye()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tntt.eye([10,20,30]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tensors with random TT cores and a given rank can be created with `torchtt.random()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tntt.random([3,4,5,6,7],[1,2,5,5,2,1]))\n",
    "print(tntt.random([(3,7),(4,6),(5,5),(6,10),(7,2)],[1,2,5,5,2,1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Random tensors with a given rank and random entries with expected value 0 and given variance can be created using `torchtt.randn()`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Variance 1.0\n",
    "x = tntt.randn([30]*5,[1,8,16,16,8,1])\n",
    "x_full = x.full()\n",
    "print('Var = ',tn.std(x_full).numpy()**2,' (has to be comparable to 1.0)')\n",
    "\n",
    "# Variance 4.0\n",
    "x = tntt.randn([30]*5,[1,8,16,16,8,1],var = 4.0)\n",
    "x_full = x.full()\n",
    "print('Var = ',tn.std(x_full).numpy()**2,' (has to be comparable to 4.0)')\n",
    "\n",
    "# Variance 0.01\n",
    "x = tntt.randn([30]*5,[1,8,16,16,8,1], var = 0.001)\n",
    "x_full = x.full()\n",
    "print('Var = ',tn.std(x_full).numpy()**2,' (has to be comparable to 0.001)')\n",
    "\n",
    "# Variance 1.0 (longer train)\n",
    "x = tntt.randn([10]*7, [1,4,4,4,4,4,4,1], var = 1.0)\n",
    "x_full = x.full()\n",
    "print('Var = ',tn.std(x_full).numpy()**2,' (has to be comparable to 1.0)')"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "df6fc3a9b7a9c6f4b0308ab6eb361a4cabbf6b5db181383d07014ff4304e5cb3"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 64-bit ('pytorch': conda)",
   "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.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
