
    !g/                        d dl Zd dlmZ d dlZd dlmZ d dlZ	d dl
mZ d dlZd dlZg dg dg dg dg dg d	g d
g dg dg dg dgZ e	j                  edd ed          j!                  d      Zg dg dg dg dg dg dg dg dg dg dg dgZ e	j                  edd ed          j!                  d      Zg dg dg dg dg dg d g d!g d"g d#g d$g d%gZ e	j                  edd ed          j!                  d      Zej(                  j*                  d&        Zd' Zej(                  j*                  d(        Zej(                  j*                  d)        Zd* Zd+ Zy),    N)	Mediation)assert_allclose)indexEstimatezLower CI boundzUpper CI boundzP-value)ACME (control)gɵ?gzn+?gH}?        )ACME (treated)g4?gVy՟?g~?r   )ADE (control)g^=1X?gUq7g$7?Q?)ADE (treated)gb?g3gMۿҤ?r   )Total effectgY-D?gGRТg$P?皙?)Prop. mediated (control)g2V?g%@M-[g켍@r   )Prop. mediated (treated)g**?g<gt@r   )ACME (average)ggyܵ?g#?g  ?r   )ADE (average)g;ۤ?g.s,&6gN\W ?r   )Prop. mediated (average)g>W[?g0!g0(h@r      )columnsr   )r   g#KXU?gcAJ?gEHξ?r   )r	   g*q㊋?g	h"?g\?r   )r
   g }?g5BX%g.?(\?)r   g'n?gj깿g	ܺ?r   )r   gP9?g+]gCg|r?{Gz?)r   gzk`?gX;sQg^J]2a@r   )r   g(yu?g:fg|}K@r   )r   g:p?gjTQ?g6;Nё?r   )r   gy7R?g+,gWv;?r   )r   g&m?ghugc@r   )r   g{O崧?g.c}k?gfI}?{Gz?)r	   gt	4ش?g|b*3?g^yȔ?r   )r
   g*A*?g7^ꚿg-AF@#?{Gz?)r   g@Û5x_?g.ojgu7Ou?r   )r   g"?g>V?gwE?r   )r   g8*5?gEn1?g"3<?r   )r   gw?gw/Q?g|.P?r   )r   gR||Bv޲?g _>Y1|?gfJo	?r   )r   gk$	?gvť*g&9`W?r   )r   g	?gW[Ɏ?gׁsF?r   c                  R   t         j                  j                  t         j                  j                  t                    } t        j                  t         j                  j                  | dd            }t        j                  |d         }t        j                  d|d      }t        j                  ||t        j                  j                  t        j                  j                   j#                                     }t        j                  |d	         }t        j                  d
|d      }t        j$                  ||      }|j&                  j)                         j+                  d      |j&                  j)                         j+                  d      g}|j&                  j)                         j+                  d	      }	t-        ||||	ddi      }
t        j.                  j1                  d       |
j3                  dd      }t        j                  |j5                         t6        z
        }t9        |dd       t        j.                  j1                  d       |
j3                  dd      }t        j                  |j5                         t:        z
        }t9        |dd       y )Nresultsframing.csv	cong_mesg*emo + treat + age + educ + gender + income	dataframereturn_typelinkfamilyemo$treat + age + educ + gender + incometreatatoldy=outcome_fit_kwargs  
parametricd   methodn_repr   ư>r)   boot)ospathdirnameabspath__file__pdread_csvjoinnpasarraypatsydmatrixsmGLMfamiliesBinomiallinksProbitOLSr   tolistr   r   randomseedfitsummaryframing_para_4231r   framing_boot_4231)cur_dirdataoutcomeoutcome_exogoutcome_modelmediatormediator_exogmediator_modeltx_posmed_posmed	para_rsltdiff	boot_rslts                 c/var/www/dash_apps/app1/venv/lib/python3.12/site-packages/statsmodels/stats/tests/test_mediation.pytest_framing_exampler_   6   s    ggoobggooh78G;;rww||GYFGDjjk*+G==!Mt.9;LFF{{##):):)A)A)C#DM zz$u+&HMM"H$-8:MVVHm4N""))+11':##**,227;=F""))+11%8G
M>67(.u~7C IINN4|37I::i''),==>DD!$'IINN4vS1I::i''),==>DD!$'    c                     t         j                  j                  t         j                  j                  t                    } t        j                  t         j                  j                  | dd            }t        j                  |d         }t        j                  d|d      }t        j                  ||t        j                  j                  t        j                  j                   j#                                     }t        j                  |d	         }t        j                  d
|d      }t        j$                  ||      }|j&                  j)                         j+                  d      |j&                  j)                         j+                  d      g}|j&                  j)                         j+                  d	      }	|j&                  j)                         j+                  d      |j&                  j)                         j+                  d      f}
|
di}t-        ||||	|      }t        j.                  j1                  d       |j3                  dd      }y )Nr   r   r   r   r   r    r"   r$   r&   r'   r(   age   
moderatorsr-   r.   r/   r0   )r6   r7   r8   r9   r:   r;   r<   r=   r>   r?   r@   rA   rB   rC   rD   rE   rF   rG   rH   r   rI   r   r   rJ   rK   rL   )rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   ixre   rZ   med_rslts                 r^   test_framing_example_moderatorrh   ]   s    ggoobggooh78G;;rww||GYFGDjjk*+G==!Mt.9;LFF{{##):):)A)A)C#DM zz$u+&HMM"H$-8:MVVHm4N""))+11':##**,227;=F""))+11%8G



%
%
'
-
-e
4



&
&
(
.
.u
5
7BrJ
M>67)+C IINN4wwl#w6Hr`   c                     t         j                  j                  t         j                  j                  t                    } t        j                  t         j                  j                  | dd            }t        j                  j                  d|t        j                  j                  t        j                  j                  j                                     }t        j                  j                  d|      }t!        ||ddd	d
i      }t"        j$                  j'                  d       |j)                  dd      }t#        j*                  |j-                         t.        z
        }t1        |dd       t"        j$                  j'                  d       |j)                  dd      }t#        j*                  |j-                         t2        z
        }t1        |dd       y )Nr   r   z6cong_mesg ~ emo + treat + age + educ + gender + incomer"   r$   z*emo ~ treat + age + educ + gender + incomer(   r&   r)   r*   r+   r-   r5   r/   r0   r   r3   r4   r.   )r6   r7   r8   r9   r:   r;   r<   r=   rB   rC   from_formularD   rE   rF   rG   rH   r   r>   rJ   rK   rL   r?   rM   rO   r   rN   )rP   rQ   rT   rW   rZ   rg   r\   s          r^   test_framing_example_formulark      sT    ggoobggooh78G;;rww||GYFGDFF''@{{##):):)A)A)C#D ( M VV(()UW[\N
M>7E(.8C IINN4wwfCw0H::h&&(+<<=DD!$'IINN4wwl#w6H::h&&(+<<=DD!$'r`   c                     t         j                  j                  t         j                  j                  t                    } t        j                  t         j                  j                  | dd            }t        j                  j                  d|t        j                  j                  t        j                  j                  j                                     }t        j                  j                  d|      }ddi}t!        ||d	d
|      }t"        j$                  j'                  d       |j)                  dd      }t#        j*                  |j-                         t.        z
        }t1        |dd       y )Nr   r   z>cong_mesg ~ emo + treat*age + emo*age + educ + gender + incomer"   r$   z(emo ~ treat*age + educ + gender + incomerb   rc   r(   r&   rd   r-   r.   r/   r0   r   r3   r4   )r6   r7   r8   r9   r:   r;   r<   r=   rB   rC   rj   rD   rE   rF   rG   rH   r   r>   rJ   rK   rL   r?   rM   framing_moderated_4231r   )rP   rQ   rT   rW   re   rZ   rg   r\   s           r^   &test_framing_example_moderator_formularn      s    ggoobggooh78G;;rww||GYFGDFF''H{{##):):)A)A)C#D ( M VV(()SUYZN"J
M>7E)+C IINN4wwl#w6H::h&&(+AABDD!$'r`   c                     t         j                  j                  d       t         j                  j                  d      } dt         j                  j	                  dd      dz
  z  }t        j
                  d      |z   |j                  z   }t         j                  j                  | |      }g d}t        ||d       t         j                  j                  d       d}t         j                  j                  |	      }t        j                  |t        j                  d
            }t        j                  g d      }|t        j                  t        j                  |dz              z  }|d   t        j                  |t        j                  d
            z  }	|	|d   t        j                  t         j                  j                  |	      t        j                  d
            z  z  }	|	|d   t         j                  j                  |d
f	      z  z  }	t        j                  |t         j                  d         |	z   }
t        j                  g d      }|t        j                  t        j                  |dz              z  }|d   t        j                  t         j                  j                  |	      t        j                  d
            z  }||d   t        j                  t         j                  j                  |	      t         j                  d         z  z  }||d   t         j                  j                  |d
f	      z  z  }|
|z   }t        j                  t        j                   |      t        j                  d
            }t        j                  t        j                  |      t         j                  d         }t#        j$                  |j'                         |j'                         |j'                         |j'                         |	j'                         d      }g d}t        t        j                  |j)                               |d       t*        j,                  j/                  dd|      }t*        j,                  j/                  dd|      }t1        ||dd      }t         j                  j                  d       |j3                  d      }|j5                         }|j6                  j3                         j8                  j;                         }g d}t        ||d       |j<                  j3                         j8                  j;                         }g d}t        ||d       t        j>                  g dg dg d g d g d!g d"g d"g dg d g d"g
      }t        |j;                         |d#$       t        |j@                  d%   d&d'd'(       |jB                  d)   }t        |d&d'd'(       y )*Ni`     g-C6?      ?)g~/[|?gV_~0?gK&<?g
|w?g(RHgHz>r4      size   )g      @   r      r   r   )r   rq   r   )g      @rw   rw   )r   r   )yxidtimerZ   )gn<vgfB¿g     X@r   gNXTÿzmed ~ xr{   )groupsrQ   zy ~ med + xrz   rZ   i r/   r2   )gjtȴ?g,'?gSA?g	?)glHEg
ۗ?g]ݱ&1@)g?gZRr?gDѰX?r   )gttd4?g!&y:?gCY?r   )gI?gV?gR?r   )guo?g`h?g >Ӧ%?r   g333333?)rtol)rx   r   gQ?g{Gz?)r   r)   r   r   )"r>   rJ   rK   randnrandeyeTmultivariate_normalr   normalouteronesr?   sqrtsumr_aranger;   	DataFrameflattenmeanrB   MixedLMrj   r   rL   rM   rT   paramsto_numpyrW   arrayilocloc)mnccovrvsrvs1nrz   xvmxrZ   eyexery   idxtimdfdmeanrW   rT   memrst	params_omp_om	params_mmp_mmres_summpms                                r^   t_est_mixedlmr      s`    IINN4		B		q!$s*+A
&&)a-!##
C
))
'
'C
0CFDCD)IINN4A 			a A	!RWWQZ	 B 
J	B"''"&&Q-
  B
Q%"((1bggaj)
)C2a5288BII,,!,4bggajAAAC2a5299##!Q#000C 
!RUU9%	&	,B 
J	B"''"&&Q-
  B
1))q)12771:>>AA"))***2BEE(OD	DDAA!!1v!.	..A
QA ((299Q<
,C ((2771:ruuX
/C	AIIKbjjl [[]CKKM!kkm- 
.B
 ?EBJJrwwy)5t<ZZ,,Yt",MNJJ++M$R+PM	=.#u	=BIINN6	c	B	B   $$&--668I;DIt$/!!%%'..779I0DIt$/ xx1111111111
 H BKKM8$7BGGENDt$?	6	7BB4d3r`   c                  t   t         j                  j                  d       d} t         j                  j                  |       }t        j                  |      }| t        j
                  t         j                  j                  |             z  }d|z  t        j
                  t         j                  j                  |             z  }||k\  j                  t              }t        j                  ||k  ||      }dD ]  }|dk(  rd|z  }n|dk(  r||z   }n|}t        j                  |       }| t        j
                  t         j                  j                  |             z  }	d|z  t        j
                  t         j                  j                  |             z  }||	k\  j                  t              }
t        j                  |	|k  |	|      }t        j                  ||
|||d	      }d
}t        j                  j                  |d|      }d}t        j                  j                  |d|      }t        ||ddddiddiddi      }|j!                  d      }|j#                         }|j$                  d   }|dk(  rt'        |ddd       |dk(  rt'        |ddd       t'        |ddd        y )Ni%	  i  rs   )fullpartialnor   rq   r   )ytimeystatusmtimemstatusexpzytime ~ exp + mtimer   )statusrQ   zmtime ~ expr   r   r   	pred_onlyTr1   lbfgs)outcome_predict_kwargsr,   mediator_fit_kwargsrw   r~   r   r   r   g?)r)   r   r   )r>   rJ   rK   r   r   loguniformastypeintwherer;   r   rB   PHRegrj   r   rL   rM   r   r   )r   r   r   mtime0ctimer   r   mtlpytime0r   r   r   fmlrT   rW   rZ   
med_resultdrr   s                       r^   	test_survr     sn   IINN4A ))



"C 
BS266"))+++344FGbffRYY..A.677E&&s+GHHVu_fe4E% &9 <VB9_vBB VVRC[rvvbii//Q/788R"&&!2!2!2!:;;F?**3/5&%8\\Eg$)g"%' ( $--c)"-M..s92.N~ug0;T/B,4g+>-5w,?A WW1W%
!VV:;:B#66\B#6B#C8M&9r`   )numpyr>   statsmodels.apiapirB   r6   statsmodels.stats.mediationr   pandasr;   numpy.testingr   r@   pytestr   r   	set_indexrO   rN   rm   markslowr_   rh   rk   rn   r   r    r`   r^   <module>r      s     	 1  )   K<<<<;GG<<G
I !BLLABA7AA'J  K<<<<;GG<<G
I !BLLABA7AA'J  K<<<<:FF<<F
H &bfbe<FFwO  "( "(L 7F ( (6 ( (.T4n69r`   