PaGOA Training Pipeline (PET/CT + EHR)

This repository trains and evaluates the PaGOA pipeline in 3 stages:

1. Stage 1: train multimodal image+EHR encoders (Cox survival objective).
2. Stage 2: export frozen embeddings and build fold-aware kNN graphs.
3. Stage 3: train graph survival model (GraphSAGE + discrete-time survival head).

Main entry points

• train/train_multicohort_image_ehr_surv.py (Stage 1)
• train/export_multicohort_embeddings.py (Stage 2a)
• graph/build_knn_graph.py (Stage 2b)
• train/train_graph_patient_surv_no_concepts.py (Stage 3)

Useful run scripts already included:

• run_parallel.sh (Stage 1, all folds)
• single_gpu_stage1_melanoma.sh (Stage 1, Melanoma only)
• run_graphs.sh (Stage 3 batch runner)
• run_coxph.sh, run_deepsurv.sh, run_deephit.sh (baselines)

Data files used by default

• pre_processing/data/all_cohorts_manifest_clean.csv
• pre_processing/data/all_cohorts_ehr_unified.csv

Quick start: full pipeline

Set base variables once:

cd /home/azureuser/PITL
export PYTHONPATH=/home/azureuser/PITL

OUT_BASE=/home/azureuser/PITL/outputs/multicohort_img_ehr_no_concepts_unified_index_resolved
MANIFEST=/home/azureuser/PITL/pre_processing/data/all_cohorts_manifest_clean.csv
EHR=/home/azureuser/PITL/pre_processing/data/all_cohorts_ehr_unified.csv

Stage 1: train encoders

Option A (recommended launcher):

bash run_parallel.sh

Option B (single cohort example):

bash single_gpu_stage1_melanoma.sh

Expected Stage-1 artifacts:

• ${OUT_BASE}/folds.json
• ${OUT_BASE}/fold_<k>/best_multicohort_img_ehr_surv.pt

Stage 2: export embeddings + build kNN graph

Run per fold and cohort (all, TCGA, Melanoma, NSCLC depending on your experiments).

Example for FOLD=0, COHORT=Melanoma:

FOLD=0
COHORT=Melanoma

python train/export_multicohort_embeddings.py \
  --manifest-csv "${MANIFEST}" \
  --ehr-csv "${EHR}" \
  --stage1-out-dir "${OUT_BASE}" \
  --ckpt-name best \
  --folds-json "${OUT_BASE}/folds.json" \
  --fold "${FOLD}" \
  --spatial-size 128 128 192 \
  --batch-size 2 \
  --device cuda:0 \
  --cohort-filter "${COHORT}" \
  --out-pt "${OUT_BASE}/fold_${FOLD}/exported_embeddings/${COHORT}/best/fold${FOLD}_embeddings.pt"

python graph/build_knn_graph.py \
  --embeddings-pt "${OUT_BASE}/fold_${FOLD}/exported_embeddings/${COHORT}/best/fold${FOLD}_embeddings.pt" \
  --folds-json "${OUT_BASE}/folds.json" \
  --fold "${FOLD}" \
  --k 10 \
  --emb-key z_img_z_ehr \
  --within-dataset \
  --manifest-csv "${MANIFEST}" \
  --out-json "${OUT_BASE}/fold_${FOLD}/graph_result/${COHORT}/best/knn_k10_within_dataset.json"

Stage 3: train graph model

Batch launcher (expects Stage-2 embeddings and kNN json already prepared):

bash run_graphs.sh

Or single run example:

python train/train_graph_patient_surv_no_concepts.py \
  --embeddings-pt "${OUT_BASE}/fold_0/exported_embeddings/Melanoma/best/fold0_embeddings.pt" \
  --knn-json "${OUT_BASE}/fold_0/graph_result/Melanoma/best/knn_k10_within_dataset.json" \
  --manifest-csv "${MANIFEST}" \
  --folds-json "${OUT_BASE}/folds.json" \
  --fold 0 \
  --n-time-bins 12 \
  --epochs 200 \
  --rank-loss-weight 0.8 \
  --weight-decay 1e-4 \
  --knn-k-eval 10 \
  --eval-edge-policy rewire_available \
  --num-graph-layers 2 \
  --device cuda:0 \
  --out-dir "${OUT_BASE}/fold_0/graph_result/Melanoma/best/rk08wd1e4_rewire_available"

Expected Stage-3 artifacts (per run directory):

• best.pt, last.pt
• train_metrics.csv
• eval_summary.csv
• optional patient_predictions.csv (if --save-patient-predictions)

Baselines

All baseline scripts consume the same exported embeddings:

• bash run_coxph.sh
• bash run_deepsurv.sh
• bash run_deephit.sh

Each script writes to outputs/baselines/... and logs progress per dataset/fold.

Notes

• run_graphs.sh does not build kNN graphs; generate them first in Stage 2.
• Most scripts hardcode /home/azureuser/PITL; update these paths if needed.