[英]Create block diagonal numpy array from a given numpy array
我有一個二維的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
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]])
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)
對於矩陣的特殊情況,簡單的切片比numpy.kron()
(最慢)快得多,並且與基於numpy.einsum()
的方法差不多(來自@Divakar答案)。 與scipy.linalg.block_diag()
相比,它對於較小的arr
表現更好,在某種程度上與塊重復次數無關。
請注意,使用block_diag_view()
可以輕松地進一步改善在較小的輸入上的block_diag_view()
的性能,但是對於較大的輸入,則性能會變差。
使用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
基准:
NUM = 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()
的方法和具有顯式循環的另一個Numba版本。)
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.