Конкатенация массивов numpy неизвестного размера вдоль произвольной оси

У меня есть два массива A и B неизвестных измерений, которые я хочу объединить по размеру N. Например:

>>> A = rand(2,2)       # just for illustration, dimensions should be unknown
>>> B = rand(2,2)       # idem
>>> N = 5

>>> C = concatenate((A, B), axis=N)
numpy.core._internal.AxisError: axis 5 is out of bounds for array of dimension 2

>>> C = stack((A, B), axis=N)
numpy.core._internal.AxisError: axis 5 is out of bounds for array of dimension 3

Связанный вопрос задается здесь. К сожалению, предлагаемые решения не работают, когда размеры неизвестны, и нам может потребоваться добавить несколько новых осей до получения минимального размера N.

То, что я сделал, состоит в том, чтобы расширить форму с помощью 1 до размера N th и затем объединить:

newshapeA = A.shape + (1,) * (N + 1 - A.ndim)
newshapeB = B.shape + (1,) * (N + 1 - B.ndim)
concatenate((A.reshape(newshapeA), B.reshape(newshapeB)), axis=N)

С помощью этого кода я должен был бы, например, объединить массив (2,2,1,3) с массивом (2,2) вдоль оси 3.

Есть ли лучшие способы достижения этого?

ps: обновлено, так как предложен первый ответ.

Ответ 1

Я не думаю, что с вашим подходом что-то не так, хотя вы можете сделать свой код немного более компактным:

newshapeA = A.shape + (1,) * (N + 1 - A.ndim)

Ответ 2

Альтернатива, используя numpy.expand_dims:

>>> import numpy as np
>>> A = np.random.rand(2,2)
>>> B = np.random.rand(2,2)
>>> N=5


>>> while A.ndim < N:
        A= np.expand_dims(A,x)
>>> while B.ndim < N:
        B= np.expand_dims(B,x)
>>> np.concatenate((A,B),axis=N-1)

Ответ 3

Это должно работать:

def atleast_nd(x, n):
    return np.array(x, ndmin=n, subok=True, copy=False)

np.concatenate((atleast_nd(a, N+1), atleast_nd(b, N+1)), axis=N)