如何在没有 for 循环的情况下填充数组答案

【问题标题】：How to fill an array without a for loop如何在没有 for 循环的情况下填充数组
【发布时间】：2017-03-03 03:54:02
【问题描述】：

我有一些代码可以生成圆柱对称曲面的坐标，坐标为 (r, theta, phi)。目前，我生成一个 phi 切片的坐标，并将其存储在一个 2xN 数组中（对于 N 个 bin），然后在一个 for 循环中，我将这个数组从 0 复制到 2pi 的每个 phi 值：

import numpy as np

# this is the number of bins that my surface is chopped into
numbins = 50

# these are the values for r
r_vals = np.linspace(0.0001, 50, numbins, endpoint = True)

# these are the values for theta
theta_vals = np.linspace(0, np.pi / 2, numbins, endpoint = True)

# I combine the r and theta values into a 2xnumbins array for one "slice" of phi
phi_slice = np.vstack([r_vals,theta_vals]).T

# this is the array which will store all of the coordinates of the surface
surface = np.zeros((numbins**2,3))
lower_bound = 0
upper_bound = numbins

# this is the size of the bin for phi
dphi = (2 * np.pi) / numbins

# this is the for loop I'd like to eliminate.
# For every value of phi, it puts a copy of the phi_slice array into
# the surface array, so that the surface is cylindrical about phi.
for phi in np.linspace(0, (2 * np.pi) - dphi, numbins):
    surface[lower_bound:upper_bound, :2] = phi_slice
    surface[lower_bound:upper_bound,2] = phi
    lower_bound += numbins
    upper_bound += numbins

我在 1e6 或 1e7 步的数值积分中调用此例程，虽然在上面的示例中 numbins 是 50，但实际上它会是数千。这个 for 循环是一个瓶颈，我真的很想消除它以加快速度。有没有一种很好的 NumPythonic 方法来做与这个 for 循环相同的事情？

【问题讨论】：

我投票决定将此问题作为离题结束，因为该问题属于 Stack Exchange 网络中的另一个站点：codereview.stackexchange.com
留在这里；从 numpy 代码中消除循环是一个常见的 SO 问题。这里有更多的 numpy 专家。
但是，如果示例是 MCVe，它会有所帮助 - 最小、完整、可验证。这样我们就可以轻松地复制粘贴、测试结果和测试更改。显示一些输入和输出。
仅通过阅读代码很难想象发生了什么。它需要更简单一些具体的值。添加一些说明。你循环了多少次？
感谢您的反馈！我已经编辑过了，现在是 MCV；让我知道我是否可以详细说明。

标签： python arrays numpy for-loop

【解决方案1】：

定时循环：

In [9]: %%timeit
   ...: lower_bound, upper_bound = 0, numbins 
   ...: for phi in np.linspace(0, (2 * np.pi) - dphi, numbins):
   ...:     surface[lower_bound:upper_bound,:2] = phi_slice
   ...:     surface[lower_bound:upper_bound,2] = phi
   ...:     lower_bound += numbins
   ...:     upper_bound += numbins

10000 loops, best of 3: 176 µs per loop

看起来不错的副手，但如果在更大的上下文中重复这确实很重要。您循环 50 次以填充 75000 个项目的数组。对于任务的大小，循环的数量并不大。

Daniel 的替代方案稍微快一点，但并不激烈

In [12]: %%timeit 
    ...: phi_slices = np.tile(phi_slice.T, numbins).T
    ...: phi_indices = np.repeat(np.linspace(0, (2 * np.pi) - dphi, numbins), numbins)
    ...: surface1 = np.c_[phi_slices, phi_indices]
    ...: 
10000 loops, best of 3: 137 µs per loop

kazemakase's还是好一点：

In [15]: %%timeit 
    ...: phis = np.repeat(np.linspace(0, (2 * np.pi) - dphi, numbins), numbins)[:, np.newaxis]
    ...: slices = np.repeat(phi_slice[np.newaxis, :, :], numbins, axis=0).reshape(-1, 2)
    ...: surface2 = np.hstack([slices, phis])
    ...: 
10000 loops, best of 3: 115 µs per loop

还有我的提名（感谢其他人帮助我了解模式）；我在作业中利用广播。

surface3 = np.zeros((numbins,numbins,3))
phis = np.linspace(0, (2 * np.pi) - dphi, numbins)
surface3[:,:,2] = phis[:,None]
surface3[:,:,:2] = phi_slice[None,:,:]
surface3.shape = (numbins**2,3)

更好一点：

In [50]: %%timeit
    ...: surface3 = np.zeros((numbins,numbins,3))
    ...: phis=np.linspace(0, (2 * np.pi) - dphi, numbins)
    ...: surface3[:,:,2]=phis[:,None]
    ...: surface3[:,:,:2]=phi_slice[None,:,:]
    ...: surface3.shape=(numbins**2,3)
    ...: 
10000 loops, best of 3: 73.3 µs per loop

编辑

替换：

surface3[:,:,:2]=phi_slice[None,:,:]

与

surface3[:,:,0]=r_vals[None,:]
surface3[:,:,1]=theta_vals[None,:]

挤出更多的时间改进，特别是如果 phi_slice 是专门为这种用途而构建的。

【讨论】：

有兴趣看看numbins 是“成千上万”时的情况如何叠加，因为你们俩最后都进行了重塑。
是的，那是事情开始变得非常缓慢的时候，因为surface 必然是numbins^2
我不认为重塑，尤其是就地重塑，会改变相对时间。这是一个微不足道的操作。

【解决方案2】：

使用np.repeat 和np.tile 构建这种分块数组

phi_slices = np.tile(phi_slice.T, numbins).T
phi_indices = np.repeat(np.linspace(0, (2 * np.pi) - dphi, numbins), numbins)
surface = np.c_[phi_slices, phi_indices]

【讨论】：

谢谢！我以前没见过 np.repeat，看起来正是我需要的

【解决方案3】：

这是一种矢量化循环的可能方式：

phis = np.repeat(np.linspace(0, (2 * np.pi) - dphi, numbins), numbins)[:, np.newaxis]
slices = np.repeat(phi_slice[np.newaxis, :, :], numbins, axis=0).reshape(-1, 2)

surface2 = np.hstack([slices, phis])

print(np.allclose(surface, surface2))
# True

这就是详细的情况：

np.repeat(np.linspace(0, (2 * np.pi) - dphi, numbins), numbins) 采用所有值为 phi 的数组并重复 每个元素 numbins 次。 [:, np.newaxis] 将结果转化为二维形状 (numbins**2, 1)。
phi_slice[np.newaxis, :, :] 将 phi_slice 带入 3D 形状 (1, numbins, 2.。该数组沿第一个轴重复，形成(numbins, numbins, 2) 形状。最后reshape(-1, 2) 将前两个维度重新组合到最终形状(numbins**2, 2)。
np.hstack([slices, phis]) 结合了最终数组的两个部分。

【讨论】：