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

у меня есть 2-мерный массив 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 и для произвольных размеров начального массива (все еще считая, что начальный массив имеет одинаковое количество строк и столбцов)?

9
arrays numpy python

2 ответов

классический случай numpy.kron -

np.kron(np.eye(n), a)

образец выполнения -

In [57]: n = 2

In [58]: np.kron(np.eye(n), a)
Out[58]: 
array([[ 5.,  7.,  0.,  0.],
       [ 6.,  3.,  0.,  0.],
       [ 0.,  0.,  5.,  7.],
       [ 0.,  0.,  6.,  3.]])

In [59]: n = 3

In [60]: np.kron(np.eye(n), a)
Out[60]: 
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.]])
12
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)
0