Skip to content

Add survival and competing-risk HTE learners#1027

Open
xushenbo wants to merge 2 commits intopy-why:mainfrom
xushenbo:censor-hte-pr
Open

Add survival and competing-risk HTE learners#1027
xushenbo wants to merge 2 commits intopy-why:mainfrom
xushenbo:censor-hte-pr

Conversation

@xushenbo
Copy link
Copy Markdown

Add survival and competing-risk HTE learners

Summary

This PR adds heterogeneous treatment effect support for censored outcomes to EconML, including both survival and competing-risk settings.

The implementation follows the censoring-unbiased transformation (CUT) framework, which converts censored time-to-event outcomes into continuous transformed outcomes so that HTE learners for continuous outcomes can be applied consistently in censored settings.

Motivation

This contribution is based on my work on orthogonal statistical learning for survival and competing-risk outcomes, including the CUT framework developed in:

  • "Estimating Heterogeneous Treatment Effects on Survival Outcomes Using Counterfactual Censoring Unbiased Transformations"

The goal is to extend EconML's reach beyond continuous and binary outcomes to time-to-event settings that are common in biostatistics and medical research.

What this PR adds

Censored-outcome functionality

  • nuisance estimation helpers for survival and competing-risk settings under econml.censor
  • RMST and RMTL transformation functions, including IPCW, BJ, AIPCW, and UIF-based transformations
  • support for separable direct and indirect effects for competing risks

Learners

  • survival and competing-risk HTE learners under econml.metalearners
  • leading learner variants rather than a single method, including:
    • direct survival learners
    • direct competing-risk learners
    • CUT-based TLearner / SLearner / XLearner / IPTW / AIPTW / MC / MCEA / U / R / IF learners
    • separable direct and indirect competing-risk learners

Forest-based methods

  • CausalSurvivalForest
  • SurvivalForest
  • GRF-style causal forest wrappers used by the censored-outcome workflow
  • these forest-based methods are translated from grf-master

Documentation and examples

  • survival estimation documentation page
  • competing-risks estimation documentation page
  • survival HTE notebook
  • competing-risks HTE notebook
  • real-data bone marrow transplant case study notebook, with data and data dictionary

Validation

I ran the following checks:

  • python3 -m unittest econml.tests.test_censor.test_crossfit_meta -v
  • python3 -m unittest econml.tests.test_censor.test_simulation -q
  • python3 -m unittest econml.tests.test_censor.test_survival_meta econml.tests.test_censor.test_competing_meta -q
  • LC_ALL=C LANG=C PYTHONPATH=/tmp/econml-docs-env python3 -m sphinx -W -E -a -b html doc build/sphinx/html

Notes

  • This PR includes notebooks and the supporting bone marrow transplant data assets used by the case study notebook.
  • The implementation is designed to follow EconML's existing structure while adding censored-outcome support as a natural extension of the package.

xushenbo and others added 2 commits April 16, 2026 08:53
Signed-off-by: Shenbo Xu <xushenbo@mit.edu>
Signed-off-by: Shenbo Xu <27264836+xushenbo@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant