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,
p_model=propensity_model,
X_pred=X,
treatment_pred=w,
)
self.propensity_model = p_model
self.propensity = p
def _treatment_estimate(self, X, w):
self._w_pred = np.zeros(len(w))
for train_index, test_index in self.cv.split(w):
X_train, X_test = X[train_index], X[test_index]
w_train, w_test = w[train_index], w[test_index]
self._w_pred[test_index], _ = compute_propensity_score(
X=X_train,
treatment=w_train,
X_pred=X_test,
treatment_pred=w_test,
calibrate_p=self.calibration,
)
self._w_pred = np.clip(self._w_pred,
a_min=self.clip_bounds[0],
a_max=self.clip_bounds[1])
def fit(self, X, treatment, y, p=None):
"""Fit the inference model.
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
p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
check_treatment_vector(treatment, self.control_name)
self.t_groups = np.unique(treatment[treatment != self.control_name])
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
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()
}
self._classes = {group: i for i, group in enumerate(self.t_groups)}
self.models_mu_c = {
group: deepcopy(self.model_mu_c)
for group in self.t_groups
}
self.models_mu_t = {
group: deepcopy(self.model_mu_t)
for group in self.t_groups
}
self.models_tau_c = {
group: deepcopy(self.model_tau_c)
for group in self.t_groups
}
self.models_tau_t = {
group: deepcopy(self.model_tau_t)
for group in self.t_groups
}
self.vars_c = {}
self.vars_t = {}
for group in self.t_groups:
mask = (treatment == group) | (treatment == self.control_name)
treatment_filt = treatment[mask]
X_filt = X[mask]
y_filt = y[mask]
w = (treatment_filt == group).astype(int)
# Train outcome models
self.models_mu_c[group].fit(X_filt[w == 0], y_filt[w == 0])
self.models_mu_t[group].fit(X_filt[w == 1], y_filt[w == 1])
# Calculate variances and treatment effects
var_c = (y_filt[w == 0] -
self.models_mu_c[group].predict(X_filt[w == 0])).var()
self.vars_c[group] = var_c
var_t = (y_filt[w == 1] -
self.models_mu_t[group].predict(X_filt[w == 1])).var()
self.vars_t[group] = var_t
# Train treatment models
d_c = self.models_mu_t[group].predict(
X_filt[w == 0]) - y_filt[w == 0]
d_t = y_filt[w == 1] - self.models_mu_c[group].predict(
X_filt[w == 1])
self.models_tau_c[group].fit(X_filt[w == 0], d_c)
self.models_tau_t[group].fit(X_filt[w == 1], d_t)
def fit(self, X, treatment, y, p=None, verbose=True):
"""Fit the treatment effect and outcome models of the R learner.
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
y (np.array or pd.Series): an outcome vector
p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
verbose (bool, optional): whether to output progress logs
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
check_treatment_vector(treatment, self.control_name)
self.t_groups = np.unique(treatment[treatment != self.control_name])
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,
cv=self.cv)
self.propensity_model = p_model
self.propensity = p
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()
}
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
}
self.vars_c = {}
self.vars_t = {}
if verbose:
logger.info('generating out-of-fold CV outcome estimates')
yhat = cross_val_predict(self.model_mu, X, y, cv=self.cv, n_jobs=-1)
for group in self.t_groups:
treatment_mask = (treatment == group) | (treatment
== self.control_name)
treatment_filt = treatment[treatment_mask]
w = (treatment_filt == group).astype(int)
X_filt = X[treatment_mask]
y_filt = y[treatment_mask]
yhat_filt = yhat[treatment_mask]
p_filt = p[group][treatment_mask]
if verbose:
logger.info(
'training the treatment effect model for {} with R-loss'.
format(group))
if self.early_stopping:
X_train_filt, X_test_filt, y_train_filt, y_test_filt, yhat_train_filt, yhat_test_filt, \
w_train, w_test, p_train_filt, p_test_filt = train_test_split(
X_filt, y_filt, yhat_filt, w, p_filt,
test_size=self.test_size, random_state=self.random_state
)
self.models_tau[group].fit(
X=X_train_filt,
y=(y_train_filt - yhat_train_filt) /
(w_train - p_train_filt),
sample_weight=(w_train - p_train_filt)**2,
eval_set=[(X_test_filt, (y_test_filt - yhat_test_filt) /
(w_test - p_test_filt))],
sample_weight_eval_set=[(w_test - p_test_filt)**2],
eval_metric=self.effect_learner_eval_metric,
early_stopping_rounds=self.early_stopping_rounds,
verbose=verbose)
else:
self.models_tau[group].fit(
X_filt, (y_filt - yhat_filt) / (w - p_filt),
sample_weight=(w - p_filt)**2,
eval_metric=self.effect_learner_eval_metric)
self.vars_c[group] = (y_filt[w == 0] - yhat_filt[w == 0]).var()
self.vars_t[group] = (y_filt[w == 1] - yhat_filt[w == 1]).var()
def fit(self, X, treatment, y, p=None, seed=None):
"""Fit the inference model.
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
p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
seed (int): random seed for cross-fitting
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
check_treatment_vector(treatment, self.control_name)
self.t_groups = np.unique(treatment[treatment != self.control_name])
self.t_groups.sort()
self._classes = {group: i for i, group in enumerate(self.t_groups)}
# The estimator splits the data into 3 partitions for cross-fit on the propensity score estimation,
# the outcome regression, and the treatment regression on the doubly robust estimates. The use of
# the partitions is rotated so we do not lose on the sample size.
cv = KFold(n_splits=3, shuffle=True, random_state=seed)
split_indices = [index for _, index in cv.split(y)]
self.models_mu_c = [
deepcopy(self.model_mu_c),
deepcopy(self.model_mu_c),
deepcopy(self.model_mu_c),
]
self.models_mu_t = {
group: [
deepcopy(self.model_mu_t),
deepcopy(self.model_mu_t),
deepcopy(self.model_mu_t),
]
for group in self.t_groups
}
self.models_tau = {
group: [
deepcopy(self.model_tau),
deepcopy(self.model_tau),
deepcopy(self.model_tau),
]
for group in self.t_groups
}
if p is None:
self.propensity = {
group: np.zeros(y.shape[0])
for group in self.t_groups
}
for ifold in range(3):
treatment_idx = split_indices[ifold]
outcome_idx = split_indices[(ifold + 1) % 3]
tau_idx = split_indices[(ifold + 2) % 3]
treatment_treat, treatment_out, treatment_tau = (
treatment[treatment_idx],
treatment[outcome_idx],
treatment[tau_idx],
)
y_out, y_tau = y[outcome_idx], y[tau_idx]
X_treat, X_out, X_tau = X[treatment_idx], X[outcome_idx], X[
tau_idx]
if p is None:
logger.info("Generating propensity score")
cur_p = dict()
for group in self.t_groups:
mask = (treatment_treat == group) | (treatment_treat
== self.control_name)
treatment_filt = treatment_treat[mask]
X_filt = X_treat[mask]
w_filt = (treatment_filt == group).astype(int)
w = (treatment_tau == group).astype(int)
cur_p[group], _ = compute_propensity_score(
X=X_filt,
treatment=w_filt,
X_pred=X_tau,
treatment_pred=w)
self.propensity[group][tau_idx] = cur_p[group]
else:
cur_p = dict()
if isinstance(p, (np.ndarray, pd.Series)):
cur_p = {self.t_groups[0]: convert_pd_to_np(p[tau_idx])}
else:
cur_p = {g: prop[tau_idx] for g, prop in p.items()}
check_p_conditions(cur_p, self.t_groups)
logger.info("Generate outcome regressions")
self.models_mu_c[ifold].fit(
X_out[treatment_out == self.control_name],
y_out[treatment_out == self.control_name],
)
for group in self.t_groups:
self.models_mu_t[group][ifold].fit(
X_out[treatment_out == group],
y_out[treatment_out == group])
logger.info("Fit pseudo outcomes from the DR formula")
for group in self.t_groups:
mask = (treatment_tau == group) | (treatment_tau
== self.control_name)
treatment_filt = treatment_tau[mask]
X_filt = X_tau[mask]
y_filt = y_tau[mask]
w_filt = (treatment_filt == group).astype(int)
p_filt = cur_p[group][mask]
mu_t = self.models_mu_t[group][ifold].predict(X_filt)
mu_c = self.models_mu_c[ifold].predict(X_filt)
dr = ((w_filt - p_filt) / p_filt / (1 - p_filt) *
(y_filt - mu_t * w_filt - mu_c *
(1 - w_filt)) + mu_t - mu_c)
self.models_tau[group][ifold].fit(X_filt, dr)
def fit(
self, X, assignment, treatment, y, p=None, pZ=None, seed=None, calibrate=True
):
"""Fit the inference model.
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
assignment (np.array or pd.Series): a (0,1)-valued assignment vector. The assignment is the
instrumental variable that does not depend on unknown confounders. The assignment status
influences treatment in a monotonic way, i.e. one can only be more likely to take the
treatment if assigned.
treatment (np.array or pd.Series): a treatment vector
y (np.array or pd.Series): an outcome vector
p (2-tuple of np.ndarray or pd.Series or dict, optional): The first (second) element corresponds to
unassigned (assigned) units. Each is an array of propensity scores of float (0,1) in the single-treatment
case; or, a dictionary of treatment groups that map to propensity vectors of float (0,1). If None will run
ElasticNetPropensityModel() to generate the propensity scores.
pZ (np.array or pd.Series, optional): an array of assignment probability of float (0,1); if None
will run ElasticNetPropensityModel() to generate the assignment probability score.
seed (int): random seed for cross-fitting
"""
X, treatment, assignment, y = convert_pd_to_np(X, treatment, assignment, y)
check_treatment_vector(treatment, self.control_name)
self.t_groups = np.unique(treatment[treatment != self.control_name])
self.t_groups.sort()
self._classes = {group: i for i, group in enumerate(self.t_groups)}
# The estimator splits the data into 3 partitions for cross-fit on the propensity score estimation,
# the outcome regression, and the treatment regression on the doubly robust estimates. The use of
# the partitions is rotated so we do not lose on the sample size. We do not cross-fit the assignment
# score estimation as the assignment process is usually simple.
cv = KFold(n_splits=3, shuffle=True, random_state=seed)
split_indices = [index for _, index in cv.split(y)]
self.models_mu_c = {
group: [
deepcopy(self.model_mu_c),
deepcopy(self.model_mu_c),
deepcopy(self.model_mu_c),
]
for group in self.t_groups
}
self.models_mu_t = {
group: [
deepcopy(self.model_mu_t),
deepcopy(self.model_mu_t),
deepcopy(self.model_mu_t),
]
for group in self.t_groups
}
self.models_tau = {
group: [
deepcopy(self.model_tau),
deepcopy(self.model_tau),
deepcopy(self.model_tau),
]
for group in self.t_groups
}
if p is None:
self.propensity_1 = {
group: np.zeros(y.shape[0]) for group in self.t_groups
} # propensity scores for those assigned
self.propensity_0 = {
group: np.zeros(y.shape[0]) for group in self.t_groups
} # propensity scores for those not assigned
if pZ is None:
self.propensity_assign, _ = compute_propensity_score(
X=X,
treatment=assignment,
X_pred=X,
treatment_pred=assignment,
calibrate_p=calibrate,
)
else:
self.propensity_assign = pZ
for ifold in range(3):
treatment_idx = split_indices[ifold]
outcome_idx = split_indices[(ifold + 1) % 3]
tau_idx = split_indices[(ifold + 2) % 3]
treatment_treat, treatment_out, treatment_tau = (
treatment[treatment_idx],
treatment[outcome_idx],
treatment[tau_idx],
)
assignment_treat, assignment_out, assignment_tau = (
assignment[treatment_idx],
assignment[outcome_idx],
assignment[tau_idx],
)
y_out, y_tau = y[outcome_idx], y[tau_idx]
X_treat, X_out, X_tau = X[treatment_idx], X[outcome_idx], X[tau_idx]
pZ_tau = self.propensity_assign[tau_idx]
if p is None:
logger.info("Generating propensity score")
cur_p_1 = dict()
cur_p_0 = dict()
for group in self.t_groups:
mask = (treatment_treat == group) | (
treatment_treat == self.control_name
)
mask_1, mask_0 = mask & (assignment_treat == 1), mask & (
assignment_treat == 0
)
cur_p_1[group], _ = compute_propensity_score(
X=X_treat[mask_1],
treatment=(treatment_treat[mask_1] == group).astype(int),
X_pred=X_tau,
treatment_pred=(treatment_tau == group).astype(int),
)
if (treatment_treat[mask_0] == group).sum() == 0:
cur_p_0[group] = np.zeros(X_tau.shape[0])
else:
cur_p_0[group], _ = compute_propensity_score(
X=X_treat[mask_0],
treatment=(treatment_treat[mask_0] == group).astype(int),
X_pred=X_tau,
treatment_pred=(treatment_tau == group).astype(int),
)
self.propensity_1[group][tau_idx] = cur_p_1[group]
self.propensity_0[group][tau_idx] = cur_p_0[group]
else:
cur_p_1 = dict()
cur_p_0 = dict()
if isinstance(p[0], (np.ndarray, pd.Series)):
cur_p_0 = {self.t_groups[0]: convert_pd_to_np(p[0][tau_idx])}
else:
cur_p_0 = {g: prop[tau_idx] for g, prop in p[0].items()}
check_p_conditions(cur_p_0, self.t_groups)
if isinstance(p[1], (np.ndarray, pd.Series)):
cur_p_1 = {self.t_groups[0]: convert_pd_to_np(p[1][tau_idx])}
else:
cur_p_1 = {g: prop[tau_idx] for g, prop in p[1].items()}
check_p_conditions(cur_p_1, self.t_groups)
logger.info("Generate outcome regressions")
for group in self.t_groups:
mask = (treatment_out == group) | (treatment_out == self.control_name)
mask_1, mask_0 = mask & (assignment_out == 1), mask & (
assignment_out == 0
)
self.models_mu_c[group][ifold].fit(X_out[mask_0], y_out[mask_0])
self.models_mu_t[group][ifold].fit(X_out[mask_1], y_out[mask_1])
logger.info("Fit pseudo outcomes from the DR formula")
for group in self.t_groups:
mask = (treatment_tau == group) | (treatment_tau == self.control_name)
treatment_filt = treatment_tau[mask]
X_filt = X_tau[mask]
y_filt = y_tau[mask]
w_filt = (treatment_filt == group).astype(int)
p_1_filt = cur_p_1[group][mask]
p_0_filt = cur_p_0[group][mask]
z_filt = assignment_tau[mask]
pZ_filt = pZ_tau[mask]
mu_t = self.models_mu_t[group][ifold].predict(X_filt)
mu_c = self.models_mu_c[group][ifold].predict(X_filt)
dr = (
z_filt * (y_filt - mu_t) / pZ_filt
- (1 - z_filt) * (y_filt - mu_c) / (1 - pZ_filt)
+ mu_t
- mu_c
)
weight = (
z_filt * (w_filt - p_1_filt) / pZ_filt
- (1 - z_filt) * (w_filt - p_0_filt) / (1 - pZ_filt)
+ p_1_filt
- p_0_filt
)
dr /= weight
self.models_tau[group][ifold].fit(X_filt, dr, sample_weight=weight**2)
