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

cesnet_datazoo/datasets/loaders.py

@@ -6,6 +6,9 @@ from tqdm import tqdm
 from cesnet_datazoo.constants import APP_COLUMN
 
 
+def collate_fn_simple(batch):
+    return batch
+
 def load_from_dataloader(dataloader: DataLoader, silent: bool = False) -> tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]:
     other_fields = []
     data_ppi = []

cesnet_datazoo/datasets/metadata/dataset_metadata.py

@@ -18,18 +18,19 @@ class DatasetMetadata():
     collection_period: str
     missing_dates_in_collection_period: list[str]
     application_count: int
-    background_traffic: list[str]
-    features_in_packet_sequences: list[str]
-    packet_histogram_features: list[str]
+    background_traffic_classes: list[str]
+    ppi_features: list[str]
     flowstats_features: list[str]
+    flowstats_features_boolean: list[str]
+    packet_histograms: list[str]
     tcp_features: list[str]
     other_fields: list[str]
     cite: str
     zenodo_url: str
     related_papers: list[str]
 
-    @field_validator("available_dataset_sizes", "missing_dates_in_collection_period", "background_traffic", "features_in_packet_sequences",
-                     "packet_histogram_features", "flowstats_features", "tcp_features", "other_fields", "related_papers", mode="before")
+    @field_validator("available_dataset_sizes", "missing_dates_in_collection_period", "background_traffic_classes", "ppi_features",
+                     "flowstats_features", "flowstats_features_boolean", "packet_histograms", "tcp_features", "other_fields", "related_papers", mode="before")
     @classmethod
     def parse_string_to_list(cls, v: str, info: ValidationInfo) -> list[str]:
         l = list(map(str.strip, v.split(","))) if v else []

cesnet_datazoo/datasets/metadata/metadata.csv

@@ -1,4 +1,4 @@
-Name,Protocol,Published in,Collected in,Collection duration,Available samples,Available dataset sizes,Collection period,Missing dates in collection period,Application count,Background traffic,Features in packet sequences,Packet histogram features,Flowstats features,TCP features,Other fields,Cite,Zenodo URL,Related papers
-CESNET-TLS22,TLS,2022,2021,2 weeks,141720670,"XS, S, M, L",4.10.2021 - 17.10.2021,,191,,"IPT, DIR, SIZE",,"BYTES, BYTES_REV, PACKETS, PACKETS_REV, PPI_LEN, PPI_ROUNDTRIPS, PPI_DURATION, DURATION","FLAG_CWR, FLAG_CWR_REV, FLAG_ECE, FLAG_ECE_REV, FLAG_URG, FLAG_URG_REV, FLAG_ACK, FLAG_ACK_REV, FLAG_PSH, FLAG_PSH_REV, FLAG_RST, FLAG_RST_REV, FLAG_SYN, FLAG_SYN_REV, FLAG_FIN, FLAG_FIN_REV",ID,https://doi.org/10.1016/j.comnet.2022.109467,https://zenodo.org/record/7965515,
-CESNET-QUIC22,QUIC,2023,2022,4 weeks,153226273,"XS, S, M, L",31.10.2022 - 27.11.2022,,102,"default-background, google-background, facebook-background","IPT, DIR, SIZE","PHIST_SRC_SIZES, PHIST_DST_SIZES, PHIST_SRC_IPT, PHIST_DST_IPT","BYTES, BYTES_REV, PACKETS, PACKETS_REV, PPI_LEN, PPI_ROUNDTRIPS, PPI_DURATION, DURATION, FLOW_ENDREASON_IDLE, FLOW_ENDREASON_ACTIVE, FLOW_ENDREASON_OTHER",,"ID, SRC_IP, DST_IP, DST_ASN, SRC_PORT, DST_PORT, PROTOCOL, QUIC_VERSION, QUIC_SNI, QUIC_USERAGENT, TIME_FIRST, TIME_LAST",https://doi.org/10.1016/j.dib.2023.108888,https://zenodo.org/record/7963302,https://doi.org/10.23919/TMA58422.2023.10199052
-CESNET-TLS-Year22,TLS,2023,2022,1 year,507739073,"XS, S, M, L",1.1.2022 - 31.12.2022,"20220128, 20220129, 20220130, 20221212, 20221213, 20221229, 20221230, 20221231",180,,"IPT, DIR, SIZE, PUSH_FLAG","PHIST_SRC_SIZES, PHIST_DST_SIZES, PHIST_SRC_IPT, PHIST_DST_IPT","BYTES, BYTES_REV, PACKETS, PACKETS_REV, PPI_LEN, PPI_ROUNDTRIPS, PPI_DURATION, DURATION, FLOW_ENDREASON_IDLE, FLOW_ENDREASON_ACTIVE, FLOW_ENDREASON_END, FLOW_ENDREASON_OTHER","FLAG_CWR, FLAG_CWR_REV, FLAG_ECE, FLAG_ECE_REV, FLAG_URG, FLAG_URG_REV, FLAG_ACK, FLAG_ACK_REV, FLAG_PSH, FLAG_PSH_REV, FLAG_RST, FLAG_RST_REV, FLAG_SYN, FLAG_SYN_REV, FLAG_FIN, FLAG_FIN_REV","ID, SRC_IP, DST_IP, DST_ASN, DST_PORT, PROTOCOL, TLS_SNI, TLS_JA3, TIME_FIRST, TIME_LAST",,,
+Name,Protocol,Published in,Collected in,Collection duration,Available samples,Available dataset sizes,Collection period,Missing dates in collection period,Application count,Background traffic classes,PPI features,Flowstats features,Flowstats features boolean,Packet histograms,TCP features,Other fields,Cite,Zenodo URL,Related papers
+CESNET-TLS22,TLS,2022,2021,2 weeks,141720670,"XS, S, M, L",4.10.2021 - 17.10.2021,,191,,"IPT, DIR, SIZE","BYTES, BYTES_REV, PACKETS, PACKETS_REV, DURATION, PPI_LEN, PPI_ROUNDTRIPS, PPI_DURATION",,,"FLAG_CWR, FLAG_CWR_REV, FLAG_ECE, FLAG_ECE_REV, FLAG_URG, FLAG_URG_REV, FLAG_ACK, FLAG_ACK_REV, FLAG_PSH, FLAG_PSH_REV, FLAG_RST, FLAG_RST_REV, FLAG_SYN, FLAG_SYN_REV, FLAG_FIN, FLAG_FIN_REV",ID,https://doi.org/10.1016/j.comnet.2022.109467,https://zenodo.org/record/7965515,
+CESNET-QUIC22,QUIC,2023,2022,4 weeks,153226273,"XS, S, M, L",31.10.2022 - 27.11.2022,,102,"default-background, google-background, facebook-background","IPT, DIR, SIZE","BYTES, BYTES_REV, PACKETS, PACKETS_REV, DURATION, PPI_LEN, PPI_ROUNDTRIPS, PPI_DURATION","FLOW_ENDREASON_IDLE, FLOW_ENDREASON_ACTIVE, FLOW_ENDREASON_OTHER","PHIST_SRC_SIZES, PHIST_DST_SIZES, PHIST_SRC_IPT, PHIST_DST_IPT",,"ID, SRC_IP, DST_IP, DST_ASN, SRC_PORT, DST_PORT, PROTOCOL, QUIC_VERSION, QUIC_SNI, QUIC_USERAGENT, TIME_FIRST, TIME_LAST",https://doi.org/10.1016/j.dib.2023.108888,https://zenodo.org/record/7963302,https://doi.org/10.23919/TMA58422.2023.10199052
+CESNET-TLS-Year22,TLS,2023,2022,1 year,507739073,"XS, S, M, L",1.1.2022 - 31.12.2022,"20220128, 20220129, 20220130, 20221212, 20221213, 20221229, 20221230, 20221231",180,,"IPT, DIR, SIZE, PUSH_FLAG","BYTES, BYTES_REV, PACKETS, PACKETS_REV, DURATION, PPI_LEN, PPI_ROUNDTRIPS, PPI_DURATION","FLOW_ENDREASON_IDLE, FLOW_ENDREASON_ACTIVE, FLOW_ENDREASON_END, FLOW_ENDREASON_OTHER","PHIST_SRC_SIZES, PHIST_DST_SIZES, PHIST_SRC_IPT, PHIST_DST_IPT","FLAG_CWR, FLAG_CWR_REV, FLAG_ECE, FLAG_ECE_REV, FLAG_URG, FLAG_URG_REV, FLAG_ACK, FLAG_ACK_REV, FLAG_PSH, FLAG_PSH_REV, FLAG_RST, FLAG_RST_REV, FLAG_SYN, FLAG_SYN_REV, FLAG_FIN, FLAG_FIN_REV","ID, SRC_IP, DST_IP, DST_ASN, DST_PORT, PROTOCOL, TLS_SNI, TLS_JA3, TIME_FIRST, TIME_LAST",,,

cesnet_datazoo/datasets/statistics.py

@@ -1,4 +1,3 @@
-import logging
 import os
 from collections import Counter
 from typing import Any, Literal
@@ -12,8 +11,9 @@ from torch.utils.data import BatchSampler, DataLoader, SequentialSampler
 from tqdm import tqdm
 
 from cesnet_datazoo.config import Protocol
-from cesnet_datazoo.constants import (APP_COLUMN, CATEGORY_COLUMN, FLOWEND_REASON_FEATURES, IPT_POS,
-                                      PHISTS_FEATURES, PPI_COLUMN, SIZE_POS)
+from cesnet_datazoo.constants import (APP_COLUMN, CATEGORY_COLUMN, IPT_POS, PPI_COLUMN, SIZE_POS,
+                                      UDP_PPI_CHANNELS)
+from cesnet_datazoo.datasets.loaders import collate_fn_simple
 from cesnet_datazoo.pytables_data.indices_setup import sort_indices
 from cesnet_datazoo.pytables_data.pytables_dataset import (PyTablesDataset, list_all_tables,
                                                            load_database, worker_init_fn)
@@ -36,17 +36,25 @@ def pick_stats_fields(batch):
         batch[CATEGORY_COLUMN],
     )
 
-def pick_extra_fields(batch, flowstats_features: list[str]):
+def pick_extra_fields(batch, packet_histograms: list[str], flow_endreason_features: list[str]):
     return (
         batch["DST_ASN"],
-        batch[PHISTS_FEATURES],
-        batch[[f for f in FLOWEND_REASON_FEATURES if f in flowstats_features]],
+        batch[packet_histograms],
+        batch[flow_endreason_features],
     )
 
-def simple_collate_fn(batch):
-    return batch
-
-def compute_dataset_statistics(database_path: str, output_dir: str, flowstats_features: list[str], protocol: Protocol, extra_fields: bool, disabled_apps: list[str], num_samples: int | Literal["all"] = 10_000_000, num_workers: int = 4, batch_size: int = 4096, silent: bool = False):
+def compute_dataset_statistics(database_path: str,
+                               tables_app_enum: dict[int, str],
+                               tables_cat_enum: dict[int, str],
+                               output_dir: str,
+                               packet_histograms: list[str],
+                               flowstats_features_boolean: list[str],
+                               protocol: Protocol, extra_fields: bool,
+                               disabled_apps: list[str],
+                               num_samples: int | Literal["all"] = 10_000_000,
+                               num_workers: int = 4,
+                               batch_size: int = 4096,
+                               silent: bool = False):
     stats_pdf_path = os.path.join(output_dir, "dataset-statistics.pdf")
     stats_csv_path = os.path.join(output_dir, "dataset-statistics.csv")
     categories_csv_path = os.path.join(output_dir, "categories.csv")
@@ -74,10 +82,22 @@ def compute_dataset_statistics(database_path: str, output_dir: str, flowstats_fe
     feature_bytes_total = []
     packet_sizes_counter = Counter()
     ipt_counter = Counter()
+    flow_endreason_features = [f for f in flowstats_features_boolean if f.startswith("FLOW_ENDREASON")]
     if not silent:
         print(f"Reading data from {database_path} for statistics")
     table_paths = list_all_tables(database_path)
-    stats_dataset = PyTablesDataset(database_path=database_path, tables_paths=table_paths, flowstats_features=flowstats_features, disabled_apps=disabled_apps, indices=None, return_all_fields=True)
+    stats_dataset = PyTablesDataset(database_path=database_path,
+                                    tables_app_enum=tables_app_enum,
+                                    tables_cat_enum=tables_cat_enum,
+                                    tables_paths=table_paths,
+                                    indices=None,
+                                    disabled_apps=disabled_apps,
+                                    return_all_fields=True,
+                                    flowstats_features=[],
+                                    flowstats_features_boolean=[],
+                                    flowstats_features_phist=[],
+                                    other_fields=[],
+                                    ppi_channels=UDP_PPI_CHANNELS,)
     if num_samples != "all":
         subset_indices = np.random.randint(low=0, high=len(stats_dataset.indices), size=num_samples)
         stats_dataset.indices = sort_indices(stats_dataset.indices[subset_indices])
@@ -87,7 +107,7 @@ def compute_dataset_statistics(database_path: str, output_dir: str, flowstats_fe
         pin_memory=False,
         num_workers=num_workers,
         worker_init_fn=worker_init_fn,
-        collate_fn=simple_collate_fn,
+        collate_fn=collate_fn_simple,
         persistent_workers=False,
         batch_size=None,
         sampler=stats_batch_sampler)
@@ -116,10 +136,10 @@ def compute_dataset_statistics(database_path: str, output_dir: str, flowstats_fe
             quic_ua_series = quic_ua_series.add(pd.Series(user_agent).str.decode("utf-8").value_counts(), fill_value=0)
             quic_version_series = quic_version_series.add(pd.Series(quic_version).value_counts(), fill_value=0)
         if extra_fields:
-            asn, phist, flowend_reason = pick_extra_fields(batch, flowstats_features=flowstats_features)
+            asn, phist, flowend_reason = pick_extra_fields(batch, packet_histograms=packet_histograms, flow_endreason_features=flow_endreason_features)
             asn_series = asn_series.add(pd.Series(asn).value_counts(), fill_value=0)
             flow_endreason_series = flow_endreason_series.add(pd.Series(structured_to_unstructured(flowend_reason).sum(axis=0)), fill_value=0)
-            df2 = pd.DataFrame(data=zip(*np.split(structured_to_unstructured(phist).sum(axis=0), 4)), columns=PHISTS_FEATURES)
+            df2 = pd.DataFrame(data=zip(*np.split(structured_to_unstructured(phist).sum(axis=0), 4)), columns=packet_histograms)
             df_phist = df_phist.add(df2, fill_value=0)
     feature_duration = np.concatenate(feature_duration)
     feature_packets_total = np.concatenate(feature_packets_total)
@@ -183,11 +203,11 @@ def compute_dataset_statistics(database_path: str, output_dir: str, flowstats_fe
         flow_endreason_df = pd.DataFrame({"COUNT": flow_endreason_series.astype("int64")})
         flow_endreason_df["PERC"] = (flow_endreason_df["COUNT"] / flow_endreason_df["COUNT"].sum() * 100).round(3)
         flow_endreason_df.index.name = "FLOW ENDREASON"
-        flow_endreason_df.index = pd.Index([f for f in FLOWEND_REASON_FEATURES if f in flowstats_features])
+        flow_endreason_df.index = pd.Index(flow_endreason_features)
         flow_endreason_df.to_csv(flow_endreason_path)
         # PHIST output
         df_phist.index.name = "BINS"
-        df_phist.columns = list(map(lambda x: x.upper().replace("_", " "), PHISTS_FEATURES))
+        df_phist.columns = list(map(lambda x: x.upper().replace("_", " "), packet_histograms))
         df_phist = df_phist.astype("int64")
         for i, column in zip((1, 3, 5, 7), df_phist.columns):
             df_phist.insert(i, column + " PERC", (df_phist[column] / df_phist[column].sum() * 100).round(3))

cesnet_datazoo/pytables_data/data_scalers.py

@@ -1,30 +1,23 @@
-from __future__ import annotations
-
+import json
 import logging
 import os
 import time
 import warnings
-from typing import TYPE_CHECKING
 
 import numpy as np
+from cesnet_models.transforms import ClipAndScaleFlowstats, ClipAndScalePPI
 from numpy.lib.recfunctions import structured_to_unstructured
 from sklearn.preprocessing import MinMaxScaler, RobustScaler, StandardScaler
-from typing_extensions import assert_never
 
-from cesnet_datazoo.config import DatasetConfig, Scaler, ScalerEnum
-from cesnet_datazoo.constants import (DIR_POS, FLOWSTATS_NO_CLIP, FLOWSTATS_TO_SCALE, IPT_POS,
-                                      PPI_COLUMN, SIZE_POS)
+from cesnet_datazoo.config import DatasetConfig
+from cesnet_datazoo.constants import DIR_POS, FLOWSTATS_NO_CLIP, IPT_POS, PPI_COLUMN, SIZE_POS
 from cesnet_datazoo.pytables_data.pytables_dataset import load_data_from_tables, load_database
-from cesnet_datazoo.utils.fileutils import pickle_dump, pickle_load
 from cesnet_datazoo.utils.random import RandomizedSection, get_fresh_random_generator
 
-if TYPE_CHECKING:
-    from cesnet_datazoo.datasets.cesnet_dataset import CesnetDataset
-
 log = logging.getLogger(__name__)
 
-def get_scaler_attrs(scaler: Scaler) -> dict[str, list[float]]:
-    assert Scaler is not None
+
+def get_scaler_attrs(scaler: StandardScaler | RobustScaler | MinMaxScaler) -> dict[str, list[float]]:
     if isinstance(scaler, StandardScaler):
         assert hasattr(scaler, "mean_") and scaler.mean_ is not None and hasattr(scaler, "scale_") and scaler.scale_ is not None
         scaler_attrs = {"mean_": scaler.mean_.tolist(), "scale_": scaler.scale_.tolist()}
@@ -36,97 +29,8 @@ def get_scaler_attrs(scaler: Scaler) -> dict[str, list[float]]:
         scaler_attrs = {"min_": scaler.min_.tolist(), "scale_": scaler.scale_.tolist()}
     return scaler_attrs
 
-def set_scaler_attrs(scaler: Scaler, scaler_attrs: dict[str, list[float]]):
-    assert Scaler is not None
-    if isinstance(scaler, StandardScaler):
-        assert "mean_" in scaler_attrs and "scale_" in scaler_attrs
-        scaler.mean_ = np.array(scaler_attrs["mean_"])
-        scaler.scale_ = np.array(scaler_attrs["scale_"])
-    elif isinstance(scaler, RobustScaler):
-        assert "center_" in scaler_attrs and "scale_" in scaler_attrs
-        scaler.center_ = np.array(scaler_attrs["center_"])
-        scaler.scale_ = np.array(scaler_attrs["scale_"])
-    elif isinstance(scaler, MinMaxScaler):
-        assert "min_" in scaler_attrs and "scale_" in scaler_attrs
-        scaler.min_ = np.array(scaler_attrs["min_"])
-        scaler.scale_ = np.array(scaler_attrs["scale_"])
-
-def save_scalers_attrs_as_dict(dataset: CesnetDataset) -> dict:
-    assert dataset.flowstats_scaler is not None or dataset.psizes_scaler is not None or dataset.ipt_scaler is not None
-    scalers_dict = {}
-    if dataset.flowstats_scaler is not None:
-        scalers_dict["flowstats_scaler_attrs"] = get_scaler_attrs(dataset.flowstats_scaler)
-    if dataset.psizes_scaler is not None:
-        scalers_dict["psizes_scaler_attrs"] = get_scaler_attrs(dataset.psizes_scaler)
-    if dataset.ipt_scaler is not None:
-        scalers_dict["ipt_scaler_attrs"] = get_scaler_attrs(dataset.ipt_scaler)
-    assert dataset.flowstats_quantiles is not None
-    scalers_dict["flowstats_quantiles"] = dataset.flowstats_quantiles.tolist()
-    return scalers_dict
-
-def fit_or_load_scalers(dataset_config: DatasetConfig, train_indices: np.ndarray) -> tuple[Scaler, Scaler, Scaler, np.ndarray]:
-    # Load the scalers from pickled files if scalers_attrs are not provided
-    if dataset_config.scalers_attrs is None:
-        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")
-        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")
-        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")
-        flowstats_quantiles_path = os.path.join(train_data_path, "stand", f"flowstats_quantiles-q{dataset_config.flowstats_clip}.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)
-            psizes_scaler = pickle_load(psizes_sizes_scaler_path)
-            ipt_scaler = pickle_load(ipt_scaler_path)
-            flowstats_quantiles = pickle_load(flowstats_quantiles_path)
-            return flowstats_scaler, psizes_scaler, ipt_scaler, flowstats_quantiles
-    # Initialize the scalers classes based on the config
-    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)
-    # Load scalers learned attributes from config if provided
-    if dataset_config.scalers_attrs is not None:
-        if "flowstats_scaler_attrs" in dataset_config.scalers_attrs:
-            if flowstats_scaler is not None:
-                set_scaler_attrs(flowstats_scaler, dataset_config.scalers_attrs["flowstats_scaler_attrs"])
-            else:
-                warnings.warn("Ignoring flowstats_scaler_attrs because flowstats_scaler is None")
-        if "psizes_scaler_attrs" in dataset_config.scalers_attrs:
-            if psizes_scaler is not None:
-                set_scaler_attrs(psizes_scaler, dataset_config.scalers_attrs["psizes_scaler_attrs"])
-            else:
-                warnings.warn("Ignoring psizes_scaler_attrs because psizes_scaler is None")
-        if "ipt_scaler_attrs" in dataset_config.scalers_attrs:
-            if ipt_scaler is not None:
-                set_scaler_attrs(ipt_scaler, dataset_config.scalers_attrs["ipt_scaler_attrs"])
-            else:
-                warnings.warn("Ignoring ipt_scaler_attrs because ipt_scaler is None")
-        assert "flowstats_quantiles" in dataset_config.scalers_attrs
-        flowstats_quantiles = np.array(dataset_config.scalers_attrs["flowstats_quantiles"])
-        return flowstats_scaler, psizes_scaler, ipt_scaler, flowstats_quantiles
-    # If the scalers are not loaded at this point, fit them
+def fit_scalers(dataset_config: DatasetConfig, train_indices: np.ndarray) -> None:
+    # Define indices for fitting scalers
     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)
@@ -136,61 +40,71 @@ def fit_or_load_scalers(dataset_config: DatasetConfig, train_indices: np.ndarray
     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,
-        psizes_scaler=psizes_scaler,
-        ipt_scaler=ipt_scaler,
-        flowstats_quantile_clip=dataset_config.flowstats_clip,
-        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(psizes_scaler, psizes_sizes_scaler_path)
-    pickle_dump(ipt_scaler, ipt_scaler_path)
-    pickle_dump(flowstats_quantiles, flowstats_quantiles_path)
-    return flowstats_scaler, psizes_scaler, ipt_scaler, flowstats_quantiles
-
-def fit_scalers(database_path: str,
-                train_tables_paths: list[str],
-                fit_scalers_indices: np.ndarray,
-                flowstats_scaler: Scaler,
-                psizes_scaler: Scaler,
-                ipt_scaler: Scaler,
-                flowstats_quantile_clip: float,
-                ipt_min: int,
-                ipt_max: int,
-                psizes_max: int) -> np.ndarray:
+    # Load data
     start_time = time.time()
-    database, tables = load_database(database_path, tables_paths=train_tables_paths)
+    database, tables = load_database(dataset_config.database_path, tables_paths=dataset_config._get_train_tables_paths())
     data = load_data_from_tables(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
+
+    clip_and_scale_ppi_transform = dataset_config.ppi_transform # TODO Fix after transforms composing is implemented
+    clip_and_scale_flowstats_transform = dataset_config.flowstats_transform
+
+    # Fit the ClipAndScalePPI transform
+    if clip_and_scale_ppi_transform is not None and clip_and_scale_ppi_transform.needs_fitting:
+        assert isinstance(clip_and_scale_ppi_transform, ClipAndScalePPI)
+        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
+        # Fit IPT scaler
+        train_ipt = data_ppi[:, IPT_POS].clip(max=clip_and_scale_ppi_transform.ipt_max, min=clip_and_scale_ppi_transform.ipt_min)
+        train_ipt[padding_mask] = np.nan # NaNs are ignored in sklearn scalers
+        if isinstance(clip_and_scale_ppi_transform.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)
+        clip_and_scale_ppi_transform.ipt_scaler.fit(train_ipt.reshape(-1, 1))
+        # Fit packet sizes scaler
+        train_psizes = data_ppi[:, SIZE_POS].clip(max=clip_and_scale_ppi_transform.psizes_max, min=clip_and_scale_ppi_transform.pszies_min)
         train_psizes[padding_mask] = np.nan
-        if isinstance(psizes_scaler, MinMaxScaler):
+        if isinstance(clip_and_scale_ppi_transform.psizes_scaler, MinMaxScaler):
             train_psizes = np.concatenate((train_psizes, [0]))
-        psizes_scaler.fit(train_psizes.reshape(-1, 1))
-    # Flow statistics
-    train_flowstats = structured_to_unstructured(data[FLOWSTATS_TO_SCALE])
-    flowstats_quantiles = np.quantile(train_flowstats, q=flowstats_quantile_clip, axis=0)
-    flowstats_quantiles[-len(FLOWSTATS_NO_CLIP):] = np.inf # disable clipping for features with "fixed" range
-    if flowstats_scaler:
+        clip_and_scale_ppi_transform.psizes_scaler.fit(train_psizes.reshape(-1, 1))
+        clip_and_scale_ppi_transform.needs_fitting = False
+
+    # Fit the ClipAndScaleFlowstats transform
+    if clip_and_scale_flowstats_transform is not None and clip_and_scale_flowstats_transform.needs_fitting:
+        assert isinstance(clip_and_scale_flowstats_transform, ClipAndScaleFlowstats)
+        train_flowstats = structured_to_unstructured(data[dataset_config.flowstats_features])
+        flowstats_quantiles = np.quantile(train_flowstats, q=clip_and_scale_flowstats_transform.quantile_clip, axis=0)
+        idx_no_clip = [dataset_config.flowstats_features.index(f) for f in FLOWSTATS_NO_CLIP]
+        flowstats_quantiles[idx_no_clip] = np.inf # Disable clipping for features with "fixed" range
         train_flowstats = train_flowstats.clip(max=flowstats_quantiles)
-        flowstats_scaler.fit(train_flowstats)
+        clip_and_scale_flowstats_transform.flowstats_scaler.fit(train_flowstats)
+        clip_and_scale_flowstats_transform.flowstats_quantiles = flowstats_quantiles.tolist()
+        clip_and_scale_flowstats_transform.needs_fitting = False
+
     log.info(f"Reading data and fitting scalers took {time.time() - start_time:.2f} seconds")
-    return flowstats_quantiles
+    train_data_path = dataset_config._get_train_data_path()
+    if clip_and_scale_ppi_transform is not None:
+        ppi_transform_path = os.path.join(train_data_path, "transforms", "ppi-transform.json")
+        ppi_transform_dict = {
+            "psizes_scaler_enum": str(clip_and_scale_ppi_transform._psizes_scaler_enum),
+            "psizes_scaler_attrs": get_scaler_attrs(clip_and_scale_ppi_transform.psizes_scaler),
+            "pszies_min": clip_and_scale_ppi_transform.pszies_min,
+            "psizes_max": clip_and_scale_ppi_transform.psizes_max,
+            "ipt_scaler_enum": str(clip_and_scale_ppi_transform._ipt_scaler_enum),
+            "ipt_scaler_attrs": get_scaler_attrs(clip_and_scale_ppi_transform.ipt_scaler),
+            "ipt_min": clip_and_scale_ppi_transform.ipt_min,
+            "ipt_max": clip_and_scale_ppi_transform.ipt_max,
+        }
+        json.dump(ppi_transform_dict, open(ppi_transform_path, "w"), indent=4)
+    if clip_and_scale_flowstats_transform is not None:
+        assert clip_and_scale_flowstats_transform.flowstats_quantiles is not None
+        flowstats_transform_path = os.path.join(train_data_path, "transforms", "flowstats-transform.json")
+        flowstats_transform_dict = {
+            "flowstats_scaler_enum": str(clip_and_scale_flowstats_transform._flowstats_scaler_enum),
+            "flowstats_scaler_attrs": get_scaler_attrs(clip_and_scale_flowstats_transform.flowstats_scaler),
+            "flowstats_quantiles": clip_and_scale_flowstats_transform.flowstats_quantiles,
+            "quantile_clip": clip_and_scale_flowstats_transform.quantile_clip,
+        }
+        json.dump(flowstats_transform_dict, open(flowstats_transform_path, "w"), indent=4)

cesnet_datazoo/pytables_data/indices_setup.py

@@ -7,13 +7,11 @@ from enum import Enum
 
 import numpy as np
 import pandas as pd
-from sklearn.preprocessing import LabelEncoder
 
 from cesnet_datazoo.config import DatasetConfig
-from cesnet_datazoo.constants import (INDICES_INDEX_POS, INDICES_LABEL_POS, INDICES_TABLE_POS,
-                                      UNKNOWN_STR_LABEL)
+from cesnet_datazoo.constants import INDICES_INDEX_POS, INDICES_LABEL_POS, INDICES_TABLE_POS
 from cesnet_datazoo.pytables_data.pytables_dataset import init_test_indices, init_train_indices
-from cesnet_datazoo.utils.fileutils import pickle_dump, pickle_load, yaml_dump, yaml_load
+from cesnet_datazoo.utils.fileutils import yaml_dump, yaml_load
 from cesnet_datazoo.utils.random import RandomizedSection, get_fresh_random_generator
 
 log = logging.getLogger(__name__)
@@ -78,59 +76,56 @@ def date_weight_sample_train_indices(dataset_config: DatasetConfig, train_indice
     sampled_train_indices = np.concatenate(sampled_indicies_per_date)
     return sampled_train_indices
 
-def indices_to_app_counts(indices: np.ndarray, database_enum: dict[int, str]) -> pd.Series:
+def indices_to_app_counts(indices: np.ndarray, tables_app_enum: dict[int, str]) -> pd.Series:
     app_counts = pd.Series(indices[:, INDICES_LABEL_POS]).value_counts()
-    app_counts.index = app_counts.index.map(lambda x: database_enum[x])
+    app_counts.index = app_counts.index.map(lambda x: tables_app_enum[x])
     return app_counts
 
-def compute_known_app_counts(dataset_indices: IndicesTuple, database_enum: dict[int, str]) -> pd.DataFrame:
-    train_app_counts = indices_to_app_counts(dataset_indices.train_indices, database_enum)
-    val_known_app_counts = indices_to_app_counts(dataset_indices.val_known_indices, database_enum)
-    test_known_app_counts = indices_to_app_counts(dataset_indices.test_known_indices, database_enum)
+def compute_known_app_counts(dataset_indices: IndicesTuple, tables_app_enum: dict[int, str]) -> pd.DataFrame:
+    train_app_counts = indices_to_app_counts(dataset_indices.train_indices, tables_app_enum)
+    val_known_app_counts = indices_to_app_counts(dataset_indices.val_known_indices, tables_app_enum)
+    test_known_app_counts = indices_to_app_counts(dataset_indices.test_known_indices, tables_app_enum)
     df = pd.DataFrame(data={"Train": train_app_counts, "Validation": val_known_app_counts, "Test": test_known_app_counts}).fillna(0).astype("int64")
     return df
 
-def compute_unknown_app_counts(dataset_indices: IndicesTuple, database_enum: dict[int, str]) -> pd.DataFrame:
-    val_unknown_app_counts = indices_to_app_counts(dataset_indices.val_unknown_indices, database_enum)
-    test_unknown_app_counts = indices_to_app_counts(dataset_indices.test_unknown_indices, database_enum)
+def compute_unknown_app_counts(dataset_indices: IndicesTuple, tables_app_enum: dict[int, str]) -> pd.DataFrame:
+    val_unknown_app_counts = indices_to_app_counts(dataset_indices.val_unknown_indices, tables_app_enum)
+    test_unknown_app_counts = indices_to_app_counts(dataset_indices.test_unknown_indices, tables_app_enum)
     df = pd.DataFrame(data={"Validation": val_unknown_app_counts, "Test": test_unknown_app_counts}).fillna(0).astype("int64")
     return df
 
-def init_or_load_train_indices(dataset_config: DatasetConfig, servicemap: pd.DataFrame, disable_indices_cache: bool) -> tuple[np.ndarray, np.ndarray, LabelEncoder, dict[int, str], dict[int, str]]:
+def init_or_load_train_indices(dataset_config: DatasetConfig, tables_app_enum: dict[int, str], servicemap: pd.DataFrame, disable_indices_cache: bool) -> tuple[np.ndarray, np.ndarray, list[str], list[str]]:
     train_data_path = dataset_config._get_train_data_path()
     init_train_data(train_data_path)
     if not os.path.isfile(os.path.join(train_data_path, TRAIN_DATA_PARAMS_FILE)):
         log.info("Processing train indices")
         train_data_params = dataset_config._get_train_data_params()
-        train_known_indices, train_unknown_indices, known_apps_database_enum, unknown_apps_database_enum = init_train_indices(
-            train_data_params=train_data_params,
-            servicemap=servicemap,
-            database_path=dataset_config.database_path,
-            rng=get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.INIT_TRAIN_INDICES))
-        encoder = LabelEncoder().fit(list(known_apps_database_enum.values()))
-        encoder.classes_ = np.append(encoder.classes_, UNKNOWN_STR_LABEL)
+        train_known_indices, train_unknown_indices, known_apps, unknown_apps = init_train_indices(train_data_params=train_data_params,
+                                                                                                  database_path=dataset_config.database_path,
+                                                                                                  tables_app_enum=tables_app_enum,
+                                                                                                  servicemap=servicemap,
+                                                                                                  rng=get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.INIT_TRAIN_INDICES))
         if not disable_indices_cache:
             yaml_dump({k: str(v) if isinstance(v, Enum) else list(v) if isinstance(v, tuple) else v for k, v in dataclasses.asdict(train_data_params).items()}, os.path.join(train_data_path, TRAIN_DATA_PARAMS_FILE))
-            yaml_dump(known_apps_database_enum, os.path.join(train_data_path, "known_apps_database_enum.yaml"))
-            yaml_dump(unknown_apps_database_enum, os.path.join(train_data_path, "unknown_apps_database_enum.yaml"))
-            pickle_dump(encoder, os.path.join(train_data_path, "encoder.pickle"))
+            yaml_dump(known_apps, os.path.join(train_data_path, "known_apps.yaml"))
+            yaml_dump(unknown_apps, os.path.join(train_data_path, "unknown_apps.yaml"))
             np.save(os.path.join(train_data_path, "train_known_indices.npy"), train_known_indices)
             np.save(os.path.join(train_data_path, "train_unknown_indices.npy"), train_unknown_indices)
     else:
-        known_apps_database_enum = yaml_load(os.path.join(train_data_path, "known_apps_database_enum.yaml"))
-        unknown_apps_database_enum = yaml_load(os.path.join(train_data_path, "unknown_apps_database_enum.yaml"))
-        encoder = pickle_load(os.path.join(train_data_path, "encoder.pickle"))
+        known_apps = yaml_load(os.path.join(train_data_path, "known_apps.yaml"))
+        unknown_apps = yaml_load(os.path.join(train_data_path, "unknown_apps.yaml"))
         train_known_indices = np.load(os.path.join(train_data_path, "train_known_indices.npy"))
         train_unknown_indices = np.load(os.path.join(train_data_path, "train_unknown_indices.npy"))
-    return train_known_indices, train_unknown_indices, encoder, known_apps_database_enum, unknown_apps_database_enum
+    return train_known_indices, train_unknown_indices, known_apps, unknown_apps
 
-def init_or_load_val_indices(dataset_config: DatasetConfig, known_apps_database_enum: dict[int, str], unknown_apps_database_enum: dict[int, str], disable_indices_cache: bool) -> tuple[np.ndarray, np.ndarray, str]:
-    val_data_params, val_data_path = dataset_config._get_val_data_params_and_path(known_apps_database_enum=known_apps_database_enum, unknown_apps_database_enum=unknown_apps_database_enum)
+def init_or_load_val_indices(dataset_config: DatasetConfig, known_apps: list[str], unknown_apps: list[str], tables_app_enum: dict[int, str], disable_indices_cache: bool) -> tuple[np.ndarray, np.ndarray, str]:
+    val_data_params, val_data_path = dataset_config._get_val_data_params_and_path(known_apps=known_apps, unknown_apps=unknown_apps)
     init_test_data(val_data_path)
     if not os.path.isfile(os.path.join(val_data_path, TEST_DATA_PARAMS_FILE)):
         log.info("Processing validation indices")
         val_known_indices, val_unknown_indices = init_test_indices(test_data_params=val_data_params,
                                                                    database_path=dataset_config.database_path,
+                                                                   tables_app_enum=tables_app_enum,
                                                                    rng=get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.INIT_VAL_INIDICES))
         if not disable_indices_cache:
             yaml_dump(dataclasses.asdict(val_data_params), os.path.join(val_data_path, TEST_DATA_PARAMS_FILE))
@@ -141,13 +136,14 @@ def init_or_load_val_indices(dataset_config: DatasetConfig, known_apps_database_
         val_unknown_indices = np.load(os.path.join(val_data_path, "val_unknown_indices.npy"))
     return val_known_indices, val_unknown_indices, val_data_path
 
-def init_or_load_test_indices(dataset_config: DatasetConfig, known_apps_database_enum: dict[int, str], unknown_apps_database_enum: dict[int, str], disable_indices_cache: bool) -> tuple[np.ndarray, np.ndarray, str]:
-    test_data_params, test_data_path = dataset_config._get_test_data_params_and_path(known_apps_database_enum=known_apps_database_enum, unknown_apps_database_enum=unknown_apps_database_enum)
+def init_or_load_test_indices(dataset_config: DatasetConfig, known_apps: list[str], unknown_apps: list[str], tables_app_enum: dict[int, str], disable_indices_cache: bool) -> tuple[np.ndarray, np.ndarray, str]:
+    test_data_params, test_data_path = dataset_config._get_test_data_params_and_path(known_apps=known_apps, unknown_apps=unknown_apps)
     init_test_data(test_data_path)
     if not os.path.isfile(os.path.join(test_data_path, TEST_DATA_PARAMS_FILE)):
         log.info("Processing test indices")
         test_known_indices, test_unknown_indices = init_test_indices(test_data_params=test_data_params,
                                                                      database_path=dataset_config.database_path,
+                                                                     tables_app_enum=tables_app_enum,
                                                                      rng=get_fresh_random_generator(dataset_config=dataset_config, section=RandomizedSection.INIT_TEST_INDICES))
         if not disable_indices_cache:
             yaml_dump(dataclasses.asdict(test_data_params), os.path.join(test_data_path, TEST_DATA_PARAMS_FILE))
@@ -160,7 +156,7 @@ def init_or_load_test_indices(dataset_config: DatasetConfig, known_apps_database
 
 def init_train_data(train_data_path: str):
     os.makedirs(train_data_path, exist_ok=True)
-    os.makedirs(os.path.join(train_data_path, "stand"), exist_ok=True)
+    os.makedirs(os.path.join(train_data_path, "transforms"), exist_ok=True)
     os.makedirs(os.path.join(train_data_path, "preload"), exist_ok=True)
 
 def init_test_data(test_data_path: str):