python - 如何做总和的平方和的总和？

Question

我有一个总和，我想加速。在一种情况下是:

S_{x,y,k,l} Fu_{ku} Fv_{lv} Fx_{kx} Fy_{ly}

在另一种情况下是:

S_{x,y} ( S_{k,l} Fu_{ku} Fv_{lv} Fx_{kx} Fy_{ly} )^2

注意:S_{indices}:是这些指数的总和

第一个案例我已经弄清楚如何使用 numpy 的 einsum 并且它导致惊人的加速 ~ x160。

此外，我曾考虑过尝试扩大正方形，但这不会是一个 killer ，因为我需要对 x,y,k,l,k,l 而不是 x,y,k,l 求和吗？

这是一个演示差异的实现以及我使用 einsum 的解决方案。

Nx = 3
Ny = 4
Nk = 5
Nl = 6
Nu = 7
Nv = 8
Fx = np.random.rand(Nx, Nk)
Fy = np.random.rand(Ny, Nl)
Fu = np.random.rand(Nu, Nk)
Fv = np.random.rand(Nv, Nl)
P = np.random.rand(Nx, Ny)
B = np.random.rand(Nk, Nl)
I1 = np.zeros([Nu, Nv])
I2 = np.zeros([Nu, Nv])
t = time.time()
for iu in range(Nu):
    for iv in range(Nv):
        for ix in range(Nx):
            for iy in range(Ny):
                S = 0.
                for ik in range(Nk):
                    for il in range(Nl):
                        S += Fu[iu,ik]*Fv[iv,il]*Fx[ix,ik]*Fy[iy,il]*P[ix,iy]*B[ik,il]
                I1[iu, iv] += S
                I2[iu, iv] += S**2.
print time.time() - t; t = time.time()
# 0.0787379741669
I1_ = np.einsum('uk, vl, xk, yl, xy, kl->uv', Fu, Fv, Fx, Fy, P, B)
print time.time() - t
# 0.00049090385437
print np.allclose(I1_, I1)
# True
# Solution by expanding the square (not ideal)
t = time.time()
I2_ = np.einsum('uk,vl,xk,yl,um,vn,xm,yn,kl,mn,xy->uv', Fu,Fv,Fx,Fy,Fu,Fv,Fx,Fy,B,B,P**2)
print time.time() - t
# 0.0226809978485 <- faster than for loop but still much slower than I1_ einsum
print np.allclose(I2_, I2)
# True

如图所示，我已经设法完成了 I1_，我已经想出了如何使用 einsum 为 I1 执行上述操作。

编辑:

我添加了如何通过扩展正方形来执行 I2_，但速度有点令人失望并且在意料之中......与 ~x160 相比，~x3.47 加速

编辑 2:

加速似乎并不一致，我之前得到了 x40 和 x1.2，但现在得到了不同的数字。无论哪种方式，差异和问题仍然存在。

编辑 3:我试图简化我实际追求的总和，但搞砸了，上面的总和提供了 @user5402 提供的出色答案。

我已经编辑了上面的代码来演示下面的总和:

I1 = S_{x,y,k,l} Fu_{ku} Fv_{lv} Fx_{kx} Fy_{ly} P_{xy} B_{kl}

I2 = S_{x,y} ( S_{k,l} Fu_{ku} Fv_{lv} Fx_{kx} Fy_{ly} P_{xy} B_{kl} )^2

Answer 1

(更新:跳到最后查看表示为一对矩阵乘法的结果。)

我认为你可以通过使用身份大大简化计算:

例如，

S_{k,l} Fu_{ku} Fv_{lv} Fx_{kx} Fy_{ly}
  = S_{k,l} Fu_{ku} Fx_{kx} Fv_{lv} Fy_{ly}            -- rearrange the factors
            \___ A ____/    \___ B ____/
  = ( S_k Fu_{ku} Fx_{kx} ) * ( S_l Fv_{lv} Fy_{ly} )  -- from the identity
  =   A_{ux}                * B_{vy}

其中 A_{ux} 只依赖于 u 和 x 而 B_{vy} 只依赖于v 和 y。

对于平方和，我们有:

S_k [ S_l Fu_{ku} Fv_{lv} Fx_{kx} Fy_{ly} ]^2
  = S_k Fu_{ku} Fx_{kx} * [ S_l Fv_{lv} Fy_{ly} ]^2
  = S_k Fu_{ku} Fx_{kx} * B_{vy}^2                 -- B is from the above calc.
  = B_{vy}^2 * S_k Fu_{ku} Fx_{kx}                 -- B_vy is free of k
  = B_{vy}^2 * A_{ux}                              -- A is from the above calc.

在 x 和 y 上继续求和时会发生类似的减少:

S_{xy} A_{ux} * B_{vy}
  = S_x A_{ux} * S_y B_{vy}                        -- from the identity
  =  C_u       *    D_v

最后对 u 和 v 求和:

S_{uv} C_u D_v = (S_u C_u) * (S_v D_v)             -- from the identity

希望这对您有所帮助。

更新:我刚刚意识到，也许对于你想要计算的平方和[ S_k S_l ... ]^2 在这种情况下，您可以这样进行:

[ S_k  S_l Fu_{ku} Fv_{lv} Fx_{kx} Fy_{ly} ]^2
  =  [ A_{ux}                * B_{vy} ]^2
  =  A_{ux}^2 * B_{vy}^2

所以当我们对以上变量求和时，我们得到:

S_{uvxy} A_{ux}^2 B_{vy}^2
  = S_{uv} ( S_{xy}  A_{ux}^2 B_{vy}^2 )
  = S_{uv} ( S_x A_{ux}^2 ) * ( S_y B_{vy}^2 )     -- from the identity
  = S_{uv}     C_u          *      D_v
  = (S_u C_u) * (S_v D_v)                          -- from the identity

更新 2:这确实归结为几个矩阵乘法。

A和B的定义:

A_{uv} = S_k Fu_{ku} Fx_{kx}
B_{vy} = S_l Fv_{lv} Fy_{ly}

也可以写成矩阵形式为:

A = (transpose Fu) . Fx             -- . = matrix multiplication
B = (transpose Fv) . Fy

以及C和D的定义:

C_u = S_x A_{ux}
D_v = S_y B_{vy}

我们看到向量 C 只是 A 的行总和，向量 D 只是 B 的行总和。因为整个求和(不是平方)的答案是:

total = (S_u C_u) * (S_v D_v)

我们看到总和就是 A 的所有矩阵元素的总和乘以 B 的所有矩阵元素的总和。

这是 numpy 代码:

from numpy import *
# ... set up Fx, Fv, Fu, Fy as above...

A = Fx.dot(Fu.transpose())
B = Fv.dot(Fy.transpose())
sum1 = sum(A) * sum(B)

A2 = square(A)
B2 = square(B)
sum2 = sum(A2) * sum(B2)

print "sum of terms:", sum1
print "sum of squares of terms:", sum2