广播 - 给定矩阵基础的 3D 系数场到矩阵的 3D 场

Question

我有一个（大）4D 数组，由矩阵字段的给定基础中的 5 个系数组成。给定 5 个基矩阵，我想有效地计算矩阵字段。

系数字段c[x,y,z,i]是第i个系数在位置x,y,z的值

矩阵字段M[x,y,z,a,b]是(3,3)位置x,y,z的矩阵

基矩阵T_1,...T_5，即(3,3)基矩阵

我可以遍历空间中的每个位置：

M[x,y,z,:,:] = T_1[:,:]*c[x,y,z,0] + T_2[:,:]*c[x,y,z,1]...T_5[:,:]*c[x,y,z,4]

但这是非常低效的。我尝试使用 np.multiply,np.sum 会导致广播错误，因为所需的产品是 3x3 矩阵的字段。

Answer 1

请记住，对于numpy，这些 4 维和 5 维数组只是这样，而不是包含 2 维矩阵等的 3 维数组。

让我们尝试以一种明确维度的方式编写您的计算：

M[x,y,z] = T_1*c[x,y,z,0] + T_2*c[x,y,z,1]...T_5*c[x,y,z,4]

M[x,y,z,:,:] = T_1[:,:]*c[x,y,z,0] + T_2[:,:]*c[x,y,z,1]...T_5[:,:]*c[x,y,z,4]

c[x,y,z,i] 是一个系数，对吧？那么M 是T_n 数组的加权和？

一种表达方式是：

T = np.stack([T_1, T_2, ...T_5], axis=0)   # 3d  (nab)
M = np.einsum('nab,xyzn->xyzab', T, c)

我们也可以在新的最后一个轴上堆叠T_i

T = np.stack([T_1, T_2 ...T_5], axis=2)   # (abn)
M = np.einsum('abn,xyzn->xyzab', T, c)

或广播乘法加和：

M = (T[None,None,None,:,:,:] * c[:,:,:,None,None,:]).sum(axis=-1)

我在写这段代码没有测试，所以可能会有错误，但我认为基本的大纲是对的。

它也可以写成dot，如果我可以将n 维度放在一个参数的最后，而将第二个维度放在另一个参数中。或tensordot。但对其他维度的广播控制较少。

对于测试计算，您还可以重新调整这些数组，以便 x,y,z 合并为一个，a,b 合并为另一个，例如

 M[xyz,:] = T_n[ab]*c[xyz,n]   # etc