diff --git a/doubleml/double_ml.py b/doubleml/double_ml.py index 2ca29b00..518f07b6 100644 --- a/doubleml/double_ml.py +++ b/doubleml/double_ml.py @@ -16,34 +16,37 @@ from ._utils_resampling import DoubleMLResampling, DoubleMLClusterResampling from ._utils import _draw_weights, _rmse, _aggregate_coefs_and_ses, _var_est -from ._utils_checks import _check_in_zero_one, _check_integer, _check_float, _check_bool, _check_is_partition, \ - _check_all_smpls, _check_smpl_split, _check_smpl_split_tpl, _check_benchmarks +from ._utils_checks import ( + _check_in_zero_one, + _check_integer, + _check_float, + _check_bool, + _check_is_partition, + _check_all_smpls, + _check_smpl_split, + _check_smpl_split_tpl, + _check_benchmarks, +) from ._utils_plots import _sensitivity_contour_plot -_implemented_data_backends = ['DoubleMLData', 'DoubleMLClusterData'] +_implemented_data_backends = ["DoubleMLData", "DoubleMLClusterData"] class DoubleML(ABC): - """Double Machine Learning. - """ - - def __init__(self, - obj_dml_data, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting): + """Double Machine Learning.""" + + def __init__(self, obj_dml_data, n_folds, n_rep, score, dml_procedure, draw_sample_splitting, apply_cross_fitting): # check and pick up obj_dml_data if not isinstance(obj_dml_data, DoubleMLBaseData): - raise TypeError('The data must be of ' + ' or '.join(_implemented_data_backends) + ' type. ' - f'{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed.') + raise TypeError( + "The data must be of " + " or ".join(_implemented_data_backends) + " type. " + f"{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed." + ) self._is_cluster_data = False if isinstance(obj_dml_data, DoubleMLClusterData): if obj_dml_data.n_cluster_vars > 2: - raise NotImplementedError('Multi-way (n_ways > 2) clustering not yet implemented.') + raise NotImplementedError("Multi-way (n_ways > 2) clustering not yet implemented.") self._is_cluster_data = True self._dml_data = obj_dml_data @@ -63,39 +66,41 @@ def __init__(self, self._sensitivity_implemented = False self._sensitivity_elements = None self._sensitivity_params = None - + # initialize external predictions self._external_predictions_implemented = False # check resampling specifications if not isinstance(n_folds, int): - raise TypeError('The number of folds must be of int type. ' - f'{str(n_folds)} of type {str(type(n_folds))} was passed.') + raise TypeError( + "The number of folds must be of int type. " f"{str(n_folds)} of type {str(type(n_folds))} was passed." + ) if n_folds < 1: - raise ValueError('The number of folds must be positive. ' - f'{str(n_folds)} was passed.') + raise ValueError("The number of folds must be positive. " f"{str(n_folds)} was passed.") if not isinstance(n_rep, int): - raise TypeError('The number of repetitions for the sample splitting must be of int type. ' - f'{str(n_rep)} of type {str(type(n_rep))} was passed.') + raise TypeError( + "The number of repetitions for the sample splitting must be of int type. " + f"{str(n_rep)} of type {str(type(n_rep))} was passed." + ) if n_rep < 1: - raise ValueError('The number of repetitions for the sample splitting must be positive. ' - f'{str(n_rep)} was passed.') + raise ValueError( + "The number of repetitions for the sample splitting must be positive. " f"{str(n_rep)} was passed." + ) if not isinstance(apply_cross_fitting, bool): - raise TypeError('apply_cross_fitting must be True or False. ' - f'Got {str(apply_cross_fitting)}.') + raise TypeError("apply_cross_fitting must be True or False. " f"Got {str(apply_cross_fitting)}.") if not isinstance(draw_sample_splitting, bool): - raise TypeError('draw_sample_splitting must be True or False. ' - f'Got {str(draw_sample_splitting)}.') + raise TypeError("draw_sample_splitting must be True or False. " f"Got {str(draw_sample_splitting)}.") # set resampling specifications if self._is_cluster_data: if (n_folds == 1) or (not apply_cross_fitting): - raise NotImplementedError('No cross-fitting (`apply_cross_fitting = False`) ' - 'is not yet implemented with clustering.') + raise NotImplementedError( + "No cross-fitting (`apply_cross_fitting = False`) " "is not yet implemented with clustering." + ) self._n_folds_per_cluster = n_folds - self._n_folds = n_folds ** self._dml_data.n_cluster_vars + self._n_folds = n_folds**self._dml_data.n_cluster_vars else: self._n_folds = n_folds self._n_rep = n_rep @@ -104,21 +109,20 @@ def __init__(self, self._strata = None # check and set dml_procedure and score - if (not isinstance(dml_procedure, str)) | (dml_procedure not in ['dml1', 'dml2']): - raise ValueError('dml_procedure must be "dml1" or "dml2". ' - f'Got {str(dml_procedure)}.') + if (not isinstance(dml_procedure, str)) | (dml_procedure not in ["dml1", "dml2"]): + raise ValueError('dml_procedure must be "dml1" or "dml2". ' f"Got {str(dml_procedure)}.") self._dml_procedure = dml_procedure self._score = score if (self.n_folds == 1) & self.apply_cross_fitting: - warnings.warn('apply_cross_fitting is set to False. Cross-fitting is not supported for n_folds = 1.') + warnings.warn("apply_cross_fitting is set to False. Cross-fitting is not supported for n_folds = 1.") self._apply_cross_fitting = False if not self.apply_cross_fitting: - assert self.n_folds <= 2, 'Estimation without cross-fitting not supported for n_folds > 2.' - if self.dml_procedure == 'dml2': + assert self.n_folds <= 2, "Estimation without cross-fitting not supported for n_folds > 2." + if self.dml_procedure == "dml2": # redirect to dml1 which works out-of-the-box; dml_procedure is of no relevance without cross-fitting - self._dml_procedure = 'dml1' + self._dml_procedure = "dml1" # perform sample splitting self._smpls = None @@ -127,8 +131,21 @@ def __init__(self, self.draw_sample_splitting() # initialize arrays according to obj_dml_data and the resampling settings - self._psi, self._psi_deriv, self._psi_elements,\ +<<<<<<< HEAD + self._psi, self._psi_deriv, self._psi_elements, \ self._coef, self._se, self._all_coef, self._all_se, self._all_dml1_coef = self._initialize_arrays() +======= + ( + self._psi, + self._psi_deriv, + self._psi_elements, + self._coef, + self._se, + self._all_coef, + self._all_se, + self._all_dml1_coef, + ) = self._initialize_arrays() +>>>>>>> 5d59ac29a02b6034a9e550709069398ea177a30d # also initialize bootstrap arrays with the default number of bootstrap replications self._n_rep_boot, self._boot_coef, self._boot_t_stat = self._initialize_boot_arrays(n_rep_boot=500) @@ -140,34 +157,44 @@ def __init__(self, def __str__(self): class_name = self.__class__.__name__ - header = f'================== {class_name} Object ==================\n' + header = f"================== {class_name} Object ==================\n" data_summary = self._dml_data._data_summary_str() - score_info = f'Score function: {str(self.score)}\n' \ - f'DML algorithm: {self.dml_procedure}\n' - learner_info = '' + score_info = f"Score function: {str(self.score)}\n" f"DML algorithm: {self.dml_procedure}\n" + learner_info = "" for key, value in self.learner.items(): - learner_info += f'Learner {key}: {str(value)}\n' + learner_info += f"Learner {key}: {str(value)}\n" if self.rmses is not None: - learner_info += 'Out-of-sample Performance:\n' + learner_info += "Out-of-sample Performance:\n" for learner in self.params_names: - learner_info += f'Learner {learner} RMSE: {self.rmses[learner]}\n' + learner_info += f"Learner {learner} RMSE: {self.rmses[learner]}\n" if self._is_cluster_data: - resampling_info = f'No. folds per cluster: {self._n_folds_per_cluster}\n' \ - f'No. folds: {self.n_folds}\n' \ - f'No. repeated sample splits: {self.n_rep}\n' \ - f'Apply cross-fitting: {self.apply_cross_fitting}\n' + resampling_info = ( + f"No. folds per cluster: {self._n_folds_per_cluster}\n" + f"No. folds: {self.n_folds}\n" + f"No. repeated sample splits: {self.n_rep}\n" + f"Apply cross-fitting: {self.apply_cross_fitting}\n" + ) else: - resampling_info = f'No. folds: {self.n_folds}\n' \ - f'No. repeated sample splits: {self.n_rep}\n' \ - f'Apply cross-fitting: {self.apply_cross_fitting}\n' + resampling_info = ( + f"No. folds: {self.n_folds}\n" + f"No. repeated sample splits: {self.n_rep}\n" + f"Apply cross-fitting: {self.apply_cross_fitting}\n" + ) fit_summary = str(self.summary) - res = header + \ - '\n------------------ Data summary ------------------\n' + data_summary + \ - '\n------------------ Score & algorithm ------------------\n' + score_info + \ - '\n------------------ Machine learner ------------------\n' + learner_info + \ - '\n------------------ Resampling ------------------\n' + resampling_info + \ - '\n------------------ Fit summary ------------------\n' + fit_summary + res = ( + header + + "\n------------------ Data summary ------------------\n" + + data_summary + + "\n------------------ Score & algorithm ------------------\n" + + score_info + + "\n------------------ Machine learner ------------------\n" + + learner_info + + "\n------------------ Resampling ------------------\n" + + resampling_info + + "\n------------------ Fit summary ------------------\n" + + fit_summary + ) return res @property @@ -291,8 +318,14 @@ def get_params(self, learner): """ valid_learner = self.params_names if (not isinstance(learner, str)) | (learner not in valid_learner): - raise ValueError('Invalid nuisance learner ' + str(learner) + '. ' + - 'Valid nuisance learner ' + ' or '.join(valid_learner) + '.') + raise ValueError( + "Invalid nuisance learner " + + str(learner) + + ". " + + "Valid nuisance learner " + + " or ".join(valid_learner) + + "." + ) return self._params[learner] # The private function _get_params delivers the single treatment, single (cross-fitting) sample subselection. @@ -309,10 +342,12 @@ def smpls(self): """ if self._smpls is None: if self._is_cluster_data: - err_msg = 'Sample splitting not specified. Draw samples via .draw_sample splitting().' + err_msg = "Sample splitting not specified. Draw samples via .draw_sample splitting()." else: - err_msg = ('Sample splitting not specified. Either draw samples via .draw_sample splitting() ' + - 'or set external samples via .set_sample_splitting().') + err_msg = ( + "Sample splitting not specified. Either draw samples via .draw_sample splitting() " + + "or set external samples via .set_sample_splitting()." + ) raise ValueError(err_msg) return self._smpls @@ -323,7 +358,7 @@ def smpls_cluster(self): """ if self._is_cluster_data: if self._smpls_cluster is None: - raise ValueError('Sample splitting not specified. Draw samples via .draw_sample splitting().') + raise ValueError("Sample splitting not specified. Draw samples via .draw_sample splitting().") return self._smpls_cluster @property @@ -451,16 +486,12 @@ def summary(self): """ A summary for the estimated causal effect after calling :meth:`fit`. """ - col_names = ['coef', 'std err', 't', 'P>|t|'] + col_names = ["coef", "std err", "t", "P>|t|"] if np.isnan(self.coef).all(): df_summary = pd.DataFrame(columns=col_names) else: - summary_stats = np.transpose(np.vstack( - [self.coef, self.se, - self.t_stat, self.pval])) - df_summary = pd.DataFrame(summary_stats, - columns=col_names, - index=self._dml_data.d_cols) + summary_stats = np.transpose(np.vstack([self.coef, self.se, self.t_stat, self.pval])) + df_summary = pd.DataFrame(summary_stats, columns=col_names, index=self._dml_data.d_cols) ci = self.confint() df_summary = df_summary.join(ci) return df_summary @@ -522,16 +553,16 @@ def fit(self, n_jobs_cv=None, store_predictions=True, external_predictions=None, if n_jobs_cv is not None: if not isinstance(n_jobs_cv, int): - raise TypeError('The number of CPUs used to fit the learners must be of int type. ' - f'{str(n_jobs_cv)} of type {str(type(n_jobs_cv))} was passed.') + raise TypeError( + "The number of CPUs used to fit the learners must be of int type. " + f"{str(n_jobs_cv)} of type {str(type(n_jobs_cv))} was passed." + ) if not isinstance(store_predictions, bool): - raise TypeError('store_predictions must be True or False. ' - f'Got {str(store_predictions)}.') + raise TypeError("store_predictions must be True or False. " f"Got {str(store_predictions)}.") if not isinstance(store_models, bool): - raise TypeError('store_models must be True or False. ' - f'Got {str(store_models)}.') + raise TypeError("store_models must be True or False. " f"Got {str(store_models)}.") # check if external predictions are implemented if self._external_predictions_implemented: @@ -550,9 +581,9 @@ def fit(self, n_jobs_cv=None, store_predictions=True, external_predictions=None, self._initialize_models() if self._sensitivity_implemented: - self._sensitivity_elements = self._initialize_sensitivity_elements((self._dml_data.n_obs, - self.n_rep, - self._dml_data.n_coefs)) + self._sensitivity_elements = self._initialize_sensitivity_elements( + (self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs) + ) for i_rep in range(self.n_rep): self._i_rep = i_rep @@ -577,34 +608,35 @@ def fit(self, n_jobs_cv=None, store_predictions=True, external_predictions=None, ext_prediction_dict[learner] = None # ml estimation of nuisance models and computation of score elements - score_elements, preds = self._nuisance_est(self.__smpls, n_jobs_cv, - external_predictions=ext_prediction_dict, - return_models=store_models) + score_elements, preds = self._nuisance_est( + self.__smpls, n_jobs_cv, external_predictions=ext_prediction_dict, return_models=store_models + ) self._set_score_elements(score_elements, self._i_rep, self._i_treat) # calculate rmses and store predictions and targets of the nuisance models - self._calc_rmses(preds['predictions'], preds['targets']) + self._calc_rmses(preds["predictions"], preds["targets"]) if store_predictions: - self._store_predictions_and_targets(preds['predictions'], preds['targets']) + self._store_predictions_and_targets(preds["predictions"], preds["targets"]) if store_models: - self._store_models(preds['models']) + self._store_models(preds["models"]) # estimate the causal parameter - self._all_coef[self._i_treat, self._i_rep], dml1_coefs = \ - self._est_causal_pars(self._get_score_elements(self._i_rep, self._i_treat)) - if self.dml_procedure == 'dml1': + self._all_coef[self._i_treat, self._i_rep], dml1_coefs = self._est_causal_pars( + self._get_score_elements(self._i_rep, self._i_treat) + ) + if self.dml_procedure == "dml1": self._all_dml1_coef[self._i_treat, self._i_rep, :] = dml1_coefs # compute score (depends on the estimated causal parameter) self._psi[:, self._i_rep, self._i_treat] = self._compute_score( - self._get_score_elements(self._i_rep, self._i_treat), - self._all_coef[self._i_treat, self._i_rep]) + self._get_score_elements(self._i_rep, self._i_treat), self._all_coef[self._i_treat, self._i_rep] + ) # compute score (can depend on the estimated causal parameter) self._psi_deriv[:, self._i_rep, self._i_treat] = self._compute_score_deriv( - self._get_score_elements(self._i_rep, self._i_treat), - self._all_coef[self._i_treat, self._i_rep]) + self._get_score_elements(self._i_rep, self._i_treat), self._all_coef[self._i_treat, self._i_rep] + ) # compute standard errors for causal parameter self._all_se[self._i_treat, self._i_rep] = self._se_causal_pars() @@ -615,7 +647,7 @@ def fit(self, n_jobs_cv=None, store_predictions=True, external_predictions=None, pass # check if callable score if callable(self.score): - warnings.warn('Sensitivity analysis not implemented for callable scores.') + warnings.warn("Sensitivity analysis not implemented for callable scores.") else: # compute sensitivity analysis elements element_dict = self._sensitivity_element_est(preds) @@ -626,7 +658,7 @@ def fit(self, n_jobs_cv=None, store_predictions=True, external_predictions=None, return self - def bootstrap(self, method='normal', n_rep_boot=500): + def bootstrap(self, method="normal", n_rep_boot=500): """ Multiplier bootstrap for DoubleML models. @@ -644,20 +676,20 @@ def bootstrap(self, method='normal', n_rep_boot=500): self : object """ if np.isnan(self.coef).all(): - raise ValueError('Apply fit() before bootstrap().') + raise ValueError("Apply fit() before bootstrap().") - if (not isinstance(method, str)) | (method not in ['Bayes', 'normal', 'wild']): - raise ValueError('Method must be "Bayes", "normal" or "wild". ' - f'Got {str(method)}.') + if (not isinstance(method, str)) | (method not in ["Bayes", "normal", "wild"]): + raise ValueError('Method must be "Bayes", "normal" or "wild". ' f"Got {str(method)}.") if not isinstance(n_rep_boot, int): - raise TypeError('The number of bootstrap replications must be of int type. ' - f'{str(n_rep_boot)} of type {str(type(n_rep_boot))} was passed.') + raise TypeError( + "The number of bootstrap replications must be of int type. " + f"{str(n_rep_boot)} of type {str(type(n_rep_boot))} was passed." + ) if n_rep_boot < 1: - raise ValueError('The number of bootstrap replications must be positive. ' - f'{str(n_rep_boot)} was passed.') + raise ValueError("The number of bootstrap replications must be positive. " f"{str(n_rep_boot)} was passed.") if self._is_cluster_data: - raise NotImplementedError('bootstrap not yet implemented with clustering.') + raise NotImplementedError("bootstrap not yet implemented with clustering.") self._n_rep_boot, self._boot_coef, self._boot_t_stat = self._initialize_boot_arrays(n_rep_boot) @@ -678,8 +710,10 @@ def bootstrap(self, method='normal', n_rep_boot=500): self._i_treat = i_d i_start = self._i_rep * self.n_rep_boot i_end = (self._i_rep + 1) * self.n_rep_boot - self._boot_coef[self._i_treat, i_start:i_end], self._boot_t_stat[self._i_treat, i_start:i_end] =\ - self._compute_bootstrap(weights) + ( + self._boot_coef[self._i_treat, i_start:i_end], + self._boot_t_stat[self._i_treat, i_start:i_end], + ) = self._compute_bootstrap(weights) self._boot_method = method return self @@ -705,36 +739,33 @@ def confint(self, joint=False, level=0.95): """ if not isinstance(joint, bool): - raise TypeError('joint must be True or False. ' - f'Got {str(joint)}.') + raise TypeError("joint must be True or False. " f"Got {str(joint)}.") if not isinstance(level, float): - raise TypeError('The confidence level must be of float type. ' - f'{str(level)} of type {str(type(level))} was passed.') + raise TypeError( + "The confidence level must be of float type. " f"{str(level)} of type {str(type(level))} was passed." + ) if (level <= 0) | (level >= 1): - raise ValueError('The confidence level must be in (0,1). ' - f'{str(level)} was passed.') + raise ValueError("The confidence level must be in (0,1). " f"{str(level)} was passed.") - a = (1 - level) - ab = np.array([a / 2, 1. - a / 2]) + a = 1 - level + ab = np.array([a / 2, 1.0 - a / 2]) if joint: if np.isnan(self.boot_coef).all(): - raise ValueError('Apply fit() & bootstrap() before confint(joint=True).') + raise ValueError("Apply fit() & bootstrap() before confint(joint=True).") sim = np.amax(np.abs(self.boot_t_stat), 0) hatc = np.quantile(sim, 1 - a) ci = np.vstack((self.coef - self.se * hatc, self.coef + self.se * hatc)).T else: if np.isnan(self.coef).all(): - raise ValueError('Apply fit() before confint().') + raise ValueError("Apply fit() before confint().") fac = norm.ppf(ab) ci = np.vstack((self.coef + self.se * fac[0], self.coef + self.se * fac[1])).T - df_ci = pd.DataFrame(ci, - columns=['{:.1f} %'.format(i * 100) for i in ab], - index=self._dml_data.d_cols) + df_ci = pd.DataFrame(ci, columns=["{:.1f} %".format(i * 100) for i in ab], index=self._dml_data.d_cols) return df_ci - def p_adjust(self, method='romano-wolf'): + def p_adjust(self, method="romano-wolf"): """ Multiple testing adjustment for DoubleML models. @@ -752,13 +783,14 @@ def p_adjust(self, method='romano-wolf'): A data frame with adjusted p-values. """ if np.isnan(self.coef).all(): - raise ValueError('Apply fit() before p_adjust().') + raise ValueError("Apply fit() before p_adjust().") if not isinstance(method, str): - raise TypeError('The p_adjust method must be of str type. ' - f'{str(method)} of type {str(type(method))} was passed.') + raise TypeError( + "The p_adjust method must be of str type. " f"{str(method)} of type {str(type(method))} was passed." + ) - if method.lower() in ['rw', 'romano-wolf']: + if method.lower() in ["rw", "romano-wolf"]: if np.isnan(self.boot_coef).all(): raise ValueError(f'Apply fit() & bootstrap() before p_adjust("{method}").') @@ -775,8 +807,7 @@ def p_adjust(self, method='romano-wolf'): sim = np.max(boot_t_stats, axis=0) pinit[i_d] = np.minimum(1, np.mean(sim >= np.abs(t_stat[stepdown_ind][i_d]))) else: - sim = np.max(np.delete(boot_t_stats, stepdown_ind[:i_d], axis=0), - axis=0) + sim = np.max(np.delete(boot_t_stats, stepdown_ind[:i_d], axis=0), axis=0) pinit[i_d] = np.minimum(1, np.mean(sim >= np.abs(t_stat[stepdown_ind][i_d]))) for i_d in range(self._dml_data.n_treat): @@ -789,22 +820,22 @@ def p_adjust(self, method='romano-wolf'): else: _, p_val, _, _ = multipletests(self.pval, method=method) - p_val = pd.DataFrame(np.vstack((self.coef, p_val)).T, - columns=['coef', 'pval'], - index=self._dml_data.d_cols) + p_val = pd.DataFrame(np.vstack((self.coef, p_val)).T, columns=["coef", "pval"], index=self._dml_data.d_cols) return p_val - def tune(self, - param_grids, - tune_on_folds=False, - scoring_methods=None, # if None the estimator's score method is used - n_folds_tune=5, - search_mode='grid_search', - n_iter_randomized_search=100, - n_jobs_cv=None, - set_as_params=True, - return_tune_res=False): + def tune( + self, + param_grids, + tune_on_folds=False, + scoring_methods=None, # if None the estimator's score method is used + n_folds_tune=5, + search_mode="grid_search", + n_iter_randomized_search=100, + n_jobs_cv=None, + set_as_params=True, + return_tune_res=False, + ): """ Hyperparameter-tuning for DoubleML models. @@ -863,14 +894,22 @@ def tune(self, """ if (not isinstance(param_grids, dict)) | (not all(k in param_grids for k in self.learner_names)): - raise ValueError('Invalid param_grids ' + str(param_grids) + '. ' - 'param_grids must be a dictionary with keys ' + ' and '.join(self.learner_names) + '.') + raise ValueError( + "Invalid param_grids " + str(param_grids) + ". " + "param_grids must be a dictionary with keys " + " and ".join(self.learner_names) + "." + ) if scoring_methods is not None: if (not isinstance(scoring_methods, dict)) | (not all(k in self.learner_names for k in scoring_methods)): - raise ValueError('Invalid scoring_methods ' + str(scoring_methods) + '. ' + - 'scoring_methods must be a dictionary. ' + - 'Valid keys are ' + ' and '.join(self.learner_names) + '.') + raise ValueError( + "Invalid scoring_methods " + + str(scoring_methods) + + ". " + + "scoring_methods must be a dictionary. " + + "Valid keys are " + + " and ".join(self.learner_names) + + "." + ) if not all(k in scoring_methods for k in self.learner_names): # if there are learners for which no scoring_method was set, we fall back to None, i.e., default scoring for learner in self.learner_names: @@ -878,40 +917,43 @@ def tune(self, scoring_methods[learner] = None if not isinstance(tune_on_folds, bool): - raise TypeError('tune_on_folds must be True or False. ' - f'Got {str(tune_on_folds)}.') + raise TypeError("tune_on_folds must be True or False. " f"Got {str(tune_on_folds)}.") if not isinstance(n_folds_tune, int): - raise TypeError('The number of folds used for tuning must be of int type. ' - f'{str(n_folds_tune)} of type {str(type(n_folds_tune))} was passed.') + raise TypeError( + "The number of folds used for tuning must be of int type. " + f"{str(n_folds_tune)} of type {str(type(n_folds_tune))} was passed." + ) if n_folds_tune < 2: - raise ValueError('The number of folds used for tuning must be at least two. ' - f'{str(n_folds_tune)} was passed.') + raise ValueError("The number of folds used for tuning must be at least two. " f"{str(n_folds_tune)} was passed.") - if (not isinstance(search_mode, str)) | (search_mode not in ['grid_search', 'randomized_search']): - raise ValueError('search_mode must be "grid_search" or "randomized_search". ' - f'Got {str(search_mode)}.') + if (not isinstance(search_mode, str)) | (search_mode not in ["grid_search", "randomized_search"]): + raise ValueError('search_mode must be "grid_search" or "randomized_search". ' f"Got {str(search_mode)}.") if not isinstance(n_iter_randomized_search, int): - raise TypeError('The number of parameter settings sampled for the randomized search must be of int type. ' - f'{str(n_iter_randomized_search)} of type ' - f'{str(type(n_iter_randomized_search))} was passed.') + raise TypeError( + "The number of parameter settings sampled for the randomized search must be of int type. " + f"{str(n_iter_randomized_search)} of type " + f"{str(type(n_iter_randomized_search))} was passed." + ) if n_iter_randomized_search < 2: - raise ValueError('The number of parameter settings sampled for the randomized search must be at least two. ' - f'{str(n_iter_randomized_search)} was passed.') + raise ValueError( + "The number of parameter settings sampled for the randomized search must be at least two. " + f"{str(n_iter_randomized_search)} was passed." + ) if n_jobs_cv is not None: if not isinstance(n_jobs_cv, int): - raise TypeError('The number of CPUs used to fit the learners must be of int type. ' - f'{str(n_jobs_cv)} of type {str(type(n_jobs_cv))} was passed.') + raise TypeError( + "The number of CPUs used to fit the learners must be of int type. " + f"{str(n_jobs_cv)} of type {str(type(n_jobs_cv))} was passed." + ) if not isinstance(set_as_params, bool): - raise TypeError('set_as_params must be True or False. ' - f'Got {str(set_as_params)}.') + raise TypeError("set_as_params must be True or False. " f"Got {str(set_as_params)}.") if not isinstance(return_tune_res, bool): - raise TypeError('return_tune_res must be True or False. ' - f'Got {str(return_tune_res)}.') + raise TypeError("return_tune_res must be True or False. " f"Got {str(return_tune_res)}.") if tune_on_folds: tuning_res = [[None] * self.n_rep] * self._dml_data.n_treat @@ -930,14 +972,18 @@ def tune(self, self._i_rep = i_rep # tune hyperparameters - res = self._nuisance_tuning(self.__smpls, - param_grids, scoring_methods, - n_folds_tune, - n_jobs_cv, - search_mode, n_iter_randomized_search) + res = self._nuisance_tuning( + self.__smpls, + param_grids, + scoring_methods, + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) tuning_res[i_rep][i_d] = res - nuisance_params.append(res['params']) + nuisance_params.append(res["params"]) if set_as_params: for nuisance_model in nuisance_params[0].keys(): @@ -947,16 +993,14 @@ def tune(self, else: smpls = [(np.arange(self._dml_data.n_obs), np.arange(self._dml_data.n_obs))] # tune hyperparameters - res = self._nuisance_tuning(smpls, - param_grids, scoring_methods, - n_folds_tune, - n_jobs_cv, - search_mode, n_iter_randomized_search) + res = self._nuisance_tuning( + smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ) tuning_res[i_d] = res if set_as_params: - for nuisance_model in res['params'].keys(): - params = res['params'][nuisance_model] + for nuisance_model in res["params"].keys(): + params = res["params"][nuisance_model] self.set_ml_nuisance_params(nuisance_model, self._dml_data.d_cols[i_d], params[0]) if return_tune_res: @@ -986,12 +1030,19 @@ def set_ml_nuisance_params(self, learner, treat_var, params): """ valid_learner = self.params_names if learner not in valid_learner: - raise ValueError('Invalid nuisance learner ' + learner + '. ' + - 'Valid nuisance learner ' + ' or '.join(valid_learner) + '.') + raise ValueError( + "Invalid nuisance learner " + learner + ". " + "Valid nuisance learner " + " or ".join(valid_learner) + "." + ) if treat_var not in self._dml_data.d_cols: - raise ValueError('Invalid treatment variable ' + treat_var + '. ' + - 'Valid treatment variable ' + ' or '.join(self._dml_data.d_cols) + '.') + raise ValueError( + "Invalid treatment variable " + + treat_var + + ". " + + "Valid treatment variable " + + " or ".join(self._dml_data.d_cols) + + "." + ) if params is None: all_params = [None] * self.n_rep @@ -1022,24 +1073,25 @@ def _nuisance_est(self, smpls, n_jobs_cv, return_models, external_predictions): pass @abstractmethod - def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): + def _nuisance_tuning( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): pass @staticmethod def _check_learner(learner, learner_name, regressor, classifier): - err_msg_prefix = f'Invalid learner provided for {learner_name}: ' - warn_msg_prefix = f'Learner provided for {learner_name} is probably invalid: ' + err_msg_prefix = f"Invalid learner provided for {learner_name}: " + warn_msg_prefix = f"Learner provided for {learner_name} is probably invalid: " if isinstance(learner, type): - raise TypeError(err_msg_prefix + 'provide an instance of a learner instead of a class.') + raise TypeError(err_msg_prefix + "provide an instance of a learner instead of a class.") - if not hasattr(learner, 'fit'): - raise TypeError(err_msg_prefix + f'{str(learner)} has no method .fit().') - if not hasattr(learner, 'set_params'): - raise TypeError(err_msg_prefix + f'{str(learner)} has no method .set_params().') - if not hasattr(learner, 'get_params'): - raise TypeError(err_msg_prefix + f'{str(learner)} has no method .get_params().') + if not hasattr(learner, "fit"): + raise TypeError(err_msg_prefix + f"{str(learner)} has no method .fit().") + if not hasattr(learner, "set_params"): + raise TypeError(err_msg_prefix + f"{str(learner)} has no method .set_params().") + if not hasattr(learner, "get_params"): + raise TypeError(err_msg_prefix + f"{str(learner)} has no method .get_params().") if regressor & classifier: if is_classifier(learner): @@ -1047,70 +1099,91 @@ def _check_learner(learner, learner_name, regressor, classifier): elif is_regressor(learner): learner_is_classifier = False else: - warnings.warn(warn_msg_prefix + f'{str(learner)} is (probably) neither a regressor nor a classifier. ' + - 'Method predict is used for prediction.') + warnings.warn( + warn_msg_prefix + + f"{str(learner)} is (probably) neither a regressor nor a classifier. " + + "Method predict is used for prediction." + ) learner_is_classifier = False elif classifier: if not is_classifier(learner): - warnings.warn(warn_msg_prefix + f'{str(learner)} is (probably) no classifier.') + warnings.warn(warn_msg_prefix + f"{str(learner)} is (probably) no classifier.") learner_is_classifier = True else: assert regressor # classifier, regressor or both must be True if not is_regressor(learner): - warnings.warn(warn_msg_prefix + f'{str(learner)} is (probably) no regressor.') + warnings.warn(warn_msg_prefix + f"{str(learner)} is (probably) no regressor.") learner_is_classifier = False # check existence of the prediction method if learner_is_classifier: - if not hasattr(learner, 'predict_proba'): - raise TypeError(err_msg_prefix + f'{str(learner)} has no method .predict_proba().') + if not hasattr(learner, "predict_proba"): + raise TypeError(err_msg_prefix + f"{str(learner)} has no method .predict_proba().") else: - if not hasattr(learner, 'predict'): - raise TypeError(err_msg_prefix + f'{str(learner)} has no method .predict().') + if not hasattr(learner, "predict"): + raise TypeError(err_msg_prefix + f"{str(learner)} has no method .predict().") return learner_is_classifier def _check_external_predictions(self, external_predictions): if external_predictions is not None: if not isinstance(external_predictions, dict): - raise TypeError('external_predictions must be a dictionary. ' - f'{str(external_predictions)} of type {str(type(external_predictions))} was passed.') + raise TypeError( + "external_predictions must be a dictionary. " + f"{str(external_predictions)} of type {str(type(external_predictions))} was passed." + ) supplied_treatments = list(external_predictions.keys()) valid_treatments = self._dml_data.d_cols if not set(supplied_treatments).issubset(valid_treatments): - raise ValueError('Invalid external_predictions. ' - f'Invalid treatment variable in {str(supplied_treatments)}. ' - 'Valid treatment variables ' + ' or '.join(valid_treatments) + '.') + raise ValueError( + "Invalid external_predictions. " + f"Invalid treatment variable in {str(supplied_treatments)}. " + "Valid treatment variables " + " or ".join(valid_treatments) + "." + ) for treatment in supplied_treatments: if not isinstance(external_predictions[treatment], dict): - raise TypeError('external_predictions must be a nested dictionary. ' - f'For treatment {str(treatment)} a value of type ' - f'{str(type(external_predictions[treatment]))} was passed.') + raise TypeError( + "external_predictions must be a nested dictionary. " + f"For treatment {str(treatment)} a value of type " + f"{str(type(external_predictions[treatment]))} was passed." + ) supplied_learners = list(external_predictions[treatment].keys()) valid_learners = self.params_names if not set(supplied_learners).issubset(valid_learners): - raise ValueError('Invalid external_predictions. ' - f'Invalid nuisance learner for treatment {str(treatment)} in {str(supplied_learners)}. ' - 'Valid nuisance learners ' + ' or '.join(valid_learners) + '.') + raise ValueError( + "Invalid external_predictions. " + f"Invalid nuisance learner for treatment {str(treatment)} in {str(supplied_learners)}. " + "Valid nuisance learners " + " or ".join(valid_learners) + "." + ) for learner in supplied_learners: if not isinstance(external_predictions[treatment][learner], np.ndarray): - raise TypeError('Invalid external_predictions. ' - 'The values of the nested list must be a numpy array. ' - 'Invalid predictions for treatment ' + str(treatment) + - ' and learner ' + str(learner) + '. ' + - f'Object of type {str(type(external_predictions[treatment][learner]))} was passed.') + raise TypeError( + "Invalid external_predictions. " + "The values of the nested list must be a numpy array. " + "Invalid predictions for treatment " + + str(treatment) + + " and learner " + + str(learner) + + ". " + + f"Object of type {str(type(external_predictions[treatment][learner]))} was passed." + ) expected_shape = (self._dml_data.n_obs, self.n_rep) if external_predictions[treatment][learner].shape != expected_shape: - raise ValueError('Invalid external_predictions. ' - f'The supplied predictions have to be of shape {str(expected_shape)}. ' - 'Invalid predictions for treatment ' + str(treatment) + - ' and learner ' + str(learner) + '. ' + - f'Predictions of shape {str(external_predictions[treatment][learner].shape)} passed.') + raise ValueError( + "Invalid external_predictions. " + f"The supplied predictions have to be of shape {str(expected_shape)}. " + "Invalid predictions for treatment " + + str(treatment) + + " and learner " + + str(learner) + + ". " + + f"Predictions of shape {str(external_predictions[treatment][learner].shape)} passed." + ) def _initialize_arrays(self): # scores @@ -1125,7 +1198,7 @@ def _initialize_arrays(self): all_coef = np.full((self._dml_data.n_coefs, self.n_rep), np.nan) all_se = np.full((self._dml_data.n_coefs, self.n_rep), np.nan) - if self.dml_procedure == 'dml1': + if self.dml_procedure == "dml1": if self.apply_cross_fitting: all_dml1_coef = np.full((self._dml_data.n_coefs, self.n_rep, self.n_folds), np.nan) else: @@ -1141,18 +1214,22 @@ def _initialize_boot_arrays(self, n_rep_boot): return n_rep_boot, boot_coef, boot_t_stat def _initialize_predictions_and_targets(self): - self._predictions = {learner: np.full((self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs), np.nan) - for learner in self.params_names} - self._nuisance_targets = {learner: np.full((self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs), np.nan) - for learner in self.params_names} + self._predictions = { + learner: np.full((self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs), np.nan) + for learner in self.params_names + } + self._nuisance_targets = { + learner: np.full((self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs), np.nan) + for learner in self.params_names + } def _initialize_rmses(self): - self._rmses = {learner: np.full((self.n_rep, self._dml_data.n_coefs), np.nan) - for learner in self.params_names} + self._rmses = {learner: np.full((self.n_rep, self._dml_data.n_coefs), np.nan) for learner in self.params_names} def _initialize_models(self): - self._models = {learner: {treat_var: [None] * self.n_rep for treat_var in self._dml_data.d_cols} - for learner in self.params_names} + self._models = { + learner: {treat_var: [None] * self.n_rep for treat_var in self._dml_data.d_cols} for learner in self.params_names + } def _store_predictions_and_targets(self, preds, targets): for learner in self.params_names: @@ -1220,27 +1297,28 @@ def evaluate_learners(self, learners=None, metric=_rmse): # check metric if not callable(metric): - raise TypeError('metric should be a callable. ' - '%r was passed.' % metric) + raise TypeError("metric should be a callable. " "%r was passed." % metric) if all(learner in self.params_names for learner in learners): if self.nuisance_targets is None: - raise ValueError('Apply fit() before evaluate_learners().') + raise ValueError("Apply fit() before evaluate_learners().") else: - dist = {learner: np.full((self.n_rep, self._dml_data.n_coefs), np.nan) - for learner in learners} + dist = {learner: np.full((self.n_rep, self._dml_data.n_coefs), np.nan) for learner in learners} for learner in learners: for rep in range(self.n_rep): for coef_idx in range(self._dml_data.n_coefs): - res = metric(y_pred=self.predictions[learner][:, rep, coef_idx].reshape(1, -1), - y_true=self.nuisance_targets[learner][:, rep, coef_idx].reshape(1, -1)) + res = metric( + y_pred=self.predictions[learner][:, rep, coef_idx].reshape(1, -1), + y_true=self.nuisance_targets[learner][:, rep, coef_idx].reshape(1, -1), + ) if not np.isfinite(res): - raise ValueError(f'Evaluation from learner {str(learner)} is not finite.') + raise ValueError(f"Evaluation from learner {str(learner)} is not finite.") dist[learner][rep, coef_idx] = res return dist else: - raise ValueError(f'The learners have to be a subset of {str(self.params_names)}. ' - f'Learners {str(learners)} provided.') + raise ValueError( + f"The learners have to be a subset of {str(self.params_names)}. " f"Learners {str(learners)} provided." + ) def draw_sample_splitting(self): """ @@ -1254,19 +1332,23 @@ def draw_sample_splitting(self): self : object """ if self._is_cluster_data: - obj_dml_resampling = DoubleMLClusterResampling(n_folds=self._n_folds_per_cluster, - n_rep=self.n_rep, - n_obs=self._dml_data.n_obs, - apply_cross_fitting=self.apply_cross_fitting, - n_cluster_vars=self._dml_data.n_cluster_vars, - cluster_vars=self._dml_data.cluster_vars) + obj_dml_resampling = DoubleMLClusterResampling( + n_folds=self._n_folds_per_cluster, + n_rep=self.n_rep, + n_obs=self._dml_data.n_obs, + apply_cross_fitting=self.apply_cross_fitting, + n_cluster_vars=self._dml_data.n_cluster_vars, + cluster_vars=self._dml_data.cluster_vars, + ) self._smpls, self._smpls_cluster = obj_dml_resampling.split_samples() else: - obj_dml_resampling = DoubleMLResampling(n_folds=self.n_folds, - n_rep=self.n_rep, - n_obs=self._dml_data.n_obs, - apply_cross_fitting=self.apply_cross_fitting, - stratify=self._strata) + obj_dml_resampling = DoubleMLResampling( + n_folds=self.n_folds, + n_rep=self.n_rep, + n_obs=self._dml_data.n_obs, + apply_cross_fitting=self.apply_cross_fitting, + stratify=self._strata, + ) self._smpls = obj_dml_resampling.split_samples() return self @@ -1330,15 +1412,18 @@ def set_sample_splitting(self, all_smpls): >>> dml_plr_obj.set_sample_splitting(smpls) """ if self._is_cluster_data: - raise NotImplementedError('Externally setting the sample splitting for DoubleML is ' - 'not yet implemented with clustering.') + raise NotImplementedError( + "Externally setting the sample splitting for DoubleML is " "not yet implemented with clustering." + ) if isinstance(all_smpls, tuple): if not len(all_smpls) == 2: - raise ValueError('Invalid partition provided. ' - 'Tuple for train_ind and test_ind must consist of exactly two elements.') + raise ValueError( + "Invalid partition provided. " "Tuple for train_ind and test_ind must consist of exactly two elements." + ) all_smpls = _check_smpl_split_tpl(all_smpls, self._dml_data.n_obs) - if (_check_is_partition([all_smpls], self._dml_data.n_obs) & - _check_is_partition([(all_smpls[1], all_smpls[0])], self._dml_data.n_obs)): + if _check_is_partition([all_smpls], self._dml_data.n_obs) & _check_is_partition( + [(all_smpls[1], all_smpls[0])], self._dml_data.n_obs + ): self._n_rep = 1 self._n_folds = 1 self._apply_cross_fitting = False @@ -1350,18 +1435,20 @@ def set_sample_splitting(self, all_smpls): self._smpls = _check_all_smpls([[all_smpls]], self._dml_data.n_obs, check_intersect=True) else: if not isinstance(all_smpls, list): - raise TypeError('all_smpls must be of list or tuple type. ' - f'{str(all_smpls)} of type {str(type(all_smpls))} was passed.') + raise TypeError( + "all_smpls must be of list or tuple type. " f"{str(all_smpls)} of type {str(type(all_smpls))} was passed." + ) all_tuple = all([isinstance(tpl, tuple) for tpl in all_smpls]) if all_tuple: if not all([len(tpl) == 2 for tpl in all_smpls]): - raise ValueError('Invalid partition provided. ' - 'All tuples for train_ind and test_ind must consist of exactly two elements.') + raise ValueError( + "Invalid partition provided. " + "All tuples for train_ind and test_ind must consist of exactly two elements." + ) self._n_rep = 1 all_smpls = _check_smpl_split(all_smpls, self._dml_data.n_obs) if _check_is_partition(all_smpls, self._dml_data.n_obs): - if ((len(all_smpls) == 1) & - _check_is_partition([(all_smpls[0][1], all_smpls[0][0])], self._dml_data.n_obs)): + if (len(all_smpls) == 1) & _check_is_partition([(all_smpls[0][1], all_smpls[0][0])], self._dml_data.n_obs): self._n_folds = 1 self._apply_cross_fitting = False self._smpls = [all_smpls] @@ -1371,34 +1458,42 @@ def set_sample_splitting(self, all_smpls): self._smpls = _check_all_smpls([all_smpls], self._dml_data.n_obs, check_intersect=True) else: if not len(all_smpls) == 1: - raise ValueError('Invalid partition provided. ' - 'Tuples for more than one fold provided that don\'t form a partition.') + raise ValueError( + "Invalid partition provided. " + "Tuples for more than one fold provided that don't form a partition." + ) self._n_folds = 2 self._apply_cross_fitting = False self._smpls = _check_all_smpls([all_smpls], self._dml_data.n_obs, check_intersect=True) else: all_list = all([isinstance(smpl, list) for smpl in all_smpls]) if not all_list: - raise ValueError('Invalid partition provided. ' - 'all_smpls is a list where neither all elements are tuples ' - 'nor all elements are lists.') + raise ValueError( + "Invalid partition provided. " + "all_smpls is a list where neither all elements are tuples " + "nor all elements are lists." + ) all_tuple = all([all([isinstance(tpl, tuple) for tpl in smpl]) for smpl in all_smpls]) if not all_tuple: - raise TypeError('For repeated sample splitting all_smpls must be list of lists of tuples.') + raise TypeError("For repeated sample splitting all_smpls must be list of lists of tuples.") all_pairs = all([all([len(tpl) == 2 for tpl in smpl]) for smpl in all_smpls]) if not all_pairs: - raise ValueError('Invalid partition provided. ' - 'All tuples for train_ind and test_ind must consist of exactly two elements.') + raise ValueError( + "Invalid partition provided. " + "All tuples for train_ind and test_ind must consist of exactly two elements." + ) n_folds_each_smpl = np.array([len(smpl) for smpl in all_smpls]) if not np.all(n_folds_each_smpl == n_folds_each_smpl[0]): - raise ValueError('Invalid partition provided. ' - 'Different number of folds for repeated sample splitting.') + raise ValueError("Invalid partition provided. " "Different number of folds for repeated sample splitting.") all_smpls = _check_all_smpls(all_smpls, self._dml_data.n_obs) smpls_are_partitions = [_check_is_partition(smpl, self._dml_data.n_obs) for smpl in all_smpls] if all(smpls_are_partitions): - if ((len(all_smpls) == 1) & (len(all_smpls[0]) == 1) & - _check_is_partition([(all_smpls[0][0][1], all_smpls[0][0][0])], self._dml_data.n_obs)): + if ( + (len(all_smpls) == 1) + & (len(all_smpls[0]) == 1) + & _check_is_partition([(all_smpls[0][0][1], all_smpls[0][0][0])], self._dml_data.n_obs) + ): self._n_rep = 1 self._n_folds = 1 self._apply_cross_fitting = False @@ -1410,16 +1505,26 @@ def set_sample_splitting(self, all_smpls): self._smpls = _check_all_smpls(all_smpls, self._dml_data.n_obs, check_intersect=True) else: if not n_folds_each_smpl[0] == 1: - raise ValueError('Invalid partition provided. ' - 'Tuples for more than one fold provided ' - 'but at least one does not form a partition.') + raise ValueError( + "Invalid partition provided. " + "Tuples for more than one fold provided " + "but at least one does not form a partition." + ) self._n_rep = len(all_smpls) self._n_folds = 2 self._apply_cross_fitting = False self._smpls = _check_all_smpls(all_smpls, self._dml_data.n_obs, check_intersect=True) - self._psi, self._psi_deriv, self._psi_elements, \ - self._coef, self._se, self._all_coef, self._all_se, self._all_dml1_coef = self._initialize_arrays() + ( + self._psi, + self._psi_deriv, + self._psi_elements, + self._coef, + self._se, + self._all_coef, + self._all_se, + self._all_dml1_coef, + ) = self._initialize_arrays() self._initialize_ml_nuisance_params() return self @@ -1428,22 +1533,22 @@ def _est_causal_pars(self, psi_elements): dml_procedure = self.dml_procedure smpls = self.__smpls - if dml_procedure == 'dml1': + if dml_procedure == "dml1": # Note that len(smpls) is only not equal to self.n_folds if self.apply_cross_fitting = False dml1_coefs = np.zeros(len(smpls)) for idx, (_, test_index) in enumerate(smpls): dml1_coefs[idx] = self._est_coef(psi_elements, inds=test_index) coef = np.mean(dml1_coefs) else: - assert dml_procedure == 'dml2' + assert dml_procedure == "dml2" dml1_coefs = None if not self._is_cluster_data: coef = self._est_coef(psi_elements) else: - scaling_factor = [1.] * len(smpls) + scaling_factor = [1.0] * len(smpls) for i_fold, (_, test_index) in enumerate(smpls): test_cluster_inds = self.__smpls_cluster[i_fold][1] - scaling_factor[i_fold] = 1./np.prod(np.array([len(inds) for inds in test_cluster_inds])) + scaling_factor[i_fold] = 1.0 / np.prod(np.array([len(inds) for inds in test_cluster_inds])) coef = self._est_coef(psi_elements, smpls=smpls, scaling_factor=scaling_factor) return coef, dml1_coefs @@ -1459,14 +1564,16 @@ def _se_causal_pars(self): smpls_cluster = self.__smpls_cluster n_folds_per_cluster = self._n_folds_per_cluster - sigma2_hat, var_scaling_factor = _var_est(psi=self.__psi, - psi_deriv=self.__psi_deriv, - apply_cross_fitting=self.apply_cross_fitting, - smpls=self.__smpls, - is_cluster_data=self._is_cluster_data, - cluster_vars=cluster_vars, - smpls_cluster=smpls_cluster, - n_folds_per_cluster=n_folds_per_cluster) + sigma2_hat, var_scaling_factor = _var_est( + psi=self.__psi, + psi_deriv=self.__psi_deriv, + apply_cross_fitting=self.apply_cross_fitting, + smpls=self.__smpls, + is_cluster_data=self._is_cluster_data, + cluster_vars=cluster_vars, + smpls_cluster=smpls_cluster, + n_folds_per_cluster=n_folds_per_cluster, + ) self._var_scaling_factor = var_scaling_factor se = np.sqrt(sigma2_hat) @@ -1480,20 +1587,21 @@ def _est_causal_pars_and_se(self): self._i_treat = i_d # estimate the causal parameter - self._all_coef[self._i_treat, self._i_rep], dml1_coefs = \ - self._est_causal_pars(self._get_score_elements(self._i_rep, self._i_treat)) - if self.dml_procedure == 'dml1': + self._all_coef[self._i_treat, self._i_rep], dml1_coefs = self._est_causal_pars( + self._get_score_elements(self._i_rep, self._i_treat) + ) + if self.dml_procedure == "dml1": self._all_dml1_coef[self._i_treat, self._i_rep, :] = dml1_coefs # compute score (depends on the estimated causal parameter) self._psi[:, self._i_rep, self._i_treat] = self._compute_score( - self._get_score_elements(self._i_rep, self._i_treat), - self._all_coef[self._i_treat, self._i_rep]) + self._get_score_elements(self._i_rep, self._i_treat), self._all_coef[self._i_treat, self._i_rep] + ) # compute score (can depend on the estimated causal parameter) self._psi_deriv[:, self._i_rep, self._i_treat] = self._compute_score_deriv( - self._get_score_elements(self._i_rep, self._i_treat), - self._all_coef[self._i_treat, self._i_rep]) + self._get_score_elements(self._i_rep, self._i_treat), self._all_coef[self._i_treat, self._i_rep] + ) # compute standard errors for causal parameter self._all_se[self._i_treat, self._i_rep] = self._se_causal_pars() @@ -1541,12 +1649,15 @@ def _get_score_elements(self, i_rep, i_treat): def _set_score_elements(self, psi_elements, i_rep, i_treat): if not isinstance(psi_elements, dict): - raise TypeError('_ml_nuisance_and_score_elements must return score elements in a dict. ' - f'Got type {str(type(psi_elements))}.') + raise TypeError( + "_ml_nuisance_and_score_elements must return score elements in a dict. " f"Got type {str(type(psi_elements))}." + ) if not (set(self._score_element_names) == set(psi_elements.keys())): - raise ValueError('_ml_nuisance_and_score_elements returned incomplete score elements. ' - 'Expected dict with keys: ' + ' and '.join(set(self._score_element_names)) + '.' - 'Got dict with keys: ' + ' and '.join(set(psi_elements.keys())) + '.') + raise ValueError( + "_ml_nuisance_and_score_elements returned incomplete score elements. " + "Expected dict with keys: " + " and ".join(set(self._score_element_names)) + "." + "Got dict with keys: " + " and ".join(set(psi_elements.keys())) + "." + ) for key in self._score_element_names: self.psi_elements[key][:, i_rep, i_treat] = psi_elements[key] return @@ -1562,14 +1673,16 @@ def _sensitivity_element_est(self, preds): @property def _sensitivity_element_names(self): - return ['sigma2', 'nu2', 'psi_sigma2', 'psi_nu2'] + return ["sigma2", "nu2", "psi_sigma2", "psi_nu2"] # the dimensions will usually be (n_obs, n_rep, n_coefs) to be equal to the score dimensions psi def _initialize_sensitivity_elements(self, score_dim): - sensitivity_elements = {'sigma2': np.full((1, score_dim[1], score_dim[2]), np.nan), - 'nu2': np.full((1, score_dim[1], score_dim[2]), np.nan), - 'psi_sigma2': np.full(score_dim, np.nan), - 'psi_nu2': np.full(score_dim, np.nan)} + sensitivity_elements = { + "sigma2": np.full((1, score_dim[1], score_dim[2]), np.nan), + "nu2": np.full((1, score_dim[1], score_dim[2]), np.nan), + "psi_sigma2": np.full(score_dim, np.nan), + "psi_nu2": np.full(score_dim, np.nan), + } return sensitivity_elements def _get_sensitivity_elements(self, i_rep, i_treat): @@ -1578,45 +1691,50 @@ def _get_sensitivity_elements(self, i_rep, i_treat): def _set_sensitivity_elements(self, sensitivity_elements, i_rep, i_treat): if not isinstance(sensitivity_elements, dict): - raise TypeError('_sensitivity_element_est must return sensitivity elements in a dict. ' - f'Got type {str(type(sensitivity_elements))}.') + raise TypeError( + "_sensitivity_element_est must return sensitivity elements in a dict. " + f"Got type {str(type(sensitivity_elements))}." + ) if not (set(self._sensitivity_element_names) == set(sensitivity_elements.keys())): - raise ValueError('_sensitivity_element_est returned incomplete sensitivity elements. ' - 'Expected dict with keys: ' + ' and '.join(set(self._sensitivity_element_names)) + '. ' - 'Got dict with keys: ' + ' and '.join(set(sensitivity_elements.keys())) + '.') + raise ValueError( + "_sensitivity_element_est returned incomplete sensitivity elements. " + "Expected dict with keys: " + " and ".join(set(self._sensitivity_element_names)) + ". " + "Got dict with keys: " + " and ".join(set(sensitivity_elements.keys())) + "." + ) for key in self._sensitivity_element_names: self.sensitivity_elements[key][:, i_rep, i_treat] = sensitivity_elements[key] return def _calc_sensitivity_analysis(self, cf_y, cf_d, rho, level): if not self.apply_cross_fitting: - raise NotImplementedError('Sensitivity analysis not yet implemented without cross-fitting.') + raise NotImplementedError("Sensitivity analysis not yet implemented without cross-fitting.") if self._sensitivity_elements is None: - raise NotImplementedError(f'Sensitivity analysis not yet implemented for {str(type(self))}.') + raise NotImplementedError(f"Sensitivity analysis not yet implemented for {str(type(self))}.") # checks - _check_in_zero_one(cf_y, 'cf_y', include_one=False) - _check_in_zero_one(cf_d, 'cf_d', include_one=False) + _check_in_zero_one(cf_y, "cf_y", include_one=False) + _check_in_zero_one(cf_d, "cf_d", include_one=False) if not isinstance(rho, float): - raise TypeError(f'rho must be of float type. ' - f'{str(rho)} of type {str(type(rho))} was passed.') - _check_in_zero_one(abs(rho), 'The absolute value of rho') - _check_in_zero_one(level, 'The confidence level', include_zero=False, include_one=False) + raise TypeError(f"rho must be of float type. " f"{str(rho)} of type {str(type(rho))} was passed.") + _check_in_zero_one(abs(rho), "The absolute value of rho") + _check_in_zero_one(level, "The confidence level", include_zero=False, include_one=False) # set elements for readability - sigma2 = self.sensitivity_elements['sigma2'] - nu2 = self.sensitivity_elements['nu2'] - psi_sigma = self.sensitivity_elements['psi_sigma2'] - psi_nu = self.sensitivity_elements['psi_nu2'] + sigma2 = self.sensitivity_elements["sigma2"] + nu2 = self.sensitivity_elements["nu2"] + psi_sigma = self.sensitivity_elements["psi_sigma2"] + psi_nu = self.sensitivity_elements["psi_nu2"] psi_scaled = np.divide(self.psi, np.mean(self.psi_deriv, axis=0)) if (np.any(sigma2 < 0)) | (np.any(nu2 < 0)): - raise ValueError('sensitivity_elements sigma2 and nu2 have to be positive. ' - f"Got sigma2 {str(sigma2)} and nu2 {str(nu2)}. " - 'Most likely this is due to low quality learners (especially propensity scores).') + raise ValueError( + "sensitivity_elements sigma2 and nu2 have to be positive. " + f"Got sigma2 {str(sigma2)} and nu2 {str(nu2)}. " + "Most likely this is due to low quality learners (especially propensity scores)." + ) # elementwise operations - confounding_strength = np.multiply(np.abs(rho), np.sqrt(np.multiply(cf_y, np.divide(cf_d, 1.0-cf_d)))) + confounding_strength = np.multiply(np.abs(rho), np.sqrt(np.multiply(cf_y, np.divide(cf_d, 1.0 - cf_d)))) S = np.sqrt(np.multiply(sigma2, nu2)) # sigma2 and nu2 are of shape (1, n_rep, n_coefs), whereas the all_coefs is of shape (n_coefs, n_reps) @@ -1645,22 +1763,26 @@ def _calc_sensitivity_analysis(self, cf_y, cf_d, rho, level): smpls_cluster = self.__smpls_cluster n_folds_per_cluster = self._n_folds_per_cluster - sigma2_lower_hat, _ = _var_est(psi=psi_lower[:, i_rep, i_d], - psi_deriv=np.ones_like(psi_lower[:, i_rep, i_d]), - apply_cross_fitting=self.apply_cross_fitting, - smpls=self.__smpls, - is_cluster_data=self._is_cluster_data, - cluster_vars=cluster_vars, - smpls_cluster=smpls_cluster, - n_folds_per_cluster=n_folds_per_cluster) - sigma2_upper_hat, _ = _var_est(psi=psi_upper[:, i_rep, i_d], - psi_deriv=np.ones_like(psi_upper[:, i_rep, i_d]), - apply_cross_fitting=self.apply_cross_fitting, - smpls=self.__smpls, - is_cluster_data=self._is_cluster_data, - cluster_vars=cluster_vars, - smpls_cluster=smpls_cluster, - n_folds_per_cluster=n_folds_per_cluster) + sigma2_lower_hat, _ = _var_est( + psi=psi_lower[:, i_rep, i_d], + psi_deriv=np.ones_like(psi_lower[:, i_rep, i_d]), + apply_cross_fitting=self.apply_cross_fitting, + smpls=self.__smpls, + is_cluster_data=self._is_cluster_data, + cluster_vars=cluster_vars, + smpls_cluster=smpls_cluster, + n_folds_per_cluster=n_folds_per_cluster, + ) + sigma2_upper_hat, _ = _var_est( + psi=psi_upper[:, i_rep, i_d], + psi_deriv=np.ones_like(psi_upper[:, i_rep, i_d]), + apply_cross_fitting=self.apply_cross_fitting, + smpls=self.__smpls, + is_cluster_data=self._is_cluster_data, + cluster_vars=cluster_vars, + smpls_cluster=smpls_cluster, + n_folds_per_cluster=n_folds_per_cluster, + ) all_sigma_lower[self._i_treat, self._i_rep] = np.sqrt(sigma2_lower_hat) all_sigma_upper[self._i_treat, self._i_rep] = np.sqrt(sigma2_upper_hat) @@ -1673,38 +1795,33 @@ def _calc_sensitivity_analysis(self, cf_y, cf_d, rho, level): ci_lower = theta_lower - np.multiply(quant, sigma_lower) ci_upper = theta_upper + np.multiply(quant, sigma_upper) - theta_dict = {'lower': theta_lower, - 'upper': theta_upper} + theta_dict = {"lower": theta_lower, "upper": theta_upper} - se_dict = {'lower': sigma_lower, - 'upper': sigma_upper} + se_dict = {"lower": sigma_lower, "upper": sigma_upper} - ci_dict = {'lower': ci_lower, - 'upper': ci_upper} + ci_dict = {"lower": ci_lower, "upper": ci_upper} - res_dict = {'theta': theta_dict, - 'se': se_dict, - 'ci': ci_dict} + res_dict = {"theta": theta_dict, "se": se_dict, "ci": ci_dict} return res_dict def _calc_robustness_value(self, null_hypothesis, level, rho, idx_treatment): _check_float(null_hypothesis, "null_hypothesis") - _check_integer(idx_treatment, "idx_treatment", lower_bound=0, upper_bound=self._dml_data.n_treat-1) + _check_integer(idx_treatment, "idx_treatment", lower_bound=0, upper_bound=self._dml_data.n_treat - 1) # check which side is relvant - bound = 'upper' if (null_hypothesis > self.coef[idx_treatment]) else 'lower' + bound = "upper" if (null_hypothesis > self.coef[idx_treatment]) else "lower" # minimize the square to find boundary solutions def rv_fct(value, param): - res = self._calc_sensitivity_analysis(cf_y=value, - cf_d=value, - rho=rho, - level=level)[param][bound][idx_treatment] - null_hypothesis + res = ( + self._calc_sensitivity_analysis(cf_y=value, cf_d=value, rho=rho, level=level)[param][bound][idx_treatment] + - null_hypothesis + ) return np.square(res) - rv = minimize_scalar(rv_fct, bounds=(0, 0.9999), method='bounded', args=('theta', )).x - rva = minimize_scalar(rv_fct, bounds=(0, 0.9999), method='bounded', args=('ci', )).x + rv = minimize_scalar(rv_fct, bounds=(0, 0.9999), method="bounded", args=("theta",)).x + rva = minimize_scalar(rv_fct, bounds=(0, 0.9999), method="bounded", args=("ci",)).x return rv, rva @@ -1753,31 +1870,32 @@ def sensitivity_analysis(self, cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95, null_h if null_hypothesis.shape == (self._dml_data.n_treat,): null_hypothesis_vec = null_hypothesis else: - raise ValueError("null_hypothesis is numpy.ndarray but does not have the required " - f"shape ({self._dml_data.n_treat},). " - f'Array of shape {str(null_hypothesis.shape)} was passed.') + raise ValueError( + "null_hypothesis is numpy.ndarray but does not have the required " + f"shape ({self._dml_data.n_treat},). " + f"Array of shape {str(null_hypothesis.shape)} was passed." + ) else: - raise TypeError("null_hypothesis has to be of type float or np.ndarry. " - f"{str(null_hypothesis)} of type {str(type(null_hypothesis))} was passed.") + raise TypeError( + "null_hypothesis has to be of type float or np.ndarry. " + f"{str(null_hypothesis)} of type {str(type(null_hypothesis))} was passed." + ) # compute robustess values with respect to null_hypothesis rv = np.full(shape=self._dml_data.n_treat, fill_value=np.nan) rva = np.full(shape=self._dml_data.n_treat, fill_value=np.nan) for i_treat in range(self._dml_data.n_treat): - rv[i_treat], rva[i_treat] = self._calc_robustness_value(null_hypothesis=null_hypothesis_vec[i_treat], - level=level, rho=rho, idx_treatment=i_treat) + rv[i_treat], rva[i_treat] = self._calc_robustness_value( + null_hypothesis=null_hypothesis_vec[i_treat], level=level, rho=rho, idx_treatment=i_treat + ) - sensitivity_dict['rv'] = rv - sensitivity_dict['rva'] = rva + sensitivity_dict["rv"] = rv + sensitivity_dict["rva"] = rva # add all input parameters - input_params = {'cf_y': cf_y, - 'cf_d': cf_d, - 'rho': rho, - 'level': level, - 'null_hypothesis': null_hypothesis_vec} - sensitivity_dict['input'] = input_params + input_params = {"cf_y": cf_y, "cf_d": cf_d, "rho": rho, "level": level, "null_hypothesis": null_hypothesis_vec} + sensitivity_dict["input"] = input_params self._sensitivity_params = sensitivity_dict return self @@ -1792,47 +1910,67 @@ def sensitivity_summary(self): res : str Summary for the sensitivity analysis. """ - header = '================== Sensitivity Analysis ==================\n' + header = "================== Sensitivity Analysis ==================\n" if self.sensitivity_params is None: - res = header + 'Apply sensitivity_analysis() to generate sensitivity_summary.' + res = header + "Apply sensitivity_analysis() to generate sensitivity_summary." else: sig_level = f'Significance Level: level={self.sensitivity_params["input"]["level"]}\n' - scenario_params = f'Sensitivity parameters: cf_y={self.sensitivity_params["input"]["cf_y"]}; ' \ - f'cf_d={self.sensitivity_params["input"]["cf_d"]}, ' \ - f'rho={self.sensitivity_params["input"]["rho"]}' - - theta_and_ci_col_names = ['CI lower', 'theta lower', ' theta', 'theta upper', 'CI upper'] - theta_and_ci = np.transpose(np.vstack((self._sensitivity_params['ci']['lower'], - self._sensitivity_params['theta']['lower'], - self.coef, - self._sensitivity_params['theta']['upper'], - self._sensitivity_params['ci']['upper']))) - df_theta_and_ci = pd.DataFrame(theta_and_ci, - columns=theta_and_ci_col_names, - index=self._dml_data.d_cols) + scenario_params = ( + f'Sensitivity parameters: cf_y={self.sensitivity_params["input"]["cf_y"]}; ' + f'cf_d={self.sensitivity_params["input"]["cf_d"]}, ' + f'rho={self.sensitivity_params["input"]["rho"]}' + ) + + theta_and_ci_col_names = ["CI lower", "theta lower", " theta", "theta upper", "CI upper"] + theta_and_ci = np.transpose( + np.vstack( + ( + self._sensitivity_params["ci"]["lower"], + self._sensitivity_params["theta"]["lower"], + self.coef, + self._sensitivity_params["theta"]["upper"], + self._sensitivity_params["ci"]["upper"], + ) + ) + ) + df_theta_and_ci = pd.DataFrame(theta_and_ci, columns=theta_and_ci_col_names, index=self._dml_data.d_cols) theta_and_ci_summary = str(df_theta_and_ci) - rvs_col_names = ['H_0', 'RV (%)', 'RVa (%)'] - rvs = np.transpose(np.vstack((self._sensitivity_params['rv'], - self._sensitivity_params['rva']))) * 100 + rvs_col_names = ["H_0", "RV (%)", "RVa (%)"] + rvs = np.transpose(np.vstack((self._sensitivity_params["rv"], self._sensitivity_params["rva"]))) * 100 - df_rvs = pd.DataFrame(np.column_stack((self.sensitivity_params["input"]["null_hypothesis"], rvs)), - columns=rvs_col_names, - index=self._dml_data.d_cols) + df_rvs = pd.DataFrame( + np.column_stack((self.sensitivity_params["input"]["null_hypothesis"], rvs)), + columns=rvs_col_names, + index=self._dml_data.d_cols, + ) rvs_summary = str(df_rvs) - res = header + \ - '\n------------------ Scenario ------------------\n' + \ - sig_level + scenario_params + '\n' + \ - '\n------------------ Bounds with CI ------------------\n' + \ - theta_and_ci_summary + '\n' + \ - '\n------------------ Robustness Values ------------------\n' + \ - rvs_summary + res = ( + header + + "\n------------------ Scenario ------------------\n" + + sig_level + + scenario_params + + "\n" + + "\n------------------ Bounds with CI ------------------\n" + + theta_and_ci_summary + + "\n" + + "\n------------------ Robustness Values ------------------\n" + + rvs_summary + ) return res - def sensitivity_plot(self, idx_treatment=0, value='theta', include_scenario=True, benchmarks=None, - fill=True, grid_bounds=(0.15, 0.15), grid_size=100): + def sensitivity_plot( + self, + idx_treatment=0, + value="theta", + include_scenario=True, + benchmarks=None, + fill=True, + grid_bounds=(0.15, 0.15), + grid_size=100, + ): """ Contour plot of the sensivity with respect to latent/confounding variables. @@ -1873,27 +2011,27 @@ def sensitivity_plot(self, idx_treatment=0, value='theta', include_scenario=True Plotly figure of the sensitivity contours. """ if self.sensitivity_params is None: - raise ValueError('Apply sensitivity_analysis() to include senario in sensitivity_plot. ' - 'The values of rho and the level are used for the scenario.') - _check_integer(idx_treatment, "idx_treatment", lower_bound=0, upper_bound=self._dml_data.n_treat-1) + raise ValueError( + "Apply sensitivity_analysis() to include senario in sensitivity_plot. " + "The values of rho and the level are used for the scenario." + ) + _check_integer(idx_treatment, "idx_treatment", lower_bound=0, upper_bound=self._dml_data.n_treat - 1) if not isinstance(value, str): - raise TypeError('value must be a string. ' - f'{str(value)} of type {type(value)} was passed.') - valid_values = ['theta', 'ci'] + raise TypeError("value must be a string. " f"{str(value)} of type {type(value)} was passed.") + valid_values = ["theta", "ci"] if value not in valid_values: - raise ValueError('Invalid value ' + value + '. ' + - 'Valid values ' + ' or '.join(valid_values) + '.') - _check_bool(include_scenario, 'include_scenario') + raise ValueError("Invalid value " + value + ". " + "Valid values " + " or ".join(valid_values) + ".") + _check_bool(include_scenario, "include_scenario") _check_benchmarks(benchmarks) - _check_bool(fill, 'fill') + _check_bool(fill, "fill") _check_in_zero_one(grid_bounds[0], "grid_bounds", include_zero=False, include_one=False) _check_in_zero_one(grid_bounds[1], "grid_bounds", include_zero=False, include_one=False) _check_integer(grid_size, "grid_size", lower_bound=10) - null_hypothesis = self.sensitivity_params['input']['null_hypothesis'][idx_treatment] + null_hypothesis = self.sensitivity_params["input"]["null_hypothesis"][idx_treatment] unadjusted_theta = self.coef[idx_treatment] # check which side is relvant - bound = 'upper' if (null_hypothesis > unadjusted_theta) else 'lower' + bound = "upper" if (null_hypothesis > unadjusted_theta) else "lower" # create evaluation grid cf_d_vec = np.linspace(0, grid_bounds[0], grid_size) @@ -1903,19 +2041,21 @@ def sensitivity_plot(self, idx_treatment=0, value='theta', include_scenario=True contour_values = np.full(shape=(grid_size, grid_size), fill_value=np.nan) for i_cf_d_grid, cf_d_grid in enumerate(cf_d_vec): for i_cf_y_grid, cf_y_grid in enumerate(cf_y_vec): - sens_dict = self._calc_sensitivity_analysis(cf_y=cf_y_grid, - cf_d=cf_d_grid, - rho=self.sensitivity_params['input']['rho'], - level=self.sensitivity_params['input']['level']) + sens_dict = self._calc_sensitivity_analysis( + cf_y=cf_y_grid, + cf_d=cf_d_grid, + rho=self.sensitivity_params["input"]["rho"], + level=self.sensitivity_params["input"]["level"], + ) contour_values[i_cf_d_grid, i_cf_y_grid] = sens_dict[value][bound][idx_treatment] # get the correct unadjusted value for confidence bands - if value == 'theta': + if value == "theta": unadjusted_value = unadjusted_theta else: - assert value == 'ci' - ci = self.confint(level=self.sensitivity_params['input']['level']) - if bound == 'upper': + assert value == "ci" + ci = self.confint(level=self.sensitivity_params["input"]["level"]) + if bound == "upper": unadjusted_value = ci.iloc[idx_treatment, 1] else: unadjusted_value = ci.iloc[idx_treatment, 0] @@ -1923,25 +2063,29 @@ def sensitivity_plot(self, idx_treatment=0, value='theta', include_scenario=True # compute the values for the benchmarks benchmark_dict = copy.deepcopy(benchmarks) if benchmarks is not None: - n_benchmarks = len(benchmarks['name']) + n_benchmarks = len(benchmarks["name"]) benchmark_values = np.full(shape=(n_benchmarks,), fill_value=np.nan) - for benchmark_idx in range(len(benchmarks['name'])): - sens_dict_bench = self._calc_sensitivity_analysis(cf_y=benchmarks['cf_y'][benchmark_idx], - cf_d=benchmarks['cf_y'][benchmark_idx], - rho=self.sensitivity_params['input']['rho'], - level=self.sensitivity_params['input']['level']) + for benchmark_idx in range(len(benchmarks["name"])): + sens_dict_bench = self._calc_sensitivity_analysis( + cf_y=benchmarks["cf_y"][benchmark_idx], + cf_d=benchmarks["cf_y"][benchmark_idx], + rho=self.sensitivity_params["input"]["rho"], + level=self.sensitivity_params["input"]["level"], + ) benchmark_values[benchmark_idx] = sens_dict_bench[value][bound][idx_treatment] - benchmark_dict['value'] = benchmark_values - fig = _sensitivity_contour_plot(x=cf_d_vec, - y=cf_y_vec, - contour_values=contour_values, - unadjusted_value=unadjusted_value, - scenario_x=self.sensitivity_params['input']['cf_d'], - scenario_y=self.sensitivity_params['input']['cf_y'], - scenario_value=self.sensitivity_params[value][bound][idx_treatment], - include_scenario=include_scenario, - benchmarks=benchmark_dict, - fill=fill) + benchmark_dict["value"] = benchmark_values + fig = _sensitivity_contour_plot( + x=cf_d_vec, + y=cf_y_vec, + contour_values=contour_values, + unadjusted_value=unadjusted_value, + scenario_x=self.sensitivity_params["input"]["cf_d"], + scenario_y=self.sensitivity_params["input"]["cf_y"], + scenario_value=self.sensitivity_params[value][bound][idx_treatment], + include_scenario=include_scenario, + benchmarks=benchmark_dict, + fill=fill, + ) return fig def sensitivity_benchmark(self, benchmarking_set): @@ -1957,15 +2101,18 @@ def sensitivity_benchmark(self, benchmarking_set): # input checks if self._sensitivity_elements is None: - raise NotImplementedError(f'Sensitivity analysis not yet implemented for {str(type(self))}.') + raise NotImplementedError(f"Sensitivity analysis not yet implemented for {str(type(self))}.") if not isinstance(benchmarking_set, list): - raise TypeError('benchmarking_set must be a list. ' - f'{str(benchmarking_set)} of type {type(benchmarking_set)} was passed.') + raise TypeError( + "benchmarking_set must be a list. " f"{str(benchmarking_set)} of type {type(benchmarking_set)} was passed." + ) if len(benchmarking_set) == 0: - raise ValueError('benchmarking_set must not be empty.') + raise ValueError("benchmarking_set must not be empty.") if not set(benchmarking_set) <= set(x_list_long): - raise ValueError(f"benchmarking_set must be a subset of features {str(self._dml_data.x_cols)}. " - f'{str(benchmarking_set)} was passed.') + raise ValueError( + f"benchmarking_set must be a subset of features {str(self._dml_data.x_cols)}. " + f"{str(benchmarking_set)} was passed." + ) # refit short form of the model x_list_short = [x for x in x_list_long if x not in benchmarking_set] @@ -1975,10 +2122,10 @@ def sensitivity_benchmark(self, benchmarking_set): # save elements for readability var_y = np.var(self._dml_data.y) - var_y_residuals_long = np.squeeze(self.sensitivity_elements['sigma2'], axis=0) - nu2_long = np.squeeze(self.sensitivity_elements['nu2'], axis=0) - var_y_residuals_short = np.squeeze(dml_short.sensitivity_elements['sigma2'], axis=0) - nu2_short = np.squeeze(dml_short.sensitivity_elements['nu2'], axis=0) + var_y_residuals_long = np.squeeze(self.sensitivity_elements["sigma2"], axis=0) + nu2_long = np.squeeze(self.sensitivity_elements["nu2"], axis=0) + var_y_residuals_short = np.squeeze(dml_short.sensitivity_elements["sigma2"], axis=0) + nu2_short = np.squeeze(dml_short.sensitivity_elements["nu2"], axis=0) # compute nonparametric R2 R2_y_long = 1.0 - np.divide(var_y_residuals_long, var_y) @@ -2000,11 +2147,9 @@ def sensitivity_benchmark(self, benchmarking_set): var_riesz = nu2_long - nu2_short denom = np.sqrt(np.multiply(var_g, var_riesz), out=np.zeros_like(var_g), where=(var_g > 0) & (var_riesz > 0)) rho_sign = np.sign(all_delta_theta) - rho_values = np.clip(np.divide(np.absolute(all_delta_theta), - denom, - out=np.ones_like(all_delta_theta), - where=denom != 0), - 0.0, 1.0) + rho_values = np.clip( + np.divide(np.absolute(all_delta_theta), denom, out=np.ones_like(all_delta_theta), where=denom != 0), 0.0, 1.0 + ) all_rho_benchmark = np.multiply(rho_values, rho_sign) rho_benchmark = np.median(all_rho_benchmark, axis=0) benchmark_dict = { diff --git a/doubleml/double_ml_did.py b/doubleml/double_ml_did.py index 1add5e0d..707bd621 100644 --- a/doubleml/double_ml_did.py +++ b/doubleml/double_ml_did.py @@ -85,62 +85,67 @@ class DoubleMLDID(LinearScoreMixin, DoubleML): coef std err t P>|t| 2.5 % 97.5 % d -2.685104 1.798071 -1.493325 0.135352 -6.209257 0.83905 """ - def __init__(self, - obj_dml_data, - ml_g, - ml_m=None, - n_folds=5, - n_rep=1, - score='observational', - in_sample_normalization=True, - dml_procedure='dml2', - trimming_rule='truncate', - trimming_threshold=1e-2, - draw_sample_splitting=True, - apply_cross_fitting=True): - super().__init__(obj_dml_data, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + + def __init__( + self, + obj_dml_data, + ml_g, + ml_m=None, + n_folds=5, + n_rep=1, + score="observational", + in_sample_normalization=True, + dml_procedure="dml2", + trimming_rule="truncate", + trimming_threshold=1e-2, + draw_sample_splitting=True, + apply_cross_fitting=True, + ): + super().__init__(obj_dml_data, n_folds, n_rep, score, dml_procedure, draw_sample_splitting, apply_cross_fitting) self._check_data(self._dml_data) - valid_scores = ['observational', 'experimental'] + valid_scores = ["observational", "experimental"] _check_score(self.score, valid_scores, allow_callable=False) self._in_sample_normalization = in_sample_normalization if not isinstance(self.in_sample_normalization, bool): - raise TypeError('in_sample_normalization indicator has to be boolean. ' + - f'Object of type {str(type(self.in_sample_normalization))} passed.') + raise TypeError( + "in_sample_normalization indicator has to be boolean. " + + f"Object of type {str(type(self.in_sample_normalization))} passed." + ) # set stratication for resampling self._strata = self._dml_data.d # check learners - ml_g_is_classifier = self._check_learner(ml_g, 'ml_g', regressor=True, classifier=True) - if self.score == 'observational': - _ = self._check_learner(ml_m, 'ml_m', regressor=False, classifier=True) - self._learner = {'ml_g': ml_g, 'ml_m': ml_m} + ml_g_is_classifier = self._check_learner(ml_g, "ml_g", regressor=True, classifier=True) + if self.score == "observational": + _ = self._check_learner(ml_m, "ml_m", regressor=False, classifier=True) + self._learner = {"ml_g": ml_g, "ml_m": ml_m} else: - assert self.score == 'experimental' + assert self.score == "experimental" if ml_m is not None: - warnings.warn(('A learner ml_m has been provided for score = "experimental" but will be ignored. ' - 'A learner ml_m is not required for estimation.')) - self._learner = {'ml_g': ml_g} + warnings.warn( + ( + 'A learner ml_m has been provided for score = "experimental" but will be ignored. ' + "A learner ml_m is not required for estimation." + ) + ) + self._learner = {"ml_g": ml_g} if ml_g_is_classifier: if obj_dml_data.binary_outcome: - self._predict_method = {'ml_g': 'predict_proba'} + self._predict_method = {"ml_g": "predict_proba"} else: - raise ValueError(f'The ml_g learner {str(ml_g)} was identified as classifier ' - 'but the outcome variable is not binary with values 0 and 1.') + raise ValueError( + f"The ml_g learner {str(ml_g)} was identified as classifier " + "but the outcome variable is not binary with values 0 and 1." + ) else: - self._predict_method = {'ml_g': 'predict'} + self._predict_method = {"ml_g": "predict"} - if 'ml_m' in self._learner: - self._predict_method['ml_m'] = 'predict_proba' + if "ml_m" in self._learner: + self._predict_method["ml_m"] = "predict_proba" self._initialize_ml_nuisance_params() self._trimming_rule = trimming_rule @@ -148,7 +153,7 @@ def __init__(self, _check_trimming(self._trimming_rule, self._trimming_threshold) self._sensitivity_implemented = True - + self._external_predictions_implemented = True @property @@ -173,105 +178,102 @@ def trimming_threshold(self): return self._trimming_threshold def _initialize_ml_nuisance_params(self): - valid_learner = ['ml_g0', 'ml_g1', 'ml_m'] - self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} - for learner in valid_learner} + valid_learner = ["ml_g0", "ml_g1", "ml_m"] + self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} for learner in valid_learner} def _check_data(self, obj_dml_data): if not isinstance(obj_dml_data, DoubleMLData): - raise TypeError('For repeated outcomes the data must be of DoubleMLData type. ' - f'{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed.') + raise TypeError( + "For repeated outcomes the data must be of DoubleMLData type. " + f"{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed." + ) if obj_dml_data.z_cols is not None: - raise ValueError('Incompatible data. ' + - ' and '.join(obj_dml_data.z_cols) + - ' have been set as instrumental variable(s). ' - 'At the moment there are not DiD models with instruments implemented.') - one_treat = (obj_dml_data.n_treat == 1) - binary_treat = (type_of_target(obj_dml_data.d) == 'binary') + raise ValueError( + "Incompatible data. " + " and ".join(obj_dml_data.z_cols) + " have been set as instrumental variable(s). " + "At the moment there are not DiD models with instruments implemented." + ) + one_treat = obj_dml_data.n_treat == 1 + binary_treat = type_of_target(obj_dml_data.d) == "binary" zero_one_treat = np.all((np.power(obj_dml_data.d, 2) - obj_dml_data.d) == 0) if not (one_treat & binary_treat & zero_one_treat): - raise ValueError('Incompatible data. ' - 'To fit an DID model with DML ' - 'exactly one binary variable with values 0 and 1 ' - 'needs to be specified as treatment variable.') + raise ValueError( + "Incompatible data. " + "To fit an DID model with DML " + "exactly one binary variable with values 0 and 1 " + "needs to be specified as treatment variable." + ) return def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=False): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) # nuisance g # get train indices for d == 0 smpls_d0, smpls_d1 = _get_cond_smpls(smpls, d) # nuisance g for d==0 - if external_predictions['ml_g0'] is not None: - g_hat0 = {'preds': external_predictions['ml_g0'], - 'targets': None, - 'models': None} + if external_predictions["ml_g0"] is not None: + g_hat0 = {"preds": external_predictions["ml_g0"], "targets": None, "models": None} else: g_hat0 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls=smpls_d0, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g0'), method=self._predict_method['ml_g'], - return_models=return_models) + est_params=self._get_params('ml_g0'), method=self._predict_method['ml_g'], + return_models=return_models) - _check_finite_predictions(g_hat0['preds'], self._learner['ml_g'], 'ml_g', smpls) + _check_finite_predictions(g_hat0["preds"], self._learner["ml_g"], "ml_g", smpls) # adjust target values to consider only compatible subsamples - g_hat0['targets'] = g_hat0['targets'].astype(float) - g_hat0['targets'][d == 1] = np.nan + g_hat0["targets"] = g_hat0["targets"].astype(float) + g_hat0["targets"][d == 1] = np.nan # nuisance g for d==1 - if external_predictions['ml_g1'] is not None: - g_hat1 = {'preds': external_predictions['ml_g1'], - 'targets': None, - 'models': None} + if external_predictions["ml_g1"] is not None: + g_hat1 = {"preds": external_predictions["ml_g1"], "targets": None, "models": None} else: g_hat1 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls=smpls_d1, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g1'), method=self._predict_method['ml_g'], - return_models=return_models) + est_params=self._get_params('ml_g1'), method=self._predict_method['ml_g'], + return_models=return_models) - _check_finite_predictions(g_hat1['preds'], self._learner['ml_g'], 'ml_g', smpls) + _check_finite_predictions(g_hat1["preds"], self._learner["ml_g"], "ml_g", smpls) # adjust target values to consider only compatible subsamples - g_hat1['targets'] = g_hat1['targets'].astype(float) - g_hat1['targets'][d == 0] = np.nan + g_hat1["targets"] = g_hat1["targets"].astype(float) + g_hat1["targets"][d == 0] = np.nan # only relevant for observational setting - m_hat = {'preds': None, 'targets': None, 'models': None} - if self.score == 'observational': + m_hat = {"preds": None, "targets": None, "models": None} + if self.score == "observational": # nuisance m if external_predictions['ml_m'] is not None: m_hat = {'preds': external_predictions['ml_m'], - 'targets': None, - 'models': None} + 'targets': None, + 'models': None} else: - m_hat = _dml_cv_predict(self._learner['ml_m'], x, d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_m'), method=self._predict_method['ml_m'], - return_models=return_models) - _check_finite_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls) - _check_is_propensity(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls, eps=1e-12) - m_hat['preds'] = _trimm(m_hat['preds'], self.trimming_rule, self.trimming_threshold) + m_hat = _dml_cv_predict( + self._learner["ml_m"], + x, + d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_m"), + method=self._predict_method["ml_m"], + return_models=return_models, + ) + _check_finite_predictions(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls) + _check_is_propensity(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls, eps=1e-12) + m_hat["preds"] = _trimm(m_hat["preds"], self.trimming_rule, self.trimming_threshold) # nuisance estimates of the uncond. treatment prob. - p_hat = np.full_like(d, np.nan, dtype='float64') + p_hat = np.full_like(d, np.nan, dtype="float64") for train_index, test_index in smpls: p_hat[test_index] = np.mean(d[train_index]) - psi_a, psi_b = self._score_elements(y, d, g_hat0['preds'], g_hat1['preds'], m_hat['preds'], p_hat) - - psi_elements = {'psi_a': psi_a, - 'psi_b': psi_b} - preds = {'predictions': {'ml_g0': g_hat0['preds'], - 'ml_g1': g_hat1['preds'], - 'ml_m': m_hat['preds']}, - 'targets': {'ml_g0': g_hat0['targets'], - 'ml_g1': g_hat1['targets'], - 'ml_m': m_hat['targets']}, - 'models': {'ml_g0': g_hat0['models'], - 'ml_g1': g_hat1['models'], - 'ml_m': m_hat['models'] - } - } + psi_a, psi_b = self._score_elements(y, d, g_hat0["preds"], g_hat1["preds"], m_hat["preds"], p_hat) + + psi_elements = {"psi_a": psi_a, "psi_b": psi_b} + preds = { + "predictions": {"ml_g0": g_hat0["preds"], "ml_g1": g_hat1["preds"], "ml_m": m_hat["preds"]}, + "targets": {"ml_g0": g_hat0["targets"], "ml_g1": g_hat1["targets"], "ml_m": m_hat["targets"]}, + "models": {"ml_g0": g_hat0["models"], "ml_g1": g_hat1["models"], "ml_m": m_hat["models"]}, + } return psi_elements, preds @@ -279,35 +281,35 @@ def _score_elements(self, y, d, g_hat0, g_hat1, m_hat, p_hat): # calc residuals resid_d0 = y - g_hat0 - if self.score == 'observational': + if self.score == "observational": if self.in_sample_normalization: weight_psi_a = np.divide(d, np.mean(d)) - propensity_weight = np.multiply(1.0-d, np.divide(m_hat, 1.0-m_hat)) + propensity_weight = np.multiply(1.0 - d, np.divide(m_hat, 1.0 - m_hat)) weight_resid_d0 = np.divide(d, np.mean(d)) - np.divide(propensity_weight, np.mean(propensity_weight)) else: weight_psi_a = np.divide(d, p_hat) - weight_resid_d0 = np.divide(d-m_hat, np.multiply(p_hat, 1.0-m_hat)) + weight_resid_d0 = np.divide(d - m_hat, np.multiply(p_hat, 1.0 - m_hat)) psi_b_1 = np.zeros_like(y) else: - assert self.score == 'experimental' + assert self.score == "experimental" if self.in_sample_normalization: weight_psi_a = np.ones_like(y) weight_g0 = np.divide(d, np.mean(d)) - 1.0 weight_g1 = 1.0 - np.divide(d, np.mean(d)) - weight_resid_d0 = np.divide(d, np.mean(d)) - np.divide(1.0-d, np.mean(1.0-d)) + weight_resid_d0 = np.divide(d, np.mean(d)) - np.divide(1.0 - d, np.mean(1.0 - d)) else: weight_psi_a = np.ones_like(y) weight_g0 = np.divide(d, p_hat) - 1.0 weight_g1 = 1.0 - np.divide(d, p_hat) - weight_resid_d0 = np.divide(d-p_hat, np.multiply(p_hat, 1.0-p_hat)) + weight_resid_d0 = np.divide(d - p_hat, np.multiply(p_hat, 1.0 - p_hat)) - psi_b_1 = np.multiply(weight_g0, g_hat0) + np.multiply(weight_g1, g_hat1) + psi_b_1 = np.multiply(weight_g0, g_hat0) + np.multiply(weight_g1, g_hat1) # set score elements psi_a = -1.0 * weight_psi_a - psi_b = psi_b_1 + np.multiply(weight_resid_d0, resid_d0) + psi_b = psi_b_1 + np.multiply(weight_resid_d0, resid_d0) return psi_a, psi_b @@ -315,90 +317,106 @@ def _sensitivity_element_est(self, preds): y = self._dml_data.y d = self._dml_data.d - m_hat = preds['predictions']['ml_m'] - g_hat0 = preds['predictions']['ml_g0'] - g_hat1 = preds['predictions']['ml_g1'] + m_hat = preds["predictions"]["ml_m"] + g_hat0 = preds["predictions"]["ml_g0"] + g_hat1 = preds["predictions"]["ml_g1"] - g_hat = np.multiply(d, g_hat1) + np.multiply(1.0-d, g_hat0) + g_hat = np.multiply(d, g_hat1) + np.multiply(1.0 - d, g_hat0) sigma2_score_element = np.square(y - g_hat) sigma2 = np.mean(sigma2_score_element) psi_sigma2 = sigma2_score_element - sigma2 # calc m(W,alpha) and Riesz representer p_hat = np.mean(d) - if self.score == 'observational': - propensity_weight_d0 = np.divide(m_hat, 1.0-m_hat) + if self.score == "observational": + propensity_weight_d0 = np.divide(m_hat, 1.0 - m_hat) if self.in_sample_normalization: - weight_d0 = np.multiply(1.0-d, propensity_weight_d0) + weight_d0 = np.multiply(1.0 - d, propensity_weight_d0) mean_weight_d0 = np.mean(weight_d0) - m_alpha = np.multiply(np.divide(d, p_hat), - np.divide(1.0, p_hat) + np.divide(propensity_weight_d0, mean_weight_d0)) + m_alpha = np.multiply( + np.divide(d, p_hat), np.divide(1.0, p_hat) + np.divide(propensity_weight_d0, mean_weight_d0) + ) rr = np.divide(d, p_hat) - np.divide(weight_d0, mean_weight_d0) else: m_alpha = np.multiply(np.divide(d, np.square(p_hat)), (1.0 + propensity_weight_d0)) - rr = np.divide(d, p_hat) - np.multiply(np.divide(1.0-d, p_hat), propensity_weight_d0) + rr = np.divide(d, p_hat) - np.multiply(np.divide(1.0 - d, p_hat), propensity_weight_d0) else: - assert self.score == 'experimental' + assert self.score == "experimental" # the same with or without self-normalization - m_alpha = np.divide(1.0, p_hat) + np.divide(1.0, 1.0-p_hat) - rr = np.divide(d, p_hat) - np.divide(1.0-d, 1.0-p_hat) + m_alpha = np.divide(1.0, p_hat) + np.divide(1.0, 1.0 - p_hat) + rr = np.divide(d, p_hat) - np.divide(1.0 - d, 1.0 - p_hat) nu2_score_element = np.multiply(2.0, m_alpha) - np.square(rr) nu2 = np.mean(nu2_score_element) psi_nu2 = nu2_score_element - nu2 - element_dict = {'sigma2': sigma2, - 'nu2': nu2, - 'psi_sigma2': psi_sigma2, - 'psi_nu2': psi_nu2} + element_dict = {"sigma2": sigma2, "nu2": nu2, "psi_sigma2": psi_sigma2, "psi_nu2": psi_nu2} return element_dict - def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + def _nuisance_tuning( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) # get train indices for d == 0 and d == 1 smpls_d0, smpls_d1 = _get_cond_smpls(smpls, d) if scoring_methods is None: - scoring_methods = {'ml_g': None, - 'ml_m': None} + scoring_methods = {"ml_g": None, "ml_m": None} train_inds = [train_index for (train_index, _) in smpls] train_inds_d0 = [train_index for (train_index, _) in smpls_d0] train_inds_d1 = [train_index for (train_index, _) in smpls_d1] - g0_tune_res = _dml_tune(y, x, train_inds_d0, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - g1_tune_res = _dml_tune(y, x, train_inds_d1, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + g0_tune_res = _dml_tune( + y, + x, + train_inds_d0, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + g1_tune_res = _dml_tune( + y, + x, + train_inds_d1, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) g0_best_params = [xx.best_params_ for xx in g0_tune_res] g1_best_params = [xx.best_params_ for xx in g1_tune_res] m_tune_res = list() - if self.score == 'observational': - m_tune_res = _dml_tune(d, x, train_inds, - self._learner['ml_m'], param_grids['ml_m'], scoring_methods['ml_m'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + if self.score == "observational": + m_tune_res = _dml_tune( + d, + x, + train_inds, + self._learner["ml_m"], + param_grids["ml_m"], + scoring_methods["ml_m"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) m_best_params = [xx.best_params_ for xx in m_tune_res] - params = {'ml_g0': g0_best_params, - 'ml_g1': g1_best_params, - 'ml_m': m_best_params} - tune_res = {'g0_tune': g0_tune_res, - 'g1_tune': g1_tune_res, - 'm_tune': m_tune_res} + params = {"ml_g0": g0_best_params, "ml_g1": g1_best_params, "ml_m": m_best_params} + tune_res = {"g0_tune": g0_tune_res, "g1_tune": g1_tune_res, "m_tune": m_tune_res} else: - params = {'ml_g0': g0_best_params, - 'ml_g1': g1_best_params} - tune_res = {'g0_tune': g0_tune_res, - 'g1_tune': g1_tune_res} + params = {"ml_g0": g0_best_params, "ml_g1": g1_best_params} + tune_res = {"g0_tune": g0_tune_res, "g1_tune": g1_tune_res} - res = {'params': params, - 'tune_res': tune_res} + res = {"params": params, "tune_res": tune_res} return res diff --git a/doubleml/double_ml_did_cs.py b/doubleml/double_ml_did_cs.py index 2cdbb003..5f13a602 100644 --- a/doubleml/double_ml_did_cs.py +++ b/doubleml/double_ml_did_cs.py @@ -85,62 +85,67 @@ class DoubleMLDIDCS(LinearScoreMixin, DoubleML): coef std err t P>|t| 2.5 % 97.5 % d -6.604603 8.725802 -0.756905 0.449107 -23.706862 10.497655 """ - def __init__(self, - obj_dml_data, - ml_g, - ml_m=None, - n_folds=5, - n_rep=1, - score='observational', - in_sample_normalization=True, - dml_procedure='dml2', - trimming_rule='truncate', - trimming_threshold=1e-2, - draw_sample_splitting=True, - apply_cross_fitting=True): - super().__init__(obj_dml_data, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + + def __init__( + self, + obj_dml_data, + ml_g, + ml_m=None, + n_folds=5, + n_rep=1, + score="observational", + in_sample_normalization=True, + dml_procedure="dml2", + trimming_rule="truncate", + trimming_threshold=1e-2, + draw_sample_splitting=True, + apply_cross_fitting=True, + ): + super().__init__(obj_dml_data, n_folds, n_rep, score, dml_procedure, draw_sample_splitting, apply_cross_fitting) self._check_data(self._dml_data) - valid_scores = ['observational', 'experimental'] + valid_scores = ["observational", "experimental"] _check_score(self.score, valid_scores, allow_callable=False) self._in_sample_normalization = in_sample_normalization if not isinstance(self.in_sample_normalization, bool): - raise TypeError('in_sample_normalization indicator has to be boolean. ' + - f'Object of type {str(type(self.in_sample_normalization))} passed.') + raise TypeError( + "in_sample_normalization indicator has to be boolean. " + + f"Object of type {str(type(self.in_sample_normalization))} passed." + ) # set stratication for resampling self._strata = self._dml_data.d.reshape(-1, 1) + 2 * self._dml_data.t.reshape(-1, 1) # check learners - ml_g_is_classifier = self._check_learner(ml_g, 'ml_g', regressor=True, classifier=True) - if self.score == 'observational': - _ = self._check_learner(ml_m, 'ml_m', regressor=False, classifier=True) - self._learner = {'ml_g': ml_g, 'ml_m': ml_m} + ml_g_is_classifier = self._check_learner(ml_g, "ml_g", regressor=True, classifier=True) + if self.score == "observational": + _ = self._check_learner(ml_m, "ml_m", regressor=False, classifier=True) + self._learner = {"ml_g": ml_g, "ml_m": ml_m} else: - assert self.score == 'experimental' + assert self.score == "experimental" if ml_m is not None: - warnings.warn(('A learner ml_m has been provided for score = "experimental" but will be ignored. ' - 'A learner ml_m is not required for estimation.')) - self._learner = {'ml_g': ml_g} + warnings.warn( + ( + 'A learner ml_m has been provided for score = "experimental" but will be ignored. ' + "A learner ml_m is not required for estimation." + ) + ) + self._learner = {"ml_g": ml_g} if ml_g_is_classifier: if obj_dml_data.binary_outcome: - self._predict_method = {'ml_g': 'predict_proba'} + self._predict_method = {"ml_g": "predict_proba"} else: - raise ValueError(f'The ml_g learner {str(ml_g)} was identified as classifier ' - 'but the outcome variable is not binary with values 0 and 1.') + raise ValueError( + f"The ml_g learner {str(ml_g)} was identified as classifier " + "but the outcome variable is not binary with values 0 and 1." + ) else: - self._predict_method = {'ml_g': 'predict'} + self._predict_method = {"ml_g": "predict"} - if 'ml_m' in self._learner: - self._predict_method['ml_m'] = 'predict_proba' + if "ml_m" in self._learner: + self._predict_method["ml_m"] = "predict_proba" self._initialize_ml_nuisance_params() self._trimming_rule = trimming_rule @@ -148,9 +153,8 @@ def __init__(self, _check_trimming(self._trimming_rule, self._trimming_threshold) self._sensitivity_implemented = True - + self._external_predictions_implemented = True - @property def in_sample_normalization(self): @@ -174,72 +178,69 @@ def trimming_threshold(self): return self._trimming_threshold def _initialize_ml_nuisance_params(self): - valid_learner = ['ml_g_d0_t0', 'ml_g_d0_t1', - 'ml_g_d1_t0', 'ml_g_d1_t1', 'ml_m'] - self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} - for learner in valid_learner} + valid_learner = ["ml_g_d0_t0", "ml_g_d0_t1", "ml_g_d1_t0", "ml_g_d1_t1", "ml_m"] + self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} for learner in valid_learner} def _check_data(self, obj_dml_data): if not isinstance(obj_dml_data, DoubleMLData): - raise TypeError('For repeated cross sections the data must be of DoubleMLData type. ' - f'{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed.') + raise TypeError( + "For repeated cross sections the data must be of DoubleMLData type. " + f"{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed." + ) if obj_dml_data.z_cols is not None: - raise ValueError('Incompatible data. ' + - ' and '.join(obj_dml_data.z_cols) + - ' have been set as instrumental variable(s). ' - 'At the moment there are no DiD models with instruments implemented.') - one_treat = (obj_dml_data.n_treat == 1) - binary_treat = (type_of_target(obj_dml_data.d) == 'binary') - zero_one_treat = np.all( - (np.power(obj_dml_data.d, 2) - obj_dml_data.d) == 0) + raise ValueError( + "Incompatible data. " + " and ".join(obj_dml_data.z_cols) + " have been set as instrumental variable(s). " + "At the moment there are no DiD models with instruments implemented." + ) + one_treat = obj_dml_data.n_treat == 1 + binary_treat = type_of_target(obj_dml_data.d) == "binary" + zero_one_treat = np.all((np.power(obj_dml_data.d, 2) - obj_dml_data.d) == 0) if not (one_treat & binary_treat & zero_one_treat): - raise ValueError('Incompatible data. ' - 'To fit an DIDCS model with DML ' - 'exactly one binary variable with values 0 and 1 ' - 'needs to be specified as treatment variable.') + raise ValueError( + "Incompatible data. " + "To fit an DIDCS model with DML " + "exactly one binary variable with values 0 and 1 " + "needs to be specified as treatment variable." + ) - binary_time = (type_of_target(obj_dml_data.t) == 'binary') - zero_one_time = np.all( - (np.power(obj_dml_data.t, 2) - obj_dml_data.t) == 0) + binary_time = type_of_target(obj_dml_data.t) == "binary" + zero_one_time = np.all((np.power(obj_dml_data.t, 2) - obj_dml_data.t) == 0) if not (binary_time & zero_one_time): - raise ValueError('Incompatible data. ' - 'To fit an DIDCS model with DML ' - 'exactly one binary variable with values 0 and 1 ' - 'needs to be specified as time variable.') + raise ValueError( + "Incompatible data. " + "To fit an DIDCS model with DML " + "exactly one binary variable with values 0 and 1 " + "needs to be specified as time variable." + ) return def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=False): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) - x, t = check_X_y(x, self._dml_data.t, - force_all_finite=False) + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) + x, t = check_X_y(x, self._dml_data.t, force_all_finite=False) # THIS DIFFERS FROM THE PAPER due to stratified splitting this should be the same for each fold # nuisance estimates of the uncond. treatment prob. - p_hat = np.full_like(d, np.nan, dtype='float64') + p_hat = np.full_like(d, np.nan, dtype="float64") for train_index, test_index in smpls: p_hat[test_index] = np.mean(d[train_index]) # nuisance estimates of the uncond. time prob. - lambda_hat = np.full_like(t, np.nan, dtype='float64') + lambda_hat = np.full_like(t, np.nan, dtype="float64") for train_index, test_index in smpls: lambda_hat[test_index] = np.mean(t[train_index]) # nuisance g smpls_d0_t0, smpls_d0_t1, smpls_d1_t0, smpls_d1_t1 = _get_cond_smpls_2d(smpls, d, t) - if external_predictions['ml_g_d0_t0'] is not None: - g_hat_d0_t0 = {'preds': external_predictions['ml_g_d0_t0'], - 'targets': None, - 'models': None} + if external_predictions["ml_g_d0_t0"] is not None: + g_hat_d0_t0 = {"preds": external_predictions["ml_g_d0_t0"], "targets": None, "models": None} else: g_hat_d0_t0 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls_d0_t0, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g_d0_t0'), method=self._predict_method['ml_g'], - return_models=return_models) - + est_params=self._get_params('ml_g_d0_t0'), method=self._predict_method['ml_g'], + return_models=return_models) + g_hat_d0_t0['targets'] = g_hat_d0_t0['targets'].astype(float) g_hat_d0_t0['targets'][np.invert((d == 0) & (t == 0))] = np.nan if external_predictions['ml_g_d0_t1'] is not None: @@ -248,8 +249,8 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa 'models': None} else: g_hat_d0_t1 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls_d0_t1, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g_d0_t1'), method=self._predict_method['ml_g'], - return_models=return_models) + est_params=self._get_params('ml_g_d0_t1'), method=self._predict_method['ml_g'], + return_models=return_models) g_hat_d0_t1['targets'] = g_hat_d0_t1['targets'].astype(float) g_hat_d0_t1['targets'][np.invert((d == 0) & (t == 1))] = np.nan if external_predictions['ml_g_d1_t0'] is not None: @@ -258,8 +259,8 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa 'models': None} else: g_hat_d1_t0 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls_d1_t0, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g_d1_t0'), method=self._predict_method['ml_g'], - return_models=return_models) + est_params=self._get_params('ml_g_d1_t0'), method=self._predict_method['ml_g'], + return_models=return_models) g_hat_d1_t0['targets'] = g_hat_d1_t0['targets'].astype(float) g_hat_d1_t0['targets'][np.invert((d == 1) & (t == 0))] = np.nan if external_predictions['ml_g_d1_t1'] is not None: @@ -268,58 +269,73 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa 'models': None} else: g_hat_d1_t1 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls_d1_t1, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g_d1_t1'), method=self._predict_method['ml_g'], - return_models=return_models) + est_params=self._get_params('ml_g_d1_t1'), method=self._predict_method['ml_g'], + return_models=return_models) g_hat_d1_t1['targets'] = g_hat_d1_t1['targets'].astype(float) g_hat_d1_t1['targets'][np.invert((d == 1) & (t == 1))] = np.nan # only relevant for observational or experimental setting - m_hat = {'preds': None, 'targets': None, 'models': None} - if self.score == 'observational': + m_hat = {"preds": None, "targets": None, "models": None} + if self.score == "observational": # nuisance m - if external_predictions['ml_m'] is not None: - m_hat = {'preds': external_predictions['ml_m'], - 'targets': None, - 'models': None} + if external_predictions["ml_m"] is not None: + m_hat = {"preds": external_predictions["ml_m"], "targets": None, "models": None} else: - m_hat = _dml_cv_predict(self._learner['ml_m'], x, d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_m'), method=self._predict_method['ml_m'], - return_models=return_models) - _check_finite_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls) - _check_is_propensity(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls, eps=1e-12) - m_hat['preds'] = _trimm(m_hat['preds'], self.trimming_rule, self.trimming_threshold) - - psi_a, psi_b = self._score_elements(y, d, t, - g_hat_d0_t0['preds'], g_hat_d0_t1['preds'], - g_hat_d1_t0['preds'], g_hat_d1_t1['preds'], - m_hat['preds'], p_hat, lambda_hat) - - psi_elements = {'psi_a': psi_a, - 'psi_b': psi_b} - preds = {'predictions': {'ml_g_d0_t0': g_hat_d0_t0['preds'], - 'ml_g_d0_t1': g_hat_d0_t1['preds'], - 'ml_g_d1_t0': g_hat_d1_t0['preds'], - 'ml_g_d1_t1': g_hat_d1_t1['preds'], - 'ml_m': m_hat['preds']}, - 'targets': {'ml_g_d0_t0': g_hat_d0_t0['targets'], - 'ml_g_d0_t1': g_hat_d0_t1['targets'], - 'ml_g_d1_t0': g_hat_d1_t0['targets'], - 'ml_g_d1_t1': g_hat_d1_t1['targets'], - 'ml_m': m_hat['targets']}, - 'models': {'ml_g_d0_t0': g_hat_d0_t0['models'], - 'ml_g_d0_t1': g_hat_d0_t1['models'], - 'ml_g_d1_t0': g_hat_d1_t0['models'], - 'ml_g_d1_t1': g_hat_d1_t1['models'], - 'ml_m': m_hat['models']} - } + m_hat = _dml_cv_predict( + self._learner["ml_m"], + x, + d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_m"), + method=self._predict_method["ml_m"], + return_models=return_models, + ) + _check_finite_predictions(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls) + _check_is_propensity(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls, eps=1e-12) + m_hat["preds"] = _trimm(m_hat["preds"], self.trimming_rule, self.trimming_threshold) + + psi_a, psi_b = self._score_elements( + y, + d, + t, + g_hat_d0_t0["preds"], + g_hat_d0_t1["preds"], + g_hat_d1_t0["preds"], + g_hat_d1_t1["preds"], + m_hat["preds"], + p_hat, + lambda_hat, + ) + + psi_elements = {"psi_a": psi_a, "psi_b": psi_b} + preds = { + "predictions": { + "ml_g_d0_t0": g_hat_d0_t0["preds"], + "ml_g_d0_t1": g_hat_d0_t1["preds"], + "ml_g_d1_t0": g_hat_d1_t0["preds"], + "ml_g_d1_t1": g_hat_d1_t1["preds"], + "ml_m": m_hat["preds"], + }, + "targets": { + "ml_g_d0_t0": g_hat_d0_t0["targets"], + "ml_g_d0_t1": g_hat_d0_t1["targets"], + "ml_g_d1_t0": g_hat_d1_t0["targets"], + "ml_g_d1_t1": g_hat_d1_t1["targets"], + "ml_m": m_hat["targets"], + }, + "models": { + "ml_g_d0_t0": g_hat_d0_t0["models"], + "ml_g_d0_t1": g_hat_d0_t1["models"], + "ml_g_d1_t0": g_hat_d1_t0["models"], + "ml_g_d1_t1": g_hat_d1_t1["models"], + "ml_m": m_hat["models"], + }, + } return psi_elements, preds - def _score_elements(self, y, d, t, - g_hat_d0_t0, g_hat_d0_t1, - g_hat_d1_t0, g_hat_d1_t1, - m_hat, p_hat, lambda_hat): - + def _score_elements(self, y, d, t, g_hat_d0_t0, g_hat_d0_t1, g_hat_d1_t0, g_hat_d1_t1, m_hat, p_hat, lambda_hat): # calculate residuals resid_d0_t0 = y - g_hat_d0_t0 resid_d0_t1 = y - g_hat_d0_t1 @@ -327,11 +343,11 @@ def _score_elements(self, y, d, t, resid_d1_t1 = y - g_hat_d1_t1 d1t1 = np.multiply(d, t) - d1t0 = np.multiply(d, 1.0-t) - d0t1 = np.multiply(1.0-d, t) - d0t0 = np.multiply(1.0-d, 1.0-t) + d1t0 = np.multiply(d, 1.0 - t) + d0t1 = np.multiply(1.0 - d, t) + d0t0 = np.multiply(1.0 - d, 1.0 - t) - if self.score == 'observational': + if self.score == "observational": if self.in_sample_normalization: weight_psi_a = np.divide(d, np.mean(d)) weight_g_d1_t1 = weight_psi_a @@ -342,7 +358,7 @@ def _score_elements(self, y, d, t, weight_resid_d1_t1 = np.divide(d1t1, np.mean(d1t1)) weight_resid_d1_t0 = -1.0 * np.divide(d1t0, np.mean(d1t0)) - prop_weighting = np.divide(m_hat, 1.0-m_hat) + prop_weighting = np.divide(m_hat, 1.0 - m_hat) unscaled_d0_t1 = np.multiply(d0t1, prop_weighting) weight_resid_d0_t1 = -1.0 * np.divide(unscaled_d0_t1, np.mean(unscaled_d0_t1)) @@ -356,15 +372,13 @@ def _score_elements(self, y, d, t, weight_g_d0_t0 = weight_psi_a weight_resid_d1_t1 = np.divide(d1t1, np.multiply(p_hat, lambda_hat)) - weight_resid_d1_t0 = -1.0 * np.divide(d1t0, np.multiply(p_hat, 1.0-lambda_hat)) + weight_resid_d1_t0 = -1.0 * np.divide(d1t0, np.multiply(p_hat, 1.0 - lambda_hat)) - prop_weighting = np.divide(m_hat, 1.0-m_hat) - weight_resid_d0_t1 = -1.0 * np.multiply(np.divide(d0t1, np.multiply(p_hat, lambda_hat)), - prop_weighting) - weight_resid_d0_t0 = np.multiply(np.divide(d0t0, np.multiply(p_hat, 1.0-lambda_hat)), - prop_weighting) + prop_weighting = np.divide(m_hat, 1.0 - m_hat) + weight_resid_d0_t1 = -1.0 * np.multiply(np.divide(d0t1, np.multiply(p_hat, lambda_hat)), prop_weighting) + weight_resid_d0_t0 = np.multiply(np.divide(d0t0, np.multiply(p_hat, 1.0 - lambda_hat)), prop_weighting) else: - assert self.score == 'experimental' + assert self.score == "experimental" if self.in_sample_normalization: weight_psi_a = np.ones_like(y) weight_g_d1_t1 = weight_psi_a @@ -384,22 +398,26 @@ def _score_elements(self, y, d, t, weight_g_d0_t0 = weight_psi_a weight_resid_d1_t1 = np.divide(d1t1, np.multiply(p_hat, lambda_hat)) - weight_resid_d1_t0 = -1.0 * np.divide(d1t0, np.multiply(p_hat, 1.0-lambda_hat)) - weight_resid_d0_t1 = -1.0 * np.divide(d0t1, np.multiply(1.0-p_hat, lambda_hat)) - weight_resid_d0_t0 = np.divide(d0t0, np.multiply(1.0-p_hat, 1.0-lambda_hat)) + weight_resid_d1_t0 = -1.0 * np.divide(d1t0, np.multiply(p_hat, 1.0 - lambda_hat)) + weight_resid_d0_t1 = -1.0 * np.divide(d0t1, np.multiply(1.0 - p_hat, lambda_hat)) + weight_resid_d0_t0 = np.divide(d0t0, np.multiply(1.0 - p_hat, 1.0 - lambda_hat)) # set score elements psi_a = -1.0 * weight_psi_a # psi_b - psi_b_1 = np.multiply(weight_g_d1_t1, g_hat_d1_t1) + \ - np.multiply(weight_g_d1_t0, g_hat_d1_t0) + \ - np.multiply(weight_g_d0_t0, g_hat_d0_t0) + \ - np.multiply(weight_g_d0_t1, g_hat_d0_t1) - psi_b_2 = np.multiply(weight_resid_d1_t1, resid_d1_t1) + \ - np.multiply(weight_resid_d1_t0, resid_d1_t0) + \ - np.multiply(weight_resid_d0_t0, resid_d0_t0) + \ - np.multiply(weight_resid_d0_t1, resid_d0_t1) + psi_b_1 = ( + np.multiply(weight_g_d1_t1, g_hat_d1_t1) + + np.multiply(weight_g_d1_t0, g_hat_d1_t0) + + np.multiply(weight_g_d0_t0, g_hat_d0_t0) + + np.multiply(weight_g_d0_t1, g_hat_d0_t1) + ) + psi_b_2 = ( + np.multiply(weight_resid_d1_t1, resid_d1_t1) + + np.multiply(weight_resid_d1_t0, resid_d1_t0) + + np.multiply(weight_resid_d0_t0, resid_d0_t0) + + np.multiply(weight_resid_d0_t1, resid_d0_t1) + ) psi_b = psi_b_1 + psi_b_2 @@ -410,19 +428,23 @@ def _sensitivity_element_est(self, preds): d = self._dml_data.d t = self._dml_data.t - m_hat = preds['predictions']['ml_m'] - g_hat_d0_t0 = preds['predictions']['ml_g_d0_t0'] - g_hat_d0_t1 = preds['predictions']['ml_g_d0_t1'] - g_hat_d1_t0 = preds['predictions']['ml_g_d1_t0'] - g_hat_d1_t1 = preds['predictions']['ml_g_d1_t1'] + m_hat = preds["predictions"]["ml_m"] + g_hat_d0_t0 = preds["predictions"]["ml_g_d0_t0"] + g_hat_d0_t1 = preds["predictions"]["ml_g_d0_t1"] + g_hat_d1_t0 = preds["predictions"]["ml_g_d1_t0"] + g_hat_d1_t1 = preds["predictions"]["ml_g_d1_t1"] - d0t0 = np.multiply(1.0-d, 1.0-t) - d0t1 = np.multiply(1.0-d, t) - d1t0 = np.multiply(d, 1.0-t) + d0t0 = np.multiply(1.0 - d, 1.0 - t) + d0t1 = np.multiply(1.0 - d, t) + d1t0 = np.multiply(d, 1.0 - t) d1t1 = np.multiply(d, t) - g_hat = np.multiply(d0t0, g_hat_d0_t0) + np.multiply(d0t1, g_hat_d0_t1) + \ - np.multiply(d1t0, g_hat_d1_t0) + np.multiply(d1t1, g_hat_d1_t1) + g_hat = ( + np.multiply(d0t0, g_hat_d0_t0) + + np.multiply(d0t1, g_hat_d0_t1) + + np.multiply(d1t0, g_hat_d1_t0) + + np.multiply(d1t1, g_hat_d1_t1) + ) sigma2_score_element = np.square(y - g_hat) sigma2 = np.mean(sigma2_score_element) psi_sigma2 = sigma2_score_element - sigma2 @@ -430,74 +452,80 @@ def _sensitivity_element_est(self, preds): # calc m(W,alpha) and Riesz representer p_hat = np.mean(d) lambda_hat = np.mean(t) - if self.score == 'observational': - propensity_weight_d0 = np.divide(m_hat, 1.0-m_hat) + if self.score == "observational": + propensity_weight_d0 = np.divide(m_hat, 1.0 - m_hat) if self.in_sample_normalization: weight_d0t1 = np.multiply(d0t1, propensity_weight_d0) weight_d0t0 = np.multiply(d0t0, propensity_weight_d0) mean_weight_d0t1 = np.mean(weight_d0t1) mean_weight_d0t0 = np.mean(weight_d0t0) - m_alpha = np.multiply(np.divide(d, p_hat), - np.divide(1.0, np.mean(d1t1)) + - np.divide(1.0, np.mean(d1t0)) + - np.divide(propensity_weight_d0, mean_weight_d0t1) + - np.divide(propensity_weight_d0, mean_weight_d0t0)) - - rr = np.divide(d1t1, np.mean(d1t1)) - \ - np.divide(d1t0, np.mean(d1t0)) - \ - np.divide(weight_d0t1, mean_weight_d0t1) + \ - np.divide(weight_d0t0, mean_weight_d0t0) + m_alpha = np.multiply( + np.divide(d, p_hat), + np.divide(1.0, np.mean(d1t1)) + + np.divide(1.0, np.mean(d1t0)) + + np.divide(propensity_weight_d0, mean_weight_d0t1) + + np.divide(propensity_weight_d0, mean_weight_d0t0), + ) + + rr = ( + np.divide(d1t1, np.mean(d1t1)) + - np.divide(d1t0, np.mean(d1t0)) + - np.divide(weight_d0t1, mean_weight_d0t1) + + np.divide(weight_d0t0, mean_weight_d0t0) + ) else: - m_alpha_1 = np.divide(1.0, lambda_hat) + np.divide(1.0, 1.0-lambda_hat) + m_alpha_1 = np.divide(1.0, lambda_hat) + np.divide(1.0, 1.0 - lambda_hat) m_alpha = np.multiply(np.divide(d, np.square(p_hat)), np.multiply(m_alpha_1, 1.0 + propensity_weight_d0)) - rr_1 = np.divide(t, np.multiply(p_hat, lambda_hat)) + np.divide(1.0-t, np.multiply(p_hat, 1.0-lambda_hat)) - rr_2 = d + np.multiply(1.0-d, propensity_weight_d0) + rr_1 = np.divide(t, np.multiply(p_hat, lambda_hat)) + np.divide(1.0 - t, np.multiply(p_hat, 1.0 - lambda_hat)) + rr_2 = d + np.multiply(1.0 - d, propensity_weight_d0) rr = np.multiply(rr_1, rr_2) else: - assert self.score == 'experimental' + assert self.score == "experimental" if self.in_sample_normalization: - m_alpha = np.divide(1.0, np.mean(d1t1)) + \ - np.divide(1.0, np.mean(d1t0)) + \ - np.divide(1.0, np.mean(d0t1)) + \ - np.divide(1.0, np.mean(d0t0)) - rr = np.divide(d1t1, np.mean(d1t1)) - \ - np.divide(d1t0, np.mean(d1t0)) - \ - np.divide(d0t1, np.mean(d0t1)) + \ - np.divide(d0t0, np.mean(d0t0)) + m_alpha = ( + np.divide(1.0, np.mean(d1t1)) + + np.divide(1.0, np.mean(d1t0)) + + np.divide(1.0, np.mean(d0t1)) + + np.divide(1.0, np.mean(d0t0)) + ) + rr = ( + np.divide(d1t1, np.mean(d1t1)) + - np.divide(d1t0, np.mean(d1t0)) + - np.divide(d0t1, np.mean(d0t1)) + + np.divide(d0t0, np.mean(d0t0)) + ) else: - m_alpha = np.divide(1.0, np.multiply(p_hat, lambda_hat)) + \ - np.divide(1.0, np.multiply(p_hat, 1.0-lambda_hat)) + \ - np.divide(1.0, np.multiply(1.0-p_hat, lambda_hat)) + \ - np.divide(1.0, np.multiply(1.0-p_hat, 1.0-lambda_hat)) - rr = np.divide(d1t1, np.multiply(p_hat, lambda_hat)) - \ - np.divide(d1t0, np.multiply(p_hat, 1.0-lambda_hat)) - \ - np.divide(d0t1, np.multiply(1.0-p_hat, lambda_hat)) + \ - np.divide(d0t0, np.multiply(1.0-p_hat, 1.0-lambda_hat)) + m_alpha = ( + np.divide(1.0, np.multiply(p_hat, lambda_hat)) + + np.divide(1.0, np.multiply(p_hat, 1.0 - lambda_hat)) + + np.divide(1.0, np.multiply(1.0 - p_hat, lambda_hat)) + + np.divide(1.0, np.multiply(1.0 - p_hat, 1.0 - lambda_hat)) + ) + rr = ( + np.divide(d1t1, np.multiply(p_hat, lambda_hat)) + - np.divide(d1t0, np.multiply(p_hat, 1.0 - lambda_hat)) + - np.divide(d0t1, np.multiply(1.0 - p_hat, lambda_hat)) + + np.divide(d0t0, np.multiply(1.0 - p_hat, 1.0 - lambda_hat)) + ) nu2_score_element = np.multiply(2.0, m_alpha) - np.square(rr) nu2 = np.mean(nu2_score_element) psi_nu2 = nu2_score_element - nu2 - element_dict = {'sigma2': sigma2, - 'nu2': nu2, - 'psi_sigma2': psi_sigma2, - 'psi_nu2': psi_nu2} + element_dict = {"sigma2": sigma2, "nu2": nu2, "psi_sigma2": psi_sigma2, "psi_nu2": psi_nu2} return element_dict - def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) - x, t = check_X_y(x, self._dml_data.t, - force_all_finite=False) + def _nuisance_tuning( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) + x, t = check_X_y(x, self._dml_data.t, force_all_finite=False) if scoring_methods is None: - scoring_methods = {'ml_g': None, - 'ml_m': None} + scoring_methods = {"ml_g": None, "ml_m": None} # nuisance training sets conditional on d and t smpls_d0_t0, smpls_d0_t1, smpls_d1_t0, smpls_d1_t1 = _get_cond_smpls_2d(smpls, d, t) @@ -507,56 +535,108 @@ def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_ train_inds_d1_t0 = [train_index for (train_index, _) in smpls_d1_t0] train_inds_d1_t1 = [train_index for (train_index, _) in smpls_d1_t1] - g_d0_t0_tune_res = _dml_tune(y, x, train_inds_d0_t0, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - - g_d0_t1_tune_res = _dml_tune(y, x, train_inds_d0_t1, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - - g_d1_t0_tune_res = _dml_tune(y, x, train_inds_d1_t0, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - - g_d1_t1_tune_res = _dml_tune(y, x, train_inds_d1_t1, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + g_d0_t0_tune_res = _dml_tune( + y, + x, + train_inds_d0_t0, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + + g_d0_t1_tune_res = _dml_tune( + y, + x, + train_inds_d0_t1, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + + g_d1_t0_tune_res = _dml_tune( + y, + x, + train_inds_d1_t0, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + + g_d1_t1_tune_res = _dml_tune( + y, + x, + train_inds_d1_t1, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) m_tune_res = list() - if self.score == 'observational': - m_tune_res = _dml_tune(d, x, train_inds, - self._learner['ml_m'], param_grids['ml_m'], scoring_methods['ml_m'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + if self.score == "observational": + m_tune_res = _dml_tune( + d, + x, + train_inds, + self._learner["ml_m"], + param_grids["ml_m"], + scoring_methods["ml_m"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) g_d0_t0_best_params = [xx.best_params_ for xx in g_d0_t0_tune_res] g_d0_t1_best_params = [xx.best_params_ for xx in g_d0_t1_tune_res] g_d1_t0_best_params = [xx.best_params_ for xx in g_d1_t0_tune_res] g_d1_t1_best_params = [xx.best_params_ for xx in g_d1_t1_tune_res] - if self.score == 'observational': + if self.score == "observational": m_best_params = [xx.best_params_ for xx in m_tune_res] - params = {'ml_g_d0_t0': g_d0_t0_best_params, - 'ml_g_d0_t1': g_d0_t1_best_params, - 'ml_g_d1_t0': g_d1_t0_best_params, - 'ml_g_d1_t1': g_d1_t1_best_params, - 'ml_m': m_best_params} - tune_res = {'g_d0_t0_tune': g_d0_t0_tune_res, - 'g_d0_t1_tune': g_d0_t1_tune_res, - 'g_d1_t0_tune': g_d1_t0_tune_res, - 'g_d1_t1_tune': g_d1_t1_tune_res, - 'm_tune': m_tune_res} + params = { + "ml_g_d0_t0": g_d0_t0_best_params, + "ml_g_d0_t1": g_d0_t1_best_params, + "ml_g_d1_t0": g_d1_t0_best_params, + "ml_g_d1_t1": g_d1_t1_best_params, + "ml_m": m_best_params, + } + tune_res = { + "g_d0_t0_tune": g_d0_t0_tune_res, + "g_d0_t1_tune": g_d0_t1_tune_res, + "g_d1_t0_tune": g_d1_t0_tune_res, + "g_d1_t1_tune": g_d1_t1_tune_res, + "m_tune": m_tune_res, + } else: - params = {'ml_g_d0_t0': g_d0_t0_best_params, - 'ml_g_d0_t1': g_d0_t1_best_params, - 'ml_g_d1_t0': g_d1_t0_best_params, - 'ml_g_d1_t1': g_d1_t1_best_params} - tune_res = {'g_d0_t0_tune': g_d0_t0_tune_res, - 'g_d0_t1_tune': g_d0_t1_tune_res, - 'g_d1_t0_tune': g_d1_t0_tune_res, - 'g_d1_t1_tune': g_d1_t1_tune_res} - - res = {'params': params, - 'tune_res': tune_res} + params = { + "ml_g_d0_t0": g_d0_t0_best_params, + "ml_g_d0_t1": g_d0_t1_best_params, + "ml_g_d1_t0": g_d1_t0_best_params, + "ml_g_d1_t1": g_d1_t1_best_params, + } + tune_res = { + "g_d0_t0_tune": g_d0_t0_tune_res, + "g_d0_t1_tune": g_d0_t1_tune_res, + "g_d1_t0_tune": g_d1_t0_tune_res, + "g_d1_t1_tune": g_d1_t1_tune_res, + } + + res = {"params": params, "tune_res": tune_res} return res diff --git a/doubleml/double_ml_iivm.py b/doubleml/double_ml_iivm.py index 5e3cb073..d09c0bf9 100644 --- a/doubleml/double_ml_iivm.py +++ b/doubleml/double_ml_iivm.py @@ -124,76 +124,78 @@ class DoubleMLIIVM(LinearScoreMixin, DoubleML): \\theta_0 = \\frac{\\mathbb{E}[g_0(1, X)] - \\mathbb{E}[g_0(0,X)]}{\\mathbb{E}[r_0(1, X)] - \\mathbb{E}[r_0(0,X)]}. """ - def __init__(self, - obj_dml_data, - ml_g, - ml_m, - ml_r, - n_folds=5, - n_rep=1, - score='LATE', - subgroups=None, - dml_procedure='dml2', - normalize_ipw=False, - trimming_rule='truncate', - trimming_threshold=1e-2, - draw_sample_splitting=True, - apply_cross_fitting=True): - super().__init__(obj_dml_data, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + + def __init__( + self, + obj_dml_data, + ml_g, + ml_m, + ml_r, + n_folds=5, + n_rep=1, + score="LATE", + subgroups=None, + dml_procedure="dml2", + normalize_ipw=False, + trimming_rule="truncate", + trimming_threshold=1e-2, + draw_sample_splitting=True, + apply_cross_fitting=True, + ): + super().__init__(obj_dml_data, n_folds, n_rep, score, dml_procedure, draw_sample_splitting, apply_cross_fitting) self._check_data(self._dml_data) - valid_scores = ['LATE'] + valid_scores = ["LATE"] _check_score(self.score, valid_scores, allow_callable=True) # set stratication for resampling self._strata = self._dml_data.d.reshape(-1, 1) + 2 * self._dml_data.z.reshape(-1, 1) - ml_g_is_classifier = self._check_learner(ml_g, 'ml_g', regressor=True, classifier=True) - _ = self._check_learner(ml_m, 'ml_m', regressor=False, classifier=True) - _ = self._check_learner(ml_r, 'ml_r', regressor=False, classifier=True) - self._learner = {'ml_g': ml_g, 'ml_m': ml_m, 'ml_r': ml_r} + ml_g_is_classifier = self._check_learner(ml_g, "ml_g", regressor=True, classifier=True) + _ = self._check_learner(ml_m, "ml_m", regressor=False, classifier=True) + _ = self._check_learner(ml_r, "ml_r", regressor=False, classifier=True) + self._learner = {"ml_g": ml_g, "ml_m": ml_m, "ml_r": ml_r} self._normalize_ipw = normalize_ipw if ml_g_is_classifier: if obj_dml_data.binary_outcome: - self._predict_method = {'ml_g': 'predict_proba', 'ml_m': 'predict_proba', 'ml_r': 'predict_proba'} + self._predict_method = {"ml_g": "predict_proba", "ml_m": "predict_proba", "ml_r": "predict_proba"} else: - raise ValueError(f'The ml_g learner {str(ml_g)} was identified as classifier ' - 'but the outcome variable is not binary with values 0 and 1.') + raise ValueError( + f"The ml_g learner {str(ml_g)} was identified as classifier " + "but the outcome variable is not binary with values 0 and 1." + ) else: - self._predict_method = {'ml_g': 'predict', 'ml_m': 'predict_proba', 'ml_r': 'predict_proba'} + self._predict_method = {"ml_g": "predict", "ml_m": "predict_proba", "ml_r": "predict_proba"} self._initialize_ml_nuisance_params() if not isinstance(self.normalize_ipw, bool): - raise TypeError('Normalization indicator has to be boolean. ' + - f'Object of type {str(type(self.normalize_ipw))} passed.') + raise TypeError( + "Normalization indicator has to be boolean. " + f"Object of type {str(type(self.normalize_ipw))} passed." + ) self._trimming_rule = trimming_rule self._trimming_threshold = trimming_threshold _check_trimming(self._trimming_rule, self._trimming_threshold) if subgroups is None: # this is the default for subgroups; via None to prevent a mutable default argument - subgroups = {'always_takers': True, 'never_takers': True} + subgroups = {"always_takers": True, "never_takers": True} else: if not isinstance(subgroups, dict): - raise TypeError('Invalid subgroups ' + str(subgroups) + '. ' + - 'subgroups must be of type dictionary.') - if (not all(k in subgroups for k in ['always_takers', 'never_takers']))\ - | (not all(k in ['always_takers', 'never_takers'] for k in subgroups)): - raise ValueError('Invalid subgroups ' + str(subgroups) + '. ' + - 'subgroups must be a dictionary with keys always_takers and never_takers.') - if not isinstance(subgroups['always_takers'], bool): - raise TypeError("subgroups['always_takers'] must be True or False. " - f'Got {str(subgroups["always_takers"])}.') - if not isinstance(subgroups['never_takers'], bool): - raise TypeError("subgroups['never_takers'] must be True or False. " - f'Got {str(subgroups["never_takers"])}.') + raise TypeError("Invalid subgroups " + str(subgroups) + ". " + "subgroups must be of type dictionary.") + if (not all(k in subgroups for k in ["always_takers", "never_takers"])) | ( + not all(k in ["always_takers", "never_takers"] for k in subgroups) + ): + raise ValueError( + "Invalid subgroups " + + str(subgroups) + + ". " + + "subgroups must be a dictionary with keys always_takers and never_takers." + ) + if not isinstance(subgroups["always_takers"], bool): + raise TypeError("subgroups['always_takers'] must be True or False. " f'Got {str(subgroups["always_takers"])}.') + if not isinstance(subgroups["never_takers"], bool): + raise TypeError("subgroups['never_takers'] must be True or False. " f'Got {str(subgroups["never_takers"])}.') self.subgroups = subgroups - + self._external_predictions_implemented = True @property @@ -218,29 +220,33 @@ def trimming_threshold(self): return self._trimming_threshold def _initialize_ml_nuisance_params(self): - valid_learner = ['ml_g0', 'ml_g1', 'ml_m', 'ml_r0', 'ml_r1'] - self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} - for learner in valid_learner} + valid_learner = ["ml_g0", "ml_g1", "ml_m", "ml_r0", "ml_r1"] + self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} for learner in valid_learner} def _check_data(self, obj_dml_data): if not isinstance(obj_dml_data, DoubleMLData): - raise TypeError('The data must be of DoubleMLData type. ' - f'{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed.') - one_treat = (obj_dml_data.n_treat == 1) - binary_treat = (type_of_target(obj_dml_data.d) == 'binary') + raise TypeError( + "The data must be of DoubleMLData type. " f"{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed." + ) + one_treat = obj_dml_data.n_treat == 1 + binary_treat = type_of_target(obj_dml_data.d) == "binary" zero_one_treat = np.all((np.power(obj_dml_data.d, 2) - obj_dml_data.d) == 0) if not (one_treat & binary_treat & zero_one_treat): - raise ValueError('Incompatible data. ' - 'To fit an IIVM model with DML ' - 'exactly one binary variable with values 0 and 1 ' - 'needs to be specified as treatment variable.') - one_instr = (obj_dml_data.n_instr == 1) - err_msg = ('Incompatible data. ' - 'To fit an IIVM model with DML ' - 'exactly one binary variable with values 0 and 1 ' - 'needs to be specified as instrumental variable.') + raise ValueError( + "Incompatible data. " + "To fit an IIVM model with DML " + "exactly one binary variable with values 0 and 1 " + "needs to be specified as treatment variable." + ) + one_instr = obj_dml_data.n_instr == 1 + err_msg = ( + "Incompatible data. " + "To fit an IIVM model with DML " + "exactly one binary variable with values 0 and 1 " + "needs to be specified as instrumental variable." + ) if one_instr: - binary_instr = (type_of_target(obj_dml_data.z) == 'binary') + binary_instr = type_of_target(obj_dml_data.z) == "binary" zero_one_instr = np.all((np.power(obj_dml_data.z, 2) - obj_dml_data.z) == 0) if not (one_instr & binary_instr & zero_one_instr): raise ValueError(err_msg) @@ -249,134 +255,138 @@ def _check_data(self, obj_dml_data): return def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=False): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, z = check_X_y(x, np.ravel(self._dml_data.z), - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, z = check_X_y(x, np.ravel(self._dml_data.z), force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) # get train indices for z == 0 and z == 1 smpls_z0, smpls_z1 = _get_cond_smpls(smpls, z) # nuisance g - if external_predictions['ml_g0'] is not None: - g_hat0 = {'preds': external_predictions['ml_g0'], - 'targets': None, - 'models': None} + if external_predictions["ml_g0"] is not None: + g_hat0 = {"preds": external_predictions["ml_g0"], "targets": None, "models": None} else: g_hat0 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls=smpls_z0, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g0'), method=self._predict_method['ml_g'], - return_models=return_models) + est_params=self._get_params('ml_g0'), method=self._predict_method['ml_g'], + return_models=return_models) _check_finite_predictions(g_hat0['preds'], self._learner['ml_g'], 'ml_g', smpls) # adjust target values to consider only compatible subsamples - g_hat0['targets'] = g_hat0['targets'].astype(float) - g_hat0['targets'][z == 1] = np.nan + g_hat0["targets"] = g_hat0["targets"].astype(float) + g_hat0["targets"][z == 1] = np.nan if self._dml_data.binary_outcome: - binary_preds = (type_of_target(g_hat0['preds']) == 'binary') - zero_one_preds = np.all((np.power(g_hat0['preds'], 2) - g_hat0['preds']) == 0) + binary_preds = type_of_target(g_hat0["preds"]) == "binary" + zero_one_preds = np.all((np.power(g_hat0["preds"], 2) - g_hat0["preds"]) == 0) if binary_preds & zero_one_preds: - raise ValueError(f'For the binary outcome variable {self._dml_data.y_col}, ' - f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also ' - 'observed to be binary with values 0 and 1. Make sure that for classifiers ' - 'probabilities and not labels are predicted.') - - _check_is_propensity(g_hat0['preds'], self._learner['ml_g'], 'ml_g', smpls, eps=1e-12) - if external_predictions['ml_g1'] is not None: - g_hat1 = {'preds': external_predictions['ml_g1'], - 'targets': None, - 'models': None} + raise ValueError( + f"For the binary outcome variable {self._dml_data.y_col}, " + f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also ' + "observed to be binary with values 0 and 1. Make sure that for classifiers " + "probabilities and not labels are predicted." + ) + + _check_is_propensity(g_hat0["preds"], self._learner["ml_g"], "ml_g", smpls, eps=1e-12) + if external_predictions["ml_g1"] is not None: + g_hat1 = {"preds": external_predictions["ml_g1"], "targets": None, "models": None} else: g_hat1 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls=smpls_z1, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g1'), method=self._predict_method['ml_g'], - return_models=return_models) + est_params=self._get_params('ml_g1'), method=self._predict_method['ml_g'], + return_models=return_models) _check_finite_predictions(g_hat1['preds'], self._learner['ml_g'], 'ml_g', smpls) # adjust target values to consider only compatible subsamples - g_hat1['targets'] = g_hat1['targets'].astype(float) - g_hat1['targets'][z == 0] = np.nan + g_hat1["targets"] = g_hat1["targets"].astype(float) + g_hat1["targets"][z == 0] = np.nan if self._dml_data.binary_outcome: - binary_preds = (type_of_target(g_hat1['preds']) == 'binary') - zero_one_preds = np.all((np.power(g_hat1['preds'], 2) - g_hat1['preds']) == 0) + binary_preds = type_of_target(g_hat1["preds"]) == "binary" + zero_one_preds = np.all((np.power(g_hat1["preds"], 2) - g_hat1["preds"]) == 0) if binary_preds & zero_one_preds: - raise ValueError(f'For the binary outcome variable {self._dml_data.y_col}, ' - f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also ' - 'observed to be binary with values 0 and 1. Make sure that for classifiers ' - 'probabilities and not labels are predicted.') + raise ValueError( + f"For the binary outcome variable {self._dml_data.y_col}, " + f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also ' + "observed to be binary with values 0 and 1. Make sure that for classifiers " + "probabilities and not labels are predicted." + ) - _check_is_propensity(g_hat1['preds'], self._learner['ml_g'], 'ml_g', smpls, eps=1e-12) + _check_is_propensity(g_hat1["preds"], self._learner["ml_g"], "ml_g", smpls, eps=1e-12) # nuisance m - if external_predictions['ml_m'] is not None: - m_hat = {'preds': external_predictions['ml_m'], - 'targets': None, - 'models': None} + if external_predictions["ml_m"] is not None: + m_hat = {"preds": external_predictions["ml_m"], "targets": None, "models": None} else: - m_hat = _dml_cv_predict(self._learner['ml_m'], x, z, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_m'), method=self._predict_method['ml_m'], - return_models=return_models) - _check_finite_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls) - _check_is_propensity(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls, eps=1e-12) + m_hat = _dml_cv_predict( + self._learner["ml_m"], + x, + z, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_m"), + method=self._predict_method["ml_m"], + return_models=return_models, + ) + _check_finite_predictions(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls) + _check_is_propensity(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls, eps=1e-12) # nuisance r - r0 = external_predictions['ml_r0'] is not None - if self.subgroups['always_takers']: + r0 = external_predictions["ml_r0"] is not None + if self.subgroups["always_takers"]: if r0: - r_hat0 = {'preds': external_predictions['ml_r0'], - 'targets': None, - 'models': None} + r_hat0 = {"preds": external_predictions["ml_r0"], "targets": None, "models": None} else: r_hat0 = _dml_cv_predict(self._learner['ml_r'], x, d, smpls=smpls_z0, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_r0'), method=self._predict_method['ml_r'], - return_models=return_models) + est_params=self._get_params('ml_r0'), method=self._predict_method['ml_r'], + return_models=return_models) else: - r_hat0 = {'preds': np.zeros_like(d), 'targets': np.zeros_like(d), 'models': None} + r_hat0 = {"preds": np.zeros_like(d), "targets": np.zeros_like(d), "models": None} if not r0: - _check_finite_predictions(r_hat0['preds'], self._learner['ml_r'], 'ml_r', smpls) + _check_finite_predictions(r_hat0["preds"], self._learner["ml_r"], "ml_r", smpls) # adjust target values to consider only compatible subsamples - r_hat0['targets'] = r_hat0['targets'].astype(float) - r_hat0['targets'][z == 1] = np.nan + r_hat0["targets"] = r_hat0["targets"].astype(float) + r_hat0["targets"][z == 1] = np.nan - r1 = external_predictions['ml_r1'] is not None - if self.subgroups['never_takers']: + r1 = external_predictions["ml_r1"] is not None + if self.subgroups["never_takers"]: if r1: - r_hat1 = {'preds': external_predictions['ml_r1'], - 'targets': None, - 'models': None} + r_hat1 = {"preds": external_predictions["ml_r1"], "targets": None, "models": None} else: r_hat1 = _dml_cv_predict(self._learner['ml_r'], x, d, smpls=smpls_z1, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_r1'), method=self._predict_method['ml_r'], - return_models=return_models) + est_params=self._get_params('ml_r1'), method=self._predict_method['ml_r'], + return_models=return_models) else: - r_hat1 = {'preds': np.ones_like(d), 'targets': np.ones_like(d), 'models': None} + r_hat1 = {"preds": np.ones_like(d), "targets": np.ones_like(d), "models": None} if not r1: - _check_finite_predictions(r_hat1['preds'], self._learner['ml_r'], 'ml_r', smpls) + _check_finite_predictions(r_hat1["preds"], self._learner["ml_r"], "ml_r", smpls) # adjust target values to consider only compatible subsamples - r_hat1['targets'] = r_hat1['targets'].astype(float) - r_hat1['targets'][z == 0] = np.nan - - psi_a, psi_b = self._score_elements(y, z, d, - g_hat0['preds'], g_hat1['preds'], m_hat['preds'], - r_hat0['preds'], r_hat1['preds'], smpls) - psi_elements = {'psi_a': psi_a, - 'psi_b': psi_b} - preds = {'predictions': {'ml_g0': g_hat0['preds'], - 'ml_g1': g_hat1['preds'], - 'ml_m': m_hat['preds'], - 'ml_r0': r_hat0['preds'], - 'ml_r1': r_hat1['preds']}, - 'targets': {'ml_g0': g_hat0['targets'], - 'ml_g1': g_hat1['targets'], - 'ml_m': m_hat['targets'], - 'ml_r0': r_hat0['targets'], - 'ml_r1': r_hat1['targets']}, - 'models': {'ml_g0': g_hat0['models'], - 'ml_g1': g_hat1['models'], - 'ml_m': m_hat['models'], - 'ml_r0': r_hat0['models'], - 'ml_r1': r_hat1['models']} - } + r_hat1["targets"] = r_hat1["targets"].astype(float) + r_hat1["targets"][z == 0] = np.nan + + psi_a, psi_b = self._score_elements( + y, z, d, g_hat0["preds"], g_hat1["preds"], m_hat["preds"], r_hat0["preds"], r_hat1["preds"], smpls + ) + psi_elements = {"psi_a": psi_a, "psi_b": psi_b} + preds = { + "predictions": { + "ml_g0": g_hat0["preds"], + "ml_g1": g_hat1["preds"], + "ml_m": m_hat["preds"], + "ml_r0": r_hat0["preds"], + "ml_r1": r_hat1["preds"], + }, + "targets": { + "ml_g0": g_hat0["targets"], + "ml_g1": g_hat1["targets"], + "ml_m": m_hat["targets"], + "ml_r0": r_hat0["targets"], + "ml_r1": r_hat1["targets"], + }, + "models": { + "ml_g0": g_hat0["models"], + "ml_g1": g_hat1["models"], + "ml_m": m_hat["models"], + "ml_r0": r_hat0["models"], + "ml_r1": r_hat1["models"], + }, + } return psi_elements, preds @@ -390,71 +400,118 @@ def _score_elements(self, y, z, d, g_hat0, g_hat1, m_hat, r_hat0, r_hat1, smpls) m_hat = _trimm(m_hat, self.trimming_rule, self.trimming_threshold) if self.normalize_ipw: - if self.dml_procedure == 'dml1': + if self.dml_procedure == "dml1": for _, test_index in smpls: m_hat[test_index] = _normalize_ipw(m_hat[test_index], d[test_index]) else: m_hat = _normalize_ipw(m_hat, d) if isinstance(self.score, str): - assert self.score == 'LATE' - psi_b = g_hat1 - g_hat0 \ - + np.divide(np.multiply(z, u_hat1), m_hat) \ - - np.divide(np.multiply(1.0-z, u_hat0), 1.0 - m_hat) - psi_a = -1*(r_hat1 - r_hat0 - + np.divide(np.multiply(z, w_hat1), m_hat) - - np.divide(np.multiply(1.0-z, w_hat0), 1.0 - m_hat)) + assert self.score == "LATE" + psi_b = ( + g_hat1 + - g_hat0 + + np.divide(np.multiply(z, u_hat1), m_hat) + - np.divide(np.multiply(1.0 - z, u_hat0), 1.0 - m_hat) + ) + psi_a = -1 * ( + r_hat1 + - r_hat0 + + np.divide(np.multiply(z, w_hat1), m_hat) + - np.divide(np.multiply(1.0 - z, w_hat0), 1.0 - m_hat) + ) else: assert callable(self.score) - psi_a, psi_b = self.score(y=y, z=z, d=d, - g_hat0=g_hat0, g_hat1=g_hat1, m_hat=m_hat, r_hat0=r_hat0, r_hat1=r_hat1, - smpls=smpls) + psi_a, psi_b = self.score( + y=y, z=z, d=d, g_hat0=g_hat0, g_hat1=g_hat1, m_hat=m_hat, r_hat0=r_hat0, r_hat1=r_hat1, smpls=smpls + ) return psi_a, psi_b - def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, z = check_X_y(x, np.ravel(self._dml_data.z), - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + def _nuisance_tuning( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, z = check_X_y(x, np.ravel(self._dml_data.z), force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) # get train indices for z == 0 and z == 1 smpls_z0, smpls_z1 = _get_cond_smpls(smpls, z) if scoring_methods is None: - scoring_methods = {'ml_g': None, - 'ml_m': None, - 'ml_r': None} + scoring_methods = {"ml_g": None, "ml_m": None, "ml_r": None} train_inds = [train_index for (train_index, _) in smpls] train_inds_z0 = [train_index for (train_index, _) in smpls_z0] train_inds_z1 = [train_index for (train_index, _) in smpls_z1] - g0_tune_res = _dml_tune(y, x, train_inds_z0, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - g1_tune_res = _dml_tune(y, x, train_inds_z1, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - m_tune_res = _dml_tune(z, x, train_inds, - self._learner['ml_m'], param_grids['ml_m'], scoring_methods['ml_m'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - - if self.subgroups['always_takers']: - r0_tune_res = _dml_tune(d, x, train_inds_z0, - self._learner['ml_r'], param_grids['ml_r'], scoring_methods['ml_r'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + g0_tune_res = _dml_tune( + y, + x, + train_inds_z0, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + g1_tune_res = _dml_tune( + y, + x, + train_inds_z1, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + m_tune_res = _dml_tune( + z, + x, + train_inds, + self._learner["ml_m"], + param_grids["ml_m"], + scoring_methods["ml_m"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + + if self.subgroups["always_takers"]: + r0_tune_res = _dml_tune( + d, + x, + train_inds_z0, + self._learner["ml_r"], + param_grids["ml_r"], + scoring_methods["ml_r"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) r0_best_params = [xx.best_params_ for xx in r0_tune_res] else: r0_tune_res = None r0_best_params = [None] * len(smpls) - if self.subgroups['never_takers']: - r1_tune_res = _dml_tune(d, x, train_inds_z1, - self._learner['ml_r'], param_grids['ml_r'], scoring_methods['ml_r'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + if self.subgroups["never_takers"]: + r1_tune_res = _dml_tune( + d, + x, + train_inds_z1, + self._learner["ml_r"], + param_grids["ml_r"], + scoring_methods["ml_r"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) r1_best_params = [xx.best_params_ for xx in r1_tune_res] else: r1_tune_res = None @@ -464,20 +521,23 @@ def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_ g1_best_params = [xx.best_params_ for xx in g1_tune_res] m_best_params = [xx.best_params_ for xx in m_tune_res] - params = {'ml_g0': g0_best_params, - 'ml_g1': g1_best_params, - 'ml_m': m_best_params, - 'ml_r0': r0_best_params, - 'ml_r1': r1_best_params} - - tune_res = {'g0_tune': g0_tune_res, - 'g1_tune': g1_tune_res, - 'm_tune': m_tune_res, - 'r0_tune': r0_tune_res, - 'r1_tune': r1_tune_res} - - res = {'params': params, - 'tune_res': tune_res} + params = { + "ml_g0": g0_best_params, + "ml_g1": g1_best_params, + "ml_m": m_best_params, + "ml_r0": r0_best_params, + "ml_r1": r1_best_params, + } + + tune_res = { + "g0_tune": g0_tune_res, + "g1_tune": g1_tune_res, + "m_tune": m_tune_res, + "r0_tune": r0_tune_res, + "r1_tune": r1_tune_res, + } + + res = {"params": params, "tune_res": tune_res} return res diff --git a/doubleml/double_ml_irm.py b/doubleml/double_ml_irm.py index 0c049b66..78815410 100644 --- a/doubleml/double_ml_irm.py +++ b/doubleml/double_ml_irm.py @@ -112,56 +112,56 @@ class DoubleMLIRM(LinearScoreMixin, DoubleML): \\theta_0 = \\mathbb{E}[g_0(1, X) - g_0(0,X) | D=1]. """ - def __init__(self, - obj_dml_data, - ml_g, - ml_m, - n_folds=5, - n_rep=1, - score='ATE', - dml_procedure='dml2', - normalize_ipw=False, - trimming_rule='truncate', - trimming_threshold=1e-2, - draw_sample_splitting=True, - apply_cross_fitting=True): - super().__init__(obj_dml_data, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + + def __init__( + self, + obj_dml_data, + ml_g, + ml_m, + n_folds=5, + n_rep=1, + score="ATE", + dml_procedure="dml2", + normalize_ipw=False, + trimming_rule="truncate", + trimming_threshold=1e-2, + draw_sample_splitting=True, + apply_cross_fitting=True, + ): + super().__init__(obj_dml_data, n_folds, n_rep, score, dml_procedure, draw_sample_splitting, apply_cross_fitting) self._check_data(self._dml_data) - valid_scores = ['ATE', 'ATTE'] + valid_scores = ["ATE", "ATTE"] _check_score(self.score, valid_scores, allow_callable=True) # set stratication for resampling self._strata = self._dml_data.d - ml_g_is_classifier = self._check_learner(ml_g, 'ml_g', regressor=True, classifier=True) - _ = self._check_learner(ml_m, 'ml_m', regressor=False, classifier=True) - self._learner = {'ml_g': ml_g, 'ml_m': ml_m} + ml_g_is_classifier = self._check_learner(ml_g, "ml_g", regressor=True, classifier=True) + _ = self._check_learner(ml_m, "ml_m", regressor=False, classifier=True) + self._learner = {"ml_g": ml_g, "ml_m": ml_m} self._normalize_ipw = normalize_ipw if ml_g_is_classifier: if obj_dml_data.binary_outcome: - self._predict_method = {'ml_g': 'predict_proba', 'ml_m': 'predict_proba'} + self._predict_method = {"ml_g": "predict_proba", "ml_m": "predict_proba"} else: - raise ValueError(f'The ml_g learner {str(ml_g)} was identified as classifier ' - 'but the outcome variable is not binary with values 0 and 1.') + raise ValueError( + f"The ml_g learner {str(ml_g)} was identified as classifier " + "but the outcome variable is not binary with values 0 and 1." + ) else: - self._predict_method = {'ml_g': 'predict', 'ml_m': 'predict_proba'} + self._predict_method = {"ml_g": "predict", "ml_m": "predict_proba"} self._initialize_ml_nuisance_params() if not isinstance(self.normalize_ipw, bool): - raise TypeError('Normalization indicator has to be boolean. ' + - f'Object of type {str(type(self.normalize_ipw))} passed.') + raise TypeError( + "Normalization indicator has to be boolean. " + f"Object of type {str(type(self.normalize_ipw))} passed." + ) self._trimming_rule = trimming_rule self._trimming_threshold = trimming_threshold _check_trimming(self._trimming_rule, self._trimming_threshold) self._sensitivity_implemented = True - + self._external_predictions_implemented = True @property @@ -186,37 +186,36 @@ def trimming_threshold(self): return self._trimming_threshold def _initialize_ml_nuisance_params(self): - valid_learner = ['ml_g0', 'ml_g1', 'ml_m'] - self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} - for learner in valid_learner} + valid_learner = ["ml_g0", "ml_g1", "ml_m"] + self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} for learner in valid_learner} def _check_data(self, obj_dml_data): if not isinstance(obj_dml_data, DoubleMLData): - raise TypeError('The data must be of DoubleMLData type. ' - f'{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed.') + raise TypeError( + "The data must be of DoubleMLData type. " f"{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed." + ) if obj_dml_data.z_cols is not None: - raise ValueError('Incompatible data. ' + - ' and '.join(obj_dml_data.z_cols) + - ' have been set as instrumental variable(s). ' - 'To fit an interactive IV regression model use DoubleMLIIVM instead of DoubleMLIRM.') - one_treat = (obj_dml_data.n_treat == 1) - binary_treat = (type_of_target(obj_dml_data.d) == 'binary') + raise ValueError( + "Incompatible data. " + " and ".join(obj_dml_data.z_cols) + " have been set as instrumental variable(s). " + "To fit an interactive IV regression model use DoubleMLIIVM instead of DoubleMLIRM." + ) + one_treat = obj_dml_data.n_treat == 1 + binary_treat = type_of_target(obj_dml_data.d) == "binary" zero_one_treat = np.all((np.power(obj_dml_data.d, 2) - obj_dml_data.d) == 0) if not (one_treat & binary_treat & zero_one_treat): - raise ValueError('Incompatible data. ' - 'To fit an IRM model with DML ' - 'exactly one binary variable with values 0 and 1 ' - 'needs to be specified as treatment variable.') + raise ValueError( + "Incompatible data. " + "To fit an IRM model with DML " + "exactly one binary variable with values 0 and 1 " + "needs to be specified as treatment variable." + ) return def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=False): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) # get train indices for d == 0 and d == 1 smpls_d0, smpls_d1 = _get_cond_smpls(smpls, d) - g0_external = external_predictions['ml_g0'] is not None g1_external = external_predictions['ml_g1'] is not None m_external = external_predictions['ml_m'] is not None @@ -224,91 +223,94 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa # nuisance g if g0_external: # use external predictions - g_hat0 = {'preds': external_predictions['ml_g0'], - 'targets': None, - 'models': None} + g_hat0 = {"preds": external_predictions["ml_g0"], "targets": None, "models": None} else: - g_hat0 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls=smpls_d0, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g0'), method=self._predict_method['ml_g'], - return_models=return_models) - _check_finite_predictions(g_hat0['preds'], self._learner['ml_g'], 'ml_g', smpls) - g_hat0['targets'] = _cond_targets(g_hat0['targets'], cond_sample=(d == 0)) + g_hat0 = _dml_cv_predict( + self._learner["ml_g"], + x, + y, + smpls=smpls_d0, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_g0"), + method=self._predict_method["ml_g"], + return_models=return_models, + ) + _check_finite_predictions(g_hat0["preds"], self._learner["ml_g"], "ml_g", smpls) + g_hat0["targets"] = _cond_targets(g_hat0["targets"], cond_sample=(d == 0)) if self._dml_data.binary_outcome: - binary_preds = (type_of_target(g_hat0['preds']) == 'binary') - zero_one_preds = np.all((np.power(g_hat0['preds'], 2) - g_hat0['preds']) == 0) + binary_preds = type_of_target(g_hat0["preds"]) == "binary" + zero_one_preds = np.all((np.power(g_hat0["preds"], 2) - g_hat0["preds"]) == 0) if binary_preds & zero_one_preds: - raise ValueError(f'For the binary outcome variable {self._dml_data.y_col}, ' - f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also ' - 'observed to be binary with values 0 and 1. Make sure that for classifiers ' - 'probabilities and not labels are predicted.') + raise ValueError( + f"For the binary outcome variable {self._dml_data.y_col}, " + f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also ' + "observed to be binary with values 0 and 1. Make sure that for classifiers " + "probabilities and not labels are predicted." + ) if g1_external: # use external predictions - g_hat1 = {'preds': external_predictions['ml_g1'], - 'targets': None, - 'models': None} + g_hat1 = {"preds": external_predictions["ml_g1"], "targets": None, "models": None} else: g_hat1 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls=smpls_d1, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g1'), method=self._predict_method['ml_g'], - return_models=return_models) + est_params=self._get_params('ml_g1'), method=self._predict_method['ml_g'], + return_models=return_models) _check_finite_predictions(g_hat1['preds'], self._learner['ml_g'], 'ml_g', smpls) # adjust target values to consider only compatible subsamples - g_hat1['targets'] = _cond_targets(g_hat1['targets'], cond_sample=(d == 1)) + g_hat1["targets"] = _cond_targets(g_hat1["targets"], cond_sample=(d == 1)) if self._dml_data.binary_outcome: - binary_preds = (type_of_target(g_hat1['preds']) == 'binary') - zero_one_preds = np.all((np.power(g_hat1['preds'], 2) - g_hat1['preds']) == 0) + binary_preds = type_of_target(g_hat1["preds"]) == "binary" + zero_one_preds = np.all((np.power(g_hat1["preds"], 2) - g_hat1["preds"]) == 0) if binary_preds & zero_one_preds: - raise ValueError(f'For the binary outcome variable {self._dml_data.y_col}, ' - f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also ' - 'observed to be binary with values 0 and 1. Make sure that for classifiers ' - 'probabilities and not labels are predicted.') + raise ValueError( + f"For the binary outcome variable {self._dml_data.y_col}, " + f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also ' + "observed to be binary with values 0 and 1. Make sure that for classifiers " + "probabilities and not labels are predicted." + ) # nuisance m if m_external: # use external predictions - m_hat = {'preds': external_predictions['ml_m'], - 'targets': None, - 'models': None} + m_hat = {"preds": external_predictions["ml_m"], "targets": None, "models": None} else: - m_hat = _dml_cv_predict(self._learner['ml_m'], x, d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_m'), method=self._predict_method['ml_m'], - return_models=return_models) - _check_finite_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls) - _check_is_propensity(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls, eps=1e-12) - m_hat['preds'] = _trimm(m_hat['preds'], self.trimming_rule, self.trimming_threshold) - - psi_a, psi_b = self._score_elements(y, d, - g_hat0['preds'], g_hat1['preds'], m_hat['preds'], - smpls) - psi_elements = {'psi_a': psi_a, - 'psi_b': psi_b} - preds = {'predictions': {'ml_g0': g_hat0['preds'], - 'ml_g1': g_hat1['preds'], - 'ml_m': m_hat['preds']}, - 'targets': {'ml_g0': g_hat0['targets'], - 'ml_g1': g_hat1['targets'], - 'ml_m': m_hat['targets']}, - 'models': {'ml_g0': g_hat0['models'], - 'ml_g1': g_hat1['models'], - 'ml_m': m_hat['models']} - } + m_hat = _dml_cv_predict( + self._learner["ml_m"], + x, + d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_m"), + method=self._predict_method["ml_m"], + return_models=return_models, + ) + _check_finite_predictions(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls) + _check_is_propensity(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls, eps=1e-12) + m_hat["preds"] = _trimm(m_hat["preds"], self.trimming_rule, self.trimming_threshold) + + psi_a, psi_b = self._score_elements(y, d, g_hat0["preds"], g_hat1["preds"], m_hat["preds"], smpls) + psi_elements = {"psi_a": psi_a, "psi_b": psi_b} + preds = { + "predictions": {"ml_g0": g_hat0["preds"], "ml_g1": g_hat1["preds"], "ml_m": m_hat["preds"]}, + "targets": {"ml_g0": g_hat0["targets"], "ml_g1": g_hat1["targets"], "ml_m": m_hat["targets"]}, + "models": {"ml_g0": g_hat0["models"], "ml_g1": g_hat1["models"], "ml_m": m_hat["models"]}, + } return psi_elements, preds def _score_elements(self, y, d, g_hat0, g_hat1, m_hat, smpls): - # fraction of treated for ATTE p_hat = None - if self.score == 'ATTE': - p_hat = np.full_like(d, np.nan, dtype='float64') + if self.score == "ATTE": + p_hat = np.full_like(d, np.nan, dtype="float64") for _, test_index in smpls: p_hat[test_index] = np.mean(d[test_index]) m_hat_adj = copy.deepcopy(m_hat) if self.normalize_ipw: - if self.dml_procedure == 'dml1': + if self.dml_procedure == "dml1": for _, test_index in smpls: m_hat_adj[test_index] = _normalize_ipw(m_hat[test_index], d[test_index]) else: @@ -317,26 +319,27 @@ def _score_elements(self, y, d, g_hat0, g_hat1, m_hat, smpls): # compute residuals u_hat0 = y - g_hat0 u_hat1 = None - if self.score == 'ATE': + if self.score == "ATE": u_hat1 = y - g_hat1 if isinstance(self.score, str): - if self.score == 'ATE': - psi_b = g_hat1 - g_hat0 \ - + np.divide(np.multiply(d, u_hat1), m_hat_adj) \ - - np.divide(np.multiply(1.0-d, u_hat0), 1.0 - m_hat_adj) + if self.score == "ATE": + psi_b = ( + g_hat1 + - g_hat0 + + np.divide(np.multiply(d, u_hat1), m_hat_adj) + - np.divide(np.multiply(1.0 - d, u_hat0), 1.0 - m_hat_adj) + ) psi_a = np.full_like(m_hat_adj, -1.0) else: - assert self.score == 'ATTE' - psi_b = np.divide(np.multiply(d, u_hat0), p_hat) \ - - np.divide(np.multiply(m_hat_adj, np.multiply(1.0-d, u_hat0)), - np.multiply(p_hat, (1.0 - m_hat_adj))) - psi_a = - np.divide(d, p_hat) + assert self.score == "ATTE" + psi_b = np.divide(np.multiply(d, u_hat0), p_hat) - np.divide( + np.multiply(m_hat_adj, np.multiply(1.0 - d, u_hat0)), np.multiply(p_hat, (1.0 - m_hat_adj)) + ) + psi_a = -np.divide(d, p_hat) else: assert callable(self.score) - psi_a, psi_b = self.score(y=y, d=d, - g_hat0=g_hat0, g_hat1=g_hat1, m_hat=m_hat_adj, - smpls=smpls) + psi_a, psi_b = self.score(y=y, d=d, g_hat0=g_hat0, g_hat1=g_hat1, m_hat=m_hat_adj, smpls=smpls) return psi_a, psi_b @@ -345,78 +348,95 @@ def _sensitivity_element_est(self, preds): y = self._dml_data.y d = self._dml_data.d - m_hat = preds['predictions']['ml_m'] - g_hat0 = preds['predictions']['ml_g0'] - g_hat1 = preds['predictions']['ml_g1'] + m_hat = preds["predictions"]["ml_m"] + g_hat0 = preds["predictions"]["ml_g0"] + g_hat1 = preds["predictions"]["ml_g1"] # use weights make this extendable - if self.score == 'ATE': + if self.score == "ATE": weights = np.ones_like(d) weights_bar = np.ones_like(d) else: - assert self.score == 'ATTE' + assert self.score == "ATTE" weights = np.divide(d, np.mean(d)) weights_bar = np.divide(m_hat, np.mean(d)) - sigma2_score_element = np.square(y - np.multiply(d, g_hat1) - np.multiply(1.0-d, g_hat0)) + sigma2_score_element = np.square(y - np.multiply(d, g_hat1) - np.multiply(1.0 - d, g_hat0)) sigma2 = np.mean(sigma2_score_element) psi_sigma2 = sigma2_score_element - sigma2 # calc m(W,alpha) and Riesz representer - m_alpha = np.multiply(weights, np.multiply(weights_bar, (np.divide(1.0, m_hat) + np.divide(1.0, 1.0-m_hat)))) - rr = np.multiply(weights_bar, (np.divide(d, m_hat) - np.divide(1.0-d, 1.0-m_hat))) + m_alpha = np.multiply(weights, np.multiply(weights_bar, (np.divide(1.0, m_hat) + np.divide(1.0, 1.0 - m_hat)))) + rr = np.multiply(weights_bar, (np.divide(d, m_hat) - np.divide(1.0 - d, 1.0 - m_hat))) nu2_score_element = np.multiply(2.0, m_alpha) - np.square(rr) nu2 = np.mean(nu2_score_element) psi_nu2 = nu2_score_element - nu2 - element_dict = {'sigma2': sigma2, - 'nu2': nu2, - 'psi_sigma2': psi_sigma2, - 'psi_nu2': psi_nu2} + element_dict = {"sigma2": sigma2, "nu2": nu2, "psi_sigma2": psi_sigma2, "psi_nu2": psi_nu2} return element_dict - def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + def _nuisance_tuning( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) # get train indices for d == 0 and d == 1 smpls_d0, smpls_d1 = _get_cond_smpls(smpls, d) if scoring_methods is None: - scoring_methods = {'ml_g': None, - 'ml_m': None} + scoring_methods = {"ml_g": None, "ml_m": None} train_inds = [train_index for (train_index, _) in smpls] train_inds_d0 = [train_index for (train_index, _) in smpls_d0] train_inds_d1 = [train_index for (train_index, _) in smpls_d1] - g0_tune_res = _dml_tune(y, x, train_inds_d0, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + g0_tune_res = _dml_tune( + y, + x, + train_inds_d0, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) g1_tune_res = list() - g1_tune_res = _dml_tune(y, x, train_inds_d1, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - - m_tune_res = _dml_tune(d, x, train_inds, - self._learner['ml_m'], param_grids['ml_m'], scoring_methods['ml_m'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + g1_tune_res = _dml_tune( + y, + x, + train_inds_d1, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + + m_tune_res = _dml_tune( + d, + x, + train_inds, + self._learner["ml_m"], + param_grids["ml_m"], + scoring_methods["ml_m"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) g0_best_params = [xx.best_params_ for xx in g0_tune_res] g1_best_params = [xx.best_params_ for xx in g1_tune_res] m_best_params = [xx.best_params_ for xx in m_tune_res] - params = {'ml_g0': g0_best_params, - 'ml_g1': g1_best_params, - 'ml_m': m_best_params} - tune_res = {'g0_tune': g0_tune_res, - 'g1_tune': g1_tune_res, - 'm_tune': m_tune_res} + params = {"ml_g0": g0_best_params, "ml_g1": g1_best_params, "ml_m": m_best_params} + tune_res = {"g0_tune": g0_tune_res, "g1_tune": g1_tune_res, "m_tune": m_tune_res} - res = {'params': params, - 'tune_res': tune_res} + res = {"params": params, "tune_res": tune_res} return res @@ -435,17 +455,15 @@ def cate(self, basis): model : :class:`doubleML.DoubleMLBLP` Best linear Predictor model. """ - valid_score = ['ATE'] + valid_score = ["ATE"] if self.score not in valid_score: - raise ValueError('Invalid score ' + self.score + '. ' + - 'Valid score ' + ' or '.join(valid_score) + '.') + raise ValueError("Invalid score " + self.score + ". " + "Valid score " + " or ".join(valid_score) + ".") if self.n_rep != 1: - raise NotImplementedError('Only implemented for one repetition. ' + - f'Number of repetitions is {str(self.n_rep)}.') + raise NotImplementedError("Only implemented for one repetition. " + f"Number of repetitions is {str(self.n_rep)}.") # define the orthogonal signal - orth_signal = self.psi_elements['psi_b'].reshape(-1) + orth_signal = self.psi_elements["psi_b"].reshape(-1) # fit the best linear predictor model = DoubleMLBLP(orth_signal, basis=basis).fit() @@ -467,31 +485,30 @@ def gate(self, groups): model : :class:`doubleML.DoubleMLBLPGATE` Best linear Predictor model for Group Effects. """ - valid_score = ['ATE'] + valid_score = ["ATE"] if self.score not in valid_score: - raise ValueError('Invalid score ' + self.score + '. ' + - 'Valid score ' + ' or '.join(valid_score) + '.') + raise ValueError("Invalid score " + self.score + ". " + "Valid score " + " or ".join(valid_score) + ".") if self.n_rep != 1: - raise NotImplementedError('Only implemented for one repetition. ' + - f'Number of repetitions is {str(self.n_rep)}.') + raise NotImplementedError("Only implemented for one repetition. " + f"Number of repetitions is {str(self.n_rep)}.") if not isinstance(groups, pd.DataFrame): - raise TypeError('Groups must be of DataFrame type. ' - f'Groups of type {str(type(groups))} was passed.') + raise TypeError("Groups must be of DataFrame type. " f"Groups of type {str(type(groups))} was passed.") if not all(groups.dtypes == bool) or all(groups.dtypes == int): if groups.shape[1] == 1: - groups = pd.get_dummies(groups, prefix='Group', prefix_sep='_') + groups = pd.get_dummies(groups, prefix="Group", prefix_sep="_") else: - raise TypeError('Columns of groups must be of bool type or int type (dummy coded). ' - 'Alternatively, groups should only contain one column.') + raise TypeError( + "Columns of groups must be of bool type or int type (dummy coded). " + "Alternatively, groups should only contain one column." + ) if any(groups.sum(0) <= 5): - warnings.warn('At least one group effect is estimated with less than 6 observations.') + warnings.warn("At least one group effect is estimated with less than 6 observations.") # define the orthogonal signal - orth_signal = self.psi_elements['psi_b'].reshape(-1) + orth_signal = self.psi_elements["psi_b"].reshape(-1) # fit the best linear predictor for GATE (different confint() method) model = DoubleMLBLP(orth_signal, basis=groups, is_gate=True).fit() @@ -522,22 +539,19 @@ def policy_tree(self, features, depth=2, **tree_params): model : :class:`doubleML.DoubleMLPolicyTree` Policy tree model. """ - valid_score = ['ATE'] + valid_score = ["ATE"] if self.score not in valid_score: - raise ValueError('Invalid score ' + self.score + '. ' + - 'Valid score ' + ' or '.join(valid_score) + '.') + raise ValueError("Invalid score " + self.score + ". " + "Valid score " + " or ".join(valid_score) + ".") if self.n_rep != 1: - raise NotImplementedError('Only implemented for one repetition. ' + - f'Number of repetitions is {str(self.n_rep)}.') + raise NotImplementedError("Only implemented for one repetition. " + f"Number of repetitions is {str(self.n_rep)}.") _check_integer(depth, "Depth", 0) if not isinstance(features, pd.DataFrame): - raise TypeError('Covariates must be of DataFrame type. ' - f'Covariates of type {str(type(features))} was passed.') + raise TypeError("Covariates must be of DataFrame type. " f"Covariates of type {str(type(features))} was passed.") - orth_signal = self.psi_elements['psi_b'].reshape(-1) + orth_signal = self.psi_elements["psi_b"].reshape(-1) model = DoubleMLPolicyTree(orth_signal, depth=depth, features=features, **tree_params).fit() diff --git a/doubleml/double_ml_lpq.py b/doubleml/double_ml_lpq.py index 6efd375f..2faea0fa 100644 --- a/doubleml/double_ml_lpq.py +++ b/doubleml/double_ml_lpq.py @@ -188,7 +188,7 @@ def __init__( stratify=strata, ) self._smpls = obj_dml_resampling.split_samples() - + self._external_predictions_implemented = True @property @@ -385,9 +385,9 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa # preliminary propensity for z ml_m_z_prelim = clone(fitted_models["ml_m_z"][i_fold]) - m_z_hat_prelim = _dml_cv_predict(ml_m_z_prelim, x_train_1, z_train_1, method="predict_proba", smpls=smpls_prelim)[ - "preds" - ] + m_z_hat_prelim = _dml_cv_predict( + ml_m_z_prelim, x_train_1, z_train_1, method="predict_proba", smpls=smpls_prelim + )["preds"] m_z_hat_prelim = _trimm(m_z_hat_prelim, self.trimming_rule, self.trimming_threshold) if self._normalize_ipw: diff --git a/doubleml/double_ml_pliv.py b/doubleml/double_ml_pliv.py index b7f6259c..43e3d701 100644 --- a/doubleml/double_ml_pliv.py +++ b/doubleml/double_ml_pliv.py @@ -103,186 +103,197 @@ class DoubleMLPLIV(LinearScoreMixin, DoubleML): :math:`V` are stochastic errors. """ - def __init__(self, - obj_dml_data, - ml_l, - ml_m, - ml_r, - ml_g=None, - n_folds=5, - n_rep=1, - score='partialling out', - dml_procedure='dml2', - draw_sample_splitting=True, - apply_cross_fitting=True): - super().__init__(obj_dml_data, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + def __init__( + self, + obj_dml_data, + ml_l, + ml_m, + ml_r, + ml_g=None, + n_folds=5, + n_rep=1, + score="partialling out", + dml_procedure="dml2", + draw_sample_splitting=True, + apply_cross_fitting=True, + ): + super().__init__(obj_dml_data, n_folds, n_rep, score, dml_procedure, draw_sample_splitting, apply_cross_fitting) self._check_data(self._dml_data) self.partialX = True self.partialZ = False self._check_score(self.score) - _ = self._check_learner(ml_l, 'ml_l', regressor=True, classifier=False) - _ = self._check_learner(ml_m, 'ml_m', regressor=True, classifier=False) - _ = self._check_learner(ml_r, 'ml_r', regressor=True, classifier=False) - self._learner = {'ml_l': ml_l, 'ml_m': ml_m, 'ml_r': ml_r} + _ = self._check_learner(ml_l, "ml_l", regressor=True, classifier=False) + _ = self._check_learner(ml_m, "ml_m", regressor=True, classifier=False) + _ = self._check_learner(ml_r, "ml_r", regressor=True, classifier=False) + self._learner = {"ml_l": ml_l, "ml_m": ml_m, "ml_r": ml_r} if ml_g is not None: - if (isinstance(self.score, str) & (self.score == 'IV-type')) | callable(self.score): - _ = self._check_learner(ml_g, 'ml_g', regressor=True, classifier=False) - self._learner['ml_g'] = ml_g + if (isinstance(self.score, str) & (self.score == "IV-type")) | callable(self.score): + _ = self._check_learner(ml_g, "ml_g", regressor=True, classifier=False) + self._learner["ml_g"] = ml_g else: - assert (isinstance(self.score, str) & (self.score == 'partialling out')) - warnings.warn(('A learner ml_g has been provided for score = "partialling out" but will be ignored. "' - 'A learner ml_g is not required for estimation.')) - elif isinstance(self.score, str) & (self.score == 'IV-type'): + assert isinstance(self.score, str) & (self.score == "partialling out") + warnings.warn( + ( + 'A learner ml_g has been provided for score = "partialling out" but will be ignored. "' + "A learner ml_g is not required for estimation." + ) + ) + elif isinstance(self.score, str) & (self.score == "IV-type"): raise ValueError("For score = 'IV-type', learners ml_l, ml_m, ml_r and ml_g need to be specified.") - self._predict_method = {'ml_l': 'predict', 'ml_m': 'predict', 'ml_r': 'predict'} - if 'ml_g' in self._learner: - self._predict_method['ml_g'] = 'predict' + self._predict_method = {"ml_l": "predict", "ml_m": "predict", "ml_r": "predict"} + if "ml_g" in self._learner: + self._predict_method["ml_g"] = "predict" self._initialize_ml_nuisance_params() - + self._external_predictions_implemented = True @classmethod - def _partialX(cls, - obj_dml_data, - ml_l, - ml_m, - ml_r, - ml_g=None, - n_folds=5, - n_rep=1, - score='partialling out', - dml_procedure='dml2', - draw_sample_splitting=True, - apply_cross_fitting=True): - obj = cls(obj_dml_data, - ml_l, - ml_m, - ml_r, - ml_g, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + def _partialX( + cls, + obj_dml_data, + ml_l, + ml_m, + ml_r, + ml_g=None, + n_folds=5, + n_rep=1, + score="partialling out", + dml_procedure="dml2", + draw_sample_splitting=True, + apply_cross_fitting=True, + ): + obj = cls( + obj_dml_data, + ml_l, + ml_m, + ml_r, + ml_g, + n_folds, + n_rep, + score, + dml_procedure, + draw_sample_splitting, + apply_cross_fitting, + ) obj._check_data(obj._dml_data) obj.partialX = True obj.partialZ = False obj._check_score(obj.score) - _ = obj._check_learner(ml_l, 'ml_l', regressor=True, classifier=False) - _ = obj._check_learner(ml_m, 'ml_m', regressor=True, classifier=False) - _ = obj._check_learner(ml_r, 'ml_r', regressor=True, classifier=False) - obj._learner = {'ml_l': ml_l, 'ml_m': ml_m, 'ml_r': ml_r} - obj._predict_method = {'ml_l': 'predict', 'ml_m': 'predict', 'ml_r': 'predict'} + _ = obj._check_learner(ml_l, "ml_l", regressor=True, classifier=False) + _ = obj._check_learner(ml_m, "ml_m", regressor=True, classifier=False) + _ = obj._check_learner(ml_r, "ml_r", regressor=True, classifier=False) + obj._learner = {"ml_l": ml_l, "ml_m": ml_m, "ml_r": ml_r} + obj._predict_method = {"ml_l": "predict", "ml_m": "predict", "ml_r": "predict"} obj._initialize_ml_nuisance_params() return obj @classmethod - def _partialZ(cls, - obj_dml_data, - ml_r, - n_folds=5, - n_rep=1, - score='partialling out', - dml_procedure='dml2', - draw_sample_splitting=True, - apply_cross_fitting=True): + def _partialZ( + cls, + obj_dml_data, + ml_r, + n_folds=5, + n_rep=1, + score="partialling out", + dml_procedure="dml2", + draw_sample_splitting=True, + apply_cross_fitting=True, + ): # to pass the checks for the learners, we temporarily set ml_l and ml_m to DummyRegressor() - obj = cls(obj_dml_data, - DummyRegressor(), - DummyRegressor(), - ml_r, - None, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + obj = cls( + obj_dml_data, + DummyRegressor(), + DummyRegressor(), + ml_r, + None, + n_folds, + n_rep, + score, + dml_procedure, + draw_sample_splitting, + apply_cross_fitting, + ) obj._check_data(obj._dml_data) obj.partialX = False obj.partialZ = True obj._check_score(obj.score) - _ = obj._check_learner(ml_r, 'ml_r', regressor=True, classifier=False) - obj._learner = {'ml_r': ml_r} - obj._predict_method = {'ml_r': 'predict'} + _ = obj._check_learner(ml_r, "ml_r", regressor=True, classifier=False) + obj._learner = {"ml_r": ml_r} + obj._predict_method = {"ml_r": "predict"} obj._initialize_ml_nuisance_params() return obj @classmethod - def _partialXZ(cls, - obj_dml_data, - ml_l, - ml_m, - ml_r, - n_folds=5, - n_rep=1, - score='partialling out', - dml_procedure='dml2', - draw_sample_splitting=True, - apply_cross_fitting=True): - obj = cls(obj_dml_data, - ml_l, - ml_m, - ml_r, - None, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + def _partialXZ( + cls, + obj_dml_data, + ml_l, + ml_m, + ml_r, + n_folds=5, + n_rep=1, + score="partialling out", + dml_procedure="dml2", + draw_sample_splitting=True, + apply_cross_fitting=True, + ): + obj = cls( + obj_dml_data, + ml_l, + ml_m, + ml_r, + None, + n_folds, + n_rep, + score, + dml_procedure, + draw_sample_splitting, + apply_cross_fitting, + ) obj._check_data(obj._dml_data) obj.partialX = True obj.partialZ = True obj._check_score(obj.score) - _ = obj._check_learner(ml_l, 'ml_l', regressor=True, classifier=False) - _ = obj._check_learner(ml_m, 'ml_m', regressor=True, classifier=False) - _ = obj._check_learner(ml_r, 'ml_r', regressor=True, classifier=False) - obj._learner = {'ml_l': ml_l, 'ml_m': ml_m, 'ml_r': ml_r} - obj._predict_method = {'ml_l': 'predict', 'ml_m': 'predict', 'ml_r': 'predict'} + _ = obj._check_learner(ml_l, "ml_l", regressor=True, classifier=False) + _ = obj._check_learner(ml_m, "ml_m", regressor=True, classifier=False) + _ = obj._check_learner(ml_r, "ml_r", regressor=True, classifier=False) + obj._learner = {"ml_l": ml_l, "ml_m": ml_m, "ml_r": ml_r} + obj._predict_method = {"ml_l": "predict", "ml_m": "predict", "ml_r": "predict"} obj._initialize_ml_nuisance_params() return obj def _initialize_ml_nuisance_params(self): if self.partialX & (not self.partialZ) & (self._dml_data.n_instr > 1): - param_names = ['ml_l', 'ml_r'] + ['ml_m_' + z_col for z_col in self._dml_data.z_cols] + param_names = ["ml_l", "ml_r"] + ["ml_m_" + z_col for z_col in self._dml_data.z_cols] else: param_names = self._learner.keys() - self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} - for learner in param_names} + self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} for learner in param_names} def _check_score(self, score): if isinstance(score, str): if self.partialX & (not self.partialZ) & (self._dml_data.n_instr == 1): - valid_score = ['partialling out', 'IV-type'] + valid_score = ["partialling out", "IV-type"] else: - valid_score = ['partialling out'] + valid_score = ["partialling out"] if score not in valid_score: - raise ValueError('Invalid score ' + score + '. ' + - 'Valid score ' + ' or '.join(valid_score) + '.') + raise ValueError("Invalid score " + score + ". " + "Valid score " + " or ".join(valid_score) + ".") else: if not callable(score): - raise TypeError('score should be either a string or a callable. ' - '%r was passed.' % score) + raise TypeError("score should be either a string or a callable. " "%r was passed." % score) return score def _check_data(self, obj_dml_data): if not isinstance(obj_dml_data, DoubleMLData): - raise TypeError('The data must be of DoubleMLData type. ' - f'{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed.') + raise TypeError( + "The data must be of DoubleMLData type. " f"{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed." + ) if obj_dml_data.n_instr == 0: - raise ValueError('Incompatible data. ' + - 'At least one variable must be set as instrumental variable. ' - 'To fit a partially linear regression model without instrumental variable(s) ' - 'use DoubleMLPLR instead of DoubleMLPLIV.') + raise ValueError( + "Incompatible data. " + "At least one variable must be set as instrumental variable. " + "To fit a partially linear regression model without instrumental variable(s) " + "use DoubleMLPLR instead of DoubleMLPLIV." + ) return def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=False): @@ -291,133 +302,197 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa elif (not self.partialX) & self.partialZ: psi_elements, preds = self._nuisance_est_partial_z(smpls, n_jobs_cv, return_models) else: - assert (self.partialX & self.partialZ) + assert self.partialX & self.partialZ psi_elements, preds = self._nuisance_est_partial_xz(smpls, n_jobs_cv, return_models) return psi_elements, preds - def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): + def _nuisance_tuning( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): if self.partialX & (not self.partialZ): - res = self._nuisance_tuning_partial_x(smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search) + res = self._nuisance_tuning_partial_x( + smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ) elif (not self.partialX) & self.partialZ: - res = self._nuisance_tuning_partial_z(smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search) + res = self._nuisance_tuning_partial_z( + smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ) else: - assert (self.partialX & self.partialZ) - res = self._nuisance_tuning_partial_xz(smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search) + assert self.partialX & self.partialZ + res = self._nuisance_tuning_partial_xz( + smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ) return res def _nuisance_est_partial_x(self, smpls, n_jobs_cv, external_predictions, return_models=False): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) # nuisance l +<<<<<<< HEAD if external_predictions['ml_l'] is not None: l_hat = {'preds': external_predictions['ml_l'], - 'targets': None, - 'models': None} + 'targets': None, + 'models': None} +======= + if external_predictions["ml_l"] is not None: + l_hat = {"preds": external_predictions["ml_l"], "targets": None, "models": None} +>>>>>>> 5d59ac29a02b6034a9e550709069398ea177a30d else: - l_hat = _dml_cv_predict(self._learner['ml_l'], x, y, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_l'), method=self._predict_method['ml_l'], - return_models=return_models) - _check_finite_predictions(l_hat['preds'], self._learner['ml_l'], 'ml_l', smpls) - - predictions = {'ml_l': l_hat['preds']} - targets = {'ml_l': l_hat['targets']} - models = {'ml_l': l_hat['models']} + l_hat = _dml_cv_predict( + self._learner["ml_l"], + x, + y, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_l"), + method=self._predict_method["ml_l"], + return_models=return_models, + ) + _check_finite_predictions(l_hat["preds"], self._learner["ml_l"], "ml_l", smpls) + + predictions = {"ml_l": l_hat["preds"]} + targets = {"ml_l": l_hat["targets"]} + models = {"ml_l": l_hat["models"]} # nuisance m if self._dml_data.n_instr == 1: # one instrument: just identified +<<<<<<< HEAD x, z = check_X_y(x, np.ravel(self._dml_data.z), - force_all_finite=False) + force_all_finite=False) if external_predictions['ml_m'] is not None: m_hat = {'preds': external_predictions['ml_m'], - 'targets': None, - 'models': None} + 'targets': None, + 'models': None} +======= + x, z = check_X_y(x, np.ravel(self._dml_data.z), force_all_finite=False) + if external_predictions["ml_m"] is not None: + m_hat = {"preds": external_predictions["ml_m"], "targets": None, "models": None} +>>>>>>> 5d59ac29a02b6034a9e550709069398ea177a30d else: - m_hat = _dml_cv_predict(self._learner['ml_m'], x, z, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_m'), method=self._predict_method['ml_m'], - return_models=return_models) - predictions['ml_m'] = m_hat['preds'] - targets['ml_m'] = m_hat['targets'] - models['ml_m'] = m_hat['models'] + m_hat = _dml_cv_predict( + self._learner["ml_m"], + x, + z, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_m"), + method=self._predict_method["ml_m"], + return_models=return_models, + ) + predictions["ml_m"] = m_hat["preds"] + targets["ml_m"] = m_hat["targets"] + models["ml_m"] = m_hat["models"] else: # several instruments: 2SLS - m_hat = {'preds': np.full((self._dml_data.n_obs, self._dml_data.n_instr), np.nan), - 'targets': [None] * self._dml_data.n_instr, - 'models': [None] * self._dml_data.n_instr} + m_hat = { + "preds": np.full((self._dml_data.n_obs, self._dml_data.n_instr), np.nan), + "targets": [None] * self._dml_data.n_instr, + "models": [None] * self._dml_data.n_instr, + } for i_instr in range(self._dml_data.n_instr): z = self._dml_data.z +<<<<<<< HEAD x, this_z = check_X_y(x, z[:, i_instr], - force_all_finite=False) + force_all_finite=False) if external_predictions['ml_m_' + self._dml_data.z_cols[i_instr]] is not None: m_hat['preds'][:, i_instr] = external_predictions['ml_m_' + self._dml_data.z_cols[i_instr]] - predictions['ml_m_' + self._dml_data.z_cols[i_instr]] = external_predictions['ml_m_' + self._dml_data.z_cols[i_instr]] + predictions['ml_m_' + self._dml_data.z_cols[i_instr]] = external_predictions[ + 'ml_m_' + self._dml_data.z_cols[i_instr] + ] targets['ml_m_' + self._dml_data.z_cols[i_instr]] = None models['ml_m_' + self._dml_data.z_cols[i_instr]] = None else: res_cv_predict = _dml_cv_predict(self._learner['ml_m'], x, this_z, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_m_' + self._dml_data.z_cols[i_instr]), - method=self._predict_method['ml_m'], return_models=return_models) - - m_hat['preds'][:, i_instr] = res_cv_predict['preds'] - - predictions['ml_m_' + self._dml_data.z_cols[i_instr]] = res_cv_predict['preds'] - targets['ml_m_' + self._dml_data.z_cols[i_instr]] = res_cv_predict['targets'] - models['ml_m_' + self._dml_data.z_cols[i_instr]] = res_cv_predict['models'] - - _check_finite_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls) + est_params=self._get_params('ml_m_' + self._dml_data.z_cols[i_instr]), + method=self._predict_method['ml_m'], return_models=return_models) +======= + x, this_z = check_X_y(x, z[:, i_instr], force_all_finite=False) + if external_predictions["ml_m_" + self._dml_data.z_cols[i_instr]] is not None: + m_hat["preds"][:, i_instr] = external_predictions["ml_m_" + self._dml_data.z_cols[i_instr]] + predictions["ml_m_" + self._dml_data.z_cols[i_instr]] = external_predictions[ + "ml_m_" + self._dml_data.z_cols[i_instr] + ] + targets["ml_m_" + self._dml_data.z_cols[i_instr]] = None + models["ml_m_" + self._dml_data.z_cols[i_instr]] = None + else: + res_cv_predict = _dml_cv_predict( + self._learner["ml_m"], + x, + this_z, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_m_" + self._dml_data.z_cols[i_instr]), + method=self._predict_method["ml_m"], + return_models=return_models, + ) +>>>>>>> 5d59ac29a02b6034a9e550709069398ea177a30d + + m_hat["preds"][:, i_instr] = res_cv_predict["preds"] + + predictions["ml_m_" + self._dml_data.z_cols[i_instr]] = res_cv_predict["preds"] + targets["ml_m_" + self._dml_data.z_cols[i_instr]] = res_cv_predict["targets"] + models["ml_m_" + self._dml_data.z_cols[i_instr]] = res_cv_predict["models"] + + _check_finite_predictions(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls) # nuisance r - if external_predictions['ml_r'] is not None: - r_hat = {'preds': external_predictions['ml_r'], - 'targets': None, - 'models': None} + if external_predictions["ml_r"] is not None: + r_hat = {"preds": external_predictions["ml_r"], "targets": None, "models": None} else: - r_hat = _dml_cv_predict(self._learner['ml_r'], x, d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_r'), method=self._predict_method['ml_r'], - return_models=return_models) - _check_finite_predictions(r_hat['preds'], self._learner['ml_r'], 'ml_r', smpls) - predictions['ml_r'] = r_hat['preds'] - targets['ml_r'] = r_hat['targets'] - models['ml_r'] = r_hat['models'] - - g_hat = {'preds': None, 'targets': None, 'models': None} - if (self._dml_data.n_instr == 1) & ('ml_g' in self._learner): + r_hat = _dml_cv_predict( + self._learner["ml_r"], + x, + d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_r"), + method=self._predict_method["ml_r"], + return_models=return_models, + ) + _check_finite_predictions(r_hat["preds"], self._learner["ml_r"], "ml_r", smpls) + predictions["ml_r"] = r_hat["preds"] + targets["ml_r"] = r_hat["targets"] + models["ml_r"] = r_hat["models"] + + g_hat = {"preds": None, "targets": None, "models": None} + if (self._dml_data.n_instr == 1) & ("ml_g" in self._learner): # an estimate of g is obtained for the IV-type score and callable scores # get an initial estimate for theta using the partialling out score +<<<<<<< HEAD if external_predictions['ml_g'] is not None: g_hat = {'preds': external_predictions['ml_g'], - 'targets': None, - 'models': None} + 'targets': None, + 'models': None} +======= + if external_predictions["ml_g"] is not None: + g_hat = {"preds": external_predictions["ml_g"], "targets": None, "models": None} +>>>>>>> 5d59ac29a02b6034a9e550709069398ea177a30d else: - psi_a = -np.multiply(d - r_hat['preds'], z - m_hat['preds']) - psi_b = np.multiply(z - m_hat['preds'], y - l_hat['preds']) + psi_a = -np.multiply(d - r_hat["preds"], z - m_hat["preds"]) + psi_b = np.multiply(z - m_hat["preds"], y - l_hat["preds"]) theta_initial = -np.nanmean(psi_b) / np.nanmean(psi_a) # nuisance g - g_hat = _dml_cv_predict(self._learner['ml_g'], x, y - theta_initial * d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g'), method=self._predict_method['ml_g'], - return_models=return_models) - _check_finite_predictions(g_hat['preds'], self._learner['ml_g'], 'ml_g', smpls) - - predictions['ml_g'] = g_hat['preds'] - targets['ml_g'] = g_hat['targets'] - models['ml_g'] = g_hat['models'] - psi_a, psi_b = self._score_elements(y, z, d, - l_hat['preds'], m_hat['preds'], r_hat['preds'], g_hat['preds'], - smpls) - psi_elements = {'psi_a': psi_a, - 'psi_b': psi_b} - predictions = {'predictions': predictions, - 'targets': targets, - 'models': models - } + g_hat = _dml_cv_predict( + self._learner["ml_g"], + x, + y - theta_initial * d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_g"), + method=self._predict_method["ml_g"], + return_models=return_models, + ) + _check_finite_predictions(g_hat["preds"], self._learner["ml_g"], "ml_g", smpls) + + predictions["ml_g"] = g_hat["preds"] + targets["ml_g"] = g_hat["targets"] + models["ml_g"] = g_hat["models"] + psi_a, psi_b = self._score_elements(y, z, d, l_hat["preds"], m_hat["preds"], r_hat["preds"], g_hat["preds"], smpls) + psi_elements = {"psi_a": psi_a, "psi_b": psi_b} + predictions = {"predictions": predictions, "targets": targets, "models": models} return psi_elements, predictions @@ -425,7 +500,7 @@ def _score_elements(self, y, z, d, l_hat, m_hat, r_hat, g_hat, smpls): # compute residuals u_hat = y - l_hat w_hat = d - r_hat - v_hat = z- m_hat + v_hat = z - m_hat r_hat_tilde = None if self._dml_data.n_instr > 1: @@ -437,168 +512,210 @@ def _score_elements(self, y, z, d, l_hat, m_hat, r_hat, g_hat, smpls): if isinstance(self.score, str): if self._dml_data.n_instr == 1: - if self.score == 'partialling out': + if self.score == "partialling out": psi_a = -np.multiply(w_hat, v_hat) psi_b = np.multiply(v_hat, u_hat) else: - assert self.score == 'IV-type' + assert self.score == "IV-type" psi_a = -np.multiply(v_hat, d) psi_b = np.multiply(v_hat, y - g_hat) else: - assert self.score == 'partialling out' + assert self.score == "partialling out" psi_a = -np.multiply(w_hat, r_hat_tilde) psi_b = np.multiply(r_hat_tilde, u_hat) else: assert callable(self.score) if self._dml_data.n_instr > 1: - raise NotImplementedError('Callable score not implemented for DoubleMLPLIV.partialX ' - 'with several instruments.') + raise NotImplementedError( + "Callable score not implemented for DoubleMLPLIV.partialX " "with several instruments." + ) else: assert self._dml_data.n_instr == 1 - psi_a, psi_b = self.score(y=y, z=z, d=d, - l_hat=l_hat, m_hat=m_hat, r_hat=r_hat, g_hat=g_hat, - smpls=smpls) + psi_a, psi_b = self.score(y=y, z=z, d=d, l_hat=l_hat, m_hat=m_hat, r_hat=r_hat, g_hat=g_hat, smpls=smpls) return psi_a, psi_b def _nuisance_est_partial_z(self, smpls, n_jobs_cv, return_models=False): y = self._dml_data.y - xz, d = check_X_y(np.hstack((self._dml_data.x, self._dml_data.z)), - self._dml_data.d, - force_all_finite=False) + xz, d = check_X_y(np.hstack((self._dml_data.x, self._dml_data.z)), self._dml_data.d, force_all_finite=False) # nuisance m - r_hat = _dml_cv_predict(self._learner['ml_r'], xz, d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_r'), method=self._predict_method['ml_r'], - return_models=return_models) - _check_finite_predictions(r_hat['preds'], self._learner['ml_r'], 'ml_r', smpls) + r_hat = _dml_cv_predict( + self._learner["ml_r"], + xz, + d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_r"), + method=self._predict_method["ml_r"], + return_models=return_models, + ) + _check_finite_predictions(r_hat["preds"], self._learner["ml_r"], "ml_r", smpls) if isinstance(self.score, str): - assert self.score == 'partialling out' - psi_a = -np.multiply(r_hat['preds'], d) - psi_b = np.multiply(r_hat['preds'], y) + assert self.score == "partialling out" + psi_a = -np.multiply(r_hat["preds"], d) + psi_b = np.multiply(r_hat["preds"], y) else: assert callable(self.score) - raise NotImplementedError('Callable score not implemented for DoubleMLPLIV.partialZ.') + raise NotImplementedError("Callable score not implemented for DoubleMLPLIV.partialZ.") - psi_elements = {'psi_a': psi_a, - 'psi_b': psi_b} - preds = {'predictions': {'ml_r': r_hat['preds']}, - 'targets': {'ml_r': r_hat['targets']}, - 'models': {'ml_r': r_hat['models']}} + psi_elements = {"psi_a": psi_a, "psi_b": psi_b} + preds = { + "predictions": {"ml_r": r_hat["preds"]}, + "targets": {"ml_r": r_hat["targets"]}, + "models": {"ml_r": r_hat["models"]}, + } return psi_elements, preds def _nuisance_est_partial_xz(self, smpls, n_jobs_cv, return_models=False): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - xz, d = check_X_y(np.hstack((self._dml_data.x, self._dml_data.z)), - self._dml_data.d, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + xz, d = check_X_y(np.hstack((self._dml_data.x, self._dml_data.z)), self._dml_data.d, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) # nuisance l - l_hat = _dml_cv_predict(self._learner['ml_l'], x, y, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_l'), method=self._predict_method['ml_l'], - return_models=return_models) - _check_finite_predictions(l_hat['preds'], self._learner['ml_l'], 'ml_l', smpls) + l_hat = _dml_cv_predict( + self._learner["ml_l"], + x, + y, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_l"), + method=self._predict_method["ml_l"], + return_models=return_models, + ) + _check_finite_predictions(l_hat["preds"], self._learner["ml_l"], "ml_l", smpls) # nuisance m - m_hat = _dml_cv_predict(self._learner['ml_m'], xz, d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_m'), return_train_preds=True, - method=self._predict_method['ml_m'], return_models=return_models) - _check_finite_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls) + m_hat = _dml_cv_predict( + self._learner["ml_m"], + xz, + d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_m"), + return_train_preds=True, + method=self._predict_method["ml_m"], + return_models=return_models, + ) + _check_finite_predictions(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls) # nuisance r - m_hat_tilde = _dml_cv_predict(self._learner['ml_r'], x, m_hat['train_preds'], smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_r'), method=self._predict_method['ml_r'], - return_models=return_models) - _check_finite_predictions(m_hat_tilde['preds'], self._learner['ml_r'], 'ml_r', smpls) + m_hat_tilde = _dml_cv_predict( + self._learner["ml_r"], + x, + m_hat["train_preds"], + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_r"), + method=self._predict_method["ml_r"], + return_models=return_models, + ) + _check_finite_predictions(m_hat_tilde["preds"], self._learner["ml_r"], "ml_r", smpls) # compute residuals - u_hat = y - l_hat['preds'] - w_hat = d - m_hat_tilde['preds'] + u_hat = y - l_hat["preds"] + w_hat = d - m_hat_tilde["preds"] if isinstance(self.score, str): - assert self.score == 'partialling out' - psi_a = -np.multiply(w_hat, (m_hat['preds']-m_hat_tilde['preds'])) - psi_b = np.multiply((m_hat['preds']-m_hat_tilde['preds']), u_hat) + assert self.score == "partialling out" + psi_a = -np.multiply(w_hat, (m_hat["preds"] - m_hat_tilde["preds"])) + psi_b = np.multiply((m_hat["preds"] - m_hat_tilde["preds"]), u_hat) else: assert callable(self.score) - raise NotImplementedError('Callable score not implemented for DoubleMLPLIV.partialXZ.') - - psi_elements = {'psi_a': psi_a, - 'psi_b': psi_b} - preds = {'predictions': {'ml_l': l_hat['preds'], - 'ml_m': m_hat['preds'], - 'ml_r': m_hat_tilde['preds']}, - 'targets': {'ml_l': l_hat['targets'], - 'ml_m': m_hat['targets'], - 'ml_r': m_hat_tilde['targets']}, - 'models': {'ml_l': l_hat['models'], - 'ml_m': m_hat['models'], - 'ml_r': m_hat_tilde['models']} - } + raise NotImplementedError("Callable score not implemented for DoubleMLPLIV.partialXZ.") + + psi_elements = {"psi_a": psi_a, "psi_b": psi_b} + preds = { + "predictions": {"ml_l": l_hat["preds"], "ml_m": m_hat["preds"], "ml_r": m_hat_tilde["preds"]}, + "targets": {"ml_l": l_hat["targets"], "ml_m": m_hat["targets"], "ml_r": m_hat_tilde["targets"]}, + "models": {"ml_l": l_hat["models"], "ml_m": m_hat["models"], "ml_r": m_hat_tilde["models"]}, + } return psi_elements, preds - def _nuisance_tuning_partial_x(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + def _nuisance_tuning_partial_x( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) if scoring_methods is None: - scoring_methods = {'ml_l': None, - 'ml_m': None, - 'ml_r': None, - 'ml_g': None} + scoring_methods = {"ml_l": None, "ml_m": None, "ml_r": None, "ml_g": None} train_inds = [train_index for (train_index, _) in smpls] - l_tune_res = _dml_tune(y, x, train_inds, - self._learner['ml_l'], param_grids['ml_l'], scoring_methods['ml_l'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + l_tune_res = _dml_tune( + y, + x, + train_inds, + self._learner["ml_l"], + param_grids["ml_l"], + scoring_methods["ml_l"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) if self._dml_data.n_instr > 1: # several instruments: 2SLS m_tune_res = {instr_var: list() for instr_var in self._dml_data.z_cols} z = self._dml_data.z for i_instr in range(self._dml_data.n_instr): - x, this_z = check_X_y(x, z[:, i_instr], - force_all_finite=False) - m_tune_res[self._dml_data.z_cols[i_instr]] = _dml_tune(this_z, x, train_inds, - self._learner['ml_m'], param_grids['ml_m'], - scoring_methods['ml_m'], - n_folds_tune, n_jobs_cv, search_mode, - n_iter_randomized_search) + x, this_z = check_X_y(x, z[:, i_instr], force_all_finite=False) + m_tune_res[self._dml_data.z_cols[i_instr]] = _dml_tune( + this_z, + x, + train_inds, + self._learner["ml_m"], + param_grids["ml_m"], + scoring_methods["ml_m"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) else: # one instrument: just identified - x, z = check_X_y(x, np.ravel(self._dml_data.z), - force_all_finite=False) - m_tune_res = _dml_tune(z, x, train_inds, - self._learner['ml_m'], param_grids['ml_m'], scoring_methods['ml_m'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - - r_tune_res = _dml_tune(d, x, train_inds, - self._learner['ml_r'], param_grids['ml_r'], scoring_methods['ml_r'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + x, z = check_X_y(x, np.ravel(self._dml_data.z), force_all_finite=False) + m_tune_res = _dml_tune( + z, + x, + train_inds, + self._learner["ml_m"], + param_grids["ml_m"], + scoring_methods["ml_m"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + + r_tune_res = _dml_tune( + d, + x, + train_inds, + self._learner["ml_r"], + param_grids["ml_r"], + scoring_methods["ml_r"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) l_best_params = [xx.best_params_ for xx in l_tune_res] r_best_params = [xx.best_params_ for xx in r_tune_res] if self._dml_data.n_instr > 1: - params = {'ml_l': l_best_params, - 'ml_r': r_best_params} + params = {"ml_l": l_best_params, "ml_r": r_best_params} for instr_var in self._dml_data.z_cols: - params['ml_m_' + instr_var] = [xx.best_params_ for xx in m_tune_res[instr_var]] - tune_res = {'l_tune': l_tune_res, - 'm_tune': m_tune_res, - 'r_tune': r_tune_res} + params["ml_m_" + instr_var] = [xx.best_params_ for xx in m_tune_res[instr_var]] + tune_res = {"l_tune": l_tune_res, "m_tune": m_tune_res, "r_tune": r_tune_res} else: m_best_params = [xx.best_params_ for xx in m_tune_res] # an ML model for g is obtained for the IV-type score and callable scores - if 'ml_g' in self._learner: + if "ml_g" in self._learner: # construct an initial theta estimate from the tuned models using the partialling out score l_hat = np.full_like(y, np.nan) m_hat = np.full_like(z, np.nan) @@ -610,110 +727,131 @@ def _nuisance_tuning_partial_x(self, smpls, param_grids, scoring_methods, n_fold psi_a = -np.multiply(d - r_hat, z - m_hat) psi_b = np.multiply(z - m_hat, y - l_hat) theta_initial = -np.nanmean(psi_b) / np.nanmean(psi_a) - g_tune_res = _dml_tune(y - theta_initial * d, x, train_inds, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + g_tune_res = _dml_tune( + y - theta_initial * d, + x, + train_inds, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) g_best_params = [xx.best_params_ for xx in g_tune_res] - params = {'ml_l': l_best_params, - 'ml_m': m_best_params, - 'ml_r': r_best_params, - 'ml_g': g_best_params} - tune_res = {'l_tune': l_tune_res, - 'm_tune': m_tune_res, - 'r_tune': r_tune_res, - 'g_tune': g_tune_res} + params = {"ml_l": l_best_params, "ml_m": m_best_params, "ml_r": r_best_params, "ml_g": g_best_params} + tune_res = {"l_tune": l_tune_res, "m_tune": m_tune_res, "r_tune": r_tune_res, "g_tune": g_tune_res} else: - params = {'ml_l': l_best_params, - 'ml_m': m_best_params, - 'ml_r': r_best_params} - tune_res = {'l_tune': l_tune_res, - 'm_tune': m_tune_res, - 'r_tune': r_tune_res} + params = {"ml_l": l_best_params, "ml_m": m_best_params, "ml_r": r_best_params} + tune_res = {"l_tune": l_tune_res, "m_tune": m_tune_res, "r_tune": r_tune_res} - res = {'params': params, - 'tune_res': tune_res} + res = {"params": params, "tune_res": tune_res} return res - def _nuisance_tuning_partial_z(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): - xz, d = check_X_y(np.hstack((self._dml_data.x, self._dml_data.z)), - self._dml_data.d, - force_all_finite=False) + def _nuisance_tuning_partial_z( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): + xz, d = check_X_y(np.hstack((self._dml_data.x, self._dml_data.z)), self._dml_data.d, force_all_finite=False) if scoring_methods is None: - scoring_methods = {'ml_r': None} + scoring_methods = {"ml_r": None} train_inds = [train_index for (train_index, _) in smpls] - m_tune_res = _dml_tune(d, xz, train_inds, - self._learner['ml_r'], param_grids['ml_r'], scoring_methods['ml_r'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + m_tune_res = _dml_tune( + d, + xz, + train_inds, + self._learner["ml_r"], + param_grids["ml_r"], + scoring_methods["ml_r"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) m_best_params = [xx.best_params_ for xx in m_tune_res] - params = {'ml_r': m_best_params} + params = {"ml_r": m_best_params} - tune_res = {'r_tune': m_tune_res} + tune_res = {"r_tune": m_tune_res} - res = {'params': params, - 'tune_res': tune_res} + res = {"params": params, "tune_res": tune_res} return res - def _nuisance_tuning_partial_xz(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - xz, d = check_X_y(np.hstack((self._dml_data.x, self._dml_data.z)), - self._dml_data.d, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + def _nuisance_tuning_partial_xz( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + xz, d = check_X_y(np.hstack((self._dml_data.x, self._dml_data.z)), self._dml_data.d, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) if scoring_methods is None: - scoring_methods = {'ml_l': None, - 'ml_m': None, - 'ml_r': None} + scoring_methods = {"ml_l": None, "ml_m": None, "ml_r": None} train_inds = [train_index for (train_index, _) in smpls] - l_tune_res = _dml_tune(y, x, train_inds, - self._learner['ml_l'], param_grids['ml_l'], scoring_methods['ml_l'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - m_tune_res = _dml_tune(d, xz, train_inds, - self._learner['ml_m'], param_grids['ml_m'], scoring_methods['ml_m'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + l_tune_res = _dml_tune( + y, + x, + train_inds, + self._learner["ml_l"], + param_grids["ml_l"], + scoring_methods["ml_l"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + m_tune_res = _dml_tune( + d, + xz, + train_inds, + self._learner["ml_m"], + param_grids["ml_m"], + scoring_methods["ml_m"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) r_tune_res = list() for idx, (train_index, _) in enumerate(smpls): m_hat = m_tune_res[idx].predict(xz[train_index, :]) r_tune_resampling = KFold(n_splits=n_folds_tune, shuffle=True) - if search_mode == 'grid_search': - r_grid_search = GridSearchCV(self._learner['ml_r'], param_grids['ml_r'], - scoring=scoring_methods['ml_r'], - cv=r_tune_resampling, n_jobs=n_jobs_cv) + if search_mode == "grid_search": + r_grid_search = GridSearchCV( + self._learner["ml_r"], + param_grids["ml_r"], + scoring=scoring_methods["ml_r"], + cv=r_tune_resampling, + n_jobs=n_jobs_cv, + ) else: - assert search_mode == 'randomized_search' - r_grid_search = RandomizedSearchCV(self._learner['ml_r'], param_grids['ml_r'], - scoring=scoring_methods['ml_r'], - cv=r_tune_resampling, n_jobs=n_jobs_cv, - n_iter=n_iter_randomized_search) + assert search_mode == "randomized_search" + r_grid_search = RandomizedSearchCV( + self._learner["ml_r"], + param_grids["ml_r"], + scoring=scoring_methods["ml_r"], + cv=r_tune_resampling, + n_jobs=n_jobs_cv, + n_iter=n_iter_randomized_search, + ) r_tune_res.append(r_grid_search.fit(x[train_index, :], m_hat)) l_best_params = [xx.best_params_ for xx in l_tune_res] m_best_params = [xx.best_params_ for xx in m_tune_res] r_best_params = [xx.best_params_ for xx in r_tune_res] - params = {'ml_l': l_best_params, - 'ml_m': m_best_params, - 'ml_r': r_best_params} + params = {"ml_l": l_best_params, "ml_m": m_best_params, "ml_r": r_best_params} - tune_res = {'l_tune': l_tune_res, - 'm_tune': m_tune_res, - 'r_tune': r_tune_res} + tune_res = {"l_tune": l_tune_res, "m_tune": m_tune_res, "r_tune": r_tune_res} - res = {'params': params, - 'tune_res': tune_res} + res = {"params": params, "tune_res": tune_res} return res diff --git a/doubleml/double_ml_plr.py b/doubleml/double_ml_plr.py index 89044fb8..93145b5b 100644 --- a/doubleml/double_ml_plr.py +++ b/doubleml/double_ml_plr.py @@ -96,75 +96,76 @@ class DoubleMLPLR(LinearScoreMixin, DoubleML): and :math:`\\zeta` and :math:`V` are stochastic errors. """ - def __init__(self, - obj_dml_data, - ml_l, - ml_m, - ml_g=None, - n_folds=5, - n_rep=1, - score='partialling out', - dml_procedure='dml2', - draw_sample_splitting=True, - apply_cross_fitting=True): - super().__init__(obj_dml_data, - n_folds, - n_rep, - score, - dml_procedure, - draw_sample_splitting, - apply_cross_fitting) + def __init__( + self, + obj_dml_data, + ml_l, + ml_m, + ml_g=None, + n_folds=5, + n_rep=1, + score="partialling out", + dml_procedure="dml2", + draw_sample_splitting=True, + apply_cross_fitting=True, + ): + super().__init__(obj_dml_data, n_folds, n_rep, score, dml_procedure, draw_sample_splitting, apply_cross_fitting) self._check_data(self._dml_data) - valid_scores = ['IV-type', 'partialling out'] + valid_scores = ["IV-type", "partialling out"] _check_score(self.score, valid_scores, allow_callable=True) - _ = self._check_learner(ml_l, 'ml_l', regressor=True, classifier=False) - ml_m_is_classifier = self._check_learner(ml_m, 'ml_m', regressor=True, classifier=True) - self._learner = {'ml_l': ml_l, 'ml_m': ml_m} + _ = self._check_learner(ml_l, "ml_l", regressor=True, classifier=False) + ml_m_is_classifier = self._check_learner(ml_m, "ml_m", regressor=True, classifier=True) + self._learner = {"ml_l": ml_l, "ml_m": ml_m} if ml_g is not None: - if (isinstance(self.score, str) & (self.score == 'IV-type')) | callable(self.score): - _ = self._check_learner(ml_g, 'ml_g', regressor=True, classifier=False) - self._learner['ml_g'] = ml_g + if (isinstance(self.score, str) & (self.score == "IV-type")) | callable(self.score): + _ = self._check_learner(ml_g, "ml_g", regressor=True, classifier=False) + self._learner["ml_g"] = ml_g else: - assert (isinstance(self.score, str) & (self.score == 'partialling out')) - warnings.warn(('A learner ml_g has been provided for score = "partialling out" but will be ignored. "' - 'A learner ml_g is not required for estimation.')) - elif isinstance(self.score, str) & (self.score == 'IV-type'): - warnings.warn(("For score = 'IV-type', learners ml_l and ml_g should be specified. " - "Set ml_g = clone(ml_l).")) - self._learner['ml_g'] = clone(ml_l) - - self._predict_method = {'ml_l': 'predict'} - if 'ml_g' in self._learner: - self._predict_method['ml_g'] = 'predict' + assert isinstance(self.score, str) & (self.score == "partialling out") + warnings.warn( + ( + 'A learner ml_g has been provided for score = "partialling out" but will be ignored. "' + "A learner ml_g is not required for estimation." + ) + ) + elif isinstance(self.score, str) & (self.score == "IV-type"): + warnings.warn(("For score = 'IV-type', learners ml_l and ml_g should be specified. " "Set ml_g = clone(ml_l).")) + self._learner["ml_g"] = clone(ml_l) + + self._predict_method = {"ml_l": "predict"} + if "ml_g" in self._learner: + self._predict_method["ml_g"] = "predict" if ml_m_is_classifier: if self._dml_data.binary_treats.all(): - self._predict_method['ml_m'] = 'predict_proba' + self._predict_method["ml_m"] = "predict_proba" else: - raise ValueError(f'The ml_m learner {str(ml_m)} was identified as classifier ' - 'but at least one treatment variable is not binary with values 0 and 1.') + raise ValueError( + f"The ml_m learner {str(ml_m)} was identified as classifier " + "but at least one treatment variable is not binary with values 0 and 1." + ) else: - self._predict_method['ml_m'] = 'predict' + self._predict_method["ml_m"] = "predict" self._initialize_ml_nuisance_params() self._sensitivity_implemented = True self._external_predictions_implemented = True def _initialize_ml_nuisance_params(self): - self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} - for learner in self._learner} + self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols} for learner in self._learner} def _check_data(self, obj_dml_data): if not isinstance(obj_dml_data, DoubleMLData): - raise TypeError('The data must be of DoubleMLData type. ' - f'{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed.') + raise TypeError( + "The data must be of DoubleMLData type. " f"{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed." + ) if obj_dml_data.z_cols is not None: - raise ValueError('Incompatible data. ' + - ' and '.join(obj_dml_data.z_cols) + - ' have been set as instrumental variable(s). ' - 'To fit a partially linear IV regression model use DoubleMLPLIV instead of DoubleMLPLR.') + raise ValueError( + "Incompatible data. " + " and ".join(obj_dml_data.z_cols) + " have been set as instrumental variable(s). " + "To fit a partially linear IV regression model use DoubleMLPLIV instead of DoubleMLPLR." + ) return def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=False): @@ -182,43 +183,55 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa # nuisance l if l_external: - l_hat = {'preds': external_predictions['ml_l'], - 'targets': None, - 'models': None} + l_hat = {"preds": external_predictions["ml_l"], "targets": None, "models": None} else: - l_hat = _dml_cv_predict(self._learner['ml_l'], x, y, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_l'), method=self._predict_method['ml_l'], - return_models=return_models) - _check_finite_predictions(l_hat['preds'], self._learner['ml_l'], 'ml_l', smpls) + l_hat = _dml_cv_predict( + self._learner["ml_l"], + x, + y, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_l"), + method=self._predict_method["ml_l"], + return_models=return_models, + ) + _check_finite_predictions(l_hat["preds"], self._learner["ml_l"], "ml_l", smpls) # nuisance m if m_external: - m_hat = {'preds': external_predictions['ml_m'], - 'targets': None, - 'models': None} + m_hat = {"preds": external_predictions["ml_m"], "targets": None, "models": None} else: - m_hat = _dml_cv_predict(self._learner['ml_m'], x, d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_m'), method=self._predict_method['ml_m'], - return_models=return_models) - _check_finite_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls) - if self._check_learner(self._learner['ml_m'], 'ml_m', regressor=True, classifier=True): - _check_is_propensity(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls, eps=1e-12) + m_hat = _dml_cv_predict( + self._learner["ml_m"], + x, + d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_m"), + method=self._predict_method["ml_m"], + return_models=return_models, + ) + _check_finite_predictions(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls) + if self._check_learner(self._learner["ml_m"], "ml_m", regressor=True, classifier=True): + _check_is_propensity(m_hat["preds"], self._learner["ml_m"], "ml_m", smpls, eps=1e-12) if self._dml_data.binary_treats[self._dml_data.d_cols[self._i_treat]]: - binary_preds = (type_of_target(m_hat['preds']) == 'binary') - zero_one_preds = np.all((np.power(m_hat['preds'], 2) - m_hat['preds']) == 0) + binary_preds = type_of_target(m_hat["preds"]) == "binary" + zero_one_preds = np.all((np.power(m_hat["preds"], 2) - m_hat["preds"]) == 0) if binary_preds & zero_one_preds: - raise ValueError(f'For the binary treatment variable {self._dml_data.d_cols[self._i_treat]}, ' - f'predictions obtained with the ml_m learner {str(self._learner["ml_m"])} are also ' - 'observed to be binary with values 0 and 1. Make sure that for classifiers ' - 'probabilities and not labels are predicted.') + raise ValueError( + f"For the binary treatment variable {self._dml_data.d_cols[self._i_treat]}, " + f'predictions obtained with the ml_m learner {str(self._learner["ml_m"])} are also ' + "observed to be binary with values 0 and 1. Make sure that for classifiers " + "probabilities and not labels are predicted." + ) # an estimate of g is obtained for the IV-type score and callable scores - g_hat = {'preds': None, 'targets': None, 'models': None} - if 'ml_g' in self._learner: + g_hat = {"preds": None, "targets": None, "models": None} + if "ml_g" in self._learner: # get an initial estimate for theta using the partialling out score - psi_a = -np.multiply(d - m_hat['preds'], d - m_hat['preds']) - psi_b = np.multiply(d - m_hat['preds'], y - l_hat['preds']) + psi_a = -np.multiply(d - m_hat["preds"], d - m_hat["preds"]) + psi_b = np.multiply(d - m_hat["preds"], y - l_hat["preds"]) theta_initial = -np.nanmean(psi_b) / np.nanmean(psi_a) # nuisance g if g_external: @@ -226,23 +239,25 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa 'targets': None, 'models': None} else: - g_hat = _dml_cv_predict(self._learner['ml_g'], x, y - theta_initial*d, smpls=smpls, n_jobs=n_jobs_cv, - est_params=self._get_params('ml_g'), method=self._predict_method['ml_g'], - return_models=return_models) - _check_finite_predictions(g_hat['preds'], self._learner['ml_g'], 'ml_g', smpls) - - psi_a, psi_b = self._score_elements(y, d, l_hat['preds'], m_hat['preds'], g_hat['preds'], smpls) - psi_elements = {'psi_a': psi_a, - 'psi_b': psi_b} - preds = {'predictions': {'ml_l': l_hat['preds'], - 'ml_m': m_hat['preds'], - 'ml_g': g_hat['preds']}, - 'targets': {'ml_l': l_hat['targets'], - 'ml_m': m_hat['targets'], - 'ml_g': g_hat['targets']}, - 'models': {'ml_l': l_hat['models'], - 'ml_m': m_hat['models'], - 'ml_g': g_hat['models']}} + g_hat = _dml_cv_predict( + self._learner["ml_g"], + x, + y - theta_initial * d, + smpls=smpls, + n_jobs=n_jobs_cv, + est_params=self._get_params("ml_g"), + method=self._predict_method["ml_g"], + return_models=return_models, + ) + _check_finite_predictions(g_hat["preds"], self._learner["ml_g"], "ml_g", smpls) + + psi_a, psi_b = self._score_elements(y, d, l_hat["preds"], m_hat["preds"], g_hat["preds"], smpls) + psi_elements = {"psi_a": psi_a, "psi_b": psi_b} + preds = { + "predictions": {"ml_l": l_hat["preds"], "ml_m": m_hat["preds"], "ml_g": g_hat["preds"]}, + "targets": {"ml_l": l_hat["targets"], "ml_m": m_hat["targets"], "ml_g": g_hat["targets"]}, + "models": {"ml_l": l_hat["models"], "ml_m": m_hat["models"], "ml_g": g_hat["models"]}, + } return psi_elements, preds @@ -252,18 +267,16 @@ def _score_elements(self, y, d, l_hat, m_hat, g_hat, smpls): v_hat = d - m_hat if isinstance(self.score, str): - if self.score == 'IV-type': - psi_a = - np.multiply(v_hat, d) + if self.score == "IV-type": + psi_a = -np.multiply(v_hat, d) psi_b = np.multiply(v_hat, y - g_hat) else: - assert self.score == 'partialling out' + assert self.score == "partialling out" psi_a = -np.multiply(v_hat, v_hat) psi_b = np.multiply(v_hat, u_hat) else: assert callable(self.score) - psi_a, psi_b = self.score(y=y, d=d, - l_hat=l_hat, m_hat=m_hat, g_hat=g_hat, - smpls=smpls) + psi_a, psi_b = self.score(y=y, d=d, l_hat=l_hat, m_hat=m_hat, g_hat=g_hat, smpls=smpls) return psi_a, psi_b @@ -272,54 +285,66 @@ def _sensitivity_element_est(self, preds): y = self._dml_data.y d = self._dml_data.d - m_hat = preds['predictions']['ml_m'] + m_hat = preds["predictions"]["ml_m"] theta = self.all_coef[self._i_treat, self._i_rep] - if self.score == 'partialling out': - l_hat = preds['predictions']['ml_l'] - sigma2_score_element = np.square(y - l_hat - np.multiply(theta, d-m_hat)) + if self.score == "partialling out": + l_hat = preds["predictions"]["ml_l"] + sigma2_score_element = np.square(y - l_hat - np.multiply(theta, d - m_hat)) else: - assert self.score == 'IV-type' - g_hat = preds['predictions']['ml_g'] + assert self.score == "IV-type" + g_hat = preds["predictions"]["ml_g"] sigma2_score_element = np.square(y - g_hat - np.multiply(theta, d)) sigma2 = np.mean(sigma2_score_element) psi_sigma2 = sigma2_score_element - sigma2 nu2 = np.divide(1.0, np.mean(np.square(d - m_hat))) - psi_nu2 = nu2 - np.multiply(np.square(d-m_hat), np.square(nu2)) + psi_nu2 = nu2 - np.multiply(np.square(d - m_hat), np.square(nu2)) - element_dict = {'sigma2': sigma2, - 'nu2': nu2, - 'psi_sigma2': psi_sigma2, - 'psi_nu2': psi_nu2} + element_dict = {"sigma2": sigma2, "nu2": nu2, "psi_sigma2": psi_sigma2, "psi_nu2": psi_nu2} return element_dict - def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, - search_mode, n_iter_randomized_search): - x, y = check_X_y(self._dml_data.x, self._dml_data.y, - force_all_finite=False) - x, d = check_X_y(x, self._dml_data.d, - force_all_finite=False) + def _nuisance_tuning( + self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search + ): + x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False) + x, d = check_X_y(x, self._dml_data.d, force_all_finite=False) if scoring_methods is None: - scoring_methods = {'ml_l': None, - 'ml_m': None, - 'ml_g': None} + scoring_methods = {"ml_l": None, "ml_m": None, "ml_g": None} train_inds = [train_index for (train_index, _) in smpls] - l_tune_res = _dml_tune(y, x, train_inds, - self._learner['ml_l'], param_grids['ml_l'], scoring_methods['ml_l'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) - m_tune_res = _dml_tune(d, x, train_inds, - self._learner['ml_m'], param_grids['ml_m'], scoring_methods['ml_m'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + l_tune_res = _dml_tune( + y, + x, + train_inds, + self._learner["ml_l"], + param_grids["ml_l"], + scoring_methods["ml_l"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) + m_tune_res = _dml_tune( + d, + x, + train_inds, + self._learner["ml_m"], + param_grids["ml_m"], + scoring_methods["ml_m"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) l_best_params = [xx.best_params_ for xx in l_tune_res] m_best_params = [xx.best_params_ for xx in m_tune_res] # an ML model for g is obtained for the IV-type score and callable scores - if 'ml_g' in self._learner: + if "ml_g" in self._learner: # construct an initial theta estimate from the tuned models using the partialling out score l_hat = np.full_like(y, np.nan) m_hat = np.full_like(d, np.nan) @@ -329,24 +354,26 @@ def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_ psi_a = -np.multiply(d - m_hat, d - m_hat) psi_b = np.multiply(d - m_hat, y - l_hat) theta_initial = -np.nanmean(psi_b) / np.nanmean(psi_a) - g_tune_res = _dml_tune(y - theta_initial*d, x, train_inds, - self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'], - n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search) + g_tune_res = _dml_tune( + y - theta_initial * d, + x, + train_inds, + self._learner["ml_g"], + param_grids["ml_g"], + scoring_methods["ml_g"], + n_folds_tune, + n_jobs_cv, + search_mode, + n_iter_randomized_search, + ) g_best_params = [xx.best_params_ for xx in g_tune_res] - params = {'ml_l': l_best_params, - 'ml_m': m_best_params, - 'ml_g': g_best_params} - tune_res = {'l_tune': l_tune_res, - 'm_tune': m_tune_res, - 'g_tune': g_tune_res} + params = {"ml_l": l_best_params, "ml_m": m_best_params, "ml_g": g_best_params} + tune_res = {"l_tune": l_tune_res, "m_tune": m_tune_res, "g_tune": g_tune_res} else: - params = {'ml_l': l_best_params, - 'ml_m': m_best_params} - tune_res = {'l_tune': l_tune_res, - 'm_tune': m_tune_res} + params = {"ml_l": l_best_params, "ml_m": m_best_params} + tune_res = {"l_tune": l_tune_res, "m_tune": m_tune_res} - res = {'params': params, - 'tune_res': tune_res} + res = {"params": params, "tune_res": tune_res} return res diff --git a/doubleml/double_ml_pq.py b/doubleml/double_ml_pq.py index e7f42eae..f035dafe 100644 --- a/doubleml/double_ml_pq.py +++ b/doubleml/double_ml_pq.py @@ -1,5 +1,4 @@ import numpy as np -import copy from sklearn.base import clone from sklearn.utils import check_X_y from sklearn.model_selection import StratifiedKFold, train_test_split @@ -182,7 +181,10 @@ def __init__( stratify=self._dml_data.d, ) self._smpls = obj_dml_resampling.split_samples() - +<<<<<<< HEAD +======= + +>>>>>>> 5d59ac29a02b6034a9e550709069398ea177a30d self._external_predictions_implemented = True @property diff --git a/doubleml/double_ml_qte.py b/doubleml/double_ml_qte.py index 8f2286d1..924ca2dd 100644 --- a/doubleml/double_ml_qte.py +++ b/doubleml/double_ml_qte.py @@ -92,23 +92,25 @@ class DoubleMLQTE: 0.50 0.449150 0.192539 2.332782 0.019660 0.071782 0.826519 0.75 0.709606 0.193308 3.670867 0.000242 0.330731 1.088482 """ - def __init__(self, - obj_dml_data, - ml_g, - ml_m=None, - quantiles=0.5, - n_folds=5, - n_rep=1, - score='PQ', - dml_procedure='dml2', - normalize_ipw=True, - kde=None, - trimming_rule='truncate', - trimming_threshold=1e-2, - draw_sample_splitting=True): + def __init__( + self, + obj_dml_data, + ml_g, + ml_m=None, + quantiles=0.5, + n_folds=5, + n_rep=1, + score="PQ", + dml_procedure="dml2", + normalize_ipw=True, + kde=None, + trimming_rule="truncate", + trimming_threshold=1e-2, + draw_sample_splitting=True, + ): self._dml_data = obj_dml_data - self._quantiles = np.asarray(quantiles).reshape((-1, )) + self._quantiles = np.asarray(quantiles).reshape((-1,)) self._check_quantile() self._n_quantiles = len(self._quantiles) @@ -116,8 +118,7 @@ def __init__(self, self._kde = _default_kde else: if not callable(kde): - raise TypeError('kde should be either a callable or None. ' - '%r was passed.' % kde) + raise TypeError("kde should be either a callable or None. " "%r was passed." % kde) self._kde = kde self._normalize_ipw = normalize_ipw @@ -127,7 +128,7 @@ def __init__(self, # check score self._score = score - valid_scores = ['PQ', 'LPQ', 'CVaR'] + valid_scores = ["PQ", "LPQ", "CVaR"] _check_score(self.score, valid_scores, allow_callable=False) # check data @@ -143,8 +144,9 @@ def __init__(self, self._check_quantile() if not isinstance(self.normalize_ipw, bool): - raise TypeError('Normalization indicator has to be boolean. ' + - f'Object of type {str(type(self.normalize_ipw))} passed.') + raise TypeError( + "Normalization indicator has to be boolean. " + f"Object of type {str(type(self.normalize_ipw))} passed." + ) # todo add crossfitting = False self._apply_cross_fitting = True @@ -154,14 +156,14 @@ def __init__(self, if draw_sample_splitting: self.draw_sample_splitting() - self._learner = {'ml_g': clone(ml_g), 'ml_m': clone(ml_m)} - self._predict_method = {'ml_g': 'predict_proba', 'ml_m': 'predict_proba'} + self._learner = {"ml_g": clone(ml_g), "ml_m": clone(ml_m)} + self._predict_method = {"ml_g": "predict_proba", "ml_m": "predict_proba"} # initialize all models self._modellist_0, self._modellist_1 = self._initialize_models() # initialize arrays according to obj_dml_data and the resampling settings - self._psi0, self._psi1, self._psi0_deriv, self._psi1_deriv,\ + self._psi0, self._psi1, self._psi0_deriv, self._psi1_deriv, \ self._coef, self._se, self._all_coef, self._all_se, self._all_dml1_coef = self._initialize_arrays() # also initialize bootstrap arrays with the default number of bootstrap replications @@ -169,10 +171,9 @@ def __init__(self, def __str__(self): class_name = self.__class__.__name__ - header = f'================== {class_name} Object ==================\n' + header = f"================== {class_name} Object ==================\n" fit_summary = str(self.summary) - res = header + \ - '\n------------------ Fit summary ------------------\n' + fit_summary + res = header + "\n------------------ Fit summary ------------------\n" + fit_summary return res @property @@ -195,8 +196,10 @@ def smpls(self): The partition used for cross-fitting. """ if self._smpls is None: - err_msg = ('Sample splitting not specified. Draw samples via .draw_sample splitting(). ' + - 'External samples not implemented yet.') + err_msg = ( + "Sample splitting not specified. Draw samples via .draw_sample splitting(). " + + "External samples not implemented yet." + ) raise ValueError(err_msg) return self._smpls @@ -348,16 +351,12 @@ def summary(self): """ A summary for the estimated causal effect after calling :meth:`fit`. """ - col_names = ['coef', 'std err', 't', 'P>|t|'] + col_names = ["coef", "std err", "t", "P>|t|"] if np.isnan(self.coef).all(): df_summary = pd.DataFrame(columns=col_names) else: - summary_stats = np.transpose(np.vstack( - [self.coef, self.se, - self.t_stat, self.pval])) - df_summary = pd.DataFrame(summary_stats, - columns=col_names, - index=self.quantiles) + summary_stats = np.transpose(np.vstack([self.coef, self.se, self.t_stat, self.pval])) + df_summary = pd.DataFrame(summary_stats, columns=col_names, index=self.quantiles) ci = self.confint() df_summary = df_summary.join(ci) return df_summary @@ -414,7 +413,7 @@ def fit(self, n_jobs_models=None, n_jobs_cv=None, store_predictions=True, store_ ------- self : object """ - + if external_predictions is not None: raise NotImplementedError(f"External predictions not implemented for {self.__class__.__name__}.") @@ -423,7 +422,6 @@ def fit(self, n_jobs_models=None, n_jobs_cv=None, store_predictions=True, store_ fitted_models = parallel(delayed(self._fit_quantile)(i_quant, n_jobs_cv, store_predictions, store_models) for i_quant in range(self.n_quantiles)) - # combine the estimates and scores for i_quant in range(self.n_quantiles): self._i_quant = i_quant @@ -432,8 +430,9 @@ def fit(self, n_jobs_models=None, n_jobs_cv=None, store_predictions=True, store_ self._modellist_1[self._i_quant] = fitted_models[self._i_quant][1] # treatment Effects - self._all_coef[self._i_quant, :] = self.modellist_1[self._i_quant].all_coef - \ - self.modellist_0[self._i_quant].all_coef + self._all_coef[self._i_quant, :] = ( + self.modellist_1[self._i_quant].all_coef - self.modellist_0[self._i_quant].all_coef + ) # save scores and derivatives self._psi0[:, :, self._i_quant] = np.squeeze(self.modellist_0[i_quant].psi, 2) @@ -453,7 +452,7 @@ def fit(self, n_jobs_models=None, n_jobs_cv=None, store_predictions=True, store_ return self - def bootstrap(self, method='normal', n_rep_boot=500): + def bootstrap(self, method="normal", n_rep_boot=500): """ Multiplier bootstrap for DoubleML models. @@ -471,18 +470,18 @@ def bootstrap(self, method='normal', n_rep_boot=500): self : object """ if np.isnan(self.coef).all(): - raise ValueError('Apply fit() before bootstrap().') + raise ValueError("Apply fit() before bootstrap().") - if (not isinstance(method, str)) | (method not in ['Bayes', 'normal', 'wild']): - raise ValueError('Method must be "Bayes", "normal" or "wild". ' - f'Got {str(method)}.') + if (not isinstance(method, str)) | (method not in ["Bayes", "normal", "wild"]): + raise ValueError('Method must be "Bayes", "normal" or "wild". ' f"Got {str(method)}.") if not isinstance(n_rep_boot, int): - raise TypeError('The number of bootstrap replications must be of int type. ' - f'{str(n_rep_boot)} of type {str(type(n_rep_boot))} was passed.') + raise TypeError( + "The number of bootstrap replications must be of int type. " + f"{str(n_rep_boot)} of type {str(type(n_rep_boot))} was passed." + ) if n_rep_boot < 1: - raise ValueError('The number of bootstrap replications must be positive. ' - f'{str(n_rep_boot)} was passed.') + raise ValueError("The number of bootstrap replications must be positive. " f"{str(n_rep_boot)} was passed.") self._n_rep_boot, self._boot_coef, self._boot_t_stat = self._initialize_boot_arrays(n_rep_boot) @@ -503,8 +502,10 @@ def bootstrap(self, method='normal', n_rep_boot=500): self._i_quant = i_quant i_start = self._i_rep * self.n_rep_boot i_end = (self._i_rep + 1) * self.n_rep_boot - self._boot_coef[self._i_quant, i_start:i_end], self._boot_t_stat[self._i_quant, i_start:i_end] =\ - self._compute_bootstrap(weights) + ( + self._boot_coef[self._i_quant, i_start:i_end], + self._boot_t_stat[self._i_quant, i_start:i_end], + ) = self._compute_bootstrap(weights) return self def draw_sample_splitting(self): @@ -518,11 +519,13 @@ def draw_sample_splitting(self): ------- self : object """ - obj_dml_resampling = DoubleMLResampling(n_folds=self.n_folds, - n_rep=self.n_rep, - n_obs=self._dml_data.n_obs, - apply_cross_fitting=self.apply_cross_fitting, - stratify=self._dml_data.d) + obj_dml_resampling = DoubleMLResampling( + n_folds=self.n_folds, + n_rep=self.n_rep, + n_obs=self._dml_data.n_obs, + apply_cross_fitting=self.apply_cross_fitting, + stratify=self._dml_data.d, + ) self._smpls = obj_dml_resampling.split_samples() return self @@ -568,37 +571,33 @@ def confint(self, joint=False, level=0.95): """ if not isinstance(joint, bool): - raise TypeError('joint must be True or False. ' - f'Got {str(joint)}.') + raise TypeError("joint must be True or False. " f"Got {str(joint)}.") if not isinstance(level, float): - raise TypeError('The confidence level must be of float type. ' - f'{str(level)} of type {str(type(level))} was passed.') + raise TypeError( + "The confidence level must be of float type. " f"{str(level)} of type {str(type(level))} was passed." + ) if (level <= 0) | (level >= 1): - raise ValueError('The confidence level must be in (0,1). ' - f'{str(level)} was passed.') + raise ValueError("The confidence level must be in (0,1). " f"{str(level)} was passed.") - a = (1 - level) - ab = np.array([a / 2, 1. - a / 2]) + a = 1 - level + ab = np.array([a / 2, 1.0 - a / 2]) if joint: if np.isnan(self.boot_coef).all(): - raise ValueError('Apply fit() & bootstrap() before confint(joint=True).') + raise ValueError("Apply fit() & bootstrap() before confint(joint=True).") sim = np.amax(np.abs(self.boot_t_stat), 0) hatc = np.quantile(sim, 1 - a) ci = np.vstack((self.coef - self.se * hatc, self.coef + self.se * hatc)).T else: if np.isnan(self.coef).all(): - raise ValueError('Apply fit() before confint().') + raise ValueError("Apply fit() before confint().") fac = norm.ppf(ab) ci = np.vstack((self.coef + self.se * fac[0], self.coef + self.se * fac[1])).T - df_ci = pd.DataFrame(ci, - columns=['{:.1f} %'.format(i * 100) for i in ab], - index=self._quantiles) + df_ci = pd.DataFrame(ci, columns=["{:.1f} %".format(i * 100) for i in ab], index=self._quantiles) return df_ci def _fit_quantile(self, i_quant, n_jobs_cv=None, store_predictions=True, store_models=False): - model_0 = self.modellist_0[i_quant] model_1 = self.modellist_1[i_quant] @@ -617,8 +616,14 @@ def _agg_cross_fit(self): # TODO: In the edge case of repeated no-cross-fitting, the test sets might have different size and therefore # it would note be valid to always use the same self._var_scaling_factor xx = np.tile(self.coef.reshape(-1, 1), self.n_rep) - self.se = np.sqrt(np.divide(np.median(np.multiply(np.power(self._all_se, 2), self._var_scaling_factor) + - np.power(self._all_coef - xx, 2), 1), self._var_scaling_factor)) + self.se = np.sqrt( + np.divide( + np.median( + np.multiply(np.power(self._all_se, 2), self._var_scaling_factor) + np.power(self._all_coef - xx, 2), 1 + ), + self._var_scaling_factor, + ) + ) def _var_est(self): """ @@ -654,7 +659,7 @@ def _initialize_arrays(self): all_coef = np.full((self.n_quantiles, self.n_rep), np.nan) all_se = np.full((self.n_quantiles, self.n_rep), np.nan) - if self.dml_procedure == 'dml1': + if self.dml_procedure == "dml1": if self.apply_cross_fitting: all_dml1_coef = np.full((self.n_quantiles, self.n_rep, self.n_folds), np.nan) else: @@ -671,15 +676,15 @@ def _initialize_boot_arrays(self, n_rep_boot): def _check_data(self, obj_dml_data): if not isinstance(obj_dml_data, DoubleMLData): - raise TypeError('The data must be of DoubleMLData type. ' - f'{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed.') + raise TypeError( + "The data must be of DoubleMLData type. " f"{str(obj_dml_data)} of type {str(type(obj_dml_data))} was passed." + ) _check_zero_one_treatment(self) return def _check_quantile(self): if np.any(self.quantiles <= 0) | np.any(self.quantiles >= 1): - raise ValueError('Quantiles have be between 0 or 1. ' + - f'Quantiles {str(self.quantiles)} passed.') + raise ValueError("Quantiles have be between 0 or 1. " + f"Quantiles {str(self.quantiles)} passed.") def _initialize_models(self): modellist_0 = [None] * self.n_quantiles @@ -687,92 +692,104 @@ def _initialize_models(self): for i_quant in range(self.n_quantiles): self._i_quant = i_quant # initialize models for both potential quantiles - if self.score == 'PQ': - model_0 = DoubleMLPQ(self._dml_data, - self._learner['ml_g'], - self._learner['ml_m'], - quantile=self._quantiles[i_quant], - treatment=0, - n_folds=self.n_folds, - n_rep=self.n_rep, - dml_procedure=self.dml_procedure, - trimming_rule=self.trimming_rule, - trimming_threshold=self.trimming_threshold, - kde=self.kde, - normalize_ipw=self.normalize_ipw, - draw_sample_splitting=False, - apply_cross_fitting=self._apply_cross_fitting) - model_1 = DoubleMLPQ(self._dml_data, - self._learner['ml_g'], - self._learner['ml_m'], - quantile=self._quantiles[i_quant], - treatment=1, - n_folds=self.n_folds, - n_rep=self.n_rep, - dml_procedure=self.dml_procedure, - trimming_rule=self.trimming_rule, - trimming_threshold=self.trimming_threshold, - kde=self.kde, - normalize_ipw=self.normalize_ipw, - draw_sample_splitting=False, - apply_cross_fitting=self._apply_cross_fitting) - elif self.score == 'LPQ': - model_0 = DoubleMLLPQ(self._dml_data, - self._learner['ml_g'], - self._learner['ml_m'], - quantile=self._quantiles[i_quant], - treatment=0, - n_folds=self.n_folds, - n_rep=self.n_rep, - dml_procedure=self.dml_procedure, - trimming_rule=self.trimming_rule, - trimming_threshold=self.trimming_threshold, - kde=self.kde, - normalize_ipw=self.normalize_ipw, - draw_sample_splitting=False, - apply_cross_fitting=self._apply_cross_fitting) - model_1 = DoubleMLLPQ(self._dml_data, - self._learner['ml_g'], - self._learner['ml_m'], - quantile=self._quantiles[i_quant], - treatment=1, - n_folds=self.n_folds, - n_rep=self.n_rep, - dml_procedure=self.dml_procedure, - trimming_rule=self.trimming_rule, - trimming_threshold=self.trimming_threshold, - kde=self.kde, - normalize_ipw=self.normalize_ipw, - draw_sample_splitting=False, - apply_cross_fitting=self._apply_cross_fitting) - - elif self.score == 'CVaR': - model_0 = DoubleMLCVAR(self._dml_data, - self._learner['ml_g'], - self._learner['ml_m'], - quantile=self._quantiles[i_quant], - treatment=0, - n_folds=self.n_folds, - n_rep=self.n_rep, - dml_procedure=self.dml_procedure, - trimming_rule=self.trimming_rule, - trimming_threshold=self.trimming_threshold, - normalize_ipw=self.normalize_ipw, - draw_sample_splitting=False, - apply_cross_fitting=self._apply_cross_fitting) - model_1 = DoubleMLCVAR(self._dml_data, - self._learner['ml_g'], - self._learner['ml_m'], - quantile=self._quantiles[i_quant], - treatment=1, - n_folds=self.n_folds, - n_rep=self.n_rep, - dml_procedure=self.dml_procedure, - trimming_rule=self.trimming_rule, - trimming_threshold=self.trimming_threshold, - normalize_ipw=self.normalize_ipw, - draw_sample_splitting=False, - apply_cross_fitting=self._apply_cross_fitting) + if self.score == "PQ": + model_0 = DoubleMLPQ( + self._dml_data, + self._learner["ml_g"], + self._learner["ml_m"], + quantile=self._quantiles[i_quant], + treatment=0, + n_folds=self.n_folds, + n_rep=self.n_rep, + dml_procedure=self.dml_procedure, + trimming_rule=self.trimming_rule, + trimming_threshold=self.trimming_threshold, + kde=self.kde, + normalize_ipw=self.normalize_ipw, + draw_sample_splitting=False, + apply_cross_fitting=self._apply_cross_fitting, + ) + model_1 = DoubleMLPQ( + self._dml_data, + self._learner["ml_g"], + self._learner["ml_m"], + quantile=self._quantiles[i_quant], + treatment=1, + n_folds=self.n_folds, + n_rep=self.n_rep, + dml_procedure=self.dml_procedure, + trimming_rule=self.trimming_rule, + trimming_threshold=self.trimming_threshold, + kde=self.kde, + normalize_ipw=self.normalize_ipw, + draw_sample_splitting=False, + apply_cross_fitting=self._apply_cross_fitting, + ) + elif self.score == "LPQ": + model_0 = DoubleMLLPQ( + self._dml_data, + self._learner["ml_g"], + self._learner["ml_m"], + quantile=self._quantiles[i_quant], + treatment=0, + n_folds=self.n_folds, + n_rep=self.n_rep, + dml_procedure=self.dml_procedure, + trimming_rule=self.trimming_rule, + trimming_threshold=self.trimming_threshold, + kde=self.kde, + normalize_ipw=self.normalize_ipw, + draw_sample_splitting=False, + apply_cross_fitting=self._apply_cross_fitting, + ) + model_1 = DoubleMLLPQ( + self._dml_data, + self._learner["ml_g"], + self._learner["ml_m"], + quantile=self._quantiles[i_quant], + treatment=1, + n_folds=self.n_folds, + n_rep=self.n_rep, + dml_procedure=self.dml_procedure, + trimming_rule=self.trimming_rule, + trimming_threshold=self.trimming_threshold, + kde=self.kde, + normalize_ipw=self.normalize_ipw, + draw_sample_splitting=False, + apply_cross_fitting=self._apply_cross_fitting, + ) + + elif self.score == "CVaR": + model_0 = DoubleMLCVAR( + self._dml_data, + self._learner["ml_g"], + self._learner["ml_m"], + quantile=self._quantiles[i_quant], + treatment=0, + n_folds=self.n_folds, + n_rep=self.n_rep, + dml_procedure=self.dml_procedure, + trimming_rule=self.trimming_rule, + trimming_threshold=self.trimming_threshold, + normalize_ipw=self.normalize_ipw, + draw_sample_splitting=False, + apply_cross_fitting=self._apply_cross_fitting, + ) + model_1 = DoubleMLCVAR( + self._dml_data, + self._learner["ml_g"], + self._learner["ml_m"], + quantile=self._quantiles[i_quant], + treatment=1, + n_folds=self.n_folds, + n_rep=self.n_rep, + dml_procedure=self.dml_procedure, + trimming_rule=self.trimming_rule, + trimming_threshold=self.trimming_threshold, + normalize_ipw=self.normalize_ipw, + draw_sample_splitting=False, + apply_cross_fitting=self._apply_cross_fitting, + ) # synchronize the sample splitting model_0.set_sample_splitting(all_smpls=self.smpls) diff --git a/doubleml/tests/test_did_external_predictions.py b/doubleml/tests/test_did_external_predictions.py index 12d7e3c9..fa6f325e 100644 --- a/doubleml/tests/test_did_external_predictions.py +++ b/doubleml/tests/test_did_external_predictions.py @@ -1,12 +1,13 @@ import numpy as np import pytest import math -from sklearn.linear_model import LinearRegression, LassoCV, LogisticRegression -from doubleml import DoubleMLData, DoubleMLDID +from sklearn.linear_model import LinearRegression, LogisticRegression +from doubleml import DoubleMLDID from doubleml.datasets import make_did_SZ2020 from doubleml.utils import dummy_regressor, dummy_classifier from ._utils import draw_smpls + @pytest.fixture(scope="module", params=["observational", "experimental"]) def did_score(request): return request.param @@ -32,7 +33,7 @@ def doubleml_did_fixture(did_score, dml_procedure, n_rep): "score": did_score, "n_rep": n_rep, "dml_procedure": dml_procedure, - "draw_sample_splitting": False + "draw_sample_splitting": False, } DMLDID = DoubleMLDID(ml_g=LinearRegression(), ml_m=LogisticRegression(), **kwargs) DMLDID.set_sample_splitting(all_smpls) diff --git a/doubleml/tests/test_didcs_external_predictions.py b/doubleml/tests/test_didcs_external_predictions.py index 0eed900a..ee578afb 100644 --- a/doubleml/tests/test_didcs_external_predictions.py +++ b/doubleml/tests/test_didcs_external_predictions.py @@ -1,8 +1,8 @@ import numpy as np import pytest import math -from sklearn.linear_model import LinearRegression, LassoCV, LogisticRegression -from doubleml import DoubleMLData, DoubleMLDIDCS +from sklearn.linear_model import LinearRegression, LogisticRegression +from doubleml import DoubleMLDIDCS from doubleml.datasets import make_did_SZ2020 from doubleml.utils import dummy_regressor, dummy_classifier from ._utils import draw_smpls @@ -34,7 +34,7 @@ def doubleml_didcs_fixture(did_score, dml_procedure, n_rep): "n_rep": n_rep, "n_folds": 5, "dml_procedure": dml_procedure, - "draw_sample_splitting": False + "draw_sample_splitting": False, } DMLDIDCS = DoubleMLDIDCS(ml_g=LinearRegression(), ml_m=LogisticRegression(), **kwargs) DMLDIDCS.set_sample_splitting(all_smpls) diff --git a/doubleml/tests/test_doubleml_exceptions.py b/doubleml/tests/test_doubleml_exceptions.py index b840cc9b..98a46fc0 100644 --- a/doubleml/tests/test_doubleml_exceptions.py +++ b/doubleml/tests/test_doubleml_exceptions.py @@ -2,7 +2,7 @@ import pandas as pd import numpy as np -from doubleml import DoubleMLPLR, DoubleMLIRM, DoubleMLIIVM, DoubleMLPLIV, DoubleMLData,\ +from doubleml import DoubleMLPLR, DoubleMLIRM, DoubleMLIIVM, DoubleMLPLIV, DoubleMLData, \ DoubleMLClusterData, DoubleMLPQ, DoubleMLLPQ, DoubleMLCVAR, DoubleMLQTE, DoubleMLDID, DoubleMLDIDCS from doubleml.datasets import make_plr_CCDDHNR2018, make_irm_data, make_pliv_CHS2015, make_iivm_data, \ make_pliv_multiway_cluster_CKMS2021, make_did_SZ2020 diff --git a/doubleml/tests/test_doubleml_external_prediction_implementation.py b/doubleml/tests/test_doubleml_exceptions_external_pred.py similarity index 90% rename from doubleml/tests/test_doubleml_external_prediction_implementation.py rename to doubleml/tests/test_doubleml_exceptions_external_pred.py index 9d082859..701b696c 100644 --- a/doubleml/tests/test_doubleml_external_prediction_implementation.py +++ b/doubleml/tests/test_doubleml_exceptions_external_pred.py @@ -1,10 +1,9 @@ -import numpy as np import pytest from doubleml import DoubleMLCVAR, DoubleMLQTE, DoubleMLData from doubleml.datasets import make_irm_data from doubleml.utils import dummy_regressor, dummy_classifier -df_irm = make_irm_data(n_obs=500, dim_x=20, theta=0.5, return_type="DataFrame") +df_irm = make_irm_data(n_obs=10, dim_x=10, theta=0.5, return_type="DataFrame") ext_predictions = {"d": {}} diff --git a/doubleml/tests/test_dummy_learners.py b/doubleml/tests/test_dummy_learners.py index ee3d979a..a357345c 100644 --- a/doubleml/tests/test_dummy_learners.py +++ b/doubleml/tests/test_dummy_learners.py @@ -42,5 +42,5 @@ def test_clone(dl_fixture): try: _ = clone(dl_fixture["dummy_regressor"]) _ = clone(dl_fixture["dummy_classifier"]) - except Error as e: + except Exception as e: pytest.fail(f"clone() raised an exception:\n{str(e)}\n") diff --git a/doubleml/tests/test_iivm_external_predictions.py b/doubleml/tests/test_iivm_external_predictions.py index 40bb02db..d934eb02 100644 --- a/doubleml/tests/test_iivm_external_predictions.py +++ b/doubleml/tests/test_iivm_external_predictions.py @@ -1,7 +1,7 @@ import numpy as np import pytest import math -from sklearn.linear_model import LinearRegression, LassoCV, LogisticRegression +from sklearn.linear_model import LinearRegression, LogisticRegression from doubleml import DoubleMLIIVM, DoubleMLData from doubleml.datasets import make_iivm_data from doubleml.utils import dummy_regressor, dummy_classifier @@ -21,9 +21,7 @@ def n_rep(request): def adapted_doubleml_fixture(dml_procedure, n_rep): ext_predictions = {"d": {}} - data = make_iivm_data( - n_obs=500, dim_x=20, theta=0.5, alpha_x=1.0, return_type="DataFrame" - ) + data = make_iivm_data(n_obs=500, dim_x=20, theta=0.5, alpha_x=1.0, return_type="DataFrame") np.random.seed(3141) @@ -45,14 +43,13 @@ def adapted_doubleml_fixture(dml_procedure, n_rep): np.random.seed(3141) DMLIIVM.fit(store_predictions=True) - + ext_predictions["d"]["ml_g0"] = DMLIIVM.predictions["ml_g0"][:, :, 0] ext_predictions["d"]["ml_g1"] = DMLIIVM.predictions["ml_g1"][:, :, 0] ext_predictions["d"]["ml_m"] = DMLIIVM.predictions["ml_m"][:, :, 0] ext_predictions["d"]["ml_r0"] = DMLIIVM.predictions["ml_r0"][:, :, 0] ext_predictions["d"]["ml_r1"] = DMLIIVM.predictions["ml_r1"][:, :, 0] - DMLIIVM_ext = DoubleMLIIVM( ml_g=dummy_regressor(), ml_m=dummy_classifier(), ml_r=dummy_classifier(), **kwargs ) diff --git a/doubleml/tests/test_lpq_external_predictions.py b/doubleml/tests/test_lpq_external_predictions.py index 2a13b4bc..a5a9a5bb 100644 --- a/doubleml/tests/test_lpq_external_predictions.py +++ b/doubleml/tests/test_lpq_external_predictions.py @@ -4,7 +4,7 @@ from sklearn.linear_model import LogisticRegression from doubleml import DoubleMLLPQ, DoubleMLData from doubleml.datasets import make_iivm_data -from doubleml.utils import dummy_regressor, dummy_classifier +from doubleml.utils import dummy_classifier from ._utils import draw_smpls diff --git a/doubleml/tests/test_pliv_external_predictions.py b/doubleml/tests/test_pliv_external_predictions.py index cbd13dfe..a9029477 100644 --- a/doubleml/tests/test_pliv_external_predictions.py +++ b/doubleml/tests/test_pliv_external_predictions.py @@ -1,7 +1,7 @@ import numpy as np import pytest import math -from sklearn.linear_model import LinearRegression, LassoCV +from sklearn.linear_model import LinearRegression from doubleml import DoubleMLPLIV, DoubleMLData from doubleml.datasets import make_pliv_CHS2015 from doubleml.utils import dummy_regressor @@ -32,12 +32,11 @@ def adapted_doubleml_fixture(score, dml_procedure, n_rep, dim_z): # IV-type score only allows dim_z = 1, so skip testcases with dim_z > 1 for IV-type score if dim_z > 1 and score == "IV-type": pytest.skip("IV-type score only allows dim_z = 1") + res_dict = None else: ext_predictions = {"d": {}} - data = make_pliv_CHS2015( - n_obs=500, dim_x=20, alpha=0.5, dim_z=dim_z, return_type="DataFrame" - ) + data = make_pliv_CHS2015(n_obs=500, dim_x=20, alpha=0.5, dim_z=dim_z, return_type="DataFrame") np.random.seed(3141) @@ -77,16 +76,14 @@ def adapted_doubleml_fixture(score, dml_procedure, n_rep, dim_z): ml_m_key = "ml_m_" + "Z" + str(instr + 1) ext_predictions["d"][ml_m_key] = DMLPLIV.predictions[ml_m_key][:, :, 0] - DMLPLIV_ext = DoubleMLPLIV( - ml_m=dummy_regressor(), ml_l=dummy_regressor(), ml_r=dummy_regressor(), **kwargs - ) + DMLPLIV_ext = DoubleMLPLIV(ml_m=dummy_regressor(), ml_l=dummy_regressor(), ml_r=dummy_regressor(), **kwargs) np.random.seed(3141) DMLPLIV_ext.fit(external_predictions=ext_predictions) res_dict = {"coef_normal": DMLPLIV.coef, "coef_ext": DMLPLIV_ext.coef} - return res_dict + return res_dict @pytest.mark.ci diff --git a/doubleml/tests/test_pq_external_predictions.py b/doubleml/tests/test_pq_external_predictions.py index 663a0def..d7c1365e 100644 --- a/doubleml/tests/test_pq_external_predictions.py +++ b/doubleml/tests/test_pq_external_predictions.py @@ -4,7 +4,7 @@ from sklearn.linear_model import LogisticRegression from doubleml import DoubleMLPQ, DoubleMLData from doubleml.datasets import make_irm_data -from doubleml.utils import dummy_regressor, dummy_classifier +from doubleml.utils import dummy_classifier from ._utils import draw_smpls @@ -79,7 +79,8 @@ def doubleml_pq_fixture(dml_procedure, n_rep, normalize_ipw, set_ml_m_ext, set_m if set_ml_m_ext and not set_ml_g_ext: # adjust tolerance for the case that ml_m is set to external predictions - # because no preliminary results are available for ml_m, the model use the (external) final predictions for ml_m + # because no preliminary results are available for ml_m, + # the model use the (external) final predictions for ml_m for calculating the ipw estimate tol_rel = 0.1 tol_abs = 0.1 else: diff --git a/doubleml/utils/dummy_learners.py b/doubleml/utils/dummy_learners.py index 2f893fb2..31e32e25 100644 --- a/doubleml/utils/dummy_learners.py +++ b/doubleml/utils/dummy_learners.py @@ -2,6 +2,23 @@ class dummy_regressor(BaseEstimator): + """ + A dummy regressor that raises an AttributeError when attempting to access + its fit, predict, or set_params methods. + Attributes + ---------- + _estimator_type : str + Type of the estimator, set to "regressor". + Methods + ------- + fit(*args) + Raises AttributeError: "Accessed fit method of DummyRegressor!" + predict(*args) + Raises AttributeError: "Accessed predict method of DummyRegressor!" + set_params(*args) + Raises AttributeError: "Accessed set_params method of DummyRegressor!" + """ + _estimator_type = "regressor" def fit(*args): @@ -15,6 +32,21 @@ def set_params(*args): class dummy_classifier(BaseEstimator): + """ + A dummy classifier that raises an AttributeError when attempting to access + its fit, predict, set_params, or predict_proba methods. + Attributes + ---------- + _estimator_type : str + Type of the estimator, set to "classifier". + predict(*args) + Raises AttributeError: "Accessed predict method of DummyClassifier!" + set_params(*args) + Raises AttributeError: "Accessed set_params method of DummyClassifier!" + predict_proba(*args, **kwargs) + Raises AttributeError: "Accessed predict_proba method of DummyClassifier!" + """ + _estimator_type = "classifier" def fit(*args):