3

我有一个分段函数，它包含 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
4

2 回答 2

3

这太慢了吗？这可以在纯 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 操作，这仍然会很快。

于 2020-10-28T16:14:40.047 回答
3

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

于 2020-10-28T16:51:41.703 回答