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)