import numpy as np import statsmodels.api as sm import os from statsmodels.stats.mediation import Mediation import pandas as pd from numpy.testing import assert_allclose import patsy import pytest # Compare to mediation R package vignette df = [['index', 'Estimate', 'Lower CI bound', 'Upper CI bound', 'P-value'], ['ACME (control)', 0.085106, 0.029938, 0.141525, 0.00], ['ACME (treated)', 0.085674, 0.031089, 0.147762, 0.00], ['ADE (control)', 0.016938, -0.129157, 0.121945, 0.66], ['ADE (treated)', 0.017506, -0.139649, 0.130030, 0.66], ['Total effect', 0.102612, -0.036749, 0.227213, 0.20], ['Prop. mediated (control)', 0.698070, -6.901715, 2.725978, 0.20], ['Prop. mediated (treated)', 0.718648, -6.145419, 2.510750, 0.20], ['ACME (average)', 0.085390, 0.030272, 0.144768, 0.00], ['ADE (average)', 0.017222, -0.134465, 0.125987, 0.66], ['Prop. mediated (average)', 0.710900, -6.523567, 2.618364, 0.20]] framing_boot_4231 = pd.DataFrame(df[1:], columns=df[0]).set_index('index') # Compare to mediation R package vignette df = [['index', 'Estimate', 'Lower CI bound', 'Upper CI bound', 'P-value'], ['ACME (control)', 0.075529, 0.024995, 0.132408, 0.00], ['ACME (treated)', 0.076348, 0.027475, 0.130138, 0.00], ['ADE (control)', 0.021389, -0.094323, 0.139148, 0.68], ['ADE (treated)', 0.022207, -0.101239, 0.145740, 0.68], ['Total effect', 0.097736, -0.025384, 0.225386, 0.16], ['Prop. mediated (control)', 0.656820, -3.664956, 4.845269, 0.16], ['Prop. mediated (treated)', 0.687690, -3.449415, 4.469289, 0.16], ['ACME (average)', 0.075938, 0.026109, 0.129450, 0.00], ['ADE (average)', 0.021798, -0.097781, 0.142444, 0.68], ['Prop. mediated (average)', 0.669659, -3.557185, 4.657279, 0.16]] framing_para_4231 = pd.DataFrame(df[1:], columns=df[0]).set_index('index') df = [['index', 'Estimate', 'Lower CI bound', 'Upper CI bound', 'P-value'], ['ACME (control)', 0.065989, 0.003366, 0.152261, 0.04], ['ACME (treated)', 0.081424, 0.008888, 0.199853, 0.04], ['ADE (control)', 0.240392, -0.026286, 0.470918, 0.08], ['ADE (treated)', 0.255827, -0.030681, 0.491535, 0.08], ['Total effect', 0.321816, 0.037238, 0.549530, 0.00], ['Prop. mediated (control)', 0.196935, 0.015232, 1.864804, 0.04], ['Prop. mediated (treated)', 0.248896, 0.032229, 1.738846, 0.04], ['ACME (average)', 0.073707, 0.006883, 0.169923, 0.04], ['ADE (average)', 0.248109, -0.028483, 0.478978, 0.08], ['Prop. mediated (average)', 0.226799, 0.028865, 1.801825, 0.04]] framing_moderated_4231 = pd.DataFrame(df[1:], columns=df[0]).set_index('index') @pytest.mark.slow def test_framing_example(): cur_dir = os.path.dirname(os.path.abspath(__file__)) data = pd.read_csv(os.path.join(cur_dir, 'results', "framing.csv")) outcome = np.asarray(data["cong_mesg"]) outcome_exog = patsy.dmatrix("emo + treat + age + educ + gender + income", data, return_type='dataframe') outcome_model = sm.GLM( outcome, outcome_exog, family=sm.families.Binomial(link=sm.families.links.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, outcome_fit_kwargs={'atol':1e-11}) np.random.seed(4231) para_rslt = med.fit(method='parametric', n_rep=100) diff = np.asarray(para_rslt.summary() - framing_para_4231) assert_allclose(diff, 0, atol=1e-6) np.random.seed(4231) boot_rslt = med.fit(method='boot', n_rep=100) diff = np.asarray(boot_rslt.summary() - framing_boot_4231) assert_allclose(diff, 0, atol=1e-6) def test_framing_example_moderator(): # moderation without formulas, generally not useful but test anyway cur_dir = os.path.dirname(os.path.abspath(__file__)) data = pd.read_csv(os.path.join(cur_dir, 'results', "framing.csv")) outcome = np.asarray(data["cong_mesg"]) outcome_exog = patsy.dmatrix("emo + treat + age + educ + gender + income", data, return_type='dataframe') outcome_model = sm.GLM( outcome, outcome_exog, family=sm.families.Binomial(link=sm.families.links.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") ix = (outcome_exog.columns.tolist().index("age"), mediator_exog.columns.tolist().index("age")) moderators = {ix : 20} med = Mediation(outcome_model, mediator_model, tx_pos, med_pos, moderators=moderators) # Just a smoke test np.random.seed(4231) med_rslt = med.fit(method='parametric', n_rep=100) @pytest.mark.slow def test_framing_example_formula(): cur_dir = os.path.dirname(os.path.abspath(__file__)) data = pd.read_csv(os.path.join(cur_dir, 'results', "framing.csv")) outcome_model = sm.GLM.from_formula( "cong_mesg ~ emo + treat + age + educ + gender + income", data, family=sm.families.Binomial(link=sm.families.links.Probit()) ) mediator_model = sm.OLS.from_formula("emo ~ treat + age + educ + gender + income", data) med = Mediation(outcome_model, mediator_model, "treat", "emo", outcome_fit_kwargs={'atol': 1e-11}) np.random.seed(4231) med_rslt = med.fit(method='boot', n_rep=100) diff = np.asarray(med_rslt.summary() - framing_boot_4231) assert_allclose(diff, 0, atol=1e-6) np.random.seed(4231) med_rslt = med.fit(method='parametric', n_rep=100) diff = np.asarray(med_rslt.summary() - framing_para_4231) assert_allclose(diff, 0, atol=1e-6) @pytest.mark.slow def test_framing_example_moderator_formula(): cur_dir = os.path.dirname(os.path.abspath(__file__)) data = pd.read_csv(os.path.join(cur_dir, 'results', "framing.csv")) outcome_model = sm.GLM.from_formula( "cong_mesg ~ emo + treat*age + emo*age + educ + gender + income", data, family=sm.families.Binomial(link=sm.families.links.Probit()) ) 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) np.random.seed(4231) med_rslt = med.fit(method='parametric', n_rep=100) diff = np.asarray(med_rslt.summary() - framing_moderated_4231) assert_allclose(diff, 0, atol=1e-6) def t_est_mixedlm(): # check backwards compat of np.random np.random.seed(3424) mn = np.random.randn(5) c = 1e-4 * (np.random.rand(5, 5) - 0.5) cov = np.eye(5) + c + c.T rvs = np.random.multivariate_normal(mn, cov) rvs1 = [0.3357151, 1.26183927, 1.22539916, 0.85838887, -0.0493799] assert_allclose(rvs, rvs1, atol=1e-7) np.random.seed(3424) n = 200 # The exposure (not time varying) x = np.random.normal(size=n) xv = np.outer(x, np.ones(3)) # The mediator (with random intercept) mx = np.asarray([4., 4, 1]) mx /= np.sqrt(np.sum(mx**2)) med = mx[0] * np.outer(x, np.ones(3)) med += mx[1] * np.outer(np.random.normal(size=n), np.ones(3)) med += mx[2] * np.random.normal(size=(n, 3)) # The outcome (exposure and mediator effects) ey = np.outer(x, np.r_[0, 0.5, 1]) + med # Random structure of the outcome (random intercept and slope) ex = np.asarray([5., 2, 2]) ex /= np.sqrt(np.sum(ex**2)) e = ex[0] * np.outer(np.random.normal(size=n), np.ones(3)) e += ex[1] * np.outer(np.random.normal(size=n), np.r_[-1, 0, 1]) e += ex[2] * np.random.normal(size=(n, 3)) y = ey + e # Group membership idx = np.outer(np.arange(n), np.ones(3)) # Time tim = np.outer(np.ones(n), np.r_[-1, 0, 1]) df = pd.DataFrame({"y": y.flatten(), "x": xv.flatten(), "id": idx.flatten(), "time": tim.flatten(), "med": med.flatten()}) # check data is unchanged, regression numbers dmean = [-0.13643661, -0.14266871, 99.5, 0.0, -0.15102166] assert_allclose(np.asarray(df.mean()), dmean, atol=1e-7) mediator_model = sm.MixedLM.from_formula("med ~ x", groups="id", data=df) outcome_model = sm.MixedLM.from_formula("y ~ med + x", groups="id", data=df) me = Mediation(outcome_model, mediator_model, "x", "med") np.random.seed(383628) mr = me.fit(n_rep=100) st = mr.summary() # check model results unchanged, regression numbers params_om = me.outcome_model.fit().params.to_numpy() p_om = [0.08118371, 0.96107436, 0.50801102, 1.22452252] assert_allclose(params_om, p_om, atol=1e-7) params_mm = me.mediator_model.fit().params.to_numpy() p_mm = [-0.0547506, 0.67478745, 17.03184275] assert_allclose(params_mm, p_mm, atol=1e-7) # more regression numbers res_summ = np.array([ [0.64539794, 0.57652012, 0.71427576, 0.0], [0.64539794, 0.57652012, 0.71427576, 0.0], [0.59401941, 0.56963807, 0.61840074, 0.0], [0.59401941, 0.56963807, 0.61840074, 0.0], [1.23941735, 1.14615820, 1.33267651, 0.0], [0.51935169, 0.50285723, 0.53584615, 0.0], [0.51935169, 0.50285723, 0.53584615, 0.0], [0.64539794, 0.57652012, 0.71427576, 0.0], [0.59401941, 0.56963807, 0.61840074, 0.0], [0.51935169, 0.50285723, 0.53584615, 0.0] ]) assert_allclose(st.to_numpy(), res_summ, rtol=0.15) assert_allclose(st.iloc[-1, 0], 0.56, rtol=1e-2, atol=1e-2) pm = st.loc["Prop. mediated (average)", "Estimate"] assert_allclose(pm, 0.56, rtol=1e-2, atol=1e-2) def test_surv(): np.random.seed(2341) n = 1000 # Generate exposures exp = np.random.normal(size=n) # Generate mediators mn = np.exp(exp) mtime0 = -mn * np.log(np.random.uniform(size=n)) ctime = -2 * mn * np.log(np.random.uniform(size=n)) mstatus = (ctime >= mtime0).astype(int) mtime = np.where(mtime0 <= ctime, mtime0, ctime) for mt in "full", "partial", "no": # Outcome if mt == "full": lp = 0.5*mtime0 elif mt == "partial": lp = exp + mtime0 else: lp = exp # Generate outcomes mn = np.exp(-lp) ytime0 = -mn * np.log(np.random.uniform(size=n)) ctime = -2 * mn * np.log(np.random.uniform(size=n)) ystatus = (ctime >= ytime0).astype(int) ytime = np.where(ytime0 <= ctime, ytime0, ctime) df = pd.DataFrame({"ytime": ytime, "ystatus": ystatus, "mtime": mtime, "mstatus": mstatus, "exp": exp}) fml = "ytime ~ exp + mtime" outcome_model = sm.PHReg.from_formula(fml, status="ystatus", data=df) fml = "mtime ~ exp" mediator_model = sm.PHReg.from_formula(fml, status="mstatus", data=df) med = Mediation(outcome_model, mediator_model, "exp", "mtime", outcome_predict_kwargs={"pred_only": True}, outcome_fit_kwargs={"method": "lbfgs"}, mediator_fit_kwargs={"method": "lbfgs"}) med_result = med.fit(n_rep=2) dr = med_result.summary() pm = dr.loc["Prop. mediated (average)", "Estimate"] if mt == "no": assert_allclose(pm, 0, atol=0.1, rtol=0.1) elif mt == "full": assert_allclose(pm, 1, atol=0.1, rtol=0.1) else: assert_allclose(pm, 0.5, atol=0.1, rtol=0.1)