U mda@sndZddlmZddlZddlmZddZddZd d Z d d Z dd dZ ddZ dddZ dddZdS)a Working with categorical data ============================= use of dummy variables, group statistics, within and between statistics examples for efficient matrix algebra dummy versions require that the number of unique groups or categories is not too large group statistics with scipy.ndimage can handle large number of observations and groups scipy.ndimage stats is missing count new: np.bincount can also be used for calculating values per label )lrangeN)ndimagecCs2tt|d}ttj|||d}||S)Nlabelsindex)nparangemaxarrayrmean)yx labelsunique labelmeansrc/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/statsmodels/sandbox/regression/try_catdata.pylabelmeanfiltersrc Cstt|d}g}g}|jD]2}ttj|||d}|||||q"ttj||d|ddd||d}|t|t|jfS)Nrrr) rr r Tr rr appendZ histogram)r rrZ labmeansdataZlabmeansZxxrZ labelcountrrrlabelmeanfilter_nd"s  rc CsFtj|ddd\}}ttj||tt|dd}||}|S)NFTZ return_indexreturn_inverserr)runiquer rr r r )ysrunilunilinvrZarr3rrrlabelmeanfilter_str7s&rc Csvt|}tj|dd\}}t|}tj||dd|}||}tj|||ddd|}||} ||||| fS)z:uses np.bincount, assumes factors/labels are integers r)r)weightsg?)lenrrZbincount) ZfactorsvaluesnZixZrindZgcountZgmeanZmeanarrZ withinvarZ withinvararrrrr groupstatsbin?s r$cCs||dkr|}nFt|}|jdkrP|jdkrPtdd|Dddtjf}n|}tj|ddd\}}|tt||fS) zconvert labels based on multiple variables or string labels to unique index labels 0,1,2,...,nk-1 where nk is the number of distinct labels Nrr cSs g|]}d|ddqS)z@%s@Nr )tostring).0iirrr Usz!convertlabels..FTr)rr sizendimnewaxisrr r!)rindicesZylabelidxrrrrr convertlabelsLs $r.cCs4ttj|||d}ttj|||d}||fS)z)use ndimage to get fast mean and variancer)rr rr var)r rrrZ labelvarsrrrgroupsstats_1d`sr0cCs|s|jdkr8|jddkr8t|t|jd\}}}n$|}|}t||d}|jdkrx|ddtj f}||k t }|S)Nr r) r*shaper.rcopyminrr r r+Zastypeint)r nonseqZycatZuniquesZ unitranslZymindummyrrr cat2dummyfs r7c Cs|jdkr|ddtjf}t||d}|jdtd}t|j||}t||j}||}t||j||}||||fS)Nr)r5r)Zdtype)r*rr+r7sumfloatdotr) r rr5r6ZcountgrZmeangrZmeandataZ xdevmeangrZvargrrrrgroupsstats_dummyss  r;)N)r)r)__doc__Zstatsmodels.compat.pythonrnumpyrZscipyrrrrr$r.r0r7r;rrrrs