-
Notifications
You must be signed in to change notification settings - Fork 0
/
tridiagonal_solver.py
59 lines (41 loc) · 1.27 KB
/
tridiagonal_solver.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
"""
Task (b)
"""
import numpy as np
import numpy.linalg as lin
from scipy.sparse import diags
def solve_tridiagonal(a,b,f):
n = len(f) # n is the number of rows
c=np.zeros(n)
d=np.zeros(n-1)
y=np.zeros(n)
u=np.zeros(n)
c[0]=np.sqrt(a[0])
d[0]=b[0]/c[0]
for i in range(1,n-1):
c[i]= np.sqrt(a[i]-(d[i-1]**2))
d[i]=b[i]/c[i]
c[n-1]= np.sqrt(a[n-1]-(d[n-2]**2))
#Forward substitution
y[0]=f[0]/c[0]
for i in range(1,n):
y[i]=(f[i]-d[i-1]*y[i-1])/c[i]
#Backward substitution
u[n-1]=y[n-1]/c[n-1]
for i in range(n-2,-1,-1):
u[i]=(y[i]-u[i+1]*d[i])/c[i]
return u
n=3
a = np.append(2*np.ones(n-1), [1])
b = -1*np.ones(n-1)
f = np.ones(n)
k = [-1*np.ones(n-1),np.append(2*np.ones(n-1), [1]),-1*np.ones(n-1)] # sequence of arrays containing the values
#on the diagonals, corresponding to offsets
A = diags(k,[-1,0,1]).toarray() #creates tridiagonal matrix and converts it to an array
# offset sets diagonals to: k = 0 the main diagonal (default), k > 0 the kth upper diagonal,
# k < 0 the kth lower diagonalonal matrix
#print(lin.solve(A,f))
print(solve_tridiagonal(a,b,f))
""" As n increases, the entries of the
displacement vector u also increase
"""