-
Notifications
You must be signed in to change notification settings - Fork 0
/
ode_solver.py
64 lines (43 loc) · 1.46 KB
/
ode_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
60
61
62
63
64
import numpy as np
from matplotlib import pyplot as plt
from numpy import linalg as la
import math
def plot_graph(t, p, mu, start):
w = mu / 1j
# Graph
fig, axes = plt.subplots(2, 1)
fig.tight_layout()
# z=0
axes[0].plot(t, solve_linear_high_order_ode(t, p, 0, mu, start).real, color='blue')
axes[0].set_title('z=0')
# z=1
axes[1].plot(t, solve_linear_high_order_ode(t, p, 1, mu, start).real, color='blue')
axes[1].plot(t, func1(w * t), color='red')
axes[1].set_title('z=1')
plt.show()
def solve_linear_high_order_ode(t, p, z, mu, start): # Solves ODE
# Find roots of a polynomial
poly = np.poly1d(p)
roots = np.roots(p)
start = np.array(start)
mu1 = np.array([mu**i for i in range(len(start))])
# Matrix A
A = np.array([[i**j for i in roots] for j in range(len(start))])
b = start - (z / poly(mu)) * mu1
# Find coefficients
c = la.solve(A, b)
# Find y
y_p = (z / poly(mu)) * np.exp(mu * t)
y_h = 0
for i in range(len(roots)):
y_h += np.exp(roots[i] * t) * c[i]
y = y_p + y_h
return y
def func1(x):
return np.cos(x)
"""Tests
"""
print(plot_graph(np.linspace(0,50,1000), [1,1], 2j, [1]))
print(plot_graph(np.linspace(0,50,1000), [2,1,10], 1j, [1,0]))
print(plot_graph(np.linspace(0,200,1000), [1,0,4], 2.1j, [1,0]))
print(plot_graph(np.linspace(0,100,1000), [1,0.1,17,1,16], 1j, [0,0,0,1]))