ADMM Hazard-AFT Model

ADMM（交互方向乗数法）を用いた Fused Lasso 正則化付き Hazard-AFT モデルの実装。

概要

本リポジトリは、ADMM（Alternating Direction Method of Multipliers）を用いて fused lasso（時間方向の差分に対する L1 正則化）付きの Hazard-AFT モデルを推定するための Python 実装です。

時間区分ごとに回帰係数が変化する生存時間分析モデルを、Fused Lasso により時間方向のスパース性を導入しながら推定します。

• モデル: Hazard-based AFT（Accelerated Failure Time）
• 正則化: Fused Lasso（区間間の差分に L1 ペナルティ）
• 最適化: ADMM（Alternating Direction Method of Multipliers）
• ベースラインハザード: B-spline 基底展開

インストール

# uvを使った環境構築
uv sync

# または直接実行
uv run main.py --config config.toml --data data/simulated_data.csv

基本的な使い方

データ生成（段階的βの拡張AFT）

段階的に係数が変化する拡張AFTデータを生成します。

uv run generation/extended_aft_step_generator.py \
  --config generation/extended_aft_step_generator.config.json \
  --output data/simulated_data.csv

設定は generation/extended_aft_step_generator.config.json を参照してください。

stepwise_beta.true_time_grid はデータ生成時の真の係数関数、 analysis_time_grid は推定モデルの時間区分を定義します。両者を分けることで、 Fine-grid／Off-grid設定を生成できます。旧 stepwise_beta.time_grid も引き続き利用できます。

Lambda並列実験用の学習・独立評価データは、同じファイル名で別ディレクトリへ一括生成します。

uv run generation/generate_extended_aft_step_datasets.py \
  --output-dir data/extended_aft_step \
  --eval-output-dir data/extended_aft_step_eval

評価データは既定で学習seedに100000を加えて生成されます。

単一実験の実行

uv run main.py --config config.toml --data data/simulated_data.csv --output result.json --plot

Lambda並列実験（スパコン環境）

Lambda値を変えながら複数のデータセットで並列実験を行います。

1. Lambda値の設定

lambda_grid.json で実験する lambda値を定義：

{
  "description": "Lambda values for parallel experiments (0 + log scale: 1 to 10, 10 points)",
  "lambda_values": [0.0, 1.0, 1.333521432163324, 1.7782794100389228, 2.371373705661655, 3.1622776601683795, 4.216965034285822, 5.623413251903491, 7.498942093324558, 10.0]
}

2. ジョブ投入

# UGEアレイジョブとして投入（データ100件 × lambda10点 = 1000パターン）
qsub qsub.sh

qsub.sh は SGE_TASK_ID を使って自動的に以下を切り替えます：

• 処理するデータファイル（data/extended_aft_step/*.csv）
• 対応する独立評価データ（data/extended_aft_step_eval/*.csv）
• 使用する lambda値（lambda_grid.json から選択）

3. 結果の集計

実験完了後、結果を集計：

uv run scripts/aggregate_lambda_results.py --base-dir outputs/lambda_experiments --output outputs/lambda_summary.csv

4. 結果の可視化

uv run scripts/visualize_lambda_results.py --summary outputs/lambda_summary.csv --output-dir outputs/lambda_plots

5. Cox回帰との C-index 比較

まず Cox 回帰を各データで学習し、c_td と Harrell C-index を計算します。

uv run scripts/compute_cox_metrics.py \
  --data-dir data/extended_aft_step \
  --output outputs/cox_summary.csv \
  --lambda-summary outputs/lambda_summary.csv \
  --compare-output outputs/cox_vs_lambda_c_td.csv

次に、既存の lambda_vs_c_td.png に加えて、Cox 基準線付きの比較図を生成します。

uv run scripts/visualize_lambda_results.py \
  --summary outputs/lambda_summary.csv \
  --output-dir outputs/lambda_plots \
  --cox-summary outputs/cox_summary.csv

追加で生成されるファイル：

• lambda_vs_c_td_with_cox.png: ADMM の lambda ごとの c_td 分布に Cox の c_td 基準線（平均/IQR）を重ねた図

生成されるプロット：

• lambda_vs_objective.png: Lambda値と目的関数の関係
• lambda_distribution.png: Lambda値ごとの目的関数分布（箱ひげ図）
• lambda_vs_convergence.png: Lambda値と収束状況（primal/dual residual）
• lambda_vs_c_td_with_cox.png: Cox の c_td 基準線付き比較図（--cox-summary 指定時）

集計CSVの c_td は独立評価データ上の値です。c_td_train、c_td_test も 明示的に保存されます。BICの自由度は n_baseline_basis + n_features + n_change_points として計算されます。

ローカルでのテスト実行

# 特定のタスクIDを指定して単一実験を実行
./run_lambda_experiment.sh 1

# データインデックスとlambdaインデックスの対応：
# task_id = data_idx * n_lambda + lambda_idx + 1
# 例: task_id=1 → data=1, lambda=1
#     task_id=11 → data=2, lambda=1
#     task_id=101 → data=11, lambda=1

ディレクトリ構造

.
├── admm/                  # コアパッケージ
│   ├── model.py          # ADMMHazardAFT推定器
│   ├── solver.py         # ADMMソルバ
│   ├── objective.py      # 目的関数・勾配・ヘッセ
│   ├── baseline.py       # ベースラインハザード（B-spline）
│   ├── time_partition.py # 時間分割・η計算
│   ├── quadrature.py     # 求積ルール
│   └── logger.py         # WandBロガー
├── scripts/              # ユーティリティスクリプト
│   ├── aggregate_lambda_results.py  # 結果集計
│   ├── compute_cox_metrics.py       # Cox回帰の指標計算
│   └── visualize_lambda_results.py  # 結果可視化
├── data/
│   └── extended_aft/     # データセット（CSV）
├── outputs/
│   ├── lambda_experiments/  # Lambda実験結果
│   │   └── {data_name}/
│   │       └── lambda_{value}/
│   │           ├── config.toml
│   │           └── result.json
│   ├── lambda_summary.csv  # 集計結果
│   └── lambda_plots/       # 可視化プロット
├── config.toml           # デフォルト設定
├── lambda_grid.json      # Lambda並列実験用の値リスト
├── qsub.sh              # UGEジョブスクリプト
├── run_lambda_experiment.sh  # Lambda実験実行スクリプト
└── main.py              # CLIエントリポイント

設定ファイル（config.toml）

主要なハイパーパラメータ：

time_grid = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]  # 時間分割
lambda_fuse = 1.0        # Fused Lasso強度
rho = 1.0                # ADMMペナルティ係数
max_admm_iter = 100      # ADMM最大反復数
clip_eta = 5.0           # exp(η)の発散抑制
n_baseline_basis = 8     # B-spline基底数

[quadrature]
Q = 5                    # 求積点数
rule = "gauss_legendre"  # 求積法

1人1行の入力

FraminghamでKとlambdaを同時に選ぶスパコン実験は K・lambda同時探索の手順 を参照してください。

時間によって共変量が変わらないデータは、k 列なしの1人1行CSVで main.py に渡せます。必須列は id, time, event で、idは重複できません。 特徴量から除外する列は既存long formatと同じです。

id,time,event,AGE,SEX,BMI,SYSBP,DIABP
2448,6.0,0,-1.09396,1,0.55252,-1.25089,-0.96352

uv run main.py --config config.toml --data subjects.csv --output result.json

Kは len(time_grid) - 1 で決まり、学習時に同じ共変量を全区間へ展開します。 ADMMHazardAFT.fit() も2次元 (n, p) を受け付けます。 上記の [0, 1, 2, 3, 4, 5, 6] はK=6です。 同じ観測期間でKだけを変更する場合は、CSVの時間スケールを固定し、configの time_grid の内部の分割点だけを変更します。1人1行CSVでは横の meta.json の time_grid はconfigを上書きしません。

--eval-data と予測専用の --load-result でも1人1行CSVを使えます。 既存の k 列付きlong formatも利用できますが、学習入力のKとconfigの 区間数の不一致はエラーになります。従来の自動切り詰めは行いません。 long formatでは従来どおりCSV横の meta.json の time_grid が優先され、 予測専用実行では保存したモデルの時間分割を使います。

従来の実データCVとブートストラップCLIはlong formatを使います。 新しいK・lambda同時探索は1人1行を使います。 既存のCSVや実験結果を変換・上書きする必要はありません。

テスト

# 全テスト実行
uv run pytest tests/

# 特定テスト実行
uv run pytest tests/test_admm_fit_smoke.py -v

WandBロギング（オプション）

# 環境変数でWandBを有効化
export WANDB_PROJECT=admm-experiments
export WANDB_ENABLED=true

uv run main.py --config config.toml --data data/simulated_data.csv

実装状況

✅ 実装済み：

• ADMMソルバ本体
• 目的関数（近似対数尤度）の勾配・ヘッセ
• B-splineベースライン
• Fused Lasso正則化
• Lambda並列実験フレームワーク
• 結果集計・可視化スクリプト

🚧 今後の実装：

• Brier score など追加の評価指標
• 適応的ρ調整
• M/I-splineベースライン（積分の解析的計算）

Ctd評価（time-dependent C-index）

本実装では、以下の定義に基づく c_td を評価指標として計算します。

$$ C^{td}

\Pr!\left( S(T_i \mid X_i(t)) < S(T_i \mid X_j(t)) ;\middle|; T_i < T_j,; D_i = 1 \right) $$

• 比較対象ペア: $T_i &lt; T_j,; D_i=1$
• 判定: $S(T_i \mid X_i(t)) &lt; S(T_i \mid X_j(t))$
• 同値（tie）: 0.5 点

main.py --output の result.json では、summary.c_td として出力されます。

参考文献

• Pang et al. (2021). "Flexible Extension of the Accelerated Failure Time Model to Account for Nonlinear and Time-Dependent Effects of Covariates on the Hazard." Statistical Methods in Medical Research, 30(11), 2526–42.
• Boyd et al. (2011). "Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers." Foundations and Trends in Machine Learning, 3(1), 1–122.