-
Notifications
You must be signed in to change notification settings - Fork 0
/
CrankNicolson.py
84 lines (63 loc) · 2.39 KB
/
CrankNicolson.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
'''
Crank-Nicolson method to solve 1D reaction-diffusion equation:
u_t = D * u_xx
with Neumann boundary conditions
at x=0: u_x = sin(2*pi)
at x=L: u_x = sin(2*pi)
with L=1 and initial condition:
u(x,0) = u(x,0) = (1.0/2.0)+ np.cos(2.0*np.pi*x) - (1.0/2.0)*np.cos(3*np.pi*x)
'''
def CrankNicolson(M, lambd, T = 0.5, L = 1, k = 1):
#Parameters needed to solve the equation within the explicit method
# M = GRID POINTS on space interval
N = (M**2) #GRID POINTS on time interval
# ---- Length of the wire in x direction ----
x0, xL = 0, L
# ----- Spatial discretization step -----
dx = (xL - x0)/(M-1)
# ---- Final time ----
t0, tF = 0, T
# ----- Time step -----
dt = (tF - t0)/(N-1)
#lambd = dt*k/(2.0*dx**2)
a0 = 1 + 2*lambd
c0 = 1 - 2*lambd
xspan = np.linspace(x0, xL, M)
tspan = np.linspace(t0, tF, N)
maindiag_a0 = a0*np.ones((1,M))
offdiag_a0 = (-lambd)*np.ones((1, M-1))
maindiag_c0 = c0*np.ones((1,M))
offdiag_c0 = lambd*np.ones((1, M-1))
#Left-hand side tri-diagonal matrix
a = maindiag_a0.shape[1]
diagonalsA = [maindiag_a0, offdiag_a0, offdiag_a0]
A = sparse.diags(diagonalsA, [0,-1,1], shape=(a,a)).toarray()
A[0,1] = (-2)*lambd
A[M-1,M-2] = (-2)*lambd
#Right-hand side tri-diagonal matrix
c = maindiag_c0.shape[1]
diagonalsC = [maindiag_c0, offdiag_c0, offdiag_c0]
Arhs = sparse.diags(diagonalsC, [0,-1,1], shape=(c,c)).toarray()
Arhs[0,1] = 2*lambd
Arhs[M-1,M-2] = 2*lambd
# ----- Initializes matrix U -----
U = np.zeros((M, N))
#----- Initial condition -----
U[:,0] = (1.0/2.0)+ np.cos(2.0*np.pi*xspan) - (1.0/2.0)*np.cos(3*np.pi*xspan)
#----- Neumann boundary conditions -----
#Add one line above and one line below using finit differences
leftBC = np.arange(1, N+1)
f = np.sin(2*leftBC*np.pi)
rightBC = np.arange(1, N+1)
g = np.sin(2*rightBC*np.pi)
for k in range(1, N):
ins = np.zeros((M-2,1)).ravel()
b1 = np.asarray([4*lambd*dx*f[k], 4*lambd*dx*g[k]])
b1 = np.insert(b1, 1, ins)
b2 = np.matmul(Arhs, np.array(U[0:M, k-1]))
b = b1 + b2 # Right hand side
U[0:M, k] = np.linalg.solve(A,b) # Solve x=A\b
return (U, tspan, xspan)
#U, tspan, xspan = CrankNicolson(M = 14, lambd = 1.0/6.0)
#Uexact, x, t = ExactSolution(M = 14)
#surfaceplot(U, Uexact, tspan, xspan)