来自形状的 Numpy 广播索引

Question

我有 2 个可以相互广播的数组形状。

例如(2, 2, 1) 和 (2, 3)

我想要一个函数，它采用这些形状并给我一个迭代器，它从这些数组中返回索引，这些形状将一起广播，以及结果输出数组中的索引。

iter, output_shape = broadcast_indeces_iterator((2, 2, 1), (2, 3))
assert output_shape == (2, 2, 3)
for in1_ix, in_2_ix, out_ix in iter:
    print (in1_ix, in_2_ix, out_ix)    

输出结果：

(0, 0, 0), (0, 0), (0, 0, 0)
(0, 0, 0), (0, 1), (0, 0, 1)
(0, 0, 0), (0, 2), (0, 0, 2)
(0, 1, 0), (1, 0), (0, 1, 0)
(0, 1, 0), (1, 1), (0, 1, 1)
(0, 1, 0), (1, 2), (0, 1, 2)
(1, 0, 0), (0, 0), (1, 0, 0)
(1, 0, 0), (0, 1), (1, 0, 1)
(1, 0, 0), (0, 2), (1, 0, 2)
(1, 1, 0), (1, 0), (1, 1, 0)
(1, 1, 0), (1, 1), (1, 1, 1)
(1, 1, 0), (1, 2), (1, 1, 2)

np.broadcast 做了一些接近但想要实际创建的数组的事情。

• NumPy 用户注意：如果 np.broadcast 有一个额外的参数允许您不迭代，例如，在最后 2 个维度上，那就太好了。这也可以解决我的问题。

Answer 1

import numpy as np
x = 10*np.arange(4).reshape((2, 2, 1))
y = 100*np.arange(6).reshape((2, 3))

z = np.nditer([x, y], flags=['multi_index', 'c_index'], order='C')
for a,b in z:
    print(np.unravel_index(z.index % x.size, x.shape)
          , np.unravel_index(z.index % y.size, y.shape)
          , z.multi_index)

产量

((0, 0, 0), (0, 0), (0, 0, 0))
((0, 1, 0), (0, 1), (0, 0, 1))
((1, 0, 0), (0, 2), (0, 0, 2))
((1, 1, 0), (1, 0), (0, 1, 0))
((0, 0, 0), (1, 1), (0, 1, 1))
((0, 1, 0), (1, 2), (0, 1, 2))
((1, 0, 0), (0, 0), (1, 0, 0))
((1, 1, 0), (0, 1), (1, 0, 1))
((0, 0, 0), (0, 2), (1, 0, 2))
((0, 1, 0), (1, 0), (1, 1, 0))
((1, 0, 0), (1, 1), (1, 1, 1))
((1, 1, 0), (1, 2), (1, 1, 2))

Answer 2

彼得的问题真好。这是你的答案：

import numpy as np


def get_broadcast_shape(*shapes):
    '''
    Given a set of array shapes, return the shape of the output when arrays of those 
    shapes are broadcast together
    '''
    max_nim = max(len(s) for s in shapes)
    equal_len_shapes = np.array([(1, )*(max_nim-len(s))+s for s in shapes]) 
    max_dim_shapes = np.max(equal_len_shapes, axis = 0)
    assert np.all(np.bitwise_or(equal_len_shapes==1, equal_len_shapes == max_dim_shapes[None, :])), \
        'Shapes %s are not broadcastable together' % (shapes, )
    return tuple(max_dim_shapes)


def get_broadcast_indeces(*shapes):
    '''
    Given a set of shapes of arrays that you could broadcast together, return:
        output_shape: The shape of the resulting output array
        broadcast_shape_iterator: An iterator that returns a len(shapes)+1 tuple
            of the indeces of each input array and their corresponding index in the 
            output array
    '''
    output_shape = get_broadcast_shape(*shapes)
    base_iter = np.ndindex(output_shape)

    def broadcast_shape_iterator():
        for out_ix in base_iter:
            in_ixs = tuple(tuple(0 if s[i] == 1 else ix for i, ix in enumerate(out_ix[-len(s):])) for s in shapes)
            yield in_ixs + (out_ix, )

    return output_shape, broadcast_shape_iterator()


output_shape, ix_iter = get_broadcast_indeces((2, 2, 1), (2, 3))
assert output_shape == (2, 2, 3)
for in1_ix, in_2_ix, out_ix in ix_iter:
    print (in1_ix, in_2_ix, out_ix)

返回

((0, 0, 0), (0, 0), (0, 0, 0))
((0, 0, 0), (0, 1), (0, 0, 1))
((0, 0, 0), (0, 2), (0, 0, 2))
((0, 1, 0), (1, 0), (0, 1, 0))
((0, 1, 0), (1, 1), (0, 1, 1))
((0, 1, 0), (1, 2), (0, 1, 2))
((1, 0, 0), (0, 0), (1, 0, 0))
((1, 0, 0), (0, 1), (1, 0, 1))
((1, 0, 0), (0, 2), (1, 0, 2))
((1, 1, 0), (1, 0), (1, 1, 0))
((1, 1, 0), (1, 1), (1, 1, 1))
((1, 1, 0), (1, 2), (1, 1, 2))

如果有人知道任何解决此问题的 numpy 内置函数，那会更好。

Answer 3

这是一个开始：

array1 = np.arange(4).reshape(2,2,1)*10
array2 = np.arange(6).reshape(2,3)

I, J = np.broadcast_arrays(array1, array2)
print I.shape
K = np.empty(I.shape, dtype=int)
for ijk in np.ndindex(I.shape):
    K[ijk] = I[ijk]+J[ijk]
print K

生产

(2, 2, 3)  # broadcasted shape

[[[ 0  1  2]
  [13 14 15]]    
 [[20 21 22]
  [33 34 35]]]

I 是 (2,2,3)，但与 array1 共享其数据 - 这是一个广播视图，而不是副本（查看它的 .__array_interface__）。

您可以通过仅给出 ndindex 那些形状来迭代仅 2 个维度。

K = np.empty(I.shape, dtype=int)
for i,j in np.ndindex(I.shape[:2]):
    K[i,j,:] = I[i,j,:]+J[i,j,:]
    print K[i,j,:]

可以通过查看broadcast_arrays 和ndindex 的代码来找到基本部分来完善它。例如在https://stackoverflow.com/a/25097271/901925 中，我直接调用nditer 以生成multi_index（可以适应cython 的操作）。

xx = np.zeros(y.shape[:2])
it = np.nditer(xx,flags=['multi_index'])                               
while not it.finished:
    print y[it.multi_index],
    it.iternext()
# [242  14 211] [198   7   0] [235  60  81] [164  64 236]

---

要制作几乎没有成本的“虚拟数组”，我可以从ndindex 中获取线索，并制作以np.zeros(1) 开头的数组

def make_dummy(shape):
    x = as_strided(np.zeros(1),shape=shape, strides=np.zeros_like(shape))
    return x
array1 = make_dummy((2,2,1))
array2 = make_dummy((2,3))

我可以深入研究 np.broadcast_arrays 以了解它如何组合来自 2 个输入数组的形状以得出 I 的形状。

---

您想要的解决方案与我的解决方案之间存在差异，我已经掩盖了。

(0, 0, 0), (0, 0), (0, 0, 0)
(0, 0, 0), (0, 1), (0, 0, 1)
...
(1, 1, 0), (1, 1), (1, 1, 1)
(1, 1, 0), (1, 2), (1, 1, 2)

期望每个数组有一个不同的迭代器元组，一个范围超过(2,2,1)，另一个范围超过(2,3)，等等。

我相信numpy c 代码（至少基于nditer 的那些部分）使用的方法在(2,2,3) 上生成一个迭代器，并通过as_strided 对数组进行按摩以接受这一点范围更大。这种方式更容易实现通用的广播机制。它将广播复杂度与计算核心分离。

这是nditer的一个很好的介绍：

http://docs.scipy.org/doc/numpy/reference/arrays.nditer.html