python - 在没有循环开销的情况下在 Python 中实现交错乘积和求和

Question

我试图在 python 中实现逻辑嵌套的总和和部分积，以构建一个函数。这个想法是在没有显式循环的情况下这样做。所需的输出是一个列表(对于给定矩阵 J，由 t 索引)

公式为:

公式的简短口头描述:对于每个 t，在 0-(N-1) 范围内有 3 个索引 i、j、k。索引 i,j 构建一个矩阵(或二维数组)，其中每个元素都是一些 (J,t) 相关函数(无关紧要)在索引 k 上的乘积，不包括 i 的特定值， j.该函数只是扁平化矩阵/数组上的 numpy.sum。

现在，下面的代码按预期工作:

import numpy as np 
t_output = np.arange(0,10,100)
jmat = np.random.random(N**2).reshape(N,N)

def corr_mat(i,j,t,params):
  return np.prod(np.cos(\
    2.0 * t * np.delete(jmat[:,i] + jmat[:,j],(i,j)))) + \
      np.prod(np.cos(\
        2.0 * t * np.delete(jmat[:,i] - jmat[:,j],(i,j))))

def corr_time(t, jmat):
   return np.array([corr_mat(i,j,t,jmat) for i in xrange(N)\
     for j in xrange(N)]).reshape(N,N)

result = np.array([np.sum(corr_time(t,jmat)) for t in t_output])

但是“corr_time”函数中的嵌套 for 循环极大地减慢了执行速度。这个我试过了

import numpy as np 
t_output = np.arange(0,10,100)
jmat = np.random.random(N**2).reshape(N,N)

def corr_mat(i,j,t,params):
  return np.prod(np.cos(\
    2.0 * t * np.delete(jmat[:,i] + jmat[:,j],(i,j)))) + \
      np.prod(np.cos(\
        2.0 * t * np.delete(jmat[:,i] - jmat[:,j],(i,j))))

i,j = np.meshgrid(range(0,N), range(0,N))

result = np.array([np.sum(corr_mat(i,j,t,params)) for t in t_output])

但是函数没有正确理解 meshgrid。谁能建议我缺少什么？提前致谢。

Answer 1

预先计算 jmat 和/差数组会在性能上产生很大 (45x) 的差异。

def precalc(jmat):
  JM1 = np.zeros((N,N,N))
  JM2 = np.zeros((N,N,N))
  for i in range(N):
    for j in range(N):
      for k in range(N):
        #JJ[i,j,k]=jmat[k,i]+jmat[k,j]
        if k!=i and k!=j:
          JM1[i,j,k]=jmat[k,i]+jmat[k,j]
          JM2[i,j,k]=jmat[k,i]-jmat[k,j]
  return JM1, JM2

def corr_time1(t, JM1, JM2):
    return np.prod(np.cos(2*JM1*t),axis=-1)+np.prod(np.cos(2*JM2*t),axis=-1)

JM1, JM2 = precalc(jmat)
result = np.array([np.sum(corr_time1(t,JM1,JM2)) for t in t_output])

还有进一步改进的空间。我对 precalc 采用了蛮力方法，因为我没有找到更矢量化的方法。可能仍然有这样的解决方案。仍然只进行一次 j,i,k 迭代有很大帮助。

我们可以通过在更大的维度数组上执行 np.prod 来向量化最后一步，该数组使用了整个 t_output:

def corr_time2(t, JM1, JM2):
    return np.prod(np.cos(2*JM1[None,...]*t[:,None,None,None]),axis=-1) +\
       np.prod(np.cos(2*JM2[None,...]*t[:,None,None,None]),axis=-1)
result = np.sum(corr_time2(t_output, JM1, JM2),axis=(1,2))

在此测试用例中，节省的时间并不多，仅为 20%。我认为这是因为 t_output 只有 10 个元素。 np.arange(0,100,10)。在最后一个版本中，precalc 是最大的时间消耗者。

---

快速 precalc，加速 28 倍

def precalc1(jmat):
  # calc all the 'outer' sums/diffs, and zero the k=i,j terms
  ii = np.arange(jmat.shape[0])
  JM1 = jmat[:,:,None] + jmat[:,None,:]
  JM2 = jmat[:,:,None] - jmat[:,None,:]
  JM1[ii,ii,:] = 0
  JM2[ii,ii,:] = 0
  JM1[ii,:,ii] = 0
  JM2[ii,:,ii] = 0
  JM1 = JM1.transpose([1,2,0])
  JM2 = JM2.transpose([1,2,0])
  return JM1, JM2