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

Question

我试图模拟一个粒子在经历电排斥（或吸引力）的同时向另一个粒子飞行，称为卢瑟福散射。我已经成功地使用 for 循环和 python 列表模拟了（一些）粒子。但是，现在我想改用 numpy 数组。该模型将使用以下步骤：对于所有粒子：计算所有其他粒子的径向距离计算与所有其他粒子的角度计算 x 方向和 y 方向的净力使用 netto xForce 和 yForce 为每个粒子创建矩阵通过 a = F/mass 创建加速度（也是 x 和 y 分量）矩阵更新速度矩阵更新位置矩阵我的问题是我不知道如何使用 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, yforcedef 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 parameterst = 0N = 3masses = 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数组。（如果您稍后进入 3-D 点，这将有助于扩展性）

• 将中间变量distance, angle,定义force为NxN表示成对交互的数组

麻木的事情要知道：

• 如果数组具有相同的形状（或一致的形状，这是一个重要的话题......），您可以在数组上调用大多数数值函数

• meshgrid帮助您生成对数组进行变形Nx2以计算NxN结果所需的数组索引

• 和一个切线音符（哈哈）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，系统正确发展：

Answer 3

经典的方法是计算系统中所有粒子的电场。假设您有 3 个带正电荷的带电粒子：

particles = np.array([[1,0,0],[2,1,0],[2,2,0]]) # location of each particle

q = np.array([1,1,1]) # charge of each particle

计算每个粒子位置的电场的最简单方法是 for 循环：

def for_method(pos,q):

    """Computes electric field vectors for all particles using for-loop."""

    Evect = np.zeros( (len(pos),len(pos[0])) ) # define output electric field vector

    k =  1 / (4 * np.pi * const.epsilon_0) * np.ones((len(pos),len(pos[0]))) * 1.602e-19 # make this into matrix as matrix addition is faster

    # alternatively you can get rid of np.ones and just define this as a number

    

    for i, v0 in enumerate(pos): # s_p - selected particle | iterate over all particles | v0 reference particle

        for v, qc in zip(pos,q): # loop over all particles and calculate electric force sum | v particle being calculated for

            if all((v0 == v)):   # do not compute for the same particle

                continue

            else:

                r = v0 - v       #

                Evect[i] += r / np.linalg.norm(r) ** 3 * qc #! multiply by charge

    return Evect * k

# to find electric field at each particle`s location call

for_method(particles, q)

此函数返回与输入粒子数组具有相同形状的向量数组。要找到每个上的力，您只需将此向量乘以q电荷数组。从那里开始，您可以使用您最喜欢的 ODE 求解器轻松找到您的加速并集成系统。

性能优化和准确性

因为方法是最慢的方法。可以仅使用线性代数来计算该场，从而显着提高速度。以下代码对这个问题非常有效的 Numpy 矩阵“单线”（几乎是单线）：

def CPU_matrix_method(pos,q):

    """Classic vectorization of for Coulomb law using numpy arrays."""

    k = 1 / (4 * np.pi * const.epsilon_0) * np.ones((len(pos),3)) * 1.602e-19 # define electric constant

    dist = distance.cdist(pos,pos)  # compute distances

    return k * np.sum( (( np.tile(pos,len(pos)).reshape((len(pos),len(pos),3)) - np.tile(pos,(len(pos),1,1))) * q.reshape(len(q),1)).T * np.power(dist,-3, where = dist != 0),axis = 1).T

请注意，此代码和以下代码还返回每个粒子的电场矢量。

如果您使用 Cupy 库将其卸载到 GPU 上，您可以获得更高的性能。以下代码几乎与CPU_matrix_method相同，我只是稍微扩展了单行代码，以便您可以更好地看到发生了什么：

def GPU_matrix_method(pos,q):

    """GPU Coulomb law vectorization.

    Takes in numpy arrays, performs computations and returns cupy array"""

    # compute distance matrix between each particle

    k_cp = 1 / (4 * cp.pi * const.epsilon_0) * cp.ones((len(pos),3)) * 1.602e-19 # define electric constant, runs faster if this is matrix

    dist = cp.array(distance.cdist(pos,pos)) # could speed this up with cupy cdist function! use this: cupyx.scipy.spatial.distance.cdist

    pos, q = cp.array(pos), cp.array(q) # load inputs to GPU memory

    dist_mod = cp.power(dist,-3)        # compute inverse cube of distance

    dist_mod[dist_mod == cp.inf] = 0    # set all infinity entries to 0 (i.e. diagonal elements/ same particle-particle pairs)

    # compute by magic

    return k_cp * cp.sum((( cp.tile(pos,len(pos)).reshape((len(pos),len(pos),3)) - cp.tile(pos,(len(pos),1,1))) * q.reshape(len(q),1)).T * dist_mod, axis = 1).T

关于上述算法的准确性，如果你计算粒子阵列上的 3 种方法，你会得到相同的结果：

[[-6.37828367e-10 -7.66608512e-10  0.00000000e+00]

 [ 5.09048221e-10 -9.30757576e-10  0.00000000e+00]

 [ 1.28780145e-10  1.69736609e-09  0.00000000e+00]]

关于性能，我在 2 到 5000 个带电粒子的系统上计算了每种算法。此外，我还包括了 for_method 的 Numba 预编译版本，以使 for-loop 方法具有竞争力：

我们看到 for-loop 执行非常需要超过 400 秒来计算具有 5000 个粒子的系统。放大到底部：

这表明解决这个问题的矩阵方法要好几个数量级。准确地说，Numba for-loop 对 5000 个粒子的评估需要 18.5 秒，CPU 矩阵需要 4 秒（比 Numba 快 5 倍），GPU 矩阵*需要 0.8 秒（比 Numba 快 23 倍）。较大的阵列显示出显着差异。

* 使用的 GPU 是 Nvidia K100。