@@ -34,15 +34,16 @@ def test_default_imputation(input_data_imputation, categorical):
34
34
X = X .astype ('str' ).astype ('object' )
35
35
X [mask ] = np .nan
36
36
else :
37
- imputation_value = 0
37
+ imputation_value = min (np .unique (X )) - 1
38
+
38
39
Y = CategoricalImputation ().fit_transform (X .copy ())
39
- assert ((np .argwhere (Y == imputation_value ) == np .argwhere (mask )).all ())
40
- assert ((np .argwhere (Y != imputation_value ) == np .argwhere (np .logical_not (mask ))).all ())
40
+
41
+ assert np .array_equal (Y == imputation_value , mask )
42
+ assert np .array_equal (Y != imputation_value , ~ mask )
41
43
42
44
43
45
@pytest .mark .parametrize ('format_type' , ('numpy' , 'pandas' ))
44
46
def test_nonzero_numerical_imputation (format_type ):
45
-
46
47
# First try with an array with 0 as only valid category. The imputation should
47
48
# happen with -1
48
49
X = np .full (fill_value = np .nan , shape = (10 , 10 ))
@@ -69,8 +70,9 @@ def test_nonzero_numerical_imputation(format_type):
69
70
@pytest .mark .parametrize ('input_data_imputation' , ('numpy' ), indirect = True )
70
71
def test_default_sparse (input_data_imputation ):
71
72
X , mask = input_data_imputation
72
- X = sparse .csc_matrix (X )
73
+ X = sparse .csr_matrix (X )
73
74
Y = CategoricalImputation ().fit_transform (X )
74
75
Y = Y .todense ()
75
- assert (np .argwhere (Y == 0 ) == np .argwhere (mask )).all ()
76
- assert (np .argwhere (Y != 0 ) == np .argwhere (np .logical_not (mask ))).all ()
76
+
77
+ assert np .array_equal (Y == - 1 , mask )
78
+ assert np .array_equal (Y != - 1 , ~ mask )
0 commit comments