Source code for statsmodels.treatment.treatment_effects
"""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
[docs]
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)
"""
[docs]
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)
[docs]
@classmethod
def from_data(cls, endog, exog, treatment, model='ols', **kwds):
"""create models from data
not yet implemented
"""
raise NotImplementedError
[docs]
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
[docs]
@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
[docs]
@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
[docs]
@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
[docs]
@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
Last update:
Apr 18, 2024