
p7]c           @   s  d  d l  Z d  d l m Z m Z m Z d  d l m Z d  d l m	 Z	 d  d l
 m Z d  d l m Z m Z m Z m Z m Z m Z m Z m Z d   Z d   Z d	   Z d
   Z d   Z d   Z d   Z d   Z d   Z d   Z d   Z d   Z  d   Z! d   Z" d   Z# d S(   iN(   t   assert_equalt   assert_t   assert_allclose(   t   OLS(   t   GLM(   t   Binomial(   t
   _calc_gradt   _calc_wdesign_matt   _est_regularized_debiasedt   _join_debiasedt   _est_regularized_naivet   _est_unregularized_naivet   _join_naivet   DistributedModelc   	      c   s  | j  d } t j | |  } t j | |  } | | | } d } xj | | d k  r t | |  } t | | d  } |  | | !| | |  d d  f f V| d 7} qJ Wt | |  } t | | d |  } |  | | !| | |  d d  f f Vd S(   s   partitions datai    i   N(   t   shapet   npt   ceilt   floort   int(	   t   endogt   exogt
   partitionst   n_exogt   n_partt   remt   stpt   iit   jj(    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt	   _data_gen   s    (c          C   s   t  j j d  t  j j d d  }  t  j j d d d d } t  j j d d  } t | |   } t | | d d i   } t | t  j d	 d
 d g  d d d d d  S(   NiA t   sizei2   i   i    i   g{Gz?i   g"3@gy~g^zoL@t   atolgư>t   rtol(   i2   i   (	   R   t   randomt   seedt   normalt   randintR   R   R   t   array(   t   Xt   yt   betat   modt   grad(    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_calc_grad"   s    c          C   s0  t  j j d  t  j j d d  }  t  j j d d d d } t  j j d d  } t | |   } t | | i   } t | t  j d d d g d	 d
 d g d d d g g  d d d d t	 | |  d t
   } t | | i   } t | t  j d d d g d d d g d d d g g  d d d d d  S(   NiA R   i   i    i   g>D?gc~gYU9?g7L4HAgXE޿gi4g
g(5
?g >R   gư>R   t   familygZ'&?glTOg F6?gSпgPv0οg/m8,տgLƼg3ڪ$?gSVs(   i   i   (   R   R    R!   R"   R#   R   R   R   R$   R   R   (   R%   R&   R'   R(   t   dmat(    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_calc_wdesign_mat3   s     c    	      C   sT  t  j j d  t  j j d d  }  t  j j d d d d } t  j j d d  } t | |   } t | d d d i d d	 6} | d } | d
 } | d } | d } t t | t	   t
 | j | j  t
 | j | j  t t | t   t t | t   t
 t |  t |   t
 | d j d  t t | d t   t | |  d t   } t | d d d i d d	 6} | d } | d
 } | d } | d } t t | t	   t
 | j | j  t
 | j | j  t t | t   t t | t   t
 t |  t |   t
 | d j d  t t | d t   d  S(   NiA R   i2   i   i    i   t   fit_kwdsg      ?t   alphai   R+   (   i2   i   (   i   (   i   (   R   R    R!   R"   R#   R   R   R   t
   isinstancet   tupleR    R   t   listt   lent   floatR   R   (	   R%   R&   R'   R(   t   rest   bhatR)   t   ghat_lt   that_l(    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_est_regularized_debiasedM   s@    







c          C   s   t  j j d  t  j j d d  }  t  j j d d d d } t  j j d d  } t | |   } t | d d d i d d	 6} t | j | j  t	 | |  d
 t
   } t | d d d i d d	 6} t | j | j  d  S(   NiA R   i2   i   i    i   R.   g      ?R/   R+   (   i2   i   (   R   R    R!   R"   R#   R   R
   R    R   R   R   (   R%   R&   R'   R(   R5   (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_est_regularized_naivex   s    c          C   s   t  j j d  t  j j d d  }  t  j j d d d d } t  j j d d  } t | |   } t | d d d i d d	 6} t | j | j  t	 | |  d
 t
   } t | d d d i d d	 6} t | j | j  d  S(   NiA R   i2   i   i    i   R.   g      ?R/   R+   (   i2   i   (   R   R    R!   R"   R#   R   R   R    R   R   R   (   R%   R&   R'   R(   R5   (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_est_unregularized_naive   s    c       	   C   sk  t  j j d  t  j j d d  }  t  j j d d d d } t | |   } g  } x@ t d  D]2 } t | | d d i d d	 6} | j |  qb Wt	 |  } t
 | t  j d
 d d g  d d d d t | |  d t   } g  } x@ t d  D]2 } t | | d d i d d	 6} | j |  q Wt	 |  } t
 | t  j d d d g  d d d d d  S(   NiA R   i2   i   i    i   R.   g?R/   gNA~6rſguŌgJOcֿR   gư>R   R+   gYſgKu/3lڿg{Ic̿(   i2   i   (   R   R    R!   R"   R#   R   t   rangeR   t   appendR	   R   R$   R   R   (   R%   R&   R(   t   res_lt   iR5   t   joined(    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_join_debiased   s&    c       	   C   sk  t  j j d  t  j j d d  }  t  j j d d d d } t | |   } g  } x@ t d  D]2 } t | | d d i d d	 6} | j |  qb Wt	 |  } t
 | t  j d
 d d g  d d d d t | |  d t   } g  } x@ t d  D]2 } t | | d d i d d	 6} | j |  q Wt	 |  } t
 | t  j d d d g  d d d d d  S(   NiA R   i2   i   i    i   R.   g?R/   gYRAg        R   gư>R   R+   (   i2   i   (   R   R    R!   R"   R#   R   R<   R
   R=   R   R   R$   R   R   (   R%   R&   R(   R>   R?   R5   R@   (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_join_naive   s&    c          C   s  t  j j d  t  j j d d%  }  t  j j d d d d } t d d t } | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d d g  d d d d t d d t } | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d d g  d d d d t d d t } | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d d g  d d d d t d d t d i t   d 6} | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d d g  d d d d t d d t d i t   d 6} | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d  d! g  d d d d t d d t d i t   d 6} | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d" d# d$ g  d d d d d  S(&   NiA R   i2   i   i    i   i   t   model_classt   parallel_methodt
   sequentialR.   g      ?R/   g۟ȿg$gNx	N}ֿR   gư>R   gP5&Ŀg8Zgۃ/޿gHgIpg|}Kٿt	   init_kwdsR+   gYſgKu/3lڿg{Ic̿gܞ =¿gd׿gY9ҿg@CHgĐӿg_8Ͽ(   i2   i   (   R   R    R!   R"   R#   R   R   t   fitR   R   t   paramsR$   R   R   (   R%   R&   R(   RG   (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_fit_sequential   sH    !!!!!!c          C   s  t  j j d  t  j j d d%  }  t  j j d d d d } t d d t } | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d d g  d d d d t d d t } | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d d g  d d d d t d d t } | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d d g  d d d d t d d t d i t   d 6} | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d d g  d d d d t d d t d i t   d 6} | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d d  d! g  d d d d t d d t d i t   d 6} | j t | |  d  d	 d
 d i d d 6} t	 | j
 t  j d" d# d$ g  d d d d d  S(&   NiA R   i2   i   i    i   i   RC   RD   t   joblibR.   g      ?R/   g۟ȿg$gNx	N}ֿR   gư>R   gP5&Ŀg8Zgۃ/޿gHgIpg|}KٿRF   R+   gYſgKu/3lڿg{Ic̿gܞ =¿gd׿gY9ҿg@CHgĐӿg_8Ͽ(   i2   i   (   R   R    R!   R"   R#   R   R   RG   R   R   RH   R$   R   R   (   R%   R&   R(   RG   (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_fit_joblib  sH    !!!!!!c          C   s  t  j j d  d }  d } d } t  j j d |  } | t  j j d d |  } t  j j d |  | f  } | j |  t  j j d |   } t |  } | j t | | |  d i d d	 6} t | d
 t	 d t
 } | j t | | |  d i d d	 6}	 t | |  }
 |
 j d	 d  } t | j | j  t |	 j | j  t | d
 t	 d t
 } | j t | | |  d i d d	 6}	 t | |  }
 |
 j d	 d  } t |	 j | j  d  S(   NiA i   i
   i   R   i    i   R.   R/   t   estimation_methodt   join_methodg?(   R   R    R!   R"   R#   t   dotR   RG   R   R
   R   R   R   RH   t   fit_regularized(   t   Nt   pt   mR'   R%   R&   t   db_modt   fitOLSdbt   nv_modt   fitOLSnvt   ols_modt   fitOLS(    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_single_partition<  s.    "(	(	(c    
      C   sL  t  j j d  d }  d } d } t  j j d |  } | t  j j d d |  } t  j j d |  | f  } | j |  t  j j d |   } t |  } | j t | | |  d i d d	 6} t	 t  j
 t  j | j   d  t | d
 t d t } | j t | | |  d i d d	 6}	 t	 t  j
 t  j |	 j   d  d  S(   NiA i(   i   R   i    i   R.   g?R/   RL   RM   (   R   R    R!   R"   R#   RN   R   RG   R   R    t   sumt   isnanRH   R
   R   (
   RP   RQ   RR   R'   R%   R&   RS   RT   RU   RV   (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_larger_pc  s    "("	(c          C   s)  t  j j d  d }  d } d } t  j j d |  } | t  j j d d |  } t  j j d |  | f  } | j |  t  j j d |   } t | d i d	 d
 6} | j t | | |  d i d d 6} t	 | |  } | j
 d d  }	 d | j d k }
 d |	 j d k } t |
 |  d  S(   NiA i   i
   i   R   i    i   t	   join_kwdsgp=
ף?t	   thresholdR.   g?R/   i   (   R   R    R!   R"   R#   RN   R   RG   R   R   RO   RH   R   (   RP   RQ   RR   R'   R%   R&   RS   RT   RW   RX   t   nz_params_dbt   nz_params_ols(    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_non_zero_params{  s    "(c          C   s  t  j j d  d }  d } d } t  j j d |  } | t  j j d d |  } t  j j d |  | f  } | j |  t  j j d |   } d   } t | d	 t d
 t } | j	 | | | |  d i d d 6} t
 | |  }	 |	 j d d  }
 t | j |
 j  d  S(   NiA i   i
   i   R   i    i   c         s   sX   | j  d } t j | |  } d } x+ | | k  rS |  | f V| t |  7} q) Wd S(   s   partitions datai    N(   R   R   R   R   (   R   R   R   R   R   R   (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   _rep_data_gen  s    RL   RM   R.   g?R/   (   R   R    R!   R"   R#   RN   R   R
   R   RG   R   RO   R   RH   (   RP   RQ   RR   R'   R%   R&   Rb   RU   RV   RW   RX   (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_repeat_partition  s    "		(c       
   C   sp  t  j j d  d }  d } d } t  j j d |  } | t  j j d d |  } t  j j d |  | f  } | j |  t  j j d |   } t |  } | j t | | |  d i d	 d
 6} t  j	 j
 | j |  } t | d t d t }	 |	 j t | | |  d i d	 d
 6}
 t  j	 j
 |
 j |  } t | | k   d d t  j | j |  t  j j d |    } d | d k } t | d t d i t   d 6} | j t | | |  d i d	 d
 6} t  j	 j
 | j |  } t | d t d i t   d 6d t d t }	 |	 j t | | |  d i d	 d
 6} t  j	 j
 | j |  } t | | k   d  S(   NiA i   i
   i   R   i    i   R.   g?R/   RL   RM   i   g      ?g      ?RC   RF   R+   (   R   R    R!   R"   R#   RN   R   RG   R   t   linalgt   normRH   R
   R   R   t   expR   R   (   RP   RQ   RR   R'   R%   R&   RS   RT   t   olsdbt   n_modt   fitOLSnt   olsnt   probt   fitGLMdbt   glmdbt   fitGLMnt   glmn(    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   test_debiased_v_average  s:    "(	(4(	(($   t   numpyR   t   numpy.testingR    R   R   t#   statsmodels.regression.linear_modelR   t+   statsmodels.genmod.generalized_linear_modelR   t   statsmodels.genmod.familiesR   t'   statsmodels.base.distributed_estimationR   R   R   R	   R
   R   R   R   R   R*   R-   R9   R:   R;   RA   RB   RI   RK   RY   R\   Ra   Rc   Rp   (    (    (    sQ   lib/python2.7/site-packages/statsmodels/base/tests/test_distributed_estimation.pyt   <module>   s(   :				+					0	0	'			$