def fit(self, X, treatment, y):
"""Fit the inference model.
Args:
X (np.matrix): a feature matrix
treatment (np.array): a treatment vector
y (np.array): an outcome vector
"""
check_control_in_treatment(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)}
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, p, treatment, y, verbose=True):
"""Fit the treatment effect and outcome models of the R learner.
Args:
X (np.matrix): a feature matrix
p (np.ndarray or dict): 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)
treatment (np.array): a treatment vector
y (np.array): an outcome vector
"""
check_control_in_treatment(treatment, self.control_name)
self.t_groups = np.unique(treatment[treatment != self.control_name])
self.t_groups.sort()
check_p_conditions(p, self.t_groups)
if isinstance(p, np.ndarray):
treatment_name = self.t_groups[0]
p = {treatment_name: p}
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,
method='predict_proba',
n_jobs=-1)[:, 1]
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]
yhat_filt = yhat[mask]
p_filt = p[group][mask]
w = (treatment_filt == group).astype(int)
if verbose:
logger.info(
'training the treatment effect model for {} with R-loss'.
format(group))
self.models_tau[group].fit(X_filt,
(y_filt - yhat_filt) / (w - p_filt),
sample_weight=(w - p_filt)**2)
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, p, treatment, y, verbose=True):
"""Fit the treatment effect and outcome models of the R learner.
Args:
X (np.matrix): a feature matrix
p (np.ndarray or dict): 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)
treatment (np.array): a treatment vector
y (np.array): an outcome vector
"""
check_control_in_treatment(treatment, self.control_name)
self.t_groups = np.unique(treatment[treatment != self.control_name])
self.t_groups.sort()
check_p_conditions(p, self.t_groups)
if isinstance(p, np.ndarray):
treatment_name = self.t_groups[0]
p = {treatment_name: p}
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 estimate_ate(self, X, p, treatment, y, segment=None, return_ci=False):
"""Estimate the Average Treatment Effect (ATE).
Args:
X (np.matrix): A feature matrix
p (np.ndarray or dict): 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)
treatment (np.array): A treatment vector of int
y (np.array): an outcome vector
segment (np.array, optional): An optional segment vector of int. If given, the ATE and its CI will be
estimated for each segment.
return_ci (bool, optional): Whether to return confidence intervals
Returns:
(tuple): The ATE and its confidence interval (LB, UB) for each treatment, t and segment, s
"""
check_control_in_treatment(treatment, self.control_name)
self.t_groups = np.unique(treatment[treatment != self.control_name])
self.t_groups.sort()
check_p_conditions(p, self.t_groups)
if isinstance(p, np.ndarray):
treatment_name = self.t_groups[0]
p = {treatment_name: p}
ate = []
ate_lb = []
ate_ub = []
for i, group in enumerate(self.t_groups):
logger.info('Estimating ATE for group {}.'.format(group))
w_group = (treatment == group).astype(int)
p_group = p[group]
if self.calibrate_propensity:
logger.info('Calibrating propensity scores.')
p_group = calibrate(p_group, w_group)
yhat_c = np.zeros_like(y, dtype=float)
yhat_t = np.zeros_like(y, dtype=float)
if self.cv:
for i_fold, (i_trn, i_val) in enumerate(self.cv.split(X, y), 1):
logger.info('Training an outcome model for CV #{}'.format(i_fold))
self.model_tau.fit(np.hstack((X[i_trn], w_group[i_trn].reshape(-1, 1))), y[i_trn])
yhat_c[i_val] = self.model_tau.predict(np.hstack((X[i_val], np.zeros((len(i_val), 1)))))
yhat_t[i_val] = self.model_tau.predict(np.hstack((X[i_val], np.ones((len(i_val), 1)))))
else:
self.model_tau.fit(np.hstack((X, w_group.reshape(-1, 1))), y)
yhat_c = self.model_tau.predict(np.hstack((X, np.zeros((len(y), 1)))))
yhat_t = self.model_tau.predict(np.hstack((X, np.ones((len(y), 1)))))
if segment is None:
logger.info('Training the TMLE learner.')
_ate, se = simple_tmle(y, w_group, yhat_c, yhat_t, p_group)
_ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
_ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
else:
assert segment.shape[0] == X.shape[0] and segment.ndim == 1, 'Segment must be the 1-d np.array of int.'
segments = np.unique(segment)
_ate = []
_ate_lb = []
_ate_ub = []
for s in sorted(segments):
logger.info('Training the TMLE learner for segment {}.'.format(s))
filt = (segment == s) & (yhat_c < np.quantile(yhat_c, q=.99))
_ate_s, se = simple_tmle(y[filt], w_group[filt], yhat_c[filt], yhat_t[filt], p_group[filt])
_ate_lb_s = _ate_s - se * norm.ppf(1 - self.ate_alpha / 2)
_ate_ub_s = _ate_s + se * norm.ppf(1 - self.ate_alpha / 2)
_ate.append(_ate_s)
_ate_lb.append(_ate_lb_s)
_ate_ub.append(_ate_ub_s)
ate.append(_ate)
ate_lb.append(_ate_lb)
ate_ub.append(_ate_ub)
return np.array(ate), np.array(ate_lb), np.array(ate_ub)