<s> """<STR_LIT>""" <EOL> import numpy as np <EOL> import pandas as pd <EOL> from statsmodels . graphics . utils import maybe_name_or_idx <EOL> import statsmodels . compat . pandas as pdc <EOL> class Mediation ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , outcome_model , mediator_model , exposure , mediator = None , <EOL> moderators = None , outcome_fit_kwargs = None , mediator_fit_kwargs = None ) : <EOL> self . outcome_model = outcome_model <EOL> self . mediator_model = mediator_model <EOL> self . exposure = exposure <EOL> self . moderators = moderators if moderators is not None else { } <EOL> if mediator is None : <EOL> self . mediator = self . _guess_endog_name ( mediator_model , '<STR_LIT>' ) <EOL> else : <EOL> self . mediator = mediator <EOL> self . _outcome_fit_kwargs = ( outcome_fit_kwargs if outcome_fit_kwargs <EOL> is not None else { } ) <EOL> self . _mediator_fit_kwargs = ( mediator_fit_kwargs if mediator_fit_kwargs <EOL> is not None else { } ) <EOL> self . _outcome_exog = outcome_model . exog . copy ( ) <EOL> self . _mediator_exog = mediator_model . exog . copy ( ) <EOL> self . _exp_pos_mediator = self . _variable_pos ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . _exp_pos_outcome = self . _variable_pos ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . _med_pos_outcome = self . _variable_pos ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> def _variable_pos ( self , var , model ) : <EOL> if model == '<STR_LIT>' : <EOL> mod = self . mediator_model <EOL> else : <EOL> mod = self . outcome_model <EOL> if var == '<STR_LIT>' : <EOL> return maybe_name_or_idx ( self . mediator , mod ) [ <NUM_LIT:1> ] <EOL> exp = self . exposure <EOL> exp_is_2 = ( ( len ( exp ) == <NUM_LIT:2> ) and ( type ( exp ) != type ( '<STR_LIT>' ) ) ) <EOL> if exp_is_2 : <EOL> if model == '<STR_LIT>' : <EOL> return exp [ <NUM_LIT:0> ] <EOL> elif model == '<STR_LIT>' : <EOL> return exp [ <NUM_LIT:1> ] <EOL> else : <EOL> return maybe_name_or_idx ( exp , mod ) [ <NUM_LIT:1> ] <EOL> def _guess_endog_name ( self , model , typ ) : <EOL> if hasattr ( model , '<STR_LIT>' ) : <EOL> return model . formula . split ( "<STR_LIT>" ) [ <NUM_LIT:0> ] . strip ( ) <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' % typ ) <EOL> def _simulate_params ( self , result ) : <EOL> """<STR_LIT>""" <EOL> mn = result . params <EOL> cov = result . cov_params ( ) <EOL> return np . random . multivariate_normal ( mn , cov ) <EOL> def _get_mediator_exog ( self , exposure ) : <EOL> """<STR_LIT>""" <EOL> mediator_exog = self . _mediator_exog <EOL> if not hasattr ( self . mediator_model , '<STR_LIT>' ) : <EOL> mediator_exog [ : , self . _exp_pos_mediator ] = exposure <EOL> for ix in self . moderators : <EOL> v = self . moderators [ ix ] <EOL> mediator_exog [ : , ix [ <NUM_LIT:1> ] ] = v <EOL> else : <EOL> df = self . mediator_model . data . frame . copy ( ) <EOL> df . loc [ : , self . exposure ] = exposure <EOL> for vname in self . moderators : <EOL> v = self . moderators [ vname ] <EOL> df . loc [ : , vname ] = v <EOL> klass = self . mediator_model . __class__ <EOL> init_kwargs = self . mediator_model . _get_init_kwds ( ) <EOL> model = klass . from_formula ( data = df , ** init_kwargs ) <EOL> mediator_exog = model . exog <EOL> return mediator_exog <EOL> def _get_outcome_exog ( self , exposure , mediator ) : <EOL> """<STR_LIT>""" <EOL> outcome_exog = self . _outcome_exog <EOL> if not hasattr ( self . outcome_model , '<STR_LIT>' ) : <EOL> outcome_exog [ : , self . _med_pos_outcome ] = mediator <EOL> outcome_exog [ : , self . _exp_pos_outcome ] = exposure <EOL> for ix in self . moderators : <EOL> v = self . moderators [ ix ] <EOL> outcome_exog [ : , ix [ <NUM_LIT:0> ] ] = v <EOL> else : <EOL> df = self . outcome_model . data . frame . copy ( ) <EOL> df . loc [ : , self . exposure ] = exposure <EOL> df . loc [ : , self . mediator ] = mediator <EOL> for vname in self . moderators : <EOL> v = self . moderators [ vname ] <EOL> df . loc [ : , vname ] = v <EOL> klass = self . outcome_model . __class__ <EOL> init_kwargs = self . outcome_model . _get_init_kwds ( ) <EOL> model = klass . from_formula ( data = df , ** init_kwargs ) <EOL> outcome_exog = model . exog <EOL> return outcome_exog <EOL> def _fit_model ( self , model , fit_kwargs , boot = False ) : <EOL> klass = model . __class__ <EOL> init_kwargs = model . _get_init_kwds ( ) <EOL> endog = model . endog <EOL> exog = model . exog <EOL> if boot : <EOL> ii = np . random . randint ( <NUM_LIT:0> , len ( endog ) , len ( endog ) ) <EOL> endog = endog [ ii ] <EOL> exog = exog [ ii , : ] <EOL> outcome_model = klass ( endog , exog , ** init_kwargs ) <EOL> return outcome_model . fit ( ** fit_kwargs ) <EOL> def fit ( self , method = "<STR_LIT>" , n_rep = <NUM_LIT:1000> ) : <EOL> """<STR_LIT>""" <EOL> if method . startswith ( "<STR_LIT>" ) : <EOL> outcome_result = self . _fit_model ( self . outcome_model , self . _outcome_fit_kwargs ) <EOL> mediator_result = self . _fit_model ( self . mediator_model , self . _mediator_fit_kwargs ) <EOL> elif not method . startswith ( "<STR_LIT>" ) : <EOL> raise ( "<STR_LIT>" ) <EOL> indirect_effects = [ [ ] , [ ] ] <EOL> direct_effects = [ [ ] , [ ] ] <EOL> for iter in range ( n_rep ) : <EOL> if method == "<STR_LIT>" : <EOL> outcome_params = self . _simulate_params ( outcome_result ) <EOL> mediation_params = self . _simulate_params ( mediator_result ) <EOL> else : <EOL> outcome_result = self . _fit_model ( self . outcome_model , <EOL> self . _outcome_fit_kwargs , boot = True ) <EOL> outcome_params = outcome_result . params <EOL> mediator_result = self . _fit_model ( self . mediator_model , <EOL> self . _mediator_fit_kwargs , boot = True ) <EOL> mediation_params = mediator_result . params <EOL> predicted_outcomes = [ [ None , None ] , [ None , None ] ] <EOL> for tm in <NUM_LIT:0> , <NUM_LIT:1> : <EOL> mex = self . _get_mediator_exog ( tm ) <EOL> gen = self . mediator_model . get_distribution ( mediation_params , <EOL> mediator_result . scale , <EOL> exog = mex ) <EOL> potential_mediator = gen . rvs ( mex . shape [ <NUM_LIT:0> ] ) <EOL> for te in <NUM_LIT:0> , <NUM_LIT:1> : <EOL> oex = self . _get_outcome_exog ( te , potential_mediator ) <EOL> po = self . outcome_model . predict ( outcome_params , oex ) <EOL> predicted_outcomes [ tm ] [ te ] = po <EOL> for t in <NUM_LIT:0> , <NUM_LIT:1> : <EOL> indirect_effects [ t ] . append ( predicted_outcomes [ <NUM_LIT:1> ] [ t ] - predicted_outcomes [ <NUM_LIT:0> ] [ t ] ) <EOL> direct_effects [ t ] . append ( predicted_outcomes [ t ] [ <NUM_LIT:1> ] - predicted_outcomes [ t ] [ <NUM_LIT:0> ] ) <EOL> for t in <NUM_LIT:0> , <NUM_LIT:1> : <EOL> indirect_effects [ t ] = np . asarray ( indirect_effects [ t ] ) . T <EOL> direct_effects [ t ] = np . asarray ( direct_effects [ t ] ) . T <EOL> self . indirect_effects = indirect_effects <EOL> self . direct_effects = direct_effects <EOL> rslt = MediationResults ( self . indirect_effects , self . direct_effects ) <EOL> rslt . method = method <EOL> return rslt <EOL> def _pvalue ( vec ) : <EOL> return <NUM_LIT:2> * min ( sum ( vec > <NUM_LIT:0> ) , sum ( vec < <NUM_LIT:0> ) ) / float ( len ( vec ) ) <EOL> class MediationResults ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , indirect_effects , direct_effects ) : <EOL> self . indirect_effects = indirect_effects <EOL> self . direct_effects = direct_effects <EOL> indirect_effects_avg = [ None , None ] <EOL> direct_effects_avg = [ None , None ] <EOL> for t in <NUM_LIT:0> , <NUM_LIT:1> : <EOL> indirect_effects_avg [ t ] = indirect_effects [ t ] . mean ( <NUM_LIT:0> ) <EOL> direct_effects_avg [ t ] = direct_effects [ t ] . mean ( <NUM_LIT:0> ) <EOL> self . ACME_ctrl = indirect_effects_avg [ <NUM_LIT:0> ] <EOL> self . ACME_tx = indirect_effects_avg [ <NUM_LIT:1> ] <EOL> self . ADE_ctrl = direct_effects_avg [ <NUM_LIT:0> ] <EOL> self . ADE_tx = direct_effects_avg [ <NUM_LIT:1> ] <EOL> self . total_effect = ( self . ACME_ctrl + self . ACME_tx + self . ADE_ctrl + self . ADE_tx ) / <NUM_LIT:2> <EOL> self . prop_med_ctrl = self . ACME_ctrl / self . total_effect <EOL> self . prop_med_tx = self . ACME_tx / self . total_effect <EOL> self . prop_med_avg = ( self . prop_med_ctrl + self . prop_med_tx ) / <NUM_LIT:2> <EOL> self . ACME_avg = ( self . ACME_ctrl + self . ACME_tx ) / <NUM_LIT:2> <EOL> self . ADE_avg = ( self . ADE_ctrl + self . ADE_tx ) / <NUM_LIT:2> <EOL> def summary ( self , alpha = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> columns = [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] <EOL> index = [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , <EOL> "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , <EOL> "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] <EOL> smry = pd . DataFrame ( columns = columns , index = index ) <EOL> for i , vec in enumerate ( [ self . ACME_ctrl , self . ACME_tx , self . ADE_ctrl , self . ADE_tx , <EOL> self . total_effect , self . prop_med_ctrl , <EOL> self . prop_med_tx , self . ACME_avg , self . ADE_avg , <EOL> self . prop_med_avg ] ) : <EOL> if ( ( vec is self . prop_med_ctrl ) or ( vec is self . prop_med_tx ) or <EOL> ( vec is self . prop_med_avg ) ) : <EOL> smry . iloc [ i , <NUM_LIT:0> ] = np . median ( vec ) <EOL> else : <EOL> smry . iloc [ i , <NUM_LIT:0> ] = vec . mean ( ) <EOL> smry . iloc [ i , <NUM_LIT:1> ] = np . percentile ( vec , <NUM_LIT:100> * alpha / <NUM_LIT:2> ) <EOL> smry . iloc [ i , <NUM_LIT:2> ] = np . percentile ( vec , <NUM_LIT:100> * ( <NUM_LIT:1> - alpha / <NUM_LIT:2> ) ) <EOL> smry . iloc [ i , <NUM_LIT:3> ] = _pvalue ( vec ) <EOL> if pdc . version < '<STR_LIT>' : <EOL> smry = smry . convert_objects ( convert_numeric = True ) <EOL> else : <EOL> smry = smry . apply ( pd . to_numeric , errors = '<STR_LIT>' ) <EOL> return smry </s>
