"""
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
"""
from statsmodels.compat.pandas import Substitution
import numpy as np
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).
Parameters
----------
params : ndarray
Parameter at which the moment condition is evaluated, the average
treatment effect.
endog : ndarray
Outcome variable.
tind : ndarray
Treatment indicator, 1 for treated and 0 for untreated observations.
prob : ndarray
Estimated propensity score, probability of treatment.
weighted : bool
If True, weights are normalized so that each of the treated and
untreated weights average to one.
Returns
-------
ndarray
Moment condition, evaluated at each observation.
"""
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.
Parameters
----------
params : ndarray
Parameters, POM0 and POM1, at which the moment conditions are
evaluated.
endog : ndarray
Outcome variable.
tind : ndarray
Treatment indicator, 1 for treated and 0 for untreated observations.
prob : ndarray
Estimated propensity score, probability of treatment.
weighted : bool
If True, weights are normalized so that each of the treated and
untreated weights average to one.
Returns
-------
ndarray
Moment conditions, evaluated at each observation, 2 columns for
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.
Parameters
----------
params : ndarray
Parameters, POM0 and POM1, at which the moment conditions are
evaluated.
endog : ndarray
Outcome variable.
tind : ndarray
Treatment indicator, 1 for treated and 0 for untreated observations.
prob : ndarray
Estimated propensity score, probability of treatment.
weighted : bool
Not used. Kept for signature compatibility with other moment
condition functions.
Returns
-------
ndarray
Moment conditions, evaluated at each observation, 2 columns for
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).
Parameters
----------
tm : ndarray
Parameters, ATE and POM0, at which the moment conditions are
evaluated.
endog : ndarray
Outcome variable.
tind : ndarray
Treatment indicator, 1 for treated and 0 for untreated observations.
prob : ndarray
Estimated propensity score, probability of treatment.
weighted : bool
Not used. Kept for signature compatibility with other moment
condition functions.
Returns
-------
ndarray
Moment conditions, evaluated at each observation, 2 columns for
ATE and POM0.
"""
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):
"""
Moment condition for OLS with an optionally external exog and scale
Parameters
----------
params : ndarray
Parameters of the outcome model.
model : instance of a model class
Model with ``endog`` and, if `exog` is None, ``exog`` attributes,
and a ``predict`` method.
exog : ndarray, optional
Explanatory variables used to compute fitted values. If None, then
``model.exog`` is used instead.
scale : float, optional
If provided, the residual based moment condition is divided by
this scale.
Returns
-------
ndarray
Moment condition, evaluated at each observation, weighted by exog.
"""
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
Parameters
----------
endog : ndarray
Outcome variable.
tind : ndarray
Treatment indicator, 1 for treated and 0 for untreated observations.
prob : ndarray
Estimated propensity score, probability of treatment.
weighted : bool
If True, weights are normalized so that each of the treated and
untreated weights average to one.
probt : ndarray, optional
Additional weight, e.g. propensity score, used to target the effect
on a subgroup such as the treated or untreated population. If None,
the sample average treatment effect is computed.
Returns
-------
ate : float
Average treatment effect.
pom0 : float
Potential outcome mean for the untreated (control) group.
pom1 : float
Potential outcome mean for the treated group.
"""
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.
Parameters
----------
endog : ndarray
Outcome variable, endog of the outcome model.
res_select : results instance
Results instance of the treatment or selection model.
mom_outcome : callable
Function that computes the moment conditions of the outcome model.
exclude_tmoms : bool
If True, then the moment conditions of the treatment or selection
model are not included and `params` does not contain the
parameters of the selection model.
**kwargs
Additional keyword arguments that are attached to the instance,
e.g. `teff` and `effect_group` used by the subclasses.
"""
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.
Parameters
----------
endog : ndarray
Outcome variable, endog of the outcome model.
res_select : results instance
Results instance of the treatment or selection model.
mom_outcome : callable
Function that computes the moment conditions of the outcome model.
"""
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 ipw treatment effect and potential outcome
Notes
-----
Uses a single, pooled outcome moment condition weighted by the inverse
probability weights.
"""
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 not isinstance(effect_group, np.ndarray):
raise ValueError("incorrect option for effect_group")
# default else probt = probt
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
Notes
-----
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
Notes
-----
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
Notes
-----
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
The treatment effect instance that produced the results.
results_gmm : instance of GMMResults class
The GMM results instance used to compute the treatment effect
parameters and their covariance.
method : string
Method and estimator of treatment effect.
**kwds
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, "untreated" 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
-------
TreatmentEffectResults or tuple
Results instance if `return_results` is True, otherwise the 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
-------
TreatmentEffectResults or tuple
Results instance if `return_results` is True, otherwise the 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 does not affect GMMResults which always
corresponds to HC0 standard errors.
**kwds
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 always 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.
Parameters
----------
endog : ndarray
Outcome variable for the outcome model.
exog : ndarray
Explanatory variables for the outcome model.
treatment : ndarray
Indicator array for observations with treatment (1) or
without (0).
model : str
Name of the model class to use for the outcome model.
**kwds
Additional keyword arguments passed to model classes.
Returns
-------
TreatmentEffect
Instance of the treatment effect class created from the data.
"""
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, "untreated" 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
-------
TreatmentEffectResults or tuple
Results instance if `return_results` is True, otherwise the
tuple (ATE, POM0, POM1).
See Also
--------
TreatmentEffectResults
"""
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
--------
TreatmentEffectResults
"""
# 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
--------
TreatmentEffectResults
"""
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
--------
TreatmentEffectResults
"""
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
--------
TreatmentEffectResults
"""
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
