非理想抛球运动轨迹的四阶龙格库塔数值解法

考虑空气阻力的抛球运动轨迹估计问题，参考https://zhuanlan.zhihu.com/p/33947566。

import numpy as np
import matplotlib.pyplot as plt

B2dM = 6e-2

def RungeKutta4(x, y, t, func, *param):
    k1 = func(x, y, *param)
    k2 = func(x+0.5*t, y+0.5*t*k1, *param)
    k3 = func(x+0.5*t, y+0.5*t*k2, *param)
    k4 = func(x+t, y+t*k3, *param)
    return y + (k1 + 2. * k2 + 2. * k3 + k4) * t / 6.

def func_z(dz, z):
    return dz

def func_dz(z, dz, v2):
    return -9.8 -B2dM * v2 * dz

def func_x(dx, x):
    return dx

def func_dx(x, dx, v2):
    return -B2dM * v2 * dx

delta_t = 0.002
t, x, dx, z, dz = 0., 0., 1., 0., 5.

history = [[t, x, dx, z, dz]]
while z > -1e-6:
    t += delta_t
    v2 = np.hypot(dx, dz)
    x = RungeKutta4(dx, x, delta_t, func_x)
    dx = RungeKutta4(x, dx, delta_t, func_dx, v2)
    z = RungeKutta4(dz, z, delta_t, func_z)
    dz = RungeKutta4(z, dz, delta_t, func_dz, v2)
    history.append([t, x, dx, z, dz])

history = np.array(history)

plt.figure(figsize=(12,7))
plt.subplot(2,2,1)
plt.plot(history[:,0], history[:,3], label="z-RK4")
plt.plot(history[:,0], history[:,0] * history[0,4] - 0.5 * 9.8 * (history[:,0] * history[:,0]), label="z")
plt.legend()
plt.grid()
plt.xlabel("time(s)")
plt.ylabel("position(m)")
plt.subplot(2,2,2)
plt.plot(history[:,0], history[:,1], label="x-RK4")
plt.legend()
plt.grid()
plt.xlabel("time(s)")

plt.subplot(2,2,3)
plt.plot(history[:,0], history[:,4], label="dz-RK4")
plt.plot(history[:,0], history[0,4] - history[:,0] * 9.8, label="dz")
plt.legend()
plt.grid()
plt.xlabel("time(s)")
plt.subplot(2,2,4)
plt.plot(history[:,0], history[:,2], label="dx-RK4")
plt.legend()
plt.grid()
plt.xlabel("time(s)")
plt.ylabel("velocity(m)")
plt.show()

plt.figure()
plt.plot(history[:,1], history[:,3], label="xz-RK4")
plt.plot(history[:,0]*history[0,2], history[:,0] * history[0,4] - 0.5 * 9.8 * (history[:,0] * history[:,0]), label="xz")
plt.legend()
plt.grid()
plt.xlabel("x(m)")
plt.ylabel("z(m)")
plt.show()

最后完整的代码：

import numpy as np

def RungeKutta4(x, y, t, func, *param):
    k1 = func(x, y, *param)
    k2 = func(x+0.5*t, y+0.5*t*k1, *param)
    k3 = func(x+0.5*t, y+0.5*t*k2, *param)
    k4 = func(x+t, y+t*k3, *param)
    return y + (k1 + 2. * k2 + 2. * k3 + k4) * t / 6.

class Parabola:
    def __init__(self, DELTA_T=2e-3, B1dM=6e-2, B2dM=6e-2):
        self.B1dM = B1dM
        self.B2dM = B2dM
        self.DELTA_T = DELTA_T

    def forwardstep(self, t, xyz, dxyz):
        t += self.DELTA_T
        x, y, z = xyz
        dx, dy, dz = dxyz
        v2 = np.sqrt(dx*dx + dy*dy + dz*dz)
        x = RungeKutta4(dx, x, self.DELTA_T, self.func_x)
        dx = RungeKutta4(x, dx, self.DELTA_T, self.func_dx, v2)
        y = RungeKutta4(dy, y, self.DELTA_T, self.func_x)
        dy = RungeKutta4(y, dy, self.DELTA_T, self.func_dx, v2)
        z = RungeKutta4(dz, z, self.DELTA_T, self.func_z)
        dz = RungeKutta4(z, dz, self.DELTA_T, self.func_dz, v2)
        return t, (x, y, z), (dx, dy, dz)

    def predictdrop(self, xyz, dxyz, z_drop):
        t = 0.
        down = False
        while True:
            if down:
                if xyz[2] <= z_drop:
                    return t, xyz
            else:
                if dxyz[2] < 0.:
                    down = True
                    if xyz[2] < z_drop:
                        return None
            t, xyz, dxyz = self.forwardstep(t, xyz, dxyz)

    def predictFromSeq(self, ts, xs, ys, zs, z_drop=0.25):
        ti = ts[-9:-1]+0.5/120.
        x, y, z = xs[-5], ys[-5], zs[-5]
        vx = np.average(np.diff(xs[-9:]) / np.diff(ts[-9:]))
        vy = np.average(np.diff(ys[-9:]) / np.diff(ts[-9:]))
        zi = np.diff(zs[-9:]) / np.diff(ts[-9:])
        w = np.polyfit(ti, zi, 1)
        vz = w[1] + w[0] * ts[-5]

        sxyz = [x, y, z]
        sdxyz = [vx, vy, vz]
        pt, pxyz = self.predictdrop(sxyz, sdxyz, z_drop)
        return pt, pxyz

    def func_x(self, dx, x):
        return dx

    def func_dx(self, x, dx, v2):
        return -self.B1dM * dx -self.B2dM * v2 * dx
    
    def func_z(self, dz, z):
        return dz

    def func_dz(self, z, dz, v2):
        return -9.8 - self.B1dM * dz - self.B2dM * v2 * dz