我正在尝试编写基于矩阵的快速、优化的代码,并且最近发现 einsum 作为实现显着加速的工具。
是否可以使用它来有效地设置多维数组的对角线,或者它只能返回数据?
在我的问题中,我尝试通过对每个方阵 (N x N) 矩阵中的列求和来设置方阵数组(形状:M x N x N)的对角线。
我当前的(缓慢的、基于循环的)解决方案是:
# Build dummy array
dimx = 2 # Dimension x (likely to be < 100)
dimy = 3 # Dimension y (likely to be between 2 and 10)
M = np.random.randint(low=1, high=9, size=[dimx, dimy, dimy])
# Blank the diagonals so we can see the intended effect
np.fill_diagonal(M[0], 0)
np.fill_diagonal(M[1], 0)
# Compute diagonals based on summing columns
diags = np.einsum('ijk->ik', M)
# Set the diagonal for each matrix
# THIS IS LOW. CAN IT BE IMPROVED?
for i in range(len(M)):
np.fill_diagonal(M[i], diags[i])
# Print result
M
请问这个可以改进吗?似乎 np.fill_diagonal 不接受非方阵(因此强制我基于循环的解决方案)。也许 einsum 也能帮上忙?
一种方法是重塑2D
,按步骤设置列ncols+1
与对角线值。重塑会创建一个视图,因此我们可以直接访问这些对角线位置。因此,实施将是 -
s0,s1,s2 = M.shape
M.reshape(s0,-1)[:,::s2+1] = diags
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)