From a8c9d2d21ee4c64dc77dd7427f2c6da2793c188b Mon Sep 17 00:00:00 2001 From: Oren Tirschwell Date: Thu, 11 Jun 2026 23:37:54 -0400 Subject: [PATCH] Add support for user-supplied propensities in DML and DRLearner estimators Signed-off-by: Oren Tirschwell --- doc/spec/estimation/dml.rst | 33 ++- doc/spec/estimation/dr.rst | 35 +++ econml/_ortho_learner.py | 111 +++++++- econml/dml/_rlearner.py | 73 ++++- econml/dml/causal_forest.py | 17 +- econml/dml/dml.py | 85 +++++- econml/dr/_drlearner.py | 108 +++++++- econml/tests/test_user_propensity.py | 389 +++++++++++++++++++++++++++ 8 files changed, 806 insertions(+), 45 deletions(-) create mode 100644 econml/tests/test_user_propensity.py diff --git a/doc/spec/estimation/dml.rst b/doc/spec/estimation/dml.rst index 74a192213..1baa1b5b4 100644 --- a/doc/spec/estimation/dml.rst +++ b/doc/spec/estimation/dml.rst @@ -549,7 +549,7 @@ Usage FAQs - **What if my treatment is categorical/binary?** - You can simply set `discrete_treatment=True` in the parameters of the class. Then use any classifier for + You can simply set `discrete_treatment=True` in the parameters of the class. Then use any classifier for `model_t`, that has a `predict_proba` method: .. testcode:: @@ -561,6 +561,37 @@ Usage FAQs point = est.const_marginal_effect(X) est.effect(X, T0=t0, T1=t1) +- **What if I already know the treatment assignment probabilities (e.g. a randomized experiment)?** + + If the treatment is discrete and the assignment probabilities are known by design (e.g. data from an + A/B test or a stratified experiment with per-unit assignment probabilities), you can pass them + directly via the ``propensity`` argument of ``fit`` instead of having the library estimate a + treatment model; the first stage treatment model is then not fitted at all: + + .. testcode:: + + from econml.dml import LinearDML + p = np.full((y.shape[0], 3), 1 / 3) # known probabilities of the three treatment values + est = LinearDML(discrete_treatment=True) + est.fit(y, t, X=X, W=W, propensity=p) + point = est.const_marginal_effect(X) + + For a binary treatment, ``p`` can be a single vector with the probability of the non-control + treatment; with more treatment categories, pass one column per category (including control as the + first column), ordered to match the sorted unique treatment values (or the ``categories`` argument + if it was set). Note that since no propensity model is fitted, the ``propensity`` argument must + also be supplied when calling ``score`` on new data. The same argument is supported by the + :class:`.DRLearner` family of estimators. + + The supplied values must be the true, by-design probabilities of treatment conditional on + *everything* that influenced assignment — including design variables such as randomization + blocks or strata, even when those variables are not part of X or W. For example, with subjects + randomized within blocks, the correct value is each subject's within-block assignment + probability; the marginal treatment rate, or probabilities estimated from (X, W) alone, are not + generally sufficient when assignment depended on other outcome-relevant variables. When the + design variables are available, it is still worthwhile to also include them in W: this is not + needed for identification, but it reduces variance. + - **How can I assess the performance of the CATE model?** Each of the DML classes have an attribute `score_` after they are fitted. So one can access that diff --git a/doc/spec/estimation/dr.rst b/doc/spec/estimation/dr.rst index a99999a8c..e0495fe43 100644 --- a/doc/spec/estimation/dr.rst +++ b/doc/spec/estimation/dr.rst @@ -469,6 +469,41 @@ Usage FAQs The method allows for multiple discrete (categorical) treatments and will estimate a CATE model for each treatment. +- **What if I already know the treatment assignment probabilities (e.g. a randomized experiment)?** + + If the assignment probabilities are known by design (e.g. data from an A/B test or a stratified + experiment with per-unit assignment probabilities), you can pass them directly via the + ``propensity`` argument of ``fit`` instead of having the library estimate a propensity model; + the propensity model is then not fitted at all and the known values are used in the doubly + robust correction (they remain subject to ``min_propensity`` clipping and ``trimming_threshold`` + trimming): + + .. testcode:: + + from econml.dr import DRLearner + p = np.full(y.shape[0], 1 / 2) # known probability of treatment for each sample + est = DRLearner() + est.fit(y, T, X=X, W=W, propensity=p) + point = est.effect(X, T0=T0, T1=T1) + + For a binary treatment, ``p`` can be a single vector with the probability of the non-control + treatment; with more treatment categories, pass one column per category (including control as + the first column), ordered to match the sorted unique treatment values (or the ``categories`` + argument if it was set). Note that since no propensity model is fitted, the ``propensity`` + argument must also be supplied when calling ``score`` on new data. The same argument is + supported by the DML family of estimators when ``discrete_treatment=True``. + + The supplied values must be the true, by-design probabilities of treatment conditional on + *everything* that influenced assignment — including design variables such as randomization + blocks or strata, even when those variables are not part of X or W. For example, with subjects + randomized within blocks, the correct value is each subject's within-block assignment + probability; the marginal treatment rate, or probabilities estimated from (X, W) alone, are not + generally sufficient when assignment depended on other outcome-relevant variables. When the + design variables are available, it is still worthwhile to also include them in W: this is not + needed for identification, but it reduces variance. Also note that a non-default + ``min_propensity`` will bias the estimates if the true design probabilities fall below it, so + aggressive clipping should be avoided when the supplied probabilities are exact. + - **How can I assess the performance of the CATE model?** Each of the DRLearner classes have an attribute `score_` after they are fitted. So one can access that diff --git a/econml/_ortho_learner.py b/econml/_ortho_learner.py index cf966b631..ac92e9bd6 100644 --- a/econml/_ortho_learner.py +++ b/econml/_ortho_learner.py @@ -27,6 +27,7 @@ class in this module implements the general logic in a very versatile way from collections import namedtuple from abc import abstractmethod from typing import List, Union +from warnings import warn import numpy as np from sklearn.base import clone @@ -304,6 +305,8 @@ def predict(self, X, y, W=None): return accumulated_nuisances, model_list, fitted_inds, accumulated_scores +# NOTE: user-supplied propensities are deliberately not cached: refit_final only refits the +# final stage, which consumes the cached nuisances that were already computed from them CachedValues = namedtuple('CachedValues', ['nuisances', 'Y', 'T', 'X', 'W', 'Z', 'sample_weight', 'freq_weight', 'sample_var', 'groups']) @@ -663,9 +666,55 @@ def _check_fitted_dims_w_z(self, W, Z): def _subinds_check_none(self, var, inds): return var[inds] if var is not None else None + @property + def _supports_user_propensity(self): + """Whether this estimator supports user-supplied propensities at fit time.""" + return False + + def _expand_and_validate_propensity(self, propensity): + """ + Validate user-supplied propensities and expand them to one column per treatment category. + + A single column is allowed for binary treatments and is interpreted as the probability of the + non-control treatment. Returns an (n, n_categories) array whose columns are ordered to match + the fitted treatment categories. + """ + n_categories = len(self.transformer.categories_[0]) + propensity = np.asarray(propensity) + if propensity.ndim == 1 or (propensity.ndim == 2 and propensity.shape[1] == 1): + if n_categories != 2: + raise ValueError(f"A single propensity column was supplied, but there are {n_categories} " + "treatment categories; propensity must have one column per category " + f"(shape (n, {n_categories})), ordered to match the fitted categories " + f"{self.transformer.categories_[0]}.") + p = propensity.reshape(-1) + propensity = np.column_stack([1 - p, p]) + elif propensity.ndim == 2: + if propensity.shape[1] != n_categories: + raise ValueError(f"propensity has {propensity.shape[1]} columns, but there are {n_categories} " + "treatment categories; propensity must have one column per category " + "(including control as the first column), ordered to match the fitted " + f"categories {self.transformer.categories_[0]}.") + else: + raise ValueError("propensity must be a vector or a 2D array, " + f"but got an array with {propensity.ndim} dimensions.") + if np.any(propensity < 0) or np.any(propensity > 1): + raise ValueError("propensity values must lie in [0, 1].") + # a loose tolerance, to accommodate rounding in user-computed probabilities + # (e.g. lower-precision dtypes); actual mistakes like passing only some of the + # categories' probabilities will be off by far more than this + if not np.allclose(propensity.sum(axis=1), 1, atol=1e-3): + raise ValueError("propensity rows must sum to 1.") + if np.any((propensity == 0) | (propensity == 1)): + warn("Some user-supplied propensities are exactly 0 or 1, meaning that some treatment " + "values are impossible for some samples; effect estimates for such samples rely on " + "extrapolation. Verify that the supplied columns are ordered to match the treatment " + "categories.", UserWarning) + return propensity + def _strata(self, Y, T, X=None, W=None, Z=None, sample_weight=None, freq_weight=None, sample_var=None, groups=None, - cache_values=False, only_final=False, check_input=True): + propensity=None, cache_values=False, only_final=False, check_input=True): arrs = [] if self.discrete_outcome: arrs.append(Y) @@ -688,7 +737,7 @@ def _prefit(self, Y, T, *args, only_final=False, **kwargs): @BaseCateEstimator._wrap_fit def fit(self, Y, T, *, X=None, W=None, Z=None, sample_weight=None, freq_weight=None, sample_var=None, groups=None, - cache_values=False, inference=None, only_final=False, check_input=True): + propensity=None, cache_values=False, inference=None, only_final=False, check_input=True): """ Estimate the counterfactual model from data, i.e. estimates function :math:`\\theta(\\cdot)`. @@ -717,6 +766,20 @@ def fit(self, Y, T, *, X=None, W=None, Z=None, sample_weight=None, freq_weight=N All rows corresponding to the same group will be kept together during splitting. If groups is not None, the cv argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, the treatment/propensity + nuisance model is not fitted and these known values are used in its place. Only + supported by estimators whose first stage models treatment propensities and only + when the treatment is discrete. If a single column is passed, the treatment must be + binary and the values are interpreted as the probability of the non-control treatment; + otherwise there must be one column per treatment category (including control as the first + column), ordered to match the fitted categories (the sorted unique values of T, or the + ``categories`` initializer argument if it was set). cache_values: bool, default False Whether to cache the inputs and computed nuisances, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -740,16 +803,18 @@ def fit(self, Y, T, *, X=None, W=None, Z=None, sample_weight=None, freq_weight=N sample_var is None), "Sample variances and frequency weights must be provided together!" assert not (self.discrete_treatment and self.treatment_featurizer), "Treatment featurization " \ "is not supported when treatment is discrete" + if propensity is not None and not self._supports_user_propensity: + raise ValueError("This estimator does not support user-supplied propensities.") if check_input: - Y, T, Z, sample_weight, freq_weight, sample_var, groups = check_input_arrays( - Y, T, Z, sample_weight, freq_weight, sample_var, groups) + Y, T, Z, sample_weight, freq_weight, sample_var, groups, propensity = check_input_arrays( + Y, T, Z, sample_weight, freq_weight, sample_var, groups, propensity) X, = check_input_arrays( X, force_all_finite='allow-nan' if 'X' in self._gen_allowed_missing_vars() else True, ensure_2d=True) W, = check_input_arrays( W, force_all_finite='allow-nan' if 'W' in self._gen_allowed_missing_vars() else True, ensure_2d=True) - self._check_input_dims(Y, T, X, W, Z, sample_weight, freq_weight, sample_var, groups) + self._check_input_dims(Y, T, X, W, Z, sample_weight, freq_weight, sample_var, groups, propensity) if not only_final: @@ -788,6 +853,10 @@ def fit(self, Y, T, *, X=None, W=None, Z=None, sample_weight=None, freq_weight=N else: self.z_transformer = None + if propensity is not None: + propensity = self._expand_and_validate_propensity(propensity) + self._fit_with_user_propensity = propensity is not None + all_nuisances = [] fitted_inds = None if sample_weight is None: @@ -809,19 +878,20 @@ def fit(self, Y, T, *, X=None, W=None, Z=None, sample_weight=None, freq_weight=N self.ray_remote_func_options = self.ray_remote_func_options or {} # Define Ray remote function (Ray remote wrapper of the _fit_nuisances function) - def _fit_nuisances(Y, T, X, W, Z, sample_weight, groups): - return self._fit_nuisances(Y, T, X, W, Z, sample_weight=sample_weight, groups=groups) + def _fit_nuisances(Y, T, X, W, Z, sample_weight, groups, propensity): + return self._fit_nuisances(Y, T, X, W, Z, sample_weight=sample_weight, groups=groups, + propensity=propensity) # Create Ray remote jobs for parallel processing self.nuisances_ref = [ray.remote(_fit_nuisances).options(**self.ray_remote_func_options).remote( - Y, T, X, W, Z, sample_weight_nuisances, groups) for _ in range(self.mc_iters or 1)] + Y, T, X, W, Z, sample_weight_nuisances, groups, propensity) for _ in range(self.mc_iters or 1)] for idx in range(self.mc_iters or 1): if self.use_ray: nuisances, fitted_models, new_inds, scores = ray.get(self.nuisances_ref[idx]) else: nuisances, fitted_models, new_inds, scores = self._fit_nuisances( - Y, T, X, W, Z, sample_weight=sample_weight_nuisances, groups=groups) + Y, T, X, W, Z, sample_weight=sample_weight_nuisances, groups=groups, propensity=propensity) all_nuisances.append(nuisances) self._models_nuisance.append(fitted_models) if scores is None: @@ -923,7 +993,7 @@ def refit_final(self, inference=None): cache_values=True, inference=inference, only_final=True, check_input=False) return self - def _fit_nuisances(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None): + def _fit_nuisances(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, propensity=None): # use a binary array to get stratified split in case of discrete treatment stratify = self.discrete_treatment or self.discrete_instrument or self.discrete_outcome @@ -973,7 +1043,7 @@ def _fit_nuisances(self, Y, T, X=None, W=None, Z=None, sample_weight=None, group nuisances, fitted_models, fitted_inds, scores = _crossfit(self._ortho_learner_model_nuisance, folds, self.use_ray, self.ray_remote_func_options, Y, T, X=X, W=W, Z=Z, sample_weight=sample_weight, - groups=groups) + groups=groups, propensity=propensity) return nuisances, fitted_models, fitted_inds, scores def _fit_final(self, Y, T, X=None, W=None, Z=None, nuisances=None, sample_weight=None, @@ -1051,7 +1121,7 @@ def effect_inference(self, X=None, *, T0=0, T1=1): effect_inference.__doc__ = LinearCateEstimator.effect_inference.__doc__ - def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, scoring=None): + def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, scoring=None, propensity=None): """ Score the fitted CATE model on a new data set. @@ -1082,6 +1152,11 @@ def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, s scoring: name of an sklearn scoring function to use instead of the default, optional Supports f1_score, log_loss, mean_absolute_error, mean_squared_error, r2_score, and roc_auc_score. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment probabilities for each sample, in the same format as the + ``propensity`` argument to ``fit``. Required if the estimator was fit with user-supplied + propensities, since in that case no propensity model was fitted that could generate + them for the new data. Returns ------- @@ -1092,6 +1167,16 @@ def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, s if not hasattr(self._ortho_learner_model_final, 'score'): raise AttributeError("Final model does not have a score method!") Y, T, Z = check_input_arrays(Y, T, Z) + if propensity is not None: + if not self._supports_user_propensity: + raise ValueError("This estimator does not support user-supplied propensities.") + propensity, = check_input_arrays(propensity) + propensity = self._expand_and_validate_propensity(propensity) + if propensity.shape[0] != np.shape(Y)[0]: + raise ValueError("propensity must have the same number of rows as Y.") + elif getattr(self, '_fit_with_user_propensity', False): + raise ValueError("This estimator was fit with user-supplied propensities, so the `propensity` " + "argument must also be supplied when scoring on new data.") X, = check_input_arrays(X, force_all_finite='allow-nan' if 'X' in self._gen_allowed_missing_vars() else True) W, = check_input_arrays(W, force_all_finite='allow-nan' if 'W' in self._gen_allowed_missing_vars() else True) self._check_fitted_dims(X) @@ -1112,7 +1197,7 @@ def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, s n_selectors = len(self._models_nuisance[0]) n_splits = len(self._models_nuisance[0][0]) - kwargs = filter_none_kwargs(X=X, W=W, Z=Z, groups=groups) + kwargs = filter_none_kwargs(X=X, W=W, Z=Z, groups=groups, propensity=propensity) accumulated_nuisances = [] diff --git a/econml/dml/_rlearner.py b/econml/dml/_rlearner.py index f64be2344..21ced1f4d 100644 --- a/econml/dml/_rlearner.py +++ b/econml/dml/_rlearner.py @@ -50,29 +50,45 @@ def __init__(self, model_y: ModelSelector, model_t: ModelSelector): self._model_y = model_y self._model_t = model_t - def train(self, is_selecting, folds, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None): + def train(self, is_selecting, folds, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, + propensity=None): assert Z is None, "Cannot accept instrument!" - self._model_t.train(is_selecting, folds, X, W, T, ** - filter_none_kwargs(sample_weight=sample_weight, groups=groups)) + self._uses_user_propensity = propensity is not None + if propensity is None: + self._model_t.train(is_selecting, folds, X, W, T, ** + filter_none_kwargs(sample_weight=sample_weight, groups=groups)) self._model_y.train(is_selecting, folds, X, W, Y, ** filter_none_kwargs(sample_weight=sample_weight, groups=groups)) return self - def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, + def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, propensity=None, y_scoring=None, t_scoring=None, t_score_by_dim=False): # note that groups are not passed to score because they are only used for fitting - T_score = self._model_t.score(X, W, T, **filter_none_kwargs(sample_weight=sample_weight), - scoring=t_scoring, score_by_dim=t_score_by_dim) + if getattr(self, '_uses_user_propensity', False): + # no treatment model was fitted, so there is nothing to score + T_score = None + else: + T_score = self._model_t.score(X, W, T, **filter_none_kwargs(sample_weight=sample_weight), + scoring=t_scoring, score_by_dim=t_score_by_dim) Y_score = self._model_y.score(X, W, Y, **filter_none_kwargs(sample_weight=sample_weight), scoring=y_scoring) return Y_score, T_score - def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None): + def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, propensity=None): + if getattr(self, '_uses_user_propensity', False) and propensity is None: + raise ValueError("This model was trained with user-supplied propensities, " + "so propensities must also be supplied at predict time.") Y_pred = self._model_y.predict(X, W) - T_pred = self._model_t.predict(X, W) + if propensity is not None: + # the first stage treatment predictions are the probabilities of each non-control treatment, + # matching the one-hot encoding of T (which drops the control category) + T_pred = propensity[:, 1:] + else: + T_pred = self._model_t.predict(X, W) if (X is None) and (W is None): # In this case predict above returns a single row Y_pred = np.tile(Y_pred.reshape(1, -1), (Y.shape[0], 1)) - T_pred = np.tile(T_pred.reshape(1, -1), (T.shape[0], 1)) + if propensity is None: + T_pred = np.tile(T_pred.reshape(1, -1), (T.shape[0], 1)) Y_res = Y - Y_pred.reshape(Y.shape) T_res = T - T_pred.reshape(T.shape) return Y_res, T_res @@ -459,8 +475,14 @@ def _gen_ortho_learner_model_nuisance(self): def _gen_ortho_learner_model_final(self): return _ModelFinal(self._gen_rlearner_model_final()) + @property + def _supports_user_propensity(self): + # with a discrete treatment, the first stage treatment model predicts treatment probabilities, + # so known propensities can be used in its place + return self.discrete_treatment + def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sample_var=None, groups=None, - cache_values=False, inference=None): + propensity=None, cache_values=False, inference=None): """ Estimate the counterfactual model from data, i.e. estimates function :math:`\\theta(\\cdot)`. @@ -487,6 +509,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the treatment model is not fitted and these values are used in its place when + residualizing the treatment. Only supported when `discrete_treatment=True`. + If a single column is passed, the treatment must be binary and the values are + interpreted as the probability of the non-control treatment; otherwise there must + be one column per treatment category (including control as the first column), + ordered to match the fitted categories (the sorted unique values of T, or the + `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -497,13 +533,16 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam ------- self: _RLearner instance """ + if propensity is not None and not self.discrete_treatment: + raise ValueError("The propensity argument is only supported when discrete_treatment=True.") # Replacing fit from _OrthoLearner, to enforce Z=None and improve the docstring return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, freq_weight=freq_weight, sample_var=sample_var, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) - def score(self, Y, T, X=None, W=None, sample_weight=None, scoring=None): + def score(self, Y, T, X=None, W=None, sample_weight=None, scoring=None, propensity=None): """ Score the fitted CATE model on a new data set. @@ -526,6 +565,10 @@ def score(self, Y, T, X=None, W=None, sample_weight=None, scoring=None): Controls for each sample sample_weight:(n,) vector, optional Weights for each samples + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment probabilities for each sample, in the same format as the + `propensity` argument to `fit`. Required if the estimator was fit with user-supplied + propensities. Returns @@ -534,7 +577,7 @@ def score(self, Y, T, X=None, W=None, sample_weight=None, scoring=None): The MSE of the final CATE model on the new data. """ # Replacing score from _OrthoLearner, to enforce Z=None and improve the docstring - return super().score(Y, T, X=X, W=W, sample_weight=sample_weight, scoring=scoring) + return super().score(Y, T, X=X, W=W, sample_weight=sample_weight, propensity=propensity, scoring=scoring) @property def rlearner_model_final_(self): @@ -548,6 +591,9 @@ def models_y(self): @property def models_t(self): + if getattr(self, '_fit_with_user_propensity', False): + raise AttributeError("No treatment model was fitted because user-supplied propensities " + "were provided at fit time.") return [[mdl._model_t.best_model for mdl in mdls] for mdls in super().models_nuisance_] @property @@ -619,7 +665,8 @@ def score_nuisances(self, Y, T, X=None, W=None, Z=None, sample_weight=None, y_sc ------- score_dict : dict[str,list[float]] A dictionary where the keys indicate the Y and T scores used and the values are - lists of scores, one per CV fold model. + lists of scores, one per CV fold model. If the estimator was fit with user-supplied + propensities, no treatment model was fitted and the T scores will be None. """ Y_key = f'Y_{_RLearner.scoring_name(y_scoring)}' T_Key = f'T_{_RLearner.scoring_name(t_scoring)}' diff --git a/econml/dml/causal_forest.py b/econml/dml/causal_forest.py index e83ca4d4e..ef6f60833 100644 --- a/econml/dml/causal_forest.py +++ b/econml/dml/causal_forest.py @@ -929,7 +929,7 @@ def robustness_value(self, null_hypothesis=0, alpha=0.05, interval_type='ci'): # override only so that we can update the docstring to indicate support for `blb` def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates functions τ(·,·,·), ∂τ(·,·). @@ -949,6 +949,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the treatment model is not fitted and these values are used in its place when + residualizing the treatment. Only supported when `discrete_treatment=True`. + If a single column is passed, the treatment must be binary and the values are + interpreted as the probability of the non-control treatment; otherwise there must + be one column per treatment category (including control as the first column), + ordered to match the fitted categories (the sorted unique values of T, or the + `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -964,6 +978,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, raise ValueError("This estimator does not support X=None!") return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) diff --git a/econml/dml/dml.py b/econml/dml/dml.py index 23e47f8ec..000b1c96e 100644 --- a/econml/dml/dml.py +++ b/econml/dml/dml.py @@ -575,7 +575,7 @@ def _gen_rlearner_model_final(self): # override only so that we can update the docstring to indicate support for `LinearModelFinalInference` def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sample_var=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates functions τ(·,·,·), ∂τ(·,·). @@ -602,6 +602,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the treatment model is not fitted and these values are used in its place when + residualizing the treatment. Only supported when `discrete_treatment=True`. + If a single column is passed, the treatment must be binary and the values are + interpreted as the probability of the non-control treatment; otherwise there must + be one column per treatment category (including control as the first column), + ordered to match the fitted categories (the sorted unique values of T, or the + `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -615,6 +629,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam """ return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, freq_weight=freq_weight, sample_var=sample_var, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) @@ -915,7 +930,7 @@ def _gen_model_final(self): # override only so that we can update the docstring to indicate support for `StatsModelsInference` def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sample_var=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates functions τ(·,·,·), ∂τ(·,·). @@ -942,6 +957,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the treatment model is not fitted and these values are used in its place when + residualizing the treatment. Only supported when `discrete_treatment=True`. + If a single column is passed, the treatment must be binary and the values are + interpreted as the probability of the non-control treatment; otherwise there must + be one column per treatment category (including control as the first column), + ordered to match the fitted categories (the sorted unique values of T, or the + `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -955,6 +984,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam """ return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, freq_weight=freq_weight, sample_var=sample_var, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) @@ -1194,7 +1224,7 @@ def _gen_model_final(self): random_state=self.random_state) def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates functions τ(·,·,·), ∂τ(·,·). @@ -1214,6 +1244,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the treatment model is not fitted and these values are used in its place when + residualizing the treatment. Only supported when `discrete_treatment=True`. + If a single column is passed, the treatment must be binary and the values are + interpreted as the probability of the non-control treatment; otherwise there must + be one column per treatment category (including control as the first column), + ordered to match the fitted categories (the sorted unique values of T, or the + `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, `Inference` instance, or None @@ -1232,6 +1276,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, "We recommend using the LinearDML estimator for this low-dimensional setting.") return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) @property @@ -1420,7 +1465,7 @@ def _gen_featurizer(self): return _RandomFeatures(dim=self.dim, bw=self.bw, random_state=self.random_state) def fit(self, Y, T, X=None, W=None, *, sample_weight=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates functions τ(·,·,·), ∂τ(·,·). @@ -1440,6 +1485,20 @@ def fit(self, Y, T, X=None, W=None, *, sample_weight=None, groups=None, All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the treatment model is not fitted and these values are used in its place when + residualizing the treatment. Only supported when `discrete_treatment=True`. + If a single column is passed, the treatment must be binary and the values are + interpreted as the probability of the non-control treatment; otherwise there must + be one column per treatment category (including control as the first column), + ordered to match the fitted categories (the sorted unique values of T, or the + `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -1453,6 +1512,7 @@ def fit(self, Y, T, X=None, W=None, *, sample_weight=None, groups=None, """ return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) @property @@ -1662,7 +1722,7 @@ def _gen_rlearner_model_final(self): # override only so that we can update the docstring to indicate # support for `GenericSingleTreatmentModelFinalInference` def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sample_var=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates functions τ(·,·,·), ∂τ(·,·). @@ -1689,6 +1749,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the treatment model is not fitted and these values are used in its place when + residualizing the treatment. Only supported when `discrete_treatment=True`. + If a single column is passed, the treatment must be binary and the values are + interpreted as the probability of the non-control treatment; otherwise there must + be one column per treatment category (including control as the first column), + ordered to match the fitted categories (the sorted unique values of T, or the + `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -1702,6 +1776,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam """ return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, freq_weight=freq_weight, sample_var=sample_var, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) diff --git a/econml/dr/_drlearner.py b/econml/dr/_drlearner.py index 256680525..803279e1f 100644 --- a/econml/dr/_drlearner.py +++ b/econml/dr/_drlearner.py @@ -120,7 +120,7 @@ def __init__(self, def _combine(self, X, W): return np.hstack([arr for arr in [X, W] if arr is not None]) - def train(self, is_selecting, folds, Y, T, X=None, W=None, *, sample_weight=None, groups=None): + def train(self, is_selecting, folds, Y, T, X=None, W=None, *, sample_weight=None, groups=None, propensity=None): if Y.ndim != 1 and (Y.ndim != 2 or Y.shape[1] != 1): raise ValueError("The outcome matrix must be of shape ({0}, ) or ({0}, 1), " "instead got {1}.".format(len(X), Y.shape)) @@ -132,22 +132,33 @@ def train(self, is_selecting, folds, Y, T, X=None, W=None, *, sample_weight=None XW = self._combine(X, W) filtered_kwargs = filter_none_kwargs(sample_weight=sample_weight) - self._model_propensity.train(is_selecting, folds, XW, inverse_onehot(T), groups=groups, **filtered_kwargs) + self._uses_user_propensity = propensity is not None + if propensity is None: + self._model_propensity.train(is_selecting, folds, XW, inverse_onehot(T), groups=groups, **filtered_kwargs) self._model_regression.train(is_selecting, folds, np.hstack([XW, T]), Y, groups=groups, **filtered_kwargs) return self - def score(self, Y, T, X=None, W=None, *, sample_weight=None, groups=None): + def score(self, Y, T, X=None, W=None, *, sample_weight=None, groups=None, propensity=None): XW = self._combine(X, W) filtered_kwargs = filter_none_kwargs(sample_weight=sample_weight) - propensity_score = self._model_propensity.score(XW, inverse_onehot(T), **filtered_kwargs) + if getattr(self, '_uses_user_propensity', False): + # no propensity model was fitted, so there is nothing to score + propensity_score = None + else: + propensity_score = self._model_propensity.score(XW, inverse_onehot(T), **filtered_kwargs) regression_score = self._model_regression.score(np.hstack([XW, T]), Y, **filtered_kwargs) return propensity_score, regression_score - def predict(self, Y, T, X=None, W=None, *, sample_weight=None, groups=None): + def predict(self, Y, T, X=None, W=None, *, sample_weight=None, groups=None, propensity=None): + if getattr(self, '_uses_user_propensity', False) and propensity is None: + raise ValueError("This model was trained with user-supplied propensities, " + "so propensities must also be supplied at predict time.") XW = self._combine(X, W) - if hasattr(self._model_propensity, 'predict_proba'): + if propensity is not None: + raw_propensities = propensity + elif hasattr(self._model_propensity, 'predict_proba'): raw_propensities = self._model_propensity.predict_proba(XW) else: warn("A regressor was passed to model_propensity. " @@ -693,8 +704,14 @@ def _gen_ortho_learner_model_final(self): return _ModelFinal(self._gen_model_final(), self._gen_featurizer(), self.multitask_model_final, self.trimming_threshold) + @property + def _supports_user_propensity(self): + # the treatment is always discrete for DR learners, so known propensities can always be + # used in place of the propensity model + return True + def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sample_var=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates function :math:`\\theta(\\cdot)`. @@ -721,6 +738,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the propensity model is not fitted and these values are used in its place in the + doubly robust correction (they remain subject to `min_propensity` clipping and + `trimming_threshold` trimming). If a single column is passed, the treatment must be + binary and the values are interpreted as the probability of the non-control treatment; + otherwise there must be one column per treatment category (including control as the + first column), ordered to match the fitted categories (the sorted unique values of T, + or the `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -746,13 +777,14 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam # Replacing fit from _OrthoLearner, to enforce Z=None and improve the docstring return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, freq_weight=freq_weight, sample_var=sample_var, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) def refit_final(self, *, inference='auto'): return super().refit_final(inference=inference) refit_final.__doc__ = _OrthoLearner.refit_final.__doc__ - def score(self, Y, T, X=None, W=None, sample_weight=None): + def score(self, Y, T, X=None, W=None, sample_weight=None, propensity=None): """ Score the fitted CATE model on a new data set. @@ -775,6 +807,10 @@ def score(self, Y, T, X=None, W=None, sample_weight=None): Controls for each sample sample_weight:(n,) vector, optional Weights for each samples + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment probabilities for each sample, in the same format as the + `propensity` argument to `fit`. Required if the estimator was fit with user-supplied + propensities. Returns ------- @@ -782,7 +818,7 @@ def score(self, Y, T, X=None, W=None, sample_weight=None): The MSE of the final CATE model on the new data. """ # Replacing score from _OrthoLearner, to enforce Z=None and improve the docstring - return super().score(Y, T, X=X, W=W, sample_weight=sample_weight) + return super().score(Y, T, X=X, W=W, sample_weight=sample_weight, propensity=propensity) @property def multitask_model_cate(self): @@ -834,6 +870,9 @@ def models_propensity(self): monte carlo iterations, each element in the sublist corresponds to a crossfitting fold and is the model instance that was fitted for that training fold. """ + if getattr(self, '_fit_with_user_propensity', False): + raise AttributeError("No propensity model was fitted because user-supplied propensities " + "were provided at fit time.") return [[mdl._model_propensity.best_model for mdl in mdls] for mdls in super().models_nuisance_] @property @@ -1322,7 +1361,7 @@ def _gen_ortho_learner_model_final(self): return _ModelFinal(self._gen_model_final(), self._gen_featurizer(), False, self.trimming_threshold) def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sample_var=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates function :math:`\\theta(\\cdot)`. @@ -1349,6 +1388,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the propensity model is not fitted and these values are used in its place in the + doubly robust correction (they remain subject to `min_propensity` clipping and + `trimming_threshold` trimming). If a single column is passed, the treatment must be + binary and the values are interpreted as the probability of the non-control treatment; + otherwise there must be one column per treatment category (including control as the + first column), ordered to match the fitted categories (the sorted unique values of T, + or the `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -1363,6 +1416,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, freq_weight=None, sam # Replacing fit from DRLearner, to add statsmodels inference in docstring return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, freq_weight=freq_weight, sample_var=sample_var, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) @property @@ -1675,7 +1729,7 @@ def _gen_ortho_learner_model_final(self): return _ModelFinal(self._gen_model_final(), self._gen_featurizer(), False, self.trimming_threshold) def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates function :math:`\\theta(\\cdot)`. @@ -1695,6 +1749,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the propensity model is not fitted and these values are used in its place in the + doubly robust correction (they remain subject to `min_propensity` clipping and + `trimming_threshold` trimming). If a single column is passed, the treatment must be + binary and the values are interpreted as the probability of the non-control treatment; + otherwise there must be one column per treatment category (including control as the + first column), ordered to match the fitted categories (the sorted unique values of T, + or the `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, :class:`.Inference` instance, or None @@ -1714,6 +1782,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, "We recommend using the LinearDRLearner for this low-dimensional setting.") return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) @property @@ -2032,7 +2101,7 @@ def _gen_ortho_learner_model_final(self): return _ModelFinal(self._gen_model_final(), self._gen_featurizer(), False, self.trimming_threshold) def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, - cache_values=False, inference='auto'): + propensity=None, cache_values=False, inference='auto'): """ Estimate the counterfactual model from data, i.e. estimates functions τ(·,·,·), ∂τ(·,·). @@ -2052,6 +2121,20 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, All rows corresponding to the same group will be kept together during splitting. If groups is not None, the `cv` argument passed to this class's initializer must support a 'groups' argument to its split method. + propensity: {(n,), (n, n_categories)} array_like, optional + User-supplied treatment assignment probabilities for each sample, e.g. the known + assignment probabilities in a randomized experiment. These should be the true, + by-design probabilities of treatment conditional on everything that influenced + assignment, including any design variables (such as randomization blocks) even if + they are not part of X or W; probabilities estimated from, or marginalized over, + (X, W) alone are not generally sufficient. When provided, + the propensity model is not fitted and these values are used in its place in the + doubly robust correction (they remain subject to `min_propensity` clipping and + `trimming_threshold` trimming). If a single column is passed, the treatment must be + binary and the values are interpreted as the probability of the non-control treatment; + otherwise there must be one column per treatment category (including control as the + first column), ordered to match the fitted categories (the sorted unique values of T, + or the `categories` initializer argument if it was set). cache_values: bool, default False Whether to cache inputs and first stage results, which will allow refitting a different final model inference: str, `Inference` instance, or None @@ -2068,6 +2151,7 @@ def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None, return super().fit(Y, T, X=X, W=W, sample_weight=sample_weight, groups=groups, + propensity=propensity, cache_values=cache_values, inference=inference) def multitask_model_cate(self): diff --git a/econml/tests/test_user_propensity.py b/econml/tests/test_user_propensity.py new file mode 100644 index 000000000..37b8d07fb --- /dev/null +++ b/econml/tests/test_user_propensity.py @@ -0,0 +1,389 @@ +# Copyright (c) PyWhy contributors. All rights reserved. +# Licensed under the MIT License. + +"""Tests for user-supplied propensities in DML and DRLearner estimators.""" + +import unittest + +import numpy as np +import pytest +from sklearn.base import BaseEstimator, ClassifierMixin +from sklearn.ensemble import RandomForestRegressor +from sklearn.linear_model import LinearRegression, LogisticRegression + +from econml.dml import CausalForestDML, LinearDML, NonParamDML +from econml.dr import DRLearner, ForestDRLearner, LinearDRLearner +from econml.inference import BootstrapInference +from econml.iv.dml import OrthoIV + + +class _FailOnFitClassifier(BaseEstimator, ClassifierMixin): + """A classifier whose fit always raises, to prove that the model is bypassed.""" + + def fit(self, X, y, **kwargs): + raise AssertionError("The first stage treatment/propensity model should not be fitted " + "when user-supplied propensities are provided!") + + def predict_proba(self, X): + raise AssertionError("The first stage treatment/propensity model should not be used " + "when user-supplied propensities are provided!") + + +def _binary_dgp(n=1000, seed=123): + rng = np.random.default_rng(seed) + X = rng.normal(size=(n, 3)) + e = 0.2 + 0.6 / (1 + np.exp(-X[:, 0])) # known propensity, varies with X + T = rng.binomial(1, e) + tau = 1 + X[:, 0] + Y = tau * T + X[:, 0] + 0.5 * X[:, 1] + rng.normal(size=n) + return Y, T, X, e, tau + + +def _multivalued_dgp(n=2000, seed=456): + # per-unit-varying assignment probabilities, so that a row-misaligned or + # column-misaligned implementation could not pass the tests below + rng = np.random.default_rng(seed) + X = rng.normal(size=(n, 3)) + logits = np.column_stack([np.zeros(n), 0.5 * X[:, 0], -0.5 * X[:, 0]]) + propensity = np.exp(logits) / np.exp(logits).sum(axis=1, keepdims=True) + u = rng.random(n) + T = (u[:, None] > np.cumsum(propensity, axis=1)).sum(axis=1) + Y = (T == 1) * 1.0 + (T == 2) * 2.0 + X[:, 0] + rng.normal(size=n) + return Y, T, X, propensity + + +def _block_rct_dgp(n=8000, seed=789): + # block-randomized experiment where the block is NOT observable from X, + # so the assignment probabilities cannot be recovered by a propensity model + rng = np.random.default_rng(seed) + X = rng.normal(size=(n, 2)) + block = rng.choice(5, size=n) + e = np.array([0.1, 0.3, 0.5, 0.7, 0.9])[block] + T = rng.binomial(1, e) + tau = 1 + X[:, 0] + Y = tau * T + 4 * block + X[:, 0] + rng.normal(size=n) + return Y, T, X, e, tau + + +class TestUserPropensity(unittest.TestCase): + + def test_dml_binary_accuracy_and_residuals(self): + """With known propensities, DML should recover the ATE and use T - e as treatment residuals.""" + Y, T, X, e, tau = _binary_dgp() + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=1, random_state=0) + # cv=1 keeps samples in input order, so we can check residuals elementwise; + # _FailOnFitClassifier proves the treatment model is completely bypassed + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e, cache_values=True) + np.testing.assert_allclose(est.ate(X[:, :1]), np.mean(tau), atol=0.15) + _, T_res, _, _ = est.residuals_ + np.testing.assert_allclose(T_res.flatten(), T - e) + + def test_dml_binary_crossfit(self): + """Cross-fitting (cv>1) should also work, with residuals matching T - e up to reordering.""" + Y, T, X, e, tau = _binary_dgp() + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=3, random_state=0) + est.fit(Y, T, X=None, W=X, propensity=e, cache_values=True) + _, T_res, _, _ = est.residuals_ + np.testing.assert_allclose(np.sort(T_res.flatten()), np.sort(T - e)) + + def test_dml_single_and_two_column_propensity_equivalent(self): + """A (n,) propensity vector and the equivalent (n, 2) matrix should give identical results.""" + Y, T, X, e, _ = _binary_dgp() + results = [] + for prop in [e, np.column_stack([1 - e, e])]: + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=prop) + results.append(est.ate(X[:, :1])) + np.testing.assert_allclose(results[0], results[1]) + + def test_dr_binary_accuracy_and_bypass(self): + """With known propensities, DRLearner should recover the ATE without fitting model_propensity.""" + Y, T, X, e, tau = _binary_dgp() + est = LinearDRLearner(model_regression=RandomForestRegressor(n_estimators=50, random_state=0), + model_propensity=_FailOnFitClassifier(), cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + np.testing.assert_allclose(est.ate(X[:, :1]), np.mean(tau), atol=0.2) + # no propensity model was fitted, so its nuisance scores should be None + for scores in est.nuisance_scores_propensity: + for score in scores: + self.assertIsNone(score) + + def test_dr_multivalued_treatment(self): + """User-supplied propensities should work with more than two treatment categories.""" + Y, T, X, propensity = _multivalued_dgp() + est = LinearDRLearner(model_regression=LinearRegression(), + model_propensity=_FailOnFitClassifier(), cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=propensity) + np.testing.assert_allclose(est.ate(X[:, :1], T0=0, T1=1), 1.0, atol=0.2) + np.testing.assert_allclose(est.ate(X[:, :1], T0=0, T1=2), 2.0, atol=0.2) + + def test_dml_multivalued_treatment_residuals(self): + """Multi-valued DML treatment residuals should equal the one-hot encoding minus the propensities.""" + Y, T, X, propensity = _multivalued_dgp() + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=1, random_state=0) + est.fit(Y, T, X=None, W=X, propensity=propensity, cache_values=True) + _, T_res, _, _ = est.residuals_ + T_onehot = np.column_stack([(T == 1), (T == 2)]).astype(float) + np.testing.assert_allclose(T_res, T_onehot - propensity[:, 1:]) + + def test_hidden_block_rct(self): + """Known propensities recover the ATE in a block-randomized RCT where the blocks are not observed. + + This is the motivating use case for the feature: the assignment probabilities are known by + design but cannot be estimated from the observed covariates, so a fitted propensity model + produces badly confounded estimates while the supplied probabilities give consistent ones + with valid confidence intervals. + """ + Y, T, X, e, tau = _block_rct_dgp() + true_ate = np.mean(tau) + for cls, kwargs in [(LinearDML, {'model_y': LinearRegression(), 'model_t': LogisticRegression(), + 'discrete_treatment': True}), + (LinearDRLearner, {'model_regression': LinearRegression(), + 'model_propensity': LogisticRegression()})]: + with self.subTest(estimator=cls.__name__): + known = cls(cv=2, random_state=0, **kwargs) + known.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + lb, ub = known.ate_interval(X[:, :1], alpha=0.05) + self.assertTrue(lb <= true_ate <= ub, + f"CI ({lb}, {ub}) does not cover the true ATE {true_ate}") + naive = cls(cv=2, random_state=0, **kwargs) + naive.fit(Y, T, X=X[:, :1], W=X[:, 1:]) + known_bias = abs(known.ate(X[:, :1]) - true_ate) + naive_bias = abs(naive.ate(X[:, :1]) - true_ate) + self.assertLess(known_bias, 0.3) + self.assertGreater(naive_bias, 5 * known_bias) + + def test_categories_ordering(self): + """Propensity columns must follow the `categories` initializer argument when it is set.""" + rng = np.random.default_rng(0) + n = 4000 + X = rng.normal(size=(n, 2)) + e_a = 0.1 + 0.8 / (1 + np.exp(-2 * X[:, 0])) # probability of treatment 'a' + T = np.where(rng.random(n) < e_a, 'a', 'b') + Y = 2.0 * (T == 'a') + X[:, 0] + rng.normal(size=n) + # with categories=['b', 'a'], 'b' is the control, so the first propensity + # column is P(T='b') and the effect of 'a' relative to 'b' is 2 + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, categories=['b', 'a'], cv=2, random_state=0) + est.fit(Y, T, X=None, W=X, propensity=np.column_stack([1 - e_a, e_a])) + np.testing.assert_allclose(est.ate(T0='b', T1='a'), 2.0, atol=0.2) + # negative control: swapping the columns must give a substantially wrong answer, + # proving the column ordering is actually consumed + swapped = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, categories=['b', 'a'], cv=2, random_state=0) + swapped.fit(Y, T, X=None, W=X, propensity=np.column_stack([e_a, 1 - e_a])) + self.assertGreater(abs(swapped.ate(T0='b', T1='a') - 2.0), 0.5) + + def test_no_X_no_W(self): + """DML with X=None and W=None (a pure experiment) should accept known propensities.""" + rng = np.random.default_rng(1) + n = 2000 + block = rng.choice(2, size=n) + e = np.array([0.2, 0.8])[block] + T = rng.binomial(1, e) + Y = 1.0 * T + 2 * block + rng.normal(size=n) + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=None, W=None, propensity=e) + np.testing.assert_allclose(est.ate(), 1.0, atol=0.25) + + def test_refit_resets_state(self): + """Re-fitting the same estimator instance with/without propensities should fully reset state.""" + Y, T, X, e, _ = _binary_dgp(n=500) + est = LinearDML(model_y=LinearRegression(), model_t=LogisticRegression(), + discrete_treatment=True, cv=2, random_state=0) + # fit with user propensities, then re-fit without: score must work without them again + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:]) + est.score(Y, T, X=X[:, :1], W=X[:, 1:]) + # and the reverse: re-fitting with propensities must re-impose the score requirement + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + with pytest.raises(ValueError, match="propensity"): + est.score(Y, T, X=X[:, :1], W=X[:, 1:]) + + def test_other_estimators_smoke(self): + """CausalForestDML, NonParamDML, DRLearner and ForestDRLearner should all accept propensities.""" + Y, T, X, e, _ = _binary_dgp(n=500) + for est in [CausalForestDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, n_estimators=100, cv=2, random_state=0), + NonParamDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + model_final=RandomForestRegressor(n_estimators=20, random_state=0), + discrete_treatment=True, cv=2, random_state=0), + DRLearner(model_regression=LinearRegression(), model_propensity=_FailOnFitClassifier(), + model_final=LinearRegression(), cv=2, random_state=0), + ForestDRLearner(model_regression=LinearRegression(), model_propensity=_FailOnFitClassifier(), + n_estimators=100, cv=2, random_state=0)]: + with self.subTest(estimator=type(est).__name__): + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + est.effect(X[:5, :1]) + + def test_score_requires_propensity_when_fit_with_it(self): + """If fit used user-supplied propensities, score must also receive them.""" + Y, T, X, e, _ = _binary_dgp(n=500) + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + with pytest.raises(ValueError, match="propensity"): + est.score(Y, T, X=X[:, :1], W=X[:, 1:]) + self.assertIsInstance(est.score(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e), float) + + dr = LinearDRLearner(model_regression=LinearRegression(), + model_propensity=_FailOnFitClassifier(), cv=2, random_state=0) + dr.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + with pytest.raises(ValueError, match="propensity"): + dr.score(Y, T, X=X[:, :1], W=X[:, 1:]) + dr.score(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + + def test_score_optional_when_fit_without_it(self): + """An estimator fit normally can score with or without user-supplied propensities.""" + Y, T, X, e, _ = _binary_dgp(n=500) + est = LinearDML(model_y=LinearRegression(), model_t=LogisticRegression(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:]) + est.score(Y, T, X=X[:, :1], W=X[:, 1:]) + est.score(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + + def test_validation_errors(self): + Y, T, X, e, _ = _binary_dgp(n=500) + Y3, T3, X3, propensity3 = _multivalued_dgp(n=500) + + # continuous treatment is not supported + with pytest.raises(ValueError, match="discrete"): + LinearDML(random_state=0).fit(Y, X[:, 0], X=X[:, :1], W=X[:, 1:], propensity=e) + # single column propensity requires binary treatment + with pytest.raises(ValueError, match="column"): + LinearDRLearner(cv=2, random_state=0).fit(Y3, T3, X=X3[:, :1], W=X3[:, 1:], propensity=e) + # column count must match the number of categories + with pytest.raises(ValueError, match="column"): + LinearDRLearner(cv=2, random_state=0).fit( + Y3, T3, X=X3[:, :1], W=X3[:, 1:], propensity=propensity3[:, :2]) + # values must be valid probabilities + with pytest.raises(ValueError, match=r"\[0, 1\]"): + LinearDML(discrete_treatment=True, random_state=0).fit( + Y, T, X=X[:, :1], W=X[:, 1:], propensity=2 * e) + # rows must sum to 1 + with pytest.raises(ValueError, match="sum"): + LinearDML(discrete_treatment=True, random_state=0).fit( + Y, T, X=X[:, :1], W=X[:, 1:], propensity=np.column_stack([e, e])) + # row count must match the data + with pytest.raises(AssertionError): + LinearDML(discrete_treatment=True, random_state=0).fit( + Y, T, X=X[:, :1], W=X[:, 1:], propensity=e[:-1]) + # estimators that don't model propensities don't accept the argument + Z = np.random.default_rng(0).binomial(1, 0.5, size=Y.shape[0]) + with pytest.raises(TypeError): + OrthoIV(discrete_treatment=True, discrete_instrument=True).fit( + Y, T, Z=Z, X=X[:, :1], W=X[:, 1:], propensity=e) + + def test_mc_iters_and_refit_final(self): + """User-supplied propensities should compose with monte carlo iterations and refit_final.""" + Y, T, X, e, tau = _binary_dgp() + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, mc_iters=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e, cache_values=True) + ate_before = est.ate(X[:, :1]) + est.refit_final() + np.testing.assert_allclose(est.ate(X[:, :1]), ate_before) + np.testing.assert_allclose(ate_before, np.mean(tau), atol=0.15) + + def test_bootstrap_inference(self): + """Bootstrap inference resamples the supplied propensities along with the rest of the data.""" + Y, T, X, e, tau = _binary_dgp(n=500) + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e, + inference=BootstrapInference(n_bootstrap_samples=5, n_jobs=1)) + lb, ub = est.ate_interval(X[:, :1]) + self.assertLess(lb, ub) + + def test_clipping_and_trimming(self): + """Supplied propensities remain subject to min_propensity clipping and trimming_threshold trimming.""" + Y, T, X, e, _ = _binary_dgp(n=1000) + # push some propensities into the trimmable/clippable region + e_extreme = np.clip(e, 0.02, 0.98) + e_extreme[:100] = 0.02 + est = DRLearner(model_regression=LinearRegression(), model_propensity=_FailOnFitClassifier(), + model_final=LinearRegression(), trimming_threshold=0.1, min_propensity=0.05, + cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e_extreme) + self.assertGreater(est.n_samples_trimmed_, 0) + # without trimming, the same input fits on all samples + est2 = DRLearner(model_regression=LinearRegression(), model_propensity=_FailOnFitClassifier(), + model_final=LinearRegression(), cv=2, random_state=0) + est2.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e_extreme) + self.assertEqual(est2.n_samples_trimmed_, 0) + + def test_composition_with_groups_and_weights(self): + """The propensity array must survive fold slicing alongside groups and frequency weights.""" + Y, T, X, e, _ = _binary_dgp(n=600) + groups = np.repeat(np.arange(300), 2) + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e, groups=groups) + est.effect(X[:5, :1]) + est2 = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, random_state=0) + est2.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e, sample_weight=np.ones(600), + freq_weight=np.ones(600, dtype=int), sample_var=np.ones(600)) + est2.effect(X[:5, :1]) + + def test_accessors_after_bypassed_fit(self): + """Accessors for the bypassed model should raise an informative error, and score_nuisances works.""" + Y, T, X, e, _ = _binary_dgp(n=500) + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + with pytest.raises(AttributeError, match="user-supplied propensities"): + est.models_t + scores = est.score_nuisances(Y, T, X=X[:, :1], W=X[:, 1:]) + self.assertTrue(all(s is None for s in scores['T_default_score'])) + self.assertTrue(all(s is not None for s in scores['Y_default_score'])) + + dr = LinearDRLearner(model_regression=LinearRegression(), + model_propensity=_FailOnFitClassifier(), cv=2, random_state=0) + dr.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + with pytest.raises(AttributeError, match="user-supplied propensities"): + dr.models_propensity + + def test_degenerate_propensity_warns(self): + """Propensities of exactly 0 or 1 should produce a warning since they imply no overlap.""" + Y, T, X, e, _ = _binary_dgp(n=500) + e_degenerate = e.copy() + e_degenerate[T == 1] = np.maximum(e_degenerate[T == 1], 0.5) + e_degenerate[:5] = np.where(T[:5] == 1, 1.0, 0.0) + est = LinearDRLearner(model_regression=LinearRegression(), + model_propensity=_FailOnFitClassifier(), cv=2, random_state=0) + with pytest.warns(UserWarning, match="exactly 0 or 1"): + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e_degenerate) + + def test_score_positional_compatibility(self): + """Existing positional score() callers must be unaffected (propensity is the last parameter).""" + Y, T, X, e, _ = _binary_dgp(n=500) + est = LinearDML(model_y=LinearRegression(), model_t=LogisticRegression(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:]) + # the historical positional order is (Y, T, X, W, sample_weight, scoring) + s = est.score(Y, T, X[:, :1], X[:, 1:], None, 'mean_squared_error') + self.assertIsInstance(s, float) + + def test_nuisance_scores_t_none(self): + """When the treatment model is bypassed, its nuisance scores should be None.""" + Y, T, X, e, _ = _binary_dgp(n=500) + est = LinearDML(model_y=LinearRegression(), model_t=_FailOnFitClassifier(), + discrete_treatment=True, cv=2, random_state=0) + est.fit(Y, T, X=X[:, :1], W=X[:, 1:], propensity=e) + for scores in est.nuisance_scores_t: + for score in scores: + self.assertIsNone(score) + # the outcome model is still fitted and scored normally + for scores in est.nuisance_scores_y: + for score in scores: + self.assertIsNotNone(score) + + +if __name__ == '__main__': + unittest.main()