from __future__ import annotations import argparse import csv from pathlib import Path import numpy as np import pandas as pd from sklearn.linear_model import LogisticRegression, Ridge from sklearn.metrics import ( accuracy_score, balanced_accuracy_score, f1_score, mean_absolute_error, mean_squared_error, r2_score, roc_auc_score, ) from sklearn.pipeline import make_pipeline from sklearn.preprocessing import LabelEncoder, StandardScaler, label_binarize from pet_vlm_dataset import load_suvr_vector CLASSIFICATION_TASKS = { "clinical_label_3way": {"column": "clinical_label", "labels": ["CN", "MCI", "AD"]}, "ad_vs_cn": {"column": "clinical_label", "labels": ["CN", "AD"]}, "pmci_vs_smci": {"column": "conversion_label", "labels": ["SMCI", "PMCI"]}, } REGRESSION_TASKS = { "adas11": "adas11", "mmse": "mmse", "ravlt_immediate": "ravlt_immediate", "ldeltotal": "ldeltotal", } def load_features(manifest: Path) -> tuple[pd.DataFrame, np.ndarray]: df = pd.read_csv(manifest) vectors = [] for csv_path in df["suvr_csv_path"]: _, suvr = load_suvr_vector(csv_path) vectors.append(suvr.numpy()) return df, np.stack(vectors, axis=0) def subset_classification(df: pd.DataFrame, x: np.ndarray, column: str, labels: list[str]) -> tuple[np.ndarray, np.ndarray]: mask = df[column].isin(labels).to_numpy() return x[mask], df.loc[mask, column].astype(str).to_numpy() def subset_regression(df: pd.DataFrame, x: np.ndarray, column: str) -> tuple[np.ndarray, np.ndarray]: y = pd.to_numeric(df[column], errors="coerce") mask = y.notna().to_numpy() return x[mask], y.loc[mask].astype(float).to_numpy() def evaluate_classification( name: str, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray, y_val: np.ndarray, x_test: np.ndarray, y_test: np.ndarray, ) -> dict[str, object]: encoder = LabelEncoder() encoder.fit(np.concatenate([y_train, y_val, y_test])) y_train_i = encoder.transform(y_train) y_val_i = encoder.transform(y_val) y_test_i = encoder.transform(y_test) best_c = None best_score = -np.inf best_model = None for c in [0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0]: model = make_pipeline( StandardScaler(), LogisticRegression(C=c, max_iter=5000, class_weight="balanced"), ) model.fit(x_train, y_train_i) pred_val = model.predict(x_val) score = balanced_accuracy_score(y_val_i, pred_val) if score > best_score: best_score = score best_c = c best_model = model assert best_model is not None pred = best_model.predict(x_test) proba = best_model.predict_proba(x_test) result: dict[str, object] = { "task": name, "type": "classification", "n_train": len(y_train), "n_val": len(y_val), "n_test": len(y_test), "selected_param": best_c, "accuracy": accuracy_score(y_test_i, pred), "balanced_accuracy": balanced_accuracy_score(y_test_i, pred), "macro_f1": f1_score(y_test_i, pred, average="macro"), } if len(encoder.classes_) == 2: result["auroc"] = roc_auc_score(y_test_i, proba[:, 1]) else: y_bin = label_binarize(y_test_i, classes=np.arange(len(encoder.classes_))) result["auroc"] = roc_auc_score(y_bin, proba, average="macro", multi_class="ovr") return result def evaluate_regression( name: str, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray, y_val: np.ndarray, x_test: np.ndarray, y_test: np.ndarray, ) -> dict[str, object]: best_alpha = None best_rmse = np.inf best_model = None for alpha in [0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0, 30.0, 100.0]: model = make_pipeline(StandardScaler(), Ridge(alpha=alpha)) model.fit(x_train, y_train) pred_val = model.predict(x_val) rmse = mean_squared_error(y_val, pred_val) ** 0.5 if rmse < best_rmse: best_rmse = rmse best_alpha = alpha best_model = model assert best_model is not None pred = best_model.predict(x_test) return { "task": name, "type": "regression", "n_train": len(y_train), "n_val": len(y_val), "n_test": len(y_test), "selected_param": best_alpha, "mae": mean_absolute_error(y_test, pred), "rmse": mean_squared_error(y_test, pred) ** 0.5, "r2": r2_score(y_test, pred), "pearson": float(np.corrcoef(y_test, pred)[0, 1]) if np.std(pred) > 1e-8 else np.nan, } def main() -> None: parser = argparse.ArgumentParser(description="Run clinical downstream probes from ground-truth SUVR vectors.") parser.add_argument("--train", type=Path, default=Path("data/metadata/splits/train_clinical.csv")) parser.add_argument("--val", type=Path, default=Path("data/metadata/splits/val_clinical.csv")) parser.add_argument("--test", type=Path, default=Path("data/metadata/splits/test_clinical.csv")) parser.add_argument("--out", type=Path, default=Path("runs/clinical/suvr_clinical_probe.csv")) args = parser.parse_args() train_df, x_train_all = load_features(args.train) val_df, x_val_all = load_features(args.val) test_df, x_test_all = load_features(args.test) results = [] for name, spec in CLASSIFICATION_TASKS.items(): x_train, y_train = subset_classification(train_df, x_train_all, spec["column"], spec["labels"]) x_val, y_val = subset_classification(val_df, x_val_all, spec["column"], spec["labels"]) x_test, y_test = subset_classification(test_df, x_test_all, spec["column"], spec["labels"]) if min(len(y_train), len(y_val), len(y_test)) == 0: continue results.append(evaluate_classification(name, x_train, y_train, x_val, y_val, x_test, y_test)) for name, column in REGRESSION_TASKS.items(): x_train, y_train = subset_regression(train_df, x_train_all, column) x_val, y_val = subset_regression(val_df, x_val_all, column) x_test, y_test = subset_regression(test_df, x_test_all, column) if min(len(y_train), len(y_val), len(y_test)) == 0: continue results.append(evaluate_regression(name, x_train, y_train, x_val, y_val, x_test, y_test)) args.out.parent.mkdir(parents=True, exist_ok=True) fieldnames = sorted({key for row in results for key in row.keys()}) with args.out.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() writer.writerows(results) print(f"wrote={args.out}") for row in results: print(row) if __name__ == "__main__": main()