in causalml/inference/meta/xlearner.py [0:0]
def __init__(self,
learner=None,
control_outcome_learner=None,
treatment_outcome_learner=None,
control_effect_learner=None,
treatment_effect_learner=None,
ate_alpha=.05,
control_name=0):
"""Initialize a X-learner.
Args:
learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
groups
control_outcome_learner (optional): a model to estimate outcomes in the control group
treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
control_effect_learner (optional): a model to estimate treatment effects in the control group
treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
ate_alpha (float, optional): the confidence level alpha of the ATE estimate
control_name (str or int, optional): name of control group
"""
assert (learner is not None) or ((control_outcome_learner is not None) and
(treatment_outcome_learner is not None) and
(control_effect_learner is not None) and
(treatment_effect_learner is not None))
if control_outcome_learner is None:
self.model_mu_c = deepcopy(learner)
else:
self.model_mu_c = control_outcome_learner
if treatment_outcome_learner is None:
self.model_mu_t = deepcopy(learner)
else:
self.model_mu_t = treatment_outcome_learner
if control_effect_learner is None:
self.model_tau_c = deepcopy(learner)
else:
self.model_tau_c = control_effect_learner
if treatment_effect_learner is None:
self.model_tau_t = deepcopy(learner)
else:
self.model_tau_t = treatment_effect_learner
self.ate_alpha = ate_alpha
self.control_name = control_name
self.propensity = None
self.propensity_model = None