python - 是否可以“矢量化”连续的矩阵乘法？

Question

我正在“矢量化”我的一个宠物项目。函数“propagateXi”传播（求解）一个随时间变化的线性方程组，并解决约 20% 的项目成本。

问题实例之间的大小差异很大。“典型”集合是：N = 501，n = 4，s = 3。

这是代码的简化：

import numpy

class Propagator():

  def __init__(self,mode):
    # lots of code here
    pass

  def propagateXi(self,Xi,terms):
    """ This function propagates the linear system, called many times with different initial conditions.
        The initial conditions were previously set.
        The "arc" represents a kind of "realization"; 
        the arcs are essentially independent from one another.

        Sizes: "vector dimensions": 2*n, "time dimension": N, "arcs dimension": s 
         - Xi (N, 2*n, s) stores the (vector) state in each time and each "arc".
         - self.MainPropMat (N, 2*n, 2*n, s) stores the 2n x 2n matrices 
         - terms (N, 2*n, s) stores the "forcing terms" of the equation.

        The equation would be something like:
          \xi_{k+1, arc} = MainPropMat_{k, arc} * \xi_{k, arc} + terms_{k, arc}  \forall k, arc
    """
    
    for k in range(self.N - 1):
      Xi[k+1,:,:] = numpy.einsum('ijs,js->is', self.MainPropMat[k, :, :, :], Xi[k, :, :]) +\
                      terms[k, :, :]
    return Xi

我试图将问题重新解释为巨大的 2nN x 2nN 矩阵乘以包含初始条件和强制项的长 2nN 列。在这些条件下，代码是这样的：

import numpy

class Propagator():

  def __init__(self,mode):
    # lots of code here
    self.prepareBigMat()

  def prepareBigMat(self)
    """
    Prepare the "big matrix" for Xi propagation. Called only once.

    This function assembles the "big matrix" (s x 2*n*N x 2*n*N !) for propagating
    the Xi differential equation in a single step.
    :return: None
    """

    n2 = 2 * self.n
    BigMat = numpy.empty((self.s,self.N * n2,self.N * n2))
    I = numpy.eye(n2*self.N)
    for arc in range(self.s):
      BigMat[arc,:,:] = I
    for k in range(1,self.N):
      kn2 = k * n2
      BigMat[:, kn2:(kn2+n2) , 0:kn2] = numpy.einsum('ijs,sjk->sik',
                                            self.MainPropMat[k-1,:,:,:],
                                            BigMat[:,(kn2-n2):k*n2, 0:kn2])
    self.BigMat = BigMat

  def propagateXi(self,Xi,terms):
    
    n2 = n * 2
    # Assemble the big column
    BigCol = numpy.empty((n2 * N, s))
    # Initial conditions
    BigCol[:n2, :] = Xi[0, :, :] 
    # Forcing terms
    BigCol[n2:, :] = terms.reshape((n2 * (N - 1), s))
    # Perform the multiplication and reshaping of the Xi array
    Xi = numpy.einsum('sij,js->is', self.BigMat, BigCol).reshape((N, n2, s))
    return Xi

第二个版本与第一个版本完全相同。可悲的是，第二个版本的运行时间几乎是前一个版本的 5 倍。我的猜测是Big Matrix的组装成本太高了......

有任何想法吗？谢谢！

标签： pythonoptimization