определите массив 3d numpy с диапазоном в столбцовом формате

Question

Я хочу определить массив 3d numpy с shape = [3, 3, 5], а также со значениями в виде диапазона, начинающегося с 11 и step = 3 по столбцам. Я имею в виду:

B[:,:,0] = [[ 11,  20,  29],
            [ 14,  23,  32],
            [17,   26,  35]]

B[:,:,1] = [[ 38,  ...],
            [ 41,  ...],
            [ 44,  ...]]
...

Я новичок в numpy и сомневаюсь, что могу сделать это с np.arange или np.mgrid, может быть. но я не знаю, как это сделать.

Как это можно сделать с минимальными строками кода?

Answer 1

Вы можете рассчитать конец диапазона, умножив форму на шаг и добавив начало. Тогда это просто изменить форму и транспонировать, чтобы переместить столбец вокруг:

start = 11
step = 3
shape = [5, 3, 3]

end = np.prod(shape) * step + start

B = np.arange(start, end, step).reshape([5, 3, 3]).transpose(2, 1, 0)
B[:, :, 0]

# array([[11, 20, 29],
#        [14, 23, 32],
#        [17, 26, 35]])