a tDf@sUddlmZddlZddlZddlZddlmZmZddlmZddl m Z ddl m Z ddl ZddlmZddlmZmZmZmZmZmZdd lmZdd lmZdd lmZdd lmZdd l m!Z!m"Z"m#Z#GdddZ$dTddZ%GdddZ&ddZ'dUddZ(ddZ)ddZ*ddZ+ddZ,d d!Z-d"d#Z.dd$d%d&d'd(Z/iZ0d)e1d*<e Z2d+d,Z3e3d-Z4e3d.Z5e3d/Z6e3d0Z7e3d1Z8e3d2Z9e3d3Z:e3d4Z;e3d5Ze3d8Z?e3d9Z@e3d:ZAe3d;ZBe3d<ZCe3d=ZDe3d>ZEe3d?ZFe3d@ZGe3dAZHe3dBZIe3dCZJe3dDZKe3dEZLe3dEZMe3dFZNe3dGZOe3dHZPe3dIZQe3dJZRe3dKZSe3dLZTe3dMZUe3dNZVe3dOZWe3dPZXe3dQZYe3dRZZe3dSZ[dS)V) annotationsN)chainproduct)Integral)getitem)Lock)array_creation_dispatch)Arrayasarraybroadcast_shapes broadcast_tonormalize_chunksslices_from_chunks)arange) asarray_safe)tokenize)HighLevelGraph) derived_fromrandom_state_datatypenamec@s(eZdZdZddZddZeddZedd Ze e j j d d d`ddZ e e j j d d daddZe e j j d d dbddZe e j j d d dcddZe e j j d d ddddZe e j j d d deddZe e j j d d dfddZe e j j d d dgdd Ze e j j d d dhd"d#Ze e j j d d did$d%Ze e j j d d d d e jd&d fd'd(Ze e j j d d djd)d*Ze e j j d d dkd+d,Ze e j j d d dld-d.Ze e j j d d dmd/d0Ze e j j d d dnd1d2Ze e j j d d dod4d5Ze e j j d d dpd6d7Ze e j j d d dqd8d9Z e e j j d d drd:d;Z!e e j j d d dsdd?Z#e e j j d d d@dAZ$e e j j d d dudBdCZ%e e j j d d dvdDdEZ&e e j j d d d e j'd d fdFdGZ e e j j d d dwdHdIZ(e e j j d d dxdJdKZ)e e j j d d dydLdMZ*e e j j d d dzdNdOZ+e e j j d d d{dPdQZ,e e j j d d d|dRdSZ-e e j j d d d}dTdUZ.e e j j d d d~dVdWZ/e e j j d d ddXdYZ0e e j j d d ddZd[Z1e e j j d d dd\d]Z2e e j j d d dd^d_Z3d S) Generatora Container for the BitGenerators. ``Generator`` exposes a number of methods for generating random numbers drawn from a variety of probability distributions and serves as a replacement for ``RandomState``. The main difference between the two is that ``Generator`` relies on an additional ``BitGenerator`` to manage state and generate the random bits, which are then transformed into random values from useful distributions. The default ``BitGenerator`` used by ``Generator`` is ``PCG64``. The ``BitGenerator`` can be changed by passing an instantiated ``BitGenerator`` to ``Generator``. The function :func:`dask.array.random.default_rng` is the recommended way to instantiate a ``Generator``. .. warning:: No Compatibility Guarantee. ``Generator`` does not provide a version compatibility guarantee. In particular, as better algorithms evolve the bit stream may change. Parameters ---------- bit_generator : BitGenerator BitGenerator to use as the core generator. Notes ----- In addition to the distribution-specific arguments, each ``Generator`` method takes a keyword argument `size` that defaults to ``None``. If `size` is ``None``, then a single value is generated and returned. If `size` is an integer, then a 1-D array filled with generated values is returned. If `size` is a tuple, then an array with that shape is filled and returned. The Python stdlib module `random` contains pseudo-random number generator with a number of methods that are similar to the ones available in ``Generator``. It uses Mersenne Twister, and this bit generator can be accessed using ``MT19937``. ``Generator``, besides being Dask-aware, has the advantage that it provides a much larger number of probability distributions to choose from. All ``Generator`` methods are identical to ``np.random.Generator`` except that they also take a `chunks=` keyword argument. ``Generator`` does not guarantee parity in the generated numbers with any third party library. In particular, numbers generated by `Dask` and `NumPy` will differ even if they use the same seed. Examples -------- >>> from numpy.random import PCG64 >>> from dask.array.random import Generator >>> rng = Generator(PCG64()) >>> rng.standard_normal().compute() # doctest: +SKIP array(0.44595957) # random See Also -------- default_rng : Recommended constructor for `Generator`. np.random.Generator cCs ||_dSN)_bit_generator)self bit_generatorr^/nfs/NAS7/SABIOD/METHODE/ermites/ermites_venv/lib/python3.9/site-packages/dask/array/random.py__init__^szGenerator.__init__cCs"|jj}|d|jjjd7}|S)N()) __class____name__r)r_strrrr__str__aszGenerator.__str__cCst|jddSN.r)rrsplitrrrr _backend_namefszGenerator._backend_namecCs t|jSr) importlib import_moduler(r'rrr_backendlszGenerator._backendZ skipblocksNautocKst|d||f||d|SNbetasizechunks _wrap_funcrabr2r3kwargsrrrr0rszGenerator.betacKst|d||f||d|SNbinomialr1r4rnpr2r3r9rrrr;vszGenerator.binomialcKst|d|f||d|SN chisquarer1r4rdfr2r3r9rrrr@zszGenerator.chisquareTrc st|||\}}}} tt|} t|jt| } dt| ||} t| ggdd|DR} fddt| | | D}tj | || d}t || ||dS)Nda.random.choice-%scss|]}tt|VqdSrrangelen.0Zbdrrr z#Generator.choice..c s(i|] \}}}|t||fqSr) _choice_rng)rHkbitgenr2r7axisr>replaceshufflerr sz$Generator.choice.. dependenciesmeta) _choice_validate_paramslistr_spawn_bitgensrrFrziprfrom_collectionsr )rr7r2rPr>rOrQr3rVrTsizesbitgensnamekeysdskgraphrrNrchoice~s*  zGenerator.choice?cKst|d|f||d|SN exponentialr1r4rscaler2r3r9rrrreszGenerator.exponentialcKst|d||f||d|SNfr1r4rdfnumdfdenr2r3r9rrrrisz Generator.fcKst|d||f||d|SNgammar1r4rshapergr2r3r9rrrrnszGenerator.gammacKst|d|f||d|SN geometricr1r4rr>r2r3r9rrrrrszGenerator.geometriccKst|d||f||d|SNgumbelr1r4rlocrgr2r3r9rrrrvszGenerator.gumbelcKst|d|||f||d|SNhypergeometricr1r4rZngoodZnbadnsampler2r3r9rrrrzszGenerator.hypergeometricFcKs t|d|f|||||d|S)Nintegers)highr2dtypeendpointr3r4)rlowr~r2rrr3r9rrrr}s  zGenerator.integerscKst|d||f||d|SNlaplacer1r4rwrrrrszGenerator.laplacecKst|d||f||d|SNlogisticr1r4rwrrrrszGenerator.logisticcKst|d||f||d|SN lognormalr1r4rmeansigmar2r3r9rrrrszGenerator.lognormalcKst|d|f||d|SN logseriesr1r4rsrrrrszGenerator.logseriescKs&t|d||f||t|ffd|SN multinomialr2r3 extra_chunksr5rFrr=Zpvalsr2r3r9rrrrs zGenerator.multinomial marginalscKst|d||f|||d|S)Nmultivariate_hypergeometric)r2methodr3r4)rcolorsr|r2rr3r9rrrrsz%Generator.multivariate_hypergeometriccKst|d||f||d|SNnegative_binomialr1r4r<rrrrszGenerator.negative_binomialcKst|d||f||d|SNnoncentral_chisquarer1r4rrBnoncr2r3r9rrrrszGenerator.noncentral_chisquarecKst|d|||f||d|SN noncentral_fr1r4rrkrlrr2r3r9rrrrs zGenerator.noncentral_fcKst|d||f||d|SNnormalr1r4rwrrrr%szGenerator.normalcKst|d|f||d|SNparetor1r4rr7r2r3r9rrrr+szGenerator.paretocCs\ddlm}|jdkrtdt|tjr6t|dd}|jt |}t |j ||||S)Nr shuffle_slicecupyz`Generator.permutation` not supported for cupy-backed Generator objects. Use the 'numpy' array backend to call `dask.array.random.default_rng`, or pass in `numpy.random.PCG64()`.r.r3) dask.array.slicingrr(NotImplementedError isinstancenumbersNumberrr+rF_shufflerrxrindexrrr permutation/s     zGenerator.permutationcKst|d|f||d|SNpoissonr1r4rZlamr2r3r9rrrrBszGenerator.poissoncKst|d|f||d|SNpowerr1r4rrrrrFszGenerator.powercKst|df||||d|S)Nrandom)r2routr3r4)rr2rrr3r9rrrrJszGenerator.randomcKst|d|f||d|SNrayleighr1r4rfrrrrPszGenerator.rayleighcKst|df||d|SNstandard_cauchyr1r4rr2r3r9rrrrTszGenerator.standard_cauchycKst|df||d|SNstandard_exponentialr1r4rrrrrXszGenerator.standard_exponentialcKst|d|f||d|SNstandard_gammar1r4rrpr2r3r9rrrr^szGenerator.standard_gammacKst|df||d|SNstandard_normalr1r4rrrrrdszGenerator.standard_normalcKst|d|f||d|SN standard_tr1r4rArrrrhszGenerator.standard_tcKst|d|||f||d|SN triangularr1r4rleftmoderightr2r3r9rrrrls zGenerator.triangularcKst|d||f||d|SNuniformr1r4rrr~r2r3r9rrrrrszGenerator.uniformcKst|d||f||d|SNvonmisesr1r4rmukappar2r3r9rrrrxszGenerator.vonmisescKst|d||f||d|SNwaldr1r4rrrgr2r3r9rrrr~szGenerator.waldcKst|d|f||d|SNweibullr1r4rrrrrszGenerator.weibullcKst|d|f||d|SNzipfr1r4rrrrrszGenerator.zipf)Nr.)Nr.)Nr.)NTNrTr.)rcNr.)Nr.)rcNr.)Nr.)rtrcNr.)Nr.)rtrcNr.)rtrcNr.)rtrcNr.)Nr.)Nr.)Nrr.)Nr.)Nr.)Nr.)rtrcNr.)Nr.)rcNr.)Nr.)rcNr.)Nr.)Nr.)Nr.)Nr.)Nr.)Nr.)rtrcNr.)Nr.)Nr.)Nr.)Nr.)4r! __module__ __qualname____doc__rr#propertyr(r+rnprrr0r;r@rbrerirnrrrvrzint64r}rrrrrrrrrrrrrrfloat64rrrrrrrrrrrrrrrrrs@      $                                rcCsBt|drt|St|tr |St|dr4t|jStt|S)af Construct a new Generator with the default BitGenerator (PCG64). Parameters ---------- seed : {None, int, array_like[ints], SeedSequence, BitGenerator, Generator}, optional A seed to initialize the `BitGenerator`. If None, then fresh, unpredictable entropy will be pulled from the OS. If an ``int`` or ``array_like[ints]`` is passed, then it will be passed to `SeedSequence` to derive the initial `BitGenerator` state. One may also pass in a `SeedSequence` instance. Additionally, when passed a `BitGenerator`, it will be wrapped by `Generator`. If passed a `Generator`, it will be returned unaltered. Returns ------- Generator The initialized generator object. Notes ----- If ``seed`` is not a `BitGenerator` or a `Generator`, a new `BitGenerator` is instantiated. This function does not manage a default global instance. Examples -------- ``default_rng`` is the recommended constructor for the random number class ``Generator``. Here are several ways we can construct a random number generator using ``default_rng`` and the ``Generator`` class. Here we use ``default_rng`` to generate a random float: >>> import dask.array as da >>> rng = da.random.default_rng(12345) >>> print(rng) Generator(PCG64) >>> rfloat = rng.random().compute() >>> rfloat array(0.86999885) >>> type(rfloat) Here we use ``default_rng`` to generate 3 random integers between 0 (inclusive) and 10 (exclusive): >>> import dask.array as da >>> rng = da.random.default_rng(12345) >>> rints = rng.integers(low=0, high=10, size=3).compute() >>> rints array([2, 8, 7]) >>> type(rints[0]) Here we specify a seed so that we have reproducible results: >>> import dask.array as da >>> rng = da.random.default_rng(seed=42) >>> print(rng) Generator(PCG64) >>> arr1 = rng.random((3, 3)).compute() >>> arr1 array([[0.91674416, 0.91098667, 0.8765925 ], [0.30931841, 0.95465607, 0.17509458], [0.99662814, 0.75203348, 0.15038118]]) If we exit and restart our Python interpreter, we'll see that we generate the same random numbers again: >>> import dask.array as da >>> rng = da.random.default_rng(seed=42) >>> arr2 = rng.random((3, 3)).compute() >>> arr2 array([[0.91674416, 0.91098667, 0.8765925 ], [0.30931841, 0.95465607, 0.17509458], [0.99662814, 0.75203348, 0.15038118]]) See Also -------- np.random.default_rng Zcapsuler)hasattrrrrrZdefault_bit_generator)seedrrr default_rngsR    rc@sReZdZdZd^ddZeddZd_ddZee j j d d d`d d Z ee j j d d daddZ ee j j d d dbddZee*ee j j d d dcddZWdn1s0Yee j j d d ddddZee j j d d deddZee j j d d dfddZee j j d d dgddZee j j d d dhdd Zee j j d d did!d"Zee j j d d djd#d$Zee j j d d dkd%d&Zee j j d d dld'd(Zee j j d d dmd)d*Zee j j d d dnd+d,Zee j j d d dod-d.Zee j j d d dpd/d0Zee j j d d dqd1d2Z ee j j d d drd3d4Z!ee j j d d dsd5d6Z"ee j j d d d7d8Z#ee j j d d dtd9d:Z$ee j j d d dud;d<Z%ee j j d d dvd>d?Z&ee j j d d dwd@dAZ'ee j j d d dxdBdCZ(e(Z ee j j d d dydDdEZ)ee j j d d dzdFdGZ*ee j j d d d{dHdIZ+ee j j d d d|dJdKZ,ee j j d d d}dLdMZ-ee j j d d d~dNdOZ.ee j j d d ddPdQZ/ee j j d d ddRdSZ0ee j j d d ddTdUZ1ee j j d d ddVdWZ2ee j j d d ddXdYZ3ee j j d d ddZd[Z4ee j j d d dd\d]Z5dS) RandomStatea Mersenne Twister pseudo-random number generator This object contains state to deterministically generate pseudo-random numbers from a variety of probability distributions. It is identical to ``np.random.RandomState`` except that all functions also take a ``chunks=`` keyword argument. Parameters ---------- seed: Number Object to pass to RandomState to serve as deterministic seed RandomState: Callable[seed] -> RandomState A callable that, when provided with a ``seed`` keyword provides an object that operates identically to ``np.random.RandomState`` (the default). This might also be a function that returns a ``mkl_random``, or ``cupy.random.RandomState`` object. Examples -------- >>> import dask.array as da >>> state = da.random.RandomState(1234) # a seed >>> x = state.normal(10, 0.1, size=3, chunks=(2,)) >>> x.compute() array([10.01867852, 10.04812289, 9.89649746]) See Also -------- np.random.RandomState NcCs&tj||_|durtjn||_dSr)rrr _numpy_stater _RandomState)rrrrrrr szRandomState.__init__cCst|jdd}t|Sr$)rrr&r)r*)rr(rrrr+szRandomState._backendcCs|j|dSr)rr)rrrrrrszRandomState.seedr,r-r.cKst|d||f||d|Sr/r4r6rrrr0szRandomState.betacKst|d||f||d|Sr:r4r<rrrr;szRandomState.binomialcKst|d|f||d|Sr?r4rArrrr@"szRandomState.chisquareTcst||d|\}}}}}tt|} tt| |j} dt| ||} t| ggdd|DR} fddt| | | D} tj | | |d}t || ||dS) NrrCcss|]}tt|VqdSrrDrGrrrrI;rJz%RandomState.choice..c s$i|]\}}}|t||fqSr) _choice_rs)rHrLstater2r7r>rPrrrR<sz&RandomState.choice..rSrU) rWrXrrrFrrrZrr[r )rr7r2rPr>r3rOrVrTr\ state_datar^r_r`rarrrrb(s.    zRandomState.choicerccKst|d|f||d|Srdr4rfrrrreFszRandomState.exponentialcKst|d||f||d|Srhr4rjrrrriLsz RandomState.fcKst|d||f||d|Srmr4rorrrrnPszRandomState.gammacKst|d|f||d|Srqr4rsrrrrrVszRandomState.geometricrtcKst|d||f||d|Srur4rwrrrrvZszRandomState.gumbelcKst|d|||f||d|Sryr4r{rrrrz`szRandomState.hypergeometriccKst|d||f||d|Srr4rwrrrrmszRandomState.laplacecKst|d||f||d|Srr4rwrrrrsszRandomState.logisticcKst|d||f||d|Srr4rrrrryszRandomState.lognormalcKst|d|f||d|Srr4rsrrrrszRandomState.logseriescKs&t|d||f||t|ffd|Srrrrrrrs zRandomState.multinomialcKst|d||f||d|Srr4r<rrrrszRandomState.negative_binomialcKst|d||f||d|Srr4rrrrrsz RandomState.noncentral_chisquarecKst|d|||f||d|Srr4rrrrrs zRandomState.noncentral_fcKst|d||f||d|Srr4rwrrrrszRandomState.normalcKst|d|f||d|Srr4rrrrrszRandomState.paretocCsHddlm}t|tjr$t|dd}tt|}|j ||||S)Nrrr.r) rrrrrrrrFrrQrrrrrs     zRandomState.permutationcKst|d|f||d|Srr4rrrrrszRandomState.poissoncKst|d|f||d|Srr4rrrrrszRandomState.powerlcKst|d||f|||d|S)Nrandint)r2r3rr4)rrr~r2r3rr9rrrrszRandomState.randintcKst|d||f||d|S)Nrandom_integersr1r4rrrrrszRandomState.random_integerscKst|df||d|S)N random_sampler1r4rrrrrszRandomState.random_samplecKst|d|f||d|Srr4rfrrrrszRandomState.rayleighcKst|df||d|Srr4rrrrrszRandomState.standard_cauchycKst|df||d|Srr4rrrrrsz RandomState.standard_exponentialcKst|d|f||d|Srr4rrrrrszRandomState.standard_gammacKst|df||d|Srr4rrrrrszRandomState.standard_normalcKst|d|f||d|Srr4rArrrrszRandomState.standard_tcKst|df||d|S)Ntomaxintr1r4rrrrrszRandomState.tomaxintcKst|d|||f||d|Srr4rrrrrs zRandomState.triangularcKst|d||f||d|Srr4rrrrrszRandomState.uniformcKst|d||f||d|Srr4rrrrrszRandomState.vonmisescKst|d||f||d|Srr4rrrrrszRandomState.waldcKst|d|f||d|Srr4rrrrrszRandomState.weibullcKst|d|f||d|Srr4rrrrr szRandomState.zipf)NN)N)Nr.)Nr.)Nr.)NTNr.)rcNr.)Nr.)rcNr.)Nr.)rtrcNr.)Nr.)rtrcNr.)rtrcNr.)rtrcNr.)Nr.)Nr.)Nr.)Nr.)Nr.)rtrcNr.)Nr.)rcNr.)Nr.)NNr.r)NNr.)Nr.)rcNr.)Nr.)Nr.)Nr.)Nr.)Nr.)Nr.)Nr.)rtrcNr.)Nr.)Nr.)Nr.)Nr.)6r!rrrrrr+rrrrrr0r;r@ contextlibsuppressAttributeErrorrbrerirnrrrvrzrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrs       *                                  rcCs(t|dd}t|}|j|Sr$)rr&r)r*rr)rM backend_nameZ backend_librrr_rng_from_bitgens rcCs,|j}t|}||_t|}|j||dS)N)rO)rtyperrQ)rrrOrrrrrrs  rcs"j|}fdd|D}|S)Ncsg|]}t|qSr)r)rHrrMrr "rJz"_spawn_bitgens..) _seed_seqspawn)rMZ n_bitgensZseedsr]rrrrY s rYcCs:t|tjjr||}t|}t||}||d|i|S)z$Apply random module method with seedr2)rrrZ SeedSequencergetattr)rngfuncnamerMr2argsr9funcrrr_apply_random_func&s  rcCs2|durtj}||}t||}||d|i|S)z"Apply RandomState method with seedNr2)rrr)rrrr2rr9rrrrr _apply_random/s  rcCst|}|j||||||dS)N)r2rPr>rOrQ)rrb)rr7r2rPr>rOrQrrrrrK8srKcCst|}|j||||dS)N)r2rPr>)rrrb)rr7r2rPr>rrrrr=s rc Csg}t|trt|trD|jdkr*td|jjjdddd}nFt|t r|durd|j dgn|} |jj jdd| d}nt d|} |dkrt dnJt |}| |j}|j}|jdkrt d t|} |||d}|durt|ts6t||d }tj|dd dd s,t d t |}n | |j}|jdkrVt dt|| krlt d|||d}|durd}nt|ttfs|f}|dkrt dt||tjd}|st|ddkrd} t| ||||||||fS)Nrz3`choice` not supported for cupy-backed `Generator`.r,r)r2r>zUnknown generator classrza must be greater than 0za must be one dimensional)likegHz>)ZrtolZatolzprobabilities do not sum to 1zp must be one dimensionalza and p must have the same sizez)axis must be 0 since a is one dimensionalrz]replace=False is not currently supported for dask.array.choice with multi-chunk output arrays)rrrr(rr+rrrbrarray ValueErrorr rechunkrpZ_metandimrFappendZ __dask_keys__r rrisclosesumtuplerXr r) rr7r2rPr>rOr3rTrVZdummy_pZlen_aerr_msgrrrrWBs`                   rWr.rrc% Os<|durt|ttfs|f}tddt||D}|durJ||t|}t|||dt j d}t |}dd} i} i} g} g} t |D]\}}t|t j tfr$| |||}t|tr| ||j| |<n,t|t j rdt|}|| |<|| |<| |td d |jDq| |qi}|D]\}}t|t j tfr| |||}t|tr| ||j| |<n,t|t j rdt|}|| |<|| |<|td d |jD||<n|||<qsz_wrap_func..rrcSsDt|trt|||St|tjr8tt||StddS)Nz!Unknown object type for broadcast)rr r rrrZascontiguousarray TypeError)rrpr3rrr_broadcast_anys   z"_wrap_func.._broadcast_anyzarray-css|] }dVqdSrNrrH_rrrrIrJz_wrap_func..css|] }dVqdSrrrrrrrIrJcSsg|] }|jqSr)r)rHZ_bitgenrrrrrJz_wrap_func..z:Unknown object type: Not a Generator and Not a RandomState-cSsg|]}tt|qSrrDrGrrrrrJrcSsg|]}tt|qSrrDrGrrrrrJzUnknown object type in argszUnknown object type in kwargs)rrSrU)&rr rXrvaluesr r r getrrr enumeraterr r^rrpitemsrrrYrrFrrrrrrrrrZrupdatedictrr[)%rrr2r3rrr9ZshapesZslicesrr`lookupZ small_argsrTirresr^Z small_kwargskeyr\r]Z bitgen_tokenZ func_appliergentokenr_blocksvalsrMZslcblockargZkwrgrLrVrarrrr5s                             r5zdict[str, RandomState]_cached_statescs,fdd}ttj|_ttj|_|S)Nc sltj}t>z t|}Wn ty8tt|<}Yn0Wdn1sN0Yt||i|Sr)rbackend_cached_states_lockr(KeyErrorrr)rr9r!rattrrrwrappers  2z_make_api..wrapper)rrr!r)r-r.rr,r _make_apis r/rr0r;r@rbrerirnrrrvrzrrrrrrrrrrrrrrrrrrrrrrrrrrrr)N)r)\ __future__rrr)r itertoolsrrroperatorr threadingrnumpyrZdask.array.backendsrZdask.array.corer r r r r rZdask.array.creationrZdask.array.utilsrZ dask.baserZdask.highlevelgraphrZ dask.utilsrrrrrrrrrYrrrKrrWr5r(__annotations__r*r/rr0r;r@rbrerirnrrrvrzrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrs         p _(   G