-
Notifications
You must be signed in to change notification settings - Fork 2.9k
/
mediation.py
412 lines (336 loc) · 16 KB
/
mediation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
"""
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