如何用xarray并行化1D操作实现多变量岭回归？

问题背景 #

现有三个xarray DataArray（ds_pr_cru、ds_tas_cru、ds_nep），需在time维度上执行多变量线性岭回归：将前两个数组作为自变量（X1、X2），第三个作为因变量（Y）。已成功实现单自变量回归，但多变量版本运行报错。

错误原因分析 #

原代码的核心问题是Scikit-learn方法参数使用错误：

• ridge.fit(x1.reshape(-1,1), x2.reshape(-1,1), y.reshape(-1,1))：fit()的第二个参数应为因变量Y，第三个参数是样本权重，此处错误将X2作为Y、Y作为权重，触发"Sample weights must be 1D array or scalar"报错
• ridge.score(x1,x2,y)、ridge.predict(x1,x2)：score()和predict()的输入需是合并后的二维特征矩阵，而非分开的单个自变量
• 额外问题：stats.coef_pval并非Scikit-learn标准方法，岭回归也无原生p值计算逻辑，需调整实现逻辑

修正后的代码实现 #

1. 构造正确的多变量回归函数 #

将两个自变量合并为二维特征矩阵，再传入模型：

from sklearn import linear_model
import numpy as np
import xarray as xr

def multi_ridge_regress(x1, x2, y):
    # 合并两个自变量为(n_samples, 2)的特征矩阵
    X = np.column_stack([x1, x2])
    y_reshaped = y.reshape(-1, 1)
    
    ridge = linear_model.Ridge()
    model = ridge.fit(X, y_reshaped)
    
    # 提取结果：2个系数、R²值、n_time个预测值
    coefs = model.coef_.flatten()  # 转为1D数组：[coef_x1, coef_x2]
    r2 = model.score(X, y_reshaped)
    y_pred = model.predict(X).flatten()
    
    # 拼接结果为一维数组，适配xarray输出要求
    return np.concatenate([coefs, [r2], y_pred])

2. 调用xr.apply_ufunc执行并行计算 #

调整输出维度参数，匹配返回结果的长度（2系数 + 1R² + n_time个预测值）：

# 先获取时间维度长度，确定输出结果总长度
subset_pr = ds_pr_cru.sel(lon=slice(-60,-50), lat=slice(-20,-10))
n_time = subset_pr.time.size
output_total_size = 3 + n_time  # 2系数 + 1R² + n_time预测值

test = xr.apply_ufunc(
    multi_ridge_regress,
    ds_pr_cru.sel(lon=slice(-60,-50), lat=slice(-20,-10)),
    ds_tas_cru.sel(lon=slice(-60,-50), lat=slice(-20,-10)),
    ds_nep.sel(lon=slice(-60,-50), lat=slice(-20,-10)),
    input_core_dims=[['time'], ['time'], ['time']],
    output_core_dims=[["result"]],
    vectorize=True,
    dask="parallelized",
    output_dtypes=['float64'],
    output_sizes={"result": output_total_size},
)

3. （可选）拆分结果维度 #

若需将不同类型结果分开处理，可对返回的DataArray进行拆分：

# 拆分系数、R²值、预测值
coefs = test.isel(result=[0,1]).rename(result="parameter")
r2 = test.isel(result=2)
y_pred = test.isel(result=slice(3, None)).rename(result="time")

关键注意事项 #

• 多变量回归的特征矩阵必须是二维数组（shape=(n_samples, n_features)），不能分开传入单个自变量
• Scikit-learn模型的fit()/score()/predict()方法参数格式需严格匹配：fit(X, y)、score(X, y)、predict(X)
• 若需计算岭回归系数的p值，Scikit-learn无原生实现，可采用以下方式：
  • 对每个网格点使用bootstrap方法做统计检验
  • 使用基于自由度调整的近似t检验（仅适用于近似场景）

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