Implementation of the KSL integrator in Python

dyn_lr.py

import sys
import os
from math import exp,sqrt
import numpy as np
from numpy.linalg import svd,norm,qr
from numpy.random import rand
from scipy.linalg import expm
import time
import timestep as ts
import itertools

def gen_A(n=100,m=10):
    M0 = np.eye(m) + 0.5 * rand(m,m)
    M2 = rand(n,n)
    M1 = np.zeros((n,n))
    M1[0:m,0:m] = M0
    return M1,M2

def gen_T(n):
    T1 = rand(n,n)
    T1 = T1 - T1.T
    return T1

def gen_Q(t,T1):
    return expm(t*T1)

def first_A(t,mat):
    T1 = mat["T1"]
    T2 = mat["T2"]
    A1 = mat["A1"]
    A2 = mat["A2"]
    Q1 = gen_Q(t,T1)
    Q2 = gen_Q(t,T2)
    W = np.dot(Q1,A1 + exp(t) * A2)
    W = np.dot(W,Q2.T)
    return W

def first_grad(t,mat):
    T1 = mat["T1"]
    T2 = mat["T2"]
    A1 = mat["A1"]
    A2 = mat["A2"]
    #Q1(t) * (A1 + exp(t) * A2) * Q2(t).T
    #T1 * Q1(t) * (A1 + t * A2) * Q2(t).T
    Q1 = gen_Q(t,T1)
    Q2 = gen_Q(t,T2)
    W1 = np.dot(np.dot(T1,Q1), A1 + exp(t) * A2)
    W1 = np.dot(W1,Q2.T)
    W2 = np.dot(Q1, exp(t) * A2)
    W2 = np.dot(W2, Q2.T)
    W3 = np.dot(Q1, (A1 + exp(t) * A2))
    W3 = np.dot(W3,np.dot(Q2.T,T2.T))
    return W1 + W2 + W3
#Implement a simple implicit midpoint rule for 
#the dynamical low-rank

def lr_rhs(t,mat,X):
    U = X.U
    S = X.S
    V = X.V
    Agr = first_grad(t,mat)
    Ur = np.dot(Agr,V) 
    Vr = np.dot(Agr.T,U)
    Qu, Ru = qr(U)
    Qv, Rv = qr(V)
    Ur = np.linalg.solve(S.T,Ur.T).T
    Vr = np.linalg.solve(S,Vr.T).T
    Ur = Ur - np.dot(Qu, np.dot(Qu.T,Ur))
    Vr = Vr - np.dot(Qv, np.dot(Qv.T,Vr))
    Sr = np.dot(np.dot(U.T, Agr),V) #This is not completely correct: U' * A' * V -> 
    #res = {"U" : Ur, "V" : Vr, "S" : Sr}
    res = ts.dyn_lr(Ur,Sr,Vr)
    return res


n = 100
m = 10

#Generate the basic data
M10,M11 = gen_A(n,m)
M20,M21 = gen_A(n,m)
T1 = gen_T(n)
T2 = gen_T(n)

eps_all = [1e-3, 1e-6]
tau0 = 1e-1 
tau1 = 1e-3
#tau_all = np.linspace(np.log(tau0),np.log(tau1),20)
#tau_all = np.exp(tau_all)
""" We have to provide different tau's such that tau = 1/m, so it will exactly pass through
    some point; thus, we have to take tau of such form. Then we need some m = m(k) """
tau_all = np.linspace(np.log(tau0),np.log(tau1),8)
tau_all = np.exp(tau_all)
tau_all = 1.0/(np.round(1.0/tau_all))
#tau_all = np.concatenate((tau_all,[5e-4]))
#tau_all = [1e-3]
#import ipdb; ipdb.set_trace()
#tau_all = [tau1 + (tau0 - tau1)/(2**(k+1)-1) for k in xrange(20)]
#tau_all = [1e-1, 1e-2, 1e-3]
#eps_all = [1e-6]
#tau_all = [1e-3]
r_all = [10,20]
t_st = {}
t_kls = {}
t_kls1 = {}
t_ksl = {}
t_ksl2 = {}
er_st = {}
er_kls = {}
er_kls1 = {}
er_ksl = {}
er_ksl2 = {}
er_dif = {}
er_best = {}
sol_st_fin = {}
sol_kls_fin = {}
sol_kls1_fin = {}
sol_ksl_fin = {}
sol_best_fin = {}
#For all eps_all, tau_all, r_all we get many data: 
#Error of the approximation (for all three schemes) 
#Time of the approximation (for all three schemes)
#And the best SVD-based approximation (to see where is the bottom)
#We will store the results in multidimensional dictionaries,
#when I will find an internet connection
for eps,tau,r in itertools.product(eps_all,tau_all,r_all):
    A1 = M10 + eps * M11
    A2 = M20 + eps * M21
    mat = {"A1" : A1, "A2" : A2, "T1" : T1, "T2" : T2}
    U,S,V = svd(first_A(0,mat),full_matrices=False)
    U = U[:,0:r]
    S = S[0:r]
    S = np.diag(S)
    V = V[0:r,:]
    V = V.T
    X = ts.dyn_lr(U,S,V)
    rhs_fun = lambda t,X : lr_rhs(t,mat,X)
    fun = lambda t : first_A(t,mat)
    X1 = X.copy()
    X2 = X.copy()
    X3 = X.copy()
    X5 = X.copy()
    X6 = X.copy()
    tf = 1
    t1 = 0
    t2 = 0
    t3 = 0
    t4 = 0
    t5 = 0
    er1 = {}
    er2 = {}
    er3 = {}
    e_dif = {}
    er4 = {}
    er5 = {}
    er6 = {}
    print 'Doing eps = %3.1e tau = %3.1e rank = %d' % (eps,tau,r)
    t = 0
    while t < tf:
        dt = -time.time()
        X1 = ts.imid_lr(t,tau,X1,rhs_fun)
        dt += time.time()
        t1 += dt
        dt = -time.time()
        X2 = ts.kls_lr(t,tau,X2,fun)
        dt += time.time()
        t2 += dt
        dt = -time.time()
        X3 = ts.kls_lr1(t,tau,X3,fun)
        dt += time.time()
        t3 += dt
        dt = -time.time()
        X5 = ts.ksl(t,tau,X5,fun)
        dt += time.time()
        t4 += dt
        dt = -time.time()
        X6 = ts.ksl2(t,tau,X6,fun)
        dt += time.time()
        t5 += dt
        
        t += tau
        Ax = first_A(t,mat)
        Ap1 = np.dot(X1.U, np.dot(X1.S,X1.V.T))
        Ap2 = np.dot(X2.U, np.dot(X2.S,X2.V.T))
        Ap3 = np.dot(X3.U, np.dot(X3.S,X3.V.T))
        Ap5 = np.dot(X5.U, np.dot(X5.S,X5.V.T))        
        Ap6 = np.dot(X6.U, np.dot(X6.S,X6.V.T))
        u0,s0,v0 = svd(Ax,full_matrices = False); u0 = u0[:,0:r]; s0 = s0[0:r]; v0 = v0[0:r,:]; Ap4 = np.dot(u0,np.dot(np.diag(s0),v0))
        
        
        nrm = norm(Ax,'fro')
        er1[t] = norm(Ax - Ap1,'fro')
        er2[t] = norm(Ax - Ap2,'fro')
        er3[t] = norm(Ax - Ap3,'fro')
        e_dif[t] = norm(Ap1 - Ap2,'fro')
        er4[t] = norm(Ax - Ap4,'fro')
        er5[t] = norm(Ax - Ap5,'fro')
        er6[t] = norm(Ax - Ap6,'fro')
    v = (eps,tau,r)
    sol_st_fin[v] = Ap1
    sol_kls_fin[v] = Ap2
    sol_kls1_fin[v] = Ap3
    sol_best_fin[v] = Ap4
    sol_ksl_fin[v] = Ap5
    er_st[v] = er1
    er_kls[v] = er2
    er_kls1[v] = er3
    er_best[v] = er4
    er_dif[v] = e_dif
    er_ksl[v] = er5
    er_ksl2[v] = er6
    t_st[v] = t1
    t_kls[v] = t2
    t_kls1[v] = t3
    t_ksl[v] = t4
    t_ksl2[v] = t5
np.savez("all_results",er_st=er_st,er_kls=er_kls,er_kls1=er_kls1,er_best=er_best,er_dif=er_dif,t_st=t_st,t_kls=t_kls,t_kls1=t_kls1
        ,er_ksl = er_ksl,er_ksl2=er_ksl2,t_ksl = t_ksl, t_ksl2=t_ksl2, sol_kls_fin = sol_kls_fin,sol_st_fin = sol_st_fin,sol_ksl_fin=sol_ksl_fin,
        sol_kls1_fin=sol_kls1_fin, sol_best_fin = sol_best_fin)
        #print 'Approximation: t=%3.4f/%3.4f integr: %e kls: %e kls1 : %e dif: %e' %  (t,tf,er/norm(Ax,'fro'), 
        #       er1/norm(Ax,'fro'), er3/norm(Ax,'fro'),er2/norm(Ax,'fro'))    
        #print 'Time: stand: %3.3f kls: %3.3f kls1: %3.3f' % (t1,t2,t3) 

timestep.py

""" Pure python module that implements dynamical low-rank approximation """
from math import exp,sqrt
import numpy as np
from numpy.linalg import svd,norm,qr


class dyn_lr:
    def __init__(self,U=None,S=None,V=None):
        if U is not None:
            self.U = U.copy() 
            self.S = S.copy()
            self.V = V.copy()
    
    def copy(self):
        c = dyn_lr()
        c.U = self.U.copy()
        c.S = self.S.copy()
        c.V = self.V.copy()
        return c

    def __add__(self,other):
        c = dyn_lr()
        c.U = self.U + other.U
        c.V = self.V + other.V
        c.S = self.S + other.S
        return c
    
    def __sub__(self,other):
        c = self + (other) * (-1.0)
        return c

    def norm(self):
        nrm = norm(self.U.flatten()) ** 2
        nrm += norm(self.V.flatten()) ** 2
        nrm += norm(self.S.flatten()) ** 2
        nrm = sqrt(nrm)
        return nrm
    
    def full(self):
        return np.dot(self.U,np.dot(self.S,self.V.T))

    def __mul__(self,other):
        c = dyn_lr()
        c.U = self.U * other
        c.S = self.S * other
        c.V = self.V * other
        return c

    def __rmul__(self,other):
        c = dyn_lr()
        c.U = self.U * other
        c.S = self.S * other
        c.V = self.V * other
        return c



""" Implicit midpoint rule integrator """
    
def imid_lr(t,tau,X, rhs):
    nit = 100
    Y = X + tau * rhs(t,X)
    tol = 1e-8
    er = 2 * tol
    i = 0
    while ( er > tol and i < nit ):
        Ynew = X + tau * rhs(t+tau/2,0.5 * (Y + X))
        er = (Ynew - Y).norm()/Ynew.norm()
        Y = Ynew.copy()
        i += 1
    #print 'Done in %d iterations' % i
    return Y


""" Dynamical low-rank scheme """
def ksl(t,tau,X,Afun):
    """ Y = kls_lr(t,tau,X,Afun) 
        It takes the t,tau, initial value and a function to compute A (only!) 
    """
    A0 = Afun(t)
    A1 = Afun(t+tau)
    S = X.S 
    U = X.U
    V = X.V
    K = np.dot(U,S); K = K + np.dot((A1 - A0), V); U,S = qr(K) #The devil is here
    S = S - np.dot(U.T,np.dot(A1 - A0,V))
    L = np.dot(V,S.T); L = L + np.dot((A1 - A0).T,U); V,S = qr(L); S = S.T

    Y = dyn_lr(U,S,V)
    
    return Y

""" Symmetrized dynamical low-rank scheme """
def ksl2(t,tau,X,Afun):
    """ Y = kls_lr(t,tau,X,Afun) 
        It takes the t,tau, initial value and a function to compute A (only!) 
    """
    A0 = Afun(t)
    A1 = Afun(t + tau / 2)
    A2 = Afun(t + tau)
    S = X.S 
    U = X.U
    V = X.V
    K = np.dot(U,S); K = K + np.dot((A1 - A0), V); U,S = qr(K) #The devil is here
    S = S - np.dot(U.T,np.dot(A1 - A0,V))
    L = np.dot(V,S.T); L = L + np.dot((A2 - A0).T,U); V,S = qr(L); S = S.T
    S = S - np.dot(U.T,np.dot(A2 - A1,V))
    K = np.dot(U,S); K = K + np.dot((A2 - A1), V); U,S = qr(K) #The devil is here
    Y = dyn_lr(U,S,V)
    
    return Y

""" Dynamical low-rank scheme """
def ksl(t,tau,X,Afun):
    """ Y = kls_lr(t,tau,X,Afun) 
        It takes the t,tau, initial value and a function to compute A (only!) 
    """
    A0 = Afun(t)
    A1 = Afun(t + tau)
    A2 = Afun(t + tau/2)
    S = X.S 
    U = X.U
    V = X.V
    K = np.dot(U,S); K = K + np.dot((A1 - A0), V); U,S = qr(K) #The devil is here
    S = S - np.dot(U.T,np.dot(A1 - A0,V))
    L = np.dot(V,S.T); L = L + np.dot((A1 - A0).T,U); V,S = qr(L); S = S.T

    Y = dyn_lr(U,S,V)
    
    return Y


""" Initial KLS scheme (of the first order) """
def kls_lr1(t,tau,X,Afun):
    """ Y = kls_lr1(t,tau,X,Afun) 
        It takes the t,tau, initial value and a function to compute A (only!) 
    """
    A0 = Afun(t)
    A1 = Afun(t+tau)
    S = X.S 
    U = X.U
    V = X.V
    K = np.dot(U,S); K = K + np.dot((A1 - A0), V); U,S = qr(K) #The devil is here
    L = np.dot(V,S.T); L = L + np.dot((A1 - A0).T,U); V,S = qr(L); S = S.T
    S = S - np.dot(U.T,np.dot(A1 - A0,V))

    Y = dyn_lr(U,S,V)
    
    return Y


""" Dynamical low-rank (second order) scheme """
def kls_lr(t,tau,X,Afun):
    """ Y = kls_lr(t,tau,X,Afun) 
        It takes the t,tau, initial value and a function to compute A (only!) 
        This is a second-order scheme, but in the degenerate case in shows only
        a first-order behaviour (which is not very good). It seems, that the first-order scheme
        introduces O(tau) error to the approximation, which in turn throws you out (?!) from
        the surface
    """
    A0 = Afun(t)
    A1 = Afun(t+tau/2)
    A2 = Afun(t+tau)
    S = X.S 
    U = X.U
    V = X.V
    """ The right order is: U step, S step, V step,  V step, S step, U step (!!!)"""

    K = np.dot(U,S); K = K + np.dot((A1 - A0), V); U,S = qr(K) #The devil is here
    L = np.dot(V,S.T); L = L + np.dot((A1 - A0).T,U); V,S = qr(L); S = S.T
    S = S - np.dot(U.T,np.dot(A2 - A0,V))
    L = np.dot(V,S.T); L = L + np.dot((A2 - A1).T,U); V,S = qr(L); S = S.T
    K = np.dot(U,S); K = K + np.dot((A2 - A1), V); U,S = qr(K) #The devil is here
    
     
    Y = dyn_lr(U,S,V)
    return Y