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

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

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

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

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

(Обычно, когда я хочу перебрать многомерный объект 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 -мерного подмассива ] 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 .