Создайте единичные матрицы произвольной формы с 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

Подход № 1

Мы можем использовать np.einsum для диагонального обзора, вдохновленного this post, и, следовательно, назначить 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)

Approach # 2

Опять же, из того же связанного поста, мы также можем изменить форму на 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 в one- go, это будет -

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

---

Расширение до более высокого dims с видом

dims вдоль base_shape в основном являются копиями элементов из последних m dims. Таким образом, мы можем получить эти более высокие значения dim как представление массива с более высоким dim с помощью np.broadcast_to. Мы создадим в основном массив идентификаторов m-dim, а затем представим широковещательный просмотр в более высокие значения dim. Это будет применимо ко всем трем подходам, опубликованным ранее. Чтобы продемонстрировать, как использовать его в решении на основе 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 и, следовательно, эффективно использовать как память, так и производительность.

Answer 2

Один из подходов к созданию единичной матрицы вдоль последних двух измерений массива состоит в использовании 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.]],
        ...