B |U#@s~dddgZddlZddlZddlmZddlmZddlm Z m Z ddZ d d Z d dZ d d ZGdddeZddZdS)balanced demo_data LookupFactorN) PatsyError)C) no_picklingassert_no_picklingc s|dd}g}t|}x6|D].|}|fddtd|dDqWttj|}i}x&t||D]\}t|||<qnW|S)aTbalanced(factor_name=num_levels, [factor_name=num_levels, ..., repeat=1]) Create simple balanced factorial designs for testing. Given some factor names and the number of desired levels for each, generates a balanced factorial design in the form of a data dictionary. For example: .. ipython:: In [1]: balanced(a=2, b=3) Out[1]: {'a': ['a1', 'a1', 'a1', 'a2', 'a2', 'a2'], 'b': ['b1', 'b2', 'b3', 'b1', 'b2', 'b3']} By default it produces exactly one instance of each combination of levels, but if you want multiple replicates this can be accomplished via the `repeat` argument: .. ipython:: In [2]: balanced(a=2, b=2, repeat=2) Out[2]: {'a': ['a1', 'a1', 'a2', 'a2', 'a1', 'a1', 'a2', 'a2'], 'b': ['b1', 'b2', 'b1', 'b2', 'b1', 'b2', 'b1', 'b2']} repeatcsg|]}d|fqS)z%s%s).0i)namer .lib/python3.7/site-packages/patsy/user_util.py 1szbalanced..)popsortedappendrangezip itertoolsproductlist)kwargsr levelsnamesZ level_countvaluesdatavaluer )rrrs  &c Cstddd}|dddddddgks(t|ddd d dd d gksDttdddd }|dddddddddddddg kszt|ddd d dd d dd d dd d g kstdS) N)abr!a1a2r"b1b2Zb3)r!r"r )rAssertionError)rr r r test_balanced9s r(c Os|dd}|dd}|r*td|ft}i}xH|D]@}|ddkrT|||<q:|ddkrl||q:td |fq:Wtjt|t d }t t |d |}||} t fd |i|} tj d} x t|D]}| j| d | |<qW| S)a@demo_data(*names, nlevels=2, min_rows=5) Create simple categorical/numerical demo data. Pass in a set of variable names, and this function will return a simple data set using those variable names. Names whose first letter falls in the range "a" through "m" will be made categorical (with `nlevels` levels). Those that start with a "p" through "z" are numerical. We attempt to produce a balanced design on the categorical variables, repeating as necessary to generate at least `min_rows` data points. Categorical variables are returned as a list of strings. Numerical data is generated by sampling from a normal distribution. A fixed random seed is used, so that identical calls to demo_data() will produce identical results. Numerical data is returned in a numpy array. Example: .. ipython: In [1]: patsy.demo_data("a", "b", "x", "y") Out[1]: {'a': ['a1', 'a1', 'a2', 'a2', 'a1', 'a1', 'a2', 'a2'], 'b': ['b1', 'b2', 'b1', 'b2', 'b1', 'b2', 'b1', 'b2'], 'x': array([ 1.76405235, 0.40015721, 0.97873798, 2.2408932 , 1.86755799, -0.97727788, 0.95008842, -0.15135721]), 'y': array([-0.10321885, 0.4105985 , 0.14404357, 1.45427351, 0.76103773, 0.12167502, 0.44386323, 0.33367433])} nlevelsrmin_rowszunexpected keyword arguments %rrZabcdefghijklmnZ pqrstuvwxyzz bad name %r)dtypeg?r )size)r TypeErrorsetaddrnpZprodrrintZceilrZrandomZ RandomStaterZnormal) rrr)r*Z numericalZ categoricalrZbalanced_design_sizer Znum_rowsrrr r rrCs(!        c Csptddd}t|dddgks&t|dddddddddgksFt|dddddddddgksft|djttks~t|djdksttdd }t|dd gkstt|dt|d krd ksntttdd d dd kstttdddd d dd kstttdddd ddddks>tddl m }|t tddd|t tdddddS)Nr!r"xr#r$r%r&)yr+ )r* r )r*r)r) assert_raisesZ__123{)Zasdfasdf) rrkeysr'r,r1floatshapelen nose.toolsr:rr.)Zd1Zd2r:r r rtest_demo_datazs    *"$ rAc@sfeZdZdZdddZddZdd Zd d Zd d ZddZ ddZ ddZ ddZ ddZ eZdS)ra8A simple factor class that simply looks up a named entry in the given data. Useful for programatically constructing formulas, and as a simple example of the factor protocol. For details see :ref:`expert-model-specification`. Example:: dmatrix(ModelDesc([], [Term([LookupFactor("x")])]), {"x": [1, 2, 3]}) :arg varname: The name of this variable; used as a lookup key in the passed in data dictionary/DataFrame/whatever. :arg force_categorical: If True, then treat this factor as categorical. (Equivalent to using :func:`C` in a regular formula, but of course you can't do that with a :class:`LookupFactor`. :arg contrast: If given, the contrast to use; see :func:`C`. (Requires ``force_categorical=True``.) :arg levels: If given, the categorical levels; see :func:`C`. (Requires ``force_categorical=True``.) :arg origin: Either ``None``, or the :class:`Origin` of this factor for use in error reporting. .. versionadded:: 0.2.0 The ``force_categorical`` and related arguments. FNcCsH||_||_||_||_||_|jsD|dk r4td|dk rDtddS)Nz)contrast= requires force_categorical=Truez'levels= requires force_categorical=True)_varname_force_categorical _contrast_levelsorigin ValueError)selfZvarnameforce_categoricalcontrastrrFr r r__init__szLookupFactor.__init__cCs|jS)N)rB)rHr r rrszLookupFactor.namecCsd|jj|jfS)Nz%s(%r)) __class____name__rB)rHr r r__repr__szLookupFactor.__repr__cCs:t|to8|j|jko8|j|jko8|j|jko8|j|jkS)N) isinstancerrBrCrDrE)rHotherr r r__eq__s     zLookupFactor.__eq__cCs ||k S)Nr )rHrPr r r__ne__szLookupFactor.__ne__cCstt|j|j|j|jfS)N)hashrrBrCrDrE)rHr r r__hash__szLookupFactor.__hash__cCsdS)Nrr )rHstateZeval_envr r rmemorize_passes_neededsz#LookupFactor.memorize_passes_neededcCs dstdS)NF)r')rHrU which_passrr r rmemorize_chunkszLookupFactor.memorize_chunkcCs dstdS)NF)r')rHrUrWr r rmemorize_finishszLookupFactor.memorize_finishcCs&||j}|jr"t||j|jd}|S)N)rJr)rBrCrrDrE)rHZmemorize_staterrr r revals zLookupFactor.eval)FNNN)rM __module__ __qualname____doc__rKrrNrQrRrTrVrXrYrZr __getstate__r r r rrs cCshtd}|dkst|tdks(t|tdks8tt|ttdksPtt|ttdksht|iddidkst|iddidkstt|dkst|jdksttddd}|jdksttdd d d d }|iddddgi}|jdddgks t|jd kst|j d ks,td dl m }|t tdd d|t tdd dt tddS)Nr!r"r rzLookupFactor('a')Zasdf)rFcTZCONTRAST)r r)rIrJrr)r:Znc)rJ)r)rrr'rSrZreprrFrrJrr@r:rGr)Zl_aZ l_with_originZl_cZboxr:r r rtest_LookupFactors,   ra)__all__rZnumpyr1ZpatsyrZpatsy.categoricalrZ patsy.utilrrrr(rrAobjectrrar r r r s   ( 7N