{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cross approximation in the TT format\n",
    "\n",
    "Using the `torchtt.TT` constructor, a TT decomposition of a given tensor can be obtained. However, in the cases where the entries of the tensor are computed using a given function, building full tensors becomes unfeasible. It is possible to construct a TT decomposition using only a part of the entries of the full tensor. This is called the cross approximation method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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": [
    "### Cross interpolation of a tensor in TT format\n",
    "\n",
    "We want to approximate the tensor $\\mathsf{x}_{i_1...i_d}=\\frac{1}{2+i_1+\\cdots+i_d}$. Since the passed indices are integers of type `torch.int64`, casting is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func1 = lambda I: 1/(2+tn.sum(I+1,1).to(dtype=tn.float64))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call the `torchtt.interpolate.dmrg_cross()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = [20]*4\n",
    "x = tntt.interpolate.dmrg_cross(func1, N, eps = 1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the full tensor and compare to the reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Is = tntt.meshgrid([tn.arange(0,n,dtype=tn.float64) for n in N])\n",
    "x_ref = 1/(2+Is[0].full()+Is[1].full()+Is[2].full()+Is[3].full()+4)\n",
    "print('Relative error ',tn.linalg.norm(x.full()-x_ref)/tn.linalg.norm(x_ref))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We consider the case $d=10$, $n_i=32$. the full tensor would contain $32^{10}$ entries.\n",
    "The total number of functions calls is in this case 25000000 compared to $32^{10}$ of the total number of entries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = [32]*10\n",
    "x = tntt.interpolate.dmrg_cross(func1, N, eps = 1e-10, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The adaptive cross method used only a fraction of function calls from the original tensor.\n",
    "Check some entries (full tensor cannot be computed this time) and show the rank and the storage requirements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(x[1,2,3,4,5,6,7,8,9,11], ' reference ', func1(tn.tensor([[1,2,3,4,5,6,7,8,9,11]])))\n",
    "print(x[12,23,17,25,30,0,7,8,9,11], ' reference ', func1(tn.tensor([[12,23,17,25,30,0,7,8,9,11]])))\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Element wise application of an univariate function on a TT tensor.\n",
    "\n",
    "Let $f:\\mathbb{R}\\rightarrow\\mathbb{R}$ be a function and $\\mathsf{x}\\in\\mathbb{R}^{N_1\\times\\cdots\\times N_d}$ be a tensor with a known TT approximation. The goal is to determine the TT approximation of $\\mathsf{y}_{i_1...i_d}=f(\\mathsf{x}_{i_1...i_d})$ within a prescribed relative accuracy $\\epsilon$ (passed as argument).\n",
    "In this case the function is `torchtt.interpoalte.function_interpolate()` and takes as arguments a function handle, the tensor $\\mathsf{x}$, the accuracy epsilon, a initial tensor (starting point), number of sweeps (nswp) and the size of the rank enrichment (kick).\n",
    "Further arguments are the dtype of the result and the verbose flag.\n",
    "\n",
    "The function handle as argument gets as arguments torch vectors and has to return torch vectors of the same size.\n",
    "\n",
    "The following example computes the elemntwise natural logarithm of a tensor. The relative error of the result is also reported.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tntt.TT(x_ref)\n",
    "func = lambda t: tn.log(t)\n",
    "y = tntt.interpolate.function_interpolate(func, x, 1e-9)\n",
    "print('Relative error ',tn.linalg.norm(y.full()-func(x_ref))/tn.linalg.norm(func(x_ref)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Element wise application of an multivariate function on a TT tensor.\n",
    "\n",
    "Let $f:\\mathbb{R}\\rightarrow\\mathbb{R}$ be a function and $\\mathsf{x}^{(1)},...,\\mathsf{x}^{(d)}\\in\\mathbb{R}^{N_1\\times\\cdots\\times N_d}$ be tensors with a known TT approximation. The goal is to determine the TT approximation of $\\mathsf{y}_{i_1...i_d}=f(\\mathsf{x}_{i_1...i_d}^{(1)},...,\\mathsf{x}^{(d)}_{i_1...i_d})$ within a prescribed relative accuracy $\\epsilon$ (passed as argument). The function is the same as in the previous case `tochtt.interpoalte.function_interpolate()`, but the second argument in this case is a list of torchtt.TT tensors. The function handle takes as argument a $M\\times d$ torch.tensor and every of the $M$ lines corresponds to an evaluation of the function $f$ at a certain tensor entry. The function handle returns a torch tensor of length $M$.\n",
    "\n",
    "The following example computes the same tensor as in the previous case, but with the `tochtt.interpoalte.function_interpolate()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = tntt.interpolate.function_interpolate(func1, Is)\n",
    "print('Relative error ',tn.linalg.norm(z.full()-x_ref)/tn.linalg.norm(x_ref))"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "df6fc3a9b7a9c6f4b0308ab6eb361a4cabbf6b5db181383d07014ff4304e5cb3"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
