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