-
Notifications
You must be signed in to change notification settings - Fork 0
/
ODEs.py
127 lines (95 loc) · 3.17 KB
/
ODEs.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
# %%
import numpy as np
import matplotlib.pyplot as plt
# the ODE function we want to work with
def f(x,y):
return x**2 + 0.1 * y
# define interval borders
xmin = -2
xmax = 2
# discretization resolution
n = 15
# initial value
y0 = 2
## numerically integrate ODEs with different methods
# Euler Method for f with initial value y0
# n steps in the interval [a,b].
def euler(f,a,b,n,y0):
# init x_i's and y_i's
x = np.linspace(a,b,n)
y = np.zeros_like(x)
y[0] = y0
for i in range(0,len(x)-1):
# distance to the next point to approximate
d = [x[i+1]]-x[i]
# slope at current point
slope = f(x[i],y[i])
# to get the next point follow the slope for distance d
y[i+1] = y[i] + d * slope
return y
# Midpoint method for f with initial value y0
# n steps in the interval [a,b].
def midpoint(f,a,b,n,y0):
# init x_i's and y_i's
x = np.linspace(a,b,n)
y = np.zeros_like(x)
y[0] = y0
for i in range(0,len(x)-1):
# distance to the next point to approximate
dnext = [x[i+1]]-x[i]
# use Euler Method to calculate y of midpoint
ymid = y[i] + dnext/2 * f(x[i],y[i])
xmid = x[i] + dnext/2
# slope at midpoint between next and current point
slope = f(xmid,ymid)
# to get the next point follow the slope for distance d
y[i+1] = y[i] + dnext * slope
return y
# Modified Euler / Heun's method for f with initial value y0
# n steps in the interval [a,b].
def modeuler(f,a,b,n,y0):
# init x_i's and y_i's
x = np.linspace(a,b,n)
y = np.zeros_like(x)
y[0] = y0
for i in range(0,len(x)-1):
# distance to the next point to approximate
d = [x[i+1]]-x[i]
# slope at the current point
slope_here = f(x[i],y[i])
# use slope to find a temporary y
ytemp = y[i] + d * slope_here
# use ytemp to find slope on the other side of the discretization step
slope_next = f(x[i+1],ytemp)
slope_avg = (slope_here + slope_next)/2
# to get the next point follow the average slope for distance d
y[i+1] = y[i] + d * slope_avg
return y
## data to plot
xs = np.linspace(xmin,xmax,n,dtype=np.float64)
ys_euler = euler(f,xmin,xmax,n,y0)
ys_midpoint = midpoint(f,xmin,xmax,n,y0)
ys_modeuler = modeuler(f,xmin,xmax,n,y0)
fig,ax = plt.subplots()
ax.plot(xs, ys_euler, label='euler')
ax.plot(xs, ys_midpoint, label='midpoint' )
ax.plot(xs, ys_modeuler, '--', label='modeuler' )
## add vector field data
# generate a grid of points. (later arrow locations)
(ymin, ymax) = ax.get_ylim()
ys = np.linspace(ymin,ymax,n)
[xgrid, ygrid] = np.meshgrid(xs,ys)
# evaluate x and y direction components of the arrow vectors
# (if a curve goes through those points it has this slope)
dy = f(xgrid,ygrid)
dx = np.ones_like(dy)
# normalize arrows
r = np.power(np.add(np.power(dx,2), np.power(dy,2)),0.5)
quiveropts = dict(color='pink', units='xy', angles='xy', width=0.02)
ax.quiver(xgrid, ygrid, dx/r, dy/r, **quiveropts)
fig.suptitle(f'numerical integration methods for $x^2+0.1*y(x)$ with $x\in [{xmin},{xmax}], y_0 = {y0}$')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.grid()
ax.legend()
plt.show()