| """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 | |