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

时间:2017-03-03 03:54:03

标签: python arrays numpy for-loop

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

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方法来做同样的循环?

3 个答案:

答案 0 :(得分:2)

定时循环:

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个项目的数组。对于任务的大小,循环的数量不是很大。

丹尼尔的替代方案是快一点,但不是很激烈

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仅为此用途而构建。

答案 1 :(得分:1)

使用np.repeatnp.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]

答案 2 :(得分:1)

这是一种对循环进行矢量化的可能方法:

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

这就是详细情况:

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