可以在numpy中优化子矩阵操作吗？

Question

我正在做一个项目，我正在对一个大尺寸的方阵（180x180）的对角子矩阵进行加法，其“内核”方阵的大小约为40*40

size1 = matrix1.shape[0] #square matrix shape 180x180
sizeKernel = kernelMatrix.shape[0] # square matrix shape 40x40
newMatrix1 = matrix1.copy()
np.fill_diagonal(newMatrix1, 0)
for i in range(size1):
      newMatrix1[i:i+sizeKernel,i:i+sizeKernel] += matrix1[i][i] * kernelMatrix[:size1-i, :size1-i]

目标是将对角系数与核矩阵一起分配给其他系数（一种部分卷积）

例如，假设矩阵大小为4x4，内核大小为2x2：

a = [[1,2,1,0],
     [3,-1,0,0],
     [0,1,1,1],
     [0,2,0,2]]

kern = [[1,2],
        [0,1]]

我们通过复制a的每个系数来创建矩阵res，但对角线系数：

res = [[0,2,1,0],
       [3,0,0,0],
       [0,1,0,1],
       [0,2,0,0]]

我们将每个对角系数分配到与内核大小相同的右下角子矩阵（如果不可能，我们只取剩余的并切割内核矩阵） 对于对角线的第一个系数，它给出：

res[0:2, 0:2] = res[0:2,0:2] + a[0,0]*kern

res[0:2, 0:2] = [[0,2]   +   1 * [[1,2],
                 [3,0]]           [0,1]]
                 
res = [[1,4,1,0],
       [3,1,0,0],
       [0,1,0,1],
       [0,2,0,0]]

然后我们对其他COEF做同样的想法。对于对角线的最后一个coef，我们不采用整个核，而是采用相同的方法将矩阵的元素扩展为0

因为它不是矢量化的，所以计算它需要相当长的时间

另一种计算方法是直接对提取的对角矩阵使用scipy方法fftconvolve或convolve2d，并将结果添加到对角为0的矩阵中：

newMatrix = matrix1.copy()
matrix1Size = matrix1.shape[0]
np.fill_diagonal(newMatrix , 0)
newMatrix += scipy.signal.fftconvolve(matrix1, kernelMatrix)[:matrix1Size,:matrixSize]

但是它所花费的时间与第一种方法差不多（当然，如果我们对整个矩阵进行完全卷积，它们的计算量会更多）

有什么方法可以将子矩阵的计算矢量化吗

提前谢谢你