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