Numpy数组切片与for循环结果比较

Question

想象一个由不同密度的壳组成的球体。
我有两个数组，一个是每个壳的外半径（rad[]），另一个是每个壳的密度（den[]）。我想计算到某个特定半径的质量，这个质量用mass[]来表示。

下面这个循环的方法可以得到我想要的结果。它首先计算最里面那个壳的质量（因为内半径是零，所以它就是一个球），然后再计算每个后续壳的质量，并把它加到之前的总质量上：

mass = numpy.zeros(len(rad))                                   # create array
mass[0] = den[0]**(rad[0]**3)                                  # find inner sphere mass
for i in range(1,len(mass)):
    mass[i] = mass[i-1] + den[i]*(rad[i]**3 - rad[i-1]**3)     # Find mass out to shell i

注意：我只需要比例，所以对π的具体数值不太在意。

有没有人能解释一下，为什么下面这个切片的结果达不到同样的效果呢？

mass = numpy.zeros(len(rad))
mass[0]  = den[0]*(rad[0]**3)
mass[1:] = mass[0:-1] + den[1:]*(rad[1:]**3-rad[0:-1]**3)

Answer 1

原因在于，numpy数组中的所有元素会同时计算。在你第二行的数组 mass[0:-1] 中，它会被计算为 den[0]*(rad[0]**3)，后面跟着的就是一堆零。（即使 mass[1] 在计算完成后不再是零，这个也没关系——到那时已经太晚了）。

你提到的例子：

test = numpy.linspace(1,10,num=10)
test[1:] += test[0:-1]
# [  1.   3.   6.  10.  15.  21.  28.  36.  45.  55.]

表现得不一样，仿佛加法是逐步进行的。你例子中的不同之处在于，右边加了一个值——这个加法让它在内存中变成了一个新数组（x + y 和 x 不是同一个数组），这样 numpy 就不再把它当作是对自身的加法了。看看这个例子

test = numpy.linspace(1,10,num=10)
test[1:] += test[0:-1] + 0
# [  1.   3.   5.   7.   9.  11.  13.  15.  17.  19.]

如果你想把你的for循环做成向量化的版本，可以这样做：

mass[1:] += den[1:]*(rad[1:]**3-rad[0:-1]**3)
mass[1:] += mass[0:-1]