|
|
import numpy as np |
|
|
import pytest |
|
|
|
|
|
from sklearn.impute._base import _BaseImputer |
|
|
from sklearn.impute._iterative import _assign_where |
|
|
from sklearn.utils._mask import _get_mask |
|
|
from sklearn.utils._testing import _convert_container, assert_allclose |
|
|
|
|
|
|
|
|
@pytest.fixture |
|
|
def data(): |
|
|
X = np.random.randn(10, 2) |
|
|
X[::2] = np.nan |
|
|
return X |
|
|
|
|
|
|
|
|
class NoFitIndicatorImputer(_BaseImputer): |
|
|
def fit(self, X, y=None): |
|
|
return self |
|
|
|
|
|
def transform(self, X, y=None): |
|
|
return self._concatenate_indicator(X, self._transform_indicator(X)) |
|
|
|
|
|
|
|
|
class NoTransformIndicatorImputer(_BaseImputer): |
|
|
def fit(self, X, y=None): |
|
|
mask = _get_mask(X, value_to_mask=np.nan) |
|
|
super()._fit_indicator(mask) |
|
|
return self |
|
|
|
|
|
def transform(self, X, y=None): |
|
|
return self._concatenate_indicator(X, None) |
|
|
|
|
|
|
|
|
class NoPrecomputedMaskFit(_BaseImputer): |
|
|
def fit(self, X, y=None): |
|
|
self._fit_indicator(X) |
|
|
return self |
|
|
|
|
|
def transform(self, X): |
|
|
return self._concatenate_indicator(X, self._transform_indicator(X)) |
|
|
|
|
|
|
|
|
class NoPrecomputedMaskTransform(_BaseImputer): |
|
|
def fit(self, X, y=None): |
|
|
mask = _get_mask(X, value_to_mask=np.nan) |
|
|
self._fit_indicator(mask) |
|
|
return self |
|
|
|
|
|
def transform(self, X): |
|
|
return self._concatenate_indicator(X, self._transform_indicator(X)) |
|
|
|
|
|
|
|
|
def test_base_imputer_not_fit(data): |
|
|
imputer = NoFitIndicatorImputer(add_indicator=True) |
|
|
err_msg = "Make sure to call _fit_indicator before _transform_indicator" |
|
|
with pytest.raises(ValueError, match=err_msg): |
|
|
imputer.fit(data).transform(data) |
|
|
with pytest.raises(ValueError, match=err_msg): |
|
|
imputer.fit_transform(data) |
|
|
|
|
|
|
|
|
def test_base_imputer_not_transform(data): |
|
|
imputer = NoTransformIndicatorImputer(add_indicator=True) |
|
|
err_msg = ( |
|
|
"Call _fit_indicator and _transform_indicator in the imputer implementation" |
|
|
) |
|
|
with pytest.raises(ValueError, match=err_msg): |
|
|
imputer.fit(data).transform(data) |
|
|
with pytest.raises(ValueError, match=err_msg): |
|
|
imputer.fit_transform(data) |
|
|
|
|
|
|
|
|
def test_base_no_precomputed_mask_fit(data): |
|
|
imputer = NoPrecomputedMaskFit(add_indicator=True) |
|
|
err_msg = "precomputed is True but the input data is not a mask" |
|
|
with pytest.raises(ValueError, match=err_msg): |
|
|
imputer.fit(data) |
|
|
with pytest.raises(ValueError, match=err_msg): |
|
|
imputer.fit_transform(data) |
|
|
|
|
|
|
|
|
def test_base_no_precomputed_mask_transform(data): |
|
|
imputer = NoPrecomputedMaskTransform(add_indicator=True) |
|
|
err_msg = "precomputed is True but the input data is not a mask" |
|
|
imputer.fit(data) |
|
|
with pytest.raises(ValueError, match=err_msg): |
|
|
imputer.transform(data) |
|
|
with pytest.raises(ValueError, match=err_msg): |
|
|
imputer.fit_transform(data) |
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("X1_type", ["array", "dataframe"]) |
|
|
def test_assign_where(X1_type): |
|
|
"""Check the behaviour of the private helpers `_assign_where`.""" |
|
|
rng = np.random.RandomState(0) |
|
|
|
|
|
n_samples, n_features = 10, 5 |
|
|
X1 = _convert_container(rng.randn(n_samples, n_features), constructor_name=X1_type) |
|
|
X2 = rng.randn(n_samples, n_features) |
|
|
mask = rng.randint(0, 2, size=(n_samples, n_features)).astype(bool) |
|
|
|
|
|
_assign_where(X1, X2, mask) |
|
|
|
|
|
if X1_type == "dataframe": |
|
|
X1 = X1.to_numpy() |
|
|
assert_allclose(X1[mask], X2[mask]) |
|
|
|
