source: codes/icosagcm/devel/Python/dynamico/time_step.py @ 886

import numpy as np
from scipy import sparse as sparse
import math as math

import getargs
log_master, log_world = getargs.getLogger(__name__)
INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
INFO_ALL, DEBUG_ALL = log_world.info, log_world.debug

class Struct: pass

def Linf(data):
    return abs(data).max()
def compare(msg, data1, data2):
    print msg, Linf(data1-data2)/Linf(data1)    

############################# Tridiagonal solver #############################

class SymTriDiag:
    """ Solve -c(i-1)x(i-1) + b(i)x(i) - c(i)x(i+1) = d(i)
    Forward sweep :
    C(i) =        -c(i)       / (b(i)+c(i-1)C(i-1)), C(0)=0
    D(i) = (d(i)+c(i-1)D(i-1)) / (b(i)+c(i-1)C(i-1)), D(0)=0
    Back substituion :
    x(i) = D(i)-C(i)x(i+1), x(N+1)=0
    """
    def __init__(self, b,c):
        M, N = b.shape
        self.M, self.N, self.b, self.c = M,N,b,c
        B,C = np.zeros((M,N)), np.zeros((M,N-1))
        Bi = b[:,0]
        B[:,0]=Bi
        for i in xrange(1,N):
            Ci = -c[:,i-1]/Bi
            C[:,i-1]=Ci
            Bi = b[:,i]+c[:,i-1]*Ci
            B[:,i]=Bi
        self.B, self.C = B,C
    def dot(self,x):
        M,N,b,c = self.M, self.N, self.b, self.c
        y = np.zeros((M,N))
        y[:,0] = b[:,0]*x[:,0] - c[:,0]*x[:,1]
        for i in xrange(1,N-1):
            y[:,i] = -c[:,i-1]*x[:,i-1] + b[:,i]*x[:,i] - c[:,i]*x[:,i+1]
        y[:,N-1] = -c[:,N-2]*x[:,N-2] + b[:,N-1]*x[:,N-1]
        return y
    def solve(self,d):
        M,N,c,B,C = self.M, self.N, self.c, self.B, self.C
        D,x = np.zeros((M,N)), np.zeros((M,N))
        Di = d[:,0]/B[:,0]
        D[:,0]=Di
        for i in xrange(1,N):
            Di = (d[:,i]+c[:,i-1]*Di)/B[:,i]
            D[:,i]=Di
        xi=D[:,N-1]
        x[:,N-1]=xi
        for i in xrange(N-2,-1,-1):
            xi = D[:,i]-C[:,i]*xi
            x[:,i]=xi
        return x

########################### ARK time schemes ##########################

class SIRK:
    """ 
    Generic SIRK scheme with fast/slow Butcher tableaux af,as
    Following Weller et al. (2013)
    sj=slow(yj), fj=fast(yj)
    zj = y(n) + dt*( sum(l<j)asjl*sl + sum(l<j)afjl*fl )
    yj = zj + dt*afjj*fast(yj)
    y(n+1) = y(n)+dt*sum(wsj*sj+wsf*fj)
    Reformulation : remarking that 
    y(j)   = y(n) + dt*sum_(l<=j)[as(j,l)*s(l)   + af(j,l)*f(l) ]
    z(j+1) = y(n) + dt*sum_(l<=j)[as(j+1,l)*s(l) + af(j+1,l)*f(l) ]
           = y(j) + dt*sum_(l<=j)[cs(j,l)*s(l)   + cf(j,l)*f(l) ]
    where b(j,l)=a(j,l) ; b(nstage,l)=w(l) ; c(j) = b(j+1)-b(j) => c0=b1
    => implementation
    z(0) = y(n)
    DO j=0, nstage-1
        [y(j), fj, sj] = bwd_fast_slow(z(j),0.)
        z(j+1) = y(j) + dt*sum_(l<=j)[cs(j,l)*s(l) + cf(j,l)*f(l) ]
    END DO
    y(n+1) = z(nstage)

    where z(j) and y(j) are actually stored in a single variable zj
    and [y,f,s] = bwd_fast_slow(z,tau) 
        * solves y - tau*fast(y) = z
        * computes f=fast(y) and s=slow(y)
    """
    def cjl(self,b): # converts Butcher tableau b(j,l) to update coefficients c(j,l)
        return np.array([ b[j+1,:]-b[j,:] for j in range(self.nstage)])
    def __init__(self,precision,bwd_fast_slow, dt,nstage, bsjl,bfjl):
        if bwd_fast_slow is None:
            DEBUG('SIRK with Fortran-side time stepping')
            precision=np.float64
        else:
            DEBUG('SIRK with Python-side time stepping at precision %s' % precision)

        self.bwd_fast_slow = bwd_fast_slow
        self.dt, self.nstage = dt, nstage
        csjl = dt*self.cjl(bsjl)
        cfjl = dt*self.cjl(bfjl)
        self.csjl, self.cfjl = [x.astype(precision) for x in csjl, cfjl]
        self.tauj = np.array([dt*bfjl[j,j] for j in range(self.nstage)], dtype=precision)
        DEBUG('Types of csjl, cfjl, tauj : %s %s %s' % (self.csjl.dtype, self.cfjl.dtype, self.tauj.dtype) )
    def next(self,zj):
        csjl, cfjl, tauj = self.csjl, self.cfjl, self.tauj
        fastj, slowj = [], []
        for j in range(self.nstage): # here zj is z(j)
            zj, fj, sj = self.bwd_fast_slow(zj, tauj[j])
            fastj.append(fj)
            slowj.append(sj)
            for l in range(0,j+1): # here zj is initially y(j) and finally z(j+1)
                zj = [z + csjl[j,l]*s + cfjl[j,l]*f for z,s,f in zip(zj, slowj[l], fastj[l]) ]
#            print 'SIRK : zj has type ', [ z.dtype for z in zj]
        return zj        
    def advance(self, flow, N):
        for k in range(N):
            flow=self.next(flow)
        return flow
    
def SIRK1(bwd_fast_slow, dt): # untested
    """ Forward-Backward""" 
    wj  = np.array([1])
    bsjl = np.array([[0],wj])
    bfjl = np.array([[1],wj])
    return SIRK(bwd_fast_slow,dt,1,bsjl,bfjl)
        
def ARK2(bwd_fast_slow, dt, precision=None, a32 = (3.+2.*np.sqrt(2.))/6.):
    """ ARK2 scheme by Giraldo et al. (2013), noted ARK2(2,3,2) in Weller et al. (2013) """ 
    gamma = 1.-1./np.sqrt(2.)
    delta = 1./(2.*np.sqrt(2.))
    wj  = np.array([delta,delta,gamma])
    bsjl = np.array([[0,0,0],[2.*gamma,0,0],[1.-a32,a32,0],wj])
    bfjl = np.array([[0,0,0],[gamma,gamma,0],[delta,delta,gamma],wj])
    return SIRK(precision,bwd_fast_slow,dt,3,bsjl,bfjl)
        
def UJ3(bwd_fast_slow, dt):
    """ Strang-carryover scheme by Ullrich et al. (2013), noted UJ3(1,3,2) in Weller et al. (2013) """
    zz   = [0,0,0,0,0,0]
    wsj  = [0,1/6.,1/6.,2/3.,0,0]
    h    = [.5,0,0,0,0,0]
    wfj  = [.5,0,0,0,0,.5]
    bsjl = np.array([zz,zz,[0,1,0,0,0,0],
                     [0,.25,.25,0,0,0],wsj,wsj,wsj])
    bfjl = np.array([zz,h,h,h,h,wfj,wfj])
    return SIRK(bwd_fast_slow,dt,6,bsjl,bfjl)

######### Explicit RK schemes written as ARK with two identical Butcher tableaus ########

def RK4(bwd_fast_slow, dt, precision=None):
    """ 4-stage 4th-order RK """ 
    bjl = np.array([[0,0,0,0],
                    [.5,0,0,0],
                    [0,.5,0,0],
                    [0.,0.,1.,0],
                    [1/6.,1/3.,1/3.,1/6.]])
    return SIRK(precision, bwd_fast_slow,dt,4,bjl,bjl)

def RK1(bwd_fast_slow, dt, precision=None):
    """ 4-stage 4th-order RK """ 
    bjl = np.array([[0.],[1.]])
    return SIRK(precision, bwd_fast_slow,dt,1,bjl,bjl)

######### Simplified explicit RK schemes (2nd-order) ########

def RK_simple(bwd_fast_slow, dt, tau):
    # take coefficients of simplified RK scheme
    # y(n) := x + tau*f(y(n-1))
    # and computes the equivalent Butcher tableau
    tau=np.array(tau)
    n = tau.size # number of stages
    bjl = np.zeros((n+1,n)) # a[stage, coef]
    for i in range(n):
        bjl[i+1,i]=tau[i] 
    return SIRK(bwd_fast_slow,dt,n,bjl,bjl)

def RKn_simple(n, bwd_fast_slow, dt):
    tau = [1./(n-i) for i in range(n)]
    return RK_simple(bwd_fast_slow, dt, tau)

def RK_Kinnmark(K, bwd_fast_slow, dt):
    """ RK schemes with max stability along imaginary axis
        Kinnmark & Gray, 1984 """
    ee=np.zeros((K+1,K+1))
    ee[2,0], ee[2,1], ee[2,2] = 1,1,1
    ee[3,0], ee[3,1], ee[3,2], ee[3,3] = 1,2,2,2
    for k in range(4,K+1):
        ee[k,0]=1
        for n in range(1,k+1):
            ee[k,n]=ee[k-2,n]+2*ee[k-1,n-1]
    tau = [ee[K,K-n]/ee[K,K-n-1]/(K-1) for n in range(K)]
    return RK_simple(bwd_fast_slow, dt, tau)