Merge pull request #22 from zStupan/fix-dataset

Fix dataset

Author: zStupan

niaarm/dataset.py

@@ -39,9 +39,16 @@ def __analyse_types(self):
                 min_value = col.min()
                 max_value = col.max()
                 unique_categories = None
+            elif col.dtype == 'bool':
+                self.data[head] = self.data[head].astype(int)
+                self.transactions = self.data.values
+                dtype = 'int'
+                min_value = 0
+                max_value = 1
+                unique_categories = None
             else:
                 dtype = "cat"
-                unique_categories = sorted(col.astype(str).unique().tolist(), key=str.lower)
+                unique_categories = sorted(col.unique().tolist(), key=str.lower)
                 min_value = None
                 max_value = None
 

niaarm/tests/test_read_csv.py

@@ -28,7 +28,7 @@ def test_read_features(self):
             'float',
             'float',
             'int']
-        
+
         data = Dataset(os.path.join(os.path.dirname(__file__), 'test_data', 'Abalone.csv'))
 
         features = data.features
@@ -60,7 +60,7 @@ def test_read_features(self):
         minval = [None, 0]
         maxval = [None, 1]
         dtypes_a = ['cat', 'int']
-        
+
         data = Dataset(os.path.join(os.path.dirname(__file__), 'test_data', 'wiki_test_case.csv'))
 
         features = data.features

tests/test_niaarm.py

@@ -1,4 +1,3 @@
 from niaarm.tests.conftest import pytest_configure
 
 __all__ = ["pytest_configure"]
-