MAINT: stats.sobol_indices: transition to rng (SPEC 7)

Author: mdhaber

scipy/stats/_multicomp.py

@@ -182,7 +182,7 @@ def confidence_interval(
         return self._ci
 
 
-@_transition_to_rng('random_state')
+@_transition_to_rng('random_state', replace_doc=False)
 def dunnett(
     *samples: npt.ArrayLike,  # noqa: D417
     control: npt.ArrayLike,

scipy/stats/_sensitivity_analysis.py

@@ -13,6 +13,7 @@
 from scipy.stats._qmc import check_random_state
 from scipy.stats._resampling import BootstrapResult
 from scipy.stats import qmc, bootstrap
+from scipy._lib._util import _transition_to_rng
 
 
 if TYPE_CHECKING:
@@ -61,7 +62,7 @@ def f_ishigami(x: npt.ArrayLike) -> np.ndarray:
 def sample_A_B(
     n: IntNumber,
     dists: list[PPFDist],
-    random_state: SeedType = None
+    rng: SeedType = None
 ) -> np.ndarray:
     """Sample two matrices A and B.
 
@@ -80,7 +81,7 @@ def sample_A_B(
        :doi:`10.1016/j.cpc.2009.09.018`, 2010.
     """
     d = len(dists)
-    A_B = qmc.Sobol(d=2*d, seed=random_state, bits=64).random(n).T
+    A_B = qmc.Sobol(d=2*d, seed=rng, bits=64).random(n).T
     A_B = A_B.reshape(2, d, -1)
     try:
         for d_, dist in enumerate(dists):
@@ -246,14 +247,15 @@ def ppf(self) -> Callable[..., float]:
         ...
 
 
+@_transition_to_rng('random_state', replace_doc=False)
 def sobol_indices(
     *,
     func: Callable[[np.ndarray], npt.ArrayLike] |
           dict[Literal['f_A', 'f_B', 'f_AB'], np.ndarray],
     n: IntNumber,
     dists: list[PPFDist] | None = None,
     method: Callable | Literal['saltelli_2010'] = 'saltelli_2010',
-    random_state: SeedType = None
+    rng: SeedType = None
 ) -> SobolResult:
     r"""Global sensitivity indices of Sobol'.
 
@@ -309,11 +311,20 @@ def sobol_indices(
         The output is a tuple of the first and total indices with
         shape ``(s, d)``.
         This is an advanced feature and misuse can lead to wrong analysis.
-    random_state : {None, int, `numpy.random.Generator`}, optional
-        If `random_state` is an int or None, a new `numpy.random.Generator` is
-        created using ``np.random.default_rng(random_state)``.
-        If `random_state` is already a ``Generator`` instance, then the
-        provided instance is used.
+    rng : `numpy.random.Generator`, optional
+        Pseudorandom number generator state. When `rng` is None, a new
+        `numpy.random.Generator` is created using entropy from the
+        operating system. Types other than `numpy.random.Generator` are
+        passed to `numpy.random.default_rng` to instantiate a ``Generator``.
+
+        .. versionchanged:: 1.15.0
+
+            As part of the `SPEC-007 <https://scientific-python.org/specs/spec-0007/>`_
+            transition from use of `numpy.random.RandomState` to
+            `numpy.random.Generator`, this keyword was changed from `random_state` to `rng`.
+            For an interim period, both keywords will continue to work, although only one
+            may be specified at a time. After the interim period, function calls using the
+            `random_state` keyword will emit warnings.
 
     Returns
     -------
@@ -466,7 +477,7 @@ def sobol_indices(
     ...         uniform(loc=-np.pi, scale=2*np.pi),
     ...         uniform(loc=-np.pi, scale=2*np.pi)
     ...     ],
-    ...     random_state=rng
+    ...     rng=rng
     ... )
     >>> indices.first_order
     array([0.31637954, 0.43781162, 0.00318825])
@@ -587,7 +598,7 @@ def sobol_indices(
     conclude on high-order terms. Hence the benefit of using Sobol' indices.
 
     """
-    random_state = check_random_state(random_state)
+    rng = check_random_state(rng)
 
     n_ = int(n)
     if not (n_ & (n_ - 1) == 0) or n != n_:
@@ -637,7 +648,7 @@ def indices_method(f_A, f_B, f_AB):
         def wrapped_func(x):
             return np.atleast_2d(func(x))
 
-        A, B = sample_A_B(n=n, dists=dists, random_state=random_state)
+        A, B = sample_A_B(n=n, dists=dists, rng=rng)
         AB = sample_AB(A=A, B=B)
 
         f_A = wrapped_func(A)

scipy/stats/tests/test_sensitivity_analysis.py

@@ -86,7 +86,7 @@ def test_ishigami(self, ishigami_ref_indices, func):
         res = sobol_indices(
             func=func, n=4096,
             dists=self.dists,
-            random_state=rng
+            rng=rng
         )
 
         if func.__name__ == 'f_ishigami_vec':
@@ -141,7 +141,7 @@ def test_func_dict(self, ishigami_ref_indices):
             stats.uniform(loc=-np.pi, scale=2*np.pi)
         ]
 
-        A, B = sample_A_B(n=n, dists=dists, random_state=rng)
+        A, B = sample_A_B(n=n, dists=dists, rng=rng)
         AB = sample_AB(A=A, B=B)
 
         func = {
@@ -150,16 +150,17 @@ def test_func_dict(self, ishigami_ref_indices):
             'f_AB': f_ishigami(AB).reshape((3, 1, -1))
         }
 
+        # preserve use of old random_state during SPEC 7 transition
         res = sobol_indices(
             func=func, n=n,
             dists=dists,
-            random_state=rng
+            rng=rng
         )
         assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)
 
         res = sobol_indices(
             func=func, n=n,
-            random_state=rng
+            rng=rng
         )
         assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)
 
@@ -180,7 +181,7 @@ def jansen_sobol(f_A, f_B, f_AB):
             func=f_ishigami, n=4096,
             dists=self.dists,
             method=jansen_sobol,
-            random_state=rng
+            rng=rng
         )
 
         assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)
@@ -195,15 +196,15 @@ def jansen_sobol_typed(
             func=f_ishigami, n=8,
             dists=self.dists,
             method=jansen_sobol_typed,
-            random_state=rng
+            rng=rng
         )
 
     def test_normalization(self, ishigami_ref_indices):
         rng = np.random.default_rng(28631265345463262246170309650372465332)
         res = sobol_indices(
             func=lambda x: f_ishigami(x) + 1000, n=4096,
             dists=self.dists,
-            random_state=rng
+            rng=rng
         )
 
         assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)
@@ -220,7 +221,7 @@ def f_ishigami_vec_const(x):
         res = sobol_indices(
             func=f_ishigami_vec_const, n=4096,
             dists=self.dists,
-            random_state=rng
+            rng=rng
         )
 
         ishigami_vec_indices = [
@@ -237,7 +238,7 @@ def test_more_converged(self, ishigami_ref_indices):
         res = sobol_indices(
             func=f_ishigami, n=2**19,  # 524288
             dists=self.dists,
-            random_state=rng
+            rng=rng
         )
 
         assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-4)