diff --git a/CHANGELOG.md b/CHANGELOG.md index 9dee06ee9ca..1a8c8e024bc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Added the `dpnp.random.Generator` API with `default_rng`, `BitGenerator`, `SeedSequence` and the `Philox4x3210` bit generator, sampling on the device with the oneMKL device RNG API [#3157](https://github.com/IntelPython/dpnp/pull/3157) * Added the continuous distributions of `numpy.random.Generator`, such as `gamma`, `beta` and `vonmises`, to `dpnp.random.Generator` [#3158](https://github.com/IntelPython/dpnp/pull/3158) * Added `integers` and the discrete distributions of `numpy.random.Generator`, such as `poisson`, `binomial` and `hypergeometric`, to `dpnp.random.Generator`, along with `noncentral_chisquare` and `noncentral_f` [#3159](https://github.com/IntelPython/dpnp/pull/3159) +* Added the multivariate distributions, `permutation`, `shuffle`, `permuted`, `choice` and `bytes` to `dpnp.random.Generator` [#3160](https://github.com/IntelPython/dpnp/pull/3160) ### Changed diff --git a/doc/reference/random.rst b/doc/reference/random.rst index b6b989aa2e8..0a7f4a300f3 100644 --- a/doc/reference/random.rst +++ b/doc/reference/random.rst @@ -55,6 +55,8 @@ Differences from NumPy NumPy's rejection test rounds to accepting every sample above ``2**53``, which makes its tail too heavy for ``a`` close to 1. * ``hypergeometric`` requires ``ngood`` and ``nbad`` below ``10**9``. +* :meth:`Generator.shuffle` only accepts a :class:`dpnp.ndarray`, as in + CuPy. Random Generator ---------------- @@ -75,6 +77,19 @@ Simple random data Generator.integers Generator.random + Generator.choice + Generator.bytes + +Permutations +^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + :nosignatures: + + Generator.shuffle + Generator.permutation + Generator.permuted Distributions ^^^^^^^^^^^^^ @@ -86,6 +101,7 @@ Distributions Generator.beta Generator.binomial Generator.chisquare + Generator.dirichlet Generator.exponential Generator.f Generator.gamma @@ -96,6 +112,9 @@ Distributions Generator.logistic Generator.lognormal Generator.logseries + Generator.multinomial + Generator.multivariate_hypergeometric + Generator.multivariate_normal Generator.negative_binomial Generator.noncentral_chisquare Generator.noncentral_f diff --git a/dpnp/random/_generator.py b/dpnp/random/_generator.py index 4274b1788d1..af9fc8ef934 100644 --- a/dpnp/random/_generator.py +++ b/dpnp/random/_generator.py @@ -40,6 +40,7 @@ # pylint: disable=too-many-lines import operator +import warnings import numpy @@ -48,6 +49,7 @@ import dpnp.tensor as dpt from dpnp.exceptions import ExecutionPlacementError from dpnp.tensor import validate_usm_type +from dpnp.tensor._numpy_helper import normalize_axis_index from ._bit_generator import BitGenerator, Philox4x3210 @@ -93,6 +95,25 @@ _F32_LT_1 = (None, numpy.nextafter(numpy.float32(1), numpy.float32(0))) _F32_LT_2_63 = (None, numpy.nextafter(numpy.float32(2**63), numpy.float32(0))) +# the messages of Generator.choice +_CHOICE_MSGS = { + "a_type": "a must be a sequence or an integer, not {}", + "a_int": "a must be a positive integer unless no samples are taken", + "a_1d": None, + "a_empty": "a cannot be empty unless no samples are taken", + "p_1d": "p must be 1-dimensional", + "p_size": "a and p must have same size", + "p_nan": "Probabilities contain NaN", + "p_neg": "Probabilities are not non-negative", + "p_sum": ( + "Probabilities do not sum to 1. See Notes section of docstring for " + "more information." + ), + "too_large": ( + "Cannot take a larger sample than population when replace is False" + ), +} + def _is_array(x): return dpnp.is_supported_array_type(x) or numpy.ndim(x) > 0 @@ -573,6 +594,16 @@ def _discrete_sample( double_precision=ft == numpy.float64, ) + def _permutation_index(self, n, usm_type="device"): + """ + A random permutation of ``arange(n)``, as ``int64`` values allocated + with `usm_type`. + + """ + + keys = self._bit_generator.random_raw(max(n, 0), usm_type=usm_type) + return dpnp.argsort(keys) + def random(self, size=None, dtype=None, out=None, *, usm_type="device"): """ Return random floats in the half-open interval [0.0, 1.0). @@ -1380,6 +1411,639 @@ def hypergeometric( param_types=(_int64,) * 3, ) + def dirichlet(self, alpha, size=None, dtype=None, *, usm_type="device"): + """ + Draw samples from the Dirichlet distribution. + + For full documentation refer to + :obj:`numpy.random.Generator.dirichlet`. + + Limitations + ----------- + `dtype` and `usm_type` are not in NumPy's signature. `alpha` is + copied to the host, where it is checked, and one kernel is launched + for every component when all of `alpha` is less than 0.1. + + """ + + self._check_queues(alpha) + validate_usm_type(usm_type, allow_none=False) + dtype = self._float_dtype(dtype, "dirichlet") + queue = self.sycl_queue + + k = len(alpha) + if dpnp.is_supported_array_type(alpha): + alpha = dpnp.asnumpy(alpha) + alpha = numpy.asarray(alpha, dtype=numpy.float64) + if alpha.ndim != 1: + raise ValueError( + "setting an array element with a sequence. The requested " + "array would exceed the maximum number of dimension of 1." + ) + if numpy.any(alpha < 0): + raise ValueError("alpha < 0") + + size = _normalize_size(size) + shape = (() if size is None else size) + (k,) + + if k == 0 or not alpha.max() < 0.1: + # normalize gamma variates, NaN in alpha giving NaN + res = self._float_sample( + ri._gamma, + (), + shape, + dtype, + usm_type, + "dirichlet", + kernel_params=(alpha, 1.0), + ) + res /= res.sum(axis=-1, keepdims=True) + return res + + # With small alpha the gamma variates underflow, so break a stick + # by beta variates instead, as NumPy does. csum[j] is sum(alpha[j:]). + csum = numpy.cumsum(alpha[::-1])[::-1] + res = dpnp.zeros( + shape, dtype=dtype, usm_type=usm_type, sycl_queue=queue + ) + if csum[0] == 0: + return res + + batch = shape[:-1] + acc = dpnp.ones(batch, dtype=dtype, usm_type=usm_type, sycl_queue=queue) + last = k - 1 + for j in range(k - 1): + if csum[j + 1] == 0: + # the stick left goes to component j + last = j + break + if alpha[j] == 0: + continue + v = self._float_sample( + ri._beta, + (), + batch, + dtype, + usm_type, + "dirichlet", + kernel_params=(alpha[j], csum[j + 1]), + ) + res[..., j] = acc * v + acc *= 1 - v + res[..., last] = acc + return res + + def multinomial(self, n, pvals, size=None, *, usm_type="device"): + """ + Draw samples from a multinomial distribution. + + For full documentation refer to + :obj:`numpy.random.Generator.multinomial`. + + Limitations + ----------- + `usm_type` is not in NumPy's signature. One binomial kernel is + launched for every category but the last. + + """ + + self._check_queues(n, pvals) + validate_usm_type(usm_type, allow_none=False) + queue = self.sycl_queue + ft = self._discrete_float_type() + + on_host = not dpnp.is_supported_array_type(pvals) + if on_host: + pvals = numpy.asarray(pvals, dtype=numpy.float64) + if pvals.ndim == 0 or pvals.shape[-1] == 0: + raise ValueError( + "pvals must have at least 1 dimension and the last dimension " + "of pvals must be greater than 0." + ) + if on_host: + _check_constraint(pvals, "pvals", _CONS_BOUNDED_0_1, False) + if numpy.any(pvals[..., :-1].sum(axis=-1) > 1.0 + 1e-12): + raise ValueError( + "sum(pvals[:-1]) > 1.0" + if pvals.ndim == 1 + else "sum(pvals[...,:-1]) > 1.0" + ) + + n = _as_int64(n) + _check_constraint( + n, + "n", + _CONS_NON_NEGATIVE, + numpy.ndim(n) == 0, + ) + + d = pvals.shape[-1] + batch = numpy.broadcast_shapes(numpy.shape(n), pvals.shape[:-1]) + size = _normalize_size(size) + if size is not None: + try: + compatible = numpy.broadcast_shapes(size, batch) == size + except ValueError: + compatible = False + if not compatible: + raise ValueError( + f"Output size {size} is not compatible with broadcast " + f"dimensions of inputs {batch}." + ) + batch = size + + res = dpnp.empty( + batch + (d,), dtype=_int64, usm_type=usm_type, sycl_queue=queue + ) + + # Draw the categories in turn, each from a binomial of the trials + # left with the probability of the category conditioned on not + # being any of the previous ones + host_p = on_host and pvals.ndim == 1 + if host_p: + remaining = 1.0 + else: + pvals = dpnp.asarray(pvals, dtype=ft, sycl_queue=queue) + remaining = dpnp.ones_like(pvals[..., 0]) + dn = dpnp.asarray(n, sycl_queue=queue) if _is_array(n) else n + for j in range(d - 1): + pj = pvals[..., j] + if host_p: + prob = ( + min(max(pj / remaining, 0.0), 1.0) if remaining > 0 else 0 + ) + else: + prob = dpnp.where( + remaining > 0, dpnp.clip(pj / remaining, 0, 1), 0 + ) + x = self._discrete_sample( + ri._binomial, + (), + batch, + usm_type, + kernel_params=(dn, prob), + param_types=(_int64, None), + ) + res[..., j] = x + dn = dn - x + remaining = remaining - pj + res[..., d - 1] = dn + return res + + def multivariate_hypergeometric( + self, + colors, + nsample, + size=None, + method="marginals", + *, + usm_type="device", + ): + """ + Generate variates from a multivariate hypergeometric distribution. + + For full documentation refer to + :obj:`numpy.random.Generator.multivariate_hypergeometric`. + + Limitations + ----------- + `usm_type` is not in NumPy's signature. Both methods draw by the + marginals algorithm, launching one kernel for every color but the + last, so ``sum(colors)`` must be less than 1000000000 for the + ``"count"`` method as well. + + """ + + self._check_queues(colors) + validate_usm_type(usm_type, allow_none=False) + if method not in ("count", "marginals"): + raise ValueError('method must be "count" or "marginals".') + + try: + nsample = operator.index(nsample) + except TypeError: + raise ValueError("nsample must be an integer") from None + if nsample < 0: + raise ValueError("nsample must be nonnegative.") + if nsample > _INT64_MAX: + raise ValueError(f"nsample must not exceed {_INT64_MAX}") + + if dpnp.is_supported_array_type(colors): + colors = dpnp.asnumpy(colors) + invalid_colors = False + try: + colors = numpy.asarray(colors) + if colors.ndim != 1: + invalid_colors = True + elif colors.size > 0 and not numpy.issubdtype( + colors.dtype, numpy.integer + ): + invalid_colors = True + elif numpy.any((colors < 0) | (colors > _INT64_MAX)): + invalid_colors = True + except ValueError: + invalid_colors = True + if invalid_colors: + raise ValueError( + "colors must be a one-dimensional sequence of nonnegative " + f"integers not exceeding {_INT64_MAX}." + ) + + colors = numpy.ascontiguousarray(colors, dtype=_int64) + total = sum(int(c) for c in colors) + if total > _INT64_MAX: + raise ValueError( + "sum(colors) must not exceed the maximum value of a 64 bit " + f"signed integer ({_INT64_MAX})" + ) + if total >= _HYPERGEOM_MAX: + raise ValueError( + f'When method is "{method}", sum(colors) must be less than ' + f"{_HYPERGEOM_MAX}." + ) + if nsample > total: + raise ValueError("nsample > sum(colors)") + + size = _normalize_size(size) + batch = () if size is None else size + k = colors.size + res = dpnp.zeros( + batch + (k,), + dtype=_int64, + usm_type=usm_type, + sycl_queue=self.sycl_queue, + ) + if k == 0 or nsample == 0 or res.size == 0: + return res + + # draw the complement when it is smaller + more_than_half = nsample > total // 2 + num_to_sample = total - nsample if more_than_half else nsample + remaining = total + for j in range(k - 1): + remaining -= int(colors[j]) + x = self._discrete_sample( + ri._hypergeometric, + (), + batch, + usm_type, + kernel_params=(int(colors[j]), remaining, num_to_sample), + param_types=(_int64,) * 3, + ) + res[..., j] = x + num_to_sample = num_to_sample - x + res[..., k - 1] = num_to_sample + + if more_than_half: + dpnp.subtract( + dpnp.asarray( + colors, usm_type=usm_type, sycl_queue=res.sycl_queue + ), + res, + out=res, + ) + return res + + def multivariate_normal( + self, + mean, + cov, + size=None, + check_valid="warn", + tol=1e-8, + *, + method="svd", + dtype=None, + usm_type="device", + ): + """ + Draw random samples from a multivariate normal distribution. + + For full documentation refer to + :obj:`numpy.random.Generator.multivariate_normal`. + + Limitations + ----------- + `dtype` and `usm_type` are not in NumPy's signature. The covariance + matrix is factorized on the device by :mod:`dpnp.linalg`, in double + precision if the device supports it. Otherwise it is factorized in + single precision, and `tol` is taken to be at least ``100 * eps`` of + :obj:`dpnp.float32`. + + """ + return self._multivariate_normal( + mean, + cov, + size, + check_valid, + tol, + method, + dtype, + usm_type, + stacklevel=3, + ) + + def _multivariate_normal( + self, + mean, + cov, + size, + check_valid, + tol, + method, + dtype, + usm_type, + stacklevel, + ): + if method not in {"eigh", "svd", "cholesky"}: + raise ValueError( + "method must be one of {'eigh', 'svd', 'cholesky'}" + ) + + self._check_queues(mean, cov) + validate_usm_type(usm_type, allow_none=False) + dtype = self._float_dtype(dtype, "multivariate_normal") + queue = self.sycl_queue + mean = dpnp.asarray(mean, usm_type=usm_type, sycl_queue=queue) + cov = dpnp.asarray(cov, usm_type=usm_type, sycl_queue=queue) + if dpnp.issubdtype(mean.dtype, dpnp.complexfloating) or dpnp.issubdtype( + cov.dtype, dpnp.complexfloating + ): + raise TypeError("mean and cov must not be complex") + + if mean.ndim != 1: + raise ValueError("mean must be 1 dimensional") + if cov.ndim != 2 or cov.shape[0] != cov.shape[1]: + raise ValueError("cov must be 2 dimensional and square") + if mean.shape[0] != cov.shape[0]: + raise ValueError("mean and cov must have same length") + + cov = cov.astype(dpnp.default_float_type(sycl_queue=queue)) + if cov.dtype == dpnp.float32: + tol = max(tol, 100 * dpnp.finfo(dpnp.float32).eps) + # as in NumPy, check_valid is ignored by the cholesky method + if method != "cholesky" and check_valid not in ( + "warn", + "raise", + "ignore", + ): + raise ValueError( + "check_valid must equal 'warn', 'raise', or 'ignore'" + ) + + psd = True + if method == "svd": + u, s, vh = dpnp.linalg.svd(cov) + if check_valid != "ignore": + psd = dpnp.allclose( + dpnp.dot(vh.T * s, vh), cov, rtol=tol, atol=tol + ) + factor = u * dpnp.sqrt(s) + elif method == "eigh": + s, u = dpnp.linalg.eigh(cov) + if check_valid != "ignore": + psd = not dpnp.any(s < -tol) + factor = u * dpnp.sqrt(dpnp.abs(s)) + else: + factor = dpnp.linalg.cholesky(cov) + + if not psd: + msg = "covariance is not symmetric positive-semidefinite." + if check_valid == "warn": + warnings.warn(msg, RuntimeWarning, stacklevel=stacklevel) + else: + raise ValueError(msg) + + size = _normalize_size(size) + shape = (() if size is None else size) + mean.shape + x = self._float_sample( + ri._normal, + (), + shape, + dtype, + usm_type, + "multivariate_normal", + kernel_params=(0.0, 1.0), + ) + factor = factor.astype(dtype) + x = mean.astype(dtype) + x.reshape(-1, mean.shape[0]) @ factor.T + return x.reshape(shape) + + def permutation(self, x, axis=0): + """ + Randomly permute a sequence, or return a permuted range. + + For full documentation refer to + :obj:`numpy.random.Generator.permutation`. + + """ + + if isinstance(x, (int, numpy.integer)): + return self._permutation_index(int(x)) + + self._check_queues(x) + arr = dpnp.asarray(x, sycl_queue=self.sycl_queue) + axis = normalize_axis_index(axis, arr.ndim) + idx = self._permutation_index(arr.shape[axis], arr.usm_type) + return dpnp.take(arr, idx, axis=axis) + + def shuffle(self, x, axis=0): + """ + Modify an array in place by shuffling its contents. + + For full documentation refer to + :obj:`numpy.random.Generator.shuffle`. + + Limitations + ----------- + As in CuPy, `x` must be a :class:`dpnp.ndarray`. + + """ + + if not isinstance(x, dpnp.ndarray): + raise TypeError(f"x must be a dpnp.ndarray, got {type(x)}") + self._check_queues(x) + if x.ndim == 0: + raise TypeError("len() of unsized object") + axis = normalize_axis_index(axis, x.ndim) + idx = self._permutation_index(x.shape[axis], x.usm_type) + x[...] = dpnp.take(x, idx, axis=axis) + + def permuted(self, x, *, axis=None, out=None): + """ + Randomly permute `x` along axis `axis`. + + For full documentation refer to + :obj:`numpy.random.Generator.permuted`. + + Limitations + ----------- + `out` must be a :class:`dpnp.ndarray`. + + """ + + self._check_queues(x, out) + x = dpnp.asarray(x, sycl_queue=self.sycl_queue) + + if out is not None: + if not isinstance(out, dpnp.ndarray): + raise TypeError("out must be a dpnp.ndarray") + if out.shape != x.shape: + raise ValueError("out must have the same shape as x") + if not dpnp.can_cast(x.dtype, out.dtype, casting="safe"): + raise TypeError( + f"Cannot cast array data from {x.dtype} to {out.dtype} " + "according to the rule 'safe'" + ) + + if axis is None: + if x.ndim == 0: + raise TypeError("len() of unsized object") + idx = self._permutation_index(x.size, x.usm_type) + res = dpnp.take(dpnp.ravel(x), idx).reshape(x.shape) + else: + axis = normalize_axis_index(axis, x.ndim) + keys = self._bit_generator.random_raw(x.shape, usm_type=x.usm_type) + idx = dpnp.argsort(keys, axis=axis) + res = dpnp.take_along_axis(x, idx, axis=axis) + + if out is None: + return res + out[...] = res + return out + + def choice( + self, + a, + size=None, + replace=True, + p=None, + axis=0, + shuffle=True, + *, + usm_type="device", + ): + """ + Generate a random sample from a given array. + + For full documentation refer to :obj:`numpy.random.Generator.choice`. + + Limitations + ----------- + `usm_type` is not in NumPy's signature. Without replacement, the + sample is always in random order, so `shuffle` has no effect, and + memory and time proportional to the population size are used. `p` + is checked only when it is on the host. + + """ + + # pylint: disable=unused-argument + return self._choice(a, size, replace, p, axis, usm_type, _CHOICE_MSGS) + + # pylint: disable=too-many-branches, too-many-statements + def _choice(self, a, size, replace, p, axis, usm_type, msgs): + """ + The implementation of :meth:`choice`, raising the messages of the + table `msgs`, as those of NumPy's legacy ``RandomState.choice`` + differ. + + """ + + self._check_queues(a, p) + validate_usm_type(usm_type, allow_none=False) + queue = self.sycl_queue + no_samples = size is not None and numpy.prod(size) == 0 + + a_original = a + a = ( + dpnp.asarray(a) + if dpnp.is_supported_array_type(a) + else numpy.asarray(a) + ) + if a.ndim == 0: + try: + pop_size = operator.index(a.item()) + except TypeError as exc: + raise ValueError( + msgs["a_type"].format(type(a_original)) + ) from exc + if pop_size <= 0 and not no_samples: + raise ValueError(msgs["a_int"]) + else: + if msgs["a_1d"] is not None and a.ndim != 1: + raise ValueError(msgs["a_1d"]) + axis = normalize_axis_index(axis, a.ndim) + pop_size = a.shape[axis] + if pop_size == 0 and not no_samples: + raise ValueError(msgs["a_empty"]) + + p_host = p_check = None + if p is not None: + if dpnp.is_supported_array_type(p): + p_dev = dpnp.asarray(p) + else: + atol = numpy.sqrt(numpy.finfo(numpy.float64).eps) + if isinstance(p, numpy.ndarray) and numpy.issubdtype( + p.dtype, numpy.floating + ): + atol = max(atol, numpy.sqrt(numpy.finfo(p.dtype).eps)) + p_host = numpy.asarray(p, dtype=numpy.float64) + p_dev = p_host + if p_dev.ndim != 1: + raise ValueError(msgs["p_1d"]) + if p_dev.size != pop_size: + raise ValueError(msgs["p_size"]) + p_check = p_host + if p_check is not None: + p_sum = float(p_check.sum()) + if numpy.isnan(p_sum): + raise ValueError(msgs["p_nan"]) + if (p_check < 0).any(): + raise ValueError(msgs["p_neg"]) + if abs(p_sum - 1.0) > atol: + raise ValueError(msgs["p_sum"]) + ft = dpnp.default_float_type(sycl_queue=queue) + p_dev = dpnp.asarray( + p_dev, dtype=ft, usm_type=usm_type, sycl_queue=queue + ) + + shape = () if size is None else _normalize_size(size) + n = int(numpy.prod(shape, dtype=numpy.int64)) + if not replace: + if n > pop_size: + raise ValueError(msgs["too_large"]) + if p_check is not None and int((p_check > 0).sum()) < n: + raise ValueError("Fewer non-zero entries in p than size") + + if n == 0: + idx = dpnp.empty( + shape, dtype=_int64, usm_type=usm_type, sycl_queue=queue + ) + elif replace: + if p is not None: + cdf = dpnp.cumsum(p_dev) + cdf /= cdf[-1] + u = self.random(shape, dtype=ft, usm_type=usm_type) + idx = dpnp.searchsorted(cdf, u, side="right") + else: + idx = self.integers(0, pop_size, size=shape, usm_type=usm_type) + elif p is not None: + # The Efraimidis-Spirakis keys: taking the smallest ones gives + # the successive weighted draws of NumPy, in the order drawn + keys = ( + self.standard_exponential(pop_size, dtype=ft, usm_type=usm_type) + / p_dev + ) + idx = dpnp.argsort(keys)[:n].reshape(shape) + else: + idx = self._permutation_index(pop_size, usm_type)[:n].reshape(shape) + + if a.ndim == 0: + return idx + if not isinstance(a, dpnp.ndarray): + a = dpnp.asarray(a, usm_type=usm_type, sycl_queue=queue) + return dpnp.take(a, idx, axis=axis) + def integers( self, low, @@ -1505,6 +2169,29 @@ def _to_bound(x, shift): ri._integers, (), size, dtype, usm_type, kernel_params=(low, high) ) + def bytes(self, length): + """ + Return random bytes. + + For full documentation refer to :obj:`numpy.random.Generator.bytes`. + + Limitations + ----------- + The bytes are drawn on the device and copied to the host. + + """ + + # as in NumPy, 32-bit values in little-endian order, as many as + # (length - 1) / 4 + 1 in C division, so bytes(0) draws one + length = operator.index(length) + q = length - 1 + n_uint32 = (q // 4 if q >= 0 else -(-q // 4)) + 1 + values = dpnp.empty( + n_uint32, dtype=dpnp.uint32, sycl_queue=self.sycl_queue + ) + self._bit_generator._fill(ri._random_raw, (), values) + return dpnp.asnumpy(values).astype("= 0) & (res <= 1)).all() + assert_allclose(res.sum(axis=-1), numpy.ones(2 * n), rtol=1e-5) + + alpha = numpy.asarray(alpha) + a0 = alpha.sum() + mean = alpha / a0 + var = mean * (1 - mean) / (a0 + 1) + assert ( + abs(res.mean(axis=0) - mean) <= 5 * numpy.sqrt(var / (2 * n)) + ).all() + + @pytest.mark.parametrize( + "alpha", [[0.0], [0.0, 0.0], [0.05, 0.0], [0.0, 0.0, 0.05], [0.0, 2.0]] + ) + def test_dirichlet_degenerate(self, alpha): + expected = numpy.random.default_rng(0).dirichlet(alpha) + res = default_rng(0).dirichlet(alpha, size=3) + assert_allclose( + dpnp.asnumpy(res), numpy.tile(expected, (3, 1)), rtol=1e-15 + ) + + def test_dirichlet_shapes(self): + rng = default_rng(0) + assert rng.dirichlet([1.0, 2.0]).shape == (2,) + assert rng.dirichlet([]).shape == (0,) + assert rng.dirichlet([], size=3).shape == (3, 0) + assert rng.dirichlet(dpnp.asarray([1.0, 2.0]), size=3).shape == (3, 2) + res = dpnp.asnumpy(rng.dirichlet([1.0, numpy.nan], size=3)) + assert numpy.isnan(res).all() + + def test_multinomial(self, discrete_precision): + n = 10**5 + pvals = [0.2, 0.3, 0.0, 0.5] + res = default_rng(8).multinomial(100, pvals, size=n) + assert res.shape == (n, 4) + assert res.dtype == dpnp.int64 + res = dpnp.asnumpy(res) + assert (res.sum(axis=1) == 100).all() + assert (res[:, 2] == 0).all() + + p = numpy.asarray(pvals) + assert ( + abs(res.mean(axis=0) - 100 * p) + <= 5 * numpy.sqrt(100 * p * (1 - p) / n) + 1e-12 + ).all() + assert_allclose( + res.var(axis=0), 100 * p * (1 - p), rtol=0.05, atol=1e-12 + ) + + def test_multinomial_shapes(self, discrete_precision): + rng = default_rng(9) + assert rng.multinomial(5, [0.5, 0.5]).shape == (2,) + assert rng.multinomial(5, [0.5, 0.5], size=(3, 2)).shape == (3, 2, 2) + assert rng.multinomial([3, 5], [0.2, 0.8], size=(4, 2)).shape == ( + 4, + 2, + 2, + ) + + n = [[1], [20]] + pvals = [[0.5, 0.5], [0.2, 0.8], [0.0, 1.0]] + for pv in (pvals, dpnp.asarray(pvals)): + res = rng.multinomial(n, pv) + assert res.shape == (2, 3, 2) + res = dpnp.asnumpy(res) + assert_array_equal(res.sum(axis=-1), numpy.broadcast_to(n, (2, 3))) + assert (res[:, 2, 0] == 0).all() + + res = rng.multinomial(dpnp.asarray([4, 7]), [0.3, 0.7]) + assert_array_equal(dpnp.asnumpy(res).sum(axis=-1), [4, 7]) + assert_raises(ValueError, rng.multinomial, [3, 5], [0.2, 0.8], size=3) + + @pytest.mark.parametrize("method", ["marginals", "count"]) + @pytest.mark.parametrize("nsample", [0, 7, 25, 45, 60]) + def test_multivariate_hypergeometric( + self, method, nsample, discrete_precision + ): + n = 10**5 + colors = numpy.array([10, 20, 0, 30]) + res = default_rng(10).multivariate_hypergeometric( + colors, nsample, size=n, method=method + ) + assert res.shape == (n, 4) + assert res.dtype == dpnp.int64 + res = dpnp.asnumpy(res) + assert (res.sum(axis=1) == nsample).all() + assert ((res >= 0) & (res <= colors)).all() + + total = colors.sum() + p = colors / total + mean = nsample * p + var = mean * (1 - p) * (total - nsample) / (total - 1) + assert ( + abs(res.mean(axis=0) - mean) <= 5 * numpy.sqrt(var / n) + 1e-12 + ).all() + + def test_multivariate_hypergeometric_shapes(self): + rng = default_rng(0) + assert rng.multivariate_hypergeometric([1, 2], 2).shape == (2,) + assert rng.multivariate_hypergeometric([], 0).shape == (0,) + assert rng.multivariate_hypergeometric([], 0, size=3).shape == (3, 0) + res = rng.multivariate_hypergeometric( + dpnp.asarray([3, 4]), 7, size=(2, 3) + ) + assert res.shape == (2, 3, 2) + assert (dpnp.asnumpy(res) == [3, 4]).all() + + @pytest.mark.parametrize("dtype", float_dtypes) + @pytest.mark.parametrize("method", ["svd", "eigh", "cholesky"]) + def test_multivariate_normal(self, method, dtype): + n = 10**5 + mean = numpy.array([1.0, -2.0, 3.0]) + cov = numpy.array([[2.0, 0.5, 0.0], [0.5, 1.0, -0.3], [0.0, -0.3, 0.5]]) + res = default_rng(11).multivariate_normal( + mean, cov, size=(n, 2), method=method, dtype=dtype + ) + assert res.shape == (n, 2, 3) + assert res.dtype == dtype + res = dpnp.asnumpy(res).reshape(-1, 3).astype(numpy.float64) + assert_allclose(res.mean(axis=0), mean, atol=0.02) + assert_allclose(numpy.cov(res.T), cov, atol=0.02) + + def test_multivariate_normal_shapes(self): + rng = default_rng(0) + mean = dpnp.zeros(2) + cov = dpnp.eye(2) + assert rng.multivariate_normal(mean, cov).shape == (2,) + assert rng.multivariate_normal(mean, cov, size=3).shape == (3, 2) + assert rng.multivariate_normal([0.0], [[1.0]], size=0).shape == (0, 1) + res = rng.multivariate_normal( + [0.0, 0.0], [[1.0, 0.0], [0.0, 0.0]], size=100 + ) + assert (dpnp.asnumpy(res)[:, 1] == 0).all() + res = rng.multivariate_normal(mean, cov, size=2, usm_type="host") + assert res.usm_type == "host" + + +class TestPermutations: + def test_permutation(self): + rng = default_rng(12) + res = rng.permutation(10) + assert res.dtype == dpnp.int64 + assert_array_equal(numpy.sort(dpnp.asnumpy(res)), numpy.arange(10)) + assert rng.permutation(-3).shape == (0,) + assert rng.permutation(0).shape == (0,) + + x = dpnp.arange(12).reshape(3, 4) + res = rng.permutation(x) + assert_array_equal( + numpy.sort(dpnp.asnumpy(res), axis=0), dpnp.asnumpy(x) + ) + # the rows are moved whole + assert ( + dpnp.asnumpy(res)[:, 1:] - dpnp.asnumpy(res)[:, :1] == [1, 2, 3] + ).all() + + res = dpnp.asnumpy(rng.permutation(x, axis=-1)) + assert (res[1:] - res[:1] == [[4], [8]]).all() + + res = rng.permutation([1.5, 2.5, 3.5]) + assert isinstance(res, dpnp.ndarray) + assert_array_equal(numpy.sort(dpnp.asnumpy(res)), [1.5, 2.5, 3.5]) + + # a copy is returned + x = dpnp.arange(5) + rng.permutation(x)[...] = 0 + assert_array_equal(x, dpnp.arange(5)) + + def test_permutation_uniform(self): + # all 6 permutations of 3 values are equally likely + rng = default_rng(13) + n = 6000 + keys = [tuple(dpnp.asnumpy(rng.permutation(3))) for _ in range(n)] + _, counts = numpy.unique(keys, axis=0, return_counts=True) + assert len(counts) == 6 + assert (abs(counts - n / 6) < 5 * numpy.sqrt(n / 6)).all() + + def test_shuffle(self): + rng = default_rng(14) + x = dpnp.arange(20) + assert rng.shuffle(x) is None + assert not (x == dpnp.arange(20)).all() + assert_array_equal(numpy.sort(dpnp.asnumpy(x)), numpy.arange(20)) + + x = dpnp.arange(12).reshape(3, 4) + rng.shuffle(x, axis=1) + res = dpnp.asnumpy(x) + assert (res[1:] - res[:1] == [[4], [8]]).all() + + def test_permuted(self): + rng = default_rng(15) + x = dpnp.arange(24).reshape(4, 6) + res = dpnp.asnumpy(rng.permuted(x)) + assert res.shape == (4, 6) + assert_array_equal(numpy.sort(res, axis=None), numpy.arange(24)) + + res = dpnp.asnumpy(rng.permuted(x, axis=1)) + assert_array_equal(numpy.sort(res, axis=1), dpnp.asnumpy(x)) + # the rows are shuffled independently + assert len({tuple(r - r.min()) for r in res}) > 1 + + out = dpnp.empty((4, 6), dtype=dpnp.float32) + res = rng.permuted(x.astype(dpnp.int16), axis=0, out=out) + assert res is out + assert_array_equal( + numpy.sort(dpnp.asnumpy(out), axis=0), dpnp.asnumpy(x), strict=False + ) + + # in place + x = dpnp.arange(10) + rng.permuted(x, out=x) + assert_array_equal(numpy.sort(dpnp.asnumpy(x)), numpy.arange(10)) + + res = rng.permuted([[1, 2], [3, 4]], axis=0) + assert isinstance(res, dpnp.ndarray) + + +class TestChoice: + def test_int(self): + rng = default_rng(16) + res = rng.choice(5, size=1000) + assert res.dtype == dpnp.int64 + assert_array_equal(numpy.unique(dpnp.asnumpy(res)), numpy.arange(5)) + + res = rng.choice(5) + assert isinstance(res, dpnp.ndarray) and res.shape == () + assert rng.choice(5, size=()).shape == () + assert rng.choice(0, size=0).shape == (0,) + assert rng.choice(5, size=(2, 3)).shape == (2, 3) + + def test_without_replacement(self): + rng = default_rng(17) + res = dpnp.asnumpy(rng.choice(100, size=(10, 10), replace=False)) + assert_array_equal(numpy.sort(res, axis=None), numpy.arange(100)) + res = dpnp.asnumpy(rng.choice(1000, size=5, replace=False)) + assert len(numpy.unique(res)) == 5 + + def test_weighted(self): + n = 10**5 + p = numpy.array([0.1, 0.0, 0.3, 0.6]) + rng = default_rng(18) + for pp in (p, dpnp.asarray(p)): + res = dpnp.asnumpy(rng.choice(4, size=n, p=pp)) + freq = numpy.bincount(res, minlength=4) / n + assert (abs(freq - p) <= 5 * numpy.sqrt(p * (1 - p) / n)).all() + + def test_weighted_without_replacement(self): + n = 20000 + p = numpy.array([0.1, 0.0, 0.3, 0.6]) + rng = default_rng(19) + res = numpy.array( + [ + dpnp.asnumpy(rng.choice(4, size=2, replace=False, p=p)) + for _ in range(n) + ] + ) + assert (res != 1).all() + assert (res[:, 0] != res[:, 1]).all() + + # the first draw is from p, and the second from p conditioned on + # not being the first + freq = numpy.bincount(res[:, 0], minlength=4) / n + assert (abs(freq - p) <= 5 * numpy.sqrt(p * (1 - p) / n)).all() + pair = (res[:, 0] == 3) & (res[:, 1] == 2) + prob = 0.6 * 0.3 / 0.4 + assert abs(pair.mean() - prob) <= 5 * numpy.sqrt(prob * (1 - prob) / n) + + def test_array(self): + rng = default_rng(20) + a = dpnp.asarray([10, 20, 30]) + res = rng.choice(a, size=100) + assert set(dpnp.asnumpy(res)) <= {10, 20, 30} + res = rng.choice(a) + assert res.shape == () + + a = numpy.arange(12).reshape(3, 4) + res = dpnp.asnumpy(rng.choice(a, size=5)) + assert res.shape == (5, 4) + assert (res[:, 1:] - res[:, :1] == [1, 2, 3]).all() + res = dpnp.asnumpy(rng.choice(a, size=(2, 2), axis=1, replace=False)) + assert res.shape == (3, 2, 2) + assert_array_equal(numpy.sort(res[0], axis=None), a[0]) + assert rng.choice(a, axis=1).shape == (3,) + @pytest.mark.skipif(len(_all_devices()) < 2, reason="needs two devices") class TestDeviceIndependence: diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py index 45c423030ce..92bdeb52abc 100644 --- a/dpnp/tests/test_sycl_queue.py +++ b/dpnp/tests/test_sycl_queue.py @@ -929,6 +929,11 @@ def test_random_state(func, args, kwargs, device, usm_type): pytest.param("beta", [0.5, 2.0], {"size": (3, 4)}), pytest.param("binomial", [10, 0.3], {"size": (3, 4)}), pytest.param("chisquare", [3.0], {"size": (3, 4)}), + pytest.param( + "choice", [5], {"size": (3, 4), "p": [0.1, 0.2, 0.3, 0.2, 0.2]} + ), + pytest.param("choice", [5], {"size": 3, "replace": False}), + pytest.param("dirichlet", [[0.5, 1.5]], {"size": 3}), pytest.param("exponential", [2.0], {"size": (3, 4)}), pytest.param("f", [2.0, 3.0], {"size": (3, 4)}), pytest.param("gamma", [2.0], {"size": (3, 4)}), @@ -940,11 +945,17 @@ def test_random_state(func, args, kwargs, device, usm_type): pytest.param("logistic", [], {"size": (3, 4)}), pytest.param("lognormal", [], {"size": (3, 4)}), pytest.param("logseries", [0.5], {"size": (3, 4)}), + pytest.param("multinomial", [5, [0.2, 0.8]], {"size": 3}), + pytest.param("multivariate_hypergeometric", [[3, 4], 2], {"size": 3}), + pytest.param( + "multivariate_normal", [[0.0, 0.0], [[1.0, 0.0], [0.0, 1.0]]], {} + ), pytest.param("negative_binomial", [3, 0.4], {"size": (3, 4)}), pytest.param("noncentral_chisquare", [3.0, 1.0], {"size": (3, 4)}), pytest.param("noncentral_f", [2.0, 3.0, 1.0], {"size": (3, 4)}), pytest.param("normal", [], {"size": (3, 4)}), pytest.param("pareto", [2.0], {"size": (3, 4)}), + pytest.param("permutation", [5], {}), pytest.param("poisson", [3.0], {"size": (3, 4)}), pytest.param("power", [2.0], {"size": (3, 4)}), pytest.param("random", [], {"size": (3, 4)}), @@ -990,16 +1001,35 @@ def test_bit_generator(bit_generator, device): assert bg.random_raw(5).sycl_device == device +@pytest.mark.parametrize("func", ["permutation", "permuted", "shuffle"]) +@pytest.mark.parametrize("device", valid_dev, ids=dev_ids) +def test_generator_permutations(func, device): + sycl_queue = dpctl.SyclQueue(device) + rng = dpnp.random.default_rng(1, sycl_queue=sycl_queue) + x = dpnp.arange(10, sycl_queue=sycl_queue) + res = getattr(rng, func)(x) + if res is None: + res = x + assert_sycl_queue_equal(res.sycl_queue, sycl_queue) + + # "f", "i" and "p" stand for float, integer and probability arrays on another # queue than the generator's _diff_queue_methods = [ pytest.param("binomial", ["i", 0.5]), + pytest.param("choice", ["f"]), + pytest.param("choice", [4, None, True, "p"]), + pytest.param("dirichlet", ["f"]), pytest.param("geometric", ["p"]), pytest.param("hypergeometric", ["i", "i", 1]), pytest.param("logseries", ["p"]), + pytest.param("multinomial", [5, "p"]), + pytest.param("multivariate_normal", ["f", numpy.eye(4).tolist()]), pytest.param("negative_binomial", ["f", 0.5]), pytest.param("normal", ["f"]), + pytest.param("permutation", ["f"]), pytest.param("poisson", ["f"]), + pytest.param("shuffle", ["f"]), pytest.param("standard_gamma", ["f"]), pytest.param("triangular", [0.0, "f", 5.0]), pytest.param("uniform", ["f", 5.0]), @@ -1022,6 +1052,9 @@ def _diff_queue_args(args, device): _diff_queue_methods + [ pytest.param("integers", ["i"]), + pytest.param("multivariate_hypergeometric", ["i", 2]), + pytest.param("permuted", ["f"]), + pytest.param("permuted", [[1.0, 2.0, 3.0, 4.0]], id="permuted-out"), pytest.param("random", [4, None], id="random-out"), ], ) diff --git a/dpnp/tests/test_usm_type.py b/dpnp/tests/test_usm_type.py index 90eb7c938e1..fd48df0945c 100644 --- a/dpnp/tests/test_usm_type.py +++ b/dpnp/tests/test_usm_type.py @@ -1837,8 +1837,25 @@ def test_tensorsolve(self, usm_type, usm_type_other): @pytest.mark.parametrize( "func, args, kwargs", [ + pytest.param("choice", [5], {"size": 3}), + pytest.param( + "choice", [5], {"size": 3, "p": [0.1, 0.2, 0.3, 0.2, 0.2]} + ), + pytest.param("choice", [5], {"size": 3, "replace": False}), + pytest.param( + "choice", + [5], + {"size": 3, "replace": False, "p": [0.1, 0.2, 0.3, 0.2, 0.2]}, + ), + pytest.param("dirichlet", [[0.05, 0.05]], {"size": 3}), + pytest.param("dirichlet", [[0.5, 1.5]], {"size": 3}), pytest.param("hypergeometric", [5, 7, 4], {"size": 3}), pytest.param("integers", [[3, 5], [7, 9]], {}), + pytest.param("multinomial", [5, [0.2, 0.8]], {"size": 3}), + pytest.param("multivariate_hypergeometric", [[3, 4], 2], {"size": 3}), + pytest.param( + "multivariate_normal", [[0.0, 0.0], [[1.0, 0.0], [0.0, 1.0]]], {} + ), pytest.param("normal", [[1.0, 2.0]], {}), pytest.param("poisson", [3.0], {"size": 3}), pytest.param("random", [3], {}), @@ -1861,8 +1878,25 @@ def test_generator(func, args, kwargs, usm_type): def test_generator_array_params(usm_type): rng = dpnp.random.default_rng(1) x = dpnp.full(4, 2.0, usm_type=usm_type) + p = dpnp.full(4, 0.25, usm_type=usm_type) assert rng.normal(x, usm_type=usm_type).usm_type == usm_type + assert rng.choice(x, 3, p=p, usm_type=usm_type).usm_type == usm_type + assert rng.multinomial(5, p, usm_type=usm_type).usm_type == usm_type + + +@pytest.mark.parametrize("usm_type", list_of_usm_types) +def test_generator_permutations(usm_type): + rng = dpnp.random.default_rng(1) + x = dpnp.arange(12, usm_type=usm_type).reshape(3, 4) + + assert rng.permutation(x).usm_type == usm_type + assert rng.permutation(x, axis=1).usm_type == usm_type + assert rng.permuted(x).usm_type == usm_type + assert rng.permuted(x, axis=1).usm_type == usm_type + + rng.shuffle(x) + assert x.usm_type == usm_type @pytest.mark.parametrize("bit_generator", [dpnp.random.Philox4x3210])