Индексация NumPy кортежами и Ellipsis: как не попасть в ловушку np.ndindex

Почему кортежи из np.ndindex с Ellipsis ведут себя не так, как ждёте, и как правильно индексировать массивы NumPy: раскрытие кортежей, scores[(*pos, ...)].

Индексировать массивы NumPy с помощью кортежей удобно — пока внезапно не перестаёт быть таковым. Частая ловушка возникает, когда вы итерируетесь с np.ndindex, а затем пытаетесь использовать полученный кортеж вместе с Ellipsis. На вид это похоже на индексацию по нескольким осям, но массив трактует её иначе, что приводит к неожиданным присваиваниям или выборкам.

Исходная ситуация

for pos in np.ndindex(cfg_dims):
    ...
    scores[pos, ...] = ...

Если pos равен (0, 1), ожидается запись в scores[0, 1, ...]. Но на деле это работает так, словно вы вовсе не раскрываете кортеж, поэтому целевой срез не попадает под индексацию.

Что происходит на самом деле

Внутри квадратных скобок Python формирует единый кортеж индексов. Когда вы передаёте pos первым элементом и добавляете Ellipsis вторым, кортеж индексов остаётся вложенным, а не раскладывается по осям. Иными словами, элементы pos не воспринимаются как отдельные измерения. Поэтому индексация попадает не туда, где вы рассчитывали. Нужен «плоский» набор индексов, а не кортеж, вложенный в другой кортеж индексации.

Как исправить

Разверните кортеж, чтобы каждая координата стала отдельным индексом верхнего уровня. Для этого поместите раскрытые значения в ещё один кортеж вместе с Ellipsis.

for pos in np.ndindex(cfg_dims):
    ...
    scores[(*pos, ...)] = ...

В этом конкретном шаблоне, если последним срезом стоит Ellipsis, он не требуется. Если ваш случай соответствует такой форме, можно индексировать напрямую кортежем:

for pos in np.ndindex(cfg_dims):
    ...
    scores[pos] = ...

Есть и более краткая форма индексации, которая раскрывает кортеж прямо в скобках. Она работает в Python 3.11 и новее:

for pos in np.ndindex(cfg_dims):
    ...
    scores[*pos, ...] = ...

Почему это важно

Перебор сеток с np.ndindex — типичный приём в численных задачах, переборах параметров и вычислении поверхностей правдоподобия. Если индексация попадает не на те оси, вы можете незаметно читать или записывать значения в неверный участок массива. Такие ошибки сложно обнаружить визуально, особенно когда код выглядит безупречно. Понимание того, как правильно раскрывать кортежи индексов, помогает сохранить предсказуемость операций с массивами.

Итоги

Когда кортеж координат должен адресовать несколько осей, раскрывайте его так, чтобы каждый элемент был самостоятельным индексом. Если в этом месте нужен Ellipsis, используйте scores[(*pos, ...)]. Если Ellipsis должен стоять последним, упростите до scores[pos]. А там, где доступно, лаконичная запись scores[*pos, ...] — аккуратный вариант для Python 3.11+.