mirror of https://github.com/hero-y/BHRL
34 lines
977 B
Python
34 lines
977 B
Python
|
"""Helpers for random number generators."""
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
def ensure_rng(rng=None):
|
||
|
|
"""Coerces input into a random number generator.
|
||
|
|
|
||
|
|
If the input is None, then a global random state is returned.
|
||
|
|
|
||
|
|
If the input is a numeric value, then that is used as a seed to construct a
|
||
|
|
random state. Otherwise the input is returned as-is.
|
||
|
|
|
||
|
|
Adapted from [1]_.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
rng (int | numpy.random.RandomState | None):
|
||
|
|
if None, then defaults to the global rng. Otherwise this can be an
|
||
|
|
integer or a RandomState class
|
||
|
|
Returns:
|
||
|
|
(numpy.random.RandomState) : rng -
|
||
|
|
a numpy random number generator
|
||
|
|
|
||
|
|
References:
|
||
|
|
.. [1] https://gitlab.kitware.com/computer-vision/kwarray/blob/master/kwarray/util_random.py#L270 # noqa: E501
|
||
|
|
"""
|
||
|
|
|
||
|
|
if rng is None:
|
||
|
|
rng = np.random.mtrand._rand
|
||
|
|
elif isinstance(rng, int):
|
||
|
|
rng = np.random.RandomState(rng)
|
||
|
|
else:
|
||
|
|
rng = rng
|
||
|
|
return rng
|