""" Mediation analysis Implements algorithm 1 ('parametric inference') and algorithm 2 ('nonparametric inference') from: Imai, Keele, Tingley (2010). A general approach to causal mediation analysis. Psychological Methods 15:4, 309-334. http://imai.princeton.edu/research/files/BaronKenny.pdf The algorithms are described on page 317 of the paper. In the case of linear models with no interactions involving the mediator, the results should be similar or identical to the earlier Barron-Kenny approach. """ import numpy as np import pandas as pd from statsmodels.graphics.utils import maybe_name_or_idx class Mediation: """ Conduct a mediation analysis. Parameters ---------- outcome_model : statsmodels model Regression model for the outcome. Predictor variables include the treatment/exposure, the mediator, and any other variables of interest. mediator_model : statsmodels model Regression model for the mediator variable. Predictor variables include the treatment/exposure and any other variables of interest. exposure : str or (int, int) tuple The name or column position of the treatment/exposure variable. If positions are given, the first integer is the column position of the exposure variable in the outcome model and the second integer is the position of the exposure variable in the mediator model. If a string is given, it must be the name of the exposure variable in both regression models. mediator : {str, int} The name or column position of the mediator variable in the outcome regression model. If None, infer the name from the mediator model formula (if present). moderators : dict Map from variable names or index positions to values of moderator variables that are held fixed when calculating mediation effects. If the keys are index position they must be tuples `(i, j)` where `i` is the index in the outcome model and `j` is the index in the mediator model. Otherwise the keys must be variable names. outcome_fit_kwargs : dict-like Keyword arguments to use when fitting the outcome model. mediator_fit_kwargs : dict-like Keyword arguments to use when fitting the mediator model. outcome_predict_kwargs : dict-like Keyword arguments to use when calling predict on the outcome model. Returns a ``MediationResults`` object. Notes ----- The mediator model class must implement ``get_distribution``. Examples -------- A basic mediation analysis using formulas: >>> import statsmodels.api as sm >>> import statsmodels.genmod.families.links as links >>> probit = links.probit >>> outcome_model = sm.GLM.from_formula("cong_mesg ~ emo + treat + age + educ + gender + income", ... data, family=sm.families.Binomial(link=Probit())) >>> mediator_model = sm.OLS.from_formula("emo ~ treat + age + educ + gender + income", data) >>> med = Mediation(outcome_model, mediator_model, "treat", "emo").fit() >>> med.summary() A basic mediation analysis without formulas. This may be slightly faster than the approach using formulas. If there are any interactions involving the treatment or mediator variables this approach will not work, you must use formulas. >>> import patsy >>> outcome = np.asarray(data["cong_mesg"]) >>> outcome_exog = patsy.dmatrix("emo + treat + age + educ + gender + income", data, ... return_type='dataframe') >>> probit = sm.families.links.probit >>> outcome_model = sm.GLM(outcome, outcome_exog, family=sm.families.Binomial(link=Probit())) >>> mediator = np.asarray(data["emo"]) >>> mediator_exog = patsy.dmatrix("treat + age + educ + gender + income", data, ... return_type='dataframe') >>> mediator_model = sm.OLS(mediator, mediator_exog) >>> tx_pos = [outcome_exog.columns.tolist().index("treat"), ... mediator_exog.columns.tolist().index("treat")] >>> med_pos = outcome_exog.columns.tolist().index("emo") >>> med = Mediation(outcome_model, mediator_model, tx_pos, med_pos).fit() >>> med.summary() A moderated mediation analysis. The mediation effect is computed for people of age 20. >>> fml = "cong_mesg ~ emo + treat*age + emo*age + educ + gender + income", >>> outcome_model = sm.GLM.from_formula(fml, data, ... family=sm.families.Binomial()) >>> mediator_model = sm.OLS.from_formula("emo ~ treat*age + educ + gender + income", data) >>> moderators = {"age" : 20} >>> med = Mediation(outcome_model, mediator_model, "treat", "emo", ... moderators=moderators).fit() References ---------- Imai, Keele, Tingley (2010). A general approach to causal mediation analysis. Psychological Methods 15:4, 309-334. http://imai.princeton.edu/research/files/BaronKenny.pdf Tingley, Yamamoto, Hirose, Keele, Imai (2014). mediation : R package for causal mediation analysis. Journal of Statistical Software 59:5. http://www.jstatsoft.org/v59/i05/paper """ def __init__(self, outcome_model, mediator_model, exposure, mediator=None, moderators=None, outcome_fit_kwargs=None, mediator_fit_kwargs=None, outcome_predict_kwargs=None): self.outcome_model = outcome_model self.mediator_model = mediator_model self.exposure = exposure self.moderators = moderators if moderators is not None else {} if mediator is None: self.mediator = self._guess_endog_name(mediator_model, 'mediator') else: self.mediator = mediator self._outcome_fit_kwargs = (outcome_fit_kwargs if outcome_fit_kwargs is not None else {}) self._mediator_fit_kwargs = (mediator_fit_kwargs if mediator_fit_kwargs is not None else {}) self._outcome_predict_kwargs = (outcome_predict_kwargs if outcome_predict_kwargs is not None else {}) # We will be changing these so need to copy. self._outcome_exog = outcome_model.exog.copy() self._mediator_exog = mediator_model.exog.copy() # Position of the exposure variable in the mediator model. self._exp_pos_mediator = self._variable_pos('exposure', 'mediator') # Position of the exposure variable in the outcome model. self._exp_pos_outcome = self._variable_pos('exposure', 'outcome') # Position of the mediator variable in the outcome model. self._med_pos_outcome = self._variable_pos('mediator', 'outcome') def _variable_pos(self, var, model): if model == 'mediator': mod = self.mediator_model else: mod = self.outcome_model if var == 'mediator': return maybe_name_or_idx(self.mediator, mod)[1] exp = self.exposure exp_is_2 = ((len(exp) == 2) and not isinstance(exp, str)) if exp_is_2: if model == 'outcome': return exp[0] elif model == 'mediator': return exp[1] else: return maybe_name_or_idx(exp, mod)[1] def _guess_endog_name(self, model, typ): if hasattr(model, 'formula'): return model.formula.split("~")[0].strip() else: raise ValueError('cannot infer %s name without formula' % typ) def _simulate_params(self, result): """ Simulate model parameters from fitted sampling distribution. """ mn = result.params cov = result.cov_params() return np.random.multivariate_normal(mn, cov) def _get_mediator_exog(self, exposure): """ Return the mediator exog matrix with exposure set to the given value. Set values of moderated variables as needed. """ mediator_exog = self._mediator_exog if not hasattr(self.mediator_model, 'formula'): mediator_exog[:, self._exp_pos_mediator] = exposure for ix in self.moderators: v = self.moderators[ix] mediator_exog[:, ix[1]] = v else: # Need to regenerate the model exog df = self.mediator_model.data.frame.copy() df[self.exposure] = exposure for vname in self.moderators: v = self.moderators[vname] df.loc[:, vname] = v klass = self.mediator_model.__class__ init_kwargs = self.mediator_model._get_init_kwds() model = klass.from_formula(data=df, **init_kwargs) mediator_exog = model.exog return mediator_exog def _get_outcome_exog(self, exposure, mediator): """ Retun the exog design matrix with mediator and exposure set to the given values. Set values of moderated variables as needed. """ outcome_exog = self._outcome_exog if not hasattr(self.outcome_model, 'formula'): outcome_exog[:, self._med_pos_outcome] = mediator outcome_exog[:, self._exp_pos_outcome] = exposure for ix in self.moderators: v = self.moderators[ix] outcome_exog[:, ix[0]] = v else: # Need to regenerate the model exog df = self.outcome_model.data.frame.copy() df[self.exposure] = exposure df[self.mediator] = mediator for vname in self.moderators: v = self.moderators[vname] df[vname] = v klass = self.outcome_model.__class__ init_kwargs = self.outcome_model._get_init_kwds() model = klass.from_formula(data=df, **init_kwargs) outcome_exog = model.exog return outcome_exog def _fit_model(self, model, fit_kwargs, boot=False): klass = model.__class__ init_kwargs = model._get_init_kwds() endog = model.endog exog = model.exog if boot: ii = np.random.randint(0, len(endog), len(endog)) endog = endog[ii] exog = exog[ii, :] outcome_model = klass(endog, exog, **init_kwargs) return outcome_model.fit(**fit_kwargs) def fit(self, method="parametric", n_rep=1000): """ Fit a regression model to assess mediation. Parameters ---------- method : str Either 'parametric' or 'bootstrap'. n_rep : int The number of simulation replications. Returns a MediationResults object. """ if method.startswith("para"): # Initial fit to unperturbed data. outcome_result = self._fit_model(self.outcome_model, self._outcome_fit_kwargs) mediator_result = self._fit_model(self.mediator_model, self._mediator_fit_kwargs) elif not method.startswith("boot"): raise ValueError( "method must be either 'parametric' or 'bootstrap'" ) indirect_effects = [[], []] direct_effects = [[], []] for iter in range(n_rep): if method == "parametric": # Realization of outcome model parameters from sampling distribution outcome_params = self._simulate_params(outcome_result) # Realization of mediation model parameters from sampling distribution mediation_params = self._simulate_params(mediator_result) else: outcome_result = self._fit_model(self.outcome_model, self._outcome_fit_kwargs, boot=True) outcome_params = outcome_result.params mediator_result = self._fit_model(self.mediator_model, self._mediator_fit_kwargs, boot=True) mediation_params = mediator_result.params # predicted outcomes[tm][te] is the outcome when the # mediator is set to tm and the outcome/exposure is set to # te. predicted_outcomes = [[None, None], [None, None]] for tm in 0, 1: mex = self._get_mediator_exog(tm) kwargs = {"exog": mex} if hasattr(mediator_result, "scale"): kwargs["scale"] = mediator_result.scale gen = self.mediator_model.get_distribution(mediation_params, **kwargs) potential_mediator = gen.rvs(mex.shape[0]) for te in 0, 1: oex = self._get_outcome_exog(te, potential_mediator) po = self.outcome_model.predict(outcome_params, oex, **self._outcome_predict_kwargs) predicted_outcomes[tm][te] = po for t in 0, 1: indirect_effects[t].append(predicted_outcomes[1][t] - predicted_outcomes[0][t]) direct_effects[t].append(predicted_outcomes[t][1] - predicted_outcomes[t][0]) for t in 0, 1: indirect_effects[t] = np.asarray(indirect_effects[t]).T direct_effects[t] = np.asarray(direct_effects[t]).T self.indirect_effects = indirect_effects self.direct_effects = direct_effects rslt = MediationResults(self.indirect_effects, self.direct_effects) rslt.method = method return rslt def _pvalue(vec): return 2 * min(sum(vec > 0), sum(vec < 0)) / float(len(vec)) class MediationResults: """ A class for holding the results of a mediation analysis. The following terms are used in the summary output: ACME : average causal mediated effect ADE : average direct effect """ def __init__(self, indirect_effects, direct_effects): self.indirect_effects = indirect_effects self.direct_effects = direct_effects indirect_effects_avg = [None, None] direct_effects_avg = [None, None] for t in 0, 1: indirect_effects_avg[t] = indirect_effects[t].mean(0) direct_effects_avg[t] = direct_effects[t].mean(0) self.ACME_ctrl = indirect_effects_avg[0] self.ACME_tx = indirect_effects_avg[1] self.ADE_ctrl = direct_effects_avg[0] self.ADE_tx = direct_effects_avg[1] self.total_effect = (self.ACME_ctrl + self.ACME_tx + self.ADE_ctrl + self.ADE_tx) / 2 self.prop_med_ctrl = self.ACME_ctrl / self.total_effect self.prop_med_tx = self.ACME_tx / self.total_effect self.prop_med_avg = (self.prop_med_ctrl + self.prop_med_tx) / 2 self.ACME_avg = (self.ACME_ctrl + self.ACME_tx) / 2 self.ADE_avg = (self.ADE_ctrl + self.ADE_tx) / 2 def summary(self, alpha=0.05): """ Provide a summary of a mediation analysis. """ columns = ["Estimate", "Lower CI bound", "Upper CI bound", "P-value"] index = ["ACME (control)", "ACME (treated)", "ADE (control)", "ADE (treated)", "Total effect", "Prop. mediated (control)", "Prop. mediated (treated)", "ACME (average)", "ADE (average)", "Prop. mediated (average)"] smry = pd.DataFrame(columns=columns, index=index) for i, vec in enumerate([self.ACME_ctrl, self.ACME_tx, self.ADE_ctrl, self.ADE_tx, self.total_effect, self.prop_med_ctrl, self.prop_med_tx, self.ACME_avg, self.ADE_avg, self.prop_med_avg]): if ((vec is self.prop_med_ctrl) or (vec is self.prop_med_tx) or (vec is self.prop_med_avg)): smry.iloc[i, 0] = np.median(vec) else: smry.iloc[i, 0] = vec.mean() smry.iloc[i, 1] = np.percentile(vec, 100 * alpha / 2) smry.iloc[i, 2] = np.percentile(vec, 100 * (1 - alpha / 2)) smry.iloc[i, 3] = _pvalue(vec) smry = smry.apply(pd.to_numeric, errors='coerce') return smry