perf: batch array-valued distribution parameters into a single fill

CHANGELOG.md

@@ -7,8 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 # [dev] (MM/DD/YYYY)
 
 ### Added
+* Added tests for the array-valued parameter paths of the location and scale distributions [gh-171](https://github.com/IntelPython/mkl_random/pull/171)
 
 ### Changed
+* Sped up `normal`, `uniform`, `exponential`, `laplace`, `gumbel`, `logistic`, `rayleigh` and `lognormal` for array-valued parameters [gh-171](https://github.com/IntelPython/mkl_random/pull/171)
+* The random streams for the array-valued-parameter paths of the distributions above have changed: with a fixed seed these now produce different (but equally valid) samples. Scalar-parameter paths are unaffected. [gh-171](https://github.com/IntelPython/mkl_random/pull/171)
 * Pinned Cython in the Coverity Scan workflow so generated code stays stable between scans, and added `coverity/README.md` documenting the known Cython-boilerplate false positives and the scan review checklist [gh-164](https://github.com/IntelPython/mkl_random/pull/164)
 
 ### Fixed

mkl_random/mklrand.pyx

@@ -698,6 +698,149 @@ cdef object vec_cont2_array(
     return arr_obj
 
 
+cdef object _param_out_shape(object size, tuple param_shapes):
+    """Result shape for a parameterised draw, matching the per-element paths."""
+    cdef object out_shape
+    cdef object bshape
+
+    if size is None:
+        return np.broadcast_shapes(*param_shapes)
+
+    out_shape = tuple(size) if np.iterable(size) else (size,)
+    try:
+        bshape = np.broadcast_shapes(out_shape, *param_shapes)
+    except ValueError:
+        raise ValueError("size is not compatible with inputs")
+    if bshape != out_shape:
+        raise ValueError("size is not compatible with inputs")
+    return out_shape
+
+
+cdef object _fill_standard2(
+    irk_state *state,
+    irk_cont2_vec func,
+    object out_shape,
+    object lock
+):
+    """Fill an entire request with one call, using standard parameters."""
+    cdef cnp.ndarray array
+    cdef cnp.npy_intp n
+    cdef double *array_data
+
+    array = <cnp.ndarray>np.empty(out_shape, np.float64)
+    n = cnp.PyArray_SIZE(array)
+    if n:
+        array_data = <double *>cnp.PyArray_DATA(array)
+        with lock, nogil:
+            func(state, n, array_data, 0.0, 1.0)
+    return array
+
+
+cdef object vec_loc_scale_array(
+    irk_state *state,
+    irk_cont2_vec func,
+    object size,
+    cnp.ndarray oloc,
+    cnp.ndarray oscale,
+    object lock
+):
+    """Draw a location and scale family with array-valued parameters.
+
+    ``func(0.0, 1.0)`` yields the standardised member, so
+    ``loc + scale * standardised`` is exact and needs one call per request.
+    """
+    cdef object array
+
+    array = _fill_standard2(
+        state,
+        func,
+        _param_out_shape(
+            size, ((<object>oloc).shape, (<object>oscale).shape)
+        ),
+        lock
+    )
+    np.multiply(array, oscale, out=array)
+    np.add(array, oloc, out=array)
+    return array
+
+
+cdef object vec_scale_array(
+    irk_state *state,
+    irk_cont1_vec func,
+    object size,
+    cnp.ndarray oscale,
+    object lock
+):
+    """Draw a scale family with an array-valued scale, one call per request."""
+    cdef cnp.ndarray array
+    cdef cnp.npy_intp n
+    cdef double *array_data
+
+    array = <cnp.ndarray>np.empty(
+        _param_out_shape(size, ((<object>oscale).shape,)), np.float64
+    )
+    n = cnp.PyArray_SIZE(array)
+    if n:
+        array_data = <double *>cnp.PyArray_DATA(array)
+        with lock, nogil:
+            func(state, n, array_data, 1.0)
+    np.multiply(array, oscale, out=array)
+    return array
+
+
+cdef object vec_uniform_array(
+    irk_state *state,
+    irk_cont2_vec func,
+    object size,
+    cnp.ndarray olow,
+    cnp.ndarray ohigh,
+    object lock
+):
+    """Draw uniforms over array-valued bounds, one call per request."""
+    cdef object array
+
+    array = _fill_standard2(
+        state,
+        func,
+        _param_out_shape(
+            size, ((<object>olow).shape, (<object>ohigh).shape)
+        ),
+        lock
+    )
+    np.multiply(array, np.subtract(ohigh, olow), out=array)
+    np.add(array, olow, out=array)
+    return array
+
+
+cdef object vec_lognormal_array(
+    irk_state *state,
+    irk_cont2_vec normal_func,
+    object size,
+    cnp.ndarray omean,
+    cnp.ndarray osigma,
+    object lock
+):
+    """Draw lognormals with array-valued parameters, one call per request.
+
+    Uses the normal fill: the parameters sit inside the exponential, so no
+    affine step applies to a standardised lognormal, but exp(mean + sigma * z) does.
+    """
+    cdef object array
+
+    array = _fill_standard2(
+        state,
+        normal_func,
+        _param_out_shape(
+            size, ((<object>omean).shape, (<object>osigma).shape)
+        ),
+        lock
+    )
+    np.multiply(array, osigma, out=array)
+    np.add(array, omean, out=array)
+    np.exp(array, out=array)
+    return array
+
+
 cdef object vec_cont3_array_sc(
     irk_state *state,
     irk_cont3_vec func,
@@ -2548,7 +2691,7 @@ cdef class _MKLRandomState:
         if np.any(olow >= ohigh):
             raise ValueError("low >= high")
 
-        return vec_cont2_array(
+        return vec_uniform_array(
             self.internal_state, irk_uniform_vec, size, olow, ohigh, self.lock
         )
 
@@ -2950,23 +3093,23 @@ cdef class _MKLRandomState:
             method, [ICDF, BOXMULLER, BOXMULLER2], _method_alias_dict_gaussian
         )
         if method is ICDF:
-            return vec_cont2_array(
+            return vec_loc_scale_array(
                 self.internal_state,
                 irk_normal_vec_ICDF,
                 size,
                 oloc,
                 oscale, self.lock
             )
         elif method is BOXMULLER2:
-            return vec_cont2_array(
+            return vec_loc_scale_array(
                 self.internal_state,
                 irk_normal_vec_BM2,
                 size,
                 oloc,
                 oscale, self.lock
             )
         else:
-            return vec_cont2_array(
+            return vec_loc_scale_array(
                 self.internal_state,
                 irk_normal_vec_BM1,
                 size,
@@ -3105,7 +3248,7 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | (oscale == 0)):
             raise ValueError("scale <= 0")
-        return vec_cont1_array(
+        return vec_scale_array(
             self.internal_state, irk_exponential_vec, size, oscale, self.lock
         )
 
@@ -4550,8 +4693,9 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)):
             raise ValueError("scale <= 0")
-        return vec_cont2_array(
-            self.internal_state, irk_laplace_vec, size, oloc, oscale, self.lock
+        return vec_loc_scale_array(
+            self.internal_state, irk_laplace_vec, size, oloc, oscale,
+            self.lock
         )
 
     def gumbel(self, loc=0.0, scale=1.0, size=None):
@@ -4690,8 +4834,9 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)):
             raise ValueError("scale <= 0")
-        return vec_cont2_array(
-            self.internal_state, irk_gumbel_vec, size, oloc, oscale, self.lock
+        return vec_loc_scale_array(
+            self.internal_state, irk_gumbel_vec, size, oloc, oscale,
+            self.lock
         )
 
     def logistic(self, loc=0.0, scale=1.0, size=None):
@@ -4791,7 +4936,7 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)):
             raise ValueError("scale <= 0")
-        return vec_cont2_array(
+        return vec_loc_scale_array(
             self.internal_state,
             irk_logistic_vec,
             size,
@@ -4952,18 +5097,18 @@ cdef class _MKLRandomState:
             method, [ICDF, BOXMULLER], _method_alias_dict_gaussian_short
         )
         if method is ICDF:
-            return vec_cont2_array(
+            return vec_lognormal_array(
                 self.internal_state,
-                irk_lognormal_vec_ICDF,
+                irk_normal_vec_ICDF,
                 size,
                 omean,
                 osigma,
                 self.lock
             )
         else:
-            return vec_cont2_array(
+            return vec_lognormal_array(
                 self.internal_state,
-                irk_lognormal_vec_BM,
+                irk_normal_vec_BM2,
                 size,
                 omean,
                 osigma,
@@ -5045,7 +5190,7 @@ cdef class _MKLRandomState:
 
         if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)):
             raise ValueError("scale <= 0.0")
-        return vec_cont1_array(
+        return vec_scale_array(
             self.internal_state, irk_rayleigh_vec, size, oscale, self.lock
         )
 

mkl_random/tests/test_random.py

@@ -1163,6 +1163,95 @@ def test_uniform_array_bounds_return_ndarray():
     assert arr.shape == (2,)
 
 
+_LOC_SCALE_DISTS = [
+    ("normal", lambda r, a, b, s: r.normal(a, b, s), 2.0, 3.0),
+    ("laplace", lambda r, a, b, s: r.laplace(a, b, s), 2.0, 3.0),
+    ("gumbel", lambda r, a, b, s: r.gumbel(a, b, s), 2.0, 3.0),
+    ("logistic", lambda r, a, b, s: r.logistic(a, b, s), 2.0, 3.0),
+    ("lognormal", lambda r, a, b, s: r.lognormal(a, b, s), 0.5, 0.75),
+    ("uniform", lambda r, a, b, s: r.uniform(a, b, s), 2.0, 5.0),
+]
+
+
+@pytest.mark.parametrize(
+    "name,draw,pa,pb", _LOC_SCALE_DISTS, ids=[d[0] for d in _LOC_SCALE_DISTS]
+)
+def test_two_param_array_matches_scalar(name, draw, pa, pb):
+    # Constant-valued arrays must agree with the scalar path.
+    n = 8192
+    scalar = draw(rnd.MKLRandomState(1234), pa, pb, n)
+    arrayed = draw(
+        rnd.MKLRandomState(1234), np.full(n, pa), np.full(n, pb), None
+    )
+    assert arrayed.shape == scalar.shape
+    np.testing.assert_allclose(
+        arrayed,
+        scalar,
+        rtol=1e-9,
+        atol=1e-9 * float(np.std(scalar)),
+        err_msg=f"{name}: array-parameter path disagrees with scalar path",
+    )
+
+
+@pytest.mark.parametrize(
+    "name,draw,pa,pb", _LOC_SCALE_DISTS, ids=[d[0] for d in _LOC_SCALE_DISTS]
+)
+def test_two_param_array_applies_per_element(name, draw, pa, pb):
+    # A scale sweep must widen the spread across the result.
+    n = 60000
+    lo = np.full(n, pa)
+    hi = np.linspace(pb, pb * 4.0, n)
+    out = draw(rnd.MKLRandomState(99), lo, hi, None)
+    first, last = out[: n // 4], out[-n // 4 :]
+    assert np.std(last) > np.std(first), (
+        f"{name}: per-element parameters do not appear to be applied"
+    )
+
+
+@pytest.mark.parametrize(
+    "name,draw,p",
+    [
+        ("exponential", lambda r, a, s: r.exponential(a, s), 3.0),
+        ("rayleigh", lambda r, a, s: r.rayleigh(a, s), 3.0),
+    ],
+    ids=["exponential", "rayleigh"],
+)
+def test_one_param_array_matches_scalar(name, draw, p):
+    n = 8192
+    scalar = draw(rnd.MKLRandomState(1234), p, n)
+    arrayed = draw(rnd.MKLRandomState(1234), np.full(n, p), None)
+    assert arrayed.shape == scalar.shape
+    np.testing.assert_allclose(
+        arrayed,
+        scalar,
+        rtol=1e-9,
+        atol=1e-9 * float(np.std(scalar)),
+        err_msg=f"{name}: array-parameter path disagrees with scalar path",
+    )
+
+
+@pytest.mark.parametrize(
+    "loc_shape,scale_shape,size,expected",
+    [
+        ((7,), (), None, (7,)),
+        ((), (7,), None, (7,)),
+        ((7,), (7,), None, (7,)),
+        ((3, 1), (4,), None, (3, 4)),
+        ((4,), (4,), (3, 4), (3, 4)),
+        ((7,), (7,), 7, (7,)),
+    ],
+)
+def test_two_param_array_broadcast_shapes(loc_shape, scale_shape, size, expected):
+    loc = np.zeros(loc_shape) if loc_shape else 0.0
+    scale = np.ones(scale_shape) if scale_shape else 1.0
+    assert rnd.MKLRandomState(5).normal(loc, scale, size).shape == expected
+
+
+def test_two_param_array_size_incompatible():
+    with pytest.raises(ValueError):
+        rnd.MKLRandomState(5).normal(np.zeros(5), np.ones(5), 3)
+
+
 def test_randomdist_vonmises(randomdist):
     rnd.seed(randomdist.seed, brng=randomdist.brng)
     actual = rnd.vonmises(mu=1.23, kappa=1.54, size=(3, 2))