如何高效计算 NumPy 矩阵乘法的实部（无需虚部）-Python教程-PHP中文网

如何高效计算 NumPy 矩阵乘法的实部（无需虚部）

当对两个复数数组执行 np.matmul(a, b).real 时，直接丢弃虚部会造成约 50% 的冗余浮点运算；本文介绍基于内存视图重构与 Numba 并行优化的高效方案，可将实部计算速度提升 67%–81%，并适用于 einsum 等复杂张量收缩场景。

当对两个复数数组执行 `np.matmul(a, b).real` 时，直接丢弃虚部会造成约 50% 的冗余浮点运算；本文介绍基于内存视图重构与 numba 并行优化的高效方案，可将实部计算速度提升 67%–81%，并适用于 `einsum` 等复杂张量收缩场景。

在科学计算中，常需对复数矩阵执行乘法后仅保留实部（例如量子力学模拟、信号处理或协方差估计）。标准写法 np.matmul(a, b).real 虽简洁，但底层仍完整计算复数乘积（每个元素含 4 次浮点运算：re×re − im×im 和 re×im + im×re），而我们仅需前者——这意味着近半数算力被浪费。

核心优化思路是绕过复数运算，转为纯实数 BLAS 计算：

将复数矩阵 A ∈ ℂ^(m×k) 通过 .view(np.float64) 重解释为 ℝ^(m×2k)，其内存布局自动交错实部与虚部（如 [a₁ᵣ, a₁ᵢ, a₂ᵣ, a₂ᵢ, ...]）；
同时将 B ∈ ℂ^(k×n) 变换为 ℝ^(2k×n)，每行 i 对应 B[i].real 和 -B[i].imag，从而使 A.view(...) @ B2 直接产出 Re(A @ B)。

纯 NumPy 实现如下（适用于中等规模）：

A2 = A.view(np.float64)  # 内存零拷贝重解释
B2 = np.empty((B.shape[0] * 2, B.shape[1]), dtype=np.float64)
B2[0::2] = B.real   # 偶数行：实部
B2[1::2] = -B.imag  # 奇数行：负虚部
C_real = A2 @ B2    # 纯实数矩阵乘，结果即为 Re(A @ B)

但该方法仍生成临时数组（B.real, -B.imag）且转置开销显著。更优解是使用 Numba JIT 编译+并行化，消除中间内存分配，并充分利用多核：

import numba as nb
import numpy as np

@nb.njit('(c16[:,::1], c16[:,::1])', parallel=True, cache=True)
def matmul_real_part(a, b):
    m, k = a.shape
    k2, n = b.shape
    assert k == k2
    # 预分配 B2：2k × n 实数矩阵
    b2 = np.empty((k * 2, n), dtype=np.float64)
    # 并行填充：每行 i → [b[i].real, -b[i].imag]
    for i in nb.prange(k):
        for j in range(n):
            b2[i*2, j] = b.real[i, j]
            b2[i*2+1, j] = -b.imag[i, j]
    # 视图转换 + 实数乘法（调用高度优化的 BLAS）
    return a.view(np.float64) @ b2

进一步加速技巧（生产环境推荐）：
✅ 预分配 + 复用内存：若矩阵尺寸固定（如循环中反复计算），提前分配 out 和 B2，避免重复内存申请；
✅ 页预填充（np.full(..., 0.0) 而非 np.zeros()）：防止运行时缺页中断，提升缓存局部性；
✅ 绑定到 einsum：对 np.einsum("ab,bc,cd,da->a", ...).real，可将各复数张量按需拆解为实/虚分量，构造等价实数 einsum 表达式（如 "ab,bc,cd,da->a" → "ar,br,cr,dr->a" + "ai,bi,ci,di->a" 等组合），再用相同策略优化。

⚠️ 注意事项：

此法要求输入为 C 连续（order='C'），否则 .view() 可能失败，建议调用前加 a = np.ascontiguousarray(a)；

Numba 版本需 ≥0.58，且启用 parallel=True 时确保 OMP_NUM_THREADS 设置合理；

对小矩阵（np.matmul(...).real 更简单可靠；

极致性能场景（如 HPC）可考虑定制 BLAS 内核，但需手写汇编级优化（如 OpenBLAS 风格分块+向量化），工程成本极高。