summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--numpy/random/_generator.pyi4
-rw-r--r--numpy/random/_generator.pyx155
-rw-r--r--numpy/random/tests/test_generator_mt19937.py64
3 files changed, 178 insertions, 45 deletions
diff --git a/numpy/random/_generator.pyi b/numpy/random/_generator.pyi
index 64b683d7c..c574bef9a 100644
--- a/numpy/random/_generator.pyi
+++ b/numpy/random/_generator.pyi
@@ -623,7 +623,9 @@ class Generator:
method: Literal["svd", "eigh", "cholesky"] = ...,
) -> ndarray[Any, dtype[float64]]: ...
def multinomial(
- self, n: _ArrayLikeInt_co, pvals: _ArrayLikeFloat_co, size: Optional[_ShapeLike] = ...
+ self, n: _ArrayLikeInt_co,
+ pvals: _ArrayLikeFloat_co,
+ size: Optional[_ShapeLike] = ...
) -> ndarray[Any, dtype[int64]]: ...
def multivariate_hypergeometric(
self,
diff --git a/numpy/random/_generator.pyx b/numpy/random/_generator.pyx
index 7087b6e1d..2134f0e24 100644
--- a/numpy/random/_generator.pyx
+++ b/numpy/random/_generator.pyx
@@ -3683,24 +3683,35 @@ cdef class Generator:
----------
n : int or array-like of ints
Number of experiments.
- pvals : sequence of floats, length p
- Probabilities of each of the ``p`` different outcomes. These
- must sum to 1 (however, the last element is always assumed to
- account for the remaining probability, as long as
- ``sum(pvals[:-1]) <= 1)``.
+ pvals : array-like of floats
+ Probabilities of each of the ``p`` different outcomes with shape
+ ``(k0, k1, ..., kn, p)``. Each element ``pvals[i,j,...,:]`` must
+ sum to 1 (however, the last element is always assumed to account
+ for the remaining probability, as long as
+ ``sum(pvals[..., :-1], axis=-1) <= 1.0``. Must have at least 1
+ dimension where pvals.shape[-1] > 0.
size : int or tuple of ints, optional
Output shape. If the given shape is, e.g., ``(m, n, k)``, then
- ``m * n * k`` samples are drawn. Default is None, in which case a
- single value is returned.
+ ``m * n * k`` samples are drawn each with ``p`` elements. Default
+ is None where the output size is determined by the broadcast shape
+ of ``n`` and all by the final dimension of ``pvals``, which is
+ denoted as ``b=(b0, b1, ..., bq)``. If size is not None, then it
+ must be compatible with the broadcast shape ``b``. Specifically,
+ size must have ``q`` or more elements and size[-(q-j):] must equal
+ ``bj``.
Returns
-------
out : ndarray
- The drawn samples, of shape *size*, if that was provided. If not,
- the shape is ``(N,)``.
+ The drawn samples, of shape size, if provided. When size is
+ provided, the output shape is size + (p,) If not specified,
+ the shape is determined by the broadcast shape of ``n`` and
+ ``pvals``, ``(b0, b1, ..., bq)`` augmented with the dimension of
+ the multinomial, ``p``, so that that output shape is
+ ``(b0, b1, ..., bq, p)``.
- In other words, each entry ``out[i,j,...,:]`` is an N-dimensional
- value drawn from the distribution.
+ Each entry ``out[i,j,...,:]`` is a ``p``-dimensional value drawn
+ from the distribution.
Examples
--------
@@ -3738,6 +3749,38 @@ cdef class Generator:
>>> rng.multinomial(100, [1/7.]*5 + [2/7.])
array([11, 16, 14, 17, 16, 26]) # random
+ Simulate 10 throws of a 4-sided die and 20 throws of a 6-sided die
+
+ >>> rng.multinomial([10, 20],[[1/4]*4 + [0]*2, [1/6]*6])
+ array([[2, 1, 4, 3, 0, 0],
+ [3, 3, 3, 6, 1, 4]], dtype=int64) # random
+
+ Generate categorical random variates from two categories where the
+ first has 3 outcomes and the second has 2.
+
+ >>> rng.multinomial(1, [[.1, .5, .4 ], [.3, .7, .0]])
+ array([[0, 0, 1],
+ [0, 1, 0]], dtype=int64) # random
+
+ ``argmax(axis=-1)`` is then used to return the categories.
+
+ >>> pvals = [[.1, .5, .4 ], [.3, .7, .0]]
+ >>> rvs = rng.multinomial(1, pvals, size=(4,2))
+ >>> rvs.argmax(axis=-1)
+ array([[0, 1],
+ [2, 0],
+ [2, 1],
+ [2, 0]], dtype=int64) # random
+
+ The same output dimension can be produced using broadcasting.
+
+ >>> rvs = rng.multinomial([[1]] * 4, pvals)
+ >>> rvs.argmax(axis=-1)
+ array([[0, 1],
+ [2, 0],
+ [2, 1],
+ [2, 0]], dtype=int64) # random
+
The probability inputs should be normalized. As an implementation
detail, the value of the last entry is ignored and assumed to take
up any leftover probability mass, but this should not be relied on.
@@ -3752,47 +3795,82 @@ cdef class Generator:
>>> rng.multinomial(100, [1.0, 2.0]) # WRONG
Traceback (most recent call last):
ValueError: pvals < 0, pvals > 1 or pvals contains NaNs
-
"""
- cdef np.npy_intp d, i, sz, offset
+ cdef np.npy_intp d, i, sz, offset, pi
cdef np.ndarray parr, mnarr, on, temp_arr
cdef double *pix
+ cdef int ndim
cdef int64_t *mnix
cdef int64_t ni
cdef np.broadcast it
+ on = <np.ndarray>np.PyArray_FROM_OTF(n,
+ np.NPY_INT64,
+ np.NPY_ARRAY_ALIGNED |
+ np.NPY_ARRAY_C_CONTIGUOUS)
+ parr = <np.ndarray>np.PyArray_FROM_OTF(pvals,
+ np.NPY_DOUBLE,
+ np.NPY_ARRAY_ALIGNED |
+ np.NPY_ARRAY_C_CONTIGUOUS)
+ ndim = parr.ndim
+ d = parr.shape[ndim - 1] if ndim >= 1 else 0
+ if d == 0:
+ raise ValueError(
+ "pvals must have at least 1 dimension and the last dimension "
+ "of pvals must be greater than 0."
+ )
- d = len(pvals)
- on = <np.ndarray>np.PyArray_FROM_OTF(n, np.NPY_INT64, np.NPY_ALIGNED)
- parr = <np.ndarray>np.PyArray_FROMANY(
- pvals, np.NPY_DOUBLE, 1, 1, np.NPY_ARRAY_ALIGNED | np.NPY_ARRAY_C_CONTIGUOUS)
- pix = <double*>np.PyArray_DATA(parr)
check_array_constraint(parr, 'pvals', CONS_BOUNDED_0_1)
- if kahan_sum(pix, d-1) > (1.0 + 1e-12):
- # When floating, but not float dtype, and close, improve the error
- # 1.0001 works for float16 and float32
- if (isinstance(pvals, np.ndarray)
- and np.issubdtype(pvals.dtype, np.floating)
- and pvals.dtype != float
- and pvals.sum() < 1.0001):
- msg = ("sum(pvals[:-1].astype(np.float64)) > 1.0. The pvals "
- "array is cast to 64-bit floating point prior to "
- "checking the sum. Precision changes when casting may "
- "cause problems even if the sum of the original pvals "
- "is valid.")
- else:
- msg = "sum(pvals[:-1]) > 1.0"
- raise ValueError(msg)
+ pix = <double*>np.PyArray_DATA(parr)
+ sz = np.PyArray_SIZE(parr)
+ # Cython 0.29.20 would not correctly translate the range-based for
+ # loop to a C for loop
+ # for offset in range(<np.npy_intp>0, sz, d):
+ offset = 0
+ while offset < sz:
+ if kahan_sum(pix + offset, d-1) > (1.0 + 1e-12):
+ # When floating, but not float dtype, and close, improve the error
+ # 1.0001 works for float16 and float32
+ slice_repr = "[:-1]" if ndim == 1 else "[...,:-1]"
+ if (isinstance(pvals, np.ndarray)
+ and np.issubdtype(pvals.dtype, np.floating)
+ and pvals.dtype != float
+ and pvals.sum() < 1.0001):
+ msg = (f"sum(pvals{slice_repr}.astype(np.float64)) > 1.0."
+ " The pvals array is cast to 64-bit floating"
+ " point prior to checking the sum. Precision "
+ "changes when casting may cause problems even "
+ "if the sum of the original pvals is valid.")
+ else:
+ msg = f"sum(pvals{slice_repr}) > 1.0"
+ raise ValueError(msg)
+ offset += d
- if np.PyArray_NDIM(on) != 0: # vector
+ if np.PyArray_NDIM(on) != 0 or ndim > 1: # vector
check_array_constraint(on, 'n', CONS_NON_NEGATIVE)
+ # This provides the offsets to use in the C-contig parr when
+ # broadcasting
+ offsets = <np.ndarray>np.arange(
+ 0, np.PyArray_SIZE(parr), d, dtype=np.intp
+ ).reshape((<object>parr).shape[:ndim - 1])
if size is None:
- it = np.PyArray_MultiIterNew1(on)
+ it = np.PyArray_MultiIterNew2(on, offsets)
else:
temp = np.empty(size, dtype=np.int8)
temp_arr = <np.ndarray>temp
- it = np.PyArray_MultiIterNew2(on, temp_arr)
- validate_output_shape(it.shape, temp_arr)
+ it = np.PyArray_MultiIterNew3(on, offsets, temp_arr)
+ # Validate size and the broadcast shape
+ try:
+ size = (operator.index(size),)
+ except:
+ size = tuple(size)
+ # This test verifies that an axis with dim 1 in size has not
+ # been increased by broadcasting with the input
+ if it.shape != size:
+ raise ValueError(
+ f"Output size {size} is not compatible with "
+ f"broadcast dimensions of inputs {it.shape}."
+ )
shape = it.shape + (d,)
multin = np.zeros(shape, dtype=np.int64)
mnarr = <np.ndarray>multin
@@ -3802,7 +3880,8 @@ cdef class Generator:
with self.lock, nogil:
for i in range(sz):
ni = (<int64_t*>np.PyArray_MultiIter_DATA(it, 0))[0]
- random_multinomial(&self._bitgen, ni, &mnix[offset], pix, d, &self._binomial)
+ pi = (<np.npy_intp*>np.PyArray_MultiIter_DATA(it, 1))[0]
+ random_multinomial(&self._bitgen, ni, &mnix[offset], &pix[pi], d, &self._binomial)
offset += d
np.PyArray_MultiIter_NEXT(it)
return multin
diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py
index d057122f1..e5411b8ef 100644
--- a/numpy/random/tests/test_generator_mt19937.py
+++ b/numpy/random/tests/test_generator_mt19937.py
@@ -136,12 +136,6 @@ class TestMultinomial:
contig = random.multinomial(100, pvals=np.ascontiguousarray(pvals))
assert_array_equal(non_contig, contig)
- def test_multidimensional_pvals(self):
- assert_raises(ValueError, random.multinomial, 10, [[0, 1]])
- assert_raises(ValueError, random.multinomial, 10, [[0], [1]])
- assert_raises(ValueError, random.multinomial, 10, [[[0], [1]], [[1], [0]]])
- assert_raises(ValueError, random.multinomial, 10, np.array([[0, 1], [1, 0]]))
-
def test_multinomial_pvals_float32(self):
x = np.array([9.9e-01, 9.9e-01, 1.0e-09, 1.0e-09, 1.0e-09, 1.0e-09,
1.0e-09, 1.0e-09, 1.0e-09, 1.0e-09], dtype=np.float32)
@@ -2361,6 +2355,64 @@ class TestBroadcast:
[2, 3, 6, 4, 2, 3]], dtype=np.int64)
assert_array_equal(actual, desired)
+ random = Generator(MT19937(self.seed))
+ actual = random.multinomial([5, 20], [[1 / 6.] * 6] * 2)
+ desired = np.array([[0, 0, 2, 1, 2, 0],
+ [2, 3, 6, 4, 2, 3]], dtype=np.int64)
+ assert_array_equal(actual, desired)
+
+ random = Generator(MT19937(self.seed))
+ actual = random.multinomial([[5], [20]], [[1 / 6.] * 6] * 2)
+ desired = np.array([[[0, 0, 2, 1, 2, 0],
+ [0, 0, 2, 1, 1, 1]],
+ [[4, 2, 3, 3, 5, 3],
+ [7, 2, 2, 1, 4, 4]]], dtype=np.int64)
+ assert_array_equal(actual, desired)
+
+ @pytest.mark.parametrize("n", [10,
+ np.array([10, 10]),
+ np.array([[[10]], [[10]]])
+ ]
+ )
+ def test_multinomial_pval_broadcast(self, n):
+ random = Generator(MT19937(self.seed))
+ pvals = np.array([1 / 4] * 4)
+ actual = random.multinomial(n, pvals)
+ n_shape = tuple() if isinstance(n, int) else n.shape
+ expected_shape = n_shape + (4,)
+ assert actual.shape == expected_shape
+ pvals = np.vstack([pvals, pvals])
+ actual = random.multinomial(n, pvals)
+ expected_shape = np.broadcast_shapes(n_shape, pvals.shape[:-1]) + (4,)
+ assert actual.shape == expected_shape
+
+ pvals = np.vstack([[pvals], [pvals]])
+ actual = random.multinomial(n, pvals)
+ expected_shape = np.broadcast_shapes(n_shape, pvals.shape[:-1])
+ assert actual.shape == expected_shape + (4,)
+ actual = random.multinomial(n, pvals, size=(3, 2) + expected_shape)
+ assert actual.shape == (3, 2) + expected_shape + (4,)
+
+ with pytest.raises(ValueError):
+ # Ensure that size is not broadcast
+ actual = random.multinomial(n, pvals, size=(1,) * 6)
+
+ def test_invalid_pvals_broadcast(self):
+ random = Generator(MT19937(self.seed))
+ pvals = [[1 / 6] * 6, [1 / 4] * 6]
+ assert_raises(ValueError, random.multinomial, 1, pvals)
+ assert_raises(ValueError, random.multinomial, 6, 0.5)
+
+ def test_empty_outputs(self):
+ random = Generator(MT19937(self.seed))
+ actual = random.multinomial(np.empty((10, 0, 6), "i8"), [1 / 6] * 6)
+ assert actual.shape == (10, 0, 6, 6)
+ actual = random.multinomial(12, np.empty((10, 0, 10)))
+ assert actual.shape == (10, 0, 10)
+ actual = random.multinomial(np.empty((3, 0, 7), "i8"),
+ np.empty((3, 0, 7, 4)))
+ assert actual.shape == (3, 0, 7, 4)
+
class TestThread:
# make sure each state produces the same sequence even in threads