Cython/Numba 编译函数可以改进 numpy.max(numpy.abs(a-b)) 吗？

Question

我正在优化代码的瓶颈部分——迭代函数 a' = f(a)，其中 a 和 a' 是 N 乘 1 个向量，直到 max(abs(a' - a))足够小了。

我在 f(a) 上放置了一个 Numba 包装器，并且在我能够生产的最优化的纯 NumPy 版本上获得了很好的加速（运行时间减少了大约 50%）。

我尝试编写一个 C 兼容版本的 numpy.max(numpy.abs(aprime - a))，但事实证明这样更慢！我实际上失去了从迭代的第一部分 Numba-fying 中获得的所有收益！

Numba 或 Cython 是否有可能改进 numpy.max(numpy.abs(aprime - a))？我在下面复制我的代码以供参考，其中 a 是 P0，a' 是 Pprime:

编辑：对我来说，"flatten()" "maxabs()" 的输入似乎很重要。当我这样做时，性能并不比 NumPy 差。然后，当我按照 JoshAdel 的建议在计时括号外执行函数 "dry run" 时，"maxabs" 的循环比 numpy.max(numpy.abs() 的循环稍好).

from numba import jit
import numpy as np

### Preliminaries, to make the working example fully functional

n = 1200
Gammer = np.exp(-np.random.rand(n,n))

alpher = np.ones((n,1))
xxer = 10000*np.random.rand(n,1)

chii = 6.5
varkappa = 6.5
phi3 = 1.5
A = .5
sig = .2 

mmer = np.dot(Gammer,xxer**phi3)


totalprod = A*alpher + (1-A)*mmer
Gammerchii = Gammer**chii
Gammerrats = Gammerchii[:,0].flatten()/Gammerchii[0,:].flatten()
Gammerrats[(Gammerchii[0,:].flatten() == 0) | (Gammerchii[:,0].flatten() == 0)] = 1.
P0 = (Gammerrats*(xxer[0]/totalprod[0])*(totalprod/xxer).flatten())**(1/(1+2*chii))
P0 *= n/np.sum(P0)
### End of preliminaries

### This is the function to produce a' = f(a)
@jit
def Piteration(P0, chii, sig, n, xxer, totalprod, Gammerrats, Gammerchii):
    Mac = np.zeros((n,))
    Pprime = np.zeros((n,))
    themacpow = 1-(1/chii)*(sig/(1-sig))
    specialchiipow = 1/(1+2*chii)
    Psum = 0.

    for i in range(n):
        for j in range(n):
            Mac[j] += ((P0[i]/P0[j])**chii)*Gammerchii[i,j]*totalprod[j]

    for i in range(n):
        Pprime[i] = (Gammerrats[i]*(xxer[0]/totalprod[0])*(totalprod[i]/xxer[i])*((Mac[i]/Mac[0])**themacpow))**specialchiipow
        Psum += Pprime[i]

    Psum = n/Psum

    for i in range(n):
        Pprime[i] *= Psum

    return Pprime

### This is the function to find max(abs(aprime - a))
@jit
def maxabs(vec1,vec2,n):
    themax = 0.
    curdiff = 0.
    for i in range(n):
        curdiff = vec1[i] - vec2[i]
        if curdiff < 0:
            curdiff *= -1
        if curdiff > themax:
            themax = curdiff
    return themax

### This is the main loop
diff = 1000.
while diff > 1e-2:
    Pprime = Piteration(P0.flatten(),  chii,  sig,  n,  xxer.flatten(), totalprod.flatten(), Gammerrats.flatten(),  Gammerchii)

    diff = maxabs(P0.flatten(),Pprime.flatten(),n)
    P0 = 1.*Pprime

Answer 1

当我为形状数组 (1200,) 计算你的 maxabs 函数与 np.max(np.abs(vec1 - vec2)) 的时间时，使用 numba 0.32.0 的 numba 版本快了约 2.6 倍。

当您为代码计时时，请确保在为它计时之前运行您的函数一次，这样您就不会包括 jit 代码所需的时间，您只需支付第一次费用。一般来说，多次使用 timeit 和运行ning 可以解决这个问题。我不确定你是如何计时的，因为我发现使用 maxabs 与 numpy 调用几乎没有区别，大部分运行时间似乎都在调用 Piteration .

Cython/Numba 编译函数可以改进 numpy.max(numpy.abs(a-b)) 吗？

Can a Cython/Numba compiled function improve on numpy.max(numpy.abs(a-b))?

python

arrays

numpy

cython

numba