NumPy库中的ufunc函数——自定义ufunc函数

NumPy提供的标准的ufunc函数可以组合出复合的表达式，但是在有些情况下，自己编写的会更加方便。我们可以把自己编写的函数用frompyfunc（）转化成ufunc函数。

我们先举一个例子来展示一下怎么进行自定义

一、自定义ufunc函数实例

调用格式：

frompyfunc（func， nin， nout）

func：计算函数

nin：func（）输入参数的个数

nout；func（）输出参数的个数

import numpy as np
def num_judge(x, a):
    if x % 3 == 0:
        r = 0
    elif x % 5 == 0:
        r = 0
    else:
        r = a
    return r
x = np.linspace(0, 10, 11)
y = np.array([num_judge(t, 2) for t in x])
numb_judge = np.frompyfunc(num_judge, 2 , 1)
y = numb_judge(x, 2)
print(y)

运行结果：

[0 2 2 0 2 0 0 2 2 0 0]

这样，我们就完成了一个自定义的ufunc函数。

二、运行速度优化

那么是否自定义定义了一个ufunc函数，就可以具有ufunc函数的计算速度快的特点呢？

我们接下来对函数运行速度进行测试，还是用上边的那段代码，我们换成大一点的参数。

import numpy as np
import time

start_time = time.time()
def num_judge(x, a):
    if x % 3 == 0:
        r = 0
    elif x % 5 == 0:
        r = 0
    else:
        r = a
    return r
x = np.linspace(0, 100, 10000000)
y = np.array([num_judge(t, 2) for t in x])
numb_judge = np.frompyfunc(num_judge, 2 , 1)
y = numb_judge(x, 2)
print(y)
end_time = time.time()
print(f"代码执行时间: {end_time - start_time} 秒")

运行结果：

[0 2 2 ... 2 2 0]
代码执行时间: 4.403841018676758 秒

如此之慢的运行速度，似乎并不是我们想要让ufunc函数起到的效果。那么怎么样才可以让自定义的ufunc函数具有超快的运行速度呢？

答案是肯定的！但是我们需要调用到numba库，它允许我们使用类似于Python的语法编写高效的数组操作代码。

在这里我们先简单展示一下优化后的代码及其运行速度，关于numba库的使用方法，将在后续推文中给大家展示~

import numpy as np
from numba import njit
import time


@njit
def num_judge_fast(x, a):
    result = np.empty_like(x)  # 创建一个与x形状和类型相同的空数组来存储结果
    for i in range(x.size):  # 遍历数组的每个元素
        if x[i] % 3 == 0:
            result[i] = 0
        elif x[i] % 5 == 0:
            result[i] = 0
        else:
            result[i] = a
    return result


start_time = time.time()

# 创建 NumPy 数组
x = np.linspace(0, 100, 10000000)

# 使用优化后的函数处理整个数组
y = num_judge_fast(x, 2)

print(y)
end_time = time.time()
print(f"代码执行时间: {end_time - start_time} 秒")

运行结果：

[0. 2. 2. ... 2. 2. 0.]
代码执行时间: 0.4699270725250244 秒

运行速度提升了10倍！

---

如果您觉得有用，点赞关注不迷路哦~~~

个人主页正在陆续更新python相关知识。