Перечисление numpy.ndarray по правильному подмножеству измерений?

Question

(В этом посте пусть np будет сокращением для numpy.)

Предположим, a является ( n + k ) & # x2011; размерным np.ndarray объектом, для некоторых целых чисел n > 1 и k > 1. (IOW, n + k > 3 - значение a.ndim). Я хочу перечислить a по первым n измерениям; это означает, что на каждой итерации перечислитель / итератор создает пару, первым элементом которой является кортеж ii из n индексов, а вторым элементом является k & # x2011; размерный суб- ndarray при a[ii].

Конечно, не сложно написать код функции для этого (на самом деле, я приведу пример такой функции ниже), но я хочу знать это:

обеспечивает ли numpy какой-либо специальный синтаксис или функции для выполнения этого типа "частичного" перечисления?

(Обычно, когда я хочу перебрать многомерный объект np.ndarray, я использую np.ndenumerate, но здесь это не поможет, потому что (насколько я могу судить) np.ndenumerate будет перебирать все n + k размеры.)

Если предположить, что ответ на поставленный выше вопрос - да, тогда есть продолжение:

как насчет случая, когда размерности n для итерации не являются смежными?

(В этом случае первый элемент пары, возвращаемый перечислителем / итератором на каждой итерации, будет набором из r > n элементов, некоторые из которых будут специальное значение, обозначающее «все», например, slice(None); вторым элементом этой пары по-прежнему будет ndarray длины k .)

Спасибо!

---

Надеемся, что приведенный ниже код проясняет спецификацию проблемы. Функция partial_enumerate делает то, что я хотел бы сделать, используя любые специальные конструкции numpy, доступные для этой цели. После определения partial_enumerate приведен простой пример для случая n = k = 2.

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 координат в a, а вторая представляет форму k & # x2011; размерный подмассив 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 итерациям, каждая из которых посещает отдельную ячейку a.

Answer 1

Вы можете использовать правила простого вещания для создания декартового произведения.Функция 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_ позволяет использовать непоследовательные индексы.Теперь, когда мы применяем правила простого вещания, мы получаем декартово произведение:

>>> 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]]]])

Однако caveat emptor .Широковещательные массивы полезны для индексации, но если вы буквально хотите перечислить значений, вам может быть лучше использовать 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

Так что если вы включаете дляцикл в любом случае, тогда itertools.product, скорее всего, будет быстрее.Тем не менее, вы можете использовать описанные выше методы, чтобы получить некоторые похожие структуры данных в чистом виде:

>> 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')])

Answer 2

Я думаю, вы ищете функцию ndindex в numpy.Просто возьмите часть нужного подмассива:

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))]