python - 使用 numpy 创建任意形状的单位矩阵

Question

是否有更快/内置的方法来生成第一个维度中具有任意形状并在最后一个维度中具有恒等性的恒等矩阵？

import numpy as np

base_shape = (10, 11, 12)
n_dim = 4

# m = 2
frames2d = np.zeros(base_shape + (n_dim, n_dim))
for i in range(n_dim):
    frames2d[..., i, i] = 1

# m = 3
frames3d = np.zeros(base_shape + (n_dim, n_dim, n_dim))
for i in range(n_dim):
    frames3d[..., i, i, i] = 1

Answer 1

方法#1

我们可以利用np.einsum灵感来自 this post 的对角 View 因此为我们想要的输出分配 1s 。因此，对于 m=3 情况，在用零初始化后，我们可以简单地执行 -

diag_view = np.einsum('...iii->...i',frames3d)
diag_view[:] = 1

概括为包括这些输入参数，它将是 -

def ndeye_einsum(base_shape, n_dim, m):
    out = np.zeros(list(base_shape) +  [n_dim]*m)
    diag_view = np.einsum('...'+'i'*m+'->...i',out)
    diag_view[:] = 1
    return out

因此，要重现这些相同的数组，那就是 -

frames2d = ndeye_einsum(base_shape, n_dim, m=2)
frames3d = ndeye_einsum(base_shape, n_dim, m=3)

方法#2

同样，从同一篇链接文章中，我们还可以将形状 reshape 为 2D 并沿着列分配到步长切片数组中，如下所示 -

def ndeye_reshape(base_shape, n_dim, m):
    N = (n_dim**np.arange(m)).sum()
    out = np.zeros(list(base_shape) +  [n_dim]*m)
    out.reshape(-1,n_dim**m)[:,::N] = 1
    return out

这再次适用于 View ，因此应该与方法 #1 同等有效。

方法#3

另一种方法是使用基于整数的索引。因此，例如，一次性分配到 frames3d 中，它将是 -

I = np.arange(n_dim)
frames3d[..., I, I, I] = 1

概括后变成 -

def ndeye_ellipsis_indexer(base_shape, n_dim, m):
    I = np.arange(n_dim)
    indexer = tuple([Ellipsis]+[I]*m)
    out = np.zeros(list(base_shape) +  [n_dim]*m)
    out[indexer] = 1
    return out

<小时/>

通过 View 扩展到更高的尺寸

沿着base_shape的暗淡基本上是最后m暗淡元素的复制。因此，我们可以使用 np.broadcast_to 将这些较高的亮度作为较高亮度的数组 View 来获取。我们将基本上创建一个 m-dim 身份数组，然后将 View 广播到更高的维度。这适用于之前发布的所有三种方法。为了演示如何在基于 einsum 的解决方案上使用它，我们将 -

# Create m-dim "trailing-base" array, basically a m-dim identity array
def ndeye_einsum_trailingbase(n_dim, m):
    out = np.zeros([n_dim]*m)
    diag_view = np.einsum('i'*m+'->...i',out)
    diag_view[:] = 1
    return out

def ndeye_einsum_view(base_shape, n_dim, m):
    trail_base = ndeye_einsum_trailingbase(n_dim, m)
    return np.broadcast_to(trail_base, list(base_shape) +  [n_dim]*m)

因此，我们再次会得到，例如-

frames3d = ndeye_einsum_view(base_shape, n_dim, m=3)

这将是 m-dim 数组的 View ，因此在内存和性能方面都很高效。