速度-Verlet双阱算法Python

import numpy as np
import matplotlib.pyplot as plt

def harmonic_potential(x, k=1.0):
    """谐振子的势能函数"""
    return 0.5 * k * x**2

def velocity_verlet(positions, velocities, accelerations, dt, k=1.0):
    """速度-Verlet算法"""
    num_particles = len(positions)
    new_positions = np.zeros(num_particles)
    new_velocities = np.zeros(num_particles)
    
    for i in range(num_particles):
        # 更新位置
        new_positions[i] = positions[i] + velocities[i] * dt + 0.5 * accelerations[i] * dt**2
        
        # 计算新的加速度（这里假设势能只依赖于位置）
        new_accelerations = -k * new_positions
        
        # 更新速度
        new_velocities[i] = velocities[i] + 0.5 * (accelerations[i] + new_accelerations[i]) * dt
        
        # 更新加速度列表
        accelerations[i] = new_accelerations[i]
    
    return new_positions, new_velocities, accelerations

# 初始化位置、速度和加速度
num_particles = 1
positions = np.array([0.0])
velocities = np.array([1.0])
accelerations = -harmonic_potential(positions)

# 设置时间步长和时间范围
dt = 0.01
total_time = 10.0
num_steps = int(total_time / dt)

# 存储位置和速度的历史记录
position_history = [positions.copy()]
velocity_history = [velocities.copy()]

# 运行速度-Verlet算法
for _ in range(num_steps):
    positions, velocities, accelerations = velocity_verlet(positions, velocities, accelerations, dt)
    position_history.append(positions.copy())
    velocity_history.append(velocities.copy())

# 将历史记录转换为numpy数组
position_history = np.array(position_history)
velocity_history = np.array(velocity_history)

# 绘制位置和速度随时间的变化
plt.figure(figsize=(12, 6))

plt.subplot(1, 2, 1)
plt.plot(np.linspace(0, total_time, num_steps + 1), position_history[:, 0], label='Position')
plt.xlabel('Time')
plt.ylabel('Position')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(np.linspace(0, total_time, num_steps + 1), velocity_history[:, 0], label='Velocity')
plt.xlabel('Time')
plt.ylabel('Velocity')
plt.legend()

plt.tight_layout()
plt.show()