sklearn的IterativeImputer使用PLSRegression时报形状不匹配错误

当使用PLSRegression作为IterativeImputer的估计器时触发形状不匹配错误，根本原因是PLSRegression.predict()方法返回的是形状为(n_samples, 1)的二维数组，而IterativeImputer内部填充缺失值时，期望接收形状为(n_samples,)的一维数组，二者无法直接广播赋值，所以触发_iterative.py第348行的ValueError。

无需修改源码的临时方案 #

对PLSRegression做一层简单封装，将预测输出的二维数组压缩为一维即可，修改后的可运行代码如下：

import numpy as np

from sklearn.datasets import fetch_california_housing
from sklearn.cross_decomposition import PLSRegression
from sklearn.experimental import enable_iterative_imputer  # noqa
from sklearn.impute import IterativeImputer
from sklearn.base import RegressorMixin, BaseEstimator

# 自定义封装PLSRegression适配IterativeImputer
class PLSRegressionWrapper(BaseEstimator, RegressorMixin):
    def __init__(self, n_components=2):
        self.n_components = n_components
        self.model = PLSRegression(n_components=self.n_components)
    
    def fit(self, X, y):
        self.model.fit(X, y)
        return self
    
    def predict(self, X):
        # 将返回的二维数组压缩为一维匹配IterativeImputer的要求
        return self.model.predict(X).flatten()

rng = np.random.RandomState(42)

X_california, y_california = fetch_california_housing(return_X_y=True)
X_california = X_california[:400]
y_california = y_california[:400]

def add_missing_values(X_full, y_full):
    n_samples, n_features = X_full.shape

    # 75%的行添加缺失值
    missing_rate = 0.75
    n_missing_samples = int(n_samples * missing_rate)

    missing_samples = np.zeros(n_samples, dtype=bool)
    missing_samples[: n_missing_samples] = True

    rng.shuffle(missing_samples)
    missing_features = rng.randint(0, n_features, n_missing_samples)
    X_missing = X_full.copy()
    X_missing[missing_samples, missing_features] = np.nan
    y_missing = y_full.copy()

    return X_missing, y_missing

X_miss_california, y_miss_california = add_missing_values(
    X_california, y_california)

# 改用封装后的估计器
imputer = IterativeImputer(estimator=PLSRegressionWrapper(n_components=2))

X_imputed = imputer.fit_transform(X_miss_california)
print(X_imputed)

运行上述代码即可得到你提供的预期输出结果。

永久适配方案 #

如果你愿意修改sklearn源码适配所有同类回归器，可找到sklearn/impute/_iterative.py第348行，将：

X_filled[missing_row_mask, feat_idx] = imputed_values

修改为：

X_filled[missing_row_mask, feat_idx] = imputed_values.ravel()

修改后所有返回二维预测结果的回归器都可以直接作为IterativeImputer的估计器使用。

内容的提问来源于stack exchange，提问作者DAKSHA DINESH SHENOY