cesnet-datazoo 0.0.16__py3-none-any.whl → 0.1.0__py3-none-any.whl

cesnet_datazoo/pytables_data/pytables_dataset.py

@@ -4,62 +4,118 @@ import os
 import time
 import warnings
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import numpy as np
 import pandas as pd
 import tables as tb
 import torch
-from numpy.lib.recfunctions import drop_fields, structured_to_unstructured
-from sklearn.preprocessing import LabelEncoder, MinMaxScaler, RobustScaler, StandardScaler
+from numpy.lib.recfunctions import structured_to_unstructured
 from torch.utils.data import Dataset
 from typing_extensions import assert_never
 
-from cesnet_datazoo.config import (AppSelection, DatasetConfig, MinTrainSamplesCheck, Scaler,
-                                   ScalerEnum, TestDataParams, TrainDataParams)
-from cesnet_datazoo.constants import (APP_COLUMN, CATEGORY_COLUMN, DIR_POS, FLOWSTATS_NO_CLIP,
-                                      FLOWSTATS_TO_SCALE, INDICES_INDEX_POS, INDICES_TABLE_POS,
-                                      IPT_POS, PHIST_BIN_COUNT, PHISTS_FEATURES, PPI_COLUMN,
-                                      SIZE_POS, UNKNOWN_STR_LABEL)
+from cesnet_datazoo.config import (AppSelection, MinTrainSamplesCheck, TestDataParams,
+                                   TrainDataParams)
+from cesnet_datazoo.constants import APP_COLUMN, INDICES_INDEX_POS, INDICES_TABLE_POS, PPI_COLUMN
 from cesnet_datazoo.pytables_data.apps_split import (is_background_app,
                                                      split_apps_topx_with_provider_groups)
-from cesnet_datazoo.utils.fileutils import pickle_dump, pickle_load
-from cesnet_datazoo.utils.random import RandomizedSection, get_fresh_random_generator
 
 log = logging.getLogger(__name__)
 
 
 class PyTablesDataset(Dataset):
-    def __init__(self, database_path: str,
+    def __init__(self,
+                 database_path: str,
                  tables_paths: list[str],
                  indices: Optional[np.ndarray],
+                 tables_app_enum: dict[int, str],
+                 tables_cat_enum: dict[int, str],
                  flowstats_features: list[str],
-                 other_fields: Optional[list[str]] = None,
-                 preload: bool = False, preload_blob: Optional[str] = None,
-                 disabled_apps: Optional[list[str]] = None,
-                 return_all_fields: bool = False,):
+                 flowstats_features_boolean: list[str],
+                 flowstats_features_phist: list[str],
+                 other_fields: list[str],
+                 ppi_channels: list[int],
+                 ppi_transform: Optional[Callable] = None,
+                 flowstats_transform: Optional[Callable] = None,
+                 flowstats_phist_transform: Optional[Callable] = None,
+                 target_transform: Optional[Callable] = None,
+                 return_tensors: bool = False,
+                 return_all_fields: bool = False,
+                 preload: bool = False,
+                 preload_blob: Optional[str] = None,
+                 disabled_apps: Optional[list[str]] = None,):
         self.database_path = database_path
         self.tables_paths = tables_paths
         self.tables = {}
-        self.flowstats_features = flowstats_features
-        self.other_fields = other_fields if other_fields is not None else []
-        self.preload = preload
-        self.preload_blob = preload_blob
-        self.return_all_fields = return_all_fields
+        self.tables_app_enum = tables_app_enum
+        self.tables_app_arr = np.array(list(tables_app_enum.values()))
+        self.tables_cat_enum = tables_cat_enum
         if indices is None:
             self.set_all_indices(disabled_apps=disabled_apps)
         else:
             self.indices = indices
 
+        self.flowstats_features = flowstats_features
+        self.flowstats_features_boolean = flowstats_features_boolean
+        self.flowstats_features_phist = flowstats_features_phist
+        self.other_fields = other_fields
+        self.ppi_channels = ppi_channels
+        self.ppi_transform = ppi_transform
+        self.flowstats_transform = flowstats_transform
+        self.flowstats_phist_transform = flowstats_phist_transform
+        self.target_transform = target_transform
+        self.return_tensors = return_tensors
+        self.return_all_fields = return_all_fields
+
+        self.preload = preload
+        self.preload_blob = preload_blob
+
     def __getitem__(self, batch_idx):
         # log.debug(f"worker {self.worker_id}: __getitem__")
         if self.preload:
             batch_data = self.data[batch_idx]
         else:
-            batch_data = load_data_from_pytables(tables=self.tables, indices=self.indices[batch_idx], data_dtype=self.data_dtype)
+            batch_data = load_data_from_tables(tables=self.tables, indices=self.indices[batch_idx], data_dtype=self.data_dtype)
         if self.return_all_fields:
             return (batch_data, batch_idx)
-        return_data = (batch_data[self.other_fields], batch_data[PPI_COLUMN].astype("float32"), batch_data[self.flowstats_features], list(map(self.app_enum, batch_data[APP_COLUMN])))
+
+        # Prepare data
+        x_ppi = batch_data[PPI_COLUMN].astype("float32")
+        x_ppi = x_ppi[:, self.ppi_channels, :]
+        x_flowstats = structured_to_unstructured(batch_data[self.flowstats_features], dtype="float32")
+        if self.flowstats_features_boolean:
+            x_flowstats_boolean = structured_to_unstructured(batch_data[self.flowstats_features_boolean], dtype="float32")
+        else:
+            x_flowstats_boolean = np.zeros(shape=(x_flowstats.shape[0], 0), dtype="float32")
+        if self.flowstats_features_phist:
+            x_flowstats_phist = structured_to_unstructured(batch_data[self.flowstats_features_phist], dtype="float32")
+        else:
+            x_flowstats_phist = np.zeros(shape=(x_flowstats.shape[0], 0), dtype="float32")
+        # Feature transformations
+        if self.ppi_transform:
+            x_ppi = self.ppi_transform(x_ppi)
+        if self.flowstats_transform:
+            x_flowstats = self.flowstats_transform(x_flowstats)
+        if self.flowstats_phist_transform:
+            x_flowstats_phist = self.flowstats_phist_transform(x_flowstats_phist)
+        x_flowstats = np.concatenate([x_flowstats, x_flowstats_boolean, x_flowstats_phist], axis=1).astype("float32")
+        # Labels transformation
+        labels = self.tables_app_arr[batch_data[APP_COLUMN]]
+        if self.target_transform:
+            labels = self.target_transform(labels)
+        # Prepare dataframe with other fields
+        other_fields_df = pd.DataFrame(batch_data[self.other_fields]) if len(self.other_fields) > 0 else pd.DataFrame()
+        for column in other_fields_df.columns:
+            if other_fields_df[column].dtype.kind == "O":
+                other_fields_df[column] = other_fields_df[column].astype(str)
+            elif column.startswith("TIME_"):
+                other_fields_df[column] = other_fields_df[column].map(lambda x: datetime.fromtimestamp(x))
+
+        if self.return_tensors:
+            x_ppi = torch.from_numpy(x_ppi)
+            x_flowstats = torch.from_numpy(x_flowstats)
+            labels = torch.from_numpy(labels).long() # PyTorch loss functions require long type for labels
+        return_data = (other_fields_df, x_ppi, x_flowstats, labels)
         return return_data
 
     def __len__(self):
@@ -70,8 +126,6 @@ class PyTablesDataset(Dataset):
         log.debug(f"Initializing dataloader worker id {self.worker_id}")
         self.database, self.tables = load_database(database_path=self.database_path, tables_paths=self.tables_paths)
         atexit.register(self.cleanup)
-        self.app_enum = self.tables[0].get_enum(APP_COLUMN)
-        self.cat_enum = self.tables[0].get_enum(CATEGORY_COLUMN)
         self.data_dtype = self.tables[0].dtype
         if self.preload:
             data = None
@@ -82,40 +136,20 @@ class PyTablesDataset(Dataset):
                 except:
                     pass # ignore if the file is corrupted (or being written at the moment)
             if data is None:
-                data = load_data_from_pytables(tables=self.tables, indices=self.indices, data_dtype=self.data_dtype)
+                data = load_data_from_tables(tables=self.tables, indices=self.indices, data_dtype=self.data_dtype)
             self.data = data
             if self.preload_blob and not os.path.isfile(self.preload_blob):
                 np.savez_compressed(self.preload_blob, data=self.data)
         log.debug(f"Finish initialization worker id {self.worker_id}")
 
-    def get_app_enum(self) -> tb.Enum:
-        if self.app_enum:
-            return self.app_enum
-        database, tables = load_database(database_path=self.database_path, tables_paths=self.tables_paths)
-        app_enum = tables[0].get_enum(APP_COLUMN)
-        cat_enum = tables[0].get_enum(CATEGORY_COLUMN)
-        self.app_enum, self.cat_enum = app_enum, cat_enum
-        database.close()
-        return app_enum
-
-    def get_cat_enum(self) -> tb.Enum:
-        if self.cat_enum:
-            return self.cat_enum
-        database, tables = load_database(database_path=self.database_path, tables_paths=self.tables_paths)
-        app_enum = tables[0].get_enum(APP_COLUMN)
-        cat_enum = tables[0].get_enum(CATEGORY_COLUMN)
-        self.app_enum, self.cat_enum = app_enum, cat_enum
-        database.close()
-        return cat_enum
-
     def set_all_indices(self, disabled_apps: Optional[list[str]] = None):
         """
         This should be called from the main process, before dataloader workers split the work.
         Does no filter apps with not enough samples.
         """
         database, tables = load_database(database_path=self.database_path, tables_paths=self.tables_paths)
-        app_enum = tables[0].get_enum(APP_COLUMN)
-        disabled_apps_ids = list(map(lambda x: app_enum[x], disabled_apps)) if disabled_apps is not None else []
+        inverted_tables_app_enum = {v: k for k, v in self.tables_app_enum.items()}
+        disabled_apps_ids = [inverted_tables_app_enum[app] for app in disabled_apps] if disabled_apps is not None else []
         base_labels = {}
         base_indices = {}
         for i in range(len(tables)):
@@ -137,64 +171,9 @@ def worker_init_fn(worker_id):
     dataset = worker_info.dataset
     dataset.pytables_worker_init(worker_id)
 
-def pytables_collate_fn(batch: tuple,
-                        flowstats_scaler: Scaler, flowstats_quantiles: pd.Series,
-                        psizes_scaler: Scaler, psizes_max: int,
-                        ipt_scaler: Scaler, ipt_min: int, ipt_max: int,
-                        use_push_flags: bool, use_packet_histograms: bool, normalize_packet_histograms: bool, zero_ppi_start: int,
-                        encoder: LabelEncoder, known_apps: list[str], return_torch: bool = False):
-    other_fields, x_ppi, x_flowstats, labels = batch
-    x_ppi = x_ppi.transpose(0, 2, 1)
-    orig_shape = x_ppi.shape
-    ppi_channels = x_ppi.shape[-1]
-    x_ppi = x_ppi.reshape(-1, ppi_channels)
-    x_ppi[:, IPT_POS] = x_ppi[:, IPT_POS].clip(max=ipt_max, min=ipt_min)
-    x_ppi[:, SIZE_POS] = x_ppi[:, SIZE_POS].clip(max=psizes_max, min=1)
-    padding_mask = x_ppi[:, DIR_POS] == 0 # mask of zero padding
-    if ipt_scaler:
-        x_ppi[:, IPT_POS] = ipt_scaler.transform(x_ppi[:, IPT_POS].reshape(-1, 1)).reshape(-1) # type: ignore
-    if psizes_scaler:
-        x_ppi[:, SIZE_POS] = psizes_scaler.transform(x_ppi[:, SIZE_POS].reshape(-1, 1)).reshape(-1) # type: ignore
-    x_ppi[padding_mask, IPT_POS] = 0
-    x_ppi[padding_mask, SIZE_POS] = 0
-    x_ppi = x_ppi.reshape(orig_shape).transpose(0, 2, 1)
-    if not use_push_flags:
-        x_ppi = x_ppi[:, (IPT_POS, DIR_POS, SIZE_POS), :]
-    if zero_ppi_start > 0:
-        x_ppi[:,:,:zero_ppi_start] = 0
-
-    if use_packet_histograms:
-        x_phist = structured_to_unstructured(x_flowstats[PHISTS_FEATURES], dtype="float32")
-        if normalize_packet_histograms:
-            src_sizes_pkt_count = x_phist[:, :PHIST_BIN_COUNT].sum(axis=1)[:, np.newaxis]
-            dst_sizes_pkt_count = x_phist[:, PHIST_BIN_COUNT:(2*PHIST_BIN_COUNT)].sum(axis=1)[:, np.newaxis]
-            np.divide(x_phist[:, :PHIST_BIN_COUNT], src_sizes_pkt_count, out=x_phist[:, :PHIST_BIN_COUNT], where=src_sizes_pkt_count != 0)
-            np.divide(x_phist[:, PHIST_BIN_COUNT:(2*PHIST_BIN_COUNT)], dst_sizes_pkt_count, out=x_phist[:, PHIST_BIN_COUNT:(2*PHIST_BIN_COUNT)], where=dst_sizes_pkt_count != 0)
-            np.divide(x_phist[:, (2*PHIST_BIN_COUNT):(3*PHIST_BIN_COUNT)], src_sizes_pkt_count - 1, out=x_phist[:, (2*PHIST_BIN_COUNT):(3*PHIST_BIN_COUNT)], where=src_sizes_pkt_count > 1)
-            np.divide(x_phist[:, (3*PHIST_BIN_COUNT):(4*PHIST_BIN_COUNT)], dst_sizes_pkt_count - 1, out=x_phist[:, (3*PHIST_BIN_COUNT):(4*PHIST_BIN_COUNT)], where=dst_sizes_pkt_count > 1)
-        x_flowstats = structured_to_unstructured(drop_fields(x_flowstats, PHISTS_FEATURES), dtype="float32")
-        x_flowstats = np.concatenate([x_flowstats, x_phist], axis=1)
-    else:
-        x_flowstats = structured_to_unstructured(x_flowstats, dtype="float32")
-    np.clip(x_flowstats[:, :len(FLOWSTATS_TO_SCALE)], a_max=flowstats_quantiles, a_min=0, out=x_flowstats[:, :len(FLOWSTATS_TO_SCALE)])
-    if flowstats_scaler:
-        x_flowstats[:, :len(FLOWSTATS_TO_SCALE)] = flowstats_scaler.transform(x_flowstats[:, :len(FLOWSTATS_TO_SCALE)])
-
-    other_fields_df = pd.DataFrame(other_fields) if len(other_fields) > 0 else pd.DataFrame()
-    for column in other_fields_df.columns:
-        if other_fields_df[column].dtype.kind == "O":
-            other_fields_df[column] = other_fields_df[column].astype(str)
-        elif column.startswith("TIME_"):
-            other_fields_df[column] = other_fields_df[column].map(lambda x: datetime.fromtimestamp(x))
-
-    labels = encoder.transform(np.where(np.isin(labels, known_apps), labels, UNKNOWN_STR_LABEL)).astype("int64") # type: ignore
-    if return_torch:
-        return other_fields_df, torch.from_numpy(x_ppi), torch.from_numpy(x_flowstats), torch.from_numpy(labels)
-    return other_fields_df, x_ppi, x_flowstats, labels
-
-def init_train_indices(train_data_params: TrainDataParams, servicemap: pd.DataFrame, database_path: str, rng: np.random.RandomState) -> tuple[np.ndarray, np.ndarray, dict[int, str], dict[int, str]]:
+def init_train_indices(train_data_params: TrainDataParams, database_path: str, tables_app_enum: dict[int, str], servicemap: pd.DataFrame, rng: np.random.RandomState) -> tuple[np.ndarray, np.ndarray, list[str], list[str]]:
     database, train_tables = load_database(database_path, tables_paths=train_data_params.train_tables_paths)
-    app_enum = train_tables[0].get_enum(APP_COLUMN)
+    inverted_tables_app_enum = {v: k for k, v in tables_app_enum.items()}
     all_app_labels = {}
     app_counts = pd.Series(dtype="int64")
     start_time = time.time()
@@ -206,15 +185,16 @@ def init_train_indices(train_data_params: TrainDataParams, servicemap: pd.DataFr
     # Handle disabled apps and apps with less than min_samples_per_app samples
     if len(train_data_params.disabled_apps) > 0:
         log.info(f"Disabled applications in dataset config: {sorted(train_data_params.disabled_apps)}")
-    disabled_apps_ids = list(map(lambda x: app_enum[x], train_data_params.disabled_apps))
+    disabled_apps_ids = [inverted_tables_app_enum[app] for app in train_data_params.disabled_apps]
     min_samples_apps_ids = set(app_counts[app_counts<train_data_params.min_train_samples_per_app].index.tolist())
     if len(min_samples_apps_ids) > 0:
+        min_samples_apps_names = sorted([tables_app_enum[app_id] for app_id in min_samples_apps_ids])
         if train_data_params.min_train_samples_check == MinTrainSamplesCheck.WARN_AND_EXIT:
-            warnings.warn(f"Found applications with less than {train_data_params.min_train_samples_per_app} train samples: {sorted(map(app_enum, min_samples_apps_ids))}. " +
+            warnings.warn(f"Found applications with less than {train_data_params.min_train_samples_per_app} train samples: {min_samples_apps_names}. " +
                             "To disable these applications, add them to config.disabled_apps or set config.min_train_samples_check to disable-apps. To turn off this check, set config.min_train_samples_per_app to zero. Exiting")
             exit()
         elif train_data_params.min_train_samples_check == MinTrainSamplesCheck.DISABLE_APPS:
-            log.info(f"Found applications with less than {train_data_params.min_train_samples_per_app} train samples: {sorted(map(app_enum, min_samples_apps_ids))}. " +
+            log.info(f"Found applications with less than {train_data_params.min_train_samples_per_app} train samples: {min_samples_apps_names}. " +
                        "Disabling these applications")
             disabled_apps_ids.extend(min_samples_apps_ids)
     # Base indices are indices of samples that are not disabled and have enough samples
@@ -223,9 +203,9 @@ def init_train_indices(train_data_params: TrainDataParams, servicemap: pd.DataFr
         base_indices[i] = np.nonzero(np.isin(all_app_labels[i], disabled_apps_ids, invert=True))[0]
     base_labels = {table_id: arr[base_indices[table_id]] for table_id, arr in all_app_labels.items()}
     # Apps selection
-    if train_data_params.apps_selection != AppSelection.LONGTERM_FIXED:
+    if train_data_params.apps_selection != AppSelection.FIXED:
         app_counts = app_counts[[app for app in app_counts.index.tolist() if app not in disabled_apps_ids]]
-        app_counts.index = app_counts.index.map(app_enum)
+        app_counts.index = app_counts.index.map(tables_app_enum)
         app_counts = app_counts.sort_values(ascending=False).astype("int64")
         sorted_apps = app_counts.index.to_list()
         if train_data_params.apps_selection == AppSelection.ALL_KNOWN:
@@ -235,31 +215,26 @@ def init_train_indices(train_data_params: TrainDataParams, servicemap: pd.DataFr
             known_apps, unknown_apps = split_apps_topx_with_provider_groups(sorted_apps=sorted_apps, known_count=train_data_params.apps_selection_topx, servicemap=servicemap)
             if len(known_apps) < train_data_params.apps_selection_topx:
                 warnings.warn(f"The number of known applications ({len(known_apps)}) is lower than requested in config.apps_selection_topx ({train_data_params.apps_selection_topx}).")
-        elif train_data_params.apps_selection == AppSelection.EXPLICIT_UNKNOWN:
-                unknown_apps = train_data_params.apps_selection_explicit_unknown
-                missing_unknown_apps = [app for app in unknown_apps if app not in sorted_apps]
-                if len(missing_unknown_apps) > 0:
-                    raise ValueError(f"Applications configured in config.apps_selection_explicit_unknown are not present in the dataset (or might be disabled): {sorted(missing_unknown_apps)}")
+        elif train_data_params.apps_selection == AppSelection.BACKGROUND_UNKNOWN:
+                unknown_apps = train_data_params.apps_selection_background_unknown
                 known_apps = [app for app in sorted_apps if not (is_background_app(app) or app in unknown_apps)]
         else: assert_never(train_data_params.apps_selection)
-
         log.info(f"Selected {len(known_apps)} known applications and {len(unknown_apps)} unknown applications")
-        known_apps_database_enum: dict[int, str] = {int(app_enum[app]): app for app in known_apps}
-        unknown_apps_database_enum: dict[int, str] = {int(app_enum[app]): app for app in unknown_apps}
     else:
-        assert train_data_params.apps_selection_fixed_longterm is not None
-        known_apps_database_enum, unknown_apps_database_enum = train_data_params.apps_selection_fixed_longterm
-    known_apps_ids = list(known_apps_database_enum)
-    unknown_apps_ids = list(unknown_apps_database_enum)
+        known_apps = train_data_params.apps_selection_fixed_known
+        unknown_apps = train_data_params.apps_selection_fixed_unknown
+    known_apps_ids = [inverted_tables_app_enum[app] for app in known_apps]
+    unknown_apps_ids = [inverted_tables_app_enum[app] for app in unknown_apps]
 
     train_known_indices, train_unknown_indices = convert_dict_indices(base_indices=base_indices, base_labels=base_labels, known_apps_ids=known_apps_ids, unknown_apps_ids=unknown_apps_ids)
     rng.shuffle(train_known_indices)
     rng.shuffle(train_unknown_indices)
     log.info(f"Processing train indices took {time.time() - start_time:.2f} seconds"); start_time = time.time()
-    return train_known_indices, train_unknown_indices, known_apps_database_enum, unknown_apps_database_enum
+    return train_known_indices, train_unknown_indices, known_apps, unknown_apps
 
-def init_test_indices(test_data_params: TestDataParams, database_path: str, rng: np.random.RandomState) -> tuple[np.ndarray, np.ndarray]:
+def init_test_indices(test_data_params: TestDataParams, database_path: str, tables_app_enum: dict[int, str], rng: np.random.RandomState) -> tuple[np.ndarray, np.ndarray]:
     database, test_tables = load_database(database_path, tables_paths=test_data_params.test_tables_paths)
+    inverted_tables_app_enum = {v: k for k, v in tables_app_enum.items()}
     base_labels = {}
     base_indices = {}
     start_time = time.time()
@@ -268,115 +243,14 @@ def init_test_indices(test_data_params: TestDataParams, database_path: str, rng:
         log.info(f"Reading app column for test table {table_path} took {time.time() - start_time:.2f} seconds"); start_time = time.time()
         base_indices[i] = np.arange(len(test_tables[i]))
     database.close()
-    known_apps_ids = list(test_data_params.known_apps_database_enum)
-    unknown_apps_ids = list(test_data_params.unknown_apps_database_enum)
+    known_apps_ids = [inverted_tables_app_enum[app] for app in test_data_params.known_apps]
+    unknown_apps_ids = [inverted_tables_app_enum[app] for app in test_data_params.unknown_apps]
     test_known_indices, test_unknown_indices = convert_dict_indices(base_indices=base_indices, base_labels=base_labels, known_apps_ids=known_apps_ids, unknown_apps_ids=unknown_apps_ids)
     rng.shuffle(test_known_indices)
     rng.shuffle(test_unknown_indices)
     log.info(f"Processing test indices took {time.time() - start_time:.2f} seconds"); start_time = time.time()
     return test_known_indices, test_unknown_indices
 
-def fit_or_load_scalers(dataset_config: DatasetConfig, train_indices: np.ndarray) -> tuple[Scaler, pd.Series, Scaler, Scaler]:
-    train_data_path = dataset_config._get_train_data_path()
-    flowstats_scaler_path = os.path.join(train_data_path, "stand", f"flowstats_scaler-{dataset_config.flowstats_scaler}-q{dataset_config.flowstats_clip}.pickle")
-    flowstats_quantiles_path = os.path.join(train_data_path, "stand", f"flowstats_quantiles-q{dataset_config.flowstats_clip}.pickle")
-    ipt_scaler_path = os.path.join(train_data_path, "stand", f"ipt_scaler-{dataset_config.ipt_scaler}-ipt_min{dataset_config.ipt_min}-ipt_max{dataset_config.ipt_max}.pickle")
-    psizes_sizes_scaler_path = os.path.join(train_data_path, "stand", f"psizes_scaler-{dataset_config.psizes_scaler}-psizes_max{dataset_config.psizes_max}.pickle")
-    if os.path.isfile(flowstats_scaler_path) and os.path.isfile(flowstats_quantiles_path) and os.path.isfile(ipt_scaler_path) and os.path.isfile(psizes_sizes_scaler_path):
-        flowstats_scaler = pickle_load(flowstats_scaler_path)
-        flowstats_quantiles = pickle_load(flowstats_quantiles_path)
-        ipt_scaler = pickle_load(ipt_scaler_path)
-        psizes_scaler = pickle_load(psizes_sizes_scaler_path)
-    else:
-        if dataset_config.flowstats_scaler == ScalerEnum.ROBUST:
-            flowstats_scaler = RobustScaler()
-        elif dataset_config.flowstats_scaler == ScalerEnum.STANDARD:
-            flowstats_scaler = StandardScaler()
-        elif dataset_config.flowstats_scaler == ScalerEnum.MINMAX:
-            flowstats_scaler = MinMaxScaler()
-        elif dataset_config.flowstats_scaler == ScalerEnum.NO_SCALER:
-            flowstats_scaler = None
-        else: assert_never(dataset_config.flowstats_scaler)
-
-        if dataset_config.ipt_scaler == ScalerEnum.ROBUST:
-            ipt_scaler = RobustScaler()
-        elif dataset_config.ipt_scaler == ScalerEnum.STANDARD:
-            ipt_scaler = StandardScaler()
-        elif dataset_config.ipt_scaler == ScalerEnum.MINMAX:
-            ipt_scaler = MinMaxScaler()
-        elif dataset_config.ipt_scaler == ScalerEnum.NO_SCALER:
-            ipt_scaler = None
-        else: assert_never(dataset_config.ipt_scaler)
-
-        if dataset_config.psizes_scaler == ScalerEnum.ROBUST:
-            psizes_scaler = RobustScaler()
-        elif dataset_config.psizes_scaler == ScalerEnum.STANDARD:
-            psizes_scaler = StandardScaler()
-        elif dataset_config.psizes_scaler == ScalerEnum.MINMAX:
-            psizes_scaler = MinMaxScaler()
-        elif dataset_config.psizes_scaler == ScalerEnum.NO_SCALER:
-            psizes_scaler = None
-        else: assert_never(dataset_config.psizes_scaler)
-
-        if isinstance(dataset_config.fit_scalers_samples, int) and dataset_config.fit_scalers_samples > len(train_indices):
-            warnings.warn(f"The number of samples for fitting scalers ({dataset_config.fit_scalers_samples}) is larger than the number of train samples ({len(train_indices)}), using the number of train samples instead")
-            dataset_config.fit_scalers_samples = len(train_indices)
-        fit_scalers_rng = get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.FIT_SCALERS_SAMPLE)
-        if isinstance(dataset_config.fit_scalers_samples, float):
-            num_samples = int(dataset_config.fit_scalers_samples * len(train_indices))
-        else:
-            num_samples = dataset_config.fit_scalers_samples
-        fit_scalers_indices = train_indices[fit_scalers_rng.choice(len(train_indices), size=num_samples, replace=False)]
-        flowstats_quantiles = fit_scalers(
-            database_path=dataset_config.database_path,
-            train_tables_paths=dataset_config._get_train_tables_paths(),
-            fit_scalers_indices=fit_scalers_indices,
-            flowstats_scaler=flowstats_scaler,
-            flowstats_quantile_clip=dataset_config.flowstats_clip,
-            ipt_scaler=ipt_scaler,
-            psizes_scaler=psizes_scaler,
-            ipt_min=dataset_config.ipt_min,
-            ipt_max=dataset_config.ipt_max,
-            psizes_max=dataset_config.psizes_max)
-        pickle_dump(flowstats_scaler, flowstats_scaler_path)
-        pickle_dump(flowstats_quantiles, flowstats_quantiles_path)
-        pickle_dump(ipt_scaler, ipt_scaler_path)
-        pickle_dump(psizes_scaler, psizes_sizes_scaler_path)
-    return flowstats_scaler, flowstats_quantiles, ipt_scaler, psizes_scaler
-
-def fit_scalers(database_path: str, train_tables_paths: list[str], fit_scalers_indices: np.ndarray, flowstats_scaler, flowstats_quantile_clip: float, ipt_scaler, psizes_scaler, ipt_min: int, ipt_max: int, psizes_max: int) -> pd.Series:
-    start_time = time.time()
-    database, tables = load_database(database_path, tables_paths=train_tables_paths)
-    data = load_data_from_pytables(tables=tables, indices=fit_scalers_indices, data_dtype=tables[0].dtype)
-    database.close()
-    # PPI
-    data_ppi = data[PPI_COLUMN].astype("float32")
-    ppi_channels = data_ppi.shape[1]
-    data_ppi = data_ppi.transpose(0, 2, 1).reshape(-1, ppi_channels)
-    padding_mask = data_ppi[:, DIR_POS] == 0 # mask of padded packets
-    if ipt_scaler:
-        train_ipt = data_ppi[:, IPT_POS].clip(max=ipt_max, min=ipt_min)
-        train_ipt[padding_mask] = np.nan # nans are ignored in sklearn scalers
-        if isinstance(ipt_scaler, MinMaxScaler):
-            # let zero be the minimum for minmax scaling
-            train_ipt = np.concatenate((train_ipt, [0]))
-        ipt_scaler.fit(train_ipt.reshape(-1, 1))
-    if psizes_scaler:
-        train_psizes = data_ppi[:, SIZE_POS].clip(max=psizes_max, min=1)
-        train_psizes[padding_mask] = np.nan
-        if isinstance(psizes_scaler, MinMaxScaler):
-            train_psizes = np.concatenate((train_psizes, [0]))
-        psizes_scaler.fit(train_psizes.reshape(-1, 1))
-    # FLOWSTATS
-    train_flowstats = pd.DataFrame(data, columns=FLOWSTATS_TO_SCALE)
-    upper_quantiles = train_flowstats.quantile(flowstats_quantile_clip)
-    upper_quantiles[FLOWSTATS_NO_CLIP] = np.Inf # disable clipping for features with "fixed" range
-    if flowstats_scaler:
-        train_flowstats = train_flowstats.clip(upper=upper_quantiles, lower=0, axis=1).to_numpy() # type: ignore
-        flowstats_scaler.fit(train_flowstats)
-    log.info(f"Reading data and fitting scalers took {time.time() - start_time:.2f} seconds")
-    return upper_quantiles
-
 def load_database(database_path: str, tables_paths: Optional[list[str]] = None, mode: str = "r") -> tuple[tb.File, dict[int, Any]]: # dict[int, tb.Table]
     database = tb.open_file(database_path, mode=mode)
     if tables_paths is None:
@@ -410,11 +284,11 @@ def convert_dict_indices(base_indices: dict[int, np.ndarray], base_labels: dict[
         np.concatenate(list(unknown_labels_dict.values()))))
     return known_indices, unknown_indices
 
-def load_data_from_pytables(tables, indices: np.ndarray, data_dtype: np.dtype) -> np.ndarray:
+def load_data_from_tables(tables, indices: np.ndarray, data_dtype: np.dtype) -> np.ndarray:
     sorted_indices = indices[indices[:, INDICES_TABLE_POS].argsort(kind="stable")]
     unique_tables, split_bounderies = np.unique(sorted_indices[:, INDICES_TABLE_POS], return_index=True)
     indices_per_table = np.split(sorted_indices, split_bounderies[1:])
-    data = np.empty(len(indices), dtype=data_dtype)
+    data = np.zeros(len(indices), dtype=data_dtype)
     for table_id, table_indices in zip(unique_tables, indices_per_table):
         data[np.where(indices[:, INDICES_TABLE_POS] == table_id)[0]] = tables[table_id].read_coordinates(table_indices[:, INDICES_INDEX_POS])
     return data

cesnet_datazoo/utils/class_info.py

@@ -10,10 +10,11 @@ from cesnet_datazoo.constants import SERVICEMAP_CATEGORY_COLUMN, SERVICEMAP_PROV
 
 @dataclass()
 class ClassInfo:
-    target_names: list[str]
     num_classes: int
     known_apps: list[str]
     unknown_apps: list[str]
+    encoder: LabelEncoder
+    target_names: list[str]
     unknown_class_label: int
     group_matrix: np.ndarray
     has_provider: dict[str, bool]
@@ -21,9 +22,9 @@ class ClassInfo:
     provider_members: dict[str, list[str]]
     categories_mapping: dict[str, Optional[str]]
 
-def create_class_info(servicemap: Any, encoder: LabelEncoder, known_apps_database_enum: dict[int, str], unknown_apps_database_enum: dict[int, str]) -> ClassInfo:
-    known_apps = sorted(known_apps_database_enum.values())
-    unknown_apps = sorted(unknown_apps_database_enum.values())
+def create_class_info(servicemap: Any, encoder: LabelEncoder, known_apps: list[str], unknown_apps: list[str]) -> ClassInfo:
+    known_apps = sorted(known_apps)
+    unknown_apps = sorted(unknown_apps)
     target_names_arr = encoder.classes_
     assert known_apps == list(target_names_arr[:-1])
     group_matrix = np.array([[a == b or
@@ -37,10 +38,11 @@ def create_class_info(servicemap: Any, encoder: LabelEncoder, known_apps_databas
     provider_members = {p: [app for app in target_names_arr if provider_mapping[app] == p] for p in providers}
     categories_mapping = {app: servicemap.loc[app, SERVICEMAP_CATEGORY_COLUMN] if app in servicemap.index else None for app in target_names_arr}
     return ClassInfo(
-            target_names=list(target_names_arr),
             num_classes=len(known_apps),
             known_apps=known_apps,
             unknown_apps=unknown_apps,
+            encoder=encoder,
+            target_names=list(target_names_arr),
             unknown_class_label=len(known_apps),
             group_matrix=group_matrix,
             has_provider=has_provider,

cesnet_datazoo/utils/download.py

@@ -11,9 +11,14 @@ def simple_download(url: str, file_path: str):
 
 def resumable_download(url: str, file_path: str, chunk_size: int = 1024**2, silent: bool = False):
     r1 = requests.get(url, stream=True)
+    try:
+        r1.raise_for_status()
+    except requests.exceptions.HTTPError as e:
+        print("The dataset hosting server is unreachable. Please contact us at https://github.com/CESNET/cesnet-datazoo/issues.")
+        raise e
+
     redirected_url = r1.url
     content_size = int(r1.headers["Content-Length"])
-
     if os.path.exists(file_path):
         temp_size = os.path.getsize(file_path)
     else:

{cesnet_datazoo-0.0.16.dist-info → cesnet_datazoo-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cesnet-datazoo
-Version: 0.0.16
+Version: 0.1.0
 Summary: A toolkit for large network traffic datasets
 Author-email: Jan Luxemburk <luxemburk@cesnet.cz>, Karel Hynek <hynekkar@cesnet.cz>
 Maintainer-email: Jan Luxemburk <luxemburk@cesnet.cz>, Karel Hynek <hynekkar@cesnet.cz>
@@ -16,6 +16,7 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENCE
+Requires-Dist: cesnet-models
 Requires-Dist: matplotlib
 Requires-Dist: numpy
 Requires-Dist: pandas

cesnet_datazoo-0.1.0.dist-info/RECORD

@@ -0,0 +1,30 @@
+cesnet_datazoo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cesnet_datazoo/config.py,sha256=vvNyM7TCMolH-uLj3ant7rGkYb_2FPyCWlRQ3mllKWs,37427
+cesnet_datazoo/constants.py,sha256=EDeeo0xrBt_pnWf3m-ZTiC5HMvyVwcikgCZ9LwZIcAE,1276
+cesnet_datazoo/datasets/__init__.py,sha256=8ziQ3EUzUh5fMfWWXwk0cqYk0lOUNU7zbi0Gom3bLnI,443
+cesnet_datazoo/datasets/cesnet_dataset.py,sha256=zoLFduBg6ZK96zoec0kEMB1hFCGn3QOtBtYFTcCbIU0,46546
+cesnet_datazoo/datasets/datasets.py,sha256=Bn4SU1k5og6AsUlnPapFPeu4uGlpRH-IaOSafz0ZT2k,3617
+cesnet_datazoo/datasets/datasets_constants.py,sha256=1P54Ns8wCQMemdKNe8OH7cVUfkxs3vL29ugSmOLXceI,29154
+cesnet_datazoo/datasets/loaders.py,sha256=9KgRY-Y8CcgtXbgqWpAaG7gyOAsSf278w7b1eHwTSyE,1854
+cesnet_datazoo/datasets/statistics.py,sha256=wR8QISIh-KC7CQ5SjN7WoTMFaoRuq0G7pgTFGhC8ek0,15137
+cesnet_datazoo/datasets/metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cesnet_datazoo/datasets/metadata/dataset_metadata.py,sha256=Ntlp8mHUSr7g-ZTvtBVh238TswZHwGAudMuE52-OA-c,1608
+cesnet_datazoo/datasets/metadata/metadata.csv,sha256=or0CB7t06G_V1OzClqtpx7sRt_ZoQWE_f7F5SDLlPC8,2175
+cesnet_datazoo/metrics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cesnet_datazoo/metrics/classification_report.py,sha256=0JgKWyB281m3EHxI8miMKTjKg3mzyV1WTQndXg_B7i0,4040
+cesnet_datazoo/metrics/provider_metrics.py,sha256=sRg2bdRTzLLTmiVjacBtGez4LEIfr35hSvMBwW-W73U,1303
+cesnet_datazoo/pytables_data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cesnet_datazoo/pytables_data/apps_split.py,sha256=RjLFomrlBCmnBn08FDw1IzL3PuQf4914yJQzwhiXH_E,1411
+cesnet_datazoo/pytables_data/data_scalers.py,sha256=IfTymhVubjLNetjOIxDhtzkETp_1xmFXbC0rSjQHVUQ,7254
+cesnet_datazoo/pytables_data/indices_setup.py,sha256=rBW1HwebPXkwLRuKg9ILO_LfUrfnJfqQYsrIAYfXtZo,12932
+cesnet_datazoo/pytables_data/pytables_dataset.py,sha256=kCvbOgIseBdUUGz3nRr7oGsuN2JicXGlsp8-Z9n4JyM,17599
+cesnet_datazoo/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cesnet_datazoo/utils/class_info.py,sha256=zMt2ndfwvtnE5QOKS1OPbw8VUlsCCbB_SVjzyFn1Wdw,2540
+cesnet_datazoo/utils/download.py,sha256=hG5V1ZYZGtqCzlVV76NMgOZkSKOywdOFiq9Lagkgego,1441
+cesnet_datazoo/utils/fileutils.py,sha256=XA_VWDuTiCXnoOgHPUzsmbnLFgrlxOo5cvUY_OBJUR8,642
+cesnet_datazoo/utils/random.py,sha256=Dqgm_T25ljbew-OJozK90PsiXKnd4Kw6lcUexxF6vIc,575
+cesnet_datazoo-0.1.0.dist-info/LICENCE,sha256=69Wc69APiM1YKrFOIipG7jjU2lk89WQuO_U0AXKU8KE,1541
+cesnet_datazoo-0.1.0.dist-info/METADATA,sha256=c6GMIPE5rkiZtsbGNv28405o-G02J4wyvP-DJL8BfJM,12679
+cesnet_datazoo-0.1.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+cesnet_datazoo-0.1.0.dist-info/top_level.txt,sha256=bu1Z8zaI_1Id_ZaYyvJnxIBa87OSrdlZ8J2OBMggK5o,15
+cesnet_datazoo-0.1.0.dist-info/RECORD,,

{cesnet_datazoo-0.0.16.dist-info → cesnet_datazoo-0.1.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.42.0)
+Generator: bdist_wheel (0.43.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

cesnet_datazoo-0.0.16.dist-info/RECORD

@@ -1,28 +0,0 @@
-cesnet_datazoo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cesnet_datazoo/config.py,sha256=LawxdianPcNcuzxf01FTyED7PoAvQj8RMSE8QG4FZNo,37531
-cesnet_datazoo/constants.py,sha256=EliK-KvW3GXeKw00W_Pd-ypJMwvFQVqMQS9A9ULyTj4,1420
-cesnet_datazoo/datasets/__init__.py,sha256=8ziQ3EUzUh5fMfWWXwk0cqYk0lOUNU7zbi0Gom3bLnI,443
-cesnet_datazoo/datasets/cesnet_dataset.py,sha256=S0FsatG3fq21fVJctWOaLbF1ZzSvxUnzg9Hbe3TFNmo,43116
-cesnet_datazoo/datasets/datasets.py,sha256=gj7jflxqDgEfHXSFUz6JOW2x8wEUSCqVe6KapaK4IKg,2279
-cesnet_datazoo/datasets/loaders.py,sha256=HU2Au0P87BCAvdgpiwO5T0xgeQgs_gL4E1d12OP1JoQ,1803
-cesnet_datazoo/datasets/statistics.py,sha256=GoM7-vFTvqx9ym239VCZd1os2TdoxLOW7WNpNtOU7Fc,14030
-cesnet_datazoo/datasets/metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cesnet_datazoo/datasets/metadata/dataset_metadata.py,sha256=OZf-NMai2XuSg57y2IdV-804ZpPcmI9sWoDu8IO7e4Y,1567
-cesnet_datazoo/datasets/metadata/metadata.csv,sha256=Zr2hf9qpJpPE_Js9XmyaHffdho912ikdQfFVQx6q8NE,2161
-cesnet_datazoo/metrics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cesnet_datazoo/metrics/classification_report.py,sha256=0JgKWyB281m3EHxI8miMKTjKg3mzyV1WTQndXg_B7i0,4040
-cesnet_datazoo/metrics/provider_metrics.py,sha256=sRg2bdRTzLLTmiVjacBtGez4LEIfr35hSvMBwW-W73U,1303
-cesnet_datazoo/pytables_data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cesnet_datazoo/pytables_data/apps_split.py,sha256=RjLFomrlBCmnBn08FDw1IzL3PuQf4914yJQzwhiXH_E,1411
-cesnet_datazoo/pytables_data/indices_setup.py,sha256=IraCOFys0p7ZojR-0E99bKN9dTjwCfQO4L6lMqcTEFg,13070
-cesnet_datazoo/pytables_data/pytables_dataset.py,sha256=y2BXuuA73w58XITKFsPm-FS8LB76TH6prNUMsKkXNBM,26511
-cesnet_datazoo/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cesnet_datazoo/utils/class_info.py,sha256=ync9U3PWo0DloRwX3uMgKW798kC6echioEbEgrPqY4E,2567
-cesnet_datazoo/utils/download.py,sha256=QVbYKuWUO9j6VUJISPTVBXscjuTuuX-XRez7MJzG3dk,1204
-cesnet_datazoo/utils/fileutils.py,sha256=XA_VWDuTiCXnoOgHPUzsmbnLFgrlxOo5cvUY_OBJUR8,642
-cesnet_datazoo/utils/random.py,sha256=Dqgm_T25ljbew-OJozK90PsiXKnd4Kw6lcUexxF6vIc,575
-cesnet_datazoo-0.0.16.dist-info/LICENCE,sha256=69Wc69APiM1YKrFOIipG7jjU2lk89WQuO_U0AXKU8KE,1541
-cesnet_datazoo-0.0.16.dist-info/METADATA,sha256=YQBQeLwNIoHcCR3W4Dn46iWLlRMWC4c8B7U8r81gRd8,12650
-cesnet_datazoo-0.0.16.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-cesnet_datazoo-0.0.16.dist-info/top_level.txt,sha256=bu1Z8zaI_1Id_ZaYyvJnxIBa87OSrdlZ8J2OBMggK5o,15
-cesnet_datazoo-0.0.16.dist-info/RECORD,,