Как нарезать определенное измерение массива np с другой начальной точкой?

Question

Пример кода приведен ниже.

Я хочу получить dataNew(h, w, length) в соответствии с data(h, w, c) и ind(h, w). Здесь length < c, это означает, что dataNew вырезано из данных.

Здесь length и ind[i, j] обязательно соответствуют значению c.

Я понял это через циклы, но я поймал путь питона. Пожалуйста, помогите, спасибо!

import numpy as np

h, w, c = 3, 4, 5

data = np.arange(60).reshape((h, w, c))
print(data)

length = 3
ind = np.random.randint(0, 3, 12).reshape(h, w)
print(ind)

dataNew = np.empty((h, w, length), np.int16)
for i in range(h):
    for j in range(w):
        st = ind[i, j]
        dataNew[i, j] = data[i, j][st : st + length]

print(dataNew)

Answer 1

Одним из способов было бы создание индексного массива с использованием broadcasting и использование np.take_along_axis для индексации массива:

ix = ind[...,None] + np.arange(length)
np.take_along_axis(data, ix, -1)

Answer 2

Мы можем использовать np.lib.stride_tricks.as_strided на основе scikit-image's view_as_windows, чтобы получить раздвижные окна. Более подробная информация об использовании as_strided на основе view_as_windows.

from skimage.util.shape import view_as_windows

# Get all sliding windows along the last axis
w = view_as_windows(data,(1,1,length))[...,0,0,:]

# Index into windows with start indices and slice out singleton dims 
out = np.take_along_axis(w,ind[...,None,None],axis=-1)[...,0]

Последний шаг в основном заключается в использовании расширенной индексации в окнах с этими начальными индексами.Это можно сделать немного проще и легче понять.Итак, в качестве альтернативы, мы могли бы сделать -

m,n = ind.shape
I,J = np.ogrid[:m,:n]
out = w[I,J,ind]