三粒子系统梯度下降与能量计算的代码溢出问题排查

问题背景 #

需要实现三粒子系统的梯度下降算法并计算能量，系统能量公式（物理常数设为1）为：
$$
E(\vec{x}_1, \vec{x}2, \dots, \vec{x}N)=\overbrace{\sum_i ||\vec{x}i||^2}{\text{Harmonic potential}}+\overbrace{\sum{i=1}^N \sum{j=i+1}^{N\frac{1}{||\vec{x}_i-\vec{x}_j||}}}{\text{Coulomb repulsion}}
$$
其中$\vec{x}i$为粒子i的位置向量，$\sum{i=1}^N \sum{j=i+1}^{N$为所有不同粒子对的求和，$||\vec{x}||=\sqrt{x}2+y^2}$为欧几里得范数。

当前代码存在数值溢出问题，推测问题出在gradient_descent函数，且绘图仅显示3个点，怀疑是迭代早期溢出导致粒子位置异常。

原代码 #

import numpy as np
import matplotlib.pyplot as plt

def gradient_energy(configuration):
    N = configuration.shape[0]
    gradient = np.zeros_like(configuration, dtype=np.float64)

    h_pot_gradient = 2 * configuration
    gradient -= h_pot_gradient

    for i in range(N):
        for j in range(i + 1, N):
            r_ij = configuration[i] - configuration[j]
            d = np.linalg.norm(r_ij)
            c_rep_gradient = r_ij / (d ** 3)
            gradient[i] += c_rep_gradient
            gradient[j] -= c_rep_gradient

    return gradient

def gradient_descent(configuration, gamma, T, epsilon):
    for i in range(T):
        gradient = gradient_energy(configuration).astype(np.float64) #again, there was a problem with the types, so the solution ChatGPT gave is the .astype(np.flpat64) part
        frobenius_norm = np.linalg.norm(gradient)
        if frobenius_norm < epsilon:
            break
        configuration -= gamma * gradient
    return configuration

def get_energy(configuration):
    N = configuration.shape[0]
    h_pot = np.sum(np.linalg.norm(configuration, axis=1) ** 2)
    c_rep = 0.0
    for i in range(N):
        for j in range(i + 1, N):
            d = np.linalg.norm(configuration[i] - configuration[j])
            c_rep += 1 / d

    E = h_pot + c_rep
    return E

gamma = 0.1
T = 1000
epsilon = 1e-6
initial_config = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]], dtype=np.float64)
final_config = gradient_descent(initial_config, gamma, T, epsilon)

energy_final = get_energy(final_config)
print("Energy of the obtained configuration is ", energy_final)

plt.figure(figsize=(5, 5))
plt.scatter(final_config[:, 0], final_config[:, 1], s=50)
plt.gca().set_aspect('equal', adjustable='box')
plt.title("Configuration after Gradient Descent")
plt.xlabel("x")
plt.ylabel("y")
plt.show()

问题分析 #

1. 梯度符号完全错误：
   能量梯度的计算方向颠倒，导致更新步骤变成梯度上升而非下降。谐波势的梯度应为2 * configuration，库仑斥力对粒子i的梯度是-(x_i - x_j)/||x_i-x_j||³，原代码符号完全相反，粒子位置会失控发散。
2. 学习率过大：
   gamma=0.1的步长太大，迭代中粒子会快速靠近，触发1/d和1/d³的数值爆炸。
3. 未处理极小距离：
   粒子间距趋近于0时，库仑项数值趋于无穷大，无保护机制导致溢出。
4. 原地修改输入数组：
   gradient_descent直接修改输入的配置数组，会污染初始数据。

修复后的代码 #

import numpy as np
import matplotlib.pyplot as plt

def gradient_energy(configuration):
    N = configuration.shape[0]
    gradient = np.zeros_like(configuration, dtype=np.float64)
    # 谐波势梯度：∇E_harmonic = 2x_i
    gradient += 2 * configuration
    
    min_distance = 1e-3  # 设置最小距离阈值，避免数值溢出
    for i in range(N):
        for j in range(i + 1, N):
            r_ij = configuration[i] - configuration[j]
            d = np.linalg.norm(r_ij)
            # 处理极小距离，防止除以0或极小值
            if d < min_distance:
                d = min_distance
            # 库仑斥力梯度：∇E_coulomb_i = -(x_i - x_j)/d³
            c_rep_gradient_i = -r_ij / (d ** 3)
            gradient[i] += c_rep_gradient_i
            gradient[j] -= c_rep_gradient_i

    return gradient

def gradient_descent(configuration, gamma, T, epsilon):
    # 复制输入数组，避免修改原始配置
    config = configuration.copy()
    for i in range(T):
        gradient = gradient_energy(config)
        frobenius_norm = np.linalg.norm(gradient)
        if frobenius_norm < epsilon:
            print(f"迭代{i+1}次后收敛")
            break
        # 梯度下降：沿负梯度方向更新位置
        config -= gamma * gradient
    return config

def get_energy(configuration):
    N = configuration.shape[0]
    h_pot = np.sum(np.linalg.norm(configuration, axis=1) ** 2)
    c_rep = 0.0
    min_distance = 1e-3
    for i in range(N):
        for j in range(i + 1, N):
            d = np.linalg.norm(configuration[i] - configuration[j])
            if d < min_distance:
                d = min_distance
            c_rep += 1 / d
    E = h_pot + c_rep
    return E

# 调整参数：减小学习率，增加迭代次数
gamma = 0.01
T = 10000
epsilon = 1e-6
initial_config = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]], dtype=np.float64)
final_config = gradient_descent(initial_config, gamma, T, epsilon)

energy_final = get_energy(final_config)
print("最终配置的能量为:", energy_final)
print("最终粒子位置:\n", final_config)

plt.figure(figsize=(5, 5))
plt.scatter(final_config[:, 0], final_config[:, 1], s=50)
plt.gca().set_aspect('equal', adjustable='box')
plt.title("梯度下降后的三粒子配置")
plt.xlabel("x")
plt.ylabel("y")
plt.show()

修复说明 #

1. 修正梯度符号，确保gradient_energy返回能量的真实梯度，更新步骤遵循梯度下降逻辑。
2. 添加最小距离阈值，避免粒子间距过小时的数值溢出。
3. 减小学习率至0.01，同时增加最大迭代次数，保证收敛稳定性。
4. 在gradient_descent中复制输入数组，避免污染初始配置。
5. 添加收敛提示，方便观察迭代过程。

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