Создать массив диагональной матрицы блоков из заданного массива numpy

У меня есть двумерный массив numpy с равным количеством столбцов и строк. Я хотел бы разместить их в более массивном массиве с меньшими по диагонали. Должно быть возможно указать, как часто исходная матрица должна быть на диагонали. Например:

a = numpy.array([[5, 7], 
                 [6, 3]])

Итак, если бы я хотел, чтобы этот массив 2 раза по диагонали, желаемый результат был бы:

array([[5, 7, 0, 0], 
       [6, 3, 0, 0], 
       [0, 0, 5, 7], 
       [0, 0, 6, 3]])

В 3 раза:

array([[5, 7, 0, 0, 0, 0], 
       [6, 3, 0, 0, 0, 0], 
       [0, 0, 5, 7, 0, 0], 
       [0, 0, 6, 3, 0, 0],
       [0, 0, 0, 0, 5, 7],
       [0, 0, 0, 0, 6, 3]])

Есть ли быстрый способ реализовать это с помощью методов numpy и для произвольных размеров стартового массива (все еще учитывая, что начальный массив имеет одинаковое количество строк и столбцов)?

ОТВЕТЫ

Ответ 1

Подход № 1

Классический футляр numpy.kron -

np.kron(np.eye(r,dtype=int),a) # r is number of repeats

Пробный прогон -

In [184]: a
Out[184]: 
array([[1, 2, 3],
       [3, 4, 5]])

In [185]: r = 3 # number of repeats

In [186]: np.kron(np.eye(r,dtype=int),a)
Out[186]: 
array([[1, 2, 3, 0, 0, 0, 0, 0, 0],
       [3, 4, 5, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 2, 3, 0, 0, 0],
       [0, 0, 0, 3, 4, 5, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 1, 2, 3],
       [0, 0, 0, 0, 0, 0, 3, 4, 5]])

Подход № 2

Еще один эффективный с diagonal-viewed-array-assignment -

def repeat_along_diag(a, r):
    m,n = a.shape
    out = np.zeros((r,m,r,n), dtype=a.dtype)
    diag = np.einsum('ijik->ijk',out)
    diag[:] = a
    return out.reshape(-1,n*r)

Пробный прогон -

In [188]: repeat_along_diag(a,3)
Out[188]: 
array([[1, 2, 3, 0, 0, 0, 0, 0, 0],
       [3, 4, 5, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 2, 3, 0, 0, 0],
       [0, 0, 0, 3, 4, 5, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 1, 2, 3],
       [0, 0, 0, 0, 0, 0, 3, 4, 5]])

Ответ 2

import numpy as np
from scipy.linalg import block_diag

a = np.array([[5, 7], 
              [6, 3]])

n = 3

d = block_diag(*([a] * n))

d

array([[5, 7, 0, 0, 0, 0],
       [6, 3, 0, 0, 0, 0],
       [0, 0, 5, 7, 0, 0],
       [0, 0, 6, 3, 0, 0],
       [0, 0, 0, 0, 5, 7],
       [0, 0, 0, 0, 6, 3]], dtype=int32)

Но похоже, что решение np.kron немного быстрее для маленьких n.

%timeit np.kron(np.eye(n), a)
12.4 µs ± 95.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit block_diag(*([a] * n))
19.2 µs ± 34.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Однако при n = 300, например, block_diag намного быстрее:

%timeit block_diag(*([a] * n))
1.11 ms ± 32.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.kron(np.eye(n), a)
4.87 ms ± 31 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Ответ 3

Для специализированного случая матриц простое срезание происходит НАМНОГО быстрее, чем numpy.kron() (самое медленное) и в основном на уровне подхода numpy.einsum() -based (из ответа @Divakar). По сравнению с scipy.linalg.block_diag() он работает лучше для меньших arr, в некоторой степени независимо от количества повторений блоков.

Обратите внимание, что производительность block_diag_view() на меньших входах может быть легко улучшена с помощью Numba, но можно получить худшие характеристики для больших входов.

С Numba, полным явным циклом и распараллеливанием (block_diag_loop_jit()) можно было бы снова получить результаты, аналогичные block_diag_einsum(), если число повторений невелико.

В целом, наиболее эффективные решения - block_diag_einsum() и block_diag_view().

import numpy as np
import scipy as sp
import numba as nb

import scipy.linalg


NUM = 4
M = 9


def block_diag_kron(arr, num=NUM):
    return np.kron(np.eye(num), arr)


def block_diag_einsum(arr, num=NUM):
    rows, cols = arr.shape
    result = np.zeros((num, rows, num, cols), dtype=arr.dtype)
    diag = np.einsum('ijik->ijk', result)
    diag[:] = arr
    return result.reshape(rows * num, cols * num)


def block_diag_scipy(arr, num=NUM):
    return sp.linalg.block_diag(*([arr] * num))


def block_diag_view(arr, num=NUM):
    rows, cols = arr.shape
    result = np.zeros((num * rows, num * cols), dtype=arr.dtype)
    for k in range(num):
        result[k * rows:(k + 1) * rows, k * cols:(k + 1) * cols] = arr
    return result


@nb.jit
def block_diag_view_jit(arr, num=NUM):
    rows, cols = arr.shape
    result = np.zeros((num * rows, num * cols), dtype=arr.dtype)
    for k in range(num):
        result[k * rows:(k + 1) * rows, k * cols:(k + 1) * cols] = arr
    return result


@nb.jit(parallel=True)
def block_diag_loop_jit(arr, num=NUM):
    rows, cols = arr.shape
    result = np.zeros((num * rows, num * cols), dtype=arr.dtype)
    for k in nb.prange(num):
        for i in nb.prange(rows):
            for j in nb.prange(cols):
                result[i + (rows * k), j + (cols * k)] = arr[i, j]
    return result

Тесты для NUM = 4:

bm_4

Тесты для NUM = 400:

bm_400

Графики были созданы из этого шаблона с использованием следующего дополнительного кода:

def gen_input(n):
    return np.random.randint(1, M, (n, n))


def equal_output(a, b):
    return np.all(a == b)


funcs = block_diag_kron, block_diag_scipy, block_diag_view, block_diag_jit


input_sizes = tuple(int(2 ** (2 + (3 * i) / 4)) for i in range(13))
print('Input Sizes:\n', input_sizes, '\n')


runtimes, input_sizes, labels, results = benchmark(
    funcs, gen_input=gen_input, equal_output=equal_output,
    input_sizes=input_sizes)


plot_benchmarks(runtimes, input_sizes, labels, units='ms')

(ИЗМЕНЕНО, чтобы включить подход np.einsum() -based и другую версию Numba с явным циклом.)