Seed the legacy global RNG so random_state pins variability draws

ssms/basic_simulators/simulator.py

@@ -6,6 +6,8 @@
 """
 
 from copy import deepcopy
+import functools
+import inspect
 import numbers
 from threading import Lock
 
@@ -60,8 +62,8 @@ def _validate_random_state_for_c_rng(random_state: object) -> None:
 
     The Cython layer casts seeds to ``long``.  On Windows that is 32-bit signed; larger
     values raise ``OverflowError`` inside Cython.  We validate here with a clear
-    Python error.  Non-integer seeds (e.g. ``numpy.random.Generator``) are skipped;
-    those paths do not use this cast in the same way.
+    Python error.  Non-integer seeds (e.g. ``numpy.random.Generator``) skip this
+    range check and are rejected by the C-level seeding itself.
     """
     if random_state is None:
         return
@@ -664,7 +666,12 @@ def simulator(
             Integer passed to the C-level RNG seeding.  Must lie in
             ``[-2**31, 2**31 - 1]`` so it fits in a 32-bit signed C ``long`` (required
             on Windows).  ``None`` draws a seed from that range automatically.
-            Non-integer RNG objects may be supported on specific code paths.
+            Non-integer RNG objects (e.g. ``numpy.random.Generator``) are rejected
+            by the C-level seeding.  An integer seed also pins the trial-to-trial
+            variability draws (``sv``/``sz``/``st``), which are taken from a
+            generator derived from it rather than from NumPy's global RNG.  A
+            variability distribution that already binds its own ``random_state``
+            keeps that generator, so its draws follow it rather than the seed.
         return_option: str <default='full'>
             Determines what the function returns. Can be either
             'full' or 'minimal'. If 'full' the function returns
@@ -774,6 +781,25 @@ def simulator(
         theta, model_config_local, n_trials
     )
 
+    # Bind the trial-to-trial variability distributions to an explicit generator:
+    # theta's ``*_dist`` entries are scipy ``rvs`` partials, whose default draw
+    # source is NumPy's global RNG. Binding order is irrelevant - the simulator
+    # fixes the call order. The modulo is load-bearing: the validated seed range
+    # admits negative integers while ``default_rng`` takes only ``[0, 2**32)``,
+    # and the map is injective over that range.
+    if isinstance(random_state, numbers.Integral):
+        dist_rng = default_rng(int(random_state) % (2**32))
+        for key in set(model_config_local.get("simulator_param_mappings", {})) | set(
+            model_config_local.get("simulator_fixed_params", {})
+        ):
+            entry = theta.get(key)
+            if (
+                callable(entry)
+                and _accepts_random_state(entry)
+                and "random_state" not in getattr(entry, "keywords", {})
+            ):
+                theta[key] = functools.partial(entry, random_state=dist_rng)
+
     # Make boundary dictionary
     boundary_dict = make_boundary_dict(model_config_local, theta)
     # Make drift dictionary
@@ -862,3 +888,22 @@ def simulator(
         bin_simulator_output(x, nbins=256, max_t=-1, freq_cnt=True), axis=0
     )
     return x
+
+
+def _accepts_random_state(func) -> bool:
+    """Report whether ``func`` takes a ``random_state`` keyword.
+
+    A model configuration may map a parameter to any callable, and only those
+    that accept ``random_state`` can be bound to an explicit generator. scipy's
+    ``rvs`` takes it through ``**kwds``, so a variadic keyword also qualifies.
+    """
+    try:
+        params = inspect.signature(func).parameters
+    except (TypeError, ValueError):
+        return False
+    declared = params.get("random_state")
+    if declared is not None:
+        # A positional-only parameter cannot be filled by keyword: the value
+        # would land in **kwargs while the parameter kept its default, or raise.
+        return declared.kind is not inspect.Parameter.POSITIONAL_ONLY
+    return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())

tests/test_simulator.py

@@ -2,11 +2,14 @@
 import logging
 from unittest.mock import patch
 
+import functools
+
 import numpy as np
+import scipy.stats as sps
 import pandas as pd
 import pytest
 
-from ssms.basic_simulators.simulator import simulator
+from ssms.basic_simulators.simulator import _accepts_random_state, simulator
 from ssms.config import model_config
 
 logger = logging.getLogger(__name__)
@@ -330,3 +333,73 @@ def test_random_state_boundary_max_ok():
         random_state=2**31 - 1,
     )
     assert "rts" in out
+
+
+@pytest.mark.rng_validation
+@pytest.mark.parametrize("model", ["ddm", "ddm_st", "full_ddm_rv", "ddm_sdv"])
+def test_random_state_pins_variability_draws(sim_input_data, model):
+    """A repeated integer random_state reproduces the trial-to-trial variability
+    draws (sv/sz/st), not just the diffusion path, and is unaffected by unrelated
+    consumption of NumPy's global RNG between calls."""
+    theta = dict(sim_input_data[model]["theta_dict_all_scalars"])
+    # The registry defaults put sv/sz/st at 1e-3, below what a float32 RT
+    # resolves; raise them so a differing draw shows up in the output.
+    theta.update(
+        {k: v for k, v in {"sv": 0.5, "sz": 0.1, "st": 0.13}.items() if k in theta}
+    )
+
+    state = np.random.get_state()
+    try:
+        first = simulator(model=model, theta=theta, n_samples=500, random_state=7)
+        # The variability draws come from a generator derived from the seed, so
+        # the call leaves the process-global RNG exactly where it found it. A
+        # global reseed or a global draw would move it.
+        after_first = np.random.get_state()
+        np.testing.assert_array_equal(after_first[1], state[1])
+        assert after_first[2] == state[2]
+        # unrelated consumption of the global RNG must not affect the draws
+        np.random.uniform(size=1234)
+        second = simulator(model=model, theta=theta, n_samples=500, random_state=7)
+    finally:
+        np.random.set_state(state)
+
+    np.testing.assert_array_equal(first["rts"], second["rts"])
+    np.testing.assert_array_equal(first["choices"], second["choices"])
+
+
+@pytest.mark.rng_validation
+def test_accepts_random_state_rejects_positional_only():
+    """A positional-only ``random_state`` cannot be filled by keyword.
+
+    Binding one would either divert the value into ``**kwargs`` while the
+    parameter kept its default, or raise, so such callables are left alone.
+    """
+
+    def positional_only(size=1, random_state=None, /, **kwargs):
+        """Take random_state positionally only; keyword use lands in kwargs."""
+        return random_state
+
+    def keyword_ok(size=1, random_state=None):
+        """Accept random_state by keyword."""
+        return random_state
+
+    def variadic_only(**kwargs):
+        """Accept random_state only through **kwargs, as scipy rvs does."""
+        return kwargs.get("random_state")
+
+    assert not _accepts_random_state(positional_only)
+    assert _accepts_random_state(keyword_ok)
+    assert _accepts_random_state(variadic_only)
+    # the failure mode the exclusion prevents: value diverted, default kept
+    assert functools.partial(positional_only, random_state="RNG")() is None
+
+
+@pytest.mark.rng_validation
+def test_bound_random_state_is_not_overridden():
+    """A distribution that already carries its own random_state keeps it."""
+    own = np.random.default_rng(123)
+    dist = functools.partial(sps.uniform.rvs, loc=0.0, scale=1.0, random_state=own)
+    # accepted by the predicate, but already bound, so the simulator leaves it
+    assert _accepts_random_state(dist)
+    assert "random_state" in dist.keywords
+    assert dist.keywords["random_state"] is own