如何使用 numba.vectorize 函数，参数为浮点数序列

Question

我正在尝试优化一些代码，并寻找使用 numba 制作 numpy 通用函数的方法。

原始函数具有以下签名

import numpy as np
from numba import njit, vectorize
from typing import Sequence

@njit
def f(x: float, y: Sequence[float], z:float = 1e-14) -> float:
    """Computations are only exemplarily"""

    a = np.sum(y)
    if np.abs(x - a) < z:
        return 0.
    else:
        return np.abs(x - a)

按预期工作。

我想使用

numba.vectorize

创建一个通用函数 s.t.

f

可以通过

调用

x: np.ndarray | float, y: Sequence[np.ndarray | float], z: float

此外，

f

是在代码的其他部分创建的，即

y

的长度有所不同，但在定义

f

时已知。

一个明显的解决方案是创建另一个接受（可能）矢量化输入的函数 并使用

numba.prange

进行一些操作以填充结果向量。

但是我正在寻找使用

numba.vectorize

的更优雅的解决方案，并且我不想处理 广播，因为

x

可以是浮点数，或者

y

的元素也可以是浮点数。

致以诚挚的问候，

我试过了

 f_v = vectorize(
     [
         numba.float64(
             numba.float64,
             numba.float64[:],
             numba.float64,
         )
     ],
     nopython=True,
)(f)

但是编译失败，没有找到匹配的签名。

编辑：

f

对标量执行简单的算术运算。我插入了一些虚拟计算，因为原始计算相当冗长。 矢量化版本应该按组件执行这些计算并返回一个数组。

Answer 1

您可以使用 guvectorize 而不是矢量化。

@numba.guvectorize("(float64, float64[:], float64, float64[:])", "(),(n),()->()", target="parallel")
def f_vectorized(x, y, z, out):
    out[0] = f(x, y, z)


result1 = f_vectorized(1.0, np.array([1.0, 2.0]), 1e-14)
result2 = f_vectorized(np.array([1.0]), np.array([[1.0, 2.0]]), 1e-14)
result3 = f_vectorized(np.array([1.0, 2.0, 3.0]), np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), 1e-14)
result4 = f_vectorized(np.array([1.0, 2.0, 3.0]), [np.array([1.0, 2.0]), np.array([3.0, 4.0]), np.array([5.0, 6.0])], 1e-14)

但是，有两个问题需要考虑。

第一个很明显：

z

不能有默认值。 要解决这个问题，您必须创建另一个函数。

def f_vectorized_with_default_value(x, y, z=1e-14):
    return f_vectorized(x, y, z)

第二个更关键：如果

y

是一个 numpy 数组的 python 列表，它就不能很好地工作。也就是这个案例：

result4 = f_vectorized(np.array([1.0, 2.0, 3.0]), [np.array([1.0, 2.0]), np.array([3.0, 4.0]), np.array([5.0, 6.0])], 1e-14)

这是基准代码。

import time
import timeit
import warnings

import numpy as np
from typing import Sequence
import numba


@numba.njit(cache=True)
def f(x: float, y: Sequence[float], z: float = 1e-14) -> float:
    """Computations are only exemplarily"""

    a = np.sum(y)
    if np.abs(x - a) < z:
        return 0.0
    else:
        return np.abs(x - a)


@numba.njit(parallel=True, cache=True)
def f_prange(x, y, z):
    """Implementation with prange as a baseline."""
    n = len(x)
    out = np.empty(n, dtype=x.dtype)
    for i in numba.prange(n):
        out[i] = f(x[i], y[i], z)
    return out


def f_loop(x, y, z):
    """Another baseline."""
    n = len(x)
    out = np.empty(n, dtype=x.dtype)
    for i in range(n):
        out[i] = f(x[i], y[i], z)
    return out


@numba.guvectorize("(float64, float64[:], float64, float64[:])", "(),(n),()->()", target="parallel", cache=True)
def f_vectorized(x, y, z, out):
    out[0] = f(x, y, z)


def f_vectorized_with_default_value(x, y, z=1e-14):
    return f_vectorized(x, y, z)


def test(title, func, x, y, z, expected, number=10):
    assert np.array_equal(func(x, y, z), expected)
    elapsed = timeit.timeit(lambda: func(x, y, z), number=number) / number
    print(f"{title}: {elapsed}")


def main():
    rng = np.random.default_rng(0)
    n = 1000000
    m = 1000

    x = rng.random(size=(n,), dtype=np.float64)
    y_as_numpy = rng.random(size=(n, m), dtype=np.float64)
    y_as_list = [y_as_numpy[i] for i in range(n)]
    z = 1e-14

    started = time.perf_counter()
    _ = np.array(y_as_list)
    print("merge into a np array:", time.perf_counter() - started)

    started = time.perf_counter()
    y_as_typed_list = numba.typed.List(y_as_list)
    print("convert to typed list:", time.perf_counter() - started)

    expected = f(x[0], y_as_numpy[0], z)
    test("f(single)", f, x[0], y_as_numpy[0], z, expected)
    test("f_vectorized(single)", f_vectorized, x[0], y_as_numpy[0], z, expected)

    expected = f_prange(x, y_as_numpy, z)
    test("f_prange(numpy)", f_prange, x, y_as_numpy, z, expected)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", numba.NumbaPendingDeprecationWarning)
        test("f_prange(python_list)", f_prange, x, y_as_list, z, expected)
    test("f_prange(typed_list)", f_prange, x, y_as_typed_list, z, expected)
    test("f_loop(python_list)", f_loop, x, y_as_list, z, expected)
    test("f_vectorized(numpy)", f_vectorized, x, y_as_numpy, z, expected)
    test("f_vectorized(python_list)", f_vectorized, x, y_as_list, z, expected)
    test("f_vectorized(typed_list)", f_vectorized, x, y_as_typed_list, z, expected)


if __name__ == "__main__":
    main()

结果：

merge into a np array: 1.3248698
convert to typed list: 1.6472616999999996
f(single): 1.430000000013365e-06
f_vectorized(single): 1.6320000000025204e-05
f_prange(numpy): 0.18999052999999985
f_prange(python_list): 7.21668874
f_prange(typed_list): 0.20048094999999932 + (convert to typed list)
f_loop(python_list): 1.3346305299999996
f_vectorized(numpy): 0.18892754000000025
f_vectorized(python_list): 1.8655723899999999
f_vectorized(typed_list): 3.323696179999999 + (convert to typed list)

从最后两行结果可以看出，速度非常慢。它有效，但速度很慢。 事实上，在 python 循环中顺序调用

f

函数会更快（不是

numba.prange

，纯 python 循环）。

如果

y

首先是一个 2D numpy 数组，这不会打扰你。

f_vectorized

应该可以正常工作。 但如果它是一个Python列表，你可能需要找到另一个技巧。