Source code for cursus.steps.scripts.feature_selection

#!/usr/bin/env python3
"""
Feature Selection Script

A comprehensive feature selection script that implements multiple statistical and
machine learning-based feature selection methods. Follows the cursus framework's
testability patterns and provides a self-contained solution.

Author: Cursus Framework
Date: 2025-10-25
"""

import os
import sys
import argparse
import json
import logging
import traceback
import time
import shutil
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple, Union
from subprocess import check_call

# ============================================================================
# PACKAGE INSTALLATION CONFIGURATION
# ============================================================================

# Control which PyPI source to use via environment variable
# Set USE_SECURE_PYPI=true to use secure CodeArtifact PyPI
# Set USE_SECURE_PYPI=false or leave unset to use public PyPI
USE_SECURE_PYPI = os.environ.get("USE_SECURE_PYPI", "false").lower() == "true"

# Logging setup for installation
logging.basicConfig(level=logging.INFO)
logger_install = logging.getLogger(__name__)


def _get_secure_pypi_access_token() -> str:
    """
    Get CodeArtifact access token for secure PyPI.

    Returns:
        str: Authorization token for CodeArtifact

    Raises:
        Exception: If token retrieval fails
    """
    import boto3

    try:
        os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional"
        sts = boto3.client("sts", region_name="us-east-1")
        caller_identity = sts.get_caller_identity()
        assumed_role_object = sts.assume_role(
            RoleArn="arn:aws:iam::675292366480:role/SecurePyPIReadRole_"
            + caller_identity["Account"],
            RoleSessionName="SecurePypiReadRole",
        )
        credentials = assumed_role_object["Credentials"]
        code_artifact_client = boto3.client(
            "codeartifact",
            aws_access_key_id=credentials["AccessKeyId"],
            aws_secret_access_key=credentials["SecretAccessKey"],
            aws_session_token=credentials["SessionToken"],
            region_name="us-west-2",
        )
        token = code_artifact_client.get_authorization_token(
            domain="amazon", domainOwner="149122183214"
        )["authorizationToken"]

        logger_install.info("Successfully retrieved secure PyPI access token")
        return token

    except Exception as e:
        logger_install.error(f"Failed to retrieve secure PyPI access token: {e}")
        raise


[docs] def install_packages_from_public_pypi(packages: list) -> None: """ Install packages from standard public PyPI. Args: packages: List of package specifications (e.g., ["pandas==1.5.0", "numpy"]) """ logger_install.info(f"Installing {len(packages)} packages from public PyPI") logger_install.info(f"Packages: {packages}") try: check_call([sys.executable, "-m", "pip", "install", *packages]) logger_install.info("✓ Successfully installed packages from public PyPI") except Exception as e: logger_install.error(f"✗ Failed to install packages from public PyPI: {e}") raise
[docs] def install_packages_from_secure_pypi(packages: list) -> None: """ Install packages from secure CodeArtifact PyPI. Args: packages: List of package specifications (e.g., ["pandas==1.5.0", "numpy"]) """ logger_install.info(f"Installing {len(packages)} packages from secure PyPI") logger_install.info(f"Packages: {packages}") try: token = _get_secure_pypi_access_token() index_url = f"https://aws:{token}@amazon-149122183214.d.codeartifact.us-west-2.amazonaws.com/pypi/secure-pypi/simple/" check_call( [ sys.executable, "-m", "pip", "install", "--index-url", index_url, *packages, ] ) logger_install.info("✓ Successfully installed packages from secure PyPI") except Exception as e: logger_install.error(f"✗ Failed to install packages from secure PyPI: {e}") raise
[docs] def install_packages(packages: list, use_secure: bool = USE_SECURE_PYPI) -> None: """ Install packages from PyPI source based on configuration. This is the main installation function that delegates to either public or secure PyPI based on the USE_SECURE_PYPI environment variable. Args: packages: List of package specifications (e.g., ["pandas==1.5.0", "numpy"]) use_secure: If True, use secure CodeArtifact PyPI; if False, use public PyPI. Defaults to USE_SECURE_PYPI environment variable. Environment Variables: USE_SECURE_PYPI: Set to "true" to use secure PyPI, "false" for public PyPI Example: # Install from public PyPI (default) install_packages(["xgboost==1.7.3"]) # Install from secure PyPI os.environ["USE_SECURE_PYPI"] = "true" install_packages(["xgboost==1.7.3"]) """ logger_install.info("=" * 70) logger_install.info("PACKAGE INSTALLATION") logger_install.info("=" * 70) logger_install.info( f"PyPI Source: {'SECURE (CodeArtifact)' if use_secure else 'PUBLIC'}" ) logger_install.info( f"Environment Variable USE_SECURE_PYPI: {os.environ.get('USE_SECURE_PYPI', 'not set')}" ) logger_install.info(f"Number of packages: {len(packages)}") logger_install.info("=" * 70) try: if use_secure: install_packages_from_secure_pypi(packages) else: install_packages_from_public_pypi(packages) logger_install.info("=" * 70) logger_install.info("✓ PACKAGE INSTALLATION COMPLETED SUCCESSFULLY") logger_install.info("=" * 70) except Exception as e: logger_install.error("=" * 70) logger_install.error("✗ PACKAGE INSTALLATION FAILED") logger_install.error("=" * 70) raise
# ============================================================================ # INSTALL REQUIRED PACKAGES # ============================================================================ # Define required packages for this script # xgboost 1.7.3 is compatible with scikit-learn 1.2.1 (sklearn framework version 1.2-1) required_packages = [ "xgboost==1.7.3", ] # Install packages using unified installation function install_packages(required_packages) print("***********************Package Installation Complete*********************") # ============================================================================ # MAIN SCRIPT IMPORTS # ============================================================================ import pandas as pd import numpy as np from sklearn.feature_selection import ( mutual_info_classif, mutual_info_regression, SelectKBest, chi2, f_classif, f_regression, RFE, ) from sklearn.ensemble import ( RandomForestClassifier, RandomForestRegressor, ExtraTreesClassifier, ExtraTreesRegressor, ) from sklearn.svm import SVC, SVR from sklearn.linear_model import LogisticRegression, LinearRegression, Lasso, LassoCV from sklearn.preprocessing import StandardScaler from sklearn.inspection import permutation_importance # ------------------------------------------------------------------------- # Logging setup - Updated for CloudWatch compatibility # -------------------------------------------------------------------------
[docs] def setup_logging() -> logging.Logger: """Configure logging for CloudWatch compatibility""" # Remove any existing handlers root = logging.getLogger() if root.handlers: for handler in root.handlers: root.removeHandler(handler) # Configure the root logger logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", force=True, handlers=[ # StreamHandler with stdout for CloudWatch logging.StreamHandler(sys.stdout) ], ) # Configure our module's logger logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) logger.propagate = True # Allow propagation to root logger # Force flush stdout sys.stdout.flush() return logger
# Initialize logger logger = setup_logging() # ------------------------------------------------------------------------- # Container path constants # ------------------------------------------------------------------------- CONTAINER_PATHS = { "INPUT_DATA": "/opt/ml/processing/input", "OUTPUT_DATA": "/opt/ml/processing/output", } # ------------------------------------------------------------------------- # File I/O Helper Functions with Format Preservation # ------------------------------------------------------------------------- def _detect_file_format(split_dir: Path, split_name: str) -> tuple: """ Detect the format of processed data file. Returns: Tuple of (file_path, format) where format is 'csv', 'tsv', or 'parquet' """ # Try different formats in order of preference formats = [ (f"{split_name}_processed_data.csv", "csv"), (f"{split_name}_processed_data.tsv", "tsv"), (f"{split_name}_processed_data.parquet", "parquet"), ] for filename, fmt in formats: file_path = split_dir / filename if file_path.exists(): return file_path, fmt raise RuntimeError( f"No processed data file found in {split_dir}. " f"Looked for: {[f[0] for f in formats]}" ) # ------------------------------------------------------------------------- # Artifact Management Functions # -------------------------------------------------------------------------
[docs] def copy_existing_artifacts(src_dir: str, dst_dir: str) -> None: """ Copy all existing model artifacts from previous processing steps. This enables the parameter accumulator pattern where each step: 1. Copies artifacts from previous steps 2. Adds its own artifacts 3. Passes all artifacts to the next step Args: src_dir: Source directory containing existing artifacts dst_dir: Destination directory to copy artifacts to """ if not src_dir or not os.path.exists(src_dir): logger.info(f"No existing artifacts to copy from {src_dir}") return os.makedirs(dst_dir, exist_ok=True) copied_count = 0 for filename in os.listdir(src_dir): src_file = os.path.join(src_dir, filename) dst_file = os.path.join(dst_dir, filename) if os.path.isfile(src_file): shutil.copy2(src_file, dst_file) copied_count += 1 logger.info(f" Copied existing artifact: {filename}") logger.info(f"✓ Copied {copied_count} existing artifact(s) to {dst_dir}")
[docs] def load_selected_features(model_artifacts_dir: str) -> List[str]: """ Load pre-computed selected features from training job. Args: model_artifacts_dir: Directory containing selected_features.json Returns: List of selected feature names """ selected_features_file = os.path.join(model_artifacts_dir, "selected_features.json") if not os.path.exists(selected_features_file): raise FileNotFoundError( f"Selected features file not found: {selected_features_file}" ) try: with open(selected_features_file, "r") as f: data = json.load(f) selected_features = data.get("selected_features", []) logger.info( f"Loaded {len(selected_features)} pre-selected features from {selected_features_file}" ) logger.info(f"Selected features: {selected_features}") return selected_features except Exception as e: logger.error(f"Error loading selected features: {e}") raise
# ------------------------------------------------------------------------- # Data Loading Functions # -------------------------------------------------------------------------
[docs] def load_single_split_data( input_data_dir: str, job_type: str ) -> Dict[str, pd.DataFrame]: """ Load single split data for non-training job types with format detection. Args: input_data_dir: Directory containing job_type subdirectory job_type: Type of job (validation, testing, etc.) Returns: Dictionary with single split DataFrame and format metadata """ split_dir = Path(input_data_dir) / job_type if not split_dir.exists(): raise FileNotFoundError(f"Split directory not found: {split_dir}") try: # Detect format and read file file_path, detected_format = _detect_file_format(split_dir, job_type) if detected_format == "csv": df = pd.read_csv(file_path) elif detected_format == "tsv": df = pd.read_csv(file_path, sep="\t") elif detected_format == "parquet": df = pd.read_parquet(file_path) else: raise RuntimeError(f"Unsupported format: {detected_format}") logger.info(f"Loaded {job_type} split (format={detected_format}): {df.shape}") return {job_type: df, "_format": detected_format} except Exception as e: logger.error(f"Error loading {job_type} data: {e}") raise
[docs] def load_preprocessed_data(input_data_dir: str) -> Dict[str, pd.DataFrame]: """ Load train/val/test splits from tabular preprocessing output with format detection. Args: input_data_dir: Directory containing train/val/test subdirectories Returns: Dictionary with 'train', 'val', 'test' DataFrames and format metadata """ splits = {} detected_format = None for split_name in ["train", "val", "test"]: split_dir = Path(input_data_dir) / split_name if not split_dir.exists(): logger.warning(f"Split directory not found: {split_dir}") continue try: # Detect format and read file file_path, fmt = _detect_file_format(split_dir, split_name) # Store format from first split (they should all match) if detected_format is None: detected_format = fmt if fmt == "csv": df = pd.read_csv(file_path) elif fmt == "tsv": df = pd.read_csv(file_path, sep="\t") elif fmt == "parquet": df = pd.read_parquet(file_path) else: raise RuntimeError(f"Unsupported format: {fmt}") splits[split_name] = df logger.info(f"Loaded {split_name} split (format={fmt}): {df.shape}") except Exception as e: logger.error(f"Error loading {split_name} data: {e}") raise if not splits: raise FileNotFoundError("No valid data splits found in input directory") # Store detected format for use in saving splits["_format"] = detected_format return splits
[docs] def save_selected_data( splits: Dict[str, pd.DataFrame], selected_features: List[str], target_variable: str, output_dir: str, ) -> None: """ Save feature-selected splits preserving input format. Args: splits: Dictionary of DataFrames by split name (includes "_format" key) selected_features: List of selected feature names target_variable: Target column name output_dir: Output directory path """ # Ensure output directory exists os.makedirs(output_dir, exist_ok=True) # Extract format from splits dictionary output_format = splits.get("_format", "csv") # Default to CSV if not found # Features to keep (selected features + target) columns_to_keep = selected_features + [target_variable] for split_name, df in splits.items(): # Skip the format metadata key if split_name == "_format": continue # Create split subdirectory split_dir = os.path.join(output_dir, split_name) os.makedirs(split_dir, exist_ok=True) # Verify all required columns exist missing_cols = [col for col in columns_to_keep if col not in df.columns] if missing_cols: logger.error(f"Missing columns in {split_name} split: {missing_cols}") raise ValueError(f"Missing columns in {split_name} split: {missing_cols}") # Filter to selected features + target selected_df = df[columns_to_keep].copy() # Save in detected format if output_format == "csv": output_file = os.path.join(split_dir, f"{split_name}_processed_data.csv") selected_df.to_csv(output_file, index=False) elif output_format == "tsv": output_file = os.path.join(split_dir, f"{split_name}_processed_data.tsv") selected_df.to_csv(output_file, sep="\t", index=False) elif output_format == "parquet": output_file = os.path.join( split_dir, f"{split_name}_processed_data.parquet" ) selected_df.to_parquet(output_file, index=False) else: raise RuntimeError(f"Unsupported output format: {output_format}") logger.info( f"Saved {split_name} split with {len(selected_features)} features (format={output_format}): {output_file}" )
# ------------------------------------------------------------------------- # Feature Selection Methods - Statistical # -------------------------------------------------------------------------
[docs] def variance_threshold_selection( X: pd.DataFrame, threshold: float = 0.01 ) -> Dict[str, Any]: """ Remove features with low variance. Args: X: Feature matrix threshold: Variance threshold Returns: Dictionary with selected features and scores """ logger.info(f"Applying variance threshold selection with threshold={threshold}") variances = X.var() selected_features = variances[variances > threshold].index.tolist() logger.info( f"Variance threshold selected {len(selected_features)} out of {len(X.columns)} features" ) return { "method": "variance_threshold", "selected_features": selected_features, "scores": variances.to_dict(), "threshold": threshold, "n_selected": len(selected_features), }
[docs] def correlation_based_selection( X: pd.DataFrame, y: pd.Series, threshold: float = 0.95 ) -> Dict[str, Any]: """ Remove highly correlated features, keeping those with higher target correlation. Args: X: Feature matrix y: Target variable threshold: Correlation threshold Returns: Dictionary with selected features and scores """ logger.info(f"Applying correlation-based selection with threshold={threshold}") # Compute feature correlations corr_matrix = X.corr().abs() # Compute target correlations target_corr = X.corrwith(y).abs() # Find highly correlated pairs upper_tri = corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(bool)) high_corr_pairs = [ (col, row) for col in upper_tri.columns for row in upper_tri.index if not pd.isna(upper_tri.loc[row, col]) and upper_tri.loc[row, col] > threshold ] # Remove features with lower target correlation from each pair features_to_remove = set() for feat1, feat2 in high_corr_pairs: if target_corr[feat1] > target_corr[feat2]: features_to_remove.add(feat2) else: features_to_remove.add(feat1) selected_features = [f for f in X.columns if f not in features_to_remove] logger.info( f"Correlation-based selection removed {len(features_to_remove)} features, " f"selected {len(selected_features)} out of {len(X.columns)} features" ) return { "method": "correlation_based", "selected_features": selected_features, "scores": target_corr.to_dict(), "removed_features": list(features_to_remove), "high_corr_pairs": high_corr_pairs, "threshold": threshold, "n_selected": len(selected_features), }
[docs] def mutual_info_selection( X: pd.DataFrame, y: pd.Series, k: int = 10, random_state: int = 42 ) -> Dict[str, Any]: """ Select features based on mutual information with target. Args: X: Feature matrix y: Target variable k: Number of features to select random_state: Random seed Returns: Dictionary with selected features and scores """ logger.info(f"Applying mutual information selection with k={k}") # Determine if classification or regression is_classification = len(y.unique()) < 20 and y.dtype in [ "object", "category", "int64", ] try: if is_classification: mi_scores = mutual_info_classif(X, y, random_state=random_state) else: mi_scores = mutual_info_regression(X, y, random_state=random_state) # Select top k features selector = SelectKBest( score_func=lambda X, y: mi_scores, k=min(k, len(X.columns)) ) selector.fit(X, y) selected_features = X.columns[selector.get_support()].tolist() scores = dict(zip(X.columns, mi_scores)) logger.info( f"Mutual information selected {len(selected_features)} features " f"(classification={is_classification})" ) return { "method": "mutual_information", "selected_features": selected_features, "scores": scores, "k": k, "n_selected": len(selected_features), "is_classification": is_classification, } except Exception as e: logger.error(f"Error in mutual information selection: {e}") # Return empty result on error return { "method": "mutual_information", "selected_features": [], "scores": {}, "k": k, "n_selected": 0, "is_classification": is_classification, "error": str(e), }
[docs] def chi2_selection(X: pd.DataFrame, y: pd.Series, k: int = 10) -> Dict[str, Any]: """ Select features using chi-square test (for non-negative features). Args: X: Feature matrix (non-negative values) y: Target variable k: Number of features to select Returns: Dictionary with selected features and scores """ logger.info(f"Applying chi-square selection with k={k}") try: # Ensure non-negative values X_nonneg = X.copy() X_nonneg[X_nonneg < 0] = 0 # Apply chi-square test selector = SelectKBest(score_func=chi2, k=min(k, len(X.columns))) selector.fit(X_nonneg, y) selected_features = X.columns[selector.get_support()].tolist() scores = dict(zip(X.columns, selector.scores_)) logger.info(f"Chi-square selected {len(selected_features)} features") return { "method": "chi2", "selected_features": selected_features, "scores": scores, "k": k, "n_selected": len(selected_features), } except Exception as e: logger.error(f"Error in chi-square selection: {e}") return { "method": "chi2", "selected_features": [], "scores": {}, "k": k, "n_selected": 0, "error": str(e), }
[docs] def f_classif_selection(X: pd.DataFrame, y: pd.Series, k: int = 10) -> Dict[str, Any]: """ Select features using ANOVA F-test. Args: X: Feature matrix y: Target variable k: Number of features to select Returns: Dictionary with selected features and scores """ logger.info(f"Applying F-test selection with k={k}") try: # Determine if classification or regression is_classification = len(y.unique()) < 20 and y.dtype in [ "object", "category", "int64", ] score_func = f_classif if is_classification else f_regression selector = SelectKBest(score_func=score_func, k=min(k, len(X.columns))) selector.fit(X, y) selected_features = X.columns[selector.get_support()].tolist() scores = dict(zip(X.columns, selector.scores_)) logger.info( f"F-test selected {len(selected_features)} features " f"(classification={is_classification})" ) return { "method": "f_test", "selected_features": selected_features, "scores": scores, "k": k, "n_selected": len(selected_features), "is_classification": is_classification, } except Exception as e: logger.error(f"Error in F-test selection: {e}") return { "method": "f_test", "selected_features": [], "scores": {}, "k": k, "n_selected": 0, "error": str(e), }
# ------------------------------------------------------------------------- # Feature Selection Methods - ML-Based # -------------------------------------------------------------------------
[docs] def rfe_selection( X: pd.DataFrame, y: pd.Series, estimator_type: str = "rf", n_features: int = 10, random_state: int = 42, ) -> Dict[str, Any]: """ Recursive Feature Elimination with various estimators. Args: X: Feature matrix y: Target variable estimator_type: Type of estimator ('rf', 'svm', 'linear') n_features: Number of features to select random_state: Random seed Returns: Dictionary with selected features and scores """ logger.info( f"Applying RFE selection with estimator={estimator_type}, n_features={n_features}" ) try: # Determine if classification or regression is_classification = len(y.unique()) < 20 and y.dtype in [ "object", "category", "int64", ] # Select estimator if estimator_type == "rf": estimator = ( RandomForestClassifier(n_estimators=50, random_state=random_state) if is_classification else RandomForestRegressor(n_estimators=50, random_state=random_state) ) elif estimator_type == "svm": estimator = ( SVC(kernel="linear", random_state=random_state) if is_classification else SVR(kernel="linear") ) elif estimator_type == "linear": estimator = ( LogisticRegression(random_state=random_state, max_iter=1000) if is_classification else LinearRegression() ) else: raise ValueError(f"Unknown estimator type: {estimator_type}") # Apply RFE selector = RFE( estimator=estimator, n_features_to_select=min(n_features, len(X.columns)) ) selector.fit(X, y) selected_features = X.columns[selector.get_support()].tolist() rankings = dict(zip(X.columns, selector.ranking_)) logger.info(f"RFE selected {len(selected_features)} features") return { "method": f"rfe_{estimator_type}", "selected_features": selected_features, "scores": { f: 1.0 / rank for f, rank in rankings.items() }, # Convert ranking to score "rankings": rankings, "estimator_type": estimator_type, "n_features": n_features, "n_selected": len(selected_features), "is_classification": is_classification, } except Exception as e: logger.error(f"Error in RFE selection: {e}") return { "method": f"rfe_{estimator_type}", "selected_features": [], "scores": {}, "rankings": {}, "estimator_type": estimator_type, "n_features": n_features, "n_selected": 0, "is_classification": is_classification, "error": str(e), }
[docs] def feature_importance_selection( X: pd.DataFrame, y: pd.Series, method: str = "random_forest", n_features: int = 10, random_state: int = 42, ) -> Dict[str, Any]: """ Select features based on model feature importance. Args: X: Feature matrix y: Target variable method: Method for importance ('random_forest', 'xgboost', 'extra_trees') n_features: Number of features to select random_state: Random seed Returns: Dictionary with selected features and scores """ logger.info( f"Applying feature importance selection with method={method}, n_features={n_features}" ) try: # Determine if classification or regression is_classification = len(y.unique()) < 20 and y.dtype in [ "object", "category", "int64", ] # Select model if method == "random_forest": model = ( RandomForestClassifier(n_estimators=100, random_state=random_state) if is_classification else RandomForestRegressor(n_estimators=100, random_state=random_state) ) elif method == "extra_trees": model = ( ExtraTreesClassifier(n_estimators=100, random_state=random_state) if is_classification else ExtraTreesRegressor(n_estimators=100, random_state=random_state) ) elif method == "xgboost": try: import xgboost as xgb model = ( xgb.XGBClassifier(n_estimators=100, random_state=random_state) if is_classification else xgb.XGBRegressor(n_estimators=100, random_state=random_state) ) except ImportError: logger.warning("XGBoost not available, falling back to random forest") model = ( RandomForestClassifier(n_estimators=100, random_state=random_state) if is_classification else RandomForestRegressor( n_estimators=100, random_state=random_state ) ) method = "random_forest" else: raise ValueError(f"Unknown method: {method}") # Fit model and get importance model.fit(X, y) importances = model.feature_importances_ # Select top features feature_importance_pairs = list(zip(X.columns, importances)) feature_importance_pairs.sort(key=lambda x: x[1], reverse=True) selected_features = [ pair[0] for pair in feature_importance_pairs[: min(n_features, len(X.columns))] ] scores = dict(feature_importance_pairs) logger.info(f"Feature importance selected {len(selected_features)} features") return { "method": f"importance_{method}", "selected_features": selected_features, "scores": scores, "method_used": method, "n_features": n_features, "n_selected": len(selected_features), "is_classification": is_classification, } except Exception as e: logger.error(f"Error in feature importance selection: {e}") return { "method": f"importance_{method}", "selected_features": [], "scores": {}, "method_used": method, "n_features": n_features, "n_selected": 0, "is_classification": is_classification, "error": str(e), }
[docs] def lasso_selection( X: pd.DataFrame, y: pd.Series, alpha: Union[float, str] = 0.01, random_state: int = 42, ) -> Dict[str, Any]: """ Select features using LASSO regularization. Args: X: Feature matrix y: Target variable alpha: Regularization strength random_state: Random seed Returns: Dictionary with selected features and scores """ logger.info(f"Applying LASSO selection with alpha={alpha}") try: # Determine if classification or regression is_classification = len(y.unique()) < 20 and y.dtype in [ "object", "category", "int64", ] # Standardize features scaler = StandardScaler() X_scaled = pd.DataFrame( scaler.fit_transform(X), columns=X.columns, index=X.index ) if is_classification: # Use Logistic Regression with L1 penalty model = LogisticRegression( penalty="l1", C=1 / alpha, solver="liblinear", random_state=random_state, max_iter=1000, ) model.fit(X_scaled, y) coefficients = model.coef_[0] if len(model.coef_.shape) > 1 else model.coef_ else: # Use LASSO regression if alpha == "auto": model = LassoCV(cv=5, random_state=random_state) else: model = Lasso(alpha=alpha, random_state=random_state) model.fit(X_scaled, y) coefficients = model.coef_ # Select features with non-zero coefficients selected_features = [ X.columns[i] for i, coef in enumerate(coefficients) if abs(coef) > 1e-6 ] scores = dict(zip(X.columns, np.abs(coefficients))) logger.info(f"LASSO selected {len(selected_features)} features") return { "method": "lasso", "selected_features": selected_features, "scores": scores, "alpha": alpha if alpha != "auto" else getattr(model, "alpha_", alpha), "n_selected": len(selected_features), "is_classification": is_classification, } except Exception as e: logger.error(f"Error in LASSO selection: {e}") return { "method": "lasso", "selected_features": [], "scores": {}, "alpha": alpha, "n_selected": 0, "is_classification": is_classification, "error": str(e), }
[docs] def permutation_importance_selection( X: pd.DataFrame, y: pd.Series, estimator_type: str = "rf", n_features: int = 10, random_state: int = 42, ) -> Dict[str, Any]: """ Select features using permutation importance. Args: X: Feature matrix y: Target variable estimator_type: Type of estimator n_features: Number of features to select random_state: Random seed Returns: Dictionary with selected features and scores """ logger.info( f"Applying permutation importance selection with estimator={estimator_type}, n_features={n_features}" ) try: # Determine if classification or regression is_classification = len(y.unique()) < 20 and y.dtype in [ "object", "category", "int64", ] # Select estimator if estimator_type == "rf": estimator = ( RandomForestClassifier(n_estimators=50, random_state=random_state) if is_classification else RandomForestRegressor(n_estimators=50, random_state=random_state) ) else: # Default to random forest estimator = ( RandomForestClassifier(n_estimators=50, random_state=random_state) if is_classification else RandomForestRegressor(n_estimators=50, random_state=random_state) ) # Fit estimator estimator.fit(X, y) # Compute permutation importance perm_importance = permutation_importance( estimator, X, y, n_repeats=5, random_state=random_state ) # Select top features feature_importance_pairs = list( zip(X.columns, perm_importance.importances_mean) ) feature_importance_pairs.sort(key=lambda x: x[1], reverse=True) selected_features = [ pair[0] for pair in feature_importance_pairs[: min(n_features, len(X.columns))] ] scores = dict(feature_importance_pairs) logger.info( f"Permutation importance selected {len(selected_features)} features" ) return { "method": f"permutation_{estimator_type}", "selected_features": selected_features, "scores": scores, "estimator_type": estimator_type, "n_features": n_features, "n_selected": len(selected_features), "is_classification": is_classification, } except Exception as e: logger.error(f"Error in permutation importance selection: {e}") return { "method": f"permutation_{estimator_type}", "selected_features": [], "scores": {}, "estimator_type": estimator_type, "n_features": n_features, "n_selected": 0, "is_classification": is_classification, "error": str(e), }
# ------------------------------------------------------------------------- # Ensemble Selection Logic # -------------------------------------------------------------------------
[docs] def combine_selection_results( method_results: List[Dict[str, Any]], combination_strategy: str = "voting", final_k: int = 10, ) -> Dict[str, Any]: """ Combine results from multiple feature selection methods. Args: method_results: List of results from different methods combination_strategy: Strategy for combination ('voting', 'ranking', 'scoring') final_k: Final number of features to select Returns: Combined selection results """ if not method_results: return {"selected_features": [], "scores": {}, "method_contributions": {}} # Collect all features and their scores from all methods all_features = set() method_scores = {} method_selections = {} for result in method_results: method_name = result["method"] selected_features = result["selected_features"] scores = result["scores"] all_features.update(selected_features) method_scores[method_name] = scores method_selections[method_name] = set(selected_features) all_features = list(all_features) if combination_strategy == "voting": # Count how many methods selected each feature feature_votes = {} for feature in all_features: votes = sum( 1 for selections in method_selections.values() if feature in selections ) feature_votes[feature] = votes # Sort by votes and select top k sorted_features = sorted( feature_votes.items(), key=lambda x: x[1], reverse=True ) selected_features = [f for f, _ in sorted_features[:final_k]] combined_scores = feature_votes elif combination_strategy == "ranking": # Combine rankings from all methods feature_ranks = {f: [] for f in all_features} for method_name, scores in method_scores.items(): # Convert scores to rankings sorted_scores = sorted(scores.items(), key=lambda x: x[1], reverse=True) for rank, (feature, score) in enumerate(sorted_scores): if feature in all_features: feature_ranks[feature].append(rank + 1) # Average rankings (lower is better) avg_ranks = {} for feature, ranks in feature_ranks.items(): if ranks: avg_ranks[feature] = np.mean(ranks) else: avg_ranks[feature] = len(all_features) # Worst possible rank # Sort by average rank and select top k sorted_features = sorted(avg_ranks.items(), key=lambda x: x[1]) selected_features = [f for f, _ in sorted_features[:final_k]] combined_scores = { f: 1.0 / rank for f, rank in avg_ranks.items() } # Convert to scores elif combination_strategy == "scoring": # Normalize and combine scores from all methods normalized_scores = {} for method_name, scores in method_scores.items(): # Normalize scores to 0-1 range score_values = list(scores.values()) if score_values: min_score = min(score_values) max_score = max(score_values) score_range = max_score - min_score if score_range > 0: normalized_scores[method_name] = { f: (s - min_score) / score_range for f, s in scores.items() } else: normalized_scores[method_name] = {f: 1.0 for f in scores.keys()} # Average normalized scores combined_scores = {} for feature in all_features: scores_for_feature = [] for method_scores in normalized_scores.values(): if feature in method_scores: scores_for_feature.append(method_scores[feature]) if scores_for_feature: combined_scores[feature] = np.mean(scores_for_feature) else: combined_scores[feature] = 0.0 # Sort by combined score and select top k sorted_features = sorted( combined_scores.items(), key=lambda x: x[1], reverse=True ) selected_features = [f for f, _ in sorted_features[:final_k]] else: raise ValueError(f"Unknown combination strategy: {combination_strategy}") # Calculate method contributions method_contributions = {} for method_name, selections in method_selections.items(): contribution = ( len(set(selected_features) & selections) / len(selected_features) if selected_features else 0 ) method_contributions[method_name] = contribution return { "selected_features": selected_features, "scores": combined_scores, "method_contributions": method_contributions, "combination_strategy": combination_strategy, "n_methods": len(method_results), "n_selected": len(selected_features), }
# ------------------------------------------------------------------------- # Feature Selection Pipeline # -------------------------------------------------------------------------
[docs] def apply_feature_selection_pipeline( splits: Dict[str, pd.DataFrame], target_variable: str, methods: List[str], method_configs: Dict[str, Dict], ) -> Dict[str, Any]: """ Apply feature selection pipeline using training data, then apply to all splits. Args: splits: Dictionary of train/val/test DataFrames target_variable: Target column name methods: List of feature selection methods to apply method_configs: Configuration for each method Returns: Dictionary with selection results and metadata """ # Use training data for feature selection if "train" not in splits: raise ValueError("Training data not found in splits") train_df = splits["train"] # Separate features and target if target_variable not in train_df.columns: raise ValueError( f"Target variable '{target_variable}' not found in training data columns" ) feature_columns = [col for col in train_df.columns if col != target_variable] X_train = train_df[feature_columns] y_train = train_df[target_variable] logger.info(f"Starting feature selection on {len(feature_columns)} features") logger.info(f"Training data shape: {X_train.shape}") logger.info(f"Target variable: {target_variable}") # Apply each selection method method_results = [] for method in methods: logger.info(f"Applying {method} feature selection...") start_time = time.time() try: if method == "variance": result = variance_threshold_selection( X_train, **method_configs.get(method, {}) ) elif method == "correlation": result = correlation_based_selection( X_train, y_train, **method_configs.get(method, {}) ) elif method == "mutual_info": result = mutual_info_selection( X_train, y_train, **method_configs.get(method, {}) ) elif method == "chi2": result = chi2_selection( X_train, y_train, **method_configs.get(method, {}) ) elif method == "f_test": result = f_classif_selection( X_train, y_train, **method_configs.get(method, {}) ) elif method == "rfe": result = rfe_selection( X_train, y_train, **method_configs.get(method, {}) ) elif method == "importance": result = feature_importance_selection( X_train, y_train, **method_configs.get(method, {}) ) elif method == "lasso": result = lasso_selection( X_train, y_train, **method_configs.get(method, {}) ) elif method == "permutation": result = permutation_importance_selection( X_train, y_train, **method_configs.get(method, {}) ) else: logger.warning(f"Unknown method: {method}, skipping...") continue # Add processing time result["processing_time"] = time.time() - start_time method_results.append(result) logger.info( f"{method} selected {result['n_selected']} features in {result['processing_time']:.2f}s" ) except Exception as e: logger.error(f"Error in {method} selection: {e}") # Continue with other methods continue if not method_results: raise RuntimeError("No feature selection methods completed successfully") # Combine results from multiple methods final_k = method_configs.get("final_k", min(50, len(feature_columns) // 2)) combination_strategy = method_configs.get("combination_strategy", "voting") logger.info( f"Combining results using {combination_strategy} strategy, selecting top {final_k} features" ) combined_result = combine_selection_results( method_results, combination_strategy, final_k ) logger.info( f"Final selection: {len(combined_result['selected_features'])} features" ) logger.info(f"Selected features: {combined_result['selected_features']}") return { "selected_features": combined_result["selected_features"], "method_results": method_results, "combined_result": combined_result, "original_features": feature_columns, "target_variable": target_variable, "n_original_features": len(feature_columns), "n_selected_features": len(combined_result["selected_features"]), }
# ------------------------------------------------------------------------- # Output Generation # -------------------------------------------------------------------------
[docs] def save_selection_results( selection_results: Dict[str, Any], model_artifacts_dir: str ) -> None: """ Save feature selection results and metadata to model artifacts directory. Args: selection_results: Results from feature selection pipeline model_artifacts_dir: Directory for all feature selection output files """ # Ensure output directory exists os.makedirs(model_artifacts_dir, exist_ok=True) # Save selected features metadata selected_features_data = { "selected_features": selection_results["selected_features"], "selection_metadata": { "n_original_features": selection_results["n_original_features"], "n_selected_features": selection_results["n_selected_features"], "selection_ratio": selection_results["n_selected_features"] / selection_results["n_original_features"], "methods_used": [ result["method"] for result in selection_results["method_results"] ], "combination_strategy": selection_results["combined_result"][ "combination_strategy" ], "target_variable": selection_results["target_variable"], }, "method_contributions": selection_results["combined_result"][ "method_contributions" ], } with open(os.path.join(model_artifacts_dir, "selected_features.json"), "w") as f: json.dump(selected_features_data, f, indent=2, sort_keys=True) # Save detailed feature scores feature_scores_data = [] for feature in selection_results["original_features"]: row = { "feature_name": feature, "combined_score": selection_results["combined_result"]["scores"].get( feature, 0 ), "selected": feature in selection_results["selected_features"], } # Add scores from individual methods for result in selection_results["method_results"]: method_name = result["method"] row[f"{method_name}_score"] = result["scores"].get(feature, 0) feature_scores_data.append(row) # Sort by combined score feature_scores_data.sort(key=lambda x: x["combined_score"], reverse=True) # Save as CSV feature_scores_df = pd.DataFrame(feature_scores_data) feature_scores_df.to_csv( os.path.join(model_artifacts_dir, "feature_scores.csv"), index=False ) # Save selection summary report (merged into same directory) # Guard against empty combined scores (e.g., all selection methods failed or # returned no features) so aggregation statistics below do not crash. combined_scores_list = list(selection_results["combined_result"]["scores"].values()) if not combined_scores_list: combined_scores_list = [0.0] summary_report = { "selection_summary": { "total_features": selection_results["n_original_features"], "selected_features": selection_results["n_selected_features"], "selection_methods": [ result["method"] for result in selection_results["method_results"] ], "combination_strategy": selection_results["combined_result"][ "combination_strategy" ], "processing_time": sum( result.get("processing_time", 0) for result in selection_results["method_results"] ), "target_variable": selection_results["target_variable"], }, "method_performance": { result["method"]: { "n_selected": result["n_selected"], "processing_time": result.get("processing_time", 0), } for result in selection_results["method_results"] }, "feature_statistics": { "avg_score": float(np.mean(combined_scores_list)), "score_std": float(np.std(combined_scores_list)), "min_score": min(combined_scores_list), "max_score": max(combined_scores_list), }, } with open( os.path.join(model_artifacts_dir, "feature_selection_report.json"), "w" ) as f: json.dump(summary_report, f, indent=2, sort_keys=True) logger.info(f"Saved all feature selection results to {model_artifacts_dir}")
# ------------------------------------------------------------------------- # Main Function # -------------------------------------------------------------------------
[docs] def main( input_paths: Dict[str, str], output_paths: Dict[str, str], environ_vars: Dict[str, str], job_args: argparse.Namespace, ) -> None: """ Main function for feature selection processing. Args: input_paths: Dictionary of input paths with logical names - "input_data": Directory containing train/val/test splits from tabular preprocessing - "model_artifacts_input": Model artifacts from previous steps (standardized) output_paths: Dictionary of output paths with logical names - "processed_data": Directory for feature-selected train/val/test splits (XGBoost input format) - "model_artifacts_output": Model artifacts output for next steps (standardized) environ_vars: Dictionary of environment variables - "FEATURE_SELECTION_METHODS": Comma-separated list of methods - "LABEL_FIELD": Target column name (standard across framework) - "N_FEATURES_TO_SELECT": Number/percentage of features to select - "CORRELATION_THRESHOLD": Threshold for correlation filtering - "VARIANCE_THRESHOLD": Threshold for variance filtering - "RANDOM_STATE": Random seed for reproducibility - "COMBINATION_STRATEGY": Method combination strategy job_args: Command line arguments - job_type: Must be "training" to process all splits """ try: logger.info("====== STARTING FEATURE SELECTION ======") # Extract configuration from environment variables methods_str = environ_vars.get( "FEATURE_SELECTION_METHODS", "variance,correlation,mutual_info,rfe" ) methods = [method.strip() for method in methods_str.split(",")] target_variable = environ_vars.get("LABEL_FIELD") if not target_variable: raise ValueError("LABEL_FIELD environment variable must be set") n_features_to_select = int(environ_vars.get("N_FEATURES_TO_SELECT", 10)) correlation_threshold = float(environ_vars.get("CORRELATION_THRESHOLD", 0.95)) variance_threshold = float(environ_vars.get("VARIANCE_THRESHOLD", 0.01)) random_state = int(environ_vars.get("RANDOM_STATE", 42)) combination_strategy = environ_vars.get("COMBINATION_STRATEGY", "voting") logger.info(f"Configuration:") logger.info(f" Methods: {methods}") logger.info(f" Target variable: {target_variable}") logger.info(f" Number of features to select: {n_features_to_select}") logger.info(f" Combination strategy: {combination_strategy}") # Set up method configurations method_configs = { "variance": {"threshold": variance_threshold}, "correlation": {"threshold": correlation_threshold}, "mutual_info": {"k": n_features_to_select, "random_state": random_state}, "chi2": {"k": n_features_to_select}, "f_test": {"k": n_features_to_select}, "rfe": { "estimator_type": "rf", "n_features": n_features_to_select, "random_state": random_state, }, "importance": { "method": "random_forest", "n_features": n_features_to_select, "random_state": random_state, }, "lasso": {"alpha": 0.01, "random_state": random_state}, "permutation": { "estimator_type": "rf", "n_features": n_features_to_select, "random_state": random_state, }, "final_k": n_features_to_select, "combination_strategy": combination_strategy, } # Handle different job types if job_args.job_type == "training": # Training mode: Run full feature selection pipeline logger.info("Running in TRAINING mode: performing full feature selection") # Load preprocessed data (all splits) input_data_dir = input_paths["input_data"] logger.info(f"Loading data from {input_data_dir}") splits = load_preprocessed_data(input_data_dir) # Determine model artifacts output directory model_artifacts_output_dir = output_paths.get("model_artifacts_output") if not model_artifacts_output_dir: model_artifacts_output_dir = os.path.join( output_paths["processed_data"], "model_artifacts" ) os.makedirs(model_artifacts_output_dir, exist_ok=True) # Copy existing artifacts from previous steps (parameter accumulator pattern) model_artifacts_input_dir = input_paths.get("model_artifacts_input") if model_artifacts_input_dir: copy_existing_artifacts( model_artifacts_input_dir, model_artifacts_output_dir ) # Apply feature selection pipeline logger.info("Starting feature selection pipeline...") selection_results = apply_feature_selection_pipeline( splits, target_variable, methods, method_configs ) # Save feature-selected data output_data_dir = output_paths["processed_data"] logger.info(f"Saving selected data to {output_data_dir}") save_selected_data( splits, selection_results["selected_features"], target_variable, output_data_dir, ) # Save selection results and metadata to model artifacts directory logger.info(f"Saving all results to {model_artifacts_output_dir}") save_selection_results(selection_results, model_artifacts_output_dir) else: # Non-training mode: Use pre-computed selected features logger.info( f"Running in {job_args.job_type.upper()} mode: using pre-computed selected features" ) # Load pre-computed selected features if "model_artifacts_input" not in input_paths: raise ValueError( f"For non-training job type '{job_args.job_type}', model_artifacts_input input path must be provided" ) model_artifacts_input_dir = input_paths["model_artifacts_input"] logger.info( f"Loading pre-computed selected features from {model_artifacts_input_dir}" ) selected_features = load_selected_features(model_artifacts_input_dir) # Load single split data input_data_dir = input_paths["input_data"] logger.info(f"Loading {job_args.job_type} data from {input_data_dir}") splits = load_single_split_data(input_data_dir, job_args.job_type) # Determine model artifacts output directory model_artifacts_output_dir = output_paths.get("model_artifacts_output") if not model_artifacts_output_dir: model_artifacts_output_dir = os.path.join( output_paths["processed_data"], "model_artifacts" ) os.makedirs(model_artifacts_output_dir, exist_ok=True) # Copy existing artifacts from previous steps (parameter accumulator pattern) copy_existing_artifacts( model_artifacts_input_dir, model_artifacts_output_dir ) # Apply feature filtering (no computation, just filtering) logger.info( f"Applying pre-computed feature selection to {job_args.job_type} data" ) output_data_dir = output_paths["processed_data"] save_selected_data( splits, selected_features, target_variable, output_data_dir ) # Create minimal selection results for consistency selection_results = { "selected_features": selected_features, "method_results": [], # No methods run in non-training mode "combined_result": { "selected_features": selected_features, "scores": {f: 1.0 for f in selected_features}, # Dummy scores "method_contributions": {}, "combination_strategy": "pre_computed", "n_methods": 0, "n_selected": len(selected_features), }, "original_features": [], # Not available in non-training mode "target_variable": target_variable, "n_original_features": 0, # Not available in non-training mode "n_selected_features": len(selected_features), } # Metadata files are already copied via copy_existing_artifacts # Just verify they exist logger.info(f"Verifying metadata files in {model_artifacts_output_dir}") # Verify key files exist for filename in ["selected_features.json", "feature_scores.csv"]: file_path = os.path.join(model_artifacts_output_dir, filename) if os.path.exists(file_path): logger.info(f"✓ Found {filename}") else: logger.warning(f"⚠ Missing {filename}") # If report file wasn't found, create minimal one report_file = os.path.join( model_artifacts_output_dir, "feature_selection_report.json" ) if not os.path.exists(report_file): minimal_report = { "selection_summary": { "total_features": 0, # Not available in non-training mode "selected_features": len(selected_features), "selection_methods": ["pre_computed"], "combination_strategy": "pre_computed", "processing_time": 0, # No processing in non-training mode "target_variable": target_variable, "job_type": job_args.job_type, "mode": "non_training", }, "method_performance": {}, # No methods run "feature_statistics": { "selected_feature_count": len(selected_features) }, } with open(report_file, "w") as f: json.dump(minimal_report, f, indent=2, sort_keys=True) logger.info("Created minimal selection report") logger.info( f"Applied pre-computed feature selection: {len(selected_features)} features" ) logger.info("====== FEATURE SELECTION COMPLETED SUCCESSFULLY ======") except Exception as e: logger.error(f"FATAL ERROR in feature selection: {str(e)}") logger.error(traceback.format_exc()) raise
# ------------------------------------------------------------------------- # Script Entry Point # ------------------------------------------------------------------------- if __name__ == "__main__": logger.info("Feature selection script starting...") # Parse command line arguments parser = argparse.ArgumentParser(description="Feature Selection Script") parser.add_argument( "--job_type", type=str, default="training", choices=["training", "validation", "testing"], help="Type of job (training/validation/testing)", ) args = parser.parse_args() # Define input and output paths using container defaults input_paths = {"input_data": CONTAINER_PATHS["INPUT_DATA"]} # For non-training jobs, add model_artifacts_input input path if args.job_type != "training": input_paths["model_artifacts_input"] = ( CONTAINER_PATHS["INPUT_DATA"] + "/model_artifacts" ) output_paths = { "processed_data": CONTAINER_PATHS["OUTPUT_DATA"] + "/data", "model_artifacts_output": CONTAINER_PATHS["OUTPUT_DATA"] + "/model_artifacts", } # Collect environment variables environ_vars = { "FEATURE_SELECTION_METHODS": os.environ.get( "FEATURE_SELECTION_METHODS", "variance,correlation,mutual_info,rfe" ), "LABEL_FIELD": os.environ.get("LABEL_FIELD"), "N_FEATURES_TO_SELECT": os.environ.get("N_FEATURES_TO_SELECT", "10"), "CORRELATION_THRESHOLD": os.environ.get("CORRELATION_THRESHOLD", "0.95"), "VARIANCE_THRESHOLD": os.environ.get("VARIANCE_THRESHOLD", "0.01"), "RANDOM_STATE": os.environ.get("RANDOM_STATE", "42"), "COMBINATION_STRATEGY": os.environ.get("COMBINATION_STRATEGY", "voting"), } try: logger.info(f"Starting feature selection with:") logger.info(f" Input directory: {input_paths['input_data']}") logger.info(f" Output directory: {output_paths['processed_data']}") logger.info(f" Job type: {args.job_type}") # Call the main function main(input_paths, output_paths, environ_vars, args) logger.info("Feature selection script completed successfully") sys.exit(0) except Exception as e: logger.error(f"Exception during feature selection: {str(e)}") logger.error(traceback.format_exc()) sys.exit(1)