Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add propensity_learner to R-learner. #297

Merged
merged 1 commit into from Feb 3, 2021
Merged

Add propensity_learner to R-learner. #297

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
58 changes: 58 additions & 0 deletions causalml/inference/meta/base.py
View file
Open in desktop
@@ -1,7 +1,14 @@
from abc import ABCMeta, abstractclassmethod
import logging
import numpy as np
import pandas as pd

from causalml.inference.meta.explainer import Explainer
from causalml.inference.meta.utils import check_p_conditions, convert_pd_to_np
from causalml.propensity import compute_propensity_score


logger = logging.getLogger('causalml')


class BaseLearner(metaclass=ABCMeta):
Expand Down Expand Up @@ -38,6 +45,57 @@ def bootstrap(self, X, treatment, y, p=None, size=10000):
self.fit(X=X_b, treatment=treatment_b, y=y_b, p=p_b)
return self.predict(X=X, p=p)

@staticmethod
def _format_p(p, t_groups):
"""Format propensity scores into a dictionary of {treatment group: propensity scores}.

Args:
p (np.ndarray, pd.Series, or dict): propensity scores
t_groups (list): treatment group names.

Returns:
dict of {treatment group: propensity scores}
"""
check_p_conditions(p, t_groups)

if isinstance(p, (np.ndarray, pd.Series)):
treatment_name = t_groups[0]
p = {treatment_name: convert_pd_to_np(p)}
elif isinstance(p, dict):
p = {treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()}

return p

def _set_propensity_models(self, X, treatment, y):
"""Set self.propensity and self.propensity_models.

It trains propensity models for all treatment groups, save them in self.propensity_models, and
save propensity scores in self.propensity in dictionaries with treatment groups as keys.

It will use self.model_p if available to train propensity models. Otherwise, it will use a default
PropensityModel (i.e. ElasticNetPropensityModel).

Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
treatment (np.array or pd.Series): a treatment vector
y (np.array or pd.Series): an outcome vector
"""
logger.info('Generating propensity score')
p = dict()
p_model = dict()
for group in self.t_groups:
mask = (treatment == group) | (treatment == self.control_name)
treatment_filt = treatment[mask]
X_filt = X[mask]
w_filt = (treatment_filt == group).astype(int)
w = (treatment == group).astype(int)
propensity_model = self.model_p if hasattr(self, 'model_p') else None
p[group], p_model[group] = compute_propensity_score(X=X_filt, treatment=w_filt,
Copy link
Collaborator

@ppstacy ppstacy Feb 3, 2021 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @jeongyoonlee a quick question that might not be directly related to this PR. Why do we need two treatment and treatment_pred here? Aren't those the same thing based on here: https://github.com/uber/causalml/blob/master/causalml/propensity.py#L189?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this call to compute_propensity_score(), treatment is only for the filtered samples with treatment belonging to either the control group or treatment group of interest while treatment_pred is for the entire samples. These two will be the same for a single treatment, but different for the multiple treatments.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gotcha I see. Sorry a follow-up question here, why do we need to calibrate it with the entire dataset? I am asking because we are only using treatment group X + control when running the model in https://github.com/uber/causalml/blob/master/causalml/inference/meta/rlearner.py#L111-L118.

p_model=propensity_model,
X_pred=X, treatment_pred=w)
self.propensity_model = p_model
self.propensity = p

def get_importance(self, X=None, tau=None, model_tau_feature=None, features=None, method='auto', normalize=True,
test_size=0.3, random_state=None):
"""
Expand Down
138 changes: 40 additions & 98 deletions causalml/inference/meta/rlearner.py
View file
Open in desktop
Expand Up @@ -8,10 +8,11 @@
from xgboost import XGBRegressor

from causalml.inference.meta.base import BaseLearner
from causalml.inference.meta.utils import (check_treatment_vector, check_p_conditions,
from causalml.inference.meta.utils import (check_treatment_vector,
get_xgboost_objective_metric, convert_pd_to_np)
from causalml.inference.meta.explainer import Explainer
from causalml.propensity import compute_propensity_score
from causalml.propensity import compute_propensity_score, ElasticNetPropensityModel


logger = logging.getLogger('causalml')

Expand All @@ -28,6 +29,7 @@ def __init__(self,
learner=None,
outcome_learner=None,
effect_learner=None,
propensity_learner=ElasticNetPropensityModel(),
ate_alpha=.05,
control_name=0,
n_fold=5,
Expand All @@ -39,22 +41,19 @@ def __init__(self,
outcome_learner (optional): a model to estimate outcomes
effect_learner (optional): a model to estimate treatment effects. It needs to take `sample_weight` as an
input argument for `fit()`
propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
be used by default.
ate_alpha (float, optional): the confidence level alpha of the ATE estimate
control_name (str or int, optional): name of control group
n_fold (int, optional): the number of cross validation folds for outcome_learner
random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
"""
assert (learner is not None) or ((outcome_learner is not None) and (effect_learner is not None))
assert propensity_learner is not None

if outcome_learner is None:
self.model_mu = deepcopy(learner)
else:
self.model_mu = outcome_learner

if effect_learner is None:
self.model_tau = deepcopy(learner)
else:
self.model_tau = effect_learner
self.model_mu = outcome_learner if outcome_learner else deepcopy(learner)
self.model_tau = effect_learner if outcome_learner else deepcopy(learner)
self.model_p = propensity_learner

self.ate_alpha = ate_alpha
self.control_name = control_name
Expand All @@ -66,10 +65,10 @@ def __init__(self,
self.propensity_model = None

def __repr__(self):
return ('{}(model_mu={},\n'
'\tmodel_tau={})'.format(self.__class__.__name__,
self.model_mu.__repr__(),
self.model_tau.__repr__()))
return (f'{self.__class__.__name__}\n'
f'\toutcome_learner={self.model_mu.__repr__()}\n'
f'\teffect_learner={self.model_tau.__repr__()}\n'
f'\tpropensity_learner={self.model_p.__repr__()}')

def fit(self, X, treatment, y, p=None, verbose=True):
"""Fit the treatment effect and outcome models of the R learner.
Expand All @@ -89,27 +88,10 @@ def fit(self, X, treatment, y, p=None, verbose=True):
self.t_groups.sort()

if p is None:
logger.info('Generating propensity score')
p = dict()
p_model = dict()
for group in self.t_groups:
mask = (treatment == group) | (treatment == self.control_name)
treatment_filt = treatment[mask]
X_filt = X[mask]
w_filt = (treatment_filt == group).astype(int)
w = (treatment == group).astype(int)
p[group], p_model[group] = compute_propensity_score(X=X_filt, treatment=w_filt,
X_pred=X, treatment_pred=w)
self.propensity_model = p_model
self.propensity = p
self._set_propensity_models(X=X, treatment=treatment, y=y)
p = self.propensity
else:
check_p_conditions(p, self.t_groups)

if isinstance(p, (np.ndarray, pd.Series)):
treatment_name = self.t_groups[0]
p = {treatment_name: convert_pd_to_np(p)}
elif isinstance(p, dict):
p = {treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()}
p = self._format_p(p, self.t_groups)

self._classes = {group: i for i, group in enumerate(self.t_groups)}
self.models_tau = {group: deepcopy(self.model_tau) for group in self.t_groups}
Expand Down Expand Up @@ -178,17 +160,6 @@ def fit_predict(self, X, treatment, y, p=None, return_ci=False,
self.fit(X, treatment, y, p, verbose=verbose)
te = self.predict(X)

if p is None:
p = self.propensity
else:
check_p_conditions(p, self.t_groups)

if isinstance(p, (np.ndarray, pd.Series)):
treatment_name = self.t_groups[0]
p = {treatment_name: convert_pd_to_np(p)}
elif isinstance(p, dict):
p = {treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()}

if not return_ci:
return te
else:
Expand All @@ -200,6 +171,10 @@ def fit_predict(self, X, treatment, y, p=None, return_ci=False,

logger.info('Bootstrap Confidence Intervals')
for i in tqdm(range(n_bootstraps)):
if p is None:
p = self.propensity
else:
p = self._format_p(p, self.t_groups)
te_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
te_bootstraps[:, :, i] = te_b

Expand Down Expand Up @@ -231,18 +206,7 @@ def estimate_ate(self, X, treatment, y, p=None, bootstrap_ci=False, n_bootstraps
The mean and confidence interval (LB, UB) of the ATE estimate.
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
te = self.fit_predict(X, treatment, y, p)

if p is None:
p = self.propensity
else:
check_p_conditions(p, self.t_groups)

if isinstance(p, (np.ndarray, pd.Series)):
treatment_name = self.t_groups[0]
p = {treatment_name: convert_pd_to_np(p)}
elif isinstance(p, dict):
p = {treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()}
te = self.fit_predict(X, treatment, y, p, return_ci=False)

ate = np.zeros(self.t_groups.shape[0])
ate_lb = np.zeros(self.t_groups.shape[0])
Expand Down Expand Up @@ -277,6 +241,10 @@ def estimate_ate(self, X, treatment, y, p=None, bootstrap_ci=False, n_bootstraps
ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))

for n in tqdm(range(n_bootstraps)):
if p is None:
p = self.propensity
else:
p = self._format_p(p, self.t_groups)
cate_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
ate_bootstraps[:, n] = cate_b.mean()

Expand All @@ -300,6 +268,7 @@ def __init__(self,
learner=None,
outcome_learner=None,
effect_learner=None,
propensity_learner=ElasticNetPropensityModel(),
ate_alpha=.05,
control_name=0,
n_fold=5,
Expand All @@ -311,6 +280,8 @@ def __init__(self,
outcome_learner (optional): a model to estimate outcomes
effect_learner (optional): a model to estimate treatment effects. It needs to take `sample_weight` as an
input argument for `fit()`
propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
be used by default.
ate_alpha (float, optional): the confidence level alpha of the ATE estimate
control_name (str or int, optional): name of control group
n_fold (int, optional): the number of cross validation folds for outcome_learner
Expand All @@ -320,6 +291,7 @@ def __init__(self,
learner=learner,
outcome_learner=outcome_learner,
effect_learner=effect_learner,
propensity_learner=propensity_learner,
ate_alpha=ate_alpha,
control_name=control_name,
n_fold=n_fold,
Expand All @@ -334,6 +306,7 @@ class BaseRClassifier(BaseRLearner):
def __init__(self,
outcome_learner=None,
effect_learner=None,
propensity_learner=ElasticNetPropensityModel(),
ate_alpha=.05,
control_name=0,
n_fold=5,
Expand All @@ -344,6 +317,8 @@ def __init__(self,
outcome_learner: a model to estimate outcomes. Should be a classifier.
effect_learner: a model to estimate treatment effects. It needs to take `sample_weight` as an
input argument for `fit()`. Should be a regressor.
propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
be used by default.
ate_alpha (float, optional): the confidence level alpha of the ATE estimate
control_name (str or int, optional): name of control group
n_fold (int, optional): the number of cross validation folds for outcome_learner
Expand All @@ -353,6 +328,7 @@ def __init__(self,
learner=None,
outcome_learner=outcome_learner,
effect_learner=effect_learner,
propensity_learner=propensity_learner,
ate_alpha=ate_alpha,
control_name=control_name,
n_fold=n_fold,
Expand All @@ -379,27 +355,10 @@ def fit(self, X, treatment, y, p=None, verbose=True):
self.t_groups.sort()

if p is None:
logger.info('Generating propensity score')
p = dict()
p_model = dict()
for group in self.t_groups:
mask = (treatment == group) | (treatment == self.control_name)
treatment_filt = treatment[mask]
X_filt = X[mask]
w_filt = (treatment_filt == group).astype(int)
w = (treatment == group).astype(int)
p[group], p_model[group] = compute_propensity_score(X=X_filt, treatment=w_filt,
X_pred=X, treatment_pred=w)
self.propensity_model = p_model
self.propensity = p
self._set_propensity_models(X=X, treatment=treatment, y=y)
p = self.propensity
else:
check_p_conditions(p, self.t_groups)

if isinstance(p, (np.ndarray, pd.Series)):
treatment_name = self.t_groups[0]
p = {treatment_name: convert_pd_to_np(p)}
elif isinstance(p, dict):
p = {treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()}
p = self._format_p(p, self.t_groups)

self._classes = {group: i for i, group in enumerate(self.t_groups)}
self.models_tau = {group: deepcopy(self.model_tau) for group in self.t_groups}
Expand Down Expand Up @@ -504,27 +463,10 @@ def fit(self, X, treatment, y, p=None, verbose=True):
self.t_groups.sort()

if p is None:
logger.info('Generating propensity score')
p = dict()
p_model = dict()
for group in self.t_groups:
mask = (treatment == group) | (treatment == self.control_name)
treatment_filt = treatment[mask]
X_filt = X[mask]
w_filt = (treatment_filt == group).astype(int)
w = (treatment == group).astype(int)
p[group], p_model[group] = compute_propensity_score(X=X_filt, treatment=w_filt,
X_pred=X, treatment_pred=w)
self.propensity_model = p_model
self.propensity = p
self._set_propensity_models(X=X, treatment=treatment, y=y)
p = self.propensity
else:
check_p_conditions(p, self.t_groups)

if isinstance(p, (np.ndarray, pd.Series)):
treatment_name = self.t_groups[0]
p = {treatment_name: convert_pd_to_np(p)}
elif isinstance(p, dict):
p = {treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()}
p = self._format_p(p, self.t_groups)

self._classes = {group: i for i, group in enumerate(self.t_groups)}
self.models_tau = {group: deepcopy(self.model_tau) for group in self.t_groups}
Expand Down