将 (N, 2) 形状的 numpy 数组“重塑”为 (N, 2, 2) ，其中每列（大小 2）成为 diag (2,2) 块？答案

【问题标题】："reshape" numpy array of (N, 2) shape into (N, 2, 2) where each column (size 2) become a diag (2,2) block?将 (N, 2) 形状的 numpy 数组“重塑”为 (N, 2, 2) ，其中每列（大小 2）成为 diag (2,2) 块？
【发布时间】：2021-11-24 13:29:50
【问题描述】：

有没有一种有效的方法来做到这一点？例如我有

[[1, 2, 3],
 [4, 5, 6]]

我想得到：

[[[1, 0],
  [0, 4]],

 [[2, 0],
  [0, 5]],

 [[3, 0],
  [0, 6]]]

【问题讨论】：

标签： python arrays numpy reshape diagonal

【解决方案1】：

对于大型数组，我推荐np.einsum 如下：

>>> data
array([[1, 2, 3],
       [4, 5, 6]])
>>> out = np.zeros((*reversed(data.shape),2),data.dtype)
>>> np.einsum("...ii->...i",out)[...] = data.T
>>> out
array([[[1, 0],
        [0, 4]],

       [[2, 0],
        [0, 5]],

       [[3, 0],
        [0, 6]]])

einsum 创建包含对角元素的内存位置的可写跨步视图。这与 numpy 中的效率差不多。

【讨论】：

感谢您的回答。对于稀疏矩阵有没有一种有效的方法？（原则上不需要额外的内存，对吧？）
scipy 有一个稀疏矩阵类，但是它们只支持二维矩阵。
@ALex，虽然einsum 是out 数组的有效view，但这不是view 的简单data。 data 被复制到 out。

【解决方案2】：

不是一个跨步视图，但也许更容易理解的是这个 (3,2,2) 数组的“对角线”填充：

In [28]: arr = np.arange(1,7).reshape(2,3)
In [29]: res = np.zeros((3,2,2),int)
In [30]: res[:,np.arange(2),np.arange(2)].shape
Out[30]: (3, 2)
In [31]: res[:,np.arange(2),np.arange(2)]=arr.T
In [32]: res
Out[32]: 
array([[[1, 0],
        [0, 4]],

       [[2, 0],
        [0, 5]],

       [[3, 0],
        [0, 6]]])

对于这个小案例，时间并没有太大的不同。我不知道他们将如何扩展：

In [39]: timeit np.einsum("...ii->...i",out)[...] = arr.T
5.21 µs ± 5.99 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [40]: timeit res[:,np.arange(2),np.arange(2)]=arr.T
6.4 µs ± 21.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

【讨论】：