import numpy as np #%% Cell 1 — Imports import pandas as pd import numpy as np import torch import warnings import joblib import os from sklearn.preprocessing import MinMaxScaler, FunctionTransformer from sklearn.pipeline import Pipeline from sklearn.model_selection import train_test_split from torch.utils.data import TensorDataset, DataLoader warnings.filterwarnings("ignore") def log1p_base10(x): return np.log10(1 + x) #%% Cell 2 — IDSDataPipeline class class IDSDataPipeline: """End-to-end preprocessing + DataLoader builder for one attack class. Loads benign and attack CSVs, applies the shared preprocessing artifacts (column drops, flag binning, one-hot encoding), filters invalid rows, splits into train/val/test, fits the log+minmax pipeline on train only, and exposes ready-to-use PyTorch DataLoaders. Attributes ---------- train_loader, val_loader, test_loader : torch.utils.data.DataLoader Shuffled-train, non-shuffled val/test loaders over (X, y) tensors. input_dim : int Number of features fed to the model. numeric_pipeline : sklearn.pipeline.Pipeline The fitted log1p_base10 + MinMaxScaler pipeline (saved to disk). feature_names : list[str] Final ordered column names of the feature matrix. """ def __init__( self, attack: str, benign_data_path: str = 'Data/benign_only/all_days_benign.csv', attack_data_dir: str = 'Data/attacks_only', preprocessing_dir: str = '_prepcosessing_artefacts/', checkpoints_dir: str = 'checkpoints_MLP', batch_size: int = 256, test_size: float = 0.10, val_size: float = 0.10, random_state: int = 42, verbose: bool = True, ): self.attack = attack self.benign_data_path = benign_data_path self.attack_data_path = os.path.join(attack_data_dir, f"{attack}.csv") self.preprocessing_dir = preprocessing_dir self.checkpoint_dir = os.path.join(checkpoints_dir, attack) self.batch_size = batch_size self.test_size = test_size self.val_size = val_size self.random_state = random_state self.verbose = verbose os.makedirs(self.checkpoint_dir, exist_ok=True) # Filled in by .build() self.train_loader = None self.val_loader = None self.test_loader = None self.input_dim = None self.numeric_pipeline = None self.feature_names = None self.build() # ---------- internal steps ---------- def _log(self, msg): if self.verbose: print(msg) def _load_artifacts(self): self.ohe = joblib.load(os.path.join(self.preprocessing_dir, "onehot_encoder.pkl")) schema = joblib.load(os.path.join(self.preprocessing_dir, "column_schema.pkl")) self.col_to_drop = schema["col_to_drop"] self.onehot_cols = schema["onehot_cols"] self.flag_cols = schema["flag_cols"] self.flag_bin_config = schema["flag_bin_config"] self.numerical_cols = schema["numerical_cols"] self.ohe_feature_names = schema["ohe_feature_names"] self._log( f"Loaded schema: {len(self.numerical_cols)} numerical, " f"{len(self.flag_cols)} flag, {len(self.onehot_cols)} OHE source " f"-> {len(self.ohe_feature_names)} OHE features" ) def _load_data(self): df_benign = pd.read_csv(self.benign_data_path) df_attack = pd.read_csv(self.attack_data_path) df_benign.drop(columns=self.col_to_drop, inplace=True) df_attack.drop(columns=self.col_to_drop, inplace=True) if 'AttackFamily' in df_attack.columns: df_attack.drop(columns=['AttackFamily'], inplace=True) assert list(df_benign.columns) == list(df_attack.columns), \ "Benign and attack dataframes have different columns after dropping." return df_benign, df_attack @staticmethod def _bin_flag_columns(df, bin_config): df = df.copy() for col, edges in bin_config.items(): df[col] = pd.cut(df[col], bins=edges, labels=False, right=True).astype(np.int8) return df def _encode(self, df): """Apply flag binning and OHE; return (numerical, flag, ohe) concatenated df.""" df = self._bin_flag_columns(df, self.flag_bin_config) ohe_arr = self.ohe.transform(df[self.onehot_cols]) ohe_df = pd.DataFrame(ohe_arr, columns=self.ohe_feature_names, index=df.index) return pd.concat( [df[self.numerical_cols].reset_index(drop=True), df[self.flag_cols].reset_index(drop=True), ohe_df.reset_index(drop=True)], axis=1, ) def _filter_invalid(self, df): """Drop rows with negative or non-finite numerical values.""" mask = (df[self.numerical_cols] >= 0).all(axis=1) & \ np.isfinite(df[self.numerical_cols]).all(axis=1) return df[mask] def _split(self, X, y): X_trainval, X_test, y_trainval, y_test = train_test_split( X, y, test_size=self.test_size, stratify=y, random_state=self.random_state, ) # val_size is fraction of the original; convert to fraction of remaining. val_relative = self.val_size / (1.0 - self.test_size) X_train, X_val, y_train, y_val = train_test_split( X_trainval, y_trainval, test_size=val_relative, stratify=y_trainval, random_state=self.random_state, ) return X_train, X_val, X_test, y_train, y_val, y_test def _fit_numeric(self, X_train): log_transformer = FunctionTransformer( func=log1p_base10, validate=False, feature_names_out="one-to-one", ) self.numeric_pipeline = Pipeline([ ("log_transform", log_transformer), ("minmax_scaler", MinMaxScaler()), ]) self.numeric_pipeline.fit(X_train[self.numerical_cols]) def _transform_numeric(self, X): return pd.DataFrame( self.numeric_pipeline.transform(X[self.numerical_cols]), columns=self.numerical_cols, index=X.index, ) def _assemble(self, X, X_num): """Reassemble: scaled numerics + flag bins + OHE columns.""" return pd.concat( [X_num, X[self.flag_cols], X[self.ohe_feature_names]], axis=1, ) def _make_loader(self, X_df, y_series, shuffle): X_t = torch.tensor(X_df.values, dtype=torch.float32) y_t = torch.tensor(y_series.values, dtype=torch.float32).unsqueeze(1) ds = TensorDataset(X_t, y_t) return DataLoader(ds, batch_size=self.batch_size, shuffle=shuffle, drop_last=False) # ---------- orchestration ---------- def build(self): self._load_artifacts() df_benign, df_attack = self._load_data() # Encode + filter (processed versions) df_benign_proc = self._encode(df_benign) df_attack_proc = self._encode(df_attack) df_benign_proc = self._filter_invalid(df_benign_proc) df_attack_proc = self._filter_invalid(df_attack_proc) # Keep raw versions aligned with the filtered processed versions. # _encode reset the index of df_*_proc to 0..N-1 (via reset_index(drop=True) # inside the concat), so we align by position rather than by original index. df_benign_raw = df_benign.iloc[: len(df_benign_proc)].reset_index(drop=True) df_attack_raw = df_attack.iloc[: len(df_attack_proc)].reset_index(drop=True) # The above assumes _filter_invalid drops rows from the END only — which # is not generally true. Safer: align by the actual surviving positions. # Actually: track surviving positions explicitly. # Re-do the filtering with explicit index preservation. df_benign_proc_unfiltered = self._encode(df_benign) df_attack_proc_unfiltered = self._encode(df_attack) mask_benign = (df_benign_proc_unfiltered[self.numerical_cols] >= 0).all(axis=1) & \ np.isfinite(df_benign_proc_unfiltered[self.numerical_cols]).all(axis=1) mask_attack = (df_attack_proc_unfiltered[self.numerical_cols] >= 0).all(axis=1) & \ np.isfinite(df_attack_proc_unfiltered[self.numerical_cols]).all(axis=1) df_benign_proc = df_benign_proc_unfiltered[mask_benign].reset_index(drop=True) df_attack_proc = df_attack_proc_unfiltered[mask_attack].reset_index(drop=True) df_benign_raw = df_benign[mask_benign.values].reset_index(drop=True) df_attack_raw = df_attack[mask_attack.values].reset_index(drop=True) # Add labels df_benign_proc["label"] = 0 df_attack_proc["label"] = 1 df_benign_raw["label"] = 0 df_attack_raw["label"] = 1 # Combine df_full = pd.concat([df_benign_proc, df_attack_proc], axis=0, ignore_index=True) df_raw_full = pd.concat([df_benign_raw, df_attack_raw], axis=0, ignore_index=True) # Shuffle with same random_state — apply same permutation to both perm = df_full.sample(frac=1, random_state=self.random_state).index df_full = df_full.loc[perm].reset_index(drop=True) df_raw_full = df_raw_full.loc[perm].reset_index(drop=True) self._log(f"\nFull dataset shape: {df_full.shape}") self._log(f"Class balance:\n{df_full['label'].value_counts(normalize=True)}") X = df_full.drop(columns=["label"]) y = df_full["label"] # Split (returns indices we can apply to df_raw_full too) indices = np.arange(len(df_full)) train_idx, test_idx = train_test_split( indices, test_size=self.test_size, stratify=y, random_state=self.random_state ) val_relative = self.val_size / (1.0 - self.test_size) train_idx, val_idx = train_test_split( train_idx, test_size=val_relative, stratify=y.iloc[train_idx], random_state=self.random_state ) X_train, X_val, X_test = X.iloc[train_idx], X.iloc[val_idx], X.iloc[test_idx] y_train, y_val, y_test = y.iloc[train_idx], y.iloc[val_idx], y.iloc[test_idx] # The raw splits, aligned to processed: raw_train = df_raw_full.iloc[train_idx].reset_index(drop=True) raw_val = df_raw_full.iloc[val_idx].reset_index(drop=True) raw_test = df_raw_full.iloc[test_idx].reset_index(drop=True) self._log(f"\nTrain: {X_train.shape}, Val: {X_val.shape}, Test: {X_test.shape}") # Fit numeric pipeline on train only, transform all splits self._fit_numeric(X_train) X_train_num = self._transform_numeric(X_train) X_val_num = self._transform_numeric(X_val) X_test_num = self._transform_numeric(X_test) X_train_final = self._assemble(X_train, X_train_num) X_val_final = self._assemble(X_val, X_val_num) X_test_final = self._assemble(X_test, X_test_num) self.X_train = X_train_final.values self.X_val = X_val_final.values self.X_test = X_test_final.values self.y_train = y_train.values self.y_val = y_val.values self.y_test = y_test.values # NEW: raw dataframes per split, same row ordering as processed splits. self.raw_train = raw_train self.raw_val = raw_val self.raw_test = raw_test self.feature_names = X_train_final.columns.tolist() self.input_dim = X_train_final.shape[1] self._log(f"\nFinal feature count: {self.input_dim}") # Persist fitted numeric pipeline joblib.dump( self.numeric_pipeline, os.path.join(self.checkpoint_dir, "numeric_pipeline.pkl"), ) # Loaders (unchanged) self.train_loader = self._make_loader(X_train_final, y_train, shuffle=True) self.val_loader = self._make_loader(X_val_final, y_val, shuffle=False) self.test_loader = self._make_loader(X_test_final, y_test, shuffle=False) #%% Cell 3 — Usage