6.3 ufunc的out参数与原地运算

6.3 ufunc的out参数与原地运算

本节摘要：ufunc（通用函数）是 NumPy 元素级运算的统一底座，加减乘除幂开方都是它的实例。out 参数把"新分配写"改成"指定缓冲区写"，原地运算符（+= 等）是其语法糖。合理使用能把内存峰值砍半以上；滥用（同一数组同时当输入与输出、整型溢出）会引入数值错误。

ufunc 是什么

先看清"运算符只是门面"这个事实：

import numpy as np

a = np.array([1.0, 4.0, 9.0])
print(type(np.sqrt))    # <class 'numpy.ufunc'>
print(np.add is np.add) # ufunc 是单例对象

# 运算符映射到 ufunc
print(a + 1)            # 等价 np.add(a, 1)
print(a ** 0.5)         # 等价 np.power(a, 0.5)

# 每个 ufunc 有一堆方法
print(np.add.reduce([1, 2, 3, 4]))     # 10，reduce=聚合（第7章的主角）
print(np.add.outer([1, 2], [10, 20]))  # 外积成表
# [[11 21]
#  [12 22]]
​

ufunc 的意义：所有元素级运算共享同一套调用协议，因此 out、where、dtype 这些参数对所有运算通用，学一次到处用。

out 参数：指定"写到哪"

默认每次运算新分配输出。out 把这笔分配取消，写进你给的缓冲区：

import numpy as np

a = np.random.rand(5)
buf = np.empty(5)

np.sqrt(a, out=buf)          # 结果写进 buf，不再分配新数组
print(buf.round(3))

# 读入相同数据连做两步，只重复用一个缓冲
np.multiply(a, 2.0, out=buf)
np.add(buf, 1.0, out=buf)    # buf 既当输入又当输出，逐元素对应安全
print(buf.round(3))
​

buf 同时当输入与输出为什么安全？ufunc 是逐元素执行的，读到第 k 个元素立刻写回第 k 个位置，没有前后依赖。但对聚合类方法（reduce），输入输出同体就危险了，等下用例子演示。

原地运算符是 out 的语法糖，且更严格——结果必须写回左侧 dtype：

import numpy as np

x = np.array([1, 2, 3], dtype=np.int32)

x += 10           # np.add(x, 10, out=x) 的糖
print(x)          # [11 12 13]

y = np.array([1, 2], dtype=np.int32)
# y += 1.5        # UFuncTypeError！
# 原地不允许 int 数组装 float 结果，安全机制拦住了静默截断
y = y + 1.5       # 普通加法会新建 float64 数组，合法
print(y, y.dtype) # [2.5 3.5] float64
​

这个对比值得记住：+= 是"必须装得下才允许"，+ 是"装不下就换更大的房子"。

常规运算与原地运算的账本对比

案例：滚动统计的缓冲区复用

完整过程。背景：实时数据流每秒到一个 100 万元素数组，要持续更新"归一化值"，每小时 3600 次，普通写法每小时新分配 3600 块 8MB。

操作：

import numpy as np
import time

rng = np.random.default_rng(2)

# 普通写法：每次表达式都新分配
def normalize_new(frame):
    return (frame - frame.min()) / (frame.ptp() + 1e-9)

# out 写法：预分配，重复复用
lo = np.float64(0.0)
work = None
def normalize_reuse(frame, work):
    if work is None or work.shape != frame.shape:
        work = np.empty_like(frame)
    d = frame.max() - frame.min() + 1e-9
    np.subtract(frame, frame.min(), out=work)
    work /= d
    return work

frame = rng.rand(1_000_000)

t0 = time.perf_counter()
for _ in range(50):
    _ = normalize_new(frame)
print("每次新建:", round((time.perf_counter() - t0) * 20, 1), "ms/次")
# 输出示例：每次新建: 5.8 ms/次

work = np.empty_like(frame)
t0 = time.perf_counter()
for _ in range(50):
    work = normalize_reuse(frame, work)
print("缓冲复用:", round((time.perf_counter() - t0) * 20, 1), "ms/次")
# 输出示例：缓冲复用: 4.3 ms/次
​

结果与解读：耗时差距不算大（省的主要是分配器开销与页错误，约四分之一），真正收益在内存——复用版峰值恒定两块，不再制造每秒一个的新块，垃圾回收压力与内存碎片同步消失。变式：多阶段管线（减均值、除方差、截断、量化）给每阶段一个复用缓冲区，或者干脆双缓冲交替，这是嵌入式与高频场景的标准做法。

⚠️ 常见坑：int8 累加别用 +=。np.array([100], dtype=np.int8) += 100 结果是 -56 的环绕。聚合同样危险：buf 做累加器时先用 astype(np.int64) 升宽，再 np.add.reduce(x, out=...)。

本节要点回顾

• 统一底座：所有元素级运算符都是 ufunc 实例，out/where/dtype 参数全线通用
• out 三大收益：免分配、降峰值、减少回收压力；逐元素同体读写安全，聚合同体危险
• 原地即语法糖：+= 比 + 多一条铁律——结果必须装回原 dtype，装不下直接报错而非截断
• 整型溢出：原地与窄整型是环绕高发组合，累加器先升 int64
• 选择标准：原值要不要保留决定用哪种写法，内存优化不等于无脑原地

第 7 章离开元素级视角，进入降维（聚合）与组合（矩阵乘）两类高阶运算。