tt-ksl.ipynb

{
 "metadata": {
  "name": "tt-ksl"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Dynamical low-rank in TT \n",
      "\n",
      "## Input data\n",
      "\n",
      "A function that produces $A(t)$ at time $t$,  \n",
      "Initial value in the TT-format\n",
      "\n",
      "## Test tensor\n",
      "\n",
      "I used the following example:\n",
      "\n",
      "\\begin{equation}\n",
      "    A(t) = A_0 (1 - t)^2 + A_1 (1-t) + \\varepsilon A_2 (1 - \\cos t) \n",
      "\\end{equation}\n",
      "\n",
      "where $A_0(t) $ is a $30 \\times 30 \\times 30 \\times 30 $ random tensor with TT-ranks $(2,2,2)$,  \n",
      "and $A_1$  has TT-ranks $(3,3,3)$ and $A_2$ has large ranks.  \n",
      "The manifold is a TT-manifold with ranks $(5,5,5)$.\n",
      "\n",
      "The example may be a little specific (the basices are not changing),  \n",
      "so there is a plan to implement time-varying basises. \n",
      "\n",
      "## The sketch of the algorithm\n",
      "\n",
      "The idea is simple: apply the KSL scheme recursively and then gather the result;  \n",
      "The second-order version is obtained by reversing the steps.  \n",
      "\n",
      "In the TT-format, the final procedure looks as follows:  \n",
      "\n",
      "**Init**: $X = X_1(i_1) X_2(i_2) X_3(i_3)$  \n",
      "\n",
      "1. Orthogonalize $X_2$, $X_3$ from the right:\n",
      "$X = (X_1(i_1) S_1) \\widehat{X}_2(i_2) \\widehat{X}_3(i_3)$  \n",
      "\n",
      "2. Make a step in $K_1(i_1)$, compute its QR from the left:  \n",
      "$K_1(i_1) = \\widehat{X}(i_1) \\widehat{S}_1$,  \n",
      "\n",
      "3. Make a step in $\\widehat{S}_1$ **backwards**:  \n",
      "\n",
      "4. Put $\\widehat{S}_1$ in the second core, $K_2(i_2) = \\widehat{S}_1 X_2(i_2)$\n",
      "\n",
      "5. Then the cycle: **step in the core, QR, step in S backwards, put $S$ into the next core**\n",
      "\n",
      "Everything is organized into one sweep.  For the second order, we will need two sweeps.  \n",
      "The second order is confirmed numerically on the example below.\n"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "\"\"\" The test tensor \"\"\"\n",
      "from math import cos\n",
      "def Atest(t,a0,a1,a2,eps=1e-4): \n",
      "    return (a0 * (1 - t) ** 2 + a1 * (1 - t) + eps * (1 - cos(t)) * a2).full()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 17
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Note\n",
      "Below are technical \"sweeping\" subroutines for the KSL in the TT-format;  \n",
      "They are organized in the way I think it would be simple to change them  \n",
      "into the HT-format."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numpy as np\n",
      "from numpy.linalg import qr\n",
      "import sys\n",
      "sys.path.append(\"/Users/ivan/work/python/ttpy/\") \n",
      "import tt\n",
      "class node:\n",
      "    __slots__ = ['x','A','B']\n",
      "    \n",
      "    def __init__(self,x,A=None,B=None):\n",
      "        self.x = x\n",
      "        self.A = A\n",
      "        self.B = B \n",
      "\n",
      "def left_qr_step(node_left,node_right):\n",
      "    x0 = node_right.x\n",
      "    r1,n,r2 = x0.shape\n",
      "    x0 = np.reshape(x0,[r1,n*r2],order='F')\n",
      "    x0 = x0.T\n",
      "    q,r = qr(x0)\n",
      "    rnew = q.shape[1]\n",
      "    q = np.reshape(q,[n*r2,rnew],order='F')\n",
      "    q = q.T\n",
      "    q = np.reshape(q,[rnew,n,r2],order='F')\n",
      "    node_right.x = q\n",
      "    x1 = node_left.x\n",
      "    rr1,nn,rr2 = x1.shape\n",
      "    x1 = np.reshape(x1,[rr1*nn,rr2],order='F')\n",
      "    x1 = np.dot(x1,r.T)\n",
      "    x1 = np.reshape(x1,[rr1,nn,rnew],order='F')\n",
      "    node_left.x = x1\n",
      "    B = node_right.A\n",
      "    B = np.reshape(B,[B.size/(n*r2),n*r2],order='F')\n",
      "    node_left.A = np.dot(B,np.reshape(q.conj(),[rnew,n*r2],order='F').T)\n",
      "    \"\"\" we have A(i1,i2,...,id) X(id,rd) -> A(i1,i2,...,id-1,rd) X(id-1,rd-1) \"\"\"\n",
      "\n",
      "def right_conv(node_left,node_right):\n",
      "    \"\"\" This is a technical subroute that prepares the \"A\" matrices for the\n",
      "        right-to-left sweep \"\"\"\n",
      "    x0 = node_left.x\n",
      "    r1,n,r2 = x0.shape\n",
      "    x0 = np.reshape(x0,[r1*n,r2],order='F')\n",
      "    B = node_left.A\n",
      "    B = np.reshape(B,[r1*n,B.size/(r1*n)],order='F')\n",
      "    node_right.A = np.dot(x0.conj().T,B)\n",
      "    \n",
      "def right_ksl_step(node_left,node_right):\n",
      "    x0 = node_left.x\n",
      "    r1,n,r2 = x0.shape\n",
      "    A = node_left.A \n",
      "    x0 = np.reshape(x0,[r1*n,r2],order='F')\n",
      "    x0 = x0 + np.reshape(A,x0.shape,order='F')\n",
      "    if node_right is None:\n",
      "        node_left.x = np.reshape(x0,[r1,n,r2],order='F')\n",
      "        return\n",
      "    q,S = qr(x0)\n",
      "    node_left.x = np.reshape(q,[r1,n,r2],order='F')\n",
      "    A = np.reshape(A,[r1*n,r2],order='F')\n",
      "    dS = np.dot(q.conj().T,A)\n",
      "    S = S - dS\n",
      "    x1 = node_right.x\n",
      "    rr1,n,rr2 = x1.shape\n",
      "    x1 = np.reshape(x1,[rr1,n*rr2],order='F')\n",
      "    x1 = np.dot(S,x1)\n",
      "    node_right.x = np.reshape(x1,[rr1,n,rr2],order='F')\n",
      "    if node_left.B is None:\n",
      "        B = np.reshape(q,[r1*n,r2],order='F')\n",
      "    else:\n",
      "        B = node_left.B\n",
      "        B = np.dot(B,np.reshape(q,[r1,n*r2],order='F'))\n",
      "        B = np.reshape(B,[B.size/r2,r2],order='F')\n",
      "    node_right.B = B\n",
      "    M = B.size/r2\n",
      "    B1 = np.reshape(B,[M,r2],order='F')\n",
      "    A1 = node_right.A\n",
      "    A1 = np.reshape(A1,[M,A1.size/M],order='F')\n",
      "    A1 = np.dot(B1.conj().T,A1)\n",
      "    node_right.A = A1\n",
      "\n",
      "def left_ksl_step(node_left,node_right):\n",
      "    \"\"\" The reverse step \n",
      "        What we expect is that:\n",
      "        We do all the steps before using A0,A1 up to the last;\n",
      "        then we start over going backwards \n",
      "        The same ideology can hold:\n",
      "        We do the actual step in the last core, do step in S, \n",
      "        and recompute stuff for the last core \"\"\"\n",
      "    x0 = node_right.x\n",
      "    r1,n,r2 = x0.shape\n",
      "    A = node_right.A\n",
      "    x0 = np.reshape(x0,[r1,n*r2],order='F')\n",
      "    x0 = x0 + np.reshape(A,x0.shape,order='F')\n",
      "    if node_left is None:\n",
      "        node_right.x = np.reshape(x0,[r1,n,r2],order='F')\n",
      "        return\n",
      "    x0 = x0.T\n",
      "    q,S = qr(x0)\n",
      "    q = q.T \n",
      "    S = S.T\n",
      "    node_right.x = np.reshape(q,[r1,n,r2],order='F')\n",
      "    #A is r1 x n x r2, q is (n x r2) x r1, q.T is r1 x (n x r2)\n",
      "    A = np.reshape(A,[r1,n*r2],order='F')\n",
      "    dS = np.dot(A,q.conj().T)\n",
      "    S = S - dS\n",
      "    x1 = node_left.x\n",
      "    rr1,n,rr2 = x1.shape\n",
      "    x1 = np.reshape(x1,[rr1*n,rr2],order='F')\n",
      "    x1 = np.dot(x1,S)\n",
      "    node_left.x = np.reshape(x1,[rr1,n,rr2],order='F')\n",
      "    if node_right.B is None:\n",
      "        B = np.reshape(q,[r1,n*r2],order='F')\n",
      "    else:\n",
      "        B = node_right.B\n",
      "        B = np.dot(np.reshape(q,[r1*n,r2],order='F'),B)\n",
      "        B = np.reshape(B,[r1,B.size/r1],order='F')\n",
      "    node_left.B = B\n",
      "    M = B.size/r1\n",
      "    B1 = np.reshape(B,[r1,M],order='F')\n",
      "    A1 = node_left.A\n",
      "    A1 = np.reshape(A1,[A1.size/M,M],order='F')\n",
      "    A1 = np.dot(A1,B.conj().T)\n",
      "    node_left.A = A1"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 18
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## The sweeping schemes\n",
      "Two schemes, ``ksl_tt`` and ``ksl_tt_symm`` which use the technical subroutines described above  \n",
      "and provide the time-integration schemes"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "\"\"\" The time-stepping scheme \"\"\"\n",
      "def ksl_tt(t,tau,X,Afun):\n",
      "    a0 = Afun(t)\n",
      "    a1 = Afun(t + tau)\n",
      "    d = X.d #The dimension\n",
      "    X1 = tt.tensor.to_list(X)\n",
      "    nodes = [node(x=X1[i]) for i in xrange(d)]\n",
      "    nodes[d-1].A = (a1 - a0) \n",
      "    \n",
      "    for i in xrange(d-2,-1,-1):\n",
      "        left_qr_step(nodes[i],nodes[i+1])\n",
      "    \"\"\" Now we go for an ordinary sweep from left to right \"\"\"\n",
      "    for i in xrange(d-1):\n",
      "        right_ksl_step(nodes[i],nodes[i+1])\n",
      "    right_ksl_step(nodes[d-1],None) #To fix the last core\n",
      "    X1 = [nodes[i].x for i in xrange(d)]\n",
      "    return tt.tensor.from_list(X1)\n",
      "\n",
      "\"\"\" Symmetrized version \"\"\"\n",
      "def ksl_tt_symm(t,tau,X,Afun):\n",
      "    a0 = Afun(t)\n",
      "    a1 = Afun(t + tau/2)\n",
      "    a2 = Afun(t + tau)\n",
      "    \n",
      "    d = X.d #The dimension\n",
      "    X1 = tt.tensor.to_list(X)\n",
      "    nodes = [node(x=X1[i]) for i in xrange(d)]\n",
      "    nodes[d-1].A = (a1 - a0)\n",
      "    \n",
      "    for i in xrange(d-2,-1,-1):\n",
      "        left_qr_step(nodes[i],nodes[i+1])\n",
      "    nodes_tmp = deepcopy(nodes)\n",
      "    \"\"\" Now we go for an ordinary sweep from left to right \"\"\"\n",
      "    for i in xrange(d-1):\n",
      "        right_ksl_step(nodes[i],nodes[i+1])\n",
      "    right_ksl_step(nodes[d-1],None) #To fix the last core\n",
      "    X1 = [nodes[i].x for i in xrange(d)]\n",
      "    X3 = tt.tensor.from_list(X1)\n",
      "    \n",
      "    nodes[0].A = (a2 - a1) #To recompute them all\n",
      "    nodes[d-1].B = None #The first-core flag\n",
      "    for i in xrange(d-1):\n",
      "        right_conv(nodes[i],nodes[i+1]) \n",
      "    for i in xrange(d-2,-1,-1):\n",
      "        left_ksl_step(nodes[i],nodes[i+1])\n",
      "    left_ksl_step(None,nodes[0])\n",
      "    \n",
      "    X1 = [nodes[i].x for i in xrange(d)]\n",
      "    X3 = tt.tensor.from_list(X1)\n",
      "\n",
      "    return tt.tensor.from_list(X1)    "
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 19
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Initilization stuff\n",
      "Initialize $A_0, A_1, A_2$"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import sys\n",
      "sys.path.append(\"/Users/ivan/work/python/ttpy/\") \n",
      "from copy import deepcopy\n",
      "import numpy as np\n",
      "import tt\n",
      "d = 3\n",
      "a0 = tt.rand(30,d,2)\n",
      "a1 = tt.rand(30,d,3)\n",
      "a2 = tt.rand(30,d,25)\n",
      "Afun = lambda t: Atest(t,a0,a1,a2,1e-2)\n",
      "tau = 1e-2\n",
      "X = tt.tensor(Afun(0),1e-8) #Initial value"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 20
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "t = 0\n",
      "tf = 0.5\n",
      "ns = [4,5,6]\n",
      "Xs1 = {}\n",
      "Xs2 = {}\n",
      "i = 0\n",
      "for n in ns:\n",
      "    t = 0\n",
      "    tau_loc = tf*1.0/(2**n)\n",
      "    X1 = X.copy()\n",
      "    X2 = X.copy()\n",
      "    j = 0\n",
      "    for j in xrange(2**n):\n",
      "        X1 = ksl_tt(t,tau_loc,X1,Afun)\n",
      "        X2 = ksl_tt_symm(t,tau_loc,X2,Afun)\n",
      "        t += tau_loc\n",
      "    Xs1[i] = X1\n",
      "    Xs2[i] = X2\n",
      "    i = i + 1"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 21
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Computation of the Runge factors\n",
      "The actual computed Runge factors:  \n",
      "2 is for the first-order scheme,  \n",
      "4 is for the second-order scheme."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "Ax = Afun(tf).flatten('F')\n",
      "def compute_Runge(Xs):\n",
      "    X1 = Xs[0].full().flatten('F')\n",
      "    X2 = Xs[1].full().flatten('F')\n",
      "    X3 = Xs[2].full().flatten('F')\n",
      "    return norm(X1 - X2)/norm(X2-X3)\n",
      "print 'Runge factor for KSL:', compute_Runge(Xs1)\n",
      "print 'Runge factor for KSL(symm):', compute_Runge(Xs2)\n"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Runge factor for KSL: 2.00516439996\n",
        "Runge factor for KSL(symm): 3.9920514875\n"
       ]
      }
     ],
     "prompt_number": 22
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}