Update categorical transformers

rdt/transformers/categorical.py

@@ -33,6 +33,12 @@ class CategoricalTransformer(BaseTransformer):
             Defaults to ``False``.
     """
 
+    INPUT_TYPE = 'categorical'
+    OUTPUT_TYPES = {'value': 'float'}
+    DETERMINISTIC_TRANSFORM = False
+    DETERMINISTIC_REVERSE = True
+    COMPOSITION_IS_IDENTITY = True
+
     mapping = None
     intervals = None
     starts = None
@@ -93,7 +99,7 @@ def _get_intervals(data):
 
         return intervals, means, starts
 
-    def fit(self, data):
+    def _fit(self, data):
         """Fit the transformer to the data.
 
         Create the mapping dict to save the label encoding.
@@ -145,9 +151,9 @@ def _get_value(self, category):
 
     def _transform_by_row(self, data):
         """Transform the data row by row."""
-        return data.fillna(np.nan).apply(self._get_value).to_numpy()
+        return data.fillna(np.nan).apply(self._get_value)
 
-    def transform(self, data):
+    def _transform(self, data):
         """Transform categorical values to float values.
 
         Replace the categories with their float representative value.
@@ -210,7 +216,7 @@ def _reverse_transform_by_row(self, data):
         """Reverse transform the data by iterating over each row."""
         return data.apply(self._get_category_from_start).astype(self.dtype)
 
-    def reverse_transform(self, data):
+    def _reverse_transform(self, data):
         """Convert float values back to the original categorical values.
 
         Args:
@@ -259,6 +265,10 @@ class OneHotEncodingTransformer(BaseTransformer):
             transform, then an error will be raised if this is True.
     """
 
+    INPUT_TYPE = 'categorical'
+    DETERMINISTIC_TRANSFORM = True
+    DETERMINISTIC_REVERSE = True
+
     dummies = None
     _dummy_na = None
     _num_dummies = None
@@ -277,7 +287,7 @@ def _prepare_data(data):
         otherwise returns it.
 
         Args:
-            data (pandas.Series, numpy.ndarray, list or list of lists):
+            data (pandas.Series or pandas.DataFrame):
                 Data to prepare.
 
         Returns:
@@ -296,33 +306,24 @@ def _prepare_data(data):
 
         return data
 
-    def _transform(self, data):
-        if self._dummy_encoded:
-            coder = self._indexer
-            codes = pd.Categorical(data, categories=self._uniques).codes
-        else:
-            coder = self._uniques
-            codes = data
+    def get_output_types(self):
+        """Return the output types produced by this transformer.
 
-        rows = len(data)
-        dummies = np.broadcast_to(coder, (rows, self._num_dummies))
-        coded = np.broadcast_to(codes, (self._num_dummies, rows)).T
-        array = (coded == dummies).astype(int)
-
-        if self._dummy_na:
-            null = np.zeros((rows, 1), dtype=int)
-            null[pd.isnull(data)] = 1
-            array = np.append(array, null, axis=1)
+        Returns:
+            dict:
+                Mapping from the transformed column names to the produced data types.
+        """
+        output_types = {f'value{i}': 'float' for i in range(len(self.dummies))}
 
-        return array
+        return self._add_prefix(output_types)
 
-    def fit(self, data):
+    def _fit(self, data):
         """Fit the transformer to the data.
 
         Get the pandas `dummies` which will be used later on for OneHotEncoding.
 
         Args:
-            data (pandas.Series, numpy.ndarray, list or list of lists):
+            data (pandas.Series or pandas.DataFrame):
                 Data to fit the transformer to.
         """
         data = self._prepare_data(data)
@@ -340,7 +341,27 @@ def fit(self, data):
         if self._dummy_na:
             self.dummies.append(np.nan)
 
-    def transform(self, data):
+    def _transform_helper(self, data):
+        if self._dummy_encoded:
+            coder = self._indexer
+            codes = pd.Categorical(data, categories=self._uniques).codes
+        else:
+            coder = self._uniques
+            codes = data
+
+        rows = len(data)
+        dummies = np.broadcast_to(coder, (rows, self._num_dummies))
+        coded = np.broadcast_to(codes, (self._num_dummies, rows)).T
+        array = (coded == dummies).astype(int)
+
+        if self._dummy_na:
+            null = np.zeros((rows, 1), dtype=int)
+            null[pd.isnull(data)] = 1
+            array = np.append(array, null, axis=1)
+
+        return array
+
+    def _transform(self, data):
         """Replace each category with the OneHot vectors.
 
         Args:
@@ -351,7 +372,7 @@ def transform(self, data):
             numpy.ndarray:
         """
         data = self._prepare_data(data)
-        array = self._transform(data)
+        array = self._transform_helper(data)
 
         if self.error_on_unknown:
             unknown = array.sum(axis=1) == 0
@@ -361,7 +382,7 @@ def transform(self, data):
 
         return array
 
-    def reverse_transform(self, data):
+    def _reverse_transform(self, data):
         """Convert float values back to the original categorical values.
 
         Args:
@@ -371,6 +392,9 @@ def reverse_transform(self, data):
         Returns:
             pandas.Series
         """
+        if not isinstance(data, np.ndarray):
+            data = data.to_numpy()
+
         if data.ndim == 1:
             data = data.reshape(-1, 1)
 
@@ -394,10 +418,16 @@ class LabelEncodingTransformer(BaseTransformer):
             integer value.
     """
 
+    INPUT_TYPE = 'categorical'
+    OUTPUT_TYPES = {'value': 'integer'}
+    DETERMINISTIC_TRANSFORM = True
+    DETERMINISTIC_REVERSE = True
+    COMPOSITION_IS_IDENTITY = True
+
     values_to_categories = None
     categories_to_values = None
 
-    def fit(self, data):
+    def _fit(self, data):
         """Fit the transformer to the data.
 
         Generate a unique integer representation for each category and
@@ -414,7 +444,7 @@ def fit(self, data):
             for value, category in self.values_to_categories.items()
         }
 
-    def transform(self, data):
+    def _transform(self, data):
         """Replace each category with its corresponding integer value.
 
         Args:
@@ -427,9 +457,9 @@ def transform(self, data):
         if not isinstance(data, pd.Series):
             data = pd.Series(data)
 
-        return data.map(self.categories_to_values)
+        return pd.Series(data).map(self.categories_to_values)
 
-    def reverse_transform(self, data):
+    def _reverse_transform(self, data):
         """Convert float values back to the original categorical values.
 
         Args:
@@ -439,8 +469,5 @@ def reverse_transform(self, data):
         Returns:
             pandas.Series
         """
-        if isinstance(data, np.ndarray) and (data.ndim == 2):
-            data = data[:, 0]
-
         data = data.clip(min(self.values_to_categories), max(self.values_to_categories))
-        return pd.Series(data).round().map(self.values_to_categories)
+        return data.round().map(self.values_to_categories)

tests/integration/test_hyper_transformer.py

@@ -228,7 +228,7 @@ def test_dtype_category():
 
     rever = ht.reverse_transform(trans)
 
-    pd.testing.assert_frame_equal(df, rever)
+    pd.testing.assert_frame_equal(rever, df)
 
 
 def test_empty_transformers():