如何使用numpy加快计算质心的代码?

用户名

我编写了一小段代码-给定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]]
安德烈亚斯·K。
center_of_mass = (A.mass * A.position + B.mass * B.position + C.mass * C.position) / total_mass

本文收集自互联网,转载请注明来源。

如有侵权,请联系 [email protected] 删除。

编辑于
0

我来说两句

0 条评论
登录 后参与评论

相关文章