Merge pull request #115 from maxsch3/bug/114

#110 fixed and tests updated. Indices were fixed and numpy conversion…

Author: maxsch3

keras_batchflow/base/batch_transformers/base_random_cell.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 from .batch_transformer import BatchTransformer
 
 
@@ -315,16 +316,25 @@ def transform(self, batch):
                 raise KeyError(f'Error: The data passed to {type(self).__name__} is not forked, while fork parameter '
                                f'is specified. Please add multiindex level to columns of your data or use DataFork '
                                f'batch transform before.')
-            subset = batch[self._fork][self._cols]
+            if self._fork not in batch.columns.get_level_values(0):
+                raise KeyError(f"Error: fork {self._fork} specified as a parameter 'data_fork' was not found in data. "
+                               f"The following forks were found: {set(batch.columns.get_level_values(0))}. Please "
+                               f"make sure you are using DataFork that is configured to provide this a fork with the"
+                               f"name specified.")
+            # the top level of multiinedex is dropped here to avoid a hassle of handling it in methods _make_mask and
+            # _make_augmented_version. This dropped level will be added later when merged back with the batch
+            subset = batch[self._fork][self._cols].copy()
         else:
-            subset = batch[self._cols]
+            subset = batch[self._cols].copy()
         mask = self._make_mask(subset)
         augmented_batch = self._make_augmented_version(subset)
-        transformed = subset.mask(mask.astype(bool), augmented_batch.values)
+        transformed = subset.mask(mask.astype(bool), augmented_batch)
         if self._fork:
-            batch.loc(axis=1)[self._fork, self._cols] = transformed.values
+            # in order for loc to work, the top level index must be restored
+            transformed.columns = pd.MultiIndex.from_product([[self._fork], transformed.columns])
+            batch.loc[:, (self._fork, self._cols)] = transformed
         else:
-            batch[self._cols] = transformed
+            batch.loc[:, self._cols] = transformed
         return batch
 
     def inverse_transform(self, batch):

keras_batchflow/base/batch_transformers/shuffle_noise.py

@@ -33,5 +33,5 @@ class ShuffleNoise(BaseRandomCellTransform):
     """
 
     def _make_augmented_version(self, batch):
-        batch1 = batch.apply(lambda x: x.sample(frac=1).values)
-        return batch1
+        augmented_batch = batch.apply(lambda x: x.sample(frac=1).values)
+        return augmented_batch

test/test_base_random_cell_transform.py

@@ -2,7 +2,7 @@
 import pandas as pd
 import numpy as np
 from scipy.stats import binom_test, chisquare
-from keras_batchflow.base.batch_transformers import BaseRandomCellTransform
+from keras_batchflow.base.batch_transformers import BaseRandomCellTransform, BatchFork
 
 
 class LocalVersionTransform(BaseRandomCellTransform):
@@ -17,6 +17,19 @@ def _make_augmented_version(self, batch):
         return batch
 
 
+class TestTransformInt(BaseRandomCellTransform):
+    """
+    BaseRandomCellTransform raise NotImplemented error in below function which does not allow to test
+    transform functionality. I'm re-defining this method here to be able to test it.
+    """
+
+    def _make_augmented_version(self, batch):
+        batch = batch.copy()
+        for c in self._cols:
+            batch[c] = 0
+        return batch
+
+
 class TestFeatureDropout:
 
     df = None
@@ -188,6 +201,19 @@ def test_transform_fork_many_cols(self):
         assert (batch1['x']['var2'] == '').all()
         assert (batch1['x']['label'] != '').all()
 
+    def test_non_numpy_dtype(self):
+        """
+        This test is to make sure the transform does not convert data to numpy behind the scenes, causing
+        unpredictable dtype changes
+        :return:
+        """
+        data = pd.DataFrame({'var1': np.random.randint(low=0, high=10, size=100)}).astype('Int64')
+        data.iloc[0, 0] = None
+        data_forked = BatchFork().transform(data.copy())
+        ct = TestTransformInt([0., 1.], cols=['var1'], data_fork='x')
+        data_transformed = ct.transform(data_forked)
+        assert all(dt.name == 'Int64' for dt in data_transformed.dtypes)
+
 
     # def test_row_dist(self):
     #     fd = FeatureDropout(.6, 'var1', '')