pythonmathnumerical-methods
Implementing the Adams Bashforth Moulton Method in Python
I am currently working on implementing the Adams Bashforth Moulton Method for solving a pendulum problem.
My current code is as follows:
import numpy as np
#####################################################################################
# DEF func
#####################################################################################
def func(y,t):
n=y.size
f=np.zeros(6)
xp=np.array([y[3],y[4],y[5]])[np.newaxis]
mass=1.
F1=0.
F2=0.
F3=-mass*9.8
F=np.array([F1,F2,F3])[np.newaxis]
phix=2.*y[0]
phiy=2.*y[1]
phiz=2.*y[2]
G=np.array([phix,phiy,phiz])[np.newaxis]
H=2.*np.eye(3)
lambd=(mass*np.dot(xp,np.dot(H,xp.T))+np.dot(F,G.T))/np.dot(G,G.T)
f[0]=y[3]
f[1]=y[4]
f[2]=y[5]
for k in range(0,3):
f[k+3]=(F[0,k]-lambd*G[0,k])/mass
return f
def dF_matrix(y):
n=y.size
dF=np.zeros((6,6))
xp=np.array([y[3],y[4],y[5]])[np.newaxis]
mass=1.
F1=0.
F2=0.
F3=-mass*9.8
F=np.array([F1,F2,F3])[np.newaxis]
phix=2.*y[0]
phiy=2.*y[1]
phiz=2.*y[2]
G=np.array([phix,phiy,phiz])[np.newaxis]
H=2.*np.eye(3)
lambd=(mass*np.dot(xp,np.dot(H,xp.T))+np.dot(F,G.T))/np.dot(G,G.T)
dF[0,3]=1
dF[1,4]=1
dF[2,5]=1
dF[3,0]=(y[0]*F1+2*lambd)/mass
dF[3,1]=(y[0]*F2)/mass
dF[3,2]=(y[0]*F3)/mass
dF[3,3]=phix*y[3]
dF[3,4]=phix*y[4]
dF[3,5]=phix*y[5]
dF[4,0]=(y[1]*F1)/mass
dF[4,1]=(y[1]*F2+2*lambd)/mass
dF[4,2]=(y[1]*F3)/mass
dF[4,3]=phiy*y[3]
dF[4,4]=phiy*y[4]
dF[4,5]=phiy*y[5]
dF[5,0]=(y[2]*F1)/mass
dF[5,1]=(y[2]*F2)/mass
dF[5,2]=(y[2]*F3+2*lambd)/mass
dF[5,3]=phiz*y[3]
dF[5,4]=phiz*y[4]
dF[5,5]=phiz*y[5]
return dF
#####################################################################################
# PYTHON script:
# Solve ODE for pendulum problem
#####################################################################################
from scipy.integrate import odeint
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
import time
### Forward Euler Method
def forward_Euler(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0, :] = y_matrix
for i in range(len(time) - 1):
dt = time[i + 1] - time[i]
y[i + 1, :] = y[i, :] + np.asarray(function(y[i, :], time[i])) * dt
return y
### Modified Euler Method
def modified_Euler(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0, :] = y_matrix
for i in range(len(time) - 1):
dt = time[i + 1] - time[i]
k1 = np.asarray(function(y[i, :], time[i]))
k2 = np.asarray(function(y[i, : + k1], time[i]))
y[i + 1, :] = y[i, :] + (k1 + k2) / 2
return y
### Adams-Bashforth 2nd order
def Adams_Bash_2nd(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0, :] = y_matrix
for i in range(len(time) - 1):
dt = time[i + 1] - time[i]
f_1 = function(y[i, :], time[i])
f_2 = function(f_1, time[i])
y[i + 1, :] = y[i, :] + dt * f_1 + ((dt**2)/2) * f_2
return y
### Adams Bashforth Moulton Method
def Adams_Moulton(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0, :] = y_matrix
### predictor formula
for i in range(len(time) - 1):
dt = time[i + 1] - time[i]
f_1 = function(y[i, :], time[i])
f_2 = function(f_1, time[i])
y[i + 1, :] = y[i, :] + dt * f_1 + ((dt**2)/2) * f_2
### Corrector formula
for i in range(len(time-1)):
dt = time[i + 1] - time[i]
k_1 = 9 * (function(y[i, :], time[i+1]))
k_2 = 19 * (function(y[i, :], time[i]))
k_3 = 5 * (function(y[i, :], time[i-1]))
k_4 = (function(y[i, :], time[i-2]))
# === ERROR HAPPENS HERE ===
y[i + 1, :] = y + (dt/24) * (k_1 + k_2 - k_3 + k_4)
return y
# initial condition
y0=np.array([0.0,1.0,0.0,0.8,0.0,1.2])
print('ODE Python ..')
time_start = time.clock()
nt = 500
# time points
t = np.linspace(0,25,nt)
# solve ODE
y1 = odeint(func,y0,t)
time_elapsed = (time.clock() - time_start)
print('elapsed time',time_elapsed)
# compute residual:
r=y1[:,0]**2+y1[:,1]**2+y1[:,2]**2-1
rmax1=np.max(np.abs(r))
print('error',rmax1)
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.plot3D(y1[:,0],y1[:,1],y1[:,2], 'gray')
plt.show()
# print ODE Euler Method
time_start = time.clock()
nt = 500
# time points
t = np.linspace(0, 25, nt)
# solve ODE
y1 = forward_Euler(func, y0, t)
print(y1)
time_elapsed = (time.clock() - time_start)
print('elapsed time', time_elapsed)
# compute residual:
r = y1[:, 0] ** 2 + y1[:, 1] ** 2 + y1[:, 2] ** 2 - 1
rmax1 = np.max(np.abs(r))
print('error', rmax1)
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.plot3D(y1[:, 0], y1[:, 1], y1[:, 2], 'gray')
plt.show()
### Modified Euler method
time_start = time.clock()
nt = 500
# time points
t = np.linspace(0, 25, nt)
# solve ODE
y1 = forward_Euler(func, y0, t)
print(y1)
time_elapsed = (time.clock() - time_start)
print('elapsed time', time_elapsed)
# compute residual:
r = y1[:, 0] ** 2 + y1[:, 1] ** 2 + y1[:, 2] ** 2 - 1
rmax1 = np.max(np.abs(r))
print('error', rmax1)
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.plot3D(y1[:, 0], y1[:, 1], y1[:, 2], 'gray')
plt.show()
### Adams-Bashforth 2nd order
time_start = time.clock()
nt = 500
# time points
t = np.linspace(0, 25, nt)
# solve ODE
y1 = Adams_Bash_2nd(func, y0, t)
print(y1)
time_elapsed = (time.clock() - time_start)
print('elapsed time', time_elapsed)
# compute residual:
r = y1[:, 0] ** 2 + y1[:, 1] ** 2 + y1[:, 2] ** 2 - 1
rmax1 = np.max(np.abs(r))
print('error', rmax1)
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.plot3D(y1[:, 0], y1[:, 1], y1[:, 2], 'gray')
plt.show()
### Adams-Moulton 1st order
# time_start = time.clock()
nt = 500
# time points
t = np.linspace(0, 25, nt)
# solve ODE
y1 = Adams_Moulton(func, y0, t)
print(y1)
# time_elapsed = (time.clock() - time_start)
# print('elapsed time', time_elapsed)
# compute residual:
r = y1[:, 0] ** 2 + y1[:, 1] ** 2 + y1[:, 2] ** 2 - 1
rmax1 = np.max(np.abs(r))
print('error', rmax1)
fig = plt.figure()
ax = plt.axes(projection='3d')
ax.plot3D(y1[:, 0], y1[:, 1], y1[:, 2], 'gray')
plt.show()
exit()
My fucntion adam_moulton() keeps returning the following error:
Traceback (most recent call last):
File "C:/Users/Andrew/Documents/solve_pendulum.py", line 249, in <module>
y1 = Adams_Moulton(func, y0, t)
File "C:/Users/Andrew/Documents/solve_pendulum.py", line 148, in Adams_Moulton
y[i + 1, :] = y + (dt/24) * (k_1 + k_2 - k_3 + k_4)
ValueError: could not broadcast input array from shape (500,6) into shape (6)
What is wrong in my thought process for implementing this multistep method?
Thank you all for the help!
Solution
The function definition for the pendulum in Cartesian coordinates, a particle moving under gravity and the constraint |x|^2=L^2=const., can be written much shorter as
def func(y,t):
x,v = y[:3],y[3:6]
grav = np.array([0., 0., -9.8 ])
lambd = (grav.dot(x)+v.dot(v))/x.dot(x)
return np.concatenate([v, grav - lambd*x] )
# initial condition
y0=np.array([0.0,1.0,0.0,0.8,0.0,1.2])
As long as no friction term is included, the mass cancels out and need not be considered in the equations.
An inefficient implementation (no use of Jacobian and Newton-like steps) of the 4th order Adams-Moulton PECE method is
def Adams_Moulton_4th(function, y_matrix, time):
y = np.zeros((np.size(time), np.size(y_matrix)))
y[0] = y_matrix
### bootstrap steps with 4th order one-step method
dt = time[1] - time[0]
y[1] = RK4_step(function,y[0], time[0], dt, N=4)
y[2] = RK4_step(function,y[1], time[1], dt, N=4)
y[3] = RK4_step(function,y[2], time[2], dt, N=4)
f_m2 = function(y[0], time[0])
f_m1 = function(y[1], time[1])
f_0 = function(y[2], time[2])
f_1 = function(y[3], time[3])
for i in range(3,len(time) - 1):
### first shift the "virtual" function value array so that
### [f_m3, f_m2, f_m1, f_0] corresponds to [ f[i-3], f[i-2], f[i-1], f[i] ]
f_m3, f_m2, f_m1, f_0 = f_m2, f_m1, f_0, f_1
### predictor formula 4th order [ 55/24, -59/24, 37/24, -3/8 ]
y[i+1] = y[i] + (dt/24) * (55*f_0 - 59*f_m1 + 37*f_m2 - 9*f_m3)
f_1 = function(y[i+1], time[i+1])
### Corrector formula 4th order [ 3/8, 19/24, -5/24, 1/24 ]
y[i+1] = y[i] + (dt/24) * (9*f_1 + 19*f_0 - 5*f_m1 + f_m2)
f_1 = function(y[i+1], time[i+1])
return y
and results in with nt=2500 steps in a radius error of 2.39e-05 with a plot that is close to the reference solution by odeint.
The classical RK4 method used in the bootstrapping process was implemented as
def RK4_step(f,y,t,dt, N=1):
dt /= N;
for k in range(N):
k1=f(y,t)*dt; k2=f(y+k1/2,t+dt/2)*dt; k3=f(y+k2/2,t+dt/2)*dt; k4=f(y+k3,t+dt)*dt;
y, t = y+(k1+2*(k2+k3)+k4)/6, t+dt
return y