1 | import numpy as np |
---|
2 | from scipy import sparse as sparse |
---|
3 | import math as math |
---|
4 | |
---|
5 | class Struct: pass |
---|
6 | |
---|
7 | def Linf(data): |
---|
8 | return abs(data).max() |
---|
9 | def compare(msg, data1, data2): |
---|
10 | print msg, Linf(data1-data2)/Linf(data1) |
---|
11 | |
---|
12 | ############################# Tridiagonal solver ############################# |
---|
13 | |
---|
14 | class SymTriDiag: |
---|
15 | """ Solve -c(i-1)x(i-1) + b(i)x(i) - c(i)x(i+1) = d(i) |
---|
16 | Forward sweep : |
---|
17 | C(i) = -c(i) / (b(i)+c(i-1)C(i-1)), C(0)=0 |
---|
18 | D(i) = (d(i)+c(i-1)D(i-1)) / (b(i)+c(i-1)C(i-1)), D(0)=0 |
---|
19 | Back substituion : |
---|
20 | x(i) = D(i)-C(i)x(i+1), x(N+1)=0 |
---|
21 | """ |
---|
22 | def __init__(self, b,c): |
---|
23 | M, N = b.shape |
---|
24 | self.M, self.N, self.b, self.c = M,N,b,c |
---|
25 | B,C = np.zeros((M,N)), np.zeros((M,N-1)) |
---|
26 | Bi = b[:,0] |
---|
27 | B[:,0]=Bi |
---|
28 | for i in xrange(1,N): |
---|
29 | Ci = -c[:,i-1]/Bi |
---|
30 | C[:,i-1]=Ci |
---|
31 | Bi = b[:,i]+c[:,i-1]*Ci |
---|
32 | B[:,i]=Bi |
---|
33 | self.B, self.C = B,C |
---|
34 | def dot(self,x): |
---|
35 | M,N,b,c = self.M, self.N, self.b, self.c |
---|
36 | y = np.zeros((M,N)) |
---|
37 | y[:,0] = b[:,0]*x[:,0] - c[:,0]*x[:,1] |
---|
38 | for i in xrange(1,N-1): |
---|
39 | y[:,i] = -c[:,i-1]*x[:,i-1] + b[:,i]*x[:,i] - c[:,i]*x[:,i+1] |
---|
40 | y[:,N-1] = -c[:,N-2]*x[:,N-2] + b[:,N-1]*x[:,N-1] |
---|
41 | return y |
---|
42 | def solve(self,d): |
---|
43 | M,N,c,B,C = self.M, self.N, self.c, self.B, self.C |
---|
44 | D,x = np.zeros((M,N)), np.zeros((M,N)) |
---|
45 | Di = d[:,0]/B[:,0] |
---|
46 | D[:,0]=Di |
---|
47 | for i in xrange(1,N): |
---|
48 | Di = (d[:,i]+c[:,i-1]*Di)/B[:,i] |
---|
49 | D[:,i]=Di |
---|
50 | xi=D[:,N-1] |
---|
51 | x[:,N-1]=xi |
---|
52 | for i in xrange(N-2,-1,-1): |
---|
53 | xi = D[:,i]-C[:,i]*xi |
---|
54 | x[:,i]=xi |
---|
55 | return x |
---|
56 | |
---|
57 | ########################### ARK time schemes ########################## |
---|
58 | |
---|
59 | class SIRK: |
---|
60 | """ |
---|
61 | Generic SIRK scheme with fast/slow Butcher tableaux af,as |
---|
62 | Following Weller et al. (2013) |
---|
63 | sj=slow(yj), fj=fast(yj) |
---|
64 | zj = y(n) + dt*( sum(l<j)asjl*sl + sum(l<j)afjl*fl ) |
---|
65 | yj = zj + dt*afjj*fast(yj) |
---|
66 | y(n+1) = y(n)+dt*sum(wsj*sj+wsf*fj) |
---|
67 | Reformulation : remarking that |
---|
68 | y(j) = y(n) + dt*sum_(l<=j)[as(j,l)*s(l) + af(j,l)*f(l) ] |
---|
69 | z(j+1) = y(n) + dt*sum_(l<=j)[as(j+1,l)*s(l) + af(j+1,l)*f(l) ] |
---|
70 | = y(j) + dt*sum_(l<=j)[cs(j,l)*s(l) + cf(j,l)*f(l) ] |
---|
71 | where b(j,l)=a(j,l) ; b(nstage,l)=w(l) ; c(j) = b(j+1)-b(j) => c0=b1 |
---|
72 | => implementation |
---|
73 | z(0) = y(n) |
---|
74 | DO j=0, nstage-1 |
---|
75 | [y(j), fj, sj] = bwd_fast_slow(z(j),0.) |
---|
76 | z(j+1) = y(j) + dt*sum_(l<=j)[cs(j,l)*s(l) + cf(j,l)*f(l) ] |
---|
77 | END DO |
---|
78 | y(n+1) = z(nstage) |
---|
79 | |
---|
80 | where z(j) and y(j) are actually stored in a single variable zj |
---|
81 | and [y,f,s] = bwd_fast_slow(z,tau) |
---|
82 | * solves y - tau*fast(y) = z |
---|
83 | * computes f=fast(y) and s=slow(y) |
---|
84 | """ |
---|
85 | def cjl(self,b): # converts Butcher tableau b(j,l) to update coefficients c(j,l) |
---|
86 | return np.array([ b[j+1,:]-b[j,:] for j in range(self.nstage)]) |
---|
87 | def __init__(self,precision,bwd_fast_slow, dt,nstage, bsjl,bfjl): |
---|
88 | if bwd_fast_slow is None: |
---|
89 | print 'SIRK with Fortran-side time stepping' |
---|
90 | precision=np.float64 |
---|
91 | else: |
---|
92 | print 'SIRK with Python-side time stepping at precision ', precision |
---|
93 | |
---|
94 | self.bwd_fast_slow = bwd_fast_slow |
---|
95 | self.dt, self.nstage = dt, nstage |
---|
96 | csjl = dt*self.cjl(bsjl) |
---|
97 | cfjl = dt*self.cjl(bfjl) |
---|
98 | self.csjl, self.cfjl = [x.astype(precision) for x in csjl, cfjl] |
---|
99 | self.tauj = np.array([dt*bfjl[j,j] for j in range(self.nstage)], dtype=precision) |
---|
100 | print 'Types of csjl, cfjl, tauj :', self.csjl.dtype, self.cfjl.dtype, self.tauj.dtype |
---|
101 | def next(self,zj): |
---|
102 | csjl, cfjl, tauj = self.csjl, self.cfjl, self.tauj |
---|
103 | fastj, slowj = [], [] |
---|
104 | for j in range(self.nstage): # here zj is z(j) |
---|
105 | zj, fj, sj = self.bwd_fast_slow(zj, tauj[j]) |
---|
106 | fastj.append(fj) |
---|
107 | slowj.append(sj) |
---|
108 | for l in range(0,j+1): # here zj is initially y(j) and finally z(j+1) |
---|
109 | zj = [z + csjl[j,l]*s + cfjl[j,l]*f for z,s,f in zip(zj, slowj[l], fastj[l]) ] |
---|
110 | # print 'SIRK : zj has type ', [ z.dtype for z in zj] |
---|
111 | return zj |
---|
112 | def advance(self, flow, N): |
---|
113 | for k in range(N): |
---|
114 | flow=self.next(flow) |
---|
115 | return flow |
---|
116 | |
---|
117 | def SIRK1(bwd_fast_slow, dt): # untested |
---|
118 | """ Forward-Backward""" |
---|
119 | wj = np.array([1]) |
---|
120 | bsjl = np.array([[0],wj]) |
---|
121 | bfjl = np.array([[1],wj]) |
---|
122 | return SIRK(bwd_fast_slow,dt,1,bsjl,bfjl) |
---|
123 | |
---|
124 | def ARK2(bwd_fast_slow, dt, precision=None, a32 = (3.+2.*np.sqrt(2.))/6.): |
---|
125 | """ ARK2 scheme by Giraldo et al. (2013), noted ARK2(2,3,2) in Weller et al. (2013) """ |
---|
126 | gamma = 1.-1./np.sqrt(2.) |
---|
127 | delta = 1./(2.*np.sqrt(2.)) |
---|
128 | wj = np.array([delta,delta,gamma]) |
---|
129 | bsjl = np.array([[0,0,0],[2.*gamma,0,0],[1.-a32,a32,0],wj]) |
---|
130 | bfjl = np.array([[0,0,0],[gamma,gamma,0],[delta,delta,gamma],wj]) |
---|
131 | return SIRK(precision,bwd_fast_slow,dt,3,bsjl,bfjl) |
---|
132 | |
---|
133 | def UJ3(bwd_fast_slow, dt): |
---|
134 | """ Strang-carryover scheme by Ullrich et al. (2013), noted UJ3(1,3,2) in Weller et al. (2013) """ |
---|
135 | zz = [0,0,0,0,0,0] |
---|
136 | wsj = [0,1/6.,1/6.,2/3.,0,0] |
---|
137 | h = [.5,0,0,0,0,0] |
---|
138 | wfj = [.5,0,0,0,0,.5] |
---|
139 | bsjl = np.array([zz,zz,[0,1,0,0,0,0], |
---|
140 | [0,.25,.25,0,0,0],wsj,wsj,wsj]) |
---|
141 | bfjl = np.array([zz,h,h,h,h,wfj,wfj]) |
---|
142 | return SIRK(bwd_fast_slow,dt,6,bsjl,bfjl) |
---|
143 | |
---|
144 | ######### Explicit RK schemes written as ARK with two identical Butcher tableaus ######## |
---|
145 | |
---|
146 | def RK4(bwd_fast_slow, dt, precision=None): |
---|
147 | """ 4-stage 4th-order RK """ |
---|
148 | bjl = np.array([[0,0,0,0], |
---|
149 | [.5,0,0,0], |
---|
150 | [0,.5,0,0], |
---|
151 | [0.,0.,1.,0], |
---|
152 | [1/6.,1/3.,1/3.,1/6.]]) |
---|
153 | return SIRK(precision, bwd_fast_slow,dt,4,bjl,bjl) |
---|
154 | |
---|
155 | def RK1(bwd_fast_slow, dt, precision=None): |
---|
156 | """ 4-stage 4th-order RK """ |
---|
157 | bjl = np.array([[0.],[1.]]) |
---|
158 | return SIRK(precision, bwd_fast_slow,dt,1,bjl,bjl) |
---|
159 | |
---|
160 | ######### Simplified explicit RK schemes (2nd-order) ######## |
---|
161 | |
---|
162 | def RK_simple(bwd_fast_slow, dt, tau): |
---|
163 | # take coefficients of simplified RK scheme |
---|
164 | # y(n) := x + tau*f(y(n-1)) |
---|
165 | # and computes the equivalent Butcher tableau |
---|
166 | tau=np.array(tau) |
---|
167 | n = tau.size # number of stages |
---|
168 | bjl = np.zeros((n+1,n)) # a[stage, coef] |
---|
169 | for i in range(n): |
---|
170 | bjl[i+1,i]=tau[i] |
---|
171 | return SIRK(bwd_fast_slow,dt,n,bjl,bjl) |
---|
172 | |
---|
173 | def RKn_simple(n, bwd_fast_slow, dt): |
---|
174 | tau = [1./(n-i) for i in range(n)] |
---|
175 | return RK_simple(bwd_fast_slow, dt, tau) |
---|
176 | |
---|
177 | def RK_Kinnmark(K, bwd_fast_slow, dt): |
---|
178 | """ RK schemes with max stability along imaginary axis |
---|
179 | Kinnmark & Gray, 1984 """ |
---|
180 | ee=np.zeros((K+1,K+1)) |
---|
181 | ee[2,0], ee[2,1], ee[2,2] = 1,1,1 |
---|
182 | ee[3,0], ee[3,1], ee[3,2], ee[3,3] = 1,2,2,2 |
---|
183 | for k in range(4,K+1): |
---|
184 | ee[k,0]=1 |
---|
185 | for n in range(1,k+1): |
---|
186 | ee[k,n]=ee[k-2,n]+2*ee[k-1,n-1] |
---|
187 | tau = [ee[K,K-n]/ee[K,K-n-1]/(K-1) for n in range(K)] |
---|
188 | return RK_simple(bwd_fast_slow, dt, tau) |
---|