-
Notifications
You must be signed in to change notification settings - Fork 0
/
exercise_7_15.py
69 lines (53 loc) · 1.72 KB
/
exercise_7_15.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
"""Exercise 7.15
Use difference method with h = 1 to solve the problem
y'' - xy' + y = 1, y(0) = y(2) = 0
"""
from typing import Callable
import numpy as np
from scipy.linalg import solve_banded
from matplotlib import pyplot as plt
_F = Callable[[np.ndarray], float]
N = 20
"""Number of points - 1"""
def make_mat(start:float, cond1:float, end:float, cond2:float,\
r_m:_F, q_m:_F, p_m:_F, ab_mode=False):
"""Construct the matrix and vector."""
x_val = np.linspace(start, end, N + 1)[1:N]
interval = (end - start) / N
a_m = - (1 + interval*p_m(x_val)/2)
d_m = 2 + interval**2*q_m(x_val)
c_m = - (1 - interval*p_m(x_val)/2)
b_m = - interval**2*r_m(x_val)
# construct matrix
if ab_mode:
result_mat = np.zeros((3, N-1))
result_mat[0, 1:] = c_m[:N-2]
result_mat[1, :] = d_m[:]
result_mat[2, :N-2] = a_m[1:]
else:
result_mat = np.zeros((N-1, N-1))
result_mat[0, 0], result_mat[0, 1] = d_m[0], c_m[0]
for i in range(1, N-2):
result_mat[i, i-1], result_mat[i, i], result_mat[i, i+1] = a_m[i], d_m[i], c_m[i]
result_mat[N-2, N-3], result_mat[N-2, N-2] = a_m[N-2], d_m[N-2]
# construct vector
result_vec = np.zeros((N-1,))
result_vec[0] = b_m[0] - a_m[0] * cond1
result_vec[1: N-2] = b_m[1: N-2]
result_vec[N-2] = b_m[N-2] - c_m[N-2] * cond2
return result_mat, result_vec
a_mat, b_vec = make_mat(
0, 0, 2, 0,
lambda x:np.ones(x.shape),
lambda x:-np.ones(x.shape),
lambda x:x,
True
)
print(a_mat, b_vec, sep="\n")
result = solve_banded((1, 1), a_mat, b_vec)
print(result)
# Draw the result
y_val = np.zeros((N+1,))
y_val[1:N] = result
plt.plot(np.linspace(0, 2, N + 1), y_val)
plt.show()