def bootstrap(self, X, assignment, treatment, y, p, pZ, size=10000, seed=None):
"""Runs a single bootstrap. Fits on bootstrapped sample, then predicts on whole population."""
idxs = np.random.choice(np.arange(0, X.shape[0]), size=size)
X_b = X[idxs]
if isinstance(p[0], (np.ndarray, pd.Series)):
p0_b = {self.t_groups[0]: convert_pd_to_np(p[0][idxs])}
else:
p0_b = {g: prop[idxs] for g, prop in p[0].items()}
if isinstance(p[1], (np.ndarray, pd.Series)):
p1_b = {self.t_groups[0]: convert_pd_to_np(p[1][idxs])}
else:
p1_b = {g: prop[idxs] for g, prop in p[1].items()}
pZ_b = pZ[idxs]
assignment_b = assignment[idxs]
treatment_b = treatment[idxs]
y_b = y[idxs]
self.fit(
X=X_b,
assignment=assignment_b,
treatment=treatment_b,
y=y_b,
p=(p0_b, p1_b),
pZ=pZ_b,
seed=seed,
)
te_b = self.predict(X=X)
return te_b
def __init__(self,
method,
control_name,
X,
tau,
classes,
model_tau=None,
features=None,
normalize=True,
test_size=0.3,
random_state=None,
override_checks=False,
r_learners=None):
"""
The Explainer class handles all feature explanation/interpretation functions, including plotting
feature importances, shapley value distributions, and shapley value dependency plots.
Currently supported methods are:
- auto (calculates importance based on estimator's default implementation of feature importance;
estimator must be tree-based)
Note: if none provided, it uses lightgbm's LGBMRegressor as estimator, and "gain" as
importance type
- permutation (calculates importance based on mean decrease in accuracy when a feature column is permuted; estimator can be any form)
- shapley (calculates shapley values; estimator must be tree-based)
Hint: for permutation, downsample data for better performance especially if X.shape[1] is large
Args:
method (str): auto, permutation, shapley
control_name (str/int/float): name of control group
X (np.matrix): a feature matrix
tau (np.array): a treatment effect vector (estimated/actual)
classes (dict): a mapping of treatment names to indices (used for indexing tau array)
model_tau (sklearn/lightgbm/xgboost model object): a model object
features (np.array): list/array of feature names. If None, an enumerated list will be used.
normalize (bool): normalize by sum of importances if method=auto (defaults to True)
test_size (float/int): if float, represents the proportion of the dataset to include in the test split. If int, represents the absolute number of test samples (used for estimating permutation importance)
random_state (int/RandomState instance/None): random state used in permutation importance estimation
override_checks (bool): overrides self.check_conditions (e.g. if importance/shapley values are pre-computed)
r_learners (dict): a mapping of treatment group to fitted R Learners
"""
self.method = method
self.control_name = control_name
self.X = convert_pd_to_np(X)
self.tau = convert_pd_to_np(tau)
if self.tau is not None and self.tau.ndim == 1:
self.tau = self.tau.reshape(-1, 1)
self.classes = classes
self.model_tau = LGBMRegressor(
importance_type='gain') if model_tau is None else model_tau
self.features = features
self.normalize = normalize
self.test_size = test_size
self.random_state = random_state
self.override_checks = override_checks
self.r_learners = r_learners
if not self.override_checks:
self.check_conditions()
self.create_feature_names()
self.build_new_tau_models()
def fit(self, X, p, treatment, y, 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
p (np.ndarray or pd.Series 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 or pd.Series): a treatment vector
y (np.array or pd.Series): an outcome vector
"""
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()
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,
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_predict(self, X, treatment, y, p=None, return_ci=False,
n_bootstraps=1000, bootstrap_size=10000, verbose=True):
"""Fit the treatment effect and outcome models of the R learner and predict treatment effects.
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.
return_ci (bool): whether to return confidence intervals
n_bootstraps (int): number of bootstrap iterations
bootstrap_size (int): number of samples per bootstrap
verbose (bool): whether to output progress logs
Returns:
(numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
UB [n_samples, n_treatment]
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
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:
t_groups_global = self.t_groups
_classes_global = self._classes
model_mu_global = deepcopy(self.model_mu)
models_tau_global = deepcopy(self.models_tau)
te_bootstraps = np.zeros(shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps))
logger.info('Bootstrap Confidence Intervals')
for i in tqdm(range(n_bootstraps)):
te_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
te_bootstraps[:, :, i] = te_b
te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
te_upper = np.percentile(te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2)
# set member variables back to global (currently last bootstrapped outcome)
self.t_groups = t_groups_global
self._classes = _classes_global
self.model_mu = deepcopy(model_mu_global)
self.models_tau = deepcopy(models_tau_global)
return (te, te_lower, te_upper)
def predict(self, X, p, treatment=None, y=None, return_components=False, verbose=True):
"""Predict treatment effects.
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
p (np.ndarray or pd.Series 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 or pd.Series, optional): a treatment vector
y (np.array or pd.Series, optional): an outcome vector
return_components (bool, optional): whether to return outcome for treatment and control seperately
Returns:
(numpy.ndarray): Predictions of treatment effects.
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
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 = np.zeros((X.shape[0], self.t_groups.shape[0]))
dhat_cs = {}
dhat_ts = {}
for i, group in enumerate(self.t_groups):
model_tau_c = self.models_tau_c[group]
model_tau_t = self.models_tau_t[group]
dhat_cs[group] = model_tau_c.predict(X)
dhat_ts[group] = model_tau_t.predict(X)
_te = (p[group] * dhat_cs[group] + (1 - p[group]) * dhat_ts[group]).reshape(-1, 1)
te[:, i] = np.ravel(_te)
if (y is not None) and (treatment is not None) and verbose:
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)
yhat = np.zeros_like(y_filt, dtype=float)
yhat[w == 0] = self.models_mu_c[group].predict(X_filt[w == 0])
yhat[w == 1] = self.models_mu_t[group].predict(X_filt[w == 1])
logger.info('Error metrics for group {}'.format(group))
regression_metrics(y_filt, yhat, w)
if not return_components:
return te
else:
return te, dhat_cs, dhat_ts
def fit(self, X, treatment, y):
"""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
"""
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)}
self.models_c = {group: deepcopy(self.model_c) for group in self.t_groups}
self.models_t = {group: deepcopy(self.model_t) for group in self.t_groups}
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)
self.models_c[group].fit(X_filt[w == 0], y_filt[w == 0])
self.models_t[group].fit(X_filt[w == 1], y_filt[w == 1])
def fit(self, X, treatment, y, p=None):
"""
Fits CEVAE.
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
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
self.cevae = CEVAEModel(outcome_dist=self.outcome_dist,
feature_dim=X.shape[-1],
latent_dim=self.latent_dim,
hidden_dim=self.hidden_dim,
num_layers=self.num_layers)
self.cevae.fit(x=torch.tensor(X, dtype=torch.float),
t=torch.tensor(treatment, dtype=torch.float),
y=torch.tensor(y, dtype=torch.float),
num_epochs=self.num_epochs,
batch_size=self.batch_size,
learning_rate=self.learning_rate,
learning_rate_decay=self.learning_rate_decay,
weight_decay=self.weight_decay)
def fit_predict(
self,
X,
treatment,
y,
p=None,
return_ci=False,
n_bootstraps=1000,
bootstrap_size=10000,
return_components=False,
verbose=True,
):
"""Fit the inference model of the T learner and predict treatment effects.
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
return_ci (bool): whether to return confidence intervals
n_bootstraps (int): number of bootstrap iterations
bootstrap_size (int): number of samples per bootstrap
return_components (bool, optional): whether to return outcome for treatment and control seperately
verbose (str): whether to output progress logs
Returns:
(numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
UB [n_samples, n_treatment]
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
self.fit(X, treatment, y)
te = self.predict(X, treatment, y, return_components=return_components)
if not return_ci:
return te
else:
t_groups_global = self.t_groups
_classes_global = self._classes
models_c_global = deepcopy(self.models_c)
models_t_global = deepcopy(self.models_t)
te_bootstraps = np.zeros(shape=(X.shape[0], self.t_groups.shape[0],
n_bootstraps))
logger.info("Bootstrap Confidence Intervals")
for i in tqdm(range(n_bootstraps)):
te_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
te_bootstraps[:, :, i] = te_b
te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100,
axis=2)
te_upper = np.percentile(te_bootstraps,
(1 - self.ate_alpha / 2) * 100,
axis=2)
# set member variables back to global (currently last bootstrapped outcome)
self.t_groups = t_groups_global
self._classes = _classes_global
self.models_c = deepcopy(models_c_global)
self.models_t = deepcopy(models_t_global)
return (te, te_lower, te_upper)
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
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.
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:
self._set_propensity_models(X=X, treatment=treatment, y=y)
p = self.propensity
else:
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
}
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, treatment, y):
"""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
"""
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)}
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):
"""
Fits the DragonNet 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
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
y = np.hstack((y.reshape(-1, 1), treatment.reshape(-1, 1)))
self.dragonnet = self.make_dragonnet(X.shape[1])
metrics = [regression_loss, binary_classification_loss, treatment_accuracy, track_epsilon]
if self.targeted_reg:
loss = make_tarreg_loss(ratio=self.ratio, dragonnet_loss=self.loss_func)
else:
loss = self.loss_func
self.dragonnet.compile(
optimizer=Adam(lr=self.learning_rate),
loss=loss, metrics=metrics)
adam_callbacks = [
TerminateOnNaN(),
EarlyStopping(monitor='val_loss', patience=2, min_delta=0.),
ReduceLROnPlateau(monitor='loss', factor=0.5, patience=5, verbose=self.verbose, mode='auto',
min_delta=1e-8, cooldown=0, min_lr=0)
]
self.dragonnet.fit(X, y,
callbacks=adam_callbacks,
validation_split=self.val_split,
epochs=self.epochs,
batch_size=self.batch_size,
verbose=self.verbose)
sgd_callbacks = [
TerminateOnNaN(),
EarlyStopping(monitor='val_loss', patience=40, min_delta=0.),
ReduceLROnPlateau(monitor='loss', factor=0.5, patience=5, verbose=self.verbose, mode='auto',
min_delta=0., cooldown=0, min_lr=0)
]
sgd_lr = 1e-5
momentum = 0.9
self.dragonnet.compile(optimizer=SGD(lr=sgd_lr, momentum=momentum, nesterov=True), loss=loss, metrics=metrics)
self.dragonnet.fit(X, y,
callbacks=sgd_callbacks,
validation_split=self.val_split,
epochs=300,
batch_size=self.batch_size,
verbose=self.verbose)
def predict(self,
X,
treatment=None,
y=None,
return_components=False,
verbose=True):
"""Predict treatment effects.
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
treatment (np.array or pd.Series, optional): a treatment vector
y (np.array or pd.Series, optional): an outcome vector
return_components (bool, optional): whether to return outcome for treatment and control seperately
verbose (bool, optional): whether to output progress logs
Returns:
(numpy.ndarray): Predictions of treatment effects.
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
yhat_cs = {}
yhat_ts = {}
for group in self.t_groups:
model = self.models[group]
# set the treatment column to zero (the control group)
X_new = np.hstack((np.zeros((X.shape[0], 1)), X))
yhat_cs[group] = model.predict(X_new)
# set the treatment column to one (the treatment group)
X_new[:, 0] = 1
yhat_ts[group] = model.predict(X_new)
if (y is not None) and (treatment is not None) and verbose:
mask = (treatment == group) | (treatment == self.control_name)
treatment_filt = treatment[mask]
w = (treatment_filt == group).astype(int)
y_filt = y[mask]
yhat = np.zeros_like(y_filt, dtype=float)
yhat[w == 0] = yhat_cs[group][mask][w == 0]
yhat[w == 1] = yhat_ts[group][mask][w == 1]
logger.info('Error metrics for group {}'.format(group))
regression_metrics(y_filt, yhat, w)
te = np.zeros((X.shape[0], self.t_groups.shape[0]))
for i, group in enumerate(self.t_groups):
te[:, i] = yhat_ts[group] - yhat_cs[group]
if not return_components:
return te
else:
return te, yhat_cs, yhat_ts
return te
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 predict(self,
X,
treatment=None,
y=None,
p=None,
return_components=False,
verbose=True):
"""Predict treatment effects.
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
treatment (np.array or pd.Series, optional): a treatment vector
y (np.array or pd.Series, optional): an outcome vector
verbose (bool, optional): whether to output progress logs
Returns:
(numpy.ndarray): Predictions of treatment effects.
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
te = np.zeros((X.shape[0], self.t_groups.shape[0]))
yhat_cs = {}
yhat_ts = {}
for i, group in enumerate(self.t_groups):
models_tau = self.models_tau[group]
_te = np.r_[[model.predict(X)
for model in models_tau]].mean(axis=0)
te[:, i] = np.ravel(_te)
yhat_cs[group] = np.r_[[
model.predict(X) for model in self.models_mu_c
]].mean(axis=0)
yhat_ts[group] = np.r_[[
model.predict(X) for model in self.models_mu_t[group]
]].mean(axis=0)
if (y is not None) and (treatment is not None) and verbose:
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)
yhat = np.zeros_like(y_filt, dtype=float)
yhat[w == 0] = yhat_cs[group][mask][w == 0]
yhat[w == 1] = yhat_ts[group][mask][w == 1]
logger.info("Error metrics for group {}".format(group))
regression_metrics(y_filt, yhat, w)
if not return_components:
return te
else:
return te, yhat_cs, yhat_ts
def predict(self, X):
"""Predict treatment effects.
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
Returns:
(numpy.ndarray): Predictions of treatment effects.
"""
X = convert_pd_to_np(X)
te = np.zeros((X.shape[0], self.t_groups.shape[0]))
for i, group in enumerate(self.t_groups):
dhat = self.models_tau[group].predict(X)
te[:, i] = dhat
return te
def fit(self, X, treatment, y, w):
"""Fits the 2SLS 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
w (np.array or pd.Series): an instrument vector
"""
X, treatment, y, w = convert_pd_to_np(X, treatment, y, w)
exog = sm.add_constant(np.c_[X, treatment])
endog = y
instrument = sm.add_constant(np.c_[X, w])
self.iv_model = IV2SLS(endog=endog, exog=exog, instrument=instrument)
self.iv_fit = self.iv_model.fit()
def estimate_ate(self, X, treatment, y, p=None):
"""Estimate the Average Treatment Effect (ATE).
Args:
X (np.matrix, 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
Returns:
The mean and confidence interval (LB, UB) of the ATE estimate.
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
self.fit(X, treatment, y)
ate = np.zeros(self.t_groups.shape[0])
ate_lb = np.zeros(self.t_groups.shape[0])
ate_ub = np.zeros(self.t_groups.shape[0])
for i, group in enumerate(self.t_groups):
ate[i] = self.models[group].coefficients[0]
ate_lb[i] = self.models[group].conf_ints[0, 0]
ate_ub[i] = self.models[group].conf_ints[0, 1]
return ate, ate_lb, ate_ub
def fit(self, X, treatment, y):
"""Fit the inference model
Args:
X (np.matrix, 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
"""
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)}
self.models = {group: deepcopy(self.model) for group in self.t_groups}
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)
X_new = np.hstack((w.reshape((-1, 1)), X_filt))
self.models[group].fit(X_new, y_filt)
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 estimate_ate(self,
X,
treatment,
y,
p=None,
bootstrap_ci=False,
n_bootstraps=1000,
bootstrap_size=10000):
"""Estimate the Average Treatment Effect (ATE).
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.
bootstrap_ci (bool): whether run bootstrap for confidence intervals
n_bootstraps (int): number of bootstrap iterations
bootstrap_size (int): number of samples per bootstrap
Returns:
The mean and confidence interval (LB, UB) of the ATE estimate.
"""
te, dhat_cs, dhat_ts = self.fit_predict(X,
treatment,
y,
p,
return_components=True)
X, treatment, y = convert_pd_to_np(X, treatment, y)
if p is None:
p = self.propensity
else:
check_p_conditions(p, self.t_groups)
if isinstance(p, np.ndarray):
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()
}
ate = np.zeros(self.t_groups.shape[0])
ate_lb = np.zeros(self.t_groups.shape[0])
ate_ub = np.zeros(self.t_groups.shape[0])
for i, group in enumerate(self.t_groups):
_ate = te[:, i].mean()
mask = (treatment == group) | (treatment == self.control_name)
treatment_filt = treatment[mask]
w = (treatment_filt == group).astype(int)
prob_treatment = float(sum(w)) / w.shape[0]
dhat_c = dhat_cs[group][mask]
dhat_t = dhat_ts[group][mask]
p_filt = p[group][mask]
# SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.
# "Recent Developments in the Econometrics of Program Evaluation." Journal of Economic Literature
se = np.sqrt(
(self.vars_t[group] / prob_treatment + self.vars_c[group] /
(1 - prob_treatment) +
(p_filt * dhat_c + (1 - p_filt) * dhat_t).var()) / w.shape[0])
_ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
_ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
ate[i] = _ate
ate_lb[i] = _ate_lb
ate_ub[i] = _ate_ub
if not bootstrap_ci:
return ate, ate_lb, ate_ub
else:
t_groups_global = self.t_groups
_classes_global = self._classes
models_mu_c_global = deepcopy(self.models_mu_c)
models_mu_t_global = deepcopy(self.models_mu_t)
models_tau_c_global = deepcopy(self.models_tau_c)
models_tau_t_global = deepcopy(self.models_tau_t)
logger.info('Bootstrap Confidence Intervals for ATE')
ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0],
n_bootstraps))
for n in tqdm(range(n_bootstraps)):
cate_b = self.bootstrap(X,
treatment,
y,
p,
size=bootstrap_size)
ate_bootstraps[:, n] = cate_b.mean()
ate_lower = np.percentile(ate_bootstraps,
(self.ate_alpha / 2) * 100,
axis=1)
ate_upper = np.percentile(ate_bootstraps,
(1 - self.ate_alpha / 2) * 100,
axis=1)
# set member variables back to global (currently last bootstrapped outcome)
self.t_groups = t_groups_global
self._classes = _classes_global
self.models_mu_c = deepcopy(models_mu_c_global)
self.models_mu_t = deepcopy(models_mu_t_global)
self.models_tau_c = deepcopy(models_tau_c_global)
self.models_tau_t = deepcopy(models_tau_t_global)
return ate, ate_lower, ate_upper
def fit(self, X, p, treatment, y, 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
p (np.ndarray or pd.Series 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 or pd.Series): a treatment vector
y (np.array or pd.Series): an outcome vector
"""
check_treatment_vector(treatment, self.control_name)
X, treatment, y = convert_pd_to_np(X, treatment, y)
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: 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 estimate_ate(self, X, p, treatment, y, bootstrap_ci=False, n_bootstraps=1000, bootstrap_size=10000):
"""Estimate the Average Treatment Effect (ATE).
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
p (np.ndarray or pd.Series 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 or pd.Series): a treatment vector
y (np.array or pd.Series): an outcome vector
bootstrap_ci (bool): whether run bootstrap for confidence intervals
n_bootstraps (int): number of bootstrap iterations
bootstrap_size (int): number of samples per bootstrap
verbose (str): whether to output progress logs
Returns:
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, p, treatment, y)
check_p_conditions(p, self.t_groups)
if isinstance(p, np.ndarray):
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()}
ate = np.zeros(self.t_groups.shape[0])
ate_lb = np.zeros(self.t_groups.shape[0])
ate_ub = np.zeros(self.t_groups.shape[0])
for i, group in enumerate(self.t_groups):
w = (treatment == group).astype(int)
prob_treatment = float(sum(w)) / X.shape[0]
_ate = te[:, i].mean()
se = (np.sqrt((self.vars_t[group] / prob_treatment)
+ (self.vars_c[group] / (1 - prob_treatment))
+ te[:, i].var())
/ X.shape[0])
_ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
_ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
ate[i] = _ate
ate_lb[i] = _ate_lb
ate_ub[i] = _ate_ub
if not bootstrap_ci:
return ate, ate_lb, ate_ub
else:
t_groups_global = self.t_groups
_classes_global = self._classes
model_mu_global = deepcopy(self.model_mu)
models_tau_global = deepcopy(self.models_tau)
logger.info('Bootstrap Confidence Intervals for ATE')
ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
for n in tqdm(range(n_bootstraps)):
cate_b = self.bootstrap(X, p, treatment, y, size=bootstrap_size)
ate_bootstraps[:, n] = cate_b.mean()
ate_lower = np.percentile(ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1)
ate_upper = np.percentile(ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1)
# set member variables back to global (currently last bootstrapped outcome)
self.t_groups = t_groups_global
self._classes = _classes_global
self.model_mu = deepcopy(model_mu_global)
self.models_tau = deepcopy(models_tau_global)
return ate, ate_lower, ate_upper
def estimate_ate(
self,
X,
treatment=None,
y=None,
p=None,
sample_weight=None,
bootstrap_ci=False,
n_bootstraps=1000,
bootstrap_size=10000,
pretrain=False,
):
"""Estimate the Average Treatment Effect (ATE).
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
treatment (np.array or pd.Series): only needed when pretrain=False, a treatment vector
y (np.array or pd.Series):only needed when pretrain=False, 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.
sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
weight of each observation for `effect_learner`. If None, it assumes equal weight.
bootstrap_ci (bool): whether run bootstrap for confidence intervals
n_bootstraps (int): number of bootstrap iterations
bootstrap_size (int): number of samples per bootstrap
pretrain (bool): whether a model has been fit, default False.
Returns:
The mean and confidence interval (LB, UB) of the ATE estimate.
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
if pretrain:
te = self.predict(X, p)
else:
if not len(treatment) or not len(y):
raise ValueError("treatmeng and y must be provided when pretrain=False")
te = self.fit_predict(X, treatment, y, p, sample_weight, return_ci=False)
ate = np.zeros(self.t_groups.shape[0])
ate_lb = np.zeros(self.t_groups.shape[0])
ate_ub = np.zeros(self.t_groups.shape[0])
for i, group in enumerate(self.t_groups):
w = (treatment == group).astype(int)
prob_treatment = float(sum(w)) / X.shape[0]
_ate = te[:, i].mean()
se = (
np.sqrt(
(self.vars_t[group] / prob_treatment)
+ (self.vars_c[group] / (1 - prob_treatment))
+ te[:, i].var()
)
/ X.shape[0]
)
_ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
_ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
ate[i] = _ate
ate_lb[i] = _ate_lb
ate_ub[i] = _ate_ub
if not bootstrap_ci:
return ate, ate_lb, ate_ub
else:
t_groups_global = self.t_groups
_classes_global = self._classes
model_mu_global = deepcopy(self.model_mu)
models_tau_global = deepcopy(self.models_tau)
logger.info("Bootstrap Confidence Intervals for ATE")
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()
ate_lower = np.percentile(
ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1
)
ate_upper = np.percentile(
ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
)
# set member variables back to global (currently last bootstrapped outcome)
self.t_groups = t_groups_global
self._classes = _classes_global
self.model_mu = deepcopy(model_mu_global)
self.models_tau = deepcopy(models_tau_global)
return ate, ate_lower, ate_upper
def fit(self, X, treatment, y, p=None, sample_weight=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.
sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
weight of each observation for `effect_learner`. If None, it assumes equal weight.
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)
# initialize equal sample weight if it's not provided, for simplicity purpose
sample_weight = (
convert_pd_to_np(sample_weight)
if sample_weight is not None
else convert_pd_to_np(np.ones(len(y)))
)
assert len(sample_weight) == len(
y
), "Data length must be equal for sample_weight and the input data"
self.t_groups = np.unique(treatment[treatment != self.control_name])
self.t_groups.sort()
if p is None:
self._set_propensity_models(X=X, treatment=treatment, y=y)
p = self.propensity
else:
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}
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]
sample_weight_filt = sample_weight[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,
sample_weight_train_filt,
sample_weight_test_filt,
) = train_test_split(
X_filt,
y_filt,
yhat_filt,
w,
p_filt,
sample_weight_filt,
test_size=self.test_size,
random_state=self.random_state,
)
weight = sample_weight_filt
self.models_tau[group].fit(
X=X_train_filt,
y=(y_train_filt - yhat_train_filt) / (w_train - p_train_filt),
sample_weight=sample_weight_train_filt
* ((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=[
sample_weight_test_filt * ((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=sample_weight_filt * ((w - p_filt) ** 2),
eval_metric=self.effect_learner_eval_metric,
)
diff_c = y_filt[w == 0] - yhat_filt[w == 0]
diff_t = y_filt[w == 1] - yhat_filt[w == 1]
sample_weight_filt_c = sample_weight_filt[w == 0]
sample_weight_filt_t = sample_weight_filt[w == 1]
self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
def fit(self, X, treatment, y, p=None, sample_weight=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
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.
sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
weight of each observation for `effect_learner`. If None, it assumes equal weight.
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)
if sample_weight is not None:
assert len(sample_weight) == len(
y
), "Data length must be equal for sample_weight and the input data"
sample_weight = convert_pd_to_np(sample_weight)
self.t_groups = np.unique(treatment[treatment != self.control_name])
self.t_groups.sort()
if p is None:
self._set_propensity_models(X=X, treatment=treatment, y=y)
p = self.propensity
else:
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}
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:
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)
weight = (w - p_filt) ** 2
diff_c = y_filt[w == 0] - yhat_filt[w == 0]
diff_t = y_filt[w == 1] - yhat_filt[w == 1]
if sample_weight is not None:
sample_weight_filt = sample_weight[mask]
sample_weight_filt_c = sample_weight_filt[w == 0]
sample_weight_filt_t = sample_weight_filt[w == 1]
self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
weight *= sample_weight_filt # update weight
else:
self.vars_c[group] = diff_c.var()
self.vars_t[group] = diff_t.var()
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=weight
)
def fit_predict(
self,
X,
assignment,
treatment,
y,
p=None,
pZ=None,
return_ci=False,
n_bootstraps=1000,
bootstrap_size=10000,
return_components=False,
verbose=True,
seed=None,
calibrate=True,
):
"""Fit the treatment effect and outcome models of the R learner and predict treatment effects.
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.
return_ci (bool): whether to return confidence intervals
n_bootstraps (int): number of bootstrap iterations
bootstrap_size (int): number of samples per bootstrap
return_components (bool, optional): whether to return outcome for treatment and control seperately
verbose (str): whether to output progress logs
seed (int): random seed for cross-fitting
Returns:
(numpy.ndarray): Predictions of treatment effects for compliers, , i.e. those individuals
who take the treatment only if they are assigned. Output dim: [n_samples, n_treatment]
If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
UB [n_samples, n_treatment]
"""
X, assignment, treatment, y = convert_pd_to_np(X, assignment, treatment, y)
self.fit(X, assignment, treatment, y, p, seed, calibrate)
if p is None:
p = (self.propensity_0, self.propensity_1)
else:
check_p_conditions(p[0], self.t_groups)
check_p_conditions(p[1], self.t_groups)
if isinstance(p[0], (np.ndarray, pd.Series)):
treatment_name = self.t_groups[0]
p = (
{treatment_name: convert_pd_to_np(p[0])},
{treatment_name: convert_pd_to_np(p[1])},
)
elif isinstance(p[0], dict):
p = (
{
treatment_name: convert_pd_to_np(_p)
for treatment_name, _p in p[0].items()
},
{
treatment_name: convert_pd_to_np(_p)
for treatment_name, _p in p[1].items()
},
)
if pZ is None:
pZ = self.propensity_assign
te = self.predict(
X, treatment=treatment, y=y, return_components=return_components
)
if not return_ci:
return te
else:
t_groups_global = self.t_groups
_classes_global = self._classes
models_mu_c_global = deepcopy(self.models_mu_c)
models_mu_t_global = deepcopy(self.models_mu_t)
models_tau_global = deepcopy(self.models_tau)
te_bootstraps = np.zeros(
shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
)
logger.info("Bootstrap Confidence Intervals")
for i in tqdm(range(n_bootstraps)):
te_b = self.bootstrap(
X, assignment, treatment, y, p, pZ, size=bootstrap_size, seed=seed
)
te_bootstraps[:, :, i] = te_b
te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
te_upper = np.percentile(
te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
)
# set member variables back to global (currently last bootstrapped outcome)
self.t_groups = t_groups_global
self._classes = _classes_global
self.models_mu_c = deepcopy(models_mu_c_global)
self.models_mu_t = deepcopy(models_mu_t_global)
self.models_tau = deepcopy(models_tau_global)
return (te, te_lower, te_upper)
def estimate_ate(
self,
X,
treatment,
y,
p=None,
return_ci=False,
bootstrap_ci=False,
n_bootstraps=1000,
bootstrap_size=10000,
pretrain=False,
):
"""Estimate the Average Treatment Effect (ATE).
Args:
X (np.matrix, 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
return_ci (bool, optional): whether to return confidence intervals
bootstrap_ci (bool): whether to return confidence intervals
n_bootstraps (int): number of bootstrap iterations
bootstrap_size (int): number of samples per bootstrap
pretrain (bool): whether a model has been fit, default False.
Returns:
The mean and confidence interval (LB, UB) of the ATE estimate.
"""
X, treatment, y = convert_pd_to_np(X, treatment, y)
if pretrain:
te, yhat_cs, yhat_ts = self.predict(X,
treatment,
y,
return_components=True)
else:
te, yhat_cs, yhat_ts = self.fit_predict(X,
treatment,
y,
return_components=True)
ate = np.zeros(self.t_groups.shape[0])
ate_lb = np.zeros(self.t_groups.shape[0])
ate_ub = np.zeros(self.t_groups.shape[0])
for i, group in enumerate(self.t_groups):
_ate = te[:, i].mean()
mask = (treatment == group) | (treatment == self.control_name)
treatment_filt = treatment[mask]
y_filt = y[mask]
w = (treatment_filt == group).astype(int)
prob_treatment = float(sum(w)) / w.shape[0]
yhat_c = yhat_cs[group][mask]
yhat_t = yhat_ts[group][mask]
se = np.sqrt(
((y_filt[w == 0] - yhat_c[w == 0]).var() /
(1 - prob_treatment) +
(y_filt[w == 1] - yhat_t[w == 1]).var() / prob_treatment +
(yhat_t - yhat_c).var()) / y_filt.shape[0])
_ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
_ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
ate[i] = _ate
ate_lb[i] = _ate_lb
ate_ub[i] = _ate_ub
if not return_ci:
return ate
elif return_ci and not bootstrap_ci:
return ate, ate_lb, ate_ub
else:
t_groups_global = self.t_groups
_classes_global = self._classes
models_global = deepcopy(self.models)
logger.info("Bootstrap Confidence Intervals for ATE")
ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0],
n_bootstraps))
for n in tqdm(range(n_bootstraps)):
ate_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
ate_bootstraps[:, n] = ate_b.mean()
ate_lower = np.percentile(ate_bootstraps,
(self.ate_alpha / 2) * 100,
axis=1)
ate_upper = np.percentile(ate_bootstraps,
(1 - self.ate_alpha / 2) * 100,
axis=1)
# set member variables back to global (currently last bootstrapped outcome)
self.t_groups = t_groups_global
self._classes = _classes_global
self.models = deepcopy(models_global)
return ate, ate_lower, ate_upper
def estimate_ate(
self,
X,
assignment,
treatment,
y,
p=None,
pZ=None,
bootstrap_ci=False,
n_bootstraps=1000,
bootstrap_size=10000,
seed=None,
calibrate=True,
):
"""Estimate the Average Treatment Effect (ATE) for compliers.
Args:
X (np.matrix or np.array or pd.Dataframe): a feature matrix
assignment (np.array or pd.Series): an 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.
bootstrap_ci (bool): whether run bootstrap for confidence intervals
n_bootstraps (int): number of bootstrap iterations
bootstrap_size (int): number of samples per bootstrap
seed (int): random seed for cross-fitting
Returns:
The mean and confidence interval (LB, UB) of the ATE estimate.
"""
te, yhat_cs, yhat_ts = self.fit_predict(
X,
assignment,
treatment,
y,
p,
return_components=True,
seed=seed,
calibrate=calibrate,
)
X, assignment, treatment, y = convert_pd_to_np(X, assignment, treatment, y)
if p is None:
p = (self.propensity_0, self.propensity_1)
else:
check_p_conditions(p[0], self.t_groups)
check_p_conditions(p[1], self.t_groups)
if isinstance(p[0], (np.ndarray, pd.Series)):
treatment_name = self.t_groups[0]
p = (
{treatment_name: convert_pd_to_np(p[0])},
{treatment_name: convert_pd_to_np(p[1])},
)
elif isinstance(p[0], dict):
p = (
{
treatment_name: convert_pd_to_np(_p)
for treatment_name, _p in p[0].items()
},
{
treatment_name: convert_pd_to_np(_p)
for treatment_name, _p in p[1].items()
},
)
ate = np.zeros(self.t_groups.shape[0])
ate_lb = np.zeros(self.t_groups.shape[0])
ate_ub = np.zeros(self.t_groups.shape[0])
for i, group in enumerate(self.t_groups):
_ate = te[:, i].mean()
mask = (treatment == group) | (treatment == self.control_name)
mask_1, mask_0 = mask & (assignment == 1), mask & (assignment == 0)
Gamma = (treatment[mask_1] == group).mean() - (
treatment[mask_0] == group
).mean()
y_filt_1, y_filt_0 = y[mask_1], y[mask_0]
yhat_0 = yhat_cs[group][mask_0]
yhat_1 = yhat_ts[group][mask_1]
treatment_filt_1, treatment_filt_0 = treatment[mask_1], treatment[mask_0]
prob_treatment_1, prob_treatment_0 = (
p[1][group][mask_1],
p[0][group][mask_0],
)
w = (assignment[mask]).mean()
part_1 = (
(y_filt_1 - yhat_1).var()
+ _ate**2 * (treatment_filt_1 - prob_treatment_1).var()
- 2
* _ate
* (y_filt_1 * treatment_filt_1 - yhat_1 * prob_treatment_1).mean()
)
part_0 = (
(y_filt_0 - yhat_0).var()
+ _ate**2 * (treatment_filt_0 - prob_treatment_0).var()
- 2
* _ate
* (y_filt_0 * treatment_filt_0 - yhat_0 * prob_treatment_0).mean()
)
part_2 = np.mean(
(
yhat_ts[group][mask]
- yhat_cs[group][mask]
- _ate * (p[1][group][mask] - p[0][group][mask])
)
** 2
)
# SE formula is based on the lower bound formula (9) from Frölich, Markus. 2006.
# "Nonparametric IV estimation of local average treatment effects wth covariates."
# Journal of Econometrics.
se = np.sqrt((part_1 / w + part_2 / (1 - w)) + part_2) / Gamma
_ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
_ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
ate[i] = _ate
ate_lb[i] = _ate_lb
ate_ub[i] = _ate_ub
if not bootstrap_ci:
return ate, ate_lb, ate_ub
else:
t_groups_global = self.t_groups
_classes_global = self._classes
models_mu_c_global = deepcopy(self.models_mu_c)
models_mu_t_global = deepcopy(self.models_mu_t)
models_tau_global = deepcopy(self.models_tau)
logger.info("Bootstrap Confidence Intervals for ATE")
ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
for n in tqdm(range(n_bootstraps)):
cate_b = self.bootstrap(
X, assignment, treatment, y, p, pZ, size=bootstrap_size, seed=seed
)
ate_bootstraps[:, n] = cate_b.mean()
ate_lower = np.percentile(
ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1
)
ate_upper = np.percentile(
ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
)
# set member variables back to global (currently last bootstrapped outcome)
self.t_groups = t_groups_global
self._classes = _classes_global
self.models_mu_c = deepcopy(models_mu_c_global)
self.models_mu_t = deepcopy(models_mu_t_global)
self.models_tau = deepcopy(models_tau_global)
return ate, ate_lower, ate_upper
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)
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)