| """ |
| Data loading and preprocessing module. |
| |
| Handles loading from HuggingFace Hub or local CSV, data cleaning, |
| train/val/test splitting with stratification by material type. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| from pathlib import Path |
| from typing import Optional, Tuple |
|
|
| import numpy as np |
| import pandas as pd |
| from sklearn.model_selection import ( |
| GroupShuffleSplit, |
| StratifiedShuffleSplit, |
| train_test_split, |
| ) |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def load_dataset( |
| source: str, |
| local_path: Optional[str] = None, |
| random_state: int = 42, |
| ) -> pd.DataFrame: |
| """ |
| Load dataset from HuggingFace Hub or local CSV. |
| |
| Parameters |
| ---------- |
| source : str |
| HuggingFace dataset ID (e.g., 'TWLAb/femtosecond-laser-hydrogel-etching-data') |
| local_path : str, optional |
| Path to local CSV file (used if available, avoids re-downloading) |
| random_state : int |
| Random seed for reproducibility |
| |
| Returns |
| ------- |
| pd.DataFrame |
| Loaded and cleaned dataset |
| """ |
| if local_path and Path(local_path).exists(): |
| logger.info(f"Loading dataset from local path: {local_path}") |
| df = pd.read_csv(local_path) |
| else: |
| logger.info(f"Loading dataset from HuggingFace Hub: {source}") |
| try: |
| from datasets import load_dataset as hf_load |
| ds = hf_load(source, split="train") |
| df = ds.to_pandas() |
| except Exception as e: |
| logger.error(f"Failed to load from Hub: {e}") |
| raise |
|
|
| logger.info(f"Dataset loaded: {df.shape[0]} samples, {df.shape[1]} columns") |
| return df |
|
|
|
|
| def clean_dataset(df: pd.DataFrame) -> pd.DataFrame: |
| """ |
| Clean dataset: handle missing values, remove outliers, fix dtypes. |
| |
| Parameters |
| ---------- |
| df : pd.DataFrame |
| Raw dataset |
| |
| Returns |
| ------- |
| pd.DataFrame |
| Cleaned dataset |
| """ |
| initial_size = len(df) |
|
|
| |
| numeric_cols = df.select_dtypes(include=[np.number]).columns |
| df = df.dropna(subset=numeric_cols) |
|
|
| |
| df = df[df["etch_depth_um"] > 0] |
| df = df[df["etch_width_um"] > 0] |
| df = df[df["surface_roughness_Sa_um"] > 0] |
| df = df[df["power_mW"] > 0] |
| df = df[df["scan_speed_mm_s"] > 0] |
|
|
| |
| target_cols = [ |
| "etch_depth_um", "etch_width_um", "surface_roughness_Sa_um", |
| "aspect_ratio", "side_wall_angle_deg" |
| ] |
| for col in target_cols: |
| if col in df.columns: |
| mean, std = df[col].mean(), df[col].std() |
| df = df[(df[col] >= mean - 5 * std) & (df[col] <= mean + 5 * std)] |
|
|
| removed = initial_size - len(df) |
| if removed > 0: |
| logger.info(f"Removed {removed} rows during cleaning ({removed/initial_size*100:.1f}%)") |
|
|
| return df.reset_index(drop=True) |
|
|
|
|
| def split_dataset( |
| df: pd.DataFrame, |
| test_size: float = 0.15, |
| val_size: float = 0.15, |
| group_column: str = "material_type", |
| random_state: int = 42, |
| ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: |
| """ |
| Split dataset into train/validation/test with stratification by material type. |
| |
| Uses StratifiedShuffleSplit to ensure proportional material representation |
| across all splits, which is critical for generalization assessment. |
| |
| Parameters |
| ---------- |
| df : pd.DataFrame |
| Full dataset |
| test_size : float |
| Fraction for test set |
| val_size : float |
| Fraction for validation set (from remaining after test) |
| group_column : str |
| Column to stratify by |
| random_state : int |
| Random seed |
| |
| Returns |
| ------- |
| tuple of (train_df, val_df, test_df) |
| """ |
| if group_column not in df.columns: |
| logger.warning(f"Group column '{group_column}' not found. Using random split.") |
| train_val, test = train_test_split(df, test_size=test_size, random_state=random_state) |
| val_frac = val_size / (1 - test_size) |
| train, val = train_test_split(train_val, test_size=val_frac, random_state=random_state) |
| else: |
| |
| strat_col = df[group_column] |
|
|
| sss_test = StratifiedShuffleSplit(n_splits=1, test_size=test_size, random_state=random_state) |
| train_val_idx, test_idx = next(sss_test.split(df, strat_col)) |
|
|
| train_val = df.iloc[train_val_idx] |
| test = df.iloc[test_idx] |
|
|
| |
| val_frac = val_size / (1 - test_size) |
| strat_col_tv = train_val[group_column] |
| sss_val = StratifiedShuffleSplit(n_splits=1, test_size=val_frac, random_state=random_state) |
| train_idx, val_idx = next(sss_val.split(train_val, strat_col_tv)) |
|
|
| train = train_val.iloc[train_idx] |
| val = train_val.iloc[val_idx] |
|
|
| logger.info( |
| f"Split sizes - Train: {len(train)} ({len(train)/len(df)*100:.1f}%), " |
| f"Val: {len(val)} ({len(val)/len(df)*100:.1f}%), " |
| f"Test: {len(test)} ({len(test)/len(df)*100:.1f}%)" |
| ) |
|
|
| |
| if group_column in df.columns: |
| for name, subset in [("Train", train), ("Val", val), ("Test", test)]: |
| dist = subset[group_column].value_counts(normalize=True) |
| logger.debug(f"{name} material distribution:\n{dist}") |
|
|
| return train.reset_index(drop=True), val.reset_index(drop=True), test.reset_index(drop=True) |
|
|
|
|
| def get_feature_target_arrays( |
| df: pd.DataFrame, |
| feature_columns: list[str], |
| target_columns: list[str], |
| ) -> Tuple[np.ndarray, np.ndarray]: |
| """ |
| Extract feature and target arrays from DataFrame. |
| |
| Parameters |
| ---------- |
| df : pd.DataFrame |
| feature_columns : list of str |
| target_columns : list of str |
| |
| Returns |
| ------- |
| tuple of (X, y) as numpy arrays |
| """ |
| X = df[feature_columns].values.astype(np.float32) |
| y = df[target_columns].values.astype(np.float32) |
| return X, y |
|
|
