如何使用numpy对2D参数网格上的卡方计算进行向量化实现

我有3个长度为N的一维数组：x、y、y_error，采用A*sin(w*x+phi)+c作为拟合模型。我已经定义了卡方（chi-squared）函数，并使用scipy.optimize.minimize完成了模型对数据的拟合。

现在需要实现以下需求：遍历A、w、phi、c四个参数中任意两个组成的离散二维网格（例如A取值为[0,1,2,3]、w取值为[0.2,0.3,0.4]的网格），计算网格上每一组参数对对应的卡方值，未被纳入网格的另外两个参数使用固定值，最终绘制二维参数空间上的卡方曲面图。

用for循环实现该需求效率极低，我希望利用numpy数组的向量化特性实现高效计算，但尝试实现时遇到数组维度不匹配的广播错误。我的测试代码如下：

import pandas as pd
import numpy as np
import scipy.optimize as opt
import scipy.stats as stats

#TEST DATA
x=np.linspace(0,2*np.pi,20)
y=np.sin(x)+np.random.normal(loc=0,scale=0.01,size=len(x))
y_error=np.array([0.01]*len(x))

#MODEL
def sinusoidal_model(t,params):
    A,w,phi,c=params
    return(A*np.sin(w*t+phi)+c)

#CHI-SQUARED
def chi_squared(model_params, model, x, y, y_error):
    return np.sum(((y - model(x, model_params))/y_error)**2)

#INITIAL MODEL PARAMETERS
A=1
w=1
phi=0
c=0
initial_params=np.array((A,w,phi,c))

#FITTING
model=sinusoidal_model
deg_freedom = x.size - initial_params.size 
fit = opt.minimize(chi_squared, initial_params, args=(model, x, y, y_error))
fit_params= fit.x
chisq_min = chi_squared(fit_params, model, x, y, y_error)

#GENERATE COLOUR AND CONTOUR PLOTS
i_of_param_i=0
j_of_param_j=1
param_i=fit_params[i_of_param_i]
param_j=fit_params[j_of_param_j]
param_i_half_range=1
param_j_half_range=1
n_points_x=5
n_points_y=4
param_i_axis=np.linspace(param_i-param_i_half_range,param_i+param_i_half_range,n_points_x)
param_j_axis=np.linspace(param_j-param_j_half_range,param_j+param_j_half_range,n_points_y)
X,Y=np.meshgrid(param_i_axis,param_j_axis,indexing='xy')

核心利用numpy的广播机制调整数组维度，完全消除Python层循环，实现最高效的向量化计算。

具体修改 #

1. 修改卡方函数，指定求和轴 #

原卡方函数默认对所有维度求和，会把参数网格的维度也纳入计算，需要指定仅对数据点对应的维度求和：

def chi_squared(model_params, model, x, y, y_error):
    # axis=-1 指定对最后一维（长度为N的数据点维度）求和
    return np.sum(((y - model(x, model_params))/y_error)**2, axis=-1)

2. 调整数组维度适配广播 #

给一维的x、y、y_error增加两个前导维度，匹配参数网格的形状，即可完成广播运算：

# 初始化参数数组，所有参数先用拟合最优值填充
parameters = [np.full(X.shape, fit_params[i]) for i in range(len(fit_params))]
# 替换两个可变参数为网格值
parameters[i_of_param_i] = X
parameters[j_of_param_j] = Y
# 堆叠参数数组，形状变为 (n_points_y, n_points_x, 4)，最后一维为四个模型参数
parameters = np.stack(parameters, axis=-1)
# 调整数据数组形状，新增两个前导维度适配参数网格
x_ = x.reshape(1, 1, -1)
y_ = y.reshape(1, 1, -1)
y_error_ = y_error.reshape(1, 1, -1)
# 计算卡方，输出Z的形状和X、Y完全一致，可直接用于绘图
Z = chi_squared(parameters, model, x_, y_, y_error_)

方案优势 #

• 完全向量化实现，无Python层循环，计算效率远高于for循环，网格越大性能优势越明显
• 无需修改其他逻辑，仅调整i_of_param_i和j_of_param_j的取值即可切换任意两个参数作为网格变量
• 输出的Z数组直接适配matplotlib等高线、曲面绘图接口，无需额外格式转换

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