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

Question

是否有更快/内置的方法来生成在第一维具有任意形状并在最后 m 维具有恒等式的单位矩阵？

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

沿数组的最后两个维度 设置单位矩阵的一种方法是使用np.broadcast_to 并指定ndarray 应具有的结果形状（这不会推广到更高的维度):

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

frame2d = np.broadcast_to(np.eye(n_dim), a.shape+(n_dim,)*2)

print(frame2d.shape)
# (10, 11, 12, 4, 4)

print(frame2d)

array([[[[[1., 0., 0., 0.],
          [0., 1., 0., 0.],
          [0., 0., 1., 0.],
          [0., 0., 0., 1.]],

         [[1., 0., 0., 0.],
          [0., 1., 0., 0.],
          [0., 0., 1., 0.],
          [0., 0., 0., 1.]],
        ...

Answer 2

方法 #1

我们可以利用 np.einsum for a diagonal view inspired by 并因此将 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

同样，从相同的链接 post，我们还可以重塑为 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

这再次适用于视图，因此应该与方法 #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

扩展到高亮度视图

沿 base_shape 的暗淡基本上是来自最后 m 暗淡的元素的复制。因此，我们可以使用 np.broadcast_to 将那些更高的暗淡作为更高暗淡的阵列视图。我们将基本上创建一个 m-dim 标识数组，然后将视图广播到更高的 dims。这将适用于之前 post 编辑的所有三种方法。为了演示如何在基于 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 数组的视图，因此在内存和性能上都很高效。

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

Create identity matrices with arbitrary shape with numpy

python

identity

numpy

方法 #1

方法 #2

方法 #3

扩展到高亮度视图