在多维numpy数组中快速迭代向量

Question

我正在编写一些Python + NumPy + Cython代码，并尝试找到在数组上进行以下迭代的最优雅和高效的方法：

假设我有一个函数f(x, y)，它接受形状为(3,)的向量x和形状为(10,)的向量y，并返回形状为(10,)的向量。现在我有两个形状分别为sx + (3,)和sy + (10,)的数组X和Y，其中sx和sy是可以广播在一起的两个形状（即sx == sy或者当轴不同时，其中一个长度为1，在这种情况下将被重复）。我想要生成一个形状为zs + (10,)的数组Z，其中zs是sx与sy的广播形状。Z中的每个10维向量都等于在X和Y中对应位置处的向量x和y的f(x, y)。

我研究了np.nditer，虽然它与cython很配合（请参见链接页面底部），但似乎不能迭代多维数组中的向量，而不是元素。我还看了一下index grids，但问题在于当索引数量等于数组维度并且存储为cython整数而不是python元组时，cython索引才快速。任何帮助都将不胜感激！

- John von N.

---

您能否给你的需求添加一个具体的示例？现在，按照您的思路跟踪并不容易。 - Oliver W.

我本来想翻译的，但是Jaime的回答中的链接已经很清楚地解释了我所想表达的概念。 - John von N.

multi_index 可用于控制 nditer 迭代的深度。基本上，您可以构建一个迭代器来遍历 zs 的形状，并使用它来索引 X、Y 和输出。http://stackoverflow.com/a/25097271/901925 - hpaulj

Answer 1

你在描述Numpy所称的广义通用函数，或者说gufunc。正如它的名字所示，它是ufunc的扩展。你可能想从阅读以下两个页面开始：

• 编写自己的ufunc
• 从头开始构建ufunc

第二个例子使用了Cython并且涉及到了一些关于gufunc的内容。要完全了解gufunc，你需要阅读numpy C API文档中相应的部分：

• 广义通用函数API

我不知道在Cython中编写gufunc的任何示例，尽管根据上面的示例应该不太难。如果您想查看用C编写的gufunc，请查看np.linalg的源代码here，但那可能是一次令人望而生畏的经历。一段时间以前，我向当地的Python用户组介绍了使用C扩展numpy的演讲，其中大部分是关于在C中编写gufunc的，可以在here找到该演讲的幻灯片和提供新gufunc的示例Python模块。

- Jaime

---

谢谢！这似乎正是我在寻找的。我知道这必须是一个常见的问题，但我还不知道广义ufuncs。 - John von N.

Answer 2

如果您想继续使用nditer，可以使用您的示例维度来实现。 这里是纯Python代码，但使用cython实现应该不难（尽管仍然具有元组迭代器）。 我从shallow iteration with nditer中描述的ndindex中借鉴了一些思路。

思路是找到公共广播形状sz，并构建一个multi_index迭代器。

我使用as_strided将X和Y扩展为可用视图，并将适当的向量（实际上是（1，n）数组）传递给f（x，y）函数。

import numpy as np
from numpy.lib.stride_tricks import as_strided

def f(x,y):
    # sample that takes (10,) and (3,) arrays, and returns (10,) array
    assert x.shape==(1,10), x.shape
    assert y.shape==(1,3), y.shape
    z = x*10 + y.mean()
    return z

def brdcast(X, X1):
    # broadcast X to shape of X1 (keep last dim of X)
    # modeled on np.broadcast_arrays
    shape = X1.shape + (X.shape[-1],)
    strides = X1.strides + (X.strides[-1],)
    X1 = as_strided(X, shape=shape, strides=strides)
    return X1

def F(X, Y):
    X1, Y1 = np.broadcast_arrays(X[...,0], Y[...,0])
    Z = np.zeros(X1.shape + (10,))
    it = np.nditer(X1, flags=['multi_index'])

    X1 = brdcast(X, X1)
    Y1 = brdcast(Y, Y1)

    while not it.finished:
        I = it.multi_index + (None,)
        Z[I] = f(X1[I], Y1[I])
        it.iternext()
    return Z

sx = (2,3)  # works with (2,1)
sy = (1,3)
# X, Y = np.ones(sx+(10,)), np.ones(sy+(3,))

X = np.repeat(np.arange(np.prod(sx)).reshape(sx)[...,None], 10, axis=-1)
Y = np.repeat(np.arange(np.prod(sy)).reshape(sy)[...,None], 3, axis=-1)

Z = F(X,Y)
print Z.shape
print Z[...,0]

- hpaulj