如何使用 numpy 加速计算质心的代码？

Question

我编写了一小段代码 - 给定 n 个具有指定质量和矢量坐标的对象随时间推移 - 将计算质心。我认为代码看起来很笨重（它使用 3 个 for 循环），并且想知道是否有 numpy 方法来矢量化（或至少加速）这种方法。请注意，此任务可能避免使用类 Body，但在此处未显示的其他相关代码中使用。

import numpy as np

class Body():

    def __init__(self, mass, position):
        self.mass = mass
        self.position = position

    def __str__(self):
        return '\n .. mass:\n{}\n\n .. position:\n{}\n'.format(self.mass, self.position)

三个对象被初始化。

mass = 100 # same for all 3 objects
ndim = 3 # 3 dimensional space
nmoments = 10 # 10 moments in time

## initialize bodies
nelems = ndim * nmoments
x = np.arange(nelems).astype(int).reshape((nmoments, ndim))
A = Body(mass, position=x)
B = Body(mass, position=x / 2.)
C = Body(mass, position=x * 2.)
bodies = [A, B, C]
total_mass = sum([body.mass for body in bodies])

# print("\n ** A **\n{}\n".format(A))
# print("\n ** B **\n{}\n".format(B))
# print("\n ** C **\n{}\n".format(C))

## get center of mass
center_of_mass = []
for dim in range(ndim):
    coms = []
    for moment in range(nmoments):
        numerator = 0
        for body in bodies:
            numerator += body.mass * body.position[moment, dim]
        com = numerator / total_mass
        coms.append(com)
    center_of_mass.append(coms)
center_of_mass = np.array(center_of_mass).T

# print("\n .. center of mass:\n{}\n".format(center_of_mass))

作为验证代码有效，上面代码中的print语句输出如下：

 ** A **

 .. mass:
100

 .. position:
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]
 [12 13 14]
 [15 16 17]
 [18 19 20]
 [21 22 23]
 [24 25 26]
 [27 28 29]]



 ** B **

 .. mass:
100

 .. position:
[[ 0.   0.5  1. ]
 [ 1.5  2.   2.5]
 [ 3.   3.5  4. ]
 [ 4.5  5.   5.5]
 [ 6.   6.5  7. ]
 [ 7.5  8.   8.5]
 [ 9.   9.5 10. ]
 [10.5 11.  11.5]
 [12.  12.5 13. ]
 [13.5 14.  14.5]]



 ** C **

 .. mass:
100

 .. position:
[[ 0.  2.  4.]
 [ 6.  8. 10.]
 [12. 14. 16.]
 [18. 20. 22.]
 [24. 26. 28.]
 [30. 32. 34.]
 [36. 38. 40.]
 [42. 44. 46.]
 [48. 50. 52.]
 [54. 56. 58.]]



 .. center of mass:
[[ 0.          1.16666667  2.33333333]
 [ 3.5         4.66666667  5.83333333]
 [ 7.          8.16666667  9.33333333]
 [10.5        11.66666667 12.83333333]
 [14.         15.16666667 16.33333333]
 [17.5        18.66666667 19.83333333]
 [21.         22.16666667 23.33333333]
 [24.5        25.66666667 26.83333333]
 [28.         29.16666667 30.33333333]
 [31.5        32.66666667 33.83333333]]

Answer 1

使用 numpy 可以加快速度并使代码更简洁。我不是 n 体问题的专家，所以我希望能够遵循算法，结果看起来是一样的。所有的循环都隐含在 numpy 中。

# *****  From the question  *****
import numpy as np

class Body():   

    def __init__(self, mass, position):
        self.mass = mass
        self.position = position

    def __str__(self):
        return '\n .. mass:\n{}\n\n .. position:\n{}\n'.format(self.mass, self.position)

mass = 100 # same for all 3 objects
ndim = 3 # 3 dimensional space
nmoments = 10 # 10 moments in time

## initialize bodies
nelems = ndim * nmoments
x = np.arange(nelems).astype(int).reshape((nmoments, ndim))
A = Body(mass, position=x)
B = Body(mass, position=x / 2.)
C = Body(mass, position=x * 2.)
bodies = [A, B, C]

# **** End of code from the question  ****

# Fill the numpy arrays
np_mass = np.array( [ body.mass for body in bodies ])[ :,None, None ]
# the [:, None, None] turns np_mass into a 3D array for correct broadcasting

np_pos = np.array( [ body.position for body in bodies ])    # 3D 

np_mass.shape
# (3, 1, 1) # (n_bodies, 1, 1 ) - The two 'spare' dimensions force the broadcasting to be along the correct axes

np_pos.shape
# (3, 10, 3) # ( n_bodies, nmoments, ndims )

total_mass = np_mass.sum()  # Sum the three masses
numerator = (np_mass * np_pos).sum(axis=0)   # sum np_mass * np_pos along the body (0) axis. 
com = numerator / total_mass                 # divide by total_mass

# Could be a oneliner
# com = (np_mass * np_pos).sum(axis=0) / np.mass.sum()

print(com)

# array([[ 0.        ,  1.16666667,  2.33333333],
#        [ 3.5       ,  4.66666667,  5.83333333],
#        [ 7.        ,  8.16666667,  9.33333333],
#        [10.5       , 11.66666667, 12.83333333],
#        [14.        , 15.16666667, 16.33333333],
#        [17.5       , 18.66666667, 19.83333333],
#        [21.        , 22.16666667, 23.33333333],
#        [24.5       , 25.66666667, 26.83333333],
#        [28.        , 29.16666667, 30.33333333],
#        [31.5       , 32.66666667, 33.83333333]])

Answer 2

center_of_mass = (A.mass * A.position + B.mass * B.position + C.mass * C.position) / total_mass