{
 "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": 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": [
    "### 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": 2,
   "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_1 i_2 ... i_d} \\approx \\sum_{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": 3,
   "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": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT cores [tensor([[[-0.0048, -0.3454],\n",
      "         [-0.0144, -0.3288],\n",
      "         [-0.0240, -0.3122],\n",
      "         [-0.0335, -0.2956],\n",
      "         [-0.0431, -0.2791],\n",
      "         [-0.0527, -0.2625],\n",
      "         [-0.0622, -0.2459],\n",
      "         [-0.0718, -0.2293],\n",
      "         [-0.0814, -0.2127],\n",
      "         [-0.0909, -0.1961],\n",
      "         [-0.1005, -0.1796],\n",
      "         [-0.1101, -0.1630],\n",
      "         [-0.1196, -0.1464],\n",
      "         [-0.1292, -0.1298],\n",
      "         [-0.1388, -0.1132],\n",
      "         [-0.1483, -0.0966],\n",
      "         [-0.1579, -0.0800],\n",
      "         [-0.1675, -0.0635],\n",
      "         [-0.1770, -0.0469],\n",
      "         [-0.1866, -0.0303],\n",
      "         [-0.1962, -0.0137],\n",
      "         [-0.2057,  0.0029],\n",
      "         [-0.2153,  0.0195],\n",
      "         [-0.2249,  0.0360],\n",
      "         [-0.2344,  0.0526],\n",
      "         [-0.2440,  0.0692],\n",
      "         [-0.2536,  0.0858],\n",
      "         [-0.2632,  0.1024],\n",
      "         [-0.2727,  0.1190],\n",
      "         [-0.2823,  0.1355],\n",
      "         [-0.2919,  0.1521],\n",
      "         [-0.3014,  0.1687]]], dtype=torch.float64), tensor([[[ 2.4448e-01,  9.6477e-03],\n",
      "         [ 2.4521e-01,  8.3789e-03],\n",
      "         [ 2.4594e-01,  7.1101e-03],\n",
      "         [ 2.4667e-01,  5.8413e-03],\n",
      "         [ 2.4741e-01,  4.5725e-03],\n",
      "         [ 2.4814e-01,  3.3037e-03],\n",
      "         [ 2.4887e-01,  2.0348e-03],\n",
      "         [ 2.4960e-01,  7.6604e-04],\n",
      "         [ 2.5034e-01, -5.0277e-04],\n",
      "         [ 2.5107e-01, -1.7716e-03],\n",
      "         [ 2.5180e-01, -3.0404e-03],\n",
      "         [ 2.5253e-01, -4.3092e-03],\n",
      "         [ 2.5327e-01, -5.5780e-03],\n",
      "         [ 2.5400e-01, -6.8468e-03],\n",
      "         [ 2.5473e-01, -8.1156e-03],\n",
      "         [ 2.5546e-01, -9.3845e-03]],\n",
      "\n",
      "        [[-3.1962e-03,  2.5540e-01],\n",
      "         [-2.7735e-03,  2.5467e-01],\n",
      "         [-2.3509e-03,  2.5393e-01],\n",
      "         [-1.9282e-03,  2.5320e-01],\n",
      "         [-1.5056e-03,  2.5247e-01],\n",
      "         [-1.0829e-03,  2.5174e-01],\n",
      "         [-6.6029e-04,  2.5101e-01],\n",
      "         [-2.3764e-04,  2.5027e-01],\n",
      "         [ 1.8500e-04,  2.4954e-01],\n",
      "         [ 6.0765e-04,  2.4881e-01],\n",
      "         [ 1.0303e-03,  2.4808e-01],\n",
      "         [ 1.4529e-03,  2.4735e-01],\n",
      "         [ 1.8756e-03,  2.4661e-01],\n",
      "         [ 2.2982e-03,  2.4588e-01],\n",
      "         [ 2.7209e-03,  2.4515e-01],\n",
      "         [ 3.1435e-03,  2.4442e-01]]], dtype=torch.float64), tensor([[[-3.5310e-01,  7.8561e-04],\n",
      "         [-3.5323e-01,  5.6136e-04],\n",
      "         [-3.5336e-01,  3.3710e-04],\n",
      "         [-3.5349e-01,  1.1285e-04],\n",
      "         [-3.5362e-01, -1.1141e-04],\n",
      "         [-3.5375e-01, -3.3566e-04],\n",
      "         [-3.5388e-01, -5.5992e-04],\n",
      "         [-3.5401e-01, -7.8417e-04]],\n",
      "\n",
      "        [[ 2.6179e-04,  3.5401e-01],\n",
      "         [ 1.8703e-04,  3.5388e-01],\n",
      "         [ 1.1228e-04,  3.5375e-01],\n",
      "         [ 3.7521e-05,  3.5362e-01],\n",
      "         [-3.7234e-05,  3.5349e-01],\n",
      "         [-1.1199e-04,  3.5336e-01],\n",
      "         [-1.8674e-04,  3.5323e-01],\n",
      "         [-2.6150e-04,  3.5310e-01]]], dtype=torch.float64), tensor([[[ 1.5132e+06],\n",
      "         [ 1.5133e+06],\n",
      "         [ 1.5133e+06],\n",
      "         [ 1.5134e+06],\n",
      "         [ 1.5134e+06],\n",
      "         [ 1.5135e+06],\n",
      "         [ 1.5135e+06],\n",
      "         [ 1.5136e+06],\n",
      "         [ 1.5137e+06],\n",
      "         [ 1.5137e+06]],\n",
      "\n",
      "        [[ 1.4401e+02],\n",
      "         [ 1.1201e+02],\n",
      "         [ 8.0011e+01],\n",
      "         [ 4.8011e+01],\n",
      "         [ 1.6010e+01],\n",
      "         [-1.5991e+01],\n",
      "         [-4.7991e+01],\n",
      "         [-7.9992e+01],\n",
      "         [-1.1199e+02],\n",
      "         [-1.4399e+02]]], dtype=torch.float64)]\n",
      "Mode size  [32, 16, 8, 10]\n",
      "TT rank  [1, np.int64(2), np.int64(2), np.int64(2), 1]\n"
     ]
    }
   ],
   "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": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, np.int64(2), np.int64(2), np.int64(2), 1]\n"
     ]
    }
   ],
   "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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2.1452e-15, dtype=torch.float64)\n"
     ]
    }
   ],
   "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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT with sizes and ranks:\n",
      "N = [32, 16, 8, 10]\n",
      "R = [1, np.int64(2), np.int64(2), np.int64(2), 1]\n",
      "\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 180 compression 0.00439453125\n",
      "\n"
     ]
    }
   ],
   "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 is $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": 8,
   "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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1474.0000, dtype=torch.float64)\n",
      "TT with sizes and ranks:\n",
      "N = [3, 10]\n",
      "R = [1, 2, 1]\n",
      "\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 26 compression 0.8666666666666667\n",
      "\n"
     ]
    }
   ],
   "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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT with sizes and ranks:\n",
      "N = [10, 20, 30, 40]\n",
      "R = [1, 6, 6, 6, 1]\n",
      "\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 2100 compression 0.00875\n",
      "\n"
     ]
    }
   ],
   "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": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT with sizes and ranks:\n",
      "N = [10, 20, 30, 40]\n",
      "R = [1, 4, 4, 4, 1]\n",
      "\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 1000 compression 0.004166666666666667\n",
      "\n",
      "Error  tensor(9.9712e-07, dtype=torch.float64)\n"
     ]
    }
   ],
   "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": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT with sizes and ranks:\n",
      "N = [10, 20, 30, 40]\n",
      "R = [1, 2, 2, 2, 1]\n",
      "\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 300 compression 0.00125\n",
      "\n",
      "Error  tensor(0.0010, dtype=torch.float64)\n"
     ]
    }
   ],
   "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": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[1., 1., 1., 1.],\n",
      "         [1., 1., 1., 1.],\n",
      "         [1., 1., 1., 1.]],\n",
      "\n",
      "        [[1., 1., 1., 1.],\n",
      "         [1., 1., 1., 1.],\n",
      "         [1., 1., 1., 1.]]], dtype=torch.float64)\n"
     ]
    }
   ],
   "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": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0., 0., 0., 0.],\n",
      "         [0., 0., 0., 0.],\n",
      "         [0., 0., 0., 0.]],\n",
      "\n",
      "        [[0., 0., 0., 0.],\n",
      "         [0., 0., 0., 0.],\n",
      "         [0., 0., 0., 0.]]], dtype=torch.float64)\n"
     ]
    }
   ],
   "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": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT-matrix with sizes and ranks:\n",
      "M = [10, 20, 30]\n",
      "N = [10, 20, 30]\n",
      "R = [1, 1, 1, 1]\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 1400 compression 3.888888888888889e-05\n",
      "\n"
     ]
    }
   ],
   "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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT with sizes and ranks:\n",
      "N = [3, 4, 5, 6, 7]\n",
      "R = [1, 2, 5, 5, 2, 1]\n",
      "\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 245 compression 0.09722222222222222\n",
      "\n",
      "TT-matrix with sizes and ranks:\n",
      "M = [3, 4, 5, 6, 7]\n",
      "N = [7, 6, 5, 10, 2]\n",
      "R = [1, 2, 5, 5, 2, 1]\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 1535 compression 0.0001450302343159486\n",
      "\n"
     ]
    }
   ],
   "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": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Var =  0.8806055022985024  (has to be comparable to 1.0)\n",
      "Var =  3.1136021195814223  (has to be comparable to 4.0)\n",
      "Var =  0.0008934094513731358  (has to be comparable to 0.001)\n",
      "Var =  0.6374323197748297  (has to be comparable to 1.0)\n"
     ]
    }
   ],
   "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": {
  "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
}