diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py index 29c2a922..de9e3863 100644 --- a/causalml/inference/meta/rlearner.py +++ b/causalml/inference/meta/rlearner.py @@ -16,7 +16,7 @@ get_xgboost_objective_metric, get_weighted_variance, ) -from causalml.propensity import ElasticNetPropensityModel +from causalml.propensity import ElasticNetPropensityModel, compute_r_residuals logger = logging.getLogger("causalml") @@ -140,10 +140,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True): if self.effect_learner is not None else deepcopy(self.learner) ) - # Build CV splitter from stored n_fold / random_state. - self.cv = KFold( - n_splits=self.n_fold, shuffle=True, random_state=self.random_state - ) self.models_tau = {group: deepcopy(self.model_tau) for group in self.t_groups} self.vars_c = {} @@ -151,9 +147,17 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True): if verbose: logger.info("generating out-of-fold CV outcome estimates") - yhat = cross_val_predict( - self.model_mu, X, y_np, cv=self.cv, n_jobs=self.cv_n_jobs + y_residual, _ = compute_r_residuals( + X, + treatment_np, + y_np, + outcome_learner=self.model_mu, + n_folds=self.n_fold, + random_state=self.random_state, + n_jobs=self.cv_n_jobs, + compute_w_residual=False, ) + yhat = y_np - y_residual # Fit the nuisance outcome model on the full data so it can be # reused by predict(return_components=True). self.model_mu.fit(X, y_np) diff --git a/causalml/metrics/__init__.py b/causalml/metrics/__init__.py index f05ebf97..0c7a05cd 100644 --- a/causalml/metrics/__init__.py +++ b/causalml/metrics/__init__.py @@ -30,6 +30,7 @@ compute_dr_pseudo_outcomes, dr_score, plug_in_t_score, + rlearner_score, ) # noqa from .sensitivity import Sensitivity, SensitivityPlaceboTreatment # noqa from .sensitivity import ( diff --git a/causalml/metrics/cate_scoring.py b/causalml/metrics/cate_scoring.py index c09a363a..5156a71b 100644 --- a/causalml/metrics/cate_scoring.py +++ b/causalml/metrics/cate_scoring.py @@ -6,7 +6,7 @@ from scipy import stats from sklearn.model_selection import StratifiedKFold -from causalml.propensity import compute_propensity_score +from causalml.propensity import compute_propensity_score, compute_r_residuals logger = logging.getLogger("causalml") @@ -171,33 +171,19 @@ def compute_dr_pseudo_outcomes( return phi -def _score_against_pseudo_outcome( - df, - pseudo_outcome, - model_cols, - score_name, - return_ci, - n_bootstrap, - alpha, - random_state, -): - """Shared MSE-against-pseudo-outcome scorer with half-sample bootstrap confidence intervals. +def _bootstrap_loss_ci(sq_err, score_name, return_ci, n_bootstrap, alpha, random_state): + """Half-sample bootstrap confidence intervals over a precomputed squared-error frame. - This mirrors the half-sample bootstrap confidence interval pattern in - ``rate.py``'s ``rate_score()`, but resamples the already-computed (model - prediction, pseudo-outcome) pairs rather than re-fitting any nuisance models - -- nuisance fitting happens once, upstream, in ``compute_dr_pseudo_outcomes()`` - or the plug-in T-learner cross-fit. + Shared by dr_score/plug_in_t_score (sq_err = (tau_hat - pseudo_outcome) ** 2) + and rlearner_score (sq_err = (y_residual - w_residual * tau_hat) ** 2) -- the + loss formula differs per metric, but once reduced to a per-model squared + error the bootstrap procedure is identical. - Lower scores are better: this is a loss (mean squared error against a CATE - proxy), not a similarity score. + Lower scores are better in all cases: this is a loss, not a similarity score. Args: - df (pandas.DataFrame): a data frame with model CATE estimates as columns - pseudo_outcome (numpy.ndarray): the CATE proxy to score models against, - one value per row of ``df`` - model_cols (list): the columns of ``df`` holding model CATE estimates - score_name (str): name used for the returned Series/column (e.g. ``"dr_loss"``) + sq_err (pandas.DataFrame): per-row, per-model squared error + score_name (str): name used for the returned Series/column return_ci (bool): whether to return bootstrap confidence intervals n_bootstrap (int): number of half-sample bootstrap iterations alpha (float): significance level for confidence intervals @@ -205,18 +191,16 @@ def _score_against_pseudo_outcome( Returns: If return_ci=False: (pandas.Series): loss for each model column - If return_ci=True: - (pandas.DataFrame): loss, standard error, and confidence interval - bounds for each model column + If return_ci=True: (pandas.DataFrame): loss, se, and CI bounds per model column """ - sq_err = (df[model_cols].sub(pseudo_outcome, axis=0)) ** 2 + model_cols = list(sq_err.columns) loss = sq_err.mean(axis=0) loss.name = score_name if not return_ci: return loss - n = len(df) + n = len(sq_err) m = n // 2 rng = np.random.default_rng(random_state) boot_losses = {model: [] for model in model_cols} @@ -232,21 +216,49 @@ def _score_against_pseudo_outcome( point = loss[model] boot = np.array(boot_losses[model]) se = np.std(boot, ddof=1) - ci_lower = point - z_crit * se - ci_upper = point + z_crit * se results.append( { "model": model, score_name: point, "se": se, - "ci_lower": ci_lower, - "ci_upper": ci_upper, + "ci_lower": point - z_crit * se, + "ci_upper": point + z_crit * se, } ) return pd.DataFrame(results).set_index("model") +def _score_against_pseudo_outcome( + df, + pseudo_outcome, + model_cols, + score_name, + return_ci, + n_bootstrap, + alpha, + random_state, +): + """MSE-against-pseudo-outcome loss, for dr_score / plug_in_t_score. + + See _bootstrap_loss_ci for the shared bootstrap-CI mechanics. + + Args: + df (pandas.DataFrame): a data frame with model CATE estimates as columns + pseudo_outcome (numpy.ndarray): the CATE proxy to score models against + model_cols (list): the columns of df holding model CATE estimates + score_name, return_ci, n_bootstrap, alpha, random_state: see + _bootstrap_loss_ci + + Returns: + See _bootstrap_loss_ci. + """ + sq_err = (df[model_cols].sub(pseudo_outcome, axis=0)) ** 2 + return _bootstrap_loss_ci( + sq_err, score_name, return_ci, n_bootstrap, alpha, random_state + ) + + def dr_score( df, X=None, @@ -479,3 +491,107 @@ def plug_in_t_score( alpha=alpha, random_state=random_state, ) + + +def rlearner_score( + df, + X=None, + treatment_col="w", + outcome_col="y", + y_residual_col=None, + w_residual_col=None, + outcome_learner=None, + propensity_learner=None, + n_folds=5, + return_ci=False, + n_bootstrap=200, + alpha=0.05, + random_state=None, +): + """Score fitted CATE models via the R-loss (Nie & Wager, 2021). + + R-loss(tau_hat) = mean[((y - m(X)) - (w - e(X)) * tau_hat(X)) ** 2] + + where m(X) = E[Y|X] and e(X) = E[W|X] are cross-fitted nuisance regressions + (see causalml.propensity.compute_r_residuals()). This is the loss + BaseRLearner.fit() already minimizes internally to train its own effect + model; exposing it standalone gives R-loss-based comparison across + arbitrary fitted CATE models -- EconML RScorer parity. Lower is better. + + R-score complements dr_score() and plug_in_t_score() on the CATE-accuracy + axis (as opposed to rate_score()'s targeting/ranking axis); Mahajan et al. + (2024) found DR-loss dominates and plug-in-T is never dominated, with + R-loss not the standout of the three -- useful as a third opinion, + particularly for parity with EconML workflows already using RScorer. + + Residuals can be supplied directly (e.g. precomputed once with + compute_r_residuals() and reused across scoring calls) via + y_residual_col / w_residual_col, or computed internally from X, + treatment_col, and outcome_col. + + Args: + df (pandas.DataFrame): a data frame with fitted CATE model estimates as + columns, plus either y_residual_col/w_residual_col or both + outcome_col and treatment_col + X (numpy.ndarray or pandas.DataFrame, optional): feature matrix for the + R-loss nuisance models. Required unless residual columns are given + treatment_col (str, optional): treatment indicator column (0 or 1). + Ignored if residual columns are provided + outcome_col (str, optional): outcome column. Ignored if residual + columns are provided + y_residual_col (str, optional): precomputed y - m_hat(X) column + w_residual_col (str, optional): precomputed w - e_hat(X) column + outcome_learner (model, optional): model for E[Y|X]. Required unless + residual columns are provided + propensity_learner (PropensityModel, optional): passed to + compute_r_residuals(). Defaults to ElasticNetPropensityModel + n_folds (int, optional): cross-fitting folds. Default 5 + return_ci (bool, optional): whether to return bootstrap CIs. Default False + n_bootstrap (int, optional): half-sample bootstrap iterations. Default 200 + alpha (float, optional): CI significance level. Default 0.05 + random_state (int or None, optional): random seed. Default None + + Returns: + If return_ci=False: (pandas.Series): R-loss for each model column (lower is better) + If return_ci=True: (pandas.DataFrame): R-loss, se, and CI bounds per model column + """ + have_residuals = ( + y_residual_col is not None + and y_residual_col in df.columns + and w_residual_col is not None + and w_residual_col in df.columns + ) + assert have_residuals or (X is not None and outcome_learner is not None), ( + "Either `y_residual_col`/`w_residual_col` (present in df) or `X` and " + "`outcome_learner` (to compute residuals internally) must be provided." + ) + + model_cols = [ + c + for c in df.columns + if c not in (outcome_col, treatment_col, y_residual_col, w_residual_col) + ] + + if have_residuals: + y_residual = df[y_residual_col].to_numpy() + w_residual = df[w_residual_col].to_numpy() + else: + assert ( + outcome_col in df.columns and treatment_col in df.columns + ), "{} and {} must be present in df to compute R-loss residuals.".format( + outcome_col, treatment_col + ) + y_residual, w_residual = compute_r_residuals( + X=X, + treatment=df[treatment_col], + y=df[outcome_col], + outcome_learner=outcome_learner, + propensity_learner=propensity_learner, + n_folds=n_folds, + random_state=random_state, + ) + + sq_err = (df[model_cols].mul(w_residual, axis=0).sub(y_residual, axis=0)) ** 2 + return _bootstrap_loss_ci( + sq_err, "r_loss", return_ci, n_bootstrap, alpha, random_state + ) diff --git a/causalml/propensity.py b/causalml/propensity.py index fc0e0640..3932f65a 100644 --- a/causalml/propensity.py +++ b/causalml/propensity.py @@ -1,9 +1,10 @@ from abc import ABCMeta, abstractmethod +from copy import deepcopy import logging import numpy as np from sklearn.metrics import roc_auc_score as auc from sklearn.linear_model import LogisticRegressionCV -from sklearn.model_selection import StratifiedKFold, train_test_split +from sklearn.model_selection import StratifiedKFold, cross_val_predict, train_test_split from sklearn.isotonic import IsotonicRegression import xgboost as xgb @@ -230,3 +231,110 @@ def compute_propensity_score( p = p_model.predict(X_pred) return p, p_model + + +def compute_r_residuals( + X, + treatment, + y, + outcome_learner, + propensity_learner=None, + p=None, + method="predict", + n_folds=5, + random_state=None, + n_jobs=-1, + compute_w_residual=True, +): + """Cross-fitted outcome/treatment residuals for the R-loss (Nie & Wager, 2021). + + Computes out-of-fold m_hat(X) = E[Y|X] and e_hat(X) = E[W|X] via + n_folds-fold cross-fitting, stratified on treatment so every fold retains + both arms, and returns: + + y_residual = y - m_hat(X) + w_residual = w - e_hat(X) + + A candidate CATE model tau_hat is scored against these via the R-loss: + + R-loss(tau_hat) = mean[(y_residual - w_residual * tau_hat(X)) ** 2] + + This is also the quantity BaseRLearner.fit() implicitly minimizes: fitting + the per-arm effect model against target (y_residual / w_residual) with + sample_weight = w_residual ** 2 is the weighted-least-squares solution to + the same R-loss objective. + + Args: + X (numpy.ndarray or pandas.DataFrame): a feature matrix + treatment (numpy.ndarray or pandas.Series): a binary treatment indicator (0 or 1) + y (numpy.ndarray or pandas.Series): an outcome vector + outcome_learner (model): a model to estimate E[Y|X]. Must implement + fit/predict (or predict_proba if method="predict_proba") + propensity_learner (PropensityModel, optional): passed through to + compute_propensity_score(). Ignored if `p` is given. + Defaults to ElasticNetPropensityModel + p (numpy.ndarray or pandas.Series, optional): pre-computed propensity + scores. If given, propensity is not re-estimated in-fold + method (str, optional): "predict" or "predict_proba" (for classifier + outcome learners, e.g. BaseRClassifier). Only the positive-class + column is used for "predict_proba". Default "predict" + n_folds (int, optional): number of cross-fitting folds. Default 5 + random_state (int or None, optional): random seed for the fold splitter + n_jobs (int, optional): parallel jobs forwarded to cross_val_predict + for the outcome model. Default -1 + compute_w_residual (bool, optional): whether to compute and return + w_residual. If False, skips propensity estimation entirely (no + in-fold propensity model is fit) and returns w_residual=None. + Set False when only the outcome residual is needed -- e.g. + BaseRLearner.fit(), which already has propensity scores from + elsewhere and would otherwise pay for a redundant per-fold + propensity fit whose output is discarded. Default True. + + Returns: + (tuple): + - y_residual (numpy.ndarray): y - m_hat(X), out-of-fold + - w_residual (numpy.ndarray or None): w - e_hat(X), out-of-fold + (or w - p directly if `p` was supplied), or None if + compute_w_residual=False + """ + X = np.asarray(X) + treatment = np.asarray(treatment) + y = np.asarray(y, dtype=float) + + splits = list( + StratifiedKFold( + n_splits=n_folds, shuffle=True, random_state=random_state + ).split(X, treatment) + ) + + if method == "predict_proba": + yhat = cross_val_predict( + outcome_learner, X, y, cv=splits, method="predict_proba", n_jobs=n_jobs + )[:, 1] + else: + yhat = cross_val_predict(outcome_learner, X, y, cv=splits, n_jobs=n_jobs) + + y_residual = y - yhat + + if not compute_w_residual: + return y_residual, None + + if p is not None: + w_residual = treatment - np.asarray(p, dtype=float) + else: + w_residual = np.empty(len(treatment), dtype=float) + for train_idx, test_idx in splits: + p_test, _ = compute_propensity_score( + X=X[train_idx], + treatment=treatment[train_idx], + p_model=( + deepcopy(propensity_learner) + if propensity_learner is not None + else None + ), + X_pred=X[test_idx], + treatment_pred=treatment[test_idx], + ) + w_residual[test_idx] = treatment[test_idx] - p_test + + return y_residual, w_residual diff --git a/tests/test_cate_scoring.py b/tests/test_cate_scoring.py index 1e21ebb8..adb3e37f 100644 --- a/tests/test_cate_scoring.py +++ b/tests/test_cate_scoring.py @@ -7,7 +7,9 @@ compute_dr_pseudo_outcomes, dr_score, plug_in_t_score, + rlearner_score, ) +from causalml.propensity import compute_r_residuals from causalml.metrics.rate import rate_score try: @@ -311,3 +313,213 @@ def test_dr_score_and_plug_in_t_score_all_finite(synthetic_data): ) assert np.isfinite(dr.values).all() assert np.isfinite(t.values).all() + + +def test_compute_r_residuals_shape(synthetic_data): + df, X = synthetic_data + y_residual, w_residual = compute_r_residuals( + X, + df["w"], + df["y"], + outcome_learner=LinearRegression(), + n_folds=5, + random_state=RANDOM_SEED, + ) + assert y_residual.shape == (len(df),) + assert w_residual.shape == (len(df),) + assert np.isfinite(y_residual).all() + assert np.isfinite(w_residual).all() + + +def test_rlearner_score_returns_series(synthetic_data): + df, X = synthetic_data + scores = rlearner_score( + df, + X=X, + treatment_col="w", + outcome_col="y", + outcome_learner=LinearRegression(), + random_state=RANDOM_SEED, + ) + assert isinstance(scores, pd.Series) + assert set(scores.index) == {"perfect_model", "noisy_model", "bad_model"} + + +def test_rlearner_score_ranks_perfect_model_lowest(synthetic_data): + df, X = synthetic_data + scores = rlearner_score( + df, + X=X, + treatment_col="w", + outcome_col="y", + outcome_learner=LinearRegression(), + random_state=RANDOM_SEED, + ) + assert scores["perfect_model"] < scores["noisy_model"] < scores["bad_model"] + + +def test_rlearner_score_with_precomputed_residuals(synthetic_data): + df, X = synthetic_data + y_residual, w_residual = compute_r_residuals( + X, + df["w"], + df["y"], + outcome_learner=LinearRegression(), + n_folds=5, + random_state=RANDOM_SEED, + ) + df_with_residuals = df.assign(y_resid=y_residual, w_resid=w_residual) + scores = rlearner_score( + df_with_residuals, + y_residual_col="y_resid", + w_residual_col="w_resid", + random_state=RANDOM_SEED, + ) + assert scores["perfect_model"] < scores["bad_model"] + + +def test_rlearner_score_missing_X_and_residuals_raises(synthetic_data): + df, _ = synthetic_data + with pytest.raises(AssertionError): + rlearner_score(df, treatment_col="w", outcome_col="y") + + +def test_rlearner_score_return_ci_returns_dataframe(synthetic_data): + df, X = synthetic_data + result = rlearner_score( + df, + X=X, + treatment_col="w", + outcome_col="y", + outcome_learner=LinearRegression(), + return_ci=True, + n_bootstrap=50, + random_state=RANDOM_SEED, + ) + assert isinstance(result, pd.DataFrame) + assert set(result.columns) == {"r_loss", "se", "ci_lower", "ci_upper"} + + +def test_rlearner_score_ci_bounds_ordered(synthetic_data): + df, X = synthetic_data + result = rlearner_score( + df, + X=X, + treatment_col="w", + outcome_col="y", + outcome_learner=LinearRegression(), + return_ci=True, + n_bootstrap=50, + random_state=RANDOM_SEED, + ) + assert (result["ci_lower"] < result["r_loss"]).all() + assert (result["ci_upper"] > result["r_loss"]).all() + + +def test_compute_r_residuals_handles_imbalanced_treatment(): + rng = np.random.default_rng(RANDOM_SEED) + + n, p = 200, 5 + X = rng.normal(size=(n, p)) + + w = np.zeros(n, dtype=int) + w[rng.choice(n, 10, replace=False)] = 1 + + y = rng.normal(size=n) + + y_residual, w_residual = compute_r_residuals( + X, + w, + y, + outcome_learner=LinearRegression(), + n_folds=5, + random_state=RANDOM_SEED, + ) + + assert y_residual.shape == (n,) + assert w_residual.shape == (n,) + assert np.isfinite(y_residual).all() + assert np.isfinite(w_residual).all() + + +def test_rlearner_score_handles_imbalanced_treatment(): + rng = np.random.default_rng(RANDOM_SEED) + + n, p = 200, 5 + X = rng.normal(size=(n, p)) + + w = np.zeros(n, dtype=int) + w[rng.choice(n, 10, replace=False)] = 1 + + y = rng.normal(size=n) + + df = pd.DataFrame( + { + "y": y, + "w": w, + "model_1": rng.normal(size=n), + "model_2": rng.normal(size=n), + } + ) + + scores = rlearner_score( + df, + X=X, + treatment_col="w", + outcome_col="y", + outcome_learner=LinearRegression(), + n_folds=5, + random_state=RANDOM_SEED, + ) + + assert isinstance(scores, pd.Series) + assert np.isfinite(scores.values).all() + + +def test_dr_plug_in_t_and_rlearner_score_all_finite(synthetic_data): + df, X = synthetic_data + dr = dr_score( + df, + X=X, + treatment_col="w", + outcome_col="y", + learner=LinearRegression(), + random_state=RANDOM_SEED, + ) + t = plug_in_t_score( + df, + X, + treatment_col="w", + outcome_col="y", + learner=LinearRegression(), + random_state=RANDOM_SEED, + ) + r = rlearner_score( + df, + X=X, + treatment_col="w", + outcome_col="y", + outcome_learner=LinearRegression(), + random_state=RANDOM_SEED, + ) + assert np.isfinite(dr.values).all() + assert np.isfinite(t.values).all() + assert np.isfinite(r.values).all() + + +def test_compute_r_residuals_skips_propensity_when_w_residual_not_needed( + synthetic_data, +): + df, X = synthetic_data + y_residual, w_residual = compute_r_residuals( + X, + df["w"], + df["y"], + outcome_learner=LinearRegression(), + n_folds=5, + random_state=RANDOM_SEED, + compute_w_residual=False, + ) + assert y_residual.shape == (len(df),) + assert np.isfinite(y_residual).all() + assert w_residual is None