Как правильно извлечь столбец из NumPy ndarray с кортежами индексов

Как корректно извлекать столбец из NumPy ndarray: почему cube[:, idx] включает расширенную индексацию и чем заменить — распаковку кортежа и slice(None).

Извлечь один «столбец» из NumPy ndarray кажется простым, пока не нужно сделать это универсально для массивов разной формы. Если ваш формат файла принимает только по одному столбцу и требуется включать индексы координат (например, y и z) в имя столбца, жёстко фиксировать размерности не масштабируется. Сложность проявляется, когда вы пытаетесь совместить ведущий ":" с кортежем индексов и ожидаете поведения как у раздельных позиционных индексов.

Постановка задачи

Рассмотрим трёхмерный массив, где x — это строки, y — столбцы, а z — глубина. Итерируясь по верхнему x-срезу, вы получаете все пары (y, z), нужные для именования, и для каждой пары хотите забрать полный x-столбец.

import numpy as np
grid = np.arange(5000).reshape(10, 4, 125)
if grid.ndim > 1:
    for idx, _ in np.ndenumerate(grid[0, ...]):
        vec = grid[:, idx[0], idx[1]]
        label = "_" + "_".join(str(v) for v in idx)
        print("Vp" + label)

Явная индексация через idx[0] и idx[1] возвращает 500 векторов формы (10,), что как раз и нужно. Логичный следующий шаг — обобщить решение, передавая кортеж напрямую:

import numpy as np
cube = np.arange(5000).reshape(10, 4, 125)
if cube.ndim > 1:
    for idx, _ in np.ndenumerate(cube[0, ...]):
        col = cube[:, idx]  # naive generic attempt
        tag = "_" + "_".join(str(u) for u in idx)
        print("Vp" + tag)

Это «почти работает», но вместо формы (10,) вы получаете массивы вида (10, 2, 125), а затем сталкиваетесь с IndexError: index 4 is out of bounds for axis 1 with size 4.

Почему это ломается

Когда вы пишете cube[:, idx], вы индексируете кортежом, у которого первый элемент — срез, а второй элемент — сам кортеж. Такой шаблон включает расширенную индексацию NumPy. У неё другие правила выбора, чем у базовой позиционной индексации, поэтому cube[:, (y, z)] не эквивалентно cube[:, y, z]. В итоге получается неожиданная форма, а по мере продолжения итерации возникает выход за границы: логика расширенной индексации применяет кортеж не так, как вы задумали.

Решение: распаковать кортеж или собрать индекс явно

Надёжный способ совместить «все строки» с кортежем оставшихся индексов — распаковать кортеж на текущем уровне. Оператор звёздочки в кортеже индексов сохраняет стандартную семантику позиционной индексации:

import numpy as np
arr = np.arange(5000).reshape(10, 4, 125)
if arr.ndim > 1:
    for yz, _ in np.ndenumerate(arr[0, ...]):
        out_vec = arr[:, *yz]
        name = "_" + "_".join(str(k) for k in yz)
        print("Vp" + name)

Если удобнее собирать полный кортеж индексов программно, замените : объектом slice(None) и сконкатенируйте его с вашим кортежем. Это наиболее общий подход, который естественно масштабируется вместе с логикой построения размерностей:

import numpy as np
block = np.arange(5000).reshape(10, 4, 125)
if block.ndim > 1:
    for yz, _ in np.ndenumerate(block[0, ...]):
        key = (slice(None),) + yz
        column = block[key]
        tag = "_" + "_".join(str(p) for p in yz)
        print("Vp" + tag)

Чтобы показать, как включить «полный срез» внутри индекса по нескольким осям, не записывая : буквально внутри кортежа, используйте slice(None). Например:

import numpy as np
A = np.arange(27).reshape(3, 3, 3)
key = (1, slice(None), 1)
sel = A[key]

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

Когда нужно по одному столбцу передавать данные из массивов разной размерности, путаница между расширенной индексацией и базовой позиционной быстро ломает формы и приводит к скрытым ошибкам. Распаковка кортежа или явная сборка кортежа индексов с slice(None) гарантируют, что : интерпретируется как «все строки», а последующие индексы привязываются к нужным осям. Такой подход одинаково чисто работает для 1D, 2D и 3D массивов и масштабируется под тот же паттерн итерации без разветвлений по форме.

Выводы

Если у вас есть кортеж координат и нужно объединить его с ведущим выбором «все строки», не передавайте кортеж одним элементом после среза. Либо распакуйте его прямо в выражении индексации через *, либо постройте новый кортеж индексов, начинающийся с slice(None). Оба варианта сохраняют семантику базовой индексации и дают ожидаемый вектор формы (число_строк,) для каждой пары координат, упрощая последующее именование и по-столбцовый вывод.