python - 包含大型矩阵运算的优化函数

Question

我可能无法很好地解释这一点，因此我将使用一个与我的问题有些相似的示例，但这里是:

我需要(重复)计算一个复杂的运算，它是仅几个标量(例如 x1、x2 和 x3)的函数。然后计算一个矩阵，其元素由 x1、x2、x3(以及行和列位置)的函数确定，称为 X。

中间步骤A、B和C涉及几个非常大的矩阵乘积。这些乘积涉及的矩阵的值是恒定的。

假设我们的矩阵大小为:

• A:k×n
• B:n×n
• C:n×k
• X:3×k

其中 n 非常非常大，k 也非常大。

该函数执行以下操作(伪代码):

func = function(x1, x2, x3)
  X = make_X(x1, x2, x3)
return X * A * B * C * X^t

其中 * 是常规点积。输出将只是一个 3 x 3 矩阵!

我(天真地可能)认为必须有某种自动方法可以有效地将其编译为 x1、x2、x3 的 9 个函数——一个对应输出矩阵的每个元素。

我经常使用 numpy/scipy，但没有使用 sympy 或 theano 的经验，尽管它们似乎在大概范围内。关于如何解决这个问题有什么建议吗？

P.S.，解决这个问题的任何代码包最好都在 python 中，但它们不是必须的，只要它们可以从 python 调用即可。

Answer 1

用einsum术语表达问题可能会有所帮助:

np.einsum('ij,jk,kl,lm,nm->in', X, A, B, C, X)

这可以分为两个步骤:

ABC = np.einsum('jk,kl,lm->jm', A, B, C) # k by k
np.einsum('ij,jm,nm->in', X, ABC, X)

因此结果的 i,j 元素为:

R[i,j] = np.einsum('j,jm,m->', X[i,:], X[j,:])

并使用新的@运算符(在本例中只是点的运算符版本)

R = X@A@B@C@(X.T)

对于k,n=10,20，最后一个是最快的。

如果您对 x1,x2,x3 的不同组合执行此操作，但对一组 A,B,C 执行此操作，则执行 ABC=A @B@C 应该可以节省时间。但我对将 X@ABC@(X.T) 分成 9 个步骤的值(value)表示怀疑。 ABC 是 kxk，因此您已经完成了涉及 B 的较大计算。