def test_get_column_map_features_changes_cols(self):
sk = SklearnTransformerMock()
sk.features_ = [1]
t = r.SklearnColumnsWrapper(sk)
t.cols = ["amount", "color"]
t.get_column_mapping()
self.assertListEqual(t.cols, ["color"])
def test_get_column_map_features(self):
sk = SklearnTransformerMock()
sk.features_ = [1]
t = r.SklearnColumnsWrapper(sk)
t.cols = ["amount", "color"]
d = t.get_column_mapping()
d_expected = {"color": "color"}
self.assertDictEqual(d, d_expected)
def test_fit(self):
sk = SklearnTransformerMock()
sk.fit = MagicMock()
t = r.SklearnFitOneWrapper(sk, "price", "amount")
df_out = t.fit(create_df_all())
a = np.array([1.5, 1, 2.5, 2, 3.5, 3]).reshape(-1, 1)
a_out = sk.fit.call_args_list[0][0][0]
self.assertEqual(a.shape, a_out.shape)
# flatten() is not really needed.
self.assertListEqual(list(a.flatten()), list(a_out.flatten()))
def _test_get_column_map_features_name(self, col_format, expected):
sk = SklearnTransformerMock()
sk.get_feature_names = MagicMock(return_value=["0_blue", "0_red"])
if col_format is not None:
t = r.SklearnColumnsWrapper(sk, col_format=col_format)
else:
t = r.SklearnColumnsWrapper(sk)
t.cols = ["color"]
d = t.get_column_mapping()
d_expected = {"color": tuple(expected)}
self.assertDictEqual(d, d_expected)
def test_fit_transform_features_name_result(self):
sk = SklearnTransformerMock()
sk.get_feature_names = MagicMock(return_value=["0_blue", "0_red"])
return_value = np.array([[1, 2, 3], [3, 2, 1]]).T
sk.transform = MagicMock(return_value=return_value)
t = r.SklearnColumnsWrapper(sk, "color")
df_out = t.fit_transform(create_df_all())
df_expected = pd.DataFrame({
"color=blue": [1, 2, 3],
"color=red": [3, 2, 1],
"price": [1.5, 2.5, 3.5],
"amount": [1, 2, 3]
})
sk.transform.assert_called_once()
self.assertTrue(df_out.equals(df_expected))
def test_get_column_map_features_name_mix(self):
"""One transformer with get_feature_names and other without. """
sk1 = SklearnTransformerMock()
return_value = np.array(["0_blue", "0_red"])
sk1.get_feature_names = MagicMock(return_value=return_value)
sk2 = SklearnTransformerMock()
t = r.SklearnColumnWrapper(None)
t.transformers = {"color": sk1, "amount": sk2}
t.cols = ["color", "amount"]
d = t.get_column_mapping()
d_expected = {
"color": tuple(["color=blue", "color=red"]),
"amount": "amount"
}
self.assertDictEqual(d, d_expected)
def test_fit_transform(self):
sk = SklearnTransformerMock()
t = r.SklearnColumnWrapper(sk)
t.fit_transform(create_df_cat())
# Original transformer should not be fitted, only copies.
self.assertEqual(sk.n_fit, 0)
self.assertTrue("color" in t.transformers)
self.assertEqual(t.transformers["color"].n_fit, 1)
self.assertEqual(t.transformers["color"].n_transform, 1)
def _test_get_column_map_features_selected(self, selected, features):
sk1 = SklearnTransformerMock()
if features:
sk1.features_ = [0] if selected else []
sk2 = SklearnTransformerMock()
sk2.features_ = [0]
t = r.SklearnColumnWrapper(None)
t.transformers = {"amount": sk1, "color": sk2}
t.cols = ["amount", "color"]
t.get_column_mapping()
cols_expected = ["amount", "color"] if selected else ["color"]
self.assertListEqual(t.cols, cols_expected)