-
Notifications
You must be signed in to change notification settings - Fork 0
/
RK.py
131 lines (108 loc) · 3.93 KB
/
RK.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
# -*- coding: utf-8 -*-
# Author: Jiajun Ren <jiajunren0522@gmail.com>
'''
automatic Runge-Kutta method coefficient calculation
'''
import numpy as np
def runge_kutta_explicit_tableau(RK_method):
'''
Butcher tableau of the explicit Runge-Kutta methods.
different types of propagation methods: e^-iHdt \Psi
1. classical 4th order Runge Kutta
0 |
1/2 | 1/2
1/2 | 0 1/2
1 | 0 0 1
----------------------------
| 1/6 1/3 1/3 1/6
2. Heun's method
0 |
1 | 1
----------------------------
1/2 1/2
'''
assert RK_method in ["Forward_Euler","midpoint_RK2","Heun_RK2","Ralston_RK2","Kutta_RK3","C_RK4","38rule_RK4", "Fehlberg5"]
print "Runge Kutta method", RK_method
if RK_method == "Forward_Euler":
# Euler explicit
a = np.array([[0]])
b = np.array([1])
c = np.array([0])
Nstage = 1
elif RK_method in ["midpoint_RK2","Heun_RK2","Ralston_RK2"]:
if RK_method == "midpoint_RK2":
# if alpha == 1, midpoint method
alpha = 1.0
elif RK_method == "Heun_RK2":
# if alpha == 0.5, heun's method
alpha = 0.5
elif RK_method == "Ralston_RK2":
alpha = 2.0/3.0
a = np.array([[0,0],[alpha,0]])
b = np.array([1-0.5/alpha,0.5/alpha])
c = np.array([0,alpha])
Nstage = 2
elif RK_method == "Kutta_RK3":
# Kutta's third-order method
a = np.array([[0,0,0],[0.5,0,0],[-1,2,0]])
b = np.array([1.0/6.0,2.0/3.0,1.0/6.0])
c = np.array([0,0.5,1])
Nstage = 3
elif RK_method == "C_RK4":
# Classic fourth-order method
a = np.array([[0,0,0,0],[0.5,0,0,0],
[0,0.5,0,0],[0,0,1,0]])
b = np.array([1.0/6.0,1.0/3.0,1.0/3.0,1.0/6.0])
c = np.array([0,0.5,0.5,1.0])
Nstage = 4
elif RK_method == "38rule_RK4":
# 3/8 rule fourth-order method
a = np.array([[0,0,0,0],\
[1.0/3.0,0,0,0],\
[-1.0/3.0,1,0,0],\
[1,-1,1,0]])
b = np.array([1.0/8.0, 3.0/8.0, 3.0/8.0, 1.0/8.0])
c = np.array([0.0, 1.0/3.0, 2.0/3.0 ,1.0])
Nstage = 4
elif RK_method == "Fehlberg5":
a = np.array([[0,0,0,0,0,0],
[1/4., 0., 0., 0., 0., 0.],
[3./32, 9./32, 0, 0, 0, 0 ],
[1932./2197, -7200./2197, 7296./2197, 0, 0, 0],
[439./216, -8., 3680./513, -845/4104, 0., 0.],
[-8./27, 2., -3544./2565, 1859./4104, -11./40, 0.]])
b = np.array([16./135,0.,6656./12825,28561./56430,-9./50,2./55])
c = np.array([0., 1./4, 3./8, 12./13, 1., 1./2])
Nstage = 5
a = a.astype(np.float64)
b = b.astype(np.float64)
c = c.astype(np.float64)
return [a,b,c,Nstage]
def runge_kutta_explicit_coefficient(tableau):
'''
only suited for time-independent propagator
y'(t) = fy(t) f is time-independent
the final formula is
y(t+dt) = d0 y(t) + d1 fy(t) dt + d2 f^2 y(t) dt^2 + ...
0 f f^2 f^3 f^4
v0
v1
v2
v3
Though, each order has different versions of RK methods, if f is time
independent, the coefficient is the same. For example, Classical 4th order
Runge Kutta and 3/8 rule Runge Kutta has some coefficient.
'''
a, b, c, Nstage = tableau
v = np.zeros([Nstage, Nstage+1]) # each row is a point
k = np.zeros([Nstage, Nstage+1]) # slope at point
v[0,0] = 1.0
for ik in xrange(a.shape[0]):
k[ik,1:] = v[ik,:-1]
k[ik,0] = 0.0
if ik != a.shape[0]-1:
v[ik+1,:] = v[0,:]
for iv in xrange(a.shape[1]):
v[ik+1,:] += a[ik+1,iv] * k[iv,:]
d = v[0,:] + b.dot(k)
return d