o
    #e1y                     @   sJ  d dl mZ d dlZd dlZd dlZd dlmZ d dlm	Z	m
Z
mZ d dlmZmZ d dlT d dlmZ d dlmZmZ d d	lmZmZ e
rJd dlZd
d Zdd Zdd Zd5ddZd5ddZdd Zdd Zdd Z de _!dd Z"dd Z#dd  Z$d!d" Z%d#d$ Z&d%d& Z'd'd( Z(d)d* Z)d+d, Z*d-d. Z+d/d0 Z,d1d2 Z-d3d4 Z.dS )6    )print_functionN)
PatsyError)atleast_2d_column_defaulthave_pandashave_pandas_categorical)Term	INTERCEPT)*)C)balancedLookupFactor)DesignMatrix
DesignInfoc                 C   sP   t | } | jd dkrdS tj| \}}}t|dk}|| jd ks&J d S )N   r   Tg|=)r   shapenpZlinalgZsvdsum)musvZrank r   0lib/python3.10/site-packages/patsy/test_build.pyassert_full_rank   s   r   c                   C   s   t td t ddgddgddgddgg ttt ddgddgg ttt ddgddgg ttt g dg dg ttt g dg dg d	g d S )
N
   r   r         r   r      )r   r   d   )r         )r   r!      )r   r   ZeyepytestraisesAssertionErrorr   r   r   r   test_assert_full_rank#   s    r&   c                  G   s*   g }| D ]}| tdd |D  q|S )Nc                 S   s   g | ]}t |qS r   )r   ).0namer   r   r   
<listcomp>3   s    z!make_termlist.<locals>.<listcomp>)appendr   )entriestermsentryr   r   r   make_termlist0   s   r.   c                 C   s\   t |  t| jjt|ksJ |d ur| jj|ksJ | jdks#J | jd |ks,J d S )Nr   r   )r   setdesign_infor,   column_namesndimr   )Zmmexpected_ranktermlistr1   r   r   r   check_design_matrix6   s   r5   c           	         s   t | } fdd}t|g|dd}t| }|d }|d j|jjks&J |d j|jjks1J |j|d u s:J t||||d |S )Nc                   3        V  d S Nr   r   datar   r   
iter_maker@      
zmake_matrix.<locals>.iter_makerr   Zeval_envr1   )r.   design_matrix_buildersbuild_design_matricesZterm_slicesr0   r1   r5   )	r9   r3   r+   r1   r4   r:   Zdesign_infosZmatricesmatrixr   r8   r   make_matrix>   s    
rA   c                  C   sV  t ddd} tddt| d  }| d< | d d  }| d< t| ddggdd	gd
}t|ddgddgddgddggs?J t| dg dggddgd
}t|ddgddgddgddggs`J t| dddggg dd
}t|g dg dg dg dgsJ t| dg dgdgddggg dd
}t|g dg dg dg dgsJ t| dg dgdgddggg dd
}t|g dg dg dg dgsJ t| ddgdgddggg dd
}t|dd|d dgdd|d dgdd|d |d gdd|d |d ggsJ t| ddgdgddggg dd
}t|t|||| fs)J d S )Nr   )abr   r   rB   x1x2a[a1]a[a2]r=   	Intercepta[T.a2]r   rC   )za[a1]:b[b1]za[a2]:b[b1]za[a1]:b[b2]za[a2]:b[b2])r   r   r   r   )r   r   r   r   )r   r   r   r   )r   r   r   r   )rH   rI   b[T.b2]za[T.a2]:b[T.b2]r   r   r   r   )r   r   r   r   )r   r   r   r   )rH   rJ   rI   zb[T.b2]:a[T.a2])rF   rG   rD   z
a[T.a2]:x1r   )rD   rE   zx2:x1)r   r   linspacelenrA   allcloseZcolumn_stack)r9   rD   rE   r   r   r   r   test_simpleN   sV   ((


&rO   c                  C   sp   t dddd} tddt| d | d< t| dg ddgg t| d	ddgddgg t| d	dd
gddgg d S )Nr   rB   rC   cr   r   rB   xr   rC   r!   rQ   )r   r   rL   rM   rA   r8   r   r   r   test_R_bugs|   s
   rS   c                     s  t dddd} tddt| d | d< | d d | d<  fd	d
 t d}t |}tt| ddgddgddgddgddgg}d}dd l}| }|D ]f}t|}|D ]\}	}
|	|v rg|
|vrg nPqYt }|D ]"}tdd |D }dd |D } |D ]}|	t
||  qqmt|}|dv rttt| || nt| || |d7 }|d dkrtd| qQtd| | f  d S )Nr   r    )rB   rC   repeatr   r   rB   rD   rE   c                 3   sV    | s	t  V  d S | d } | dd  D ]}t t|V  t t|f| V  qd S )Nr   r   )tuplesorted)lobjsubsetall_subsetsr   r   r[      s   z/test_redundancy_thoroughly.<locals>.all_subsets)rB   rC   rD   rE   )rC   )rB   )rE   )rD   )rC   rE   )rB   rD   )rB   rC   rE   )rB   rC   rD   )rC   rD   rE   )rB   rD   rE   c                 S   s   g | ]	}| d r|qS rR   
startswithr'   tr   r   r   r)          z.test_redundancy_thoroughly.<locals>.<listcomp>c                 S   s   g | ]	}| d s|qS r\   r]   r_   r   r   r   r)      ra   )r   )r   r   z
Completed:zTook %0.2f seconds)r   r   rL   rM   listprinttimer/   rU   add	frozensetr#   r$   r   rA   )r9   	all_termsZall_termlist_templatesZ	redundantcountrd   startZtermlist_templateZtermlist_setZdispreferredZ	preferredZexpanded_termsZterm_templatenumericrestZsubset_restr3   r   rZ   r   test_redundancy_thoroughly   sR   	
rl   r   c               	   C   sV  g dg dd} t | }dd |d D |d< g dg dd}t | }dd |d D |d< tjtt| d | d d	dtfgd
}tjtt| d | d ddtfgd
}|tj}|tj}| ||||g}trt	
|}|| t	
|}	||	 |D ]%}
t|
ddgddggg dd}t|g dg dg dg dgsJ qd S )N)a1a2rm   rn   r   r   r   r   )rB   rR   c                 S   s   g | ]}| d qS )ascii)encoder'   r   r   r   r   r)          z#test_data_types.<locals>.<listcomp>rB   c                 S   s   g | ]}t |qS r   )six	text_typerr   r   r   r   r)      rs   rR   )rB   ZS2dtype)rB   ZU2r   )rF   rG   za[a1]:xza[a2]:xr=   rK   )r   r   r   r   )r   r   r   r   )r   r   r   r   )dictr   Zarrayrb   zipintviewZrecarrayr   pandas	DataFramer*   rA   rN   )Z
basic_dictZbasic_dict_bytesZbasic_dict_unicodeZstructured_array_bytesZstructured_array_unicodeZrecarray_bytesZrecarray_unicodedatasZdf_bytesZ
df_unicoder9   r   r   r   r   test_data_types   sP   





r   c                     s   dg di  fdd} t tdg| dd }t|g d }|jttjks*J t|g tjdd }|jttjks@J ttdr[t|g tjdd }|jttjks]J d S d S )NrR   r   c                   3   r6   r7   r   r   r8   r   r   r:      r;   z4test_build_design_matrices_dtype.<locals>.iter_makerr   rv   float128)	r>   r.   r?   rw   r   Zfloat64Zfloat32hasattrr   r:   buildermatr   r8   r    test_build_design_matrices_dtype   s   
r   c                     sj   dg di  fdd} t tdg| dd }t|g ddd }t|ts(J tjtt|g dd d S )	NrR   r   c                   3   r6   r7   r   r   r8   r   r   r:      r;   z$test_return_type.<locals>.iter_makerr   r@   return_typeZasdfsadf)r>   r.   r?   
isinstancer   r#   r$   r   r   r   r8   r   test_return_type   s   
r   c               	      s  g dg dd  fdd} t tddg| dd }t|gd	tjd
gtjg dtddd }|jdks8J t|g dgsCJ t|gd	tjd
gtjg dtddddd }|jdksbJ t|g dgsmJ ddl	m
} |g d}t|gd	tjgtjddgtdd|dd }|jdksJ tj|g dddtjgg tjtt|gd	tjd
gtjg dtdddd d S )Nr   )c1c2r   rR   rQ   c                   3   r6   r7   r   r   Zinitial_datar   r   r:     r;   z"test_NA_action.<locals>.iter_makerrR   rQ   r         $@g      4@)r   r   Nrv   )r   r   )      ?        r   drop)	NA_action)NAAction)ZNA_typesr   r   )r   r   r   r   raise)r>   r.   r?   r   nanasarrayobjectr   array_equalZpatsy.missingr   ZtestingZassert_array_equalr#   r$   r   )r:   r   r   r   r   r   r   r   test_NA_action
  sd   








r   c                     s   dt jdgg dd  fdd} ttddg| d	d	 }|jg d
ks&J t|g \}|jdks4J t |g dg dgsBJ d S )Nr         @)r   r   Zc3r   c                   3   r6   r7   r   r   r8   r   r   r:   ;  r;   z1test_NA_drop_preserves_levels.<locals>.iter_makerrR   rQ   r   )zc[c1]zc[c2]zc[c3]rR   )r   r   )r   r   r   r   )r   r   r   r   )r   r   r>   r.   r1   r?   r   r   )r:   r0   r   r   r8   r   test_NA_drop_preserves_levels7  s   r   c               
      st  t sd S tjg dg dg ddg dd  fdd} ttg g| d	\}ttd
tdg| d	d\}}ttddg| d	d\}ttdd
g| d	d\}ttt|g d  d d d d d ttt||g d  d
 d d d d G dd dtj}ttt|g|  t|g d g dd\}t	
|ddgddgddggsJ t||g dd\}}	t|tjsJ t|	tjsJ t	|dgdgdggsJ t	|	dgdgdggsJ t	|jg dsJ t	|	jg dsJ t	|jd
gsJ t	|	jdgsJ |jjd
gksJ |	jjdgks#J |jjd
gks-J |	jjdgks7J t||gg d d d dd\}}	t|tjsRJ t|	tjs[J t	|d!gd"gd#ggsjJ t	|	dgdgdggsyJ t	|jg dsJ t	|	jg dsJ t	|jd
gsJ t	|	jdgsJ |jjd
gksJ |	jjdgksJ |jjd
gksJ |	jjdgksJ t|gg d$ d ddd\}
t|
tjsJ t	|
g d%g d&g d'gsJ t	|
jg dsJ t|gg dg d(d dd\}t|tjsJ t	|d)d!gd*d"gd+d#ggs1J t	|jg d,s=J t|g dd\}t|tjsOJ t	|dgdgdggs^J |jtg dslJ d	d l}|jj }zd-|j_ tjtt|gdg didd W ||j_ n||j_ w t|gd.t	jd/gt	jg d0td1dd2dd3\}	|	jtdgsJ d S )4Nr   )r   r    r!   rm   rn   rm   )rR   yrB   )r         )indexc                   3   r6   r7   r   r   r8   r   r   r:   O  r;   z+test_return_type_pandas.<locals>.iter_makerr   r   rR   r<   rB   rR   rB   rR   r   c                   @   s   e Zd Zdd ZdS )z2test_return_type_pandas.<locals>.CheatingDataFramec                 S   s.   |dkrt j| |d d d S t j| |S )NrR   r   )r|   r}   __getitem__)selfkeyr   r   r   r   j  s   z>test_return_type_pandas.<locals>.CheatingDataFrame.__getitem__N)__name__
__module____qualname__r   r   r   r   r   CheatingDataFramei  s    r   )(   2   <   r   r   r   r   r   r   Z	dataframer   r   r    r!   )r"      	   )r   rR   r"   r   r   )r   )r   r   r   )r   r   r   )r   r   r   )r         r   r   r   )r   r   r   Fr   r   )Nrn   rm   rv   r   )r   r   )r   r|   r}   r>   r.   r#   r$   r   r?   r   rN   r   r   r   columnsr0   r1   Z
term_namesZequalsZIndexpatsy.buildZbuildr   r   r   )r:   Zint_builderZ	y_builderZ	x_builderZx_a_builderZx_y_builderr   r   Zy_dfZx_dfZx_a_dfZx_y_dfZint_dfpatsyZ
had_pandasr   r8   r   test_return_type_pandasG  s   

"


$$

r   c                     s  g dg dft g dg ddt g dg ddfdgdgdggddgddgddggfg d	g d
g dgddgddgddggfg} g dg dfg dt g dfg dt g dfg dg dfg}g dg dfg}tdg  fdd} fdd}| D ]\}}||| ||| ||| ||| q|D ]\}}||| ||| q|D ]\}}||| ||| qtttg dg ddddgdgg d S )Nr   )TFTrP   )rQ   rC   rB   )Zlevelsr   r   r   )r   r   r   )r   r   r   )r   r   r   )rB   rC   drR   c                    sT    fdd}zt g|d}t|d i t|di W t ty)   Y d S w )Nc                   3   s    d iV  diV  d S NrR   r   r   data1data2r   r   r:     s   
z=test_data_mismatch.<locals>.t_incremental.<locals>.iter_makerr   rR   )r>   r?   r   r%   r   r   r:   buildersr4   r   r   t_incremental  s   z)test_data_mismatch.<locals>.t_incrementalc                    s2    fdd}t g|d}ttt|d|i d S )Nc                   3   s    d iV  d S r   r   r   r   r   r   r:     s   z?test_data_mismatch.<locals>.t_setup_predict.<locals>.iter_makerr   rR   )r>   r#   r$   r   r?   r   r   r   r   t_setup_predict  s
   
z+test_data_mismatch.<locals>.t_setup_predictro   r   r   )r
   r.   r#   r$   r   rA   )Ztest_cases_twowayZtest_cases_onewayZsetup_predict_onlyr   r   rB   rC   r   r   r   test_data_mismatch  s@   $*







r   c            	         s  dg di  fdd} t t g| dd }ttt|g  t tg g| ddd }ttt|g  ttt||g  trat||gt \}}t	
|dgdgdggsZJ |jdksaJ tdg}t |t g| dd}t| \}}t	
|dgd	gd
ggsJ |jdksJ t |tg g| dd}t| \}}t| \}}t	
|dgd	gd
ggsJ t	
|dgdgdggsJ d S )NrR   r   c                   3   r6   r7   r   r   r8   r   r   r:      r;   z1test_data_independent_builder.<locals>.iter_makerr   r<   r   )r   r   r   r   )r>   r.   r#   r$   r   r?   r   r|   r}   r   rN   r   )	r:   Znull_builderZintercept_builderZint_mZnull_mZ
x_termlistr   Zx_mZintercept_mr   r8   r   test_data_independent_builder  sN   


 r   c                     s   g dg dd  fdd} t dg}t ddg}t||g| dd	}t| \}}t|d
|dgd t|d
gdgdggsAJ t|d|ddgd t|d
dgddgddggs\J d S )Nr   r   r   c                   3   r6   r7   r   r   r8   r   r   r:   1  r;   z4test_same_factor_in_two_matrices.<locals>.iter_makerrR   rB   r   r<   r   r=   r   r   zx:a[a1]zx:a[a2])r.   r>   r?   r5   r   rN   )r:   t1t2r   Zm1Zm2r   r8   r    test_same_factor_in_two_matrices/  s   
&r   c                     s4   dg di  fdd} t tttdg| d d S )NrR   r   c                   3   r6   r7   r   r   r8   r   r   r:   >  r;   z.test_eval_env_type_builder.<locals>.iter_makerZfoo)r#   r$   	TypeErrorr>   r.   )r:   r   r8   r   test_eval_env_type_builder<  s   r   c                  C   sn   dg di} dt g di}| |g}tr"dtg di}|| dd }|D ]}|D ]}||| q,q(d S )NrB   r   )rn   rm   rn   )rm   rn   rn   c                    s0    fdd}t tdgg|dd}t|| d S )Nc                   3   r6   r7   r   r   r   r   r   r:   J  r;   z/test_categorical.<locals>.t.<locals>.iter_makerrB   r   r<   )r>   r.   r?   r   r   r   r   r`   I  s   ztest_categorical.<locals>.t)r
   r   r|   ZCategoricalr*   )Zdata_stringsZ
data_categr~   Zdata_pandasr`   r   r   r   r   r   test_categoricalB  s   
r   c               	   C   s  ddl m} m} g d}tdt|iddggg dd}t|g dg d	g dg d
gs0J || fD ]'}tdt||iddggg dd}t|g dg dg dg dgs\J q5tdt||ddiddggg dd}t|g dg dg dg dgsJ tdt|idg dggg dd}t|g dg dg dg dgsJ || fD ](}tdt||idg dggg dd}t|g dg dg dg dgsJ qtdt||ddidg dggg dd}t|g dg dg dg dgsJ tdt|ddgddgddggiddggddgd}t|ddgddgddgddggs1J tdt|| ddgddgddggddgiddggddgd}t|ddgddgddgddggseJ d S ) Nr   )ContrastMatrixSum)rm   a3rm   rn   rB   r   )rF   rG   za[a3]r=   )r   r   r   )r   r   r   )r   r   r   )a[mean]a[S.a1]a[S.a2])r   r   r   )r   r   r   )r   r   r   )Zomit)r   r   a[S.a3])rH   rI   za[T.a3])rH   r   r   )rH   r   r   r"   r   r      r   r   z
a[custom0]z
a[custom1]z[foo]z[bar]za[foo]za[bar])Zpatsy.contrastsr   r   rA   r
   r   rN   )r   r   valuesr   r   r   r   r   test_contrastT  s   


r   c                     s  ddgddgddggddgd	t d
dd} fdd}t| g|dd  t gd  fdd}|dg d	td  |g d	g d	td  tjsa|td
tdtdgg d	td  || g d	td  || d d| d gg d	td  |dd
dgddg |d
dgd
dgddg tjs|td
tdgd
dgddg || d | d gd
dgddg || d dgd
dgddg |dd
dgddg |dd
gd
dgddg |tdtd
gd
dgddg || d | d gd
dgddg || d d
gd
dgddg |ddgddg |dgdgddg |tdgdgddg || d gdgddg t	t
 jd t	t jd t	t jdg t	t jtdgg tg d}|ddg}|jddgksqJ |jd u syJ d S )Nr   r   g@g	@gffffff@g@r    r!   )rR   r   zrR   r   r   c                   3   r6   r7   r   r   )all_datar   r   r:     r;   z*test_DesignInfo_subset.<locals>.iter_makerr   c                    sz     | }i }|D ]}| ||< q	t|g|d }d d |f }t| tjs3t| t|jks3J t||s;J d S )Nr   )	rY   r?   r   rt   Zstring_typesrM   r,   r   r   )Zwhich_termsZ	variablesr   Zsub_design_infoZsub_datavariableZ
sub_matrixZsub_full_matrixZall_builderr   Zfull_matrixr   r   r`     s   
z!test_DesignInfo_subset.<locals>.tz~ 0 + x + y + zz~ 0 + x + zr   z~ 0 + z + xz~ 0 + yza ~ az~ asdfZasdfrP   rQ   rB   )r.   r>   r?   slicert   ZPY3Zunicoderu   r#   r$   r   rY   KeyErrorr   r   r1   r,   )rg   r:   r`   Zmin_diZmin_di_subsetr   r   r   test_DesignInfo_subset  sX   "  $ r   r7   )/Z
__future__r   rt   Znumpyr   r#   r   r   Z
patsy.utilr   r   r   Z
patsy.descr   r   r   Zpatsy.categoricalr
   Zpatsy.user_utilr   r   Zpatsy.design_infor   r   r|   r   r&   r.   r5   rA   rO   rS   rl   Zslowr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   <module>   sF   	

.8#-~91N