沿任意轴连接未知维度的numpy数组

Question

我有两个未知维度的数组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。

我的做法是将形状扩展到第N维，并添加 1 直到要连接的数组具有相同的形状:

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)的数组沿着axis 3连接起来。

有更好的方法实现这个功能吗？

ps: 根据第一个回答的建议进行了更新。

- Miguel

Answer 1

这应该可以工作：

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)

- Eric

---

我本以为我需要手动编写一堆重塑逻辑才能实现atleast_nd函数，但这样太好了。谢谢！ - tel

Answer 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)

- Lee

---

expand_dims 的核心是一个重塑操作：a.reshape(shape[:axis] + (1,) + shape[axis:]。 - hpaulj

Answer 3

我认为你的方法没有问题，尽管你可以让你的代码更加紧凑:

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

- Jaime

---

谢谢！这样好多了。不过，我正在寻找的是一些避免显式构建新形状的解决方案。vstack和dstack仅适用于2D和3D数组。 - Miguel