python - 在 Python 中实现 Adams Bashforth Moulton 方法

Question

我目前正致力于实现 Adams Bashforth Moulton 方法来解决钟摆问题。

我目前的代码如下:

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()

我的函数 adam_moulton() 不断返回以下错误:

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)

我实现这种多步骤方法的思维过程有什么问题？

谢谢大家的帮助!

Answer 1

笛卡尔坐标中摆的函数定义，一个在重力作用下运动的粒子和约束|x|^2=L^2=const.，可以写得更短

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])

只要不包含摩擦项，质量就会抵消，不需要在方程式中考虑。

4 阶 Adams-Moulton PECE 方法的低效实现(不使用 Jacobian 和 Newton-like 步骤)是

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

并导致 nt=2500 的半径误差为 2.39e-05，其绘图接近 odeint< 的引用解.

引导过程中使用的经典RK4方法实现为

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