Merge pull request #196 from IntelPython/fix/geometric-numpy-convention

Count trials in geometric, as the docstring and NumPy do

Author: vlad-perevezentsev

CHANGELOG.md

@@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * 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)
 * `negative_binomial` now rejects `p == 0`, which MKL does not support, with a `ValueError`; the documented range of `p` is `(0, 1]` [gh-195](https://github.com/IntelPython/mkl_random/pull/195)
 * `poisson`, `negative_binomial`, `geometric` and `logseries` now reject `NaN` parameters with a `ValueError`, as NumPy does, instead of passing them to MKL [gh-195](https://github.com/IntelPython/mkl_random/pull/195)
+* `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 `poisson` with a scalar `lam == 0` returning uninitialized memory as random values, since MKL requires `lam > 0`. The same applied to array-valued `lam` with `method="PTPE"`. Both now return zeros, as the default array path and NumPy already did [gh-195](https://github.com/IntelPython/mkl_random/pull/195)
@@ -56,6 +57,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Fixed the `tomaxint` example, which compared against the Python 2 `sys.maxint` [gh-192](https://github.com/IntelPython/mkl_random/pull/192)
 * Fixed the remaining Python 2 references in the `random_sample`, `random_integers` and `randint` docstrings and in the `zipf` implementation comment: the `<type '...'>` reprs, `np.int` and `sys.maxint` [gh-192](https://github.com/IntelPython/mkl_random/pull/192)
 * Fixed the `random_integers` docstring, which described the result as `numpy.int_` rather than a C `long`, documented the wrong scalar return type, and passed a float `size` that no longer works [gh-192](https://github.com/IntelPython/mkl_random/pull/192)
+* 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)
 
 ### Removed
 * Removed the `python-gil` constraint from the conda recipes, which pinned `mkl_random` to GIL-enabled Python 3.14 builds [gh-159](https://github.com/IntelPython/mkl_random/pull/159)

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);

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_geometric_exceptions():
     for p in (0.0, -0.5, 1.5, np.nan):
         assert_raises(ValueError, rnd.geometric, p)