Python中高效计算n体引力

Question

Python中高效计算n体引力

8

我正在尝试计算三维空间中n体问题的重力加速度（使用辛欧拉法symplectic Euler）。

我对每个时间步长都有位置向量和速度向量，并使用下面的（有效的）代码来计算加速度并更新速度和位置。请注意，加速度是三维空间中的向量，而不仅仅是数量。

我想知道是否有更有效的方式使用numpy计算以避免循环。

def accelerations(positions, masses):
    '''Params:
    - positions: numpy array of size (n,3)
    - masses: numpy array of size (n,)
    Returns:
    - accelerations: numpy of size (n,3), the acceleration vectors in 3-space
    '''
    n_bodies = len(masses)
    accelerations = numpy.zeros([n_bodies,3]) # n_bodies * (x,y,z)

    # vectors from mass(i) to mass(j)
    D = numpy.zeros([n_bodies,n_bodies,3]) # n_bodies * n_bodies * (x,y,z)
    for i, j in itertools.product(range(n_bodies), range(n_bodies)):
        D[i][j] = positions[j]-positions[i]

    # Acceleration due to gravitational force between each pair of bodies
    A = numpy.zeros((n_bodies, n_bodies,3))
    for i, j in itertools.product(range(n_bodies), range(n_bodies)):
        if numpy.linalg.norm(D[i][j]) > epsilon:
            A[i][j] = gravitational_constant * masses[j] * D[i][j] \
            / numpy.linalg.norm(D[i][j])**3

    # Calculate net acceleration of each body (vectors in 3-space)
    accelerations = numpy.sum(A, axis=1) # sum of accel vectors for each body of shape (n_bodies,3)

    return accelerations

- Paul Dickinson

可能是Python中最快的成对距离度量的重复问题。 - PMende

我已将此标记为重复，但是为了回答您的问题，请参考以下链接：https://dev59.com/4mIj5IYBdhLWcg3wb0jX，以及最终的 https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.pdist.html。 - PMende

至于加速部分，您只需使用调用np.where的方法，结合一个“质量矩阵”，可以通过执行以下操作获得：masses.reshape((1, -1))*masses.reshape((-1, 1))。 - PMende

我认为pdist只是给出向量的大小。由于我想在三维空间中添加力，所以我还需要保留方向。 - Paul Dickinson

最佳算法取决于epsilon和点质量的数量。对于至少中等规模的问题，我建议使用kd-tree https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.cKDTree.html 来查找彼此在给定范围内的点质量，并仅计算它们之间的引力。这比计算无用的距离要快得多。对于其余部分，我建议使用简单的Numba方法。您能否提供合理的测试数据（位置、质量、epsilon），因为这对运行时间有很大影响？ - max9111

3个回答

5

在你原始帖子的评论后续：

from numpy.linalg import norm

def accelerations(positions, masses):
    '''Params:
    - positions: numpy array of size (n,3)
    - masses: numpy array of size (n,)
    '''
    mass_matrix = masses.reshape((1, -1, 1))*masses.reshape((-1, 1, 1))
    disps = positions.reshape((1, -1, 3)) - positions.reshape((-1, 1, 3)) # displacements
    dists = norm(disps, axis=2)
    dists[dists == 0] = 1 # Avoid divide by zero warnings
    forces = G*disps*mass_matrix/np.expand_dims(dists, 2)**3
    return forces.sum(axis=1)/masses.reshape(-1, 1)

- PMende

谢谢，pdist和np.where都是很好知道的！但这并没有完全解决问题，因为输出应该是三维向量，而不仅仅是它们的大小。这也是为什么原始代码中有一个立方体而不是平方因子的原因。 - Paul Dickinson

@PaulDickinson 啊，好的。对不起，我已经更新了答案，使其具有正确的形状。 - PMende

广播仍然存在问题。我会在一分钟内更新。 - PMende

@PaulDickinson 现在应该可以工作了。主要是通过巧妙的重塑/广播来实现的练习。 - PMende

很好，这比原来的解决方案更加优雅！ - Paul Dickinson

2

需要考虑的一些事情：

你只需要计算一半的距离；一旦你计算出了 D[i][j]，那么它就等同于 -D[j][i]。

你可以使用 df.apply(lambda x:gravitational_constant/x**3) 来进行操作。

你可以生成一个记录每对物体质量乘积的数据框。你只需要这样做一次，然后每次调用 accelearations 时都可以将其传递给它。

然后，df.product(df2).product(mass_products).sum().div(masses) 就可以给你加速度。

- Acccumulation

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Paul Panzer · Accepted Answer

这里是使用blas进行了优化的版本。blas拥有针对对称或厄米矩阵的特殊线性代数例程。这些例程使用专门的打包存储，仅保留上三角或下三角，并省略（冗余的）镜像条目。这样，blas不仅节约了大约一半的存储空间，还节约了大约一半的浮点运算次数。我添加了相当多的注释以使其易于阅读。

import numpy as np
import itertools
from scipy.linalg.blas import zhpr, dspr2, zhpmv

def acc_vect(pos, mas):
    n = mas.size
    d2 = pos@(-2*pos.T)
    diag = -0.5 * np.einsum('ii->i', d2)
    d2 += diag + diag[:, None]
    np.einsum('ii->i', d2)[...] = 1
    return np.nansum((pos[:, None, :] - pos) * (mas[:, None] * d2**-1.5)[..., None], axis=0)

def acc_blas(pos, mas):
    n = mas.size
    # trick: use complex Hermitian to get the packed anti-symmetric
    # outer difference in the imaginary part of the zhpr answer
    # don't want to sum over dimensions yet, therefore must do them one-by-one
    trck = np.zeros((3, n * (n + 1) // 2), complex)
    for a, p in zip(trck, pos.T - 1j):
        zhpr(n, -2, p, a, 1, 0, 0, 1)
        # does  a  ->  a + alpha x x^H
        # parameters: n             --  matrix dimension
        #             alpha         --  real scalar
        #             x             --  complex vector
        #             ap            --  packed Hermitian n x n matrix a
        #                               i.e. an n(n+1)/2 vector
        #             incx          --  x stride
        #             offx          --  x offset
        #             lower         --  is storage of ap lower or upper
        #             overwrite_ap  --  whether to change a inplace
    # as a by-product we get pos pos^T:
    ppT = trck.real.sum(0) + 6
    # now compute matrix of squared distances ...
    # ... using (A-B)^2 = A^2 + B^2 - 2AB
    # ... that and the outer sum X (+) X.T equals X ones^T + ones X^T
    dspr2(n, -0.5, ppT[np.r_[0, 2:n+1].cumsum()], np.ones((n,)), ppT,
          1, 0, 1, 0, 0, 1)
    # does  a  ->  a + alpha x y^T + alpha y x^T    in packed symmetric storage
    # scale anti-symmetric differences by distance^-3
    np.divide(trck.imag, ppT*np.sqrt(ppT), where=ppT.astype(bool),
              out=trck.imag)
    # it remains to scale by mass and sum
    # this can be done by matrix multiplication with the vector of masses ...
    # ... unfortunately because we need anti-symmetry we need to work
    # with Hermitian storage, i.e. complex numbers, even though the actual
    # computation is only real:
    out = np.zeros((3, n), complex)
    for a, o in zip(trck, out):
        zhpmv(n, 0.5, a, mas*-1j, 1, 0, 0, o, 1, 0, 0, 1)
        # multiplies packed Hermitian matrix by vector
    return out.real.T

def accelerations(positions, masses, epsilon=1e-6, gravitational_constant=1.0):
    '''Params:
    - positions: numpy array of size (n,3)
    - masses: numpy array of size (n,)
    '''
    n_bodies = len(masses)
    accelerations = np.zeros([n_bodies,3]) # n_bodies * (x,y,z)

    # vectors from mass(i) to mass(j)
    D = np.zeros([n_bodies,n_bodies,3]) # n_bodies * n_bodies * (x,y,z)
    for i, j in itertools.product(range(n_bodies), range(n_bodies)):
        D[i][j] = positions[j]-positions[i]

    # Acceleration due to gravitational force between each pair of bodies
    A = np.zeros((n_bodies, n_bodies,3))
    for i, j in itertools.product(range(n_bodies), range(n_bodies)):
        if np.linalg.norm(D[i][j]) > epsilon:
            A[i][j] = gravitational_constant * masses[j] * D[i][j] \
            / np.linalg.norm(D[i][j])**3

    # Calculate net accleration of each body
    accelerations = np.sum(A, axis=1) # sum of accel vectors for each body

    return accelerations

from numpy.linalg import norm

def acc_pm(positions, masses, G=1):
    '''Params:
    - positions: numpy array of size (n,3)
    - masses: numpy array of size (n,)
    '''
    mass_matrix = masses.reshape((1, -1, 1))*masses.reshape((-1, 1, 1))
    disps = positions.reshape((1, -1, 3)) - positions.reshape((-1, 1, 3)) # displacements
    dists = norm(disps, axis=2)
    dists[dists == 0] = 1 # Avoid divide by zero warnings
    forces = G*disps*mass_matrix/np.expand_dims(dists, 2)**3
    return forces.sum(axis=1)/masses.reshape(-1, 1)

n = 500
pos = np.random.random((n, 3))
mas = np.random.random((n,))

from timeit import timeit

print(f"loops:      {timeit('accelerations(pos, mas)', globals=globals(), number=1)*1000:10.3f} ms")
print(f"pmende:     {timeit('acc_pm(pos, mas)', globals=globals(), number=10)*100:10.3f} ms")
print(f"vectorized: {timeit('acc_vect(pos, mas)', globals=globals(), number=10)*100:10.3f} ms")
print(f"blas:       {timeit('acc_blas(pos, mas)', globals=globals(), number=10)*100:10.3f} ms")

A = accelerations(pos, mas)
AV = acc_vect(pos, mas)
AB = acc_blas(pos, mas)
AP = acc_pm(pos, mas)

assert np.allclose(A, AV) and np.allclose(AB, AV) and np.allclose(AP, AV)

样例运行; 与OP、我的纯numpy向量化和@P Mende的比较。

loops:        3213.130 ms
pmende:         41.480 ms
vectorized:     43.860 ms
blas:            7.726 ms

我们可以看到：

1）P Mende在矢量化方面略胜过我；

2）blas速度大约快了 ~5倍，请注意我的blas不是很好。我怀疑如果使用优化过的blas，你可能会得到更好的结果（numpy在更好的blas上也应该运行得更快，但需要注意的是）。

3）无论哪个答案都比循环要快得多。