TST: Add tests for broadcasting

randomgen/tests/_shims.pyx

@@ -5,8 +5,16 @@ from randomgen.common cimport (
     view_little_endian,
 )
 
+from randomgen.chacha import ChaCha
+
+from numpy.random cimport bitgen_t
+
 import numpy as np
 
+from cpython.pycapsule cimport PyCapsule_GetPointer, PyCapsule_IsValid
+
+from randomgen.broadcasting cimport constraint_type, cont
+
 
 def view_little_endian_shim(arr, dtype):
     return view_little_endian(arr, dtype)
@@ -21,3 +29,61 @@ def byteswap_little_endian_shim(arr):
 
 def object_to_int_shim(val, bits, name, default_bits=64, allowed_sizes=(64,)):
     return object_to_int(val, bits, name, default_bits, allowed_sizes)
+
+cdef double double0_func(void *state):
+    return 3.141592
+
+cdef double double1_func(void *state, double a):
+    return a
+
+cdef double double2_func(void *state, double a, double b):
+    return a+b
+
+cdef double double3_func(void *state, double a, double b, double c):
+    return a+b+c
+
+cdef bitgen_t _bitgen
+cdef const char *name = "BitGenerator"
+
+chacha = ChaCha()
+capsule = chacha.capsule
+_bitgen = (<bitgen_t *> PyCapsule_GetPointer(capsule, name))[0]
+
+def cont_0(size=None, out=None):
+    return cont(&double0_func, &_bitgen ,size, chacha.lock, 0,
+               0.0, "", constraint_type.CONS_NONE,
+               0.0, "", constraint_type.CONS_NONE,
+               0.0, "", constraint_type.CONS_NONE,
+               out)
+def cont_1(a, size=None, out=None):
+    return cont(&double1_func, &_bitgen ,size, chacha.lock, 1,
+               a, "a", constraint_type.CONS_NON_NEGATIVE,
+               0.0, "", constraint_type.CONS_NONE,
+               0.0, "", constraint_type.CONS_NONE,
+               out)
+
+def cont_2(a, b, size=None, out=None):
+    return cont(&double2_func, &_bitgen ,size, chacha.lock, 2,
+               a, "a", constraint_type.CONS_POSITIVE,
+               b, "b", constraint_type.CONS_POSITIVE_NOT_NAN,
+               0.0, "", constraint_type.CONS_NONE,
+               out)
+
+def cont_3(a, b, c, size=None, out=None):
+    return cont(&double3_func, &_bitgen ,size, chacha.lock, 3,
+               a, "a", constraint_type.CONS_BOUNDED_0_1,
+               b, "b", constraint_type.CONS_BOUNDED_GT_0_1,
+               c, "c", constraint_type.CONS_GT_1,
+               out)
+
+def cont_3_alt_cons(a, b, c, size=None, out=None):
+    return cont(&double3_func, &_bitgen, size, chacha.lock, 3,
+                a, "a", constraint_type.CONS_GTE_1,
+                b, "b", constraint_type.CONS_POISSON,
+                c, "c", constraint_type.LEGACY_CONS_POISSON,
+                out)
+
+# Iterations
+# Scalar or vector parameters (must be broadcastable)
+# size
+# out, if used with size, and parameters, must be compatible

randomgen/tests/test_broadcasting.py

@@ -0,0 +1,124 @@
+from collections import defaultdict
+from typing import NamedTuple
+
+import numpy as np
+from numpy.testing import assert_allclose
+import pytest
+
+from randomgen.tests._shims import cont_0, cont_1, cont_2, cont_3, cont_3_alt_cons
+
+
+class Config(NamedTuple):
+    a: float | np.ndarray | None
+    b: float | np.ndarray | None
+    c: float | np.ndarray | None
+    size: tuple[int, ...] | None
+    out: np.ndarray | None
+
+
+CONFIGS = defaultdict(list)
+
+
+def all_scalar(*args):
+    return all([(arg is None or np.isscalar(arg)) for arg in args])
+
+
+def get_broadcastable_size(a, b, c):
+    if c is not None:
+        return (a + b + c).shape
+    if b is not None:
+        return (a + b).shape
+    if a is not None:
+        return a.shape
+
+
+def count_params(a, b, c):
+    return sum((v is not None) for v in (a, b, c))
+
+
+for a in (None, 0.5, 0.5 * np.ones((1, 2)), 0.5 * np.ones((3, 2))):
+    for b in (
+        None,
+        0.5,
+        0.5 * np.ones((1, 2)),
+        0.5 * np.ones((3, 1)),
+        0.5 * np.ones((3, 2)),
+    ):
+        if a is None and b is not None:
+            continue
+        for c in (
+            None,
+            1.5,
+            1.5 * np.ones((1, 2)),
+            1.5 * np.ones((3, 1)),
+            1.5 * np.ones((3, 2)),
+        ):
+            if b is None and c is not None:
+                continue
+            for size in (True, False):
+                for out in (True, False):
+                    if size:
+                        if all_scalar(a, b, c):
+                            _size = (7, 5)
+                        else:
+                            _size = get_broadcastable_size(a, b, c)
+                    else:
+                        _size = None
+                    if out:
+                        if size:
+                            _out = np.empty(_size)
+                        elif all_scalar(a, b, c):
+                            _out = np.empty((11, 7))
+                        else:
+                            _out = np.empty(get_broadcastable_size(a, b, c))
+                    else:
+                        _out = None
+                    print(_size, _out.shape if isinstance(_out, np.ndarray) else _out)
+                    CONFIGS[count_params(a, b, c)].append(Config(a, b, c, _size, _out))
+
+
+@pytest.mark.parametrize("config", CONFIGS[0])
+def test_cont_0(config):
+    res = cont_0(size=config.size, out=config.out)
+    if isinstance(res, np.ndarray):
+        assert_allclose(res, 3.141592 * np.ones_like(res))
+    else:
+        assert_allclose(res, 3.141592)
+
+
+@pytest.mark.parametrize("config", CONFIGS[1])
+def test_cont_1(config):
+    res = cont_1(config.a, size=config.size, out=config.out)
+    if isinstance(res, np.ndarray):
+        assert_allclose(res, 0.5 * np.ones_like(res))
+    else:
+        assert_allclose(res, 0.5)
+
+
+@pytest.mark.parametrize("config", CONFIGS[2])
+def test_cont_2(config):
+    res = cont_2(config.a, config.b, size=config.size, out=config.out)
+    if isinstance(res, np.ndarray):
+        assert_allclose(res, np.ones_like(res))
+    else:
+        assert_allclose(res, 1.0)
+
+
+@pytest.mark.parametrize("config", CONFIGS[3])
+def test_cont_3(config):
+    res = cont_3(config.a, config.b, config.c, size=config.size, out=config.out)
+    if isinstance(res, np.ndarray):
+        assert_allclose(res, 2.5 * np.ones_like(res))
+    else:
+        assert_allclose(res, 2.5)
+
+
+@pytest.mark.parametrize("config", CONFIGS[3])
+def test_cont_3_alt_cons(config):
+    res = cont_3_alt_cons(
+        1.0 + config.a, config.b, config.c, size=config.size, out=config.out
+    )
+    if isinstance(res, np.ndarray):
+        assert_allclose(res, 3.5 * np.ones_like(res))
+    else:
+        assert_allclose(res, 3.5)