{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as tn\n",
    "import datetime\n",
    "import numpy as np\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": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = [10,11,12,13,14]\n",
    "Rt = [1,3,4,5,6,1]\n",
    "Rx = [1,6,6,6,6,1]\n",
    "target = tntt.randn(N,Rt).round(0)\n",
    "func = lambda x: 0.5*(x-target).norm(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = tntt.randn(N,Rx)\n",
    "x =x0.clone()\n",
    "for i in range(20):\n",
    "    # compute riemannian gradient using AD    \n",
    "    gr = tntt.manifold.riemannian_gradient(x,func)\n",
    "    \n",
    "    #stepsize length\n",
    "    alpha = 1.0\n",
    "    \n",
    "    # update step\n",
    "    x = (x-alpha*gr).round(0,Rx)    \n",
    "    print('Value ' , func(x).numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a comparison, conventional gradient descent with respect to the TT cores is performed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = x0.detach().clone()\n",
    "\n",
    "for i in range(1000):\n",
    "    tntt.grad.watch(y)\n",
    "    fval = func(y)\n",
    "    deriv = tntt.grad.grad(fval,y)    \n",
    "    alpha = 0.00001 # for stability\n",
    "    y = tntt.TT([y.cores[i].detach()-alpha*deriv[i] for i in range(len(deriv))])\n",
    "    if (i+1)%100 == 0: print(func(y))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manifold tensor completion\n",
    "\n",
    "One other task where the manifold learning can be applied is tensor completion.\n",
    "The goal for this problem is to reconstruct a tensor in the TT format given only a few entries (possible noisy)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 25\n",
    "target = tntt.randn([N]*4,[1,2,3,3,1])\n",
    "Xs = tntt.meshgrid([tn.linspace(0,1,N, dtype = tn.float64)]*4)\n",
    "target = Xs[0]+1+Xs[1]+Xs[2]+Xs[3]+Xs[0]*Xs[1]+Xs[1]*Xs[2]+tntt.TT(tn.sin(Xs[0].full()))\n",
    "target = target.round(1e-10)\n",
    "print(target.R)\n",
    "\n",
    "M = 15000 # number of observations \n",
    "indices = tn.randint(0,N,(M,4))\n",
    "\n",
    "# observations are considered to be noisy\n",
    "sigma_noise = 0.00001\n",
    "obs = tn.normal(target.apply_mask(indices), sigma_noise)\n",
    "\n",
    "# define the loss function\n",
    "loss = lambda x: (x.apply_mask(indices)-obs).norm()**2\n",
    "\n",
    "#%% Manifold learning\n",
    "print('Riemannian gradient descent\\n')\n",
    "# starting point\n",
    "x = tntt.randn([N]*4,[1,4,4,4,1])\n",
    "\n",
    "tme = datetime.datetime.now()\n",
    "# iterations\n",
    "for i in range(10000):\n",
    "    # manifold gradient \n",
    "    gr = tntt.manifold.riemannian_gradient(x,loss)\n",
    "\n",
    "    step_size = 1.0\n",
    "    R = x.R\n",
    "    # step update\n",
    "    x = (x - step_size * gr).round(0,R)\n",
    "\n",
    "    # compute loss value\n",
    "    if (i+1)%100 == 0:\n",
    "        loss_value = loss(x)\n",
    "        print('Iteration %4d loss value %e error %e tensor norm %e'%(i+1,loss_value.numpy(),(x-target).norm()/target.norm(), x.norm()**2))\n",
    "\n",
    "tme = datetime.datetime.now() - tme\n",
    "print('')\n",
    "print('Time elapsed',tme)\n",
    "print('Number of observations %d, tensor shape %s, percentage of entries observed %6.4f'%(M,str(x.N),100*M/np.prod(x.N)))\n",
    "print('Number of unknowns %d, number of observations %d, DoF/observations %.6f'%(tntt.numel(x),M,tntt.numel(x)/M))\n",
    "\n",
    "print('Rank after rounding',x.round(1e-6))"
   ]
  }
 ],
 "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.6.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
