(在这篇文章中,让np
成为numpy
的简写。)
假设a
是( n + k ) - 维np.ndarray
对象,对于某些整数 n > 1和 k > 1.(IOW, n + k > 3是a.ndim
)的值。我想在其第一个 n 维度上枚举a
;这意味着,在每次迭代时,枚举器/迭代器都会生成一对,其第一个元素是 n 索引的元组ii
,第二个元素是 k ndarray
处的维度子a[ii]
。
当然,编写一个函数来做这个并不困难(事实上,我在下面给出了这样一个函数的例子),但我想知道这个:
numpy
是否提供了执行此类“部分”枚举的特殊语法或函数?
(通常,当我想迭代一个多维np.ndarray
对象时,我使用np.ndenumerate
,但这在这里没有帮助,因为(据我所知){{1}将迭代所有 n + k 尺寸。)
假设上述问题的答案是肯定的,那么就是这样的后续行动:
迭代的 n 维度不连续的情况怎么样?
(在这种情况下,枚举器/迭代器在每次迭代时返回的对的第一个元素将是 r > n 元素的元组,其中一些这是一个表示“全部”的特殊值,例如np.ndenumerate
;该对的第二个元素仍然是长度为 k 的slice(None)
。)
谢谢!
下面的代码有希望澄清问题规范。函数ndarray
使用可用于此目的的任何特殊partial_enumerate
构造执行我想要执行的操作。在numpy
的定义之后是 n = k = 2的简单示例。
partial_enumerate
输出的每一行都是“一对元组”,其中第一个元组表示import numpy as np
import itertools as it
def partial_enumerate(nda, n):
"""Enumerate over the first N dimensions of the numpy.ndarray NDA.
Returns an iterator of pairs. The first element of each pair is a tuple
of N integers, corresponding to a partial index I into NDA; the second element
is the subarray of NDA at I.
"""
# ERROR CHECKING & HANDLING OMITTED
for ii in it.product(*[range(d) for d in nda.shape[:n]]):
yield ii, nda[ii]
a = np.zeros((2, 3, 4, 5))
for ii, vv in partial_enumerate(a, 2):
print ii, vv.shape
中 n 坐标的部分集合,第二个元组表示形状这些部分坐标处的{em> k - a
的维数子阵列; (第二对的值对于所有行都是相同的,正如数组的规律性所预期的那样):
a
相反,在这种情况下迭代(0, 0) (4, 5)
(0, 1) (4, 5)
(0, 2) (4, 5)
(1, 0) (4, 5)
(1, 1) (4, 5)
(1, 2) (4, 5)
将导致np.ndenumerate(a)
次迭代,每次迭代都会访问a.size
的单个单元格。
答案 0 :(得分:5)
您可以使用numpy广播规则生成笛卡尔积。 numpy.ix_
函数创建适当数组的列表。它等同于以下内容:
>>> def pseudo_ix_gen(*arrays):
... base_shape = [1 for arr in arrays]
... for dim, arr in enumerate(arrays):
... shape = base_shape[:]
... shape[dim] = len(arr)
... yield numpy.array(arr).reshape(shape)
...
>>> def pseudo_ix_(*arrays):
... return list(pseudo_ix_gen(*arrays))
或者,更简洁:
>>> def pseudo_ix_(*arrays):
... shapes = numpy.diagflat([len(a) - 1 for a in arrays]) + 1
... return [numpy.array(a).reshape(s) for a, s in zip(arrays, shapes)]
结果是可广播数组列表:
>>> numpy.ix_(*[[2, 4], [1, 3], [0, 2]])
[array([[[2]],
[[4]]]), array([[[1],
[3]]]), array([[[0, 2]]])]
将此与numpy.ogrid
:
>>> numpy.ogrid[0:2, 0:2, 0:2]
[array([[[0]],
[[1]]]), array([[[0],
[1]]]), array([[[0, 1]]])]
如您所见,它是相同的,但numpy.ix_
允许您使用非连续索引。现在,当我们应用numpy广播规则时,我们得到了一个笛卡尔积:
>>> list(numpy.broadcast(*numpy.ix_(*[[2, 4], [1, 3], [0, 2]])))
[(2, 1, 0), (2, 1, 2), (2, 3, 0), (2, 3, 2),
(4, 1, 0), (4, 1, 2), (4, 3, 0), (4, 3, 2)]
如果我们不是将numpy.ix_
的结果传递给numpy.broadcast
,而是使用它来索引数组,我们得到这个:
>>> a = numpy.arange(6 ** 4).reshape((6, 6, 6, 6))
>>> a[numpy.ix_(*[[2, 4], [1, 3], [0, 2]])]
array([[[[468, 469, 470, 471, 472, 473],
[480, 481, 482, 483, 484, 485]],
[[540, 541, 542, 543, 544, 545],
[552, 553, 554, 555, 556, 557]]],
[[[900, 901, 902, 903, 904, 905],
[912, 913, 914, 915, 916, 917]],
[[972, 973, 974, 975, 976, 977],
[984, 985, 986, 987, 988, 989]]]])
然而,警告经纪人。可广播数组对于索引非常有用,但如果你真的希望枚举值,那么最好使用itertools.product
:
>>> %timeit list(itertools.product(range(5), repeat=5))
10000 loops, best of 3: 196 us per loop
>>> %timeit list(numpy.broadcast(*numpy.ix_(*([range(5)] * 5))))
100 loops, best of 3: 2.74 ms per loop
因此,如果您仍然要合并for循环,那么itertools.product
可能会更快。不过,您仍然可以使用上述方法在纯粹的numpy中获得一些类似的数据结构:
>> pgrid_idx = numpy.ix_(*[[2, 4], [1, 3], [0, 2]])
>>> sub_indices = numpy.rec.fromarrays(numpy.indices((6, 6, 6)))
>>> a[pgrid_idx].reshape((8, 6))
array([[468, 469, 470, 471, 472, 473],
[480, 481, 482, 483, 484, 485],
[540, 541, 542, 543, 544, 545],
[552, 553, 554, 555, 556, 557],
[900, 901, 902, 903, 904, 905],
[912, 913, 914, 915, 916, 917],
[972, 973, 974, 975, 976, 977],
[984, 985, 986, 987, 988, 989]])
>>> sub_indices[pgrid_idx].reshape((8,))
rec.array([(2, 1, 0), (2, 1, 2), (2, 3, 0), (2, 3, 2),
(4, 1, 0), (4, 1, 2), (4, 3, 0), (4, 3, 2)],
dtype=[('f0', '<i8'), ('f1', '<i8'), ('f2', '<i8')])
答案 1 :(得分:4)
我认为您正在寻找numpy
中的ndindex功能。只需要一小部分您想要的子阵列:
from numpy import *
# Create the array
A = zeros((2,3,4,5))
# Identify the subindex you're looking for
idx = ndindex(A.shape[:2])
# Iterate through the array
[(x, A[x].shape) for x in idx]
这给出了预期的结果:
[((0, 0), (4, 5)), ((0, 1), (4, 5)), ((0, 2), (4, 5)), ((1, 0), (4, 5)), ((1, 1), (4, 5)), ((1, 2), (4, 5))]