NumPy 通用函数 (ufunc)

通用函数（ufunc）是NumPy中对数组进行逐元素运算的高性能函数。它们提供了向量化操作，避免了Python循环，并且支持广播、类型转换以及多种高级方法。理解ufunc是掌握NumPy高效计算的关键。

1. 什么是通用函数？

ufunc是“universal function”的缩写，它是一种能够对数组的每个元素进行操作的函数。大多数ufunc都是用C语言实现的，因此执行速度非常快。常见的ufunc包括算术运算、比较运算、三角函数、指数对数等。

import numpy as np

# 加法 ufunc
arr = np.array([1, 2, 3, 4])
print("np.add(arr, 2):", np.add(arr, 2))   # 等价于 arr + 2

# 正弦 ufunc
print("np.sin(arr):", np.sin(arr))

输出：

np.add(arr, 2): [3 4 5 6]
np.sin(arr): [ 0.84147098  0.90929743  0.14112001 -0.7568025 ]

2. ufunc 的属性和方法

ufunc对象本身具有一些有用的属性和方法。例如，np.add 就是一个ufunc对象。

2.1 常用属性

• ufunc.nin：输入参数个数
• ufunc.nout：输出参数个数
• ufunc.nargs：参数总数（包括输出）
• ufunc.types：支持的数据类型列表
• ufunc.identity：该操作的恒等值（例如加法为0，乘法为1）

print("np.add.nin:", np.add.nin)       # 2
print("np.add.nout:", np.add.nout)     # 1
print("np.add.identity:", np.add.identity)   # 0
print("np.multiply.identity:", np.multiply.identity)  # 1

2.2 重要方法

ufunc提供了几个强大的方法，允许在不同维度上进行聚合操作。

reduce()

沿指定轴对数组进行归约操作，例如累加、累乘。

arr = np.array([1, 2, 3, 4])
print("np.add.reduce(arr):", np.add.reduce(arr))      # 总和 10
print("np.multiply.reduce(arr):", np.multiply.reduce(arr))  # 乘积 24

# 二维数组指定轴
arr2d = np.array([[1, 2, 3], [4, 5, 6]])
print("沿列求和 (axis=0):", np.add.reduce(arr2d, axis=0))  # [5 7 9]
print("沿行求和 (axis=1):", np.add.reduce(arr2d, axis=1))  # [ 6 15]

accumulate()

累积结果，返回与输入形状相同的数组，每个元素是到当前位置的归约结果。

arr = np.array([1, 2, 3, 4])
print("np.add.accumulate(arr):", np.add.accumulate(arr))   # [1 3 6 10]
print("np.multiply.accumulate(arr):", np.multiply.accumulate(arr))  # [1 2 6 24]

reduceat()

在指定位置进行归约，返回部分归约结果。

arr = np.array([0, 1, 2, 3, 4, 5, 6, 7])
indices = [0, 2, 5]
print("np.add.reduceat(arr, indices):", np.add.reduceat(arr, indices))
# 计算三个区间的和：0-1, 2-4, 5-7 -> [0+1, 2+3+4, 5+6+7] = [1, 9, 18]

outer()

计算两个数组的所有元素对之间的运算，类似于外积。

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
print("np.add.outer(a, b):\n", np.add.outer(a, b))
# 输出：
# [[5 6 7]
#  [6 7 8]
#  [7 8 9]]

at()

对指定索引位置的元素执行就地操作，常用于更新部分元素。

arr = np.array([0, 0, 0, 0, 0])
indices = [1, 3, 3]
np.add.at(arr, indices, 5)   # 在索引1和3处各加5（索引3加了两次）
print("np.add.at 后:", arr)   # [0 5 0 10 0]

3. 可用的 ufunc 分类

NumPy提供了大量ufunc，涵盖以下常见类别：

• 算术运算：add, subtract, multiply, divide, power, mod 等
• 三角函数：sin, cos, tan, arcsin, arccos, arctan 等
• 双曲函数：sinh, cosh, tanh 等
• 指数对数：exp, log, log10, log2 等
• 比较运算：greater, less, equal, maximum, minimum 等
• 位运算：bitwise_and, bitwise_or, invert 等
• 浮点函数：floor, ceil, isnan, isfinite 等

可以通过 np.source(np.add) 查看源代码（如果可用），或访问官方文档查看完整列表。

4. 自定义 ufunc

虽然内置ufunc已足够丰富，但有时仍需自定义。NumPy提供了两种方式：

4.1 frompyfunc()

从一个Python函数创建ufunc。返回的ufunc总是返回Python对象类型，性能不如内置ufunc。

def my_func(x, y):
    return x * 2 + y

my_ufunc = np.frompyfunc(my_func, 2, 1)  # 2个输入，1个输出
a = np.array([1, 2, 3])
b = np.array([10, 20, 30])
print("my_ufunc(a, b):", my_ufunc(a, b))  # [12 24 36]

4.2 vectorize()

np.vectorize 是一个装饰器或函数，用于将普通Python函数向量化。它实际上是一个便利的包装，内部使用广播，但性能与 frompyfunc 类似（仍为Python循环）。

@np.vectorize
def my_func2(x, y):
    return x if x > y else y   # 自定义最大值

a = np.array([1, 5, 3])
b = np.array([4, 2, 6])
print("my_func2(a, b):", my_func2(a, b))  # [4 5 6]

注意：vectorize 主要用于方便，并非为了提高性能。对于高性能需求，应考虑使用C扩展或Numba。

5. ufunc 与广播

所有ufunc都支持NumPy的广播机制，允许对不同形状的数组进行逐元素运算。广播的规则在运算时自动应用。

a = np.array([[1, 2, 3], [4, 5, 6]])   # shape (2,3)
b = np.array([10, 20, 30])               # shape (3,)
print("a + b (广播):\n", np.add(a, b))   # b 广播到每一行

输出：

a + b (广播):
 [[11 22 33]
 [14 25 36]]

6. 性能对比：ufunc vs Python循环

ufunc的速度优势非常明显，尤其对于大型数组：

import time

arr = np.random.rand(1000000)

# Python 循环
start = time.time()
result = [np.sin(x) for x in arr]
print("Python循环耗时:", time.time() - start)

# ufunc
start = time.time()
result_ufunc = np.sin(arr)
print("ufunc耗时:", time.time() - start)

通常ufunc速度会快几十到上百倍。