Как сгладить массив в матрицу в Numpy?

Question

Я ищу элегантный способ сглаживания массива произвольной формы в матрицу, основанную на одном параметре, который задает размер для сохранения.Для иллюстрации я бы хотел

def my_func(input, dim):
    # code to compute output
    return output

. Например, input массив формы 2x3x4, output должен быть для dim=0 массив формы 12x2;для dim=1 массив формы 8x3;для dim=2 массив формы 6x8.Если я хочу выровнять только последнее измерение, то это легко сделать с помощью

input.reshape(-1, input.shape[-1])

Но я бы хотел добавить функциональность добавления dim (элегантно, безвсе возможные случаи + проверка условий и т. д.).Это может быть возможно сначала путем замены измерений, так чтобы интересующее измерение тянулось, а затем применялась указанная выше операция.

Любая помощь?

Answer 1

Мы можем переставить оси и изменить их форму -

# a is input array; axis is input axis/dim
np.moveaxis(a,axis,-1).reshape(-1,a.shape[axis])

Функционально это, в основном, отодвигает указанную ось назад , а затем изменяет форму, сохраняя длину этой оси, образуя вторую ось и объединяяостальные оси образуют первую ось.

Примеры прогонов -

In [32]: a = np.random.rand(2,3,4)

In [33]: axis = 0

In [34]: np.moveaxis(a,axis,-1).reshape(-1,a.shape[axis]).shape
Out[34]: (12, 2)

In [35]: axis = 1

In [36]: np.moveaxis(a,axis,-1).reshape(-1,a.shape[axis]).shape
Out[36]: (8, 3)

In [37]: axis = 2

In [38]: np.moveaxis(a,axis,-1).reshape(-1,a.shape[axis]).shape
Out[38]: (6, 4)