带有 if 语句的 Python Numba jit 函数

Question

我有一个包含 3 个部分的分段函数，我正尝试使用 Numba @jit 指令在 Python 中编写这些函数。该函数是在数组上计算的。该函数定义为:

@njit(parallel=True)
def f(x_vec):
    N=len(x_vec)
    y_vec=np.zeros(N)
    for i in prange(N):
        x=x_vec[i]
        if x<=2000:
            y=64/x
        elif x>=4000:
            y=np.log(x)
        else:
            y=np.log(1.2*x)
        y_vec[i]=y
    return y_vec

我正在使用 Numba 使这段代码非常快，并在我的 CPU 的所有 8 个线程上运行它。

现在，我的问题是，如果我想将函数的每个部分分别定义为 f1、f2 和 f3，然后将if 语句中的那些(并且仍然受益于 Numba 速度)，我该怎么做？原因是子函数可以更复杂，我不想让我的代码难以阅读。我希望它和这个一样快(或稍慢但不是很多)。

为了测试功能，我们可以使用这个数组:

Np=10000000
x_vec=100*np.power(1e8/100,np.random.rand(Np))
%timeit f(x_vec)  #0.06sec on intel core i7 3610

对于完成主义，调用了以下库:

import numpy as np
from numba import njit, prange

所以在这种情况下，函数将是:

def f1(x):
    return 64/x
def f2(x):
    return np.log(x)
def f3(x):
    return np.log(1.2*x)

实际函数是这些，它们用于层流、过渡和湍流状态的光滑管道摩擦系数:

@njit
def f1(x):
    return 64/x

@njit
def f2(x):
    #x is the Reynolds number(Re), y is the Darcy friction(f)
    #for transition, we can assume Re=4000 (max possible friction)
    y=0.02
    y=(-2/np.log(10))*np.log(2.51/(4000*np.sqrt(y)))
    return 1/(y*y)

@njit
def f3(x): #colebrook-white approximation
    #x is the Reynolds number(Re), y is the Darcy friction(f)
    y=0.02
    y=(-2/np.log(10))*np.log(2.51/(x*np.sqrt(y)))
    return 1/(y*y)

感谢大家的贡献。这是 numpy 解决方案(由于某种原因，最后的树行很慢，但不需要预热):

y = np.empty_like(x_vec)

a1=np.where(x_vec<=2000,True,False)
a3=np.where(x_vec>=4000,True,False)
a2=~(a1 | a3)

y[a1] = f1(x_vec[a1])
y[a2] = f2(x_vec[a2])
y[a3] = f3(x_vec[a3])

最快的 Numba 解决方案，允许传递函数名称并利用 prange(但受到 jit 预热的阻碍)是这个，它可以和第一个解决方案一样快(问题的顶部):

@njit(parallel=True)
def f(x_vec,f1,f2,f3):
    N = len(x_vec)
    y_vec = np.zeros(N)
    for i in prange(N):
        x=x_vec[i]
        if x<=2000:
            y=f1(x)
        elif x>=4000:
            y=f3(x)
        else:
            y=f2(x)
        y_vec[i]=y
    return y_vec

Answer 1

你可以写f()来接受函数参数，例如:

@njit
def f(arr, f1, f2, f3):
    N = len(arr)
    y_vec = np.zeros(N)
    for i in range(N):
        x = x_vec[i]
        if x <= 2000:
            y = f1(x)
        elif x >= 4000:
            y = f2(x)
        else:
            y = f3(x)
        y_vec[i] = y
    return y_vec

确保您传递的函数与 Numba 兼容。