有没有办法使用矢量化操作具有 1D 的 3D 数组？

Question

我正在尝试创建一个 3D 数组，其中 1 轴的前 n 个位置，其中 n 来自长度为 0 轴的数组。我正在处理一个大型数据集并试图加快速度。

我认为代码会更有意义。我正在尝试矢量化 for 循环。

    import numpy as np 

    data = np.zeros((3, 4, 5))
    num = np.array([2, 4, 3])

    for i in range(len(num)):
        data[:, 0][i, 0:num[i]] = 1

    print(data)

这是我要找的结果：

[[[1. 1. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[1. 1. 1. 1. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]

 [[1. 1. 1. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]
  [0. 0. 0. 0. 0.]]]

Answer 1

在下面的代码中，我介绍了三种通过运行时比较实现所需结果的方法：

你的 for 循环。
NumPy 的 apply_along_axis 函数
Numba 的 @jit 装饰器，它基本上编译你的 for 循环。

import numpy as np 
from numba import jit

def set_ones(arr_tup):
    arr_tup[0][0][:arr_tup[1]] = 1

@jit
def compiled_func(data, num):
    for i in range(len(num)): 
        data[:, 0][i, 0:num[i]] = 1

data = np.zeros((300, 4, 5))
num = np.random.randint(2, 5, 300)

data_num = np.array(list(zip(data, num)))

print(f'\n> for loop:')
%timeit for i in range(len(num)): data[:, 0][i, 0:num[i]] = 1

print(f'\n> apply_along_axis Loop:')
%timeit np.apply_along_axis(set_ones, 1, data_num)

print(f'\n> Compiled Loop:')
%timeit compiled_func(data, num)

data, num = zip(*data_num)

时间比较

> for loop:
288 µs ± 12.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

> apply_along_axis Loop:
1.15 ms ± 43.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

> Compiled Loop:
5.45 µs ± 483 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

从上面的对比我们看出：

正如@sai 正确指出的那样，NumPy's apply_along_axis function does not perform better than the regular loop, and my take on that is that due to the relatively small number of operations performed in each iteration, the speed gain achieved by using this function is dominated by the cost of its' header，所以不仅没有增益，而且实际上比常规循环慢。
第三种方式，即@jit效果最好，因为编译了for循环。所以基本上，现在 just-in-time（或 jit）编译器现在处理此代码，而不是 python 的解释器，这会导致 c - 类似循环性能。有关 Numba 的更多信息，您可以找到 here

干杯。

Answer 2

这有点 hacky 但有效：

>>> data = np.zeros((3, 4, 5))
>>> num = np.array([2, 4, 3])
>>> max_idx = min(data.shape[2], np.max(num))
>>> data[:,0,:max_idx] = np.arange(max_idx) < num[:,np.newaxis]
>>> data
array([[[1., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

       [[1., 1., 1., 1., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

       [[1., 1., 1., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]]])

它基于以下想法：

>>> np.arange(5)
array([0, 1, 2, 3, 4])
>>> (np.arange(5) < 2).astype(np.float64)
array([1., 1., 0., 0., 0.])

除 2 替换为垂直广播 num 以生成 'first N elements set to 1, rest 0'。

Answer 3

我不确定何时使用您的确切数据，但使用类似于您的 3x4x5 示例，此代码：

## make a 2d array of ones.
x = np.ones().reshape(3,1)

## mutiply those by an array of  [1,2,...,20]
## 20 because we want 4x5 max value
x = x*np.arange(1, 4*5+1)

## test where those are greater than the row number
## change the '3+1' expression to change the value of the comparison.
x = (x <= np.arange(1, 3+1).reshape(3, 1))

## booleans -> ints, reshape 2d to 3d (3x4x5) array
x = x.astype(int).reshape(-1, 4, 5)

产生这个输出：

array([[[1, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

       [[1, 1, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

       [[1, 1, 1, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]]])

它会比迭代快很多，所以我希望它能有所帮助。

为了使代码更通用一些，

## assign values for the axis lengths
ax_0_len = 3
ax_1_len = 4
ax_2_len = 5

## then replace the 3, 4, and 5 above with these variables.

有没有办法使用矢量化操作具有 1D 的 3D 数组？

Is there a way to manipulate a 3D array with a 1D using vectorization?

python

numpy

vectorization