Python 带有 if 语句的 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

这太慢了吗？这可以在纯 numpy 中完成，方法是避免循环并使用掩码进行索引：

def f(x):
    y = np.empty_like(x)
    
    mask = x <= 2000
    y[mask] = 64 / x[mask]
    
    mask = (x > 2000) & (x < 4000)
    y[mask] = np.log(1.2 * x[mask])
    
    mask = x >= 4000
    y[mask] = np.log(x[mask])

    return y

您也可以运行 “else” 情况，首先将没有任何掩码的中间部分应用到整个数组，它可能有点慢：

def f_else(x):
    y = np.log(1.2 * x)
    
    mask = x <= 2000
    y[mask] = 64 / x[mask]
    
    mask = x >= 4000
    y[mask] = np.log(x[mask])

    return y

有

Np=10000000
x_vec=100*np.power(1e8/100,np.random.rand(Np))

我得到（配备 i7-8850H 和 6 + 6VT 内核的笔记本电脑）

f1: 1 loop, best of 5: 294 ms per loop
f_else: 1 loop, best of 5: 400 ms per loop

如果您打算使用的子功能主要是numpy-operations，这仍然会很快。

Answer 2

可以写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 兼容。

Python 带有 if 语句的 Numba jit 函数

Python Numba jit function with if statement

python

jit

numpy

numba

boolean-indexing