在没有循环开销的情况下在Python中实现相互关联的产品和总和

Question

我试图在python逻辑嵌套的sum和部分产品中实现以构建一个函数。这个想法是在没有显式循环的情况下这样做。期望的输出是一个列表（对于给定的矩阵J，用t索引）

公式为：

公式的简短口头描述：对于每个t，在0-（N-1）范围内有3个指数i，j，k。索引i，j构建矩阵（或2d数组），其中每个元素是索引k上的某些（J，t）依赖函数（并不重要）的乘积，不包括特定值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])

但是＆＃34; corr_time＆＃34;中的嵌套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和/差阵列会使性能差异很大（<45}。

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