Skip to content

Commit bbdf80d

Browse files
bashtagebashtage
committed
ENH: Extend multinomial
Extend multinomial to allow broadcasting
1 parent 469bf4c commit bbdf80d

File tree

6 files changed

+97
-40
lines changed

6 files changed

+97
-40
lines changed

‎numpy/random/distributions.pxd

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,3 +144,6 @@ cdef extern from "src/distributions/distributions.h":
144144
np.npy_bool off, np.npy_bool rng, np.npy_intp cnt,
145145
bint use_masked,
146146
np.npy_bool *out) nogil
147+
148+
void random_multinomial(brng_t *brng_state, int64_t n, int64_t *mnix,
149+
double *pix, np.npy_intp d, binomial_t *binomial) nogil

‎numpy/random/generator.pyx

Lines changed: 45 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3655,7 +3655,7 @@ cdef class RandomGenerator:
36553655
36563656
Parameters
36573657
----------
3658-
n : int
3658+
n : int or array-like of ints
36593659
Number of experiments.
36603660
pvals : sequence of floats, length p
36613661
Probabilities of each of the ``p`` different outcomes. These
@@ -3694,6 +3694,18 @@ cdef class RandomGenerator:
36943694
For the first run, we threw 3 times 1, 4 times 2, etc. For the second,
36953695
we threw 2 times 1, 4 times 2, etc.
36963696
3697+
Now, do one experiment throwing the dice 10 time, and 10 times again,
3698+
and another throwing the dice 20 times, and 20 times again:
3699+
3700+
>>> np.random.multinomial([[10], [20]], [1/6.]*6, size=2)
3701+
array([[[2, 4, 0, 1, 2, 1],
3702+
[1, 3, 0, 3, 1, 2]],
3703+
[[1, 4, 4, 4, 4, 3],
3704+
[3, 3, 2, 5, 5, 2]]]) # random
3705+
3706+
The first array shows the outcomes of throwing the dice 10 times, and
3707+
the second shows the outcomes from throwing the dice 20 times.
3708+
36973709
A loaded die is more likely to land on number 6:
36983710
36993711
>>> np.random.multinomial(100, [1/7.]*5 + [2/7.])
@@ -3714,19 +3726,43 @@ cdef class RandomGenerator:
37143726
array([100, 0])
37153727
37163728
"""
3717-
cdef np.npy_intp d, i, j, dn, sz
3718-
cdef np.ndarray parr "arrayObject_parr", mnarr "arrayObject_mnarr"
3729+
3730+
cdef np.npy_intp d, i, sz, offset
3731+
cdef np.ndarray parr, mnarr, on, temp_arr
37193732
cdef double *pix
37203733
cdef int64_t *mnix
3721-
cdef double Sum
3734+
cdef int64_t ni
3735+
cdef np.broadcast it
37223736

37233737
d = len(pvals)
3738+
on = <np.ndarray>np.PyArray_FROM_OTF(n, np.NPY_INT64, np.NPY_ALIGNED)
37243739
parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
37253740
pix = <double*>np.PyArray_DATA(parr)
37263741

37273742
if kahan_sum(pix, d-1) > (1.0 + 1e-12):
37283743
raise ValueError("sum(pvals[:-1]) > 1.0")
37293744

3745+
if np.PyArray_NDIM(on) != 0: # vector
3746+
if size is None:
3747+
it = np.PyArray_MultiIterNew1(on)
3748+
else:
3749+
temp = np.empty(size, dtype=np.int8)
3750+
temp_arr = <np.ndarray>temp
3751+
it = np.PyArray_MultiIterNew2(on, temp_arr)
3752+
shape = it.shape + (d,)
3753+
multin = np.zeros(shape, dtype=np.int64)
3754+
mnarr = <np.ndarray>multin
3755+
mnix = <int64_t*>np.PyArray_DATA(mnarr)
3756+
offset = 0
3757+
sz = it.size
3758+
with self.lock, nogil:
3759+
for i in range(sz):
3760+
ni = (<int64_t*>np.PyArray_MultiIter_DATA(it, 0))[0]
3761+
random_multinomial(self._brng, ni, &mnix[offset], pix, d, self._binomial)
3762+
offset += d
3763+
np.PyArray_MultiIter_NEXT(it)
3764+
return multin
3765+
37303766
if size is None:
37313767
shape = (d,)
37323768
else:
@@ -3739,23 +3775,12 @@ cdef class RandomGenerator:
37393775
mnarr = <np.ndarray>multin
37403776
mnix = <int64_t*>np.PyArray_DATA(mnarr)
37413777
sz = np.PyArray_SIZE(mnarr)
3742-
3778+
ni = n
3779+
offset = 0
37433780
with self.lock, nogil:
3744-
i = 0
3745-
while i < sz:
3746-
Sum = 1.0
3747-
dn = n
3748-
for j in range(d-1):
3749-
mnix[i+j] = random_binomial(self._brng, pix[j]/Sum, dn,
3750-
self._binomial)
3751-
dn = dn - mnix[i+j]
3752-
if dn <= 0:
3753-
break
3754-
Sum = Sum - pix[j]
3755-
if dn > 0:
3756-
mnix[i+d-1] = dn
3757-
3758-
i = i + d
3781+
for i in range(sz // d):
3782+
random_multinomial(self._brng, ni, &mnix[offset], pix, d, self._binomial)
3783+
offset += d
37593784

37603785
return multin
37613786

‎numpy/random/mtrand.pyx

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3790,11 +3790,11 @@ cdef class RandomState:
37903790
array([100, 0])
37913791
37923792
"""
3793-
cdef np.npy_intp d, i, j, dn, sz
3794-
cdef np.ndarray parr "arrayObject_parr", mnarr "arrayObject_mnarr"
3793+
cdef np.npy_intp d, i, sz, offset
3794+
cdef np.ndarray parr, mnarr
37953795
cdef double *pix
37963796
cdef int64_t *mnix
3797-
cdef double Sum
3797+
cdef int64_t ni
37983798

37993799
d = len(pvals)
38003800
parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
@@ -3815,23 +3815,12 @@ cdef class RandomState:
38153815
mnarr = <np.ndarray>multin
38163816
mnix = <int64_t*>np.PyArray_DATA(mnarr)
38173817
sz = np.PyArray_SIZE(mnarr)
3818-
3818+
ni = n
3819+
offset = 0
38193820
with self.lock, nogil:
3820-
i = 0
3821-
while i < sz:
3822-
Sum = 1.0
3823-
dn = n
3824-
for j in range(d-1):
3825-
mnix[i+j] = random_binomial(self._brng, pix[j]/Sum, dn,
3826-
self._binomial)
3827-
dn = dn - mnix[i+j]
3828-
if dn <= 0:
3829-
break
3830-
Sum = Sum - pix[j]
3831-
if dn > 0:
3832-
mnix[i+d-1] = dn
3833-
3834-
i = i + d
3821+
for i in range(sz // d):
3822+
random_multinomial(self._brng, ni, &mnix[offset], pix, d, self._binomial)
3823+
offset += d
38353824

38363825
return multin
38373826

‎numpy/random/src/distributions/distributions.c

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1179,8 +1179,10 @@ int64_t random_hypergeometric(brng_t *brng_state, int64_t good, int64_t bad,
11791179
int64_t sample) {
11801180
if (sample > 10) {
11811181
return random_hypergeometric_hrua(brng_state, good, bad, sample);
1182-
} else {
1182+
} else if (sample > 0) {
11831183
return random_hypergeometric_hyp(brng_state, good, bad, sample);
1184+
} else {
1185+
return 0;
11841186
}
11851187
}
11861188

@@ -1809,3 +1811,21 @@ void random_bounded_bool_fill(brng_t *brng_state, npy_bool off, npy_bool rng,
18091811
out[i] = buffered_bounded_bool(brng_state, off, rng, mask, &bcnt, &buf);
18101812
}
18111813
}
1814+
1815+
void random_multinomial(brng_t *brng_state, int64_t n, int64_t *mnix,
1816+
double *pix, npy_intp d, binomial_t *binomial) {
1817+
double remaining_p = 1.0;
1818+
npy_intp j;
1819+
int64_t dn = n;
1820+
for (j = 0; j < (d - 1); j++) {
1821+
mnix[j] = random_binomial(brng_state, pix[j] / remaining_p, dn, binomial);
1822+
dn = dn - mnix[j];
1823+
if (dn <= 0) {
1824+
break;
1825+
}
1826+
remaining_p -= pix[j];
1827+
if (dn > 0) {
1828+
mnix[d - 1] = dn;
1829+
}
1830+
}
1831+
}

‎numpy/random/src/distributions/distributions.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,4 +217,7 @@ DECLDIR void random_bounded_bool_fill(brng_t *brng_state, npy_bool off,
217217
npy_bool rng, npy_intp cnt,
218218
bool use_masked, npy_bool *out);
219219

220+
DECLDIR void random_multinomial(brng_t *brng_state, int64_t n, int64_t *mnix,
221+
double *pix, npy_intp d, binomial_t *binomial);
222+
220223
#endif

‎numpy/random/tests/test_generator_mt19937.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1849,6 +1849,23 @@ def test_logseries(self):
18491849
assert_raises(ValueError, logseries, bad_p_one * 3)
18501850
assert_raises(ValueError, logseries, bad_p_two * 3)
18511851

1852+
def test_multinomial(self):
1853+
random.seed(self.seed)
1854+
actual = random.multinomial([5, 20], [1 / 6.] * 6, size=(3, 2))
1855+
desired = np.array([[[1, 1, 1, 1, 0, 1],
1856+
[4, 5, 1, 4, 3, 3]],
1857+
[[1, 1, 1, 0, 0, 2],
1858+
[2, 0, 4, 3, 7, 4]],
1859+
[[1, 2, 0, 0, 2, 2],
1860+
[3, 2, 3, 4, 2, 6]]], dtype=np.int64)
1861+
assert_array_equal(actual, desired)
1862+
1863+
random.seed(self.seed)
1864+
actual = random.multinomial([5, 20], [1 / 6.] * 6)
1865+
desired = np.array([[1, 1, 1, 1, 0, 1],
1866+
[4, 5, 1, 4, 3, 3]], dtype=np.int64)
1867+
assert_array_equal(actual, desired)
1868+
18521869

18531870
class TestThread(object):
18541871
# make sure each state produces the same sequence even in threads

0 commit comments

Comments
 (0)