diff --git a/CHANGELOG.md b/CHANGELOG.md index b05a04e..cdfaa10 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,8 +29,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Sped up `randint` with `array_like` bounds by generating words in cache-sized chunks and by picking the Lemire rejection test by range width [gh-173](https://github.com/IntelPython/mkl_random/pull/173) * Reduced the memory use and sped up integer generation with `WH`, `MCG31`, `R250` and `MRG32K3A`, which lack `viRngUniformBits` support [gh-178](https://github.com/IntelPython/mkl_random/pull/178) * Replaced the `mklrand_py_helper.h` shim with `cpython.bytes` cimports and dropped the no-op `__future__` import, since the project requires Python 3.10+ [gh-192](https://github.com/IntelPython/mkl_random/pull/192) +* `geometric` samples are now one larger than before for every `p`, so results for a fixed seed change; `geometric(p) - 1` reproduces the old values. All other streams are unaffected [gh-196](https://github.com/IntelPython/mkl_random/pull/196) ### Fixed +* Fixed `geometric` returning the number of failures before the first success instead of the number of trials, which contradicted its own docstring and NumPy. The support now starts at `1`, `P(1)` is `p` and the mean is `1 / p` [gh-196](https://github.com/IntelPython/mkl_random/pull/196) * Fixed `multinormal_cholesky` and the other functions taking `method` routing integer-like ids such as `np.int64(0)`, `np.int64(2)` or `False` to `BoxMuller`; NumPy integers now select the method they name [gh-189](https://github.com/IntelPython/mkl_random/pull/189) * Fixed `multinormal_cholesky` raising `ZeroDivisionError` for an empty `mean`; it now raises `ValueError` [gh-189](https://github.com/IntelPython/mkl_random/pull/189) * Fixed an out-of-range integer `brng` indexing `brng_list` past its end, which returned uninitialized memory as random values; it now warns and falls back to `MT19937` [gh-177](https://github.com/IntelPython/mkl_random/pull/177) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index 81c6f0e..01adcb2 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -1306,6 +1306,10 @@ void irk_multinomial_vec(irk_state *state, void irk_geometric_vec(irk_state *state, npy_intp len, int *res, const double p) { int err = 0; + /* res and len advance in the chunk loop, so the +1 pass below + needs the whole range */ + npy_intp i = 0, total = len; + int *out = res; if (len < 1) return; @@ -1323,12 +1327,18 @@ void irk_geometric_vec(irk_state *state, npy_intp len, int *res, const double p) err = viRngGeometric(VSL_RNG_METHOD_GEOMETRIC_ICDF, state->stream, len, res, p); assert(err == VSL_STATUS_OK); + + /* MKL counts the failures, the number of trials also counts the + success that ended them; overflowed values, which MKL writes as + negative, are kept */ + DIST_PRAGMA_VECTOR + for (i = 0; i < total; ++i) + out[i] += (out[i] >= 0 && out[i] < INT_MAX) ? 1 : 0; } else { if (p == 1.0) { - npy_intp i; for (i = 0; i < len; ++i) - res[i] = 0; + res[i] = 1; } else { assert(p >= 0.0); diff --git a/mkl_random/tests/test_random.py b/mkl_random/tests/test_random.py index 13eb55b..3e48f7a 100644 --- a/mkl_random/tests/test_random.py +++ b/mkl_random/tests/test_random.py @@ -922,10 +922,27 @@ def test_randomdist_gamma(randomdist): def test_randomdsit_geometric(randomdist): rnd.seed(randomdist.seed, brng=randomdist.brng) actual = rnd.geometric(0.123456789, size=(3, 2)) - desired = np.array([[0, 30], [13, 2], [4, 0]]) + desired = np.array([[1, 31], [14, 3], [5, 1]]) np.testing.assert_array_equal(actual, desired) +def test_geometric_counts_trials(): + # The number of trials, as the docstring and NumPy define it, so the + # support starts at one, P(1) is p and the mean is 1 / p + p = 0.35 + samples = rnd.MKLRandomState(1).geometric(p, 100_000) + assert samples.min() == 1 + np.testing.assert_allclose((samples == 1).mean(), p, atol=0.01) + np.testing.assert_allclose(samples.mean(), 1 / p, rtol=0.01) + + ones = np.ones(5, dtype="int32") + np.testing.assert_array_equal(rnd.geometric(1.0, 5), ones, strict=True) + np.testing.assert_array_equal(rnd.geometric(np.ones(5)), ones, strict=True) + + # MKL marks draws that overflow int with INT_MIN, which must survive the +1 + assert (rnd.geometric(1e-300, 3) == np.iinfo(np.int32).min).all() + + def test_randomdist_gumbel(randomdist): rnd.seed(randomdist.seed, randomdist.brng) actual = rnd.gumbel(loc=0.123456789, scale=2.0, size=(3, 2))