{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AMEN and DMRG for fast TT operations\n",
    "\n",
    "The torchtt package includes DMRG and AMEN schemes for fast matrix vector product and elementwise inversion in the TT format.\n",
    "\n",
    "Imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch as tn\n",
    "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\n",
    "import datetime"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Efficient matrix vector product\n",
    "\n",
    "When performing the multiplication between a a TT matrix and a TT tensor the rank of the result is the product of the ranks of the inputs. \n",
    "Therefore rank rounding has to be performed. This increases the complexity to $\\mathcal{O}(Ndr^6)$. In order to overcome this, Oseledets proposed in \"DMRG Approach to Fast Linear Algebra in the TT-Format\" the DMRG optimization scheme to reduce the complexity. This feature is implemented in torchtt by the member function fast_matvec() of the TT class. An example is showed in the following."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a random TT object and a TT matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 32 # mode size\n",
    "A = tntt.random([(n,n)]*8,[1]+7*[6]+[1]) # random array\n",
    "x = tntt.random([n]*8,[1]+7*[5]+[1]) # random tensor "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Increase the rank without adding redundant information. The multiplication performed in this case is actually equivalent to $32\\mathbf{\\mathsf{Ax}}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TT-matrix with sizes and ranks:\n",
      "M = [32, 32, 32, 32, 32, 32, 32, 32]\n",
      "N = [32, 32, 32, 32, 32, 32, 32, 32]\n",
      "R = [1, 48, 48, 48, 48, 48, 48, 48, 1]\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 14254080 compression 1.179069862577986e-17\n",
      "\n",
      "TT with sizes and ranks:\n",
      "N = [32, 32, 32, 32, 32, 32, 32, 32]\n",
      "R = [1, 60, 60, 60, 60, 60, 60, 60, 1]\n",
      "\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 695040 compression 6.32135197520256e-07\n",
      "\n"
     ]
    }
   ],
   "source": [
    "A = A + A + A + A - A + A - A + A\n",
    "x = x + x + x + x + x + x + x + x - x + x - x + x \n",
    "print(A)\n",
    "print(x)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform the TT matvec directly and round the result. The runtime is reported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time classic  0:01:16.990921\n"
     ]
    }
   ],
   "source": [
    "tme = datetime.datetime.now()\n",
    "y = (A @ x).round(1e-12) \n",
    "tme = datetime.datetime.now() - tme \n",
    "print('Time classic ', tme)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time run the fast matvec routine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time DMRG     0:00:03.134253\n"
     ]
    }
   ],
   "source": [
    "tme = datetime.datetime.now()\n",
    "yf = A.fast_matvec(x)\n",
    "tme = datetime.datetime.now() - tme \n",
    "print('Time DMRG    ', tme)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check if the error is the same (debugging purpose)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relative error  1.762527613735165e-14\n"
     ]
    }
   ],
   "source": [
    "print('Relative error ',(y-yf).norm().numpy()/y.norm().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A second routine is the `torchtt.fast_mv()`. The method is described in [https://arxiv.org/pdf/2410.19747](https://arxiv.org/pdf/2410.19747). This works well for tensors in QTT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time fast 2   0:00:01.307426\n"
     ]
    }
   ],
   "source": [
    "A = tntt.random([(2,2)]*8,[1]+7*[6]+[1]) # random array\n",
    "x = tntt.random([2]*8,[1]+7*[5]+[1]) # random tensor \n",
    "for _ in range(8): A+=A\n",
    "for _ in range(8): x+=x\n",
    "\n",
    "tme = datetime.datetime.now()\n",
    "yf2 = tntt.fast_mv(A, x)\n",
    "tme = datetime.datetime.now() - tme \n",
    "print('Time fast 2  ', tme)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Elementwise division in the TT format\n",
    "\n",
    "One other basic linear algebra function that cannot be done without optimization is the elementwise division of two tensors in the TT format.\n",
    "In contrast to the elemntwise multiplication (where the resulting TT cores can be explicitly computed), the elementwise inversion has to be solved by means of an optimization problem (the method of choice is AMEN). \n",
    "\n",
    "The operator \"/\" can be used  for elemntwise division between tensors. Moreover one can use \"/\" between a scalar and a  torchtt.TT instance."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create 2 tensors:\n",
    " - $\\mathsf{x}_{i_1i_2i_3i_4} = 2 + i_1$\n",
    " - $\\mathsf{y}_{i_1i_2i_3i_4} = i_1^2+i_2+i_3+1$\n",
    " \n",
    "and express them in the TT format.\n",
    "For both of them a TT decomposition of the elemmentwise inverse cannot be explicitly formed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = [32,50,44,64]\n",
    "I = tntt.meshgrid([tn.arange(n,dtype = tn.float64) for n in N])\n",
    "x = 2+I[0]\n",
    "x = x.round(1e-15)\n",
    "y = I[0]*I[0]+I[1]+I[2]+I[3]+1\n",
    "y = y.round(1e-15)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform $\\mathsf{z}_{\\mathbf{i}} = \\frac{\\mathsf{x}_{\\mathbf{i}}}{\\mathsf{z}_{\\mathbf{i}}}$ and report the relative error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relative error tensor(3.3539e-11, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "z = x/y\n",
    "print('Relative error', tn.linalg.norm(z.full()-x.full()/y.full())/tn.linalg.norm(z.full()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform $\\mathsf{u}_{\\mathbf{i}} = \\frac{1}{\\mathsf{z}_{\\mathbf{i}}}$ and report the relative error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relative error tensor(1.1624e-11, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "u = 1/y\n",
    "print('Relative error', tn.linalg.norm(u.full()-1/y.full())/tn.linalg.norm(u.full()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Following are also possible:\n",
    "- scalar (float, int) divided elementwise by a tensor in the TT format.\n",
    "- torch.tensor with 1 element divided elementwise by a tensor in the TT format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = 1.0/y\n",
    "a = tn.tensor(1.0)/y"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}