mboukabous/train_unsupervised

1 2"""3train_dimred_model.py4 5Trains a dimensionality reduction model (e.g., PCA, t-SNE, UMAP) on a dataset.6It can drop or select specific columns, perform label encoding on any non-numeric columns,7and optionally visualize the reduced data (2D or 3D).8 9Example Usage:10--------------11python scripts/train_dimred_model.py \12    --model_module pca \13    --data_path data/raw/breast-cancer-wisconsin-data/data.csv \14    --drop_columns "id" \15    --select_columns "radius_mean, texture_mean, perimeter_mean, area_mean" \16    --visualize17"""18 19import os20import sys21import argparse22import importlib23import pandas as pd24import numpy as np25import joblib26 27from sklearn.impute import SimpleImputer28from sklearn.preprocessing import LabelEncoder29import matplotlib.pyplot as plt30 31def main(args):32    # Move to project root if needed33    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))34    os.chdir(project_root)35    sys.path.insert(0, project_root)36 37    # Dynamically import the chosen model module (pca.py, tsne.py, umap.py, etc.)38    model_module_path = f"models.unsupervised.dimred.{args.model_module}"39    model_module = importlib.import_module(model_module_path)40 41    # Retrieve the estimator from the model file42    estimator = model_module.estimator43    default_n_components = getattr(model_module, 'default_n_components', 2)  # fallback44 45    # Prepare results directory46    if args.results_path is None:47        # e.g., 'results/PCA_DimRed'48        args.results_path = os.path.join('results', f"{estimator.__class__.__name__}_DimRed")49    os.makedirs(args.results_path, exist_ok=True)50    51    # Prepare model directory52    if args.model_path is None:53        # e.g., 'saved_model/PCA_DimRed'54        args.model_path = os.path.join('saved_models', f"{estimator.__class__.__name__}_DimRed")55    os.makedirs(args.model_path, exist_ok=True)56 57    # Load data from CSV58    df = pd.read_csv(args.data_path)59    print(f"Data loaded from {args.data_path}, initial shape: {df.shape}")60 61    # Drop empty columns62    df = df.dropna(axis='columns', how='all')63    print("After dropping empty columns:", df.shape)64 65    # Drop specified columns if any66    if args.drop_columns:67        drop_cols = [col.strip() for col in args.drop_columns.split(',') if col.strip()]68        df = df.drop(columns=drop_cols, errors='ignore')69        print(f"Dropped columns: {drop_cols} | New shape: {df.shape}")70    71    # Select specified columns if any72    if args.select_columns:73        keep_cols = [col.strip() for col in args.select_columns.split(',') if col.strip()]74        df = df[keep_cols]75        print(f"Selected columns: {keep_cols} | New shape: {df.shape}")76 77    # Label-encode non-numeric columns78    for col in df.columns:79        if df[col].dtype == 'object':80            le = LabelEncoder()81            df[col] = le.fit_transform(df[col])82 83    # Impute84    imputer = SimpleImputer(strategy='mean')  # or 'median'85    df_array = imputer.fit_transform(df)86    df_imputed = pd.DataFrame(df_array, columns=df.columns)87    print("After label-encoding and imputation:", df_imputed.shape)88    89    # Convert DataFrame to numpy array90    X = df_imputed.values91    print(f"Final data shape after dropping/selecting columns and encoding: {X.shape}")92 93    # Fit-transform the data (typical for dimensionality reduction)94    X_transformed = estimator.fit_transform(X)95    print(f"Dimensionality reduction done using {args.model_module}. Output shape: {X_transformed.shape}")96 97    # Save the model98    model_output_path = os.path.join(args.model_path, "dimred_model.pkl")99    joblib.dump(estimator, model_output_path)100    print(f"Model saved to {model_output_path}")101 102    # Save the transformed data103    transformed_path = os.path.join(args.results_path, "X_transformed.csv")104    pd.DataFrame(X_transformed).to_csv(transformed_path, index=False)105    print(f"Transformed data saved to {transformed_path}")106 107    # Visualization (only if 2D or 3D)108    if args.visualize:109        n_dims = X_transformed.shape[1]110        if n_dims == 2:111            plt.figure(figsize=(6,5))112            plt.scatter(X_transformed[:,0], X_transformed[:,1], s=30, alpha=0.7, c='blue')113            plt.title(f"{estimator.__class__.__name__} 2D Projection")114            plt.xlabel("Component 1")115            plt.ylabel("Component 2")116            plot_path = os.path.join(args.results_path, "dimred_plot_2D.png")117            plt.savefig(plot_path)118            plt.show()119            print(f"2D plot saved to {plot_path}")120        elif n_dims == 3:121            from mpl_toolkits.mplot3d import Axes3D122            fig = plt.figure()123            ax = fig.add_subplot(111, projection='3d')124            ax.scatter(X_transformed[:,0], X_transformed[:,1], X_transformed[:,2], s=30, alpha=0.7, c='blue')125            ax.set_title(f"{estimator.__class__.__name__} 3D Projection")126            ax.set_xlabel("Component 1")127            ax.set_ylabel("Component 2")128            ax.set_zlabel("Component 3")129            plot_path = os.path.join(args.results_path, "dimred_plot_3D.png")130            plt.savefig(plot_path)131            plt.show()132            print(f"3D plot saved to {plot_path}")133        else:134            print(f"Visualization only supported for 2D or 3D outputs. Got {n_dims}D, skipping.")135 136 137if __name__ == "__main__":138    parser = argparse.ArgumentParser(description="Train a dimensionality reduction model.")139    parser.add_argument('--model_module', type=str, required=True,140                        help='Name of the dimred model module (e.g. pca, tsne, umap).')141    parser.add_argument('--data_path', type=str, required=True,142                        help='Path to the CSV dataset file.')143    parser.add_argument('--model_path', type=str, default=None,144                        help='Where to save the fitted model.')145    parser.add_argument('--results_path', type=str, default=None,146                        help='Directory to store results (transformed data, plots).')147    parser.add_argument('--visualize', action='store_true',148                        help='Plot the transformed data if 2D or 3D.')149    parser.add_argument('--drop_columns', type=str, default='',150                        help='Comma-separated column names to drop from the dataset.')151    parser.add_argument('--select_columns', type=str, default='',152                        help='Comma-separated column names to keep (ignore the rest).')153 154    args = parser.parse_args()155    main(args)156