{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TT Manifold and Riemannian Optimization\n",
    "\n",
    "## Overview\n",
    "\n",
    "This notebook demonstrates **Riemannian optimization** on the tensor train (TT) manifold using `torchtt`. \n",
    "\n",
    "Tensor trains of fixed ranks form a smooth **Riemannian manifold** $\\mathcal{M}_{\\mathbf{r}}$, meaning we can perform optimization directly on this manifold rather than in the ambient space of all tensors. This approach is particularly powerful for problems where we want to constrain our solution to have a specific TT-rank structure.\n",
    "\n",
    "## Riemannian Gradient via Automatic Differentiation\n",
    "\n",
    "The key insight is that `torchtt` can compute the **Riemannian gradient** using **automatic differentiation (AD)**. Given a smooth function $f: \\mathcal{M}_{\\mathbf{r}} \\to \\mathbb{R}$ and a point $\\mathcal{X}$ on the manifold:\n",
    "\n",
    "1. The Euclidean gradient $\\nabla f(\\mathcal{X})$ is computed via PyTorch's AD\n",
    "2. This gradient is then **projected onto the tangent space** $T_{\\mathcal{X}}\\mathcal{M}_{\\mathbf{r}}$ to obtain the Riemannian gradient\n",
    "\n",
    "The function `torchtt.manifold.riemannian_gradient(x, func)` handles both steps automatically, making manifold optimization as simple as standard gradient descent.\n",
    "\n",
    "## Why Riemannian Optimization?\n",
    "\n",
    "Compared to standard gradient descent on TT cores:\n",
    "- **Preserves rank structure**: The iterates stay on the manifold of fixed-rank TTs\n",
    "- **Better conditioning**: Avoids issues with redundant parameterization\n",
    "- **Faster convergence**: Often requires fewer iterations and larger step sizes\n",
    "\n",
    "This notebook shows two examples:\n",
    "1. **Function minimization**: Finding a TT tensor that minimizes a given objective\n",
    "2. **Tensor completion**: Reconstructing a low-rank tensor from sparse observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": [
    "## Example 1: Minimizing a Quadratic Function\n",
    "\n",
    "We create a random target tensor and define a quadratic loss function $f(\\mathcal{X}) = \\frac{1}{2}\\|\\mathcal{X} - \\mathcal{T}\\|^2$ where $\\mathcal{T}$ is the target. The goal is to find a TT tensor of fixed rank that minimizes this objective.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": [
    "### Riemannian Gradient Descent\n",
    "\n",
    "The optimization loop below uses `torchtt.manifold.riemannian_gradient()` to compute the Riemannian gradient at each step. Notice how:\n",
    "- The gradient is computed automatically via AD — no manual derivation needed\n",
    "- A simple retraction (rounding) projects the updated tensor back onto the manifold\n",
    "- Convergence is extremely fast, reaching machine precision in just a few iterations\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Value  198043.7594408848\n",
      "Value  161702.02598821177\n",
      "Value  71389.60114411256\n",
      "Value  23698.57390652948\n",
      "Value  2460.4626114669804\n",
      "Value  0.9337952947342707\n",
      "Value  1.815510743155269e-09\n",
      "Value  9.409020444391612e-25\n",
      "Value  1.479598700080797e-24\n",
      "Value  1.5548277884655298e-24\n",
      "Value  2.969944199121008e-24\n",
      "Value  7.059492707508358e-25\n",
      "Value  1.7110194779310687e-24\n",
      "Value  8.010955421507427e-25\n",
      "Value  1.1671818921822796e-24\n",
      "Value  1.1501295103843825e-24\n",
      "Value  6.074575267741516e-25\n",
      "Value  6.251762467634565e-25\n",
      "Value  8.579854638729817e-25\n",
      "Value  7.778999496013356e-25\n"
     ]
    }
   ],
   "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": [
    "### Comparison: Standard Gradient Descent on TT Cores\n",
    "\n",
    "As a comparison, we perform conventional gradient descent with respect to the TT cores (the ambient space parameterization). Notice the dramatic difference:\n",
    "- **Much smaller step size** required for stability (`α = 0.00001` vs `α = 1.0`)\n",
    "- **Slower convergence**: Even after 1000 iterations, the loss is orders of magnitude higher\n",
    "- The Riemannian approach avoids the ill-conditioning inherent in the TT core parameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(199489.7611, dtype=torch.float64)\n",
      "tensor(199176.1922, dtype=torch.float64)\n",
      "tensor(193181.5858, dtype=torch.float64)\n",
      "tensor(86946.0497, dtype=torch.float64)\n",
      "tensor(18060.0565, dtype=torch.float64)\n",
      "tensor(6145.5984, dtype=torch.float64)\n",
      "tensor(3902.7198, dtype=torch.float64)\n",
      "tensor(2690.0054, dtype=torch.float64)\n",
      "tensor(1782.9748, dtype=torch.float64)\n",
      "tensor(1109.8465, dtype=torch.float64)\n"
     ]
    }
   ],
   "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": [
    "## Example 2: Manifold Tensor Completion\n",
    "\n",
    "A powerful application of Riemannian optimization is **tensor completion** — reconstructing a low-rank tensor from sparse, potentially noisy observations.\n",
    "\n",
    "Given:\n",
    "- A set of observed indices $\\Omega = \\{(i_1^{(k)}, \\ldots, i_d^{(k)})\\}_{k=1}^{M}$\n",
    "- Noisy observations $y_k = \\mathcal{T}_{i_1^{(k)}, \\ldots, i_d^{(k)}} + \\epsilon_k$\n",
    "\n",
    "We minimize the least-squares loss:\n",
    "$$f(\\mathcal{X}) = \\sum_{k=1}^{M} \\left(\\mathcal{X}_{i_1^{(k)}, \\ldots, i_d^{(k)}} - y_k\\right)^2$$\n",
    "\n",
    "The Riemannian gradient is again computed via AD using `torchtt.manifold.riemannian_gradient()`. This is particularly useful because:\n",
    "- The loss involves only sparse evaluations (not the full tensor)\n",
    "- The gradient projection onto the tangent space is handled automatically\n",
    "- We can recover tensors even when observing a tiny fraction of the entries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 3, 3, 2, 1]\n",
      "Riemannian gradient descent\n",
      "\n",
      "Iteration  100 loss value 2.835414e+02 error 7.095354e-02 tensor norm 6.438727e+06\n",
      "Iteration  200 loss value 1.564026e+02 error 5.716689e-02 tensor norm 6.570423e+06\n",
      "Iteration  300 loss value 7.340437e+01 error 4.943336e-02 tensor norm 6.582769e+06\n",
      "Iteration  400 loss value 4.401695e+01 error 4.601466e-02 tensor norm 6.586063e+06\n",
      "Iteration  500 loss value 3.480826e+01 error 4.460205e-02 tensor norm 6.586760e+06\n",
      "Iteration  600 loss value 3.336573e+01 error 4.413909e-02 tensor norm 6.587039e+06\n",
      "Iteration  700 loss value 3.175818e+01 error 4.368099e-02 tensor norm 6.587212e+06\n",
      "Iteration  800 loss value 2.895179e+01 error 4.290056e-02 tensor norm 6.587260e+06\n",
      "Iteration  900 loss value 2.442120e+01 error 4.159506e-02 tensor norm 6.586853e+06\n",
      "Iteration 1000 loss value 1.837452e+01 error 3.981241e-02 tensor norm 6.585507e+06\n",
      "Iteration 1100 loss value 1.143451e+01 error 3.807197e-02 tensor norm 6.583487e+06\n",
      "Iteration 1200 loss value 6.860995e+00 error 3.674023e-02 tensor norm 6.581799e+06\n",
      "Iteration 1300 loss value 3.933710e+00 error 3.574334e-02 tensor norm 6.580529e+06\n",
      "Iteration 1400 loss value 2.447444e+00 error 3.506999e-02 tensor norm 6.579593e+06\n",
      "Iteration 1500 loss value 1.509968e+00 error 3.457828e-02 tensor norm 6.578953e+06\n",
      "Iteration 1600 loss value 8.928927e-01 error 3.427316e-02 tensor norm 6.578572e+06\n",
      "Iteration 1700 loss value 5.786926e-01 error 3.407428e-02 tensor norm 6.578382e+06\n",
      "Iteration 1800 loss value 3.306242e-01 error 3.393152e-02 tensor norm 6.578289e+06\n",
      "Iteration 1900 loss value 1.931639e-01 error 3.383940e-02 tensor norm 6.578253e+06\n",
      "Iteration 2000 loss value 1.021285e-01 error 3.378516e-02 tensor norm 6.578264e+06\n",
      "Iteration 2100 loss value 5.196498e-02 error 3.375590e-02 tensor norm 6.578287e+06\n",
      "Iteration 2200 loss value 2.782952e-02 error 3.373978e-02 tensor norm 6.578301e+06\n",
      "Iteration 2300 loss value 1.667637e-02 error 3.372975e-02 tensor norm 6.578304e+06\n",
      "Iteration 2400 loss value 1.121847e-02 error 3.372262e-02 tensor norm 6.578300e+06\n",
      "Iteration 2500 loss value 8.206944e-03 error 3.371708e-02 tensor norm 6.578290e+06\n",
      "Iteration 2600 loss value 6.315916e-03 error 3.371260e-02 tensor norm 6.578279e+06\n",
      "Iteration 2700 loss value 5.002699e-03 error 3.370891e-02 tensor norm 6.578268e+06\n",
      "Iteration 2800 loss value 4.029615e-03 error 3.370584e-02 tensor norm 6.578257e+06\n",
      "Iteration 2900 loss value 3.280340e-03 error 3.370329e-02 tensor norm 6.578246e+06\n",
      "Iteration 3000 loss value 2.690207e-03 error 3.370117e-02 tensor norm 6.578237e+06\n",
      "Iteration 3100 loss value 2.218880e-03 error 3.369939e-02 tensor norm 6.578229e+06\n",
      "Iteration 3200 loss value 1.838926e-03 error 3.369790e-02 tensor norm 6.578221e+06\n",
      "Iteration 3300 loss value 1.530572e-03 error 3.369665e-02 tensor norm 6.578215e+06\n",
      "Iteration 3400 loss value 1.279022e-03 error 3.369559e-02 tensor norm 6.578209e+06\n",
      "Iteration 3500 loss value 1.072939e-03 error 3.369470e-02 tensor norm 6.578204e+06\n",
      "Iteration 3600 loss value 9.034909e-04 error 3.369395e-02 tensor norm 6.578199e+06\n",
      "Iteration 3700 loss value 7.637210e-04 error 3.369331e-02 tensor norm 6.578195e+06\n",
      "Iteration 3800 loss value 6.480980e-04 error 3.369276e-02 tensor norm 6.578191e+06\n",
      "Iteration 3900 loss value 5.521954e-04 error 3.369229e-02 tensor norm 6.578188e+06\n",
      "Iteration 4000 loss value 4.724499e-04 error 3.369189e-02 tensor norm 6.578185e+06\n",
      "Iteration 4100 loss value 4.059799e-04 error 3.369154e-02 tensor norm 6.578182e+06\n",
      "Iteration 4200 loss value 3.504455e-04 error 3.369124e-02 tensor norm 6.578180e+06\n",
      "Iteration 4300 loss value 3.039398e-04 error 3.369098e-02 tensor norm 6.578178e+06\n",
      "Iteration 4400 loss value 2.649045e-04 error 3.369075e-02 tensor norm 6.578176e+06\n",
      "Iteration 4500 loss value 2.320622e-04 error 3.369054e-02 tensor norm 6.578175e+06\n",
      "Iteration 4600 loss value 2.043636e-04 error 3.369037e-02 tensor norm 6.578173e+06\n",
      "Iteration 4700 loss value 1.809446e-04 error 3.369021e-02 tensor norm 6.578172e+06\n",
      "Iteration 4800 loss value 1.610926e-04 error 3.369006e-02 tensor norm 6.578171e+06\n",
      "Iteration 4900 loss value 1.442184e-04 error 3.368994e-02 tensor norm 6.578170e+06\n",
      "Iteration 5000 loss value 1.298341e-04 error 3.368982e-02 tensor norm 6.578169e+06\n",
      "Iteration 5100 loss value 1.175355e-04 error 3.368972e-02 tensor norm 6.578168e+06\n",
      "Iteration 5200 loss value 1.069866e-04 error 3.368962e-02 tensor norm 6.578168e+06\n",
      "Iteration 5300 loss value 9.790816e-05 error 3.368953e-02 tensor norm 6.578167e+06\n",
      "Iteration 5400 loss value 9.006754e-05 error 3.368945e-02 tensor norm 6.578167e+06\n",
      "Iteration 5500 loss value 8.327080e-05 error 3.368938e-02 tensor norm 6.578166e+06\n",
      "Iteration 5600 loss value 7.735601e-05 error 3.368931e-02 tensor norm 6.578166e+06\n",
      "Iteration 5700 loss value 7.218779e-05 error 3.368925e-02 tensor norm 6.578166e+06\n",
      "Iteration 5800 loss value 6.765283e-05 error 3.368919e-02 tensor norm 6.578165e+06\n",
      "Iteration 5900 loss value 6.365614e-05 error 3.368913e-02 tensor norm 6.578165e+06\n",
      "Iteration 6000 loss value 6.011804e-05 error 3.368908e-02 tensor norm 6.578165e+06\n",
      "Iteration 6100 loss value 5.697154e-05 error 3.368902e-02 tensor norm 6.578165e+06\n",
      "Iteration 6200 loss value 5.416030e-05 error 3.368898e-02 tensor norm 6.578165e+06\n",
      "Iteration 6300 loss value 5.163683e-05 error 3.368893e-02 tensor norm 6.578165e+06\n",
      "Iteration 6400 loss value 4.936106e-05 error 3.368889e-02 tensor norm 6.578165e+06\n",
      "Iteration 6500 loss value 4.729913e-05 error 3.368884e-02 tensor norm 6.578165e+06\n",
      "Iteration 6600 loss value 4.542238e-05 error 3.368880e-02 tensor norm 6.578165e+06\n",
      "Iteration 6700 loss value 4.370652e-05 error 3.368877e-02 tensor norm 6.578165e+06\n",
      "Iteration 6800 loss value 4.213089e-05 error 3.368873e-02 tensor norm 6.578165e+06\n",
      "Iteration 6900 loss value 4.067795e-05 error 3.368869e-02 tensor norm 6.578165e+06\n",
      "Iteration 7000 loss value 3.933271e-05 error 3.368866e-02 tensor norm 6.578165e+06\n",
      "Iteration 7100 loss value 3.808240e-05 error 3.368862e-02 tensor norm 6.578165e+06\n",
      "Iteration 7200 loss value 3.691605e-05 error 3.368859e-02 tensor norm 6.578166e+06\n",
      "Iteration 7300 loss value 3.582426e-05 error 3.368856e-02 tensor norm 6.578166e+06\n",
      "Iteration 7400 loss value 3.479897e-05 error 3.368853e-02 tensor norm 6.578166e+06\n",
      "Iteration 7500 loss value 3.383319e-05 error 3.368850e-02 tensor norm 6.578166e+06\n",
      "Iteration 7600 loss value 3.292090e-05 error 3.368847e-02 tensor norm 6.578166e+06\n",
      "Iteration 7700 loss value 3.205688e-05 error 3.368844e-02 tensor norm 6.578166e+06\n",
      "Iteration 7800 loss value 3.123657e-05 error 3.368842e-02 tensor norm 6.578167e+06\n",
      "Iteration 7900 loss value 3.045602e-05 error 3.368839e-02 tensor norm 6.578167e+06\n",
      "Iteration 8000 loss value 2.971177e-05 error 3.368837e-02 tensor norm 6.578167e+06\n",
      "Iteration 8100 loss value 2.900076e-05 error 3.368834e-02 tensor norm 6.578167e+06\n",
      "Iteration 8200 loss value 2.832031e-05 error 3.368832e-02 tensor norm 6.578167e+06\n",
      "Iteration 8300 loss value 2.766807e-05 error 3.368829e-02 tensor norm 6.578168e+06\n",
      "Iteration 8400 loss value 2.704194e-05 error 3.368827e-02 tensor norm 6.578168e+06\n",
      "Iteration 8500 loss value 2.644004e-05 error 3.368825e-02 tensor norm 6.578168e+06\n",
      "Iteration 8600 loss value 2.586071e-05 error 3.368822e-02 tensor norm 6.578168e+06\n",
      "Iteration 8700 loss value 2.530246e-05 error 3.368820e-02 tensor norm 6.578169e+06\n",
      "Iteration 8800 loss value 2.476395e-05 error 3.368818e-02 tensor norm 6.578169e+06\n",
      "Iteration 8900 loss value 2.424398e-05 error 3.368816e-02 tensor norm 6.578169e+06\n",
      "Iteration 9000 loss value 2.374144e-05 error 3.368814e-02 tensor norm 6.578169e+06\n",
      "Iteration 9100 loss value 2.325534e-05 error 3.368812e-02 tensor norm 6.578169e+06\n",
      "Iteration 9200 loss value 2.278479e-05 error 3.368810e-02 tensor norm 6.578170e+06\n",
      "Iteration 9300 loss value 2.232895e-05 error 3.368808e-02 tensor norm 6.578170e+06\n",
      "Iteration 9400 loss value 2.188707e-05 error 3.368806e-02 tensor norm 6.578170e+06\n",
      "Iteration 9500 loss value 2.145846e-05 error 3.368805e-02 tensor norm 6.578170e+06\n",
      "Iteration 9600 loss value 2.104246e-05 error 3.368803e-02 tensor norm 6.578170e+06\n",
      "Iteration 9700 loss value 2.063850e-05 error 3.368801e-02 tensor norm 6.578171e+06\n",
      "Iteration 9800 loss value 2.024601e-05 error 3.368799e-02 tensor norm 6.578171e+06\n",
      "Iteration 9900 loss value 1.986450e-05 error 3.368798e-02 tensor norm 6.578171e+06\n",
      "Iteration 10000 loss value 1.949349e-05 error 3.368796e-02 tensor norm 6.578171e+06\n",
      "\n",
      "Time elapsed 0:01:34.523436\n",
      "Number of observations 15000, tensor shape [25, 25, 25, 25], percentage of entries observed 3.8400\n",
      "Number of unknowns 1000, number of observations 15000, DoF/observations 0.066667\n",
      "Rank after rounding TT with sizes and ranks:\n",
      "N = [25, 25, 25, 25]\n",
      "R = [1, 4, 4, 4, 1]\n",
      "\n",
      "Device: cpu, dtype: torch.float64\n",
      "#entries 1000 compression 0.00256\n",
      "\n"
     ]
    }
   ],
   "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": {
  "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
}