使用numpy阵列的粒子间电作用力

Question

我试图模拟一个粒子在经历电斥力（或引力）时，与另一个粒子相撞，这被称为卢瑟福散射。我已经成功地使用for循环和python列表模拟了（一些）粒子。但是，现在我想改用numpy数组。模型将使用以下步骤：

1. 对于所有粒子：
   1. 计算与所有其他粒子的径向距离
   2. 计算与所有其他粒子的角度
   3. 计算x方向和y方向的净力
2. 为每个粒子创建具有netto xForce和yForce的矩阵
3. 通过a=F/质量创建加速度（以及x和y分量）矩阵
4. 更新速度矩阵
5. 更新位置矩阵

我的问题是我不知道如何使用numpy数组来计算力分量。 下面是我不可运行的代码

import numpy as np
# I used this function to calculate the force while using for-loops.
def force(x1, y1, x2, x2):
    angle =  math.atan((y2 - y1)/(x2 - x1))
    dr = ((x1-x2)**2 + (y1-y2)**2)**0.5
    force = charge2 * charge2 / dr**2 
    xforce = math.cos(angle) * force 
    yforce = math.sin(angle) * force

    # The direction of force depends on relative location
    if x1 > x2 and y1<y2:
        xforce = xforce
        yforce = yforce
    elif x1< x2 and y1< y2:
        xforce = -1 * xforce
        yforce = -1 * yforce
    elif x1 > x2 and y1 > y2:
        xforce = xforce
        yforce = yforce
    else:
        xforce = -1 * xforce
        yforce = -1* yforce

    return xforce, yforce

def update(array):

    # this for loop defeats the entire use of numpy arrays
    for particle in range(len(array[0])):
        # find distance of all particles pov from 1 particle

        # find all x-forces and y-forces on that particle

        xforce = # sum of all x-forces from all particles
        yforce = # sum of all y-forces from all particles
        force_arr[0, particle] = xforce
        force_arr[1, particle] = yforce

    return force

# begin parameters
t = 0
N = 3
masses = np.ones(N)
charges = np.ones(N)
loc_arr = np.random.rand(2, N)
speed_arr = np.random.rand(2, N)
acc_arr = np.random.rand(2, N)
force = np.random.rand(2, N)


while t < 0.5:
    force_arr = update(loc_arry)
    acc_arr = force_arr / masses
    speed_arr += acc_array
    loc_arr += speed_arr
    t += dt

    # plot animation

Answer 1

用阵列对此问题建模的一种方法可能是：

• 将点坐标定义为Nx2数组。（如果以后升级到三维点，这将有助于扩展性）
• 将中间变量distance、angle、force定义为NxN数组以表示成对交互

需要了解的重要事项：

• 如果数组具有相同的形状（或一致的形状，这是一个非常重要的主题…），则可以对数组调用大多数数值函数
• meshgrid帮助您生成调整Nx2数组形状以计算NxN结果所需的数组索引
• 切向音符（ha）arctan2()计算有符号角度，因此可以绕过复杂的“哪个象限”逻辑

例如，你可以做这样的事情。注意get_dist和get_angle中点之间的算术运算发生在最底部的维度中：

import numpy as np

# 2-D locations of particles
points = np.array([[1,0],[2,1],[2,2]])
N = len(points)  # 3

def get_dist(p1, p2):
    r = p2 - p1
    return np.sqrt(np.sum(r*r, axis=2))

def get_angle(p1, p2):
    r = p2 - p1
    return np.arctan2(r[:,:,1], r[:,:,0])

ii = np.arange(N)
ix, iy = np.meshgrid(ii, ii)

dist = get_dist(points[ix], points[iy])
angle = get_angle(points[ix], points[iy])
# ... compute force
# ... apply the force, etc.

对于上面显示的示例3点向量：

In [246]: dist
Out[246]: 
array([[0.        , 1.41421356, 2.23606798],
       [1.41421356, 0.        , 1.        ],
       [2.23606798, 1.        , 0.        ]])

In [247]: angle / np.pi     # divide by Pi to make the numbers recognizable
Out[247]: 
array([[ 0.        , -0.75      , -0.64758362],
       [ 0.25      ,  0.        , -0.5       ],
       [ 0.35241638,  0.5       ,  0.        ]])

Answer 2

这里有一个围棋，每个时间步只有一个循环，它应该适用于任意数量的维度，我也用3个维度进行了测试：

from matplotlib import pyplot as plt
import numpy as np

fig, ax = plt.subplots()

N = 4
ndim = 2
masses = np.ones(N)
charges = np.array([-1, 1, -1, 1]) * 2
# loc_arr = np.random.rand(N, ndim)
loc_arr = np.array(((-1,0), (1,0), (0,-1), (0,1)), dtype=float)
speed_arr = np.zeros((N, ndim))

# compute charge matrix, ie c1 * c2
charge_matrix = -1 * np.outer(charges, charges)

time = np.linspace(0, 0.5)
dt = np.ediff1d(time).mean()

for i, t in enumerate(time):
    # get (dx, dy) for every point
    delta = (loc_arr.T[..., np.newaxis] - loc_arr.T[:, np.newaxis]).T
    # calculate Euclidean distance
    distances = np.linalg.norm(delta, axis=-1)
    # and normalised unit vector
    unit_vector = (delta.T / distances).T
    unit_vector[np.isnan(unit_vector)] = 0 # replace NaN values with 0

    # calculate force
    force = charge_matrix / distances**2 # norm gives length of delta vector
    force[np.isinf(force)] = 0 # NaN forces are 0

    # calculate acceleration in all dimensions
    acc = (unit_vector.T * force / masses).T.sum(axis=1)
    # v = a * dt
    speed_arr += acc * dt

    # increment position, xyz = v * dt
    loc_arr += speed_arr * dt 

    # plotting
    if not i:
        color = 'k'
        zorder = 3
        ms = 3
        for i, pt in enumerate(loc_arr):
            ax.text(*pt + 0.1, s='{}q {}m'.format(charges[i], masses[i]))
    elif i == len(time)-1:
        color = 'b'
        zroder = 3
        ms = 3
    else:
        color = 'r'
        zorder = 1
        ms = 1
    ax.plot(loc_arr[:,0], loc_arr[:,1], '.', color=color, ms=ms, zorder=zorder)

ax.set_aspect('equal')

上面的示例生成，其中黑色和蓝色点分别表示开始和结束位置：

当电荷相等时charges = np.ones(N) * 2系统对称性得以保持，电荷相互排斥：

最后是一些随机初始速度speed_arr = np.random.rand(N, 2)：

编辑

对上面的代码做了一个小改动，以确保它是正确的。（我在合力上遗漏了-1，即+/+之间的力应该是负数，我总结了错误的轴，对此表示歉意。现在在masses[0] = 5的情况下，系统正确发展：