import os import pandas as pd import numpy as np from sklearn.model_selection import GroupShuffleSplit, train_test_split from sklearn.preprocessing import StandardScaler, OneHotEncoder from sklearn.impute import SimpleImputer, KNNImputer from sklearn.compose import ColumnTransformer from sklearn.pipeline import Pipeline class PPMIDataPreprocessor: """ Clean, leak-proof preprocessing pipeline for PPMI dataset. Handles patient-level splitting, imputation, encoding, scaling. """ def __init__(self): # ------------------------- # Selected Features # ------------------------- self.selected_features = [ # Demographics "age", "SEX", "EDUCYRS", "race", "BMI", # Family History "fampd", "fampd_bin", # Motor symptoms "sym_tremor", "sym_rigid", "sym_brady", "sym_posins", # Non-motor "rem", "ess", "gds", "stai", # Cognitive "moca", "clockdraw", "bjlot", # --- EXTENDED MEDICAL FEATURES --- # Olfactory "upsit", "upsit_pctl", # UPDRS (Gold standard for PD) "updrs1_score", "updrs2_score", "updrs3_score", "updrs4_score", "updrs_totscore", # DatScan (Imaging) "mean_caudate", "mean_putamen", # Biomarkers (CSF) "abeta", "tau", "ptau", # Target "COHORT", # Patient ID "PATNO" ] # Features to impute using KNN (to preserve correlations) self.knn_cols = [ "moca", "clockdraw", "bjlot", "rem", "gds", "ess", "stai", "upsit", "upsit_pctl", "updrs1_score", "updrs2_score", "updrs3_score", "updrs4_score", "updrs_totscore", "mean_caudate", "mean_putamen", "abeta", "tau", "ptau" ] # Categorical self.categorical_cols = ["SEX", "race", "fampd"] # Numeric self.numeric_cols = [ "age", "EDUCYRS", "BMI", "fampd_bin", "sym_tremor", "sym_rigid", "sym_brady", "sym_posins", # New Medical Features (added for enhancement) "updrs1_score", "updrs2_score", "updrs3_score", "updrs4_score", "updrs_totscore", "mean_caudate", "mean_putamen", "abeta", "tau", "ptau", "upsit", "upsit_pctl" ] self.preprocessor = None # ------------------------------------------------------------ def _clean_biomarker(self, x): """Clean biomarker strings like '>1700' or '<200' to floats.""" if pd.isna(x): return np.nan if isinstance(x, str): x = x.replace(">", "").replace("<", "") try: return float(x) except ValueError: return np.nan def load_data(self, file_path): """Load and filter only necessary columns.""" df = pd.read_csv(file_path, low_memory=False) # Clean biomarker columns if they exist for col in ["abeta", "tau", "ptau"]: if col in df.columns: df[col] = df[col].apply(self._clean_biomarker) # Select available features (intersection with file columns) available = [c for c in self.selected_features if c in df.columns] # Always require COHORT/PATNO if "COHORT" not in available or "PATNO" not in available: raise ValueError("Dataset missing required COHORT or PATNO columns") df = df[available].dropna(subset=["COHORT", "PATNO"]) return df # ------------------------------------------------------------ def patient_split(self, df, test_size=0.2): """Split such that the same patient never appears in both sets.""" gss = GroupShuffleSplit(test_size=test_size, n_splits=1, random_state=42) train_idx, test_idx = next(gss.split(df, groups=df["PATNO"])) train_df = df.iloc[train_idx].reset_index(drop=True) test_df = df.iloc[test_idx].reset_index(drop=True) return train_df, test_df # ------------------------------------------------------------ def build_preprocessor(self): """Builds leak-proof transformers.""" numeric_transformer = Pipeline(steps=[ ("imputer_mean", SimpleImputer(strategy="mean")), ("scaler", StandardScaler()) ]) knn_transformer = Pipeline(steps=[ ("imputer_knn", KNNImputer(n_neighbors=5)), ("scaler", StandardScaler()) ]) categorical_transformer = Pipeline(steps=[ ("imputer_mode", SimpleImputer(strategy="most_frequent")), ("encoder", OneHotEncoder(handle_unknown="ignore")) ]) self.preprocessor = ColumnTransformer( transformers=[ ("num", numeric_transformer, [c for c in self.numeric_cols if c not in self.knn_cols]), ("knn", knn_transformer, self.knn_cols), ("cat", categorical_transformer, self.categorical_cols), ], remainder="drop" ) return self.preprocessor # ------------------------------------------------------------ def prepare(self, file_path): """ COMPLETE DATA PREPARATION PIPELINE Returns X_train, X_test, y_train, y_test """ df = self.load_data(file_path) train_df, test_df = self.patient_split(df) X_train = train_df.drop(["COHORT", "PATNO"], axis=1) y_train = train_df["COHORT"] X_test = test_df.drop(["COHORT", "PATNO"], axis=1) y_test = test_df["COHORT"] # Build processor & fit only on training data pre = self.build_preprocessor() X_train_processed = pre.fit_transform(X_train) X_test_processed = pre.transform(X_test) return X_train_processed, X_test_processed, y_train.values, y_test.values class DataPreprocessor: """Backwards-compatible wrapper around the new PPMI pipeline.""" def __init__(self): self.core = PPMIDataPreprocessor() self.feature_names_ = None self.class_mapping_ = None self.preprocessor_ = None self.train_df_ = None self.test_df_ = None @staticmethod def _to_python_scalar(value): if isinstance(value, np.generic): return value.item() return value def _load_all_files(self, file_paths): if isinstance(file_paths, str): file_paths = [file_paths] dataframes = [] for path in file_paths: if not path: continue if not os.path.exists(path): print(f"[WARN] DataPreprocessor: '{path}' not found, skipping.") continue df = self.core.load_data(path) dataframes.append(df) if not dataframes: raise FileNotFoundError("No valid CSV files were provided to DataPreprocessor.") combined = pd.concat(dataframes, ignore_index=True) combined = combined.drop_duplicates().reset_index(drop=True) return combined def prepare_data(self, file_paths, test_size=0.2, use_patient_split=True): """Expose the legacy API expected by the training scripts.""" df = self._load_all_files(file_paths) if use_patient_split: train_df, test_df = self.core.patient_split(df, test_size=test_size) else: train_df, test_df = train_test_split( df, test_size=test_size, random_state=42, stratify=df["COHORT"], ) self.train_df_ = train_df.reset_index(drop=True) self.test_df_ = test_df.reset_index(drop=True) classes_sorted = np.sort(df["COHORT"].unique()) self.class_mapping_ = { self._to_python_scalar(original): int(idx) for idx, original in enumerate(classes_sorted) } X_train = train_df.drop(["COHORT", "PATNO"], axis=1) y_train = train_df["COHORT"].map(self.class_mapping_).values X_test = test_df.drop(["COHORT", "PATNO"], axis=1) y_test = test_df["COHORT"].map(self.class_mapping_).values pre = self.core.build_preprocessor() X_train_processed = pre.fit_transform(X_train) X_test_processed = pre.transform(X_test) self.preprocessor_ = pre try: self.feature_names_ = pre.get_feature_names_out(X_train.columns).tolist() except AttributeError: self.feature_names_ = None return X_train_processed, X_test_processed, y_train, y_test def get_feature_names(self): if self.feature_names_ is None: raise ValueError("Feature names are unavailable. Call prepare_data() first.") return self.feature_names_ def get_preprocessor(self): if self.preprocessor_ is None: raise ValueError("Preprocessor is unavailable. Call prepare_data() first.") return self.preprocessor_ def get_class_mapping(self): if self.class_mapping_ is None: raise ValueError("Class mapping is unavailable. Call prepare_data() first.") return self.class_mapping_ def get_split_frames(self): if self.train_df_ is None or self.test_df_ is None: raise ValueError("Split frames are unavailable. Call prepare_data() first.") return self.train_df_.copy(), self.test_df_.copy() # ------------------------------------------------------------------- # Script example (not executed when imported) # ------------------------------------------------------------------- if __name__ == "__main__": file_path = "PPMI_Curated_Data.csv" prep = PPMIDataPreprocessor() X_train, X_test, y_train, y_test = prep.prepare(file_path) print("Train:", X_train.shape) print("Test:", X_test.shape)