def test_BaseXRegressor_without_p(generate_regression_data):
y, X, treatment, tau, b, e = generate_regression_data()
learner = BaseXRegressor(learner=XGBRegressor())
# check the accuracy of the ATE estimation
ate_p, lb, ub = learner.estimate_ate(X=X, treatment=treatment, y=y)
assert (ate_p >= lb) and (ate_p <= ub)
assert ape(tau.mean(), ate_p) < ERROR_THRESHOLD
# check the accuracy of the CATE estimation with the bootstrap CI
cate_p, _, _ = learner.fit_predict(X=X,
treatment=treatment,
y=y,
return_ci=True,
n_bootstraps=10)
auuc_metrics = pd.DataFrame({
'cate_p': cate_p.flatten(),
'W': treatment,
'y': y,
'treatment_effect_col': tau
})
cumgain = get_cumgain(auuc_metrics,
outcome_col='y',
treatment_col='W',
treatment_effect_col='tau')
# Check if the cumulative gain when using the model's prediction is
# higher than it would be under random targeting
assert cumgain['cate_p'].sum() > cumgain['Random'].sum()
def test_BaseXRegressor_without_p(generate_regression_data):
y, X, treatment, tau, b, e = generate_regression_data()
learner = BaseXRegressor(learner=XGBRegressor())
# check the accuracy of the ATE estimation
ate_p, lb, ub = learner.estimate_ate(X=X, treatment=treatment, y=y)
assert (ate_p >= lb) and (ate_p <= ub)
assert ape(tau.mean(), ate_p) < ERROR_THRESHOLD
# check the accuracy of the CATE estimation with the bootstrap CI
cate_p, _, _ = learner.fit_predict(X=X, treatment=treatment, y=y, return_ci=True, n_bootstraps=10)
assert gini(tau, cate_p.flatten()) > .5
def __init__(
self,
learner=None,
outcome_learner_c=None,
outcome_learner_t=None,
effect_learner_c=None,
effect_learner_t=None,
):
"""Setup a XLearner
All learners must have ``fit(x, y)`` and ``predict(x)`` methods.
Args:
learner: default learner for all roles
outcome_learner_c: specific learner for control outcome function
outcome_learner_t: specific learner for treated outcome function
effect_learner_c: specific learner for treated effect
effect_learner_t: specific learner for control effect
"""
from causalml.inference.meta import BaseXRegressor
if (learner is not None) or ((outcome_learner_c is not None) and
(outcome_learner_t is not None) and
(effect_learner_c is not None) and
(effect_learner_t is not None)):
self.model = BaseXRegressor(
learner,
outcome_learner_c,
outcome_learner_t,
effect_learner_c,
effect_learner_t,
)
else:
# Assign default learner
learner = LassoLars()
self.model = BaseXRegressor(
learner,
outcome_learner_c,
outcome_learner_t,
effect_learner_c,
effect_learner_t,
)
def test_BaseXRegressor(generate_regression_data):
y, X, treatment, tau, b, e = generate_regression_data()
learner = BaseXRegressor(learner=XGBRegressor())
# check the accuracy of the ATE estimation
ate_p, lb, ub = learner.estimate_ate(X=X, treatment=treatment, y=y, p=e)
assert (ate_p >= lb) and (ate_p <= ub)
assert ape(tau.mean(), ate_p) < ERROR_THRESHOLD
# check pre-train model
ate_p_pt, lb_pt, ub_pt = learner.estimate_ate(X=X,
treatment=treatment,
y=y,
p=e,
pretrain=True)
assert (ate_p_pt == ate_p) and (lb_pt == lb) and (ub_pt == ub)
# check the accuracy of the CATE estimation with the bootstrap CI
cate_p, _, _ = learner.fit_predict(X=X,
treatment=treatment,
y=y,
p=e,
return_ci=True,
n_bootstraps=10)
auuc_metrics = pd.DataFrame({
"cate_p": cate_p.flatten(),
"W": treatment,
"y": y,
"treatment_effect_col": tau,
})
cumgain = get_cumgain(auuc_metrics,
outcome_col="y",
treatment_col="W",
treatment_effect_col="tau")
# Check if the cumulative gain when using the model's prediction is
# higher than it would be under random targeting
assert cumgain["cate_p"].sum() > cumgain["Random"].sum()
class XLearner:
"""Wrapper of the BaseXRegressor from causalml
Defaults to `sklearn.linear_model.LassoLars` as a base learner if not specified
otherwise. Allows to either specify one learner for all or four distinct learners
for the tasks
- outcome control
- outcome treated
- effect control
- effect treated
References:
[1] CausalML Framework `on Github <https://github.com/uber/causalml/>'_.
[2] S. R. Künzel, J. S. Sekhon, P. J. Bickel, and B. Yu,
“Meta-learners for Estimating Heterogeneous
Treatment Effects using Machine Learning,” 2019.
"""
def __init__(
self,
learner=None,
outcome_learner_c=None,
outcome_learner_t=None,
effect_learner_c=None,
effect_learner_t=None,
):
"""Setup a XLearner
All learners must have ``fit(x, y)`` and ``predict(x)`` methods.
Args:
learner: default learner for all roles
outcome_learner_c: specific learner for control outcome function
outcome_learner_t: specific learner for treated outcome function
effect_learner_c: specific learner for treated effect
effect_learner_t: specific learner for control effect
"""
from causalml.inference.meta import BaseXRegressor
if (learner is not None) or ((outcome_learner_c is not None) and
(outcome_learner_t is not None) and
(effect_learner_c is not None) and
(effect_learner_t is not None)):
self.model = BaseXRegressor(
learner,
outcome_learner_c,
outcome_learner_t,
effect_learner_c,
effect_learner_t,
)
else:
# Assign default learner
learner = LassoLars()
self.model = BaseXRegressor(
learner,
outcome_learner_c,
outcome_learner_t,
effect_learner_c,
effect_learner_t,
)
def __str__(self):
"""Simple string representation for logs and outputs"""
return "{}(outcome_c={}, outcome_t={}, effect_c={}, effect_t={})".format(
self.__class__.__name__,
self.model.model_mu_c.__class__.__name__,
self.model.model_mu_t.__class__.__name__,
self.model.model_tau_c.__class__.__name__,
self.model.model_tau_t.__class__.__name__,
)
def __repr__(self):
return self.__str__()
def fit(self, x: np.array, t: np.array, y: np.array) -> None:
"""Fits the RLearner on given samples
Args:
x: covariate matrix of shape (num_instances, num_features)
t: treatment indicator vector, shape (num_instances)
y: factual outcomes, (num_instances)
"""
self.model.fit(x, t, y)
def predict_ite(
self,
x: np.array,
t: np.array = None,
y: np.array = None,
p: Optional[np.array] = None,
return_components: bool = False,
replace_factuals: bool = False,
) -> np.array:
"""Predicts ITE for the given population
If propensities ``p`` are not given, they are estimated using the default
implementation `justcause.learners.propensities.estimate_propensities`
Args:
x: covariates
t: treatment indicator
y: factual outcomes
p: propensity scores
return_components: whether to return Y(1) and Y(0) for all instances or not
replace_factuals: whether to replace predicted outcomes with the factual
outcomes where applicable
Returns:
the ITE prediction either with or without components
"""
if p is None:
# Set default propensity, because CausalML currently requires it
p = estimate_propensities(x, t)
if return_components:
ite, y_0, y_1 = self.model.predict(x,
p,
t,
y,
return_components=True)
if t is not None and y is not None and replace_factuals:
y_0, y_1 = replace_factual_outcomes(y_0, y_1, y, t)
return ite.flatten(), y_0.flatten(), y_1.flatten()
else:
return self.model.predict(x, p, t, y,
return_components=False).flatten()
def estimate_ate(
self,
x: np.array,
t: np.array,
y: np.array,
propensities: Optional[np.array] = None,
) -> float:
"""Predicts ATE for given samples as mean of ITE predictions"""
self.fit(x, t, y)
return float(np.mean(self.predict_ite(x, t, y, propensities)))