U md @sdZddlmZddlZddZdddZd d Zd d Ze d kr~ddl m Z ddZ e eddgddge dde eddgddge dde ddddgddddgddddgddddgddddggZeedddeddde Ze eede dgdgdgdggZeedddeddde Ze eede ddddggZeedddeddde Ze eeddS)z?Quantizing a continuous distribution in 2d Author: josef-pktd )lmapNcCsD||}||d|d}||d|d}||}||||S)ahelper function for probability of a rectangle in a bivariate distribution Parameters ---------- lower : array_like tuple of lower integration bounds upper : array_like tuple of upper integration bounds cdf : callable cdf(x,y), cumulative distribution function of bivariate distribution how does this generalize to more than 2 variates ? r)loweruppercdfZprobuuZprobulZprobluZprobllrrc/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/statsmodels/sandbox/distributions/quantize.pyprob_bv_rectangles r c Cst|tjsvttj|}t|}g}tttj|t|krt |D],}dg|}t d||<| |||qFn|j d}|}t t|||}|}t |D]}tj||d}q|S)zhelper function for probability of a rectangle grid in a multivariate distribution how does this generalize to more than 2 variates ? bins : tuple tuple of bin edges, currently it is assumed that they broadcast correctly Nr)axis) isinstancenpZndarrayrasarraylenallndimonesrangesliceappendshapeprintcopydiff) Zbinsrr Zn_dimZbins_dsl cdf_valuesprobsrrr prob_mv_grids"        rc st|}t|}t|d}t|d}tjt||f}||dddf|fdd}td|dD]L}td|dD]8}||f} |d|df} t| | |||d|df<qqpt|rt |S)zquantize a continuous distribution given by a cdf Parameters ---------- binsx : array_like, 1d binedges rNcs ||fSNrxyrrrMz#prob_quantize_cdf.. r rrnanrrr isnananyAssertionError) binsxbinsyrnxnyrZcdf_funcxindyindrrrr#rprob_quantize_cdf>s      r1c Cst|}t|}t|d}t|d}tjt||f}td|dD]\}td|dD]H}||||f}||d||df} t| ||||d|df<q`qNt|rt |S)zquantize a continuous distribution given by a cdf old version without precomputing cdf values Parameters ---------- binsx : array_like, 1d binedges rr&) r+r,rr-r.rr/r0rrrrrprob_quantize_cdf_oldXs     r2__main__)assert_almost_equalcCs||Srrr rrrr$wr%r$rg?g?g?)r )__doc__Zstatsmodels.compat.pythonrnumpyr r rr1r2__name__Z numpy.testingr4Zunif_2darrayZarr1bZlinspaceZarr1aZarr2bZarr2aZarr3bZarr3arrrrs:  !