Add assessment and all extra-info experiments

extra_info_assessment_all_analysis_plan.md

@@ -0,0 +1,164 @@
+# assessment_only 与 all 实验及分析方案
+
+## 1. 实验目的
+
+补充现有四级 extra-information 证据链：
+
+1. `disease_only`：疾病事件、相对患病时间和 sex；无 CHECKUP、无 extra-info token。
+2. `smoking_alcohol_bmi`：疾病史、sex、CHECKUP、smoking/alcohol/BMI。
+3. `assessment_only`：疾病史、sex、CHECKUP、65项常规体格、肺功能、血液、尿液和生化指标。
+4. `all`：疾病史、sex、CHECKUP、全部265项体检和暴露信息。
+
+目标是区分：
+
+- 疾病史本身能够提供多少未来疾病信息；
+- 常规体检在疾病史之外增加多少信息；
+- 生活方式、社会经济、心理和环境暴露在完整体检之外增加多少信息；
+- 额外信息对疾病预测和死亡预测是否具有不同作用。
+
+## 2. 固定模型
+
+所有新增实验固定为：
+
+```text
+TrajMixer
++ all_future
++ relative
++ Weibull
++ timed disease history
++ sex
+```
+
+仅改变 `extra_info_types_file`：
+
+- `extra_info_types_assessment_only.txt`
+- `extra_info_types_all.txt`
+
+每种配置运行 seed 42、43、44。
+
+## 3. A6000 48GB 设置
+
+| 配置 | Batch size | 原因 |
+|---|---:|---|
+| assessment_only | 256 | 最多65个 extra-info token，48GB余量充足 |
+| all | 128 | 最多265个 extra-info token，relative RBF attention 显存随总序列长度平方增长 |
+
+当前训练代码为 FP32/TF32，并未使用 AMP。若 `all batch=128` 在极端长序列 batch 上出现 CUDA OOM，降为64，不自动重试，避免同一配置产生多个不完整 run。
+
+batch size 是 `all` 与其他模型之间的潜在训练差异。代码会按 batch size 自动缩放学习率，但最终报告仍需明确记录该差异。若 `all` 的结果处于模型选择临界区，再补 seed 42 的 batch-size sensitivity，而不是预先扩大实验矩阵。
+
+## 4. 主要比较
+
+### 4.1 常规体检的增量价值
+
+```text
+assessment_only − disease_only
+```
+
+回答常规器官功能检测指标在疾病序列之外提供多少信息。
+
+### 4.2 全部信息相对常规体检
+
+```text
+all − assessment_only
+```
+
+回答生活方式、社会经济、心理和环境暴露是否在常规体检之后仍有增量价值。这是新增实验中最干净的主要比较，因为两组都保留 CHECKUP，Landmark 和随访边界应一致。
+
+### 4.3 全体检相对紧凑变量集
+
+```text
+assessment_only − smoking_alcohol_bmi
+```
+
+回答65项常规体检是否优于紧凑的 smoking/alcohol/BMI 输入。
+
+### 4.4 全部信息相对紧凑变量集
+
+```text
+all − smoking_alcohol_bmi
+```
+
+衡量从当前最终模型扩展到全部 extra information 的最大增益。
+
+### 4.5 既有比较
+
+保留：
+
+```text
+smoking_alcohol_bmi − disease_only
+```
+
+与新增结果共同形成完整的信息增量路径。
+
+## 5. 评估指标
+
+疾病和死亡分开分析。
+
+### 判别能力
+
+- Landmark AUC；
+- 各 horizon AUC；
+- disease cell win rate；
+- 三个 seed 的均值、标准差和方向一致性。
+
+### 概率与似然质量
+
+- IPCW Brier；
+- 固定时点 IPCW NLL；
+- 连续时间 point-process NLL；
+- Expected/Observed ratio；
+- calibration-in-the-large；
+- calibration slope；
+- 校准曲线。
+
+### 复杂度与稳定性
+
+- 参数量；
+- 每个 epoch 运行时间；
+- 峰值显存；
+- 三个 seed 的性能波动；
+- 缺失值较多的 extra-info 类型是否造成训练不稳定。
+
+## 6. 配对和汇总方法
+
+- replicate unit 为 seed；
+- 按相同 `seed × label_code × sex × horizon` 配对；
+- 每个比较使用三个 seed 共同存在的 cell；
+- 先在 seed 内汇总，再计算三个 seed 的均值和样本标准差；
+- AUC 越高越好；
+- Brier、NLL及绝对校准偏差越低越好；
+- 不以单个 seed 或单个 horizon 决定模型。
+
+## 7. disease_only 比较的评估限制
+
+`disease_only` 按设计删除 CHECKUP，其他三组保留 CHECKUP。当前评估实现会使两类模型的随访终点和 `n_at_risk` 略有差异。
+
+因此：
+
+- `assessment_only`、`smoking_alcohol_bmi`、`all` 三者之间可以直接比较；
+- 它们与 `disease_only` 的比较应使用固定的原始随访终点、Landmark 和 censoring；
+- 在共享风险集评估完成前，不能把与 `disease_only` 的全部差异严格归因于 extra-info 数值。
+
+## 8. 决策规则
+
+1. 如果 `assessment_only` 已达到 `all` 的绝大部分性能，并且校准更稳定，优先选择 `assessment_only`，因为它与器官功能重建目标一致且解释更清楚。
+2. 如果 `all` 在疾病和死亡 AUC、Brier、NLL上均稳定优于 `assessment_only`，则将 `all` 作为性能上限模型，但不直接作为器官负担教师模型。
+3. 如果 `all` 只提高 AUC而恶化 Brier/NLL或 seed 波动明显，不升级最终模型。
+4. disease-only Timed 仍是生成纯疾病来源器官负担分数的教师模型；assessment/all 实验用于界定疾病史遗漏的信息，而不是改变该分数的纯疾病定义。
+
+## 9. 运行与后续评估
+
+训练：
+
+```bash
+bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0
+```
+
+多GPU：
+
+```bash
+bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0,1,2
+```
+
+训练完成后，现有扫描脚本分别生成 AUC 和 calibration/Brier/NLL。合并分析时，将新增 `assessment_only` 和 `all` 两个配置纳入主 Timed 表，不纳入 disease-history Ordered/Set 专表。

train_extra_info_assessment_all_multiseed_linux.sh

@@ -0,0 +1,389 @@
+#!/usr/bin/env bash
+#
+# Train the two remaining extra-information experiments:
+#   1. assessment_only: 65 routine assessment/body/laboratory variables.
+#   2. all:             all 265 assessment and exposure variables.
+#
+# Fixed model:
+#   TrajMixer + all_future + relative + Weibull + timed disease history + sex
+#
+# A6000 48 GB defaults:
+#   assessment_only batch_size=256
+#   all             batch_size=128
+#
+# Each task uses one GPU. Tasks assigned to the same GPU run sequentially;
+# different GPUs run in parallel.
+#
+# Examples:
+#   bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0
+#   bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0,1,2
+#   bash train_extra_info_assessment_all_multiseed_linux.sh \
+#       --gpus 0 --seeds 42 --types all --dry-run
+#
+
+set -uo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+cd "$SCRIPT_DIR"
+
+GPU_CSV=""
+SEED_CSV="42,43,44"
+TYPE_CSV="assessment_only,all"
+NUM_WORKERS=4
+ASSESSMENT_BATCH_SIZE=256
+ALL_BATCH_SIZE=128
+PYTHON_BIN="${PYTHON_BIN:-python}"
+CAMPAIGN_NAME="extra_info_assessment_all_multiseed"
+DRY_RUN=0
+
+ENTRYPOINT="$SCRIPT_DIR/train_all_future.py"
+ASSESSMENT_FILE="$SCRIPT_DIR/extra_info_types_assessment_only.txt"
+ALL_FILE="$SCRIPT_DIR/extra_info_types_all.txt"
+
+usage() {
+    cat <<'EOF'
+Usage:
+  bash train_extra_info_assessment_all_multiseed_linux.sh \
+      --gpus GPU_LIST [options]
+
+Required:
+  --gpus LIST                  Comma-separated GPU ids, for example 0,1,2.
+
+Options:
+  --seeds LIST                 Comma-separated seeds (default: 42,43,44).
+  --types LIST                 Subset of assessment_only,all (default: both).
+  --assessment-batch-size N    assessment_only batch size (default: 256).
+  --all-batch-size N           all batch size (default: 128 for A6000 48 GB).
+  --num-workers N              DataLoader workers per task (default: 4).
+  --python PATH                Python executable (default: $PYTHON_BIN or python).
+  --campaign NAME              Output campaign name.
+  --dry-run                    Print commands without creating files or training.
+  -h, --help                   Show this help message.
+
+Fixed experiment settings:
+  architecture                 traj_mixer_v5
+  target                       all_future
+  time mode                    relative
+  distribution                 weibull
+  disease history              timed
+  sex                          enabled by the model
+
+A6000 48 GB memory policy:
+  assessment_only              batch_size=256
+  all                          batch_size=128
+
+If all still runs out of memory because of an unusually long padded batch,
+restart that experiment with --all-batch-size 64.
+
+Outputs:
+  runs/<campaign>/seed_<seed>/traj_mixer_v5/...
+  batch_logs/<campaign>/seed_<seed>/<type>.log
+EOF
+}
+
+while (($# > 0)); do
+    case "$1" in
+        --gpus)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --gpus requires a value." >&2
+                exit 2
+            }
+            GPU_CSV="$2"
+            shift 2
+            ;;
+        --seeds)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --seeds requires a value." >&2
+                exit 2
+            }
+            SEED_CSV="$2"
+            shift 2
+            ;;
+        --types)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --types requires a value." >&2
+                exit 2
+            }
+            TYPE_CSV="$2"
+            shift 2
+            ;;
+        --assessment-batch-size)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --assessment-batch-size requires a value." >&2
+                exit 2
+            }
+            ASSESSMENT_BATCH_SIZE="$2"
+            shift 2
+            ;;
+        --all-batch-size)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --all-batch-size requires a value." >&2
+                exit 2
+            }
+            ALL_BATCH_SIZE="$2"
+            shift 2
+            ;;
+        --num-workers)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --num-workers requires a value." >&2
+                exit 2
+            }
+            NUM_WORKERS="$2"
+            shift 2
+            ;;
+        --python)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --python requires a value." >&2
+                exit 2
+            }
+            PYTHON_BIN="$2"
+            shift 2
+            ;;
+        --campaign)
+            [[ $# -ge 2 ]] || {
+                echo "ERROR: --campaign requires a value." >&2
+                exit 2
+            }
+            CAMPAIGN_NAME="$2"
+            shift 2
+            ;;
+        --dry-run)
+            DRY_RUN=1
+            shift
+            ;;
+        -h|--help)
+            usage
+            exit 0
+            ;;
+        *)
+            echo "ERROR: unknown argument: $1" >&2
+            usage >&2
+            exit 2
+            ;;
+    esac
+done
+
+[[ -n "$GPU_CSV" ]] || {
+    echo "ERROR: --gpus is required." >&2
+    usage >&2
+    exit 2
+}
+[[ -n "$SEED_CSV" ]] || {
+    echo "ERROR: --seeds must not be empty." >&2
+    exit 2
+}
+[[ -n "$TYPE_CSV" ]] || {
+    echo "ERROR: --types must not be empty." >&2
+    exit 2
+}
+[[ "$ASSESSMENT_BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
+    echo "ERROR: --assessment-batch-size must be a positive integer." >&2
+    exit 2
+}
+[[ "$ALL_BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
+    echo "ERROR: --all-batch-size must be a positive integer." >&2
+    exit 2
+}
+[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
+    echo "ERROR: --num-workers must be a non-negative integer." >&2
+    exit 2
+}
+[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
+    echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
+    exit 2
+}
+for required_file in "$ENTRYPOINT" "$ASSESSMENT_FILE" "$ALL_FILE"; do
+    [[ -f "$required_file" ]] || {
+        echo "ERROR: missing required file: $required_file" >&2
+        exit 2
+    }
+done
+command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
+    echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
+    exit 2
+}
+
+IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
+declare -A SEEN_GPUS=()
+for gpu in "${GPU_IDS[@]}"; do
+    [[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
+        echo "ERROR: invalid GPU id: $gpu" >&2
+        exit 2
+    }
+    [[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
+        echo "ERROR: duplicate GPU id: $gpu" >&2
+        exit 2
+    }
+    SEEN_GPUS["$gpu"]=1
+done
+
+IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
+declare -A SEEN_SEEDS=()
+for seed in "${SEEDS[@]}"; do
+    [[ "$seed" =~ ^[0-9]+$ ]] || {
+        echo "ERROR: invalid seed: $seed" >&2
+        exit 2
+    }
+    [[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
+        echo "ERROR: duplicate seed: $seed" >&2
+        exit 2
+    }
+    SEEN_SEEDS["$seed"]=1
+done
+
+IFS=',' read -r -a TYPES <<< "$TYPE_CSV"
+declare -A SEEN_TYPES=()
+for info_type in "${TYPES[@]}"; do
+    case "$info_type" in
+        assessment_only|all)
+            ;;
+        *)
+            echo "ERROR: invalid type: $info_type" >&2
+            echo "Expected assessment_only or all." >&2
+            exit 2
+            ;;
+    esac
+    [[ -z "${SEEN_TYPES[$info_type]+x}" ]] || {
+        echo "ERROR: duplicate type: $info_type" >&2
+        exit 2
+    }
+    SEEN_TYPES["$info_type"]=1
+done
+
+RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
+LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
+if ((!DRY_RUN)); then
+    mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
+fi
+
+declare -a JOB_SEEDS=()
+declare -a JOB_TYPES=()
+declare -a JOB_BATCH_SIZES=()
+declare -a JOB_EXTRA_FILES=()
+
+for seed in "${SEEDS[@]}"; do
+    for info_type in "${TYPES[@]}"; do
+        JOB_SEEDS+=("$seed")
+        JOB_TYPES+=("$info_type")
+        if [[ "$info_type" == "assessment_only" ]]; then
+            JOB_BATCH_SIZES+=("$ASSESSMENT_BATCH_SIZE")
+            JOB_EXTRA_FILES+=("$ASSESSMENT_FILE")
+        else
+            JOB_BATCH_SIZES+=("$ALL_BATCH_SIZE")
+            JOB_EXTRA_FILES+=("$ALL_FILE")
+        fi
+    done
+done
+
+print_command() {
+    printf '%q ' "$@"
+    printf '\n'
+}
+
+run_job() {
+    local job_index="$1"
+    local gpu="$2"
+    local seed="${JOB_SEEDS[$job_index]}"
+    local info_type="${JOB_TYPES[$job_index]}"
+    local batch_size="${JOB_BATCH_SIZES[$job_index]}"
+    local extra_file="${JOB_EXTRA_FILES[$job_index]}"
+    local seed_runs_root="$RUNS_ROOT/seed_$seed"
+    local seed_log_root="$LOG_ROOT/seed_$seed"
+    local log_file="$seed_log_root/$info_type.log"
+    local -a command=(
+        "$PYTHON_BIN"
+        -u
+        "$ENTRYPOINT"
+        --runs_root "$seed_runs_root"
+        --seed "$seed"
+        --batch_size "$batch_size"
+        --num_workers "$NUM_WORKERS"
+        --device cuda
+        --model_architecture traj_mixer_v5
+        --time_mode relative
+        --dist_mode weibull
+        --disease_history_mode timed
+        --extra_info_types_file "$extra_file"
+    )
+
+    if ((!DRY_RUN)); then
+        mkdir -p "$seed_runs_root" "$seed_log_root"
+    fi
+
+    echo "[$(date '+%F %T')] START seed=$seed type=$info_type gpu=$gpu batch=$batch_size"
+    echo "  log=$log_file"
+    if ((DRY_RUN)); then
+        printf '  CUDA_VISIBLE_DEVICES=%q ' "$gpu"
+        print_command "${command[@]}"
+        return 0
+    fi
+
+    if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
+        "${command[@]}" >"$log_file" 2>&1; then
+        echo "[$(date '+%F %T')] DONE  seed=$seed type=$info_type gpu=$gpu"
+        return 0
+    else
+        local exit_code=$?
+        echo "[$(date '+%F %T')] FAIL seed=$seed type=$info_type gpu=$gpu exit=$exit_code" >&2
+        echo "  See: $log_file" >&2
+        if [[ "$info_type" == "all" ]]; then
+            echo "  If this is CUDA OOM, retry with --all-batch-size 64." >&2
+        fi
+        return "$exit_code"
+    fi
+}
+
+worker() {
+    local slot="$1"
+    local gpu="${GPU_IDS[$slot]}"
+    local job_index
+    local failed=0
+
+    for ((job_index = slot; job_index < ${#JOB_SEEDS[@]}; job_index += ${#GPU_IDS[@]})); do
+        run_job "$job_index" "$gpu" || failed=1
+    done
+    return "$failed"
+}
+
+echo "Campaign: $CAMPAIGN_NAME"
+echo "Seeds: ${SEEDS[*]}"
+echo "Extra-info types: ${TYPES[*]}"
+echo "GPUs: ${GPU_IDS[*]}"
+echo "assessment_only batch size: $ASSESSMENT_BATCH_SIZE"
+echo "all batch size: $ALL_BATCH_SIZE"
+echo "Total tasks: ${#JOB_SEEDS[@]}"
+echo "Runs root: $RUNS_ROOT"
+echo "Log root: $LOG_ROOT"
+if command -v nvidia-smi >/dev/null 2>&1; then
+    echo "Selected GPU inventory:"
+    for gpu in "${GPU_IDS[@]}"; do
+        nvidia-smi \
+            --id="$gpu" \
+            --query-gpu=index,name,memory.total \
+            --format=csv,noheader \
+            2>/dev/null || true
+    done
+fi
+echo
+
+declare -a WORKER_PIDS=()
+for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
+    worker "$slot" &
+    WORKER_PIDS+=("$!")
+done
+
+overall_status=0
+for pid in "${WORKER_PIDS[@]}"; do
+    wait "$pid" || overall_status=1
+done
+
+if ((overall_status != 0)); then
+    echo "One or more training tasks failed. Inspect logs under: $LOG_ROOT" >&2
+    exit 1
+fi
+
+if ((DRY_RUN)); then
+    echo "Dry run completed successfully."
+else
+    echo "All assessment_only/all training tasks completed successfully."
+fi