| 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() |
|
|
