Skip to content

Commit

Permalink
MAINT: stats.sobol_indices: transition to rng (SPEC 7)
Browse files Browse the repository at this point in the history
  • Loading branch information
mdhaber committed Nov 13, 2024
1 parent 55fca04 commit 68acb0b
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 21 deletions.
2 changes: 1 addition & 1 deletion scipy/stats/_multicomp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
33 changes: 22 additions & 11 deletions scipy/stats/_sensitivity_analysis.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand All @@ -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):
Expand Down Expand Up @@ -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'.
Expand Down Expand Up @@ -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
-------
Expand Down Expand Up @@ -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])
Expand Down Expand Up @@ -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_:
Expand Down Expand Up @@ -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)
Expand Down
19 changes: 10 additions & 9 deletions scipy/stats/tests/test_sensitivity_analysis.py
Original file line number Diff line number Diff line change
Expand Up @@ -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':
Expand Down Expand Up @@ -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 = {
Expand All @@ -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)

Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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 = [
Expand All @@ -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)
Expand Down

0 comments on commit 68acb0b

Please sign in to comment.