"""Treatment effect estimators follows largely Stata's teffects in Stata 13 manual Created on Tue Jun 9 22:45:23 2015 Author: Josef Perktold License: BSD-3 currently available ATE POM_0 POM_1 res_ipw 230.688598 3172.774059 3403.462658 res_aipw -230.989201 3403.355253 3172.366052 res_aipw_wls -227.195618 3403.250651 3176.055033 res_ra -239.639211 3403.242272 3163.603060 res_ipwra -229.967078 3403.335639 3173.368561 Lots of todos, just the beginning, but most effects are available but not standard errors, and no code structure that has a useful pattern see https://github.com/statsmodels/statsmodels/issues/2443 Note: script requires cattaneo2 data file from Stata 14, hardcoded file path could be loaded with webuse """ import numpy as np from statsmodels.compat.pandas import Substitution from scipy.linalg import block_diag from statsmodels.regression.linear_model import WLS from statsmodels.sandbox.regression.gmm import GMM from statsmodels.stats.contrast import ContrastResults from statsmodels.tools.docstring import indent def _mom_ate(params, endog, tind, prob, weighted=True): """moment condition for average treatment effect This does not include a moment condition for potential outcome mean (POM). """ w1 = (tind / prob) w0 = (1. - tind) / (1. - prob) if weighted: w0 /= w0.mean() w1 /= w1.mean() wdiff = w1 - w0 return endog * wdiff - params def _mom_atm(params, endog, tind, prob, weighted=True): """moment conditions for average treatment means (POM) moment conditions are POM0 and POM1 """ w1 = (tind / prob) w0 = (1. - tind) / (1. - prob) if weighted: w1 /= w1.mean() w0 /= w0.mean() return np.column_stack((endog * w0 - params[0], endog * w1 - params[1])) def _mom_ols(params, endog, tind, prob, weighted=True): """ moment condition for average treatment mean based on OLS dummy regression moment conditions are POM0 and POM1 """ w = tind / prob + (1-tind) / (1 - prob) treat_ind = np.column_stack((1 - tind, tind)) mom = (w * (endog - treat_ind.dot(params)))[:, None] * treat_ind return mom def _mom_ols_te(tm, endog, tind, prob, weighted=True): """ moment condition for average treatment mean based on OLS dummy regression first moment is ATE second moment is POM0 (control) """ w = tind / prob + (1-tind) / (1 - prob) treat_ind = np.column_stack((tind, np.ones(len(tind)))) mom = (w * (endog - treat_ind.dot(tm)))[:, None] * treat_ind return mom def _mom_olsex(params, model=None, exog=None, scale=None): exog = exog if exog is not None else model.exog fitted = model.predict(params, exog) resid = model.endog - fitted if scale is not None: resid /= scale mom = resid[:, None] * exog return mom def ate_ipw(endog, tind, prob, weighted=True, probt=None): """average treatment effect based on basic inverse propensity weighting. """ w1 = (tind / prob) w0 = (1. - tind) / (1. - prob) if probt is not None: w1 *= probt w0 *= probt if weighted: w0 /= w0.mean() w1 /= w1.mean() wdiff = w1 - w0 return (endog * wdiff).mean(), (endog * w0).mean(), (endog * w1).mean() class _TEGMMGeneric1(GMM): """GMM class to get cov_params for treatment effects This combines moment conditions for the selection/treatment model and the outcome model to get the standard errors for the treatment effect that takes the first step estimation of the treatment model into account. this also matches standard errors of ATE and POM in Stata """ def __init__(self, endog, res_select, mom_outcome, exclude_tmoms=False, **kwargs): super().__init__(endog, None, None) self.results_select = res_select self.mom_outcome = mom_outcome self.exclude_tmoms = exclude_tmoms self.__dict__.update(kwargs) # add xnames so it's not None # we don't have exog in init in this version if self.data.xnames is None: self.data.xnames = [] # need information about decomposition of parameters if exclude_tmoms: self.k_select = 0 else: self.k_select = len(res_select.model.data.param_names) if exclude_tmoms: # fittedvalues is still linpred self.prob = self.results_select.predict() else: self.prob = None def momcond(self, params): k_outcome = len(params) - self.k_select tm = params[:k_outcome] p_tm = params[k_outcome:] tind = self.results_select.model.endog if self.exclude_tmoms: prob = self.prob else: prob = self.results_select.model.predict(p_tm) moms_list = [] mom_o = self.mom_outcome(tm, self.endog, tind, prob, weighted=True) moms_list.append(mom_o) if not self.exclude_tmoms: mom_t = self.results_select.model.score_obs(p_tm) moms_list.append(mom_t) moms = np.column_stack(moms_list) return moms class _TEGMM(GMM): """GMM class to get cov_params for treatment effects This combines moment conditions for the selection/treatment model and the outcome model to get the standard errors for the treatment effect that takes the first step estimation of the treatment model into account. this also matches standard errors of ATE and POM in Stata """ def __init__(self, endog, res_select, mom_outcome): super().__init__(endog, None, None) self.results_select = res_select self.mom_outcome = mom_outcome # add xnames so it's not None # we don't have exog in init in this version if self.data.xnames is None: self.data.xnames = [] def momcond(self, params): tm = params[:2] p_tm = params[2:] tind = self.results_select.model.endog prob = self.results_select.model.predict(p_tm) momt = self.mom_outcome(tm, self.endog, tind, prob) # weighted=True) moms = np.column_stack((momt, self.results_select.model.score_obs(p_tm))) return moms class _IPWGMM(_TEGMMGeneric1): """ GMM for aipw treatment effect and potential outcome uses unweighted outcome regression """ def momcond(self, params): # Note: momcond in original order of observations ra = self.teff res_select = ra.results_select tind = ra.treatment endog = ra.model_pool.endog effect_group = self.effect_group tm = params[:2] ps = params[2:] prob_sel = np.asarray(res_select.model.predict(ps)) prob_sel = np.clip(prob_sel, 0.01, 0.99) prob = prob_sel if effect_group == "all": probt = None elif effect_group in [1, "treated"]: probt = prob elif effect_group in [0, "untreated", "control"]: probt = 1 - prob elif isinstance(effect_group, np.ndarray): probt = probt else: raise ValueError("incorrect option for effect_group") w = tind / prob + (1 - tind) / (1 - prob) # Are we supposed to use scaled weights? doesn't cloesely match Stata # w1 = tind / prob # w2 = (1 - tind) / (1 - prob) # w = w1 / w1.sum() * tind.sum() + w2 / w2.sum() * (1 - tind).sum() if probt is not None: w *= probt treat_ind = np.column_stack((tind, np.ones(len(tind)))) mm = (w * (endog - treat_ind.dot(tm)))[:, None] * treat_ind mom_select = res_select.model.score_obs(ps) moms = np.column_stack((mm, mom_select)) return moms class _AIPWGMM(_TEGMMGeneric1): """ GMM for aipw treatment effect and potential outcome uses unweighted outcome regression """ def momcond(self, params): ra = self.teff treat_mask = ra.treat_mask res_select = ra.results_select ppom = params[1] mask = np.arange(len(params)) != 1 params = params[mask] k = ra.results0.model.exog.shape[1] pm = params[0] # ATE parameter p0 = params[1:k+1] p1 = params[k+1:2*k+1] ps = params[2*k+1:] mod0 = ra.results0.model mod1 = ra.results1.model # use reordered exog, endog so it matches sub models by group exog = ra.exog_grouped endog = ra.endog_grouped prob_sel = np.asarray(res_select.model.predict(ps)) prob_sel = np.clip(prob_sel, 0.01, 0.99) prob0 = prob_sel[~treat_mask] prob1 = prob_sel[treat_mask] prob = np.concatenate((prob0, prob1)) # outcome models by treatment unweighted fitted0 = mod0.predict(p0, exog) mom0 = _mom_olsex(p0, model=mod0) fitted1 = mod1.predict(p1, exog) mom1 = _mom_olsex(p1, model=mod1) mom_outcome = block_diag(mom0, mom1) # moments for target statistics, ATE and POM tind = ra.treatment tind = np.concatenate((tind[~treat_mask], tind[treat_mask])) correct0 = (endog - fitted0) / (1 - prob) * (1 - tind) correct1 = (endog - fitted1) / prob * tind tmean0 = fitted0 + correct0 tmean1 = fitted1 + correct1 ate = tmean1 - tmean0 mm = ate - pm mpom = tmean0 - ppom mm = np.column_stack((mm, mpom)) # Note: res_select has original data order, # mom_outcome and mm use grouped observations mom_select = res_select.model.score_obs(ps) mom_select = np.concatenate((mom_select[~treat_mask], mom_select[treat_mask]), axis=0) moms = np.column_stack((mm, mom_outcome, mom_select)) return moms class _AIPWWLSGMM(_TEGMMGeneric1): """ GMM for aipw-wls treatment effect and potential outcome uses weighted outcome regression """ def momcond(self, params): ra = self.teff treat_mask = ra.treat_mask res_select = ra.results_select ppom = params[1] mask = np.arange(len(params)) != 1 params = params[mask] k = ra.results0.model.exog.shape[1] pm = params[0] # ATE parameter p0 = params[1:k+1] p1 = params[k+1:2*k+1] ps = params[-6:] mod0 = ra.results0.model mod1 = ra.results1.model # use reordered exog, endog so it matches sub models by group exog = ra.exog_grouped endog = ra.endog_grouped # todo: need weights in outcome models prob_sel = np.asarray(res_select.model.predict(ps)) prob_sel = np.clip(prob_sel, 0.001, 0.999) prob0 = prob_sel[~treat_mask] prob1 = prob_sel[treat_mask] prob = np.concatenate((prob0, prob1)) tind = 0 ww0 = (1 - tind) / (1 - prob0) * ((1 - tind) / (1 - prob0) - 1) tind = 1 ww1 = tind / prob1 * (tind / prob1 - 1) # outcome models by treatment using IPW weights fitted0 = mod0.predict(p0, exog) mom0 = _mom_olsex(p0, model=mod0) * ww0[:, None] fitted1 = mod1.predict(p1, exog) mom1 = _mom_olsex(p1, model=mod1) * ww1[:, None] mom_outcome = block_diag(mom0, mom1) # moments for target statistics, ATE and POM tind = ra.treatment tind = np.concatenate((tind[~treat_mask], tind[treat_mask])) correct0 = (endog - fitted0) / (1 - prob) * (1 - tind) correct1 = (endog - fitted1) / prob * tind tmean0 = fitted0 + correct0 tmean1 = fitted1 + correct1 ate = tmean1 - tmean0 mm = ate - pm mpom = tmean0 - ppom mm = np.column_stack((mm, mpom)) # Note: res_select has original data order, # mom_outcome and mm use grouped observations mom_select = res_select.model.score_obs(ps) mom_select = np.concatenate((mom_select[~treat_mask], mom_select[treat_mask]), axis=0) moms = np.column_stack((mm, mom_outcome, mom_select)) return moms class _RAGMM(_TEGMMGeneric1): """GMM for regression adjustment treatment effect and potential outcome uses unweighted outcome regression """ def momcond(self, params): ra = self.teff ppom = params[1] mask = np.arange(len(params)) != 1 params = params[mask] k = ra.results0.model.exog.shape[1] pm = params[0] p0 = params[1:k+1] p1 = params[-k:] mod0 = ra.results0.model mod1 = ra.results1.model # use reordered exog, endog so it matches sub models by group exog = ra.exog_grouped fitted0 = mod0.predict(p0, exog) mom0 = _mom_olsex(p0, model=mod0) fitted1 = mod1.predict(p1, exog) mom1 = _mom_olsex(p1, model=mod1) momout = block_diag(mom0, mom1) mm = fitted1 - fitted0 - pm mpom = fitted0 - ppom mm = np.column_stack((mm, mpom)) if self.probt is not None: mm *= (self.probt / self.probt.mean())[:, None] moms = np.column_stack((mm, momout)) return moms class _IPWRAGMM(_TEGMMGeneric1): """ GMM for ipwra treatment effect and potential outcome """ def momcond(self, params): ra = self.teff treat_mask = ra.treat_mask res_select = ra.results_select ppom = params[1] mask = np.arange(len(params)) != 1 params = params[mask] k = ra.results0.model.exog.shape[1] pm = params[0] # ATE parameter p0 = params[1:k+1] p1 = params[k+1:2*k+1] ps = params[-6:] mod0 = ra.results0.model mod1 = ra.results1.model # use reordered exog so it matches sub models by group exog = ra.exog_grouped tind = np.zeros(len(treat_mask)) tind[-treat_mask.sum():] = 1 # selection probability by group, propensity score prob_sel = np.asarray(res_select.model.predict(ps)) prob_sel = np.clip(prob_sel, 0.001, 0.999) prob0 = prob_sel[~treat_mask] prob1 = prob_sel[treat_mask] effect_group = self.effect_group if effect_group == "all": w0 = 1 / (1 - prob0) w1 = 1 / prob1 sind = 1 elif effect_group in [1, "treated"]: w0 = prob0 / (1 - prob0) w1 = prob1 / prob1 # for averaging effect on treated sind = tind / tind.mean() elif effect_group in [0, "untreated", "control"]: w0 = (1 - prob0) / (1 - prob0) w1 = (1 - prob1) / prob1 sind = (1 - tind) sind /= sind.mean() else: raise ValueError("incorrect option for effect_group") # outcome models by treatment using IPW weights fitted0 = mod0.predict(p0, exog) mom0 = _mom_olsex(p0, model=mod0) * w0[:, None] fitted1 = mod1.predict(p1, exog) mom1 = _mom_olsex(p1, model=mod1) * w1[:, None] mom_outcome = block_diag(mom0, mom1) # moments for target statistics, ATE and POM mm = (fitted1 - fitted0 - pm) * sind mpom = (fitted0 - ppom) * sind mm = np.column_stack((mm, mpom)) # Note: res_select has original data order, # mom_outcome and mm use grouped observations mom_select = res_select.model.score_obs(ps) mom_select = np.concatenate((mom_select[~treat_mask], mom_select[treat_mask]), axis=0) moms = np.column_stack((mm, mom_outcome, mom_select)) return moms class TreatmentEffectResults(ContrastResults): """ Results class for treatment effect estimation Parameters ---------- teff : instance of TreatmentEffect class results_gmm : instance of GMMResults class method : string Method and estimator of treatment effect. kwds: dict Other keywords with additional information. Notes ----- This class is a subclass of ContrastResults and inherits methods like summary, summary_frame and conf_int. Attributes correspond to a z-test given by ``GMMResults.t_test``. """ def __init__(self, teff, results_gmm, method, **kwds): super().__init__() k_params = len(results_gmm.params) constraints = np.zeros((3, k_params)) constraints[0, 0] = 1 constraints[1, 1] = 1 constraints[2, :2] = [1, 1] tt = results_gmm.t_test(constraints) self.__dict__.update(tt.__dict__) self.teff = teff self.results_gmm = results_gmm self.method = method # TODO: make those explicit? self.__dict__.update(kwds) self.c_names = ["ATE", "POM0", "POM1"] doc_params_returns = """\ Parameters ---------- return_results : bool If True, then a results instance is returned. If False, just ATE, POM0 and POM1 are returned. effect_group : {"all", 0, 1} ``effectgroup`` determines for which population the effects are estimated. If effect_group is "all", then sample average treatment effect and potential outcomes are returned If effect_group is 1 or "treated", then effects on treated are returned. If effect_group is 0, "treated" or "control", then effects on untreated, i.e. control group, are returned. disp : bool Indicates whether the scipy optimizer should display the optimization results Returns ------- TreatmentEffectsResults instance or tuple (ATE, POM0, POM1) """ doc_params_returns2 = """\ Parameters ---------- return_results : bool If True, then a results instance is returned. If False, just ATE, POM0 and POM1 are returned. disp : bool Indicates whether the scipy optimizer should display the optimization results Returns ------- TreatmentEffectsResults instance or tuple (ATE, POM0, POM1) """ class TreatmentEffect: """ Estimate average treatment effect under conditional independence .. versionadded:: 0.14.0 This class estimates treatment effect and potential outcome using 5 different methods, ipw, ra, aipw, aipw-wls, ipw-ra. Standard errors and inference are based on the joint GMM representation of selection or treatment model, outcome model and effect functions. Parameters ---------- model : instance of a model class The model class should contain endog and exog for the outcome model. treatment : ndarray indicator array for observations with treatment (1) or without (0) results_select : results instance The results instance for the treatment or selection model. _cov_type : "HC0" Internal keyword. The keyword oes not affect GMMResults which always corresponds to HC0 standard errors. kwds : keyword arguments currently not used Notes ----- The outcome model is currently limited to a linear model based on OLS. Other outcome models, like Logit and Poisson, will become available in future. See `Treatment Effect notebook <../examples/notebooks/generated/treatment_effect.html>`__ for an overview. """ def __init__(self, model, treatment, results_select=None, _cov_type="HC0", **kwds): # Note _cov_type is only for preliminary estimators, # cov in GMM alwasy corresponds to HC0 self.__dict__.update(kwds) # currently not used self.treatment = np.asarray(treatment) self.treat_mask = treat_mask = (treatment == 1) if results_select is not None: self.results_select = results_select self.prob_select = results_select.predict() self.model_pool = model endog = model.endog exog = model.exog self.nobs = endog.shape[0] self._cov_type = _cov_type # no init keys are supported mod0 = model.__class__(endog[~treat_mask], exog[~treat_mask]) self.results0 = mod0.fit(cov_type=_cov_type) mod1 = model.__class__(endog[treat_mask], exog[treat_mask]) self.results1 = mod1.fit(cov_type=_cov_type) # self.predict_mean0 = self.model_pool.predict(self.results0.params # ).mean() # self.predict_mean1 = self.model_pool.predict(self.results1.params # ).mean() self.exog_grouped = np.concatenate((mod0.exog, mod1.exog), axis=0) self.endog_grouped = np.concatenate((mod0.endog, mod1.endog), axis=0) @classmethod def from_data(cls, endog, exog, treatment, model='ols', **kwds): """create models from data not yet implemented """ raise NotImplementedError def ipw(self, return_results=True, effect_group="all", disp=False): """Inverse Probability Weighted treatment effect estimation. Parameters ---------- return_results : bool If True, then a results instance is returned. If False, just ATE, POM0 and POM1 are returned. effect_group : {"all", 0, 1} ``effectgroup`` determines for which population the effects are estimated. If effect_group is "all", then sample average treatment effect and potential outcomes are returned. If effect_group is 1 or "treated", then effects on treated are returned. If effect_group is 0, "treated" or "control", then effects on untreated, i.e. control group, are returned. disp : bool Indicates whether the scipy optimizer should display the optimization results Returns ------- TreatmentEffectsResults instance or tuple (ATE, POM0, POM1) See Also -------- TreatmentEffectsResults """ endog = self.model_pool.endog tind = self.treatment prob = self.prob_select if effect_group == "all": probt = None elif effect_group in [1, "treated"]: probt = prob effect_group = 1 # standardize effect_group name elif effect_group in [0, "untreated", "control"]: probt = 1 - prob effect_group = 0 # standardize effect_group name elif isinstance(effect_group, np.ndarray): probt = effect_group effect_group = "user" # standardize effect_group name else: raise ValueError("incorrect option for effect_group") res_ipw = ate_ipw(endog, tind, prob, weighted=True, probt=probt) if not return_results: return res_ipw # gmm = _TEGMMGeneric1(endog, self.results_select, _mom_ols_te, # probt=probt) gmm = _IPWGMM(endog, self.results_select, None, teff=self, effect_group=effect_group) start_params = np.concatenate((res_ipw[:2], self.results_select.params)) res_gmm = gmm.fit(start_params=start_params, inv_weights=np.eye(len(start_params)), optim_method='nm', optim_args={"maxiter": 5000, "disp": disp}, maxiter=1, ) res = TreatmentEffectResults(self, res_gmm, "IPW", start_params=start_params, effect_group=effect_group, ) return res @Substitution(params_returns=indent(doc_params_returns, " " * 8)) def ra(self, return_results=True, effect_group="all", disp=False): """ Regression Adjustment treatment effect estimation. \n%(params_returns)s See Also -------- TreatmentEffectsResults """ # need indicator for reordered observations tind = np.zeros(len(self.treatment)) tind[-self.treatment.sum():] = 1 if effect_group == "all": probt = None elif effect_group in [1, "treated"]: probt = tind effect_group = 1 # standardize effect_group name elif effect_group in [0, "untreated", "control"]: probt = 1 - tind effect_group = 0 # standardize effect_group name elif isinstance(effect_group, np.ndarray): # TODO: do we keep this? probt = effect_group effect_group = "user" # standardize effect_group name else: raise ValueError("incorrect option for effect_group") exog = self.exog_grouped # weight or indicator for effect_group if probt is not None: cw = (probt / probt.mean()) else: cw = 1 pom0 = (self.results0.predict(exog) * cw).mean() pom1 = (self.results1.predict(exog) * cw).mean() if not return_results: return pom1 - pom0, pom0, pom1 endog = self.model_pool.endog mod_gmm = _RAGMM(endog, self.results_select, None, teff=self, probt=probt) start_params = np.concatenate(( # ate, tt0.effect, [pom1 - pom0, pom0], self.results0.params, self.results1.params)) res_gmm = mod_gmm.fit(start_params=start_params, inv_weights=np.eye(len(start_params)), optim_method='nm', optim_args={"maxiter": 5000, "disp": disp}, maxiter=1, ) res = TreatmentEffectResults(self, res_gmm, "IPW", start_params=start_params, effect_group=effect_group, ) return res @Substitution(params_returns=indent(doc_params_returns2, " " * 8)) def aipw(self, return_results=True, disp=False): """ ATE and POM from double robust augmented inverse probability weighting \n%(params_returns)s See Also -------- TreatmentEffectsResults """ nobs = self.nobs prob = self.prob_select tind = self.treatment exog = self.model_pool.exog # in original order correct0 = (self.results0.resid / (1 - prob[tind == 0])).sum() / nobs correct1 = (self.results1.resid / (prob[tind == 1])).sum() / nobs tmean0 = self.results0.predict(exog).mean() + correct0 tmean1 = self.results1.predict(exog).mean() + correct1 ate = tmean1 - tmean0 if not return_results: return ate, tmean0, tmean1 endog = self.model_pool.endog p2_aipw = np.asarray([ate, tmean0]) mag_aipw1 = _AIPWGMM(endog, self.results_select, None, teff=self) start_params = np.concatenate(( p2_aipw, self.results0.params, self.results1.params, self.results_select.params)) res_gmm = mag_aipw1.fit( start_params=start_params, inv_weights=np.eye(len(start_params)), optim_method='nm', optim_args={"maxiter": 5000, "disp": disp}, maxiter=1) res = TreatmentEffectResults(self, res_gmm, "IPW", start_params=start_params, effect_group="all", ) return res @Substitution(params_returns=indent(doc_params_returns2, " " * 8)) def aipw_wls(self, return_results=True, disp=False): """ ATE and POM from double robust augmented inverse probability weighting. This uses weighted outcome regression, while `aipw` uses unweighted outcome regression. Option for effect on treated or on untreated is not available. \n%(params_returns)s See Also -------- TreatmentEffectsResults """ nobs = self.nobs prob = self.prob_select endog = self.model_pool.endog exog = self.model_pool.exog tind = self.treatment treat_mask = self.treat_mask ww1 = tind / prob * (tind / prob - 1) mod1 = WLS(endog[treat_mask], exog[treat_mask], weights=ww1[treat_mask]) result1 = mod1.fit(cov_type='HC1') mean1_ipw2 = result1.predict(exog).mean() ww0 = (1 - tind) / (1 - prob) * ((1 - tind) / (1 - prob) - 1) mod0 = WLS(endog[~treat_mask], exog[~treat_mask], weights=ww0[~treat_mask]) result0 = mod0.fit(cov_type='HC1') mean0_ipw2 = result0.predict(exog).mean() self.results_ipwwls0 = result0 self.results_ipwwls1 = result1 correct0 = (result0.resid / (1 - prob[tind == 0])).sum() / nobs correct1 = (result1.resid / (prob[tind == 1])).sum() / nobs tmean0 = mean0_ipw2 + correct0 tmean1 = mean1_ipw2 + correct1 ate = tmean1 - tmean0 if not return_results: return ate, tmean0, tmean1 p2_aipw_wls = np.asarray([ate, tmean0]).squeeze() # GMM mod_gmm = _AIPWWLSGMM(endog, self.results_select, None, teff=self) start_params = np.concatenate(( p2_aipw_wls, result0.params, result1.params, self.results_select.params)) res_gmm = mod_gmm.fit( start_params=start_params, inv_weights=np.eye(len(start_params)), optim_method='nm', optim_args={"maxiter": 5000, "disp": disp}, maxiter=1) res = TreatmentEffectResults(self, res_gmm, "IPW", start_params=start_params, effect_group="all", ) return res @Substitution(params_returns=indent(doc_params_returns, " " * 8)) def ipw_ra(self, return_results=True, effect_group="all", disp=False): """ ATE and POM from inverse probability weighted regression adjustment. \n%(params_returns)s See Also -------- TreatmentEffectsResults """ treat_mask = self.treat_mask endog = self.model_pool.endog exog = self.model_pool.exog prob = self.prob_select prob0 = prob[~treat_mask] prob1 = prob[treat_mask] if effect_group == "all": w0 = 1 / (1 - prob0) w1 = 1 / prob1 exogt = exog elif effect_group in [1, "treated"]: w0 = prob0 / (1 - prob0) w1 = prob1 / prob1 exogt = exog[treat_mask] effect_group = 1 # standardize effect_group name elif effect_group in [0, "untreated", "control"]: w0 = (1 - prob0) / (1 - prob0) w1 = (1 - prob1) / prob1 exogt = exog[~treat_mask] effect_group = 0 # standardize effect_group name else: raise ValueError("incorrect option for effect_group") mod0 = WLS(endog[~treat_mask], exog[~treat_mask], weights=w0) result0 = mod0.fit(cov_type='HC1') # mean0_ipwra = (result0.predict(exog) * (prob / prob.mean())).mean() mean0_ipwra = result0.predict(exogt).mean() mod1 = WLS(endog[treat_mask], exog[treat_mask], weights=w1) result1 = mod1.fit(cov_type='HC1') # mean1_ipwra = (result1.predict(exog) * (prob / prob.mean())).mean() mean1_ipwra = result1.predict(exogt).mean() if not return_results: return mean1_ipwra - mean0_ipwra, mean0_ipwra, mean1_ipwra # GMM mod_gmm = _IPWRAGMM(endog, self.results_select, None, teff=self, effect_group=effect_group) start_params = np.concatenate(( [mean1_ipwra - mean0_ipwra, mean0_ipwra], result0.params, result1.params, np.asarray(self.results_select.params) )) res_gmm = mod_gmm.fit( start_params=start_params, inv_weights=np.eye(len(start_params)), optim_method='nm', optim_args={"maxiter": 2000, "disp": disp}, maxiter=1 ) res = TreatmentEffectResults(self, res_gmm, "IPW", start_params=start_params, effect_group=effect_group, ) return res