8000 ENH: Add broadcasting to multinomial · bashtage/randomgen@06cd3b4 · GitHub
[go: up one dir, main page]
Skip to content

Commit 06cd3b4

Browse files
bashtagebashtage
committed
ENH: Add broadcasting to multinomial
Add broadcasting to multinomial xref numpy/numpy#9710
1 parent 4f57c42 commit 06cd3b4

File tree

7 files changed

+94
-40
lines changed

7 files changed

+94
-40
lines changed

‎doc/source/change-log.rst

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
Change Log
22
----------
33

4+
Since v1.16.2
5+
=============
6+
- Added broadcasting to multinomial (see
7+
`NumPy issue 9710 <https://github.com/numpy/numpy/pull/9710>`_)
8+
49
v1.16.2
510
=======
611
- Updated Xoroshiro120 to use AUthor's latest parameterization

‎randomgen/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

‎randomgen/generator.pyx

Lines changed: 46 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -3702,7 +3702,7 @@ cdef class RandomGenerator:
37023702
x.shape = tuple(final_shape)
37033703
return x
37043704

3705-
def multinomial(self, np.npy_intp n, object pvals, size=None):
3705+
def multinomial(self, object n, object pvals, size=None):
37063706
"""
37073707
multinomial(n, pvals, size=None)
37083708
@@ -3718,7 +3718,7 @@ cdef class RandomGenerator:
37183718
37193719
Parameters
37203720
----------
3721-
n : int
3721+
n : int or array-like of ints
37223722
Number of experiments.
37233723
pvals : sequence of floats, length p
37243724
Probabilities of each of the ``p`` different outcomes. These
@@ -3757,6 +3757,18 @@ cdef class RandomGenerator:
37573757
For the first run, we threw 3 times 1, 4 times 2, etc. For the second,
37583758
we threw 2 times 1, 4 times 2, etc.
37593759
3760+
Now, do one experiment throwing the dice 10 time, and 10 times again,
3761+
and another throwing the dice 20 times, and 20 times again:
3762+
3763+
>>> randomgen.generator.multinomial([[10], [20]], [1/6.]*6, size=2)
3764+
array([[[2, 4, 0, 1, 2, 1],
3765+
[1, 3, 0, 3, 1, 2]],
3766+
[[1, 4, 4, 4, 4, 3],
3767+
[3, 3, 2, 5, 5, 2]]]) # random
3768+
3769+
The first array shows the outcomes of throwing the dice 10 times, and
3770+
the second shows the outcomes from throwing the dice 20 times.
3771+
37603772
A loaded die is more likely to land on number 6:
37613773
37623774
>>> randomgen.generator.multinomial(100, [1/7.]*5 + [2/7.])
@@ -3777,19 +3789,43 @@ cdef class RandomGenerator:
37773789
array([100, 0])
37783790
37793791
"""
3780-
cdef np.npy_intp d, i, j, dn, sz
3781-
cdef np.ndarray parr "arrayObject_parr", mnarr "arrayObject_mnarr"
3792+
3793+
cdef np.npy_intp d, i, sz, offset
3794+
cdef np.ndarray parr, mnarr, on, temp_arr
37823795
cdef double *pix
37833796
cdef int64_t *mnix
3784-
cdef double Sum
3797+
cdef int64_t ni
3798+
cdef np.broadcast it
37853799

37863800
d = len(pvals)
3801+
on = <np.ndarray>np.PyArray_FROM_OTF(n, np.NPY_INT64, np.NPY_ALIGNED)
37873802
parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
37883803
pix = <double*>np.PyArray_DATA(parr)
37893804

37903805
if kahan_sum(pix, d-1) > (1.0 + 1e-12):
37913806
raise ValueError("sum(pvals[:-1]) > 1.0")
37923807

3808+
if np.PyArray_NDIM(on) != 0: # vector
3809+
if size is None:
3810+
it = np.PyArray_MultiIterNew1(on)
3811+
else:
3812+
temp = np.empty(size, dtype=np.int8)
3813+
temp_arr = <np.ndarray>temp
3814+
it = np.PyArray_MultiIterNew2(on, temp_arr)
3815+
shape = it.shape + (d,)
3816+
multin = np.zeros(shape, dtype=np.int64)
3817+
mnarr = <np.ndarray>multin
3818+
mnix = <int64_t*>np.PyArray_DATA(mnarr)
3819+
offset = 0
3820+
sz = it.size
3821+
with self.lock, nogil:
3822+
for i in range(sz):
3823+
ni = (<int64_t*>np.PyArray_MultiIter_DATA(it, 0))[0]
3824+
random_multinomial(self._brng, ni, &mnix[offset], pix, d, self._binomial)
3825+
offset += d
3826+
np.PyArray_MultiIter_NEXT(it)
3827+
return multin
3828+
37933829
if size is None:
37943830
shape = (d,)
37953831
else:
@@ -3802,23 +3838,12 @@ cdef class RandomGenerator:
38023838
mnarr = <np.ndarray>multin
38033839
mnix = <int64_t*>np.PyArray_DATA(mnarr)
38043840
sz = np.PyArray_SIZE(mnarr)
3805-
3841+
ni = n
3842+
offset = 0
38063843
with self.lock, nogil:
3807-
i = 0
3808-
while i < sz:
3809-
Sum = 1.0
3810-
dn = n
3811-
for j in range(d-1):
3812-
mnix[i+j] = random_binomial(self._brng, pix[j]/Sum, dn,
3813-
self._binomial)
3814-
dn = dn - mnix[i+j]
3815-
if dn <= 0:
3816-
break
3817-
Sum = Sum - pix[j]
3818-
if dn > 0:
3819-
mnix[i+d-1] = dn
3820-
3821-
i = i + d
3844+
for i in range(sz // d):
3845+
random_multinomial(self._brng, ni, &mnix[offset], pix, d, self._binomial)
3846+
offset += d
38223847

38233848
return multin
38243849

‎randomgen/mtrand.pyx

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

38023802
d = len(pvals)
38033803
parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
@@ -3818,23 +3818,12 @@ cdef class RandomState:
38183818
mnarr = <np.ndarray>multin
38193819
mnix = <int64_t*>np.PyArray_DATA(mnarr)
38203820
sz = np.PyArray_SIZE(mnarr)
3821-
3821+
ni = n
3822+
offset = 0
38223823
with self.lock, nogil:
3823-
i = 0
3824-
while i < sz:
3825-
Sum = 1.0
3826-
dn = n
3827-
for j in range(d-1):
3828-
mnix[i+j] = random_binomial(self._brng, pix[j]/Sum, dn,
3829-
self._binomial)
3830-
dn = dn - mnix[i+j]
3831-
if dn <= 0:
3832-
break
3833-
Sum = Sum - pix[j]
3834-
if dn > 0:
3835-
mnix[i+d-1] = dn
3836-
3837-
i = i + d
3824+
for i in range(sz // d):
3825+
random_multinomial(self._brng, ni, &mnix[offset], pix, d, self._binomial)
3826+
offset += d
38383827

38393828
return multin
38403829

‎randomgen/src/distributions/distributions.c

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1809,3 +1809,21 @@ void random_bounded_bool_fill(brng_t *brng_state, npy_bool off, npy_bool rng,
18091809
out[i] = buffered_bounded_bool(brng_state, off, rng, mask, &bcnt, &buf);
18101810
}
18111811
}
1812+
1813+
void random_multinomial(brng_t *brng_state, int64_t n, int64_t *mnix,
1814+
double *pix, npy_intp d, binomial_t *binomial) {
1815+
double remaining_p = 1.0;
1816+
npy_intp j;
1817+
int64_t dn = n;
1818+
for (j = 0; j < (d - 1); j++) {
1819+
mnix[j] = random_binomial(brng_state, pix[j] / remaining_p, dn, binomial);
1820+
dn = dn - mnix[j];
1821+
if (dn <= 0) {
1822+
break;
1823+
}
1824+
remaining_p -= pix[j];
1825+
if (dn > 0) {
1826+
mnix[d - 1] = dn;
1827+
}
1828+
}
1829+
}

‎randomgen/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

‎randomgen/tests/test_generator_mt19937.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1898,6 +1898,17 @@ def test_complex_normal(self):
18981898
size=(3, 2))
18991899
assert_array_almost_equal(actual, desired, decimal=15)
19001900

1901+
def test_multinomial(self):
1902+
random.seed(self.seed)
1903+
actual = random.multinomial([5, 20], [1 / 6.] * 6, size=(3, 2))
1904+
desired = np.array([[[1, 1, 1, 1, 0, 1],
1905+
[4, 5, 1, 4, 3, 3]],
1906+
[[1, 1, 1, 0, 0, 2],
1907+
[2, 0, 4, 3, 7, 4]],
1908+
[[1, 2, 0, 0, 2, 2],
1909+
[3, 2, 3, 4, 2, 6]]], dtype=np.int64)
1910+
assert_array_equal(actual, desired)
1911+
19011912

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

0 commit comments

Comments
 (0)
0