numpy 使用np.nditer作为zip

whlutmcx 于 2023-02-08 发布在其他

关注(0)|答案(1)|浏览(130)

我尝试过像zip()一样将函数np.nditer()应用于不同维度的数组，其中迭代器应该只使用第一维度。

最小示例

a_all = np.arange(6).reshape(2,3)
idx_all = np.arange(12).reshape(2,3,2)

for a, idx in np.nditer([a_all, idx_all]):
    print((a, idx))

这将引发错误：

ValueError: operands could not be broadcast together with shapes (2,3) (2,3,2)

我的使用案例

我有两个数组，其中的数据我想互相计算，而且我还有另一个数组的索引列表，所以我尝试：

a_all = np.arange(6).reshape(2,3)
b_all = np.arange(6).reshape(2,3)
idx_all = (
    ((0,0), (0,1), (0,2)),
    ((1,0), (1,1), (1,2))
)
result = np.zeros((2,3))

for a, b, idx in np.nditer([a_all, b_all, idx_all]):
    result[idx] += a*b

它会抛出与最小示例相同的错误。
我假设问题在于np.nditer()试图迭代idx_all的所有维度，但我不知道如何将其限制在前两个维度。
我不想使用zip()，否则我需要两个循环：

for a_, b_, idx_ in zip(a_all, b_all, idx_all):
    for a, b, idx in zip(a_, b_, idx_):
        result[idx] += a*b

更合理的示例

a_all = np.random.randn(2,3)
b_all = np.random.randn(2)
idx_all = (
    ((1,1), (2,2))
)
result = np.zeros(2)

for a, b, idx, res in np.nditer([a_all, b_all, idx_all, result], op_flags=['readwrite']):
    res += a[idx] + b

numpy

来源：https://stackoverflow.com/questions/75281577/use-np-nditer-as-zip

1条答案

按热度按时间

agxfikkp1#

看看第一个例子，修正后的数组确实一起广播（如果你不明白我做了什么改变，你已经读了足够多的基本文档）。

In [14]: a_all = np.arange(6).reshape(2,3,1)
    ...: idx_all = np.arange(12).reshape(2,3,2)
    ...:
    ...: for a, idx in np.nditer([a_all, idx_all]):
    ...:     print((a, idx))
    ...:
(array(0), array(0))
(array(0), array(1))
(array(1), array(2))
(array(1), array(3))
(array(2), array(4))
(array(2), array(5))
(array(3), array(6))
(array(3), array(7))
(array(4), array(8))
(array(4), array(9))
(array(5), array(10))
(array(5), array(11))

nditer在“平面”意义上进行迭代，将单个元素数组（0d）传递到主体。它不像zip那样只在第一维（或嵌套列表的外层）上进行迭代。
np.vectorize（我也不推荐它），执行相同类型的广播，但是将python标量元素传递给函数：

In [15]: np.vectorize(lambda a,idx: print((a,idx)))(a_all,idx_all)
(0, 0)             # test run
(0, 0)
(0, 1)
(1, 2)
(1, 3)
(2, 4)
(2, 5)
(3, 6)
(3, 7)
(4, 8)
(4, 9)
(5, 10)
(5, 11)
Out[15]:
array([[[None, None],
        [None, None],
        [None, None]],

       [[None, None],
        [None, None],
        [None, None]]], dtype=object)

nditer需要和np.vectorize一样的性能免责声明。它没有帮助，至少在python代码中使用时没有帮助。在cython中，它可能很有用，如更大的nditer文档页所示。
此外，nditer输入可能很复杂，如上一个示例生成的TypeError所示。
最后一个例子：
我不得不将idx_all更改为数组，而不是元组，这样它就可以读写了。请更仔细地阅读op_flags文档。
我们仍然会得到广播错误。它没有迭代“第一层”。

In [24]: a_all = np.random.randn(2,3)
    ...: b_all = np.random.randn(2)
    ...: idx_all = (
    ...:     ((1,1), (2,2))
    ...: ); idx_all=np.array(idx_all)
    ...: result = np.zeros(2)
    ...:
    ...: for a, b, idx, res in np.nditer([a_all, b_all, idx_all, result], op_flags=['readwrite']):
    ...:     res += a[idx] + b
    ...:
ValueError: operands could not be broadcast together with shapes (2,3) (2,) (2,2) (2,)

赞(0）回复(0）举报 2023-02-08

我来回答

numpy 使用np.nditer作为zip

1条答案

相关问题

热门标签

最新问答