ds-agent-cli 0.1.0

package/src/tools/advanced_preprocessing.py

@@ -0,0 +1,549 @@
+"""
+Advanced Preprocessing Tools
+Tools for handling imbalanced data, feature scaling, and strategic data splitting.
+"""
+
+import polars as pl
+import numpy as np
+from typing import Dict, Any, List, Optional, Tuple
+from pathlib import Path
+import sys
+import os
+import joblib
+import warnings
+
+warnings.filterwarnings('ignore')
+
+# Add parent directory to path for imports
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler, LabelEncoder
+from imblearn.over_sampling import SMOTE, ADASYN, BorderlineSMOTE
+from imblearn.under_sampling import RandomUnderSampler, TomekLinks, EditedNearestNeighbours
+from imblearn.combine import SMOTETomek, SMOTEENN
+from collections import Counter
+
+from ds_agent.utils.polars_helpers import (
+    load_dataframe, save_dataframe, get_numeric_columns,
+    get_categorical_columns, split_features_target
+)
+from ds_agent.utils.validation import (
+    validate_file_exists, validate_file_format, validate_dataframe,
+    validate_column_exists
+)
+
+
+def handle_imbalanced_data(
+    file_path: str,
+    target_col: str,
+    strategy: str = "smote",
+    sampling_ratio: float = 1.0,
+    output_path: str = None,
+    random_state: int = 42
+) -> Dict[str, Any]:
+    """
+    Handle imbalanced datasets using various resampling techniques.
+    
+    Args:
+        file_path: Path to dataset
+        target_col: Target column name
+        strategy: Resampling strategy:
+            - 'smote': Synthetic Minority Over-sampling (SMOTE)
+            - 'adasyn': Adaptive Synthetic Sampling
+            - 'borderline_smote': Borderline SMOTE variant
+            - 'random_undersample': Random undersampling
+            - 'tomek': Tomek Links undersampling
+            - 'smote_tomek': Combined SMOTE + Tomek Links
+            - 'smote_enn': Combined SMOTE + Edited Nearest Neighbours
+            - 'class_weights': Return class weights (no resampling)
+        sampling_ratio: Ratio of minority to majority class (0.5 = 50%, 1.0 = 100%)
+        output_path: Path to save balanced dataset
+        random_state: Random seed
+        
+    Returns:
+        Dictionary with balancing results and class distributions
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    validate_column_exists(df, target_col)
+    
+    # Get original class distribution
+    original_dist = df[target_col].value_counts().to_dict()
+    original_counts = dict(sorted(original_dist.items()))
+    
+    print(f"📊 Original class distribution: {original_counts}")
+    
+    # Calculate imbalance ratio
+    class_counts = list(original_counts.values())
+    imbalance_ratio = max(class_counts) / min(class_counts)
+    
+    if imbalance_ratio < 1.5:
+        return {
+            'status': 'skipped',
+            'message': 'Dataset is already balanced (ratio < 1.5)',
+            'original_distribution': original_counts,
+            'imbalance_ratio': float(imbalance_ratio)
+        }
+    
+    # Prepare data
+    X, y = split_features_target(df, target_col)
+    
+    # Handle class weights strategy (no resampling)
+    if strategy == "class_weights":
+        from sklearn.utils.class_weight import compute_class_weight
+        classes = np.unique(y)
+        weights = compute_class_weight('balanced', classes=classes, y=y)
+        class_weights = dict(zip(classes, weights))
+        
+        return {
+            'status': 'success',
+            'strategy': 'class_weights',
+            'class_weights': {str(k): float(v) for k, v in class_weights.items()},
+            'original_distribution': original_counts,
+            'imbalance_ratio': float(imbalance_ratio),
+            'recommendation': 'Use class_weight parameter in your model training'
+        }
+    
+    # Create resampler based on strategy
+    sampling_strategy = sampling_ratio if sampling_ratio < 1.0 else 'auto'
+    
+    if strategy == "smote":
+        resampler = SMOTE(sampling_strategy=sampling_strategy, random_state=random_state)
+    elif strategy == "adasyn":
+        resampler = ADASYN(sampling_strategy=sampling_strategy, random_state=random_state)
+    elif strategy == "borderline_smote":
+        resampler = BorderlineSMOTE(sampling_strategy=sampling_strategy, random_state=random_state)
+    elif strategy == "random_undersample":
+        resampler = RandomUnderSampler(sampling_strategy=sampling_strategy, random_state=random_state)
+    elif strategy == "tomek":
+        resampler = TomekLinks(sampling_strategy='auto')
+    elif strategy == "smote_tomek":
+        resampler = SMOTETomek(sampling_strategy=sampling_strategy, random_state=random_state)
+    elif strategy == "smote_enn":
+        resampler = SMOTEENN(sampling_strategy=sampling_strategy, random_state=random_state)
+    else:
+        raise ValueError(f"Unsupported strategy: {strategy}")
+    
+    # Perform resampling
+    print(f"⚖️ Applying {strategy} resampling...")
+    X_resampled, y_resampled = resampler.fit_resample(X, y)
+    
+    # Get new class distribution
+    new_counts = dict(Counter(y_resampled))
+    new_counts = dict(sorted(new_counts.items()))
+    
+    print(f"✅ New class distribution: {new_counts}")
+    
+    # Calculate changes
+    total_original = sum(original_counts.values())
+    total_new = sum(new_counts.values())
+    
+    changes = {
+        str(cls): {
+            'original': original_counts.get(cls, 0),
+            'new': new_counts.get(cls, 0),
+            'change': new_counts.get(cls, 0) - original_counts.get(cls, 0)
+        }
+        for cls in set(list(original_counts.keys()) + list(new_counts.keys()))
+    }
+    
+    # Create balanced dataframe
+    feature_cols = [col for col in df.columns if col != target_col]
+    balanced_data = {col: X_resampled[:, i] for i, col in enumerate(feature_cols)}
+    balanced_data[target_col] = y_resampled
+    
+    balanced_df = pl.DataFrame(balanced_data)
+    
+    # Save if output path provided
+    if output_path:
+        save_dataframe(balanced_df, output_path)
+        print(f"💾 Balanced dataset saved to: {output_path}")
+    
+    return {
+        'status': 'success',
+        'strategy': strategy,
+        'original_distribution': original_counts,
+        'new_distribution': new_counts,
+        'changes_by_class': changes,
+        'total_samples_before': total_original,
+        'total_samples_after': total_new,
+        'sample_change': f"{'+' if total_new > total_original else ''}{total_new - total_original}",
+        'new_imbalance_ratio': float(max(new_counts.values()) / min(new_counts.values())),
+        'output_path': output_path
+    }
+
+
+def perform_feature_scaling(
+    file_path: str,
+    scaler_type: str = "standard",
+    columns: Optional[List[str]] = None,
+    output_path: Optional[str] = None,
+    scaler_save_path: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Scale features using various normalization techniques.
+    
+    Args:
+        file_path: Path to dataset
+        scaler_type: Scaling method:
+            - 'standard': StandardScaler (mean=0, std=1)
+            - 'minmax': MinMaxScaler (range 0-1)
+            - 'robust': RobustScaler (median, IQR - robust to outliers)
+            - 'power': PowerTransformer (Yeo-Johnson, makes data more Gaussian)
+            - 'quantile': QuantileTransformer (uniform or normal output distribution)
+        columns: List of columns to scale (None = all numeric columns)
+        output_path: Path to save scaled dataset
+        scaler_save_path: Path to save fitted scaler for future use
+        
+    Returns:
+        Dictionary with scaling statistics
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    # Get numeric columns if not specified
+    if columns is None:
+        columns = get_numeric_columns(df)
+        print(f"🔢 Auto-detected {len(columns)} numeric columns for scaling")
+    else:
+        for col in columns:
+            validate_column_exists(df, col)
+    
+    if not columns:
+        return {
+            'status': 'skipped',
+            'message': 'No numeric columns found to scale'
+        }
+    
+    # Create scaler
+    if scaler_type == "standard":
+        scaler = StandardScaler()
+    elif scaler_type == "minmax":
+        scaler = MinMaxScaler()
+    elif scaler_type == "robust":
+        scaler = RobustScaler()
+    elif scaler_type == "power":
+        from sklearn.preprocessing import PowerTransformer
+        scaler = PowerTransformer(method='yeo-johnson', standardize=True)
+        print("   📐 Using Yeo-Johnson PowerTransformer (makes data more Gaussian)")
+    elif scaler_type == "quantile":
+        from sklearn.preprocessing import QuantileTransformer
+        scaler = QuantileTransformer(output_distribution='normal', random_state=42, n_quantiles=min(1000, len(df)))
+        print("   📐 Using QuantileTransformer (maps to normal distribution)")
+    else:
+        raise ValueError(f"Unsupported scaler_type: {scaler_type}. Use 'standard', 'minmax', 'robust', 'power', or 'quantile'.")
+    
+    # Get original statistics
+    original_stats = {}
+    for col in columns:
+        col_data = df[col].to_numpy()
+        original_stats[col] = {
+            'mean': float(np.mean(col_data)),
+            'std': float(np.std(col_data)),
+            'min': float(np.min(col_data)),
+            'max': float(np.max(col_data)),
+            'median': float(np.median(col_data))
+        }
+    
+    # Fit and transform
+    print(f"📏 Applying {scaler_type} scaling to {len(columns)} columns...")
+    scaled_data = scaler.fit_transform(df[columns].to_numpy())
+    
+    # Create scaled dataframe
+    df_scaled = df.clone()
+    for i, col in enumerate(columns):
+        df_scaled = df_scaled.with_columns(
+            pl.Series(col, scaled_data[:, i])
+        )
+    
+    # Get new statistics
+    new_stats = {}
+    for i, col in enumerate(columns):
+        new_stats[col] = {
+            'mean': float(np.mean(scaled_data[:, i])),
+            'std': float(np.std(scaled_data[:, i])),
+            'min': float(np.min(scaled_data[:, i])),
+            'max': float(np.max(scaled_data[:, i])),
+            'median': float(np.median(scaled_data[:, i]))
+        }
+    
+    # Save scaled data
+    if output_path:
+        save_dataframe(df_scaled, output_path)
+        print(f"💾 Scaled dataset saved to: {output_path}")
+    
+    # Save scaler
+    if scaler_save_path:
+        os.makedirs(os.path.dirname(scaler_save_path), exist_ok=True)
+        joblib.dump(scaler, scaler_save_path)
+        print(f"💾 Scaler saved to: {scaler_save_path}")
+    
+    return {
+        'status': 'success',
+        'scaler_type': scaler_type,
+        'columns_scaled': columns,
+        'n_columns': len(columns),
+        'original_stats': original_stats,
+        'scaled_stats': new_stats,
+        'output_path': output_path,
+        'scaler_path': scaler_save_path
+    }
+
+
+def split_data_strategically(
+    file_path: str,
+    target_col: Optional[str] = None,
+    split_type: str = "train_test",
+    test_size: float = 0.2,
+    val_size: float = 0.1,
+    stratify: bool = True,
+    time_col: Optional[str] = None,
+    group_col: Optional[str] = None,
+    random_state: int = 42,
+    output_dir: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Perform strategic data splitting with multiple options.
+    
+    Args:
+        file_path: Path to dataset
+        target_col: Target column (for stratification)
+        split_type: Split strategy:
+            - 'train_test': Train/test split
+            - 'train_val_test': Train/validation/test split
+            - 'time_based': Time-based split (requires time_col)
+            - 'group_based': Group-based split (requires group_col, prevents leakage)
+        test_size: Test set proportion
+        val_size: Validation set proportion (for train_val_test)
+        stratify: Whether to stratify by target
+        time_col: Column to use for time-based splitting
+        group_col: Column to use for group-based splitting
+        random_state: Random seed
+        output_dir: Directory to save split datasets
+        
+    Returns:
+        Dictionary with split information and file paths
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    if target_col:
+        validate_column_exists(df, target_col)
+    
+    n_samples = len(df)
+    
+    # Time-based split
+    if split_type == "time_based":
+        if not time_col:
+            raise ValueError("time_col is required for time_based split")
+        validate_column_exists(df, time_col)
+        
+        # Sort by time
+        df = df.sort(time_col)
+        
+        # Calculate split points
+        test_idx = int(n_samples * (1 - test_size))
+        
+        if output_dir:
+            train_df = df[:test_idx]
+            test_df = df[test_idx:]
+            
+            os.makedirs(output_dir, exist_ok=True)
+            train_path = os.path.join(output_dir, "train.csv")
+            test_path = os.path.join(output_dir, "test.csv")
+            
+            save_dataframe(train_df, train_path)
+            save_dataframe(test_df, test_path)
+            
+            print(f"✅ Time-based split: train={len(train_df)}, test={len(test_df)}")
+            
+            return {
+                'status': 'success',
+                'split_type': 'time_based',
+                'train_size': len(train_df),
+                'test_size': len(test_df),
+                'train_path': train_path,
+                'test_path': test_path,
+                'time_column': time_col
+            }
+    
+    # Group-based split
+    elif split_type == "group_based":
+        if not group_col:
+            raise ValueError("group_col is required for group_based split")
+        validate_column_exists(df, group_col)
+        
+        # Get unique groups
+        unique_groups = df[group_col].unique().to_list()
+        n_groups = len(unique_groups)
+        
+        # Split groups
+        np.random.seed(random_state)
+        np.random.shuffle(unique_groups)
+        
+        test_n_groups = max(1, int(n_groups * test_size))
+        test_groups = unique_groups[:test_n_groups]
+        train_groups = unique_groups[test_n_groups:]
+        
+        train_df = df.filter(pl.col(group_col).is_in(train_groups))
+        test_df = df.filter(pl.col(group_col).is_in(test_groups))
+        
+        if output_dir:
+            os.makedirs(output_dir, exist_ok=True)
+            train_path = os.path.join(output_dir, "train.csv")
+            test_path = os.path.join(output_dir, "test.csv")
+            
+            save_dataframe(train_df, train_path)
+            save_dataframe(test_df, test_path)
+            
+            print(f"✅ Group-based split: train={len(train_df)}, test={len(test_df)}")
+            
+            return {
+                'status': 'success',
+                'split_type': 'group_based',
+                'train_size': len(train_df),
+                'test_size': len(test_df),
+                'train_groups': len(train_groups),
+                'test_groups': len(test_groups),
+                'train_path': train_path,
+                'test_path': test_path,
+                'group_column': group_col
+            }
+    
+    # Standard train/test split
+    elif split_type == "train_test":
+        X, y = split_features_target(df, target_col) if target_col else (df.to_numpy(), None)
+        
+        stratify_y = y if (stratify and target_col and len(np.unique(y)) < 20) else None
+        
+        if target_col:
+            X_train, X_test, y_train, y_test = train_test_split(
+                X, y, test_size=test_size, random_state=random_state, stratify=stratify_y
+            )
+            
+            # Reconstruct dataframes
+            feature_cols = [col for col in df.columns if col != target_col]
+            train_data = {col: X_train[:, i] for i, col in enumerate(feature_cols)}
+            train_data[target_col] = y_train
+            train_df = pl.DataFrame(train_data)
+            
+            test_data = {col: X_test[:, i] for i, col in enumerate(feature_cols)}
+            test_data[target_col] = y_test
+            test_df = pl.DataFrame(test_data)
+        else:
+            indices = np.arange(len(df))
+            train_idx, test_idx = train_test_split(
+                indices, test_size=test_size, random_state=random_state
+            )
+            train_df = df[train_idx]
+            test_df = df[test_idx]
+        
+        if output_dir:
+            os.makedirs(output_dir, exist_ok=True)
+            train_path = os.path.join(output_dir, "train.csv")
+            test_path = os.path.join(output_dir, "test.csv")
+            
+            save_dataframe(train_df, train_path)
+            save_dataframe(test_df, test_path)
+            
+            print(f"✅ Train/test split: train={len(train_df)}, test={len(test_df)}")
+            
+            return {
+                'status': 'success',
+                'split_type': 'train_test',
+                'train_size': len(train_df),
+                'test_size': len(test_df),
+                'stratified': bool(stratify_y is not None),
+                'train_path': train_path,
+                'test_path': test_path
+            }
+    
+    # Train/val/test split
+    elif split_type == "train_val_test":
+        X, y = split_features_target(df, target_col) if target_col else (df.to_numpy(), None)
+        
+        stratify_y = y if (stratify and target_col and len(np.unique(y)) < 20) else None
+        
+        # First split: train+val vs test
+        if target_col:
+            X_temp, X_test, y_temp, y_test = train_test_split(
+                X, y, test_size=test_size, random_state=random_state, stratify=stratify_y
+            )
+            
+            # Second split: train vs val
+            val_ratio = val_size / (1 - test_size)
+            stratify_temp = y_temp if stratify_y is not None else None
+            X_train, X_val, y_train, y_val = train_test_split(
+                X_temp, y_temp, test_size=val_ratio, random_state=random_state, stratify=stratify_temp
+            )
+            
+            # Reconstruct dataframes
+            feature_cols = [col for col in df.columns if col != target_col]
+            
+            train_data = {col: X_train[:, i] for i, col in enumerate(feature_cols)}
+            train_data[target_col] = y_train
+            train_df = pl.DataFrame(train_data)
+            
+            val_data = {col: X_val[:, i] for i, col in enumerate(feature_cols)}
+            val_data[target_col] = y_val
+            val_df = pl.DataFrame(val_data)
+            
+            test_data = {col: X_test[:, i] for i, col in enumerate(feature_cols)}
+            test_data[target_col] = y_test
+            test_df = pl.DataFrame(test_data)
+        else:
+            indices = np.arange(len(df))
+            temp_idx, test_idx = train_test_split(
+                indices, test_size=test_size, random_state=random_state
+            )
+            val_ratio = val_size / (1 - test_size)
+            train_idx, val_idx = train_test_split(
+                temp_idx, test_size=val_ratio, random_state=random_state
+            )
+            
+            train_df = df[train_idx]
+            val_df = df[val_idx]
+            test_df = df[test_idx]
+        
+        if output_dir:
+            os.makedirs(output_dir, exist_ok=True)
+            train_path = os.path.join(output_dir, "train.csv")
+            val_path = os.path.join(output_dir, "val.csv")
+            test_path = os.path.join(output_dir, "test.csv")
+            
+            save_dataframe(train_df, train_path)
+            save_dataframe(val_df, val_path)
+            save_dataframe(test_df, test_path)
+            
+            print(f"✅ Train/val/test split: train={len(train_df)}, val={len(val_df)}, test={len(test_df)}")
+            
+            return {
+                'status': 'success',
+                'split_type': 'train_val_test',
+                'train_size': len(train_df),
+                'val_size': len(val_df),
+                'test_size': len(test_df),
+                'stratified': bool(stratify_y is not None),
+                'train_path': train_path,
+                'val_path': val_path,
+                'test_path': test_path
+            }
+    
+    else:
+        raise ValueError(f"Unsupported split_type: {split_type}")