Python Numba jit函数与if语句

问题描述

我有一个由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,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

解决方法

这太慢了吗？可以通过避免循环并使用掩码进行索引来以纯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))

我得到了（带有6 + 6VT内核的i7-8850H笔记本电脑）

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

如果您想要的子功能主要是numpy操作，那么它仍然会很快。

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

boolean-indexing jit jit jit numba numpy python