Add DR and plug-in T scoring metrics for CATE model evaluation#930
Conversation
jeongyoonlee
left a comment
There was a problem hiding this comment.
The core AIPW/T-learner construction and the docs are solid, but one blocker before merge:
Both compute_dr_pseudo_outcomes and plug_in_t_score can crash on imbalanced treatment — plain KFold doesn't stratify.
Each fold fits per-arm outcome models on X[train_idx][w_train == 0] and X[train_idx][w_train == 1]. With KFold(shuffle=True) (not StratifiedKFold), a fold's training split can end up with zero units of one arm, so .fit() on the empty slice raises:
ValueError: Found array with 0 sample(s) (shape=(0, n_features)) while a minimum of 1 is required
Reproduced on rare-treatment data — which is exactly the observational setting these surrogate metrics target — at 3 treated / n=200 / 5 folds (crashes on 1 of 50 seeds):
w = np.zeros(200, int); w[rng.choice(200, 3, replace=False)] = 1
compute_dr_pseudo_outcomes(X, w, y, learner=LinearRegression(), n_folds=5, random_state=15)
# ValueError: Found array with 0 sample(s) ...Please switch both functions to StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=...) and split on the treatment vector (cv.split(X, treatment)), so every fold retains both arms. That's the standard splitter for cross-fitting AIPW / T-learner nuisances.
Non-blocking notes to follow separately.
3a60762 to
aa10d2a
Compare
jeongyoonlee
left a comment
There was a problem hiding this comment.
Two non-blocking nits:
-
p_valuecolumn (_score_against_pseudo_outcome):z_stat = point / setests H0: loss = 0, but the loss is an MSE that is ~never 0, so the p-value collapses to ≈0 for every model and conveys nothing. This is meaningful inrate_score(H0: RATE = 0 = no targeting value) but not for a loss — the informative test is a paired model-vs-model comparison. Suggest dropping thep_valuecolumn or replacing it with the paired test. -
metrics/__init__.pyhas a duplicatedfrom .sensitivity import Sensitivity, SensitivityPlaceboTreatment # noqaline (a second copy of an already-present import).
jeongyoonlee
left a comment
There was a problem hiding this comment.
The p_value removal missed the docstrings — they still list a p_value column for return_ci=True. Please scrub the remaining references in cate_scoring.py:
_score_against_pseudo_outcome— thereturn_ciarg ("…confidence intervals and p-values") and the Returns ("…CI bounds, and p-value (H0: loss = 0)…").dr_score— thereturn_ciarg and the Returns (same two).plug_in_t_score— thereturn_ciarg and the Returns (same two).
Also the _score_against_pseudo_outcome summary line still says it "mirrors the half-sample bootstrap CI/p-value pattern in rate.py" — drop the p-value half.
Proposed changes
this pr adds two new observed-data surrogate metrics for evaluating fitted cate models without access to true counterfactual outcomes
what's included
compute_dr_pseudo_outcomes()to construct cross-fitted doubly robust (aipw) pseudo-outcomes that can be reused across different evaluation methodsdr_score()to evaluate cate models using mean squared error against the dr pseudo-outcomesplug_in_t_score()to evaluate cate models against a cross-fitted plug-in t-learner cate proxyrate_score()to avoid fitting nuisance models multiple timesrate_score()implementationrate_score()closes #922
Types of changes
What types of changes does your code introduce to CausalML?
Put an
xin the boxes that applyChecklist
Put an
xin the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your code.Further comments
If this is a relatively large or complex change, kick off the discussion by explaining why you chose the solution you did and what alternatives you considered, etc. This PR template is adopted from appium.