def test_one_hot_numerical_nans():
"""Ensure OneHotEncodingTransformer works on numerical + nan only columns."""
data = pd.Series([1, 2, float('nan'), np.nan])
transformer = OneHotEncodingTransformer()
transformer.fit(data)
transformed = transformer.transform(data)
reverse = transformer.reverse_transform(transformed)
pd.testing.assert_series_equal(reverse, data)
def test_reverse_transform_no_nans(self):
# Setup
ohet = OneHotEncodingTransformer()
data = pd.Series(['a', 'b', 'c'])
ohet.fit(data)
# Run
transformed = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
out = ohet.reverse_transform(transformed)
# Assert
expected = pd.Series(['a', 'b', 'c'])
pd.testing.assert_series_equal(out, expected)