Count trials in geometric, as the docstring and NumPy do

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)

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_randomdist_gumbel(randomdist):
     rnd.seed(randomdist.seed, randomdist.brng)
     actual = rnd.gumbel(loc=0.123456789, scale=2.0, size=(3, 2))