174 lines
3.7 KiB
Python
Executable File
174 lines
3.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
"""Ordinary Differential Equations"""
|
|
|
|
import numpy as np
|
|
|
|
# Integrations Schemes #
|
|
########################
|
|
|
|
# Single Step
|
|
#------------
|
|
def mk_phi_euler(f):
|
|
""" Return a Phi computed for the Euler Method. """
|
|
def phi_euler(x, y, h):
|
|
return f(x[-1], y[-1])
|
|
|
|
return phi_euler
|
|
|
|
def mk_phi_euler_mod(f):
|
|
""" Return a Phi computed for the Modified Euler (Collatz) Method. """
|
|
def phi_euler_mod(x, y, h):
|
|
return f(x[-1] + 0.5*h, y[-1] + 0.5*h*f(x[-1], y[-1]))
|
|
|
|
return phi_euler_mod
|
|
|
|
def mk_phi_heun(f):
|
|
""" Return a Phi computed for the Heun Method. """
|
|
def phi_heun(x, y, h):
|
|
return (f(x[-1], y[-1]) + f(x[-1] + h, y[-1] + h * f(x[-1], y[-1])))/2
|
|
|
|
return phi_heun
|
|
|
|
def mk_phi_rk4(f):
|
|
""" Return a Phi computed for the 4th order Runge-Kutta Method. """
|
|
def phi_rk4(x, y, h):
|
|
k1 = f(x[-1], y[-1])
|
|
k2 = f(x[-1] + 0.5*h, y[-1] + 0.5*k1)
|
|
k3 = f(x[-1] + 0.5*h, y[-1] + 0.5*h*k2)
|
|
k4 = f(x[-1] + h, y[-1] + h*k3)
|
|
|
|
return (k1 + 2*k2 + 2*k3 + k4)/6
|
|
|
|
return phi_rk4
|
|
|
|
|
|
# Multi Step
|
|
#-----------
|
|
def mk_phi_AB3(f, phi_short):
|
|
steps = 3
|
|
def phi_AB3(x, y, h):
|
|
if len(x) <= steps:
|
|
return phi_short(x, y, h)
|
|
else:
|
|
return ( 23*f(x[-1], y[-1]) - 16*f(x[-2], y[-2]) + 5*f(x[-3], y[-3]) )/12
|
|
|
|
return phi_AB3
|
|
|
|
def mk_phi_AB4(f, phi_short):
|
|
steps = 4
|
|
def phi_AB4(x, y, h):
|
|
if len(x) <= steps:
|
|
return phi_short(x, y, h)
|
|
else:
|
|
return ( 55*f(x[-1], y[-1]) - 59*f(x[-2], y[-2]) + 37*f(x[-3], y[-3]) -9*f(x[-4],y[-4]))/24
|
|
|
|
return phi_AB4
|
|
|
|
# Integrator #
|
|
##############
|
|
def integrator(x, y_0, phi):
|
|
x = np.asarray(x)
|
|
if isinstance(y_0, (list,np.ndarray)):
|
|
y_0 = np.asarray(y_0)
|
|
else:
|
|
y_0 = np.array([ y_0 ])
|
|
|
|
N = len(x)
|
|
M = len(y_0)
|
|
|
|
y = np.zeros((N,M), dtype=np.float64)
|
|
y[0] = y_0
|
|
|
|
h = x[1]-x[0]
|
|
for i in range(1,N):
|
|
y[i] = y[i-1] + h*phi(x[:i], y[:i], h)
|
|
|
|
return y
|
|
|
|
|
|
# Math Test Functions #
|
|
#######################
|
|
def ODEF(x, y):
|
|
return y - x**2 + 1
|
|
|
|
def ODEF_sol(x):
|
|
return (x+1)**2 - 0.5*np.exp(x)
|
|
|
|
|
|
# Testing #
|
|
###########
|
|
def test_integrator_singlestep():
|
|
|
|
test_func = ODEF
|
|
exact_sol = ODEF_sol
|
|
|
|
N = 2e1
|
|
M = 1
|
|
|
|
# Domain on which to do the integration
|
|
x = np.linspace(0, 1, N)
|
|
|
|
# Generate Initial Value
|
|
y_0 = 0.5
|
|
|
|
schemes = [
|
|
["Euler", mk_phi_euler],
|
|
["Collatz", mk_phi_euler_mod],
|
|
["Heun", mk_phi_heun],
|
|
["RK4", mk_phi_rk4],
|
|
]
|
|
|
|
# Show Plot
|
|
from matplotlib import pyplot
|
|
pyplot.subplots()
|
|
|
|
|
|
for name, func in schemes:
|
|
pyplot.plot(x, integrator(x, y_0, func(test_func)), '--o', label=name)
|
|
|
|
pyplot.plot(x, exact_sol(x), '-', label="Exact Solution")
|
|
pyplot.xlabel("x")
|
|
pyplot.ylabel("y")
|
|
|
|
pyplot.legend()
|
|
|
|
pyplot.show()
|
|
|
|
def test_integrator_multistep():
|
|
test_func = ODEF
|
|
exact_sol = ODEF_sol
|
|
|
|
N = 2e1
|
|
M = 1
|
|
|
|
# Domain on which to do the integration
|
|
x = np.linspace(0, 1, N)
|
|
|
|
# Calculate Initial Values using RK4
|
|
y_0 = 0.5
|
|
schemes = [
|
|
["AB3", mk_phi_AB3],
|
|
["AB4", mk_phi_AB4],
|
|
]
|
|
|
|
# Show Plot
|
|
from matplotlib import pyplot
|
|
pyplot.subplots()
|
|
|
|
for name, func in schemes:
|
|
pyplot.plot(x, integrator(x, y_0, func(test_func, mk_phi_rk4(test_func))), '--o', label=name)
|
|
|
|
pyplot.plot(x, exact_sol(x), '-', label="Exact Solution")
|
|
pyplot.xlabel("x")
|
|
pyplot.ylabel("y")
|
|
|
|
pyplot.legend()
|
|
|
|
pyplot.show()
|
|
|
|
if __name__ == "__main__":
|
|
np.random.seed(0)
|
|
#test_integrator_singlestep()
|
|
test_integrator_multistep()
|