Python函数添加时间循环填充数组实现抛体运动积分计算

我想用Euler、Euler-Cromer、Midpoint三种积分方法绘制抛体运动，需要填充数组记录物体在连续时间戳下的x坐标。目前函数能计算数值，但要么只返回首个时间戳的x坐标，要么返回重复值，还会出现轴长度不匹配的错误。我需要解决时间循环填充数组的问题，最终实现三种方法的结果对比。以下是我的代码：

import math
import numpy as np

def projectile_motion(v0, theta, tau, method, air_resistance=True):

    # Convert launch angle to radians
    theta = math.radians(theta)

    # Initial conditions
    x = [np.empty((1,100))]
    y = h0
    vx = v0 * math.cos(theta)
    vy = v0 * math.sin(theta)

    for i in range(len(x)):#filling in array
        j=i+1
        x[j]=x[i]+x[i]*tau

    # Integration method
    if method == 'Euler':
        def update():
            nonlocal x, y, vx, vy
            if air_resistance:
                v = math.sqrt(vx**2 + vy**2)
                ax = -(Cd*rho*A*v*vx)/(2*m)
                ay = -g - (Cd*rho*A*v*vy)/(2*m)
                vx += ax * tau
                vy += ay * tau
            else:
                vy -= g * tau
            x += vx * tau
            y += vy * tau
    elif method == 'Euler-Cromer':
        def update():
            nonlocal x, y, vx, vy
            if air_resistance:
                v = math.sqrt(vx**2 + vy**2)
                ax = -(Cd*rho*A*v*vx)/(2*m)
                ay = -g - (Cd*rho*A*v*vy)/(2*m)
                vx += ax * tau
                vy += ay * tau
            else:
                vy -= g * tau
            x += vx * tau
            y += vy * tau
            if y < 0:
                vy = -vy
    elif method == 'Midpoint':
        def update():
            nonlocal x, y, vx, vy
            if air_resistance:
                v = math.sqrt(vx**2 + vy**2)
                ax = -(Cd*rho*A*v*vx)/(2*m)
                ay = -g - (Cd*rho*A*v*vy)/(2*m)
                vx_half = vx + ax * tau/2
                vy_half = vy + ay * tau/2
                v_half = math.sqrt(vx_half**2 + vy_half**2)
                ax_half = -(Cd*rho*A*v_half*vx_half)/(2*m)
                ay_half = -g - (Cd*rho*A*v_half*vy_half)/(2*m)
                vx += ax_half * tau
                vy += ay_half * tau
            else:
                vy -= g * tau
            x += vx * tau
            y += vy * tau
            if y < 0:
                 vy = -vy

    # Time loop
    while y >= 0:
        update()

    return x 

问题分析 #

1. 初始数组定义错误：x = [np.empty((1,100))]将二维空数组嵌套进列表，后续索引操作、赋值逻辑完全不符合需求，且未设置合理的初始x坐标。
2. 无效预填充循环：for i in range(len(x))仅执行1次（列表长度为1），且x[j]会触发索引越界错误。
3. 坐标更新逻辑混乱：x += vx * tau在混合列表与numpy数组的情况下会导致轴不匹配，且未将新坐标记录到数组/列表中。
4. 缺少物理常量定义：代码中h0, Cd, rho, A, m, g未定义，直接运行会报错。
5. 落地处理易引发无限循环：反弹时仅反转vy，未限制反弹次数或添加能量损失，可能导致程序死循环。

修复后的代码 #

import math
import numpy as np
import matplotlib.pyplot as plt

# 物理常量（可根据需求修改，或改为函数参数）
g = 9.81       # 重力加速度 (m/s²)
h0 = 0.0       # 初始高度 (m)
Cd = 0.47      # 拖曳系数
rho = 1.225    # 空气密度 (kg/m³)
A = 0.07       # 物体横截面积 (m²)
m = 0.15       # 物体质量 (kg)

def projectile_motion(v0, theta, tau, method, air_resistance=True):
    # 角度转弧度
    theta = math.radians(theta)

    # 初始化：用列表记录每一步的坐标
    x_history = [0.0]
    y_history = [h0]
    vx = v0 * math.cos(theta)
    vy = v0 * math.sin(theta)

    # Euler方法更新逻辑
    def update_euler():
        nonlocal vx, vy
        if air_resistance:
            v = math.sqrt(vx**2 + vy**2)
            ax = -(Cd * rho * A * v * vx) / (2 * m)
            ay = -g - (Cd * rho * A * v * vy) / (2 * m)
            vx += ax * tau
            vy += ay * tau
        else:
            vy -= g * tau
        # 计算新坐标
        new_x = x_history[-1] + vx * tau
        new_y = y_history[-1] + vy * tau
        return new_x, new_y

    # Euler-Cromer方法更新逻辑
    def update_euler_cromer():
        nonlocal vx, vy
        if air_resistance:
            v = math.sqrt(vx**2 + vy**2)
            ax = -(Cd * rho * A * v * vx) / (2 * m)
            ay = -g - (Cd * rho * A * v * vy) / (2 * m)
            vx += ax * tau
            vy += ay * tau
        else:
            vy -= g * tau
        new_x = x_history[-1] + vx * tau
        new_y = y_history[-1] + vy * tau
        # 落地处理：加入能量损失避免无限反弹
        if new_y < 0:
            vy = -vy * 0.8
            new_y = 0.0
        return new_x, new_y

    # Midpoint方法更新逻辑
    def update_midpoint():
        nonlocal vx, vy
        if air_resistance:
            # 当前时刻加速度
            v = math.sqrt(vx**2 + vy**2)
            ax = -(Cd * rho * A * v * vx) / (2 * m)
            ay = -g - (Cd * rho * A * v * vy) / (2 * m)
            # 半步速度
            vx_half = vx + ax * tau / 2
            vy_half = vy + ay * tau / 2
            # 半步时刻加速度
            v_half = math.sqrt(vx_half**2 + vy_half**2)
            ax_half = -(Cd * rho * A * v_half * vx_half) / (2 * m)
            ay_half = -g - (Cd * rho * A * v_half * vy_half) / (2 * m)
            # 更新速度
            vx += ax_half * tau
            vy += ay_half * tau
        else:
            vy -= g * tau
        new_x = x_history[-1] + vx * tau
        new_y = y_history[-1] + vy * tau
        # 落地处理
        if new_y < 0:
            vy = -vy * 0.8
            new_y = 0.0
        return new_x, new_y

    # 绑定对应方法的更新函数
    if method == 'Euler':
        update = update_euler
    elif method == 'Euler-Cromer':
        update = update_euler_cromer
    elif method == 'Midpoint':
        update = update_midpoint
    else:
        raise ValueError("无效方法，可选：'Euler'、'Euler-Cromer'、'Midpoint'")

    # 时间循环：限制最大反弹次数防止死循环
    bounce_count = 0
    max_bounces = 3
    while bounce_count <= max_bounces:
        new_x, new_y = update()
        x_history.append(new_x)
        y_history.append(new_y)
        if new_y == 0.0:
            bounce_count += 1

    # 返回numpy数组方便后续处理
    return np.array(x_history), np.array(y_history)

# 示例：三种方法对比绘图
if __name__ == "__main__":
    v0 = 20    # 初速度 (m/s)
    theta = 45 # 发射角度 (°)
    tau = 0.01 # 时间步长 (s)

    # 获取三种方法的结果
    x_euler, y_euler = projectile_motion(v0, theta, tau, 'Euler')
    x_ec, y_ec = projectile_motion(v0, theta, tau, 'Euler-Cromer')
    x_mid, y_mid = projectile_motion(v0, theta, tau, 'Midpoint')

    # 绘制对比图
    plt.figure(figsize=(10, 6))
    plt.plot(x_euler, y_euler, label='Euler', linestyle='--', alpha=0.8)
    plt.plot(x_ec, y_ec, label='Euler-Cromer', linestyle='-.', alpha=0.8)
    plt.plot(x_mid, y_mid, label='Midpoint', linestyle='-', alpha=0.8)
    plt.xlabel('X坐标 (m)')
    plt.ylabel('Y坐标 (m)')
    plt.title('三种数值积分方法的抛体运动轨迹对比')
    plt.legend()
    plt.grid(True)
    plt.show()

修复关键点 #

• 坐标记录方式修正：用列表x_history和y_history动态记录每一步的坐标，避免数组预分配的长度问题。
• 移除无效循环：删除原代码中错误的预填充循环，改为在时间循环中逐步添加新坐标。
• 拆分更新逻辑：将每种积分方法的更新逻辑拆分为独立函数，清晰分离不同方法的差异，避免嵌套函数的混乱。
• 优化落地处理：加入反弹能量损失并限制最大反弹次数，防止程序无限循环。
• 补充常量定义：添加必要的物理常量，确保代码可直接运行。
• 返回值优化：返回numpy数组，方便后续绘图和数值分析。

内容的提问来源于stack exchange，提问作者Kiera