NumPy 中的单位矩阵堆叠

Question

我需要一个 NumPy 中的 2n x n 矩阵，该矩阵由相互堆叠的 n x n 单位矩阵和负的 n x n 单位矩阵组成。

这是我最初的解决方案，效果很好。

def id_stack(n): 
    id_ = np.identity(n) 
    return np.vstack((id_, -id_))

id_stack(3)
# array([[ 1.,  0.,  0.],
#        [ 0.,  1.,  0.],
#        [ 0.,  0.,  1.],
#        [-1., -0., -0.],
#        [-0., -1., -0.],
#        [-0., -0., -1.]])

然后我想我可以改为设置对角线并像这样更快，这也有效。

def id_stack2(n): 
    full = np.zeros((2*n, n)) 
    rng = np.arange(n) 
    full[rng, rng] = 1 
    full[rng + n, rng] = -1     
    return full

我想知道是否有更快的方法来实现这一点，也许使用某种跨步技巧？

Answer 1

您可能从自己的示例中注意到，分配一个大缓冲区并在其中设置元素通常比分配两个较小的缓冲区和一个大缓冲区来复制它们要快。

关于 numpy 的巧妙之处在于，您可以在不分配新数组的情况下将视图获取到同一个缓冲区。例如：

 output = np.zeros((2 * n, n))

在这种情况下一个有用的视图是

flat = output.ravel()

您可以将每个 n + 1st 元素设置为 1，从第一个开始，平展视图中总共有 n 个元素，-1 也类似。这只需要对 raveled 视图进行简单的索引操作：

output[:n * n:n + 1] = 1
output[n * n::n + 1] = -1

这避免了创建全范围数组和触发高级索引语义，这些语义也更占用内存。