fix notebook and add features_names to ArrayDataset

and allow providing features names in QI and Cat features not just indexes

apt/anonymization/anonymizer.py

@@ -22,19 +22,19 @@ class Anonymize:
         """
         :param k: The privacy parameter that determines the number of records that will be indistinguishable from each
                   other (when looking at the quasi identifiers). Should be at least 2.
-        :param quasi_identifiers: The indexes of features that need to be minimized.
-        :param categorical_features: The list of categorical features indexes
+        :param quasi_identifiers: The features that need to be minimized.
+        :param categorical_features: The list of categorical features.
         :param is_regression: Boolean param indicates that is is a regression problem.
         """
         if k < 2:
             raise ValueError("k should be a positive integer with a value of 2 or higher")
         if quasi_identifiers is None or len(quasi_identifiers) < 1:
             raise ValueError("The list of quasi-identifiers cannot be empty")
-
         self.k = k
         self.quasi_identifiers = quasi_identifiers
         self.categorical_features = categorical_features
         self.is_regression = is_regression
+        self.features_names = None
 
     def anonymize(self, dataset: ArrayDataset) -> DATA_PANDAS_NUMPY_TYPE:
         """
@@ -51,6 +51,15 @@ class Anonymize:
             self._features = [i for i in range(dataset.get_samples().shape[0])]
         else:
             raise ValueError('No data provided')
+        if not set(self.quasi_identifiers).issubset(set(self.features_names)):
+            raise ValueError('Quasi identifiers should bs a subset of the supplied features or indexes in range of '
+                             'the data columns')
+        if self.categorical_features and not set(self.categorical_features).issubset(set(self.features_names)):
+            raise ValueError('Categorical features should bs a subset of the supplied features or indexes in range of '
+                             'the data columns')
+        self.quasi_identifiers = [i for i, v in enumerate(self.features_names) if v in self.quasi_identifiers]
+        if self.categorical_features:
+            self.categorical_features = [i for i, v in enumerate(self.features_names) if v in self.categorical_features]
 
         transformed = self._anonymize(dataset.get_samples().copy(), dataset.get_labels())
         if dataset.is_pandas:

apt/utils/dataset_utils.py

@@ -273,7 +273,7 @@ def get_nursery_dataset(raw: bool = True, test_set: float = 0.2, transform_socia
             raise Exception("Bad label value: %s" % value)
 
     data["label"] = data["label"].apply(modify_label)
-    data["children"] = data["children"].apply(lambda x: 4 if x == "more" else x)
+    data["children"] = data["children"].apply(lambda x: "4" if x == "more" else x)
 
     if transform_social:
 

apt/utils/datasets/datasets.py

@@ -18,7 +18,6 @@ from torch import Tensor
 
 logger = logging.getLogger(__name__)
 
-
 INPUT_DATA_ARRAY_TYPE = Union[np.ndarray, pd.DataFrame, List, Tensor]
 OUTPUT_DATA_ARRAY_TYPE = np.ndarray
 DATA_PANDAS_NUMPY_TYPE = Union[np.ndarray, pd.DataFrame]
@@ -113,7 +112,6 @@ class StoredDataset(Dataset):
         if unzip:
             StoredDataset.extract_archive(zip_path=file_path, dest_path=dest_path, remove_archive=False)
 
-
     @staticmethod
     def extract_archive(zip_path: str, dest_path=None, remove_archive=False):
         """
@@ -164,7 +162,8 @@ class StoredDataset(Dataset):
 class ArrayDataset(Dataset):
     """Dataset that is based on x and y arrays (e.g., numpy/pandas/list...)"""
 
-    def __init__(self, x: INPUT_DATA_ARRAY_TYPE, y: Optional[INPUT_DATA_ARRAY_TYPE] = None, **kwargs):
+    def __init__(self, x: INPUT_DATA_ARRAY_TYPE, y: Optional[INPUT_DATA_ARRAY_TYPE] = None, features_names=None,
+                 **kwargs):
         """
         ArrayDataset constructor.
         :param x: collection of data samples
@@ -172,10 +171,12 @@ class ArrayDataset(Dataset):
         :param kwargs: dataset parameters
         """
         self.is_pandas = False
-        self.features_names = None
+        self.features_names = features_names
         self._y = array2numpy(self, y) if y is not None else None
         self._x = array2numpy(self, x)
         if self.is_pandas:
+            if features_names and not np.array_equal(features_names, x.columns):
+                raise ValueError("The supplied features are not the same as in the data features")
             self.features_names = x.columns
 
         if y is not None and len(self._x) != len(self._y):
@@ -213,7 +214,6 @@ class PytorchData(Dataset):
         else:
             self.__getitem__ = self.get_sample_item
 
-
     def get_samples(self) -> OUTPUT_DATA_ARRAY_TYPE:
         """Return data samples as numpy array"""
         return array2numpy(self._x)
@@ -244,6 +244,7 @@ class DatasetFactory:
         :param name: dataset name
         :return:
         """
+
         def inner_wrapper(wrapped_class: Dataset) -> Any:
             if name in cls.registry:
                 logger.warning('Dataset %s already exists. Will replace it', name)