Python多NumPy数组参数通用函数的向量化实现与绘图需求

现有一个处理NumPy数组的Python函数refl_general，代码如下：

import numpy as np
import matplotlib.pyplot as plt
import scipy.constants

pi = np.pi
i = 0.+1.j
c = scipy.constants.c

lam = 0.860e-6
f_c = c/lam

def refl_general(array_r, array_L, f):
    array_phi = 2*pi*array_L*f/c
    phi_sum = np.sum(array_phi)
    exp_sum = np.exp(i*phi_sum)

    r_all = np.prod(array_r)
    r_in = array_r[0] 
    r_except_in = np.prod(array_r[1:]) 
    
    return (-r_in + r_except_in*exp_sum)/(1-r_all*exp_sum)

该函数接收反射系数数组array_r、长度数组array_L和频率f，可正常运行，例如调用方式：

refl_general(np.array([np.sqrt(0.5), np.sqrt(0.5), np.sqrt(0.5)]), np.array([1, 1, 1]), f_c)

现在需要针对L = np.linspace(lam, 3*lam, 1000)，批量计算refl_general(np.array([np.sqrt(0.5), np.sqrt(0.5), np.sqrt(0.5)]), np.array([L/3, L/3, L/3]), f_c)的结果，替代以下循环代码的功能，且不修改原函数：

array_L = np.linspace(lam, 3*lam, 1000)
array_signal = np.array([])
for L in array_L:
    signal = np.abs(refl_general(np.array([np.sqrt(0.5), np.sqrt(0.5), np.sqrt(0.5)]), np.array([L/3, L/3, L/3]), f_c))
    array_signal = np.append(array_signal, signal)
plt.plot(array_L, array_signal)

直接传入np.array([L/3, L/3, L/3])会生成二维数组，导致原函数计算错误；尝试np.frompyfunc也因二维参数问题未成功，寻求无循环的解决方案。

方法1：使用np.vectorize包装适配函数 #

np.vectorize可将适配单值输入的逻辑转为支持数组批量处理的函数，这里我们包装一个lambda函数，自动将L的每个元素转为原函数需要的[L/3, L/3, L/3]格式：

array_L = np.linspace(lam, 3*lam, 1000)
# 包装函数，将单个L值转为对应的array_L参数传入原函数
vec_refl = np.vectorize(lambda L_val: refl_general(
    np.array([np.sqrt(0.5), np.sqrt(0.5), np.sqrt(0.5)]),
    np.array([L_val/3, L_val/3, L_val/3]),
    f_c
))
# 批量计算并取绝对值
array_signal = np.abs(vec_refl(array_L))
# 绘图
plt.plot(array_L, array_signal)
plt.show()

方法2：用np.apply_along_axis逐行处理二维数组 #

先构造每行对应一个L值的二维array_L数组，再通过np.apply_along_axis逐行传入原函数计算：

array_L = np.linspace(lam, 3*lam, 1000)
# 构造二维数组：每行对应一个L值的[L/3, L/3, L/3]
array_L_2d = np.tile(array_L/3, (3, 1)).T
# 逐行应用原函数
array_signal = np.abs(np.apply_along_axis(
    lambda row: refl_general(
        np.array([np.sqrt(0.5), np.sqrt(0.5), np.sqrt(0.5)]),
        row,
        f_c
    ),
    axis=1,
    arr=array_L_2d
))
# 绘图
plt.plot(array_L, array_signal)
plt.show()

补充说明 #

• 方法1的np.vectorize底层仍是循环，但语法简洁，适合快速实现需求；
• 方法2的np.apply_along_axis同样是逐元素处理，但更贴合NumPy的向量化编程风格；
• 两种方案均无需修改原函数refl_general，完全适配其参数格式。

内容的提问来源于stack exchange，提问作者KYKY