Source code for hetero.stims.utils

from numpy.random import MT19937, Generator, RandomState, default_rng
from typing import Union
import numpy as np
import numbers

__global_rg = default_rng()



def get_seed(seed=None):
    if seed is None:
        return 2026  # a fixed seed to prevent progression of the rng
    else:
        return seed





def rand_generator(seed: Union[int, Generator, RandomState] = None) -> Generator:
    if seed is None:
        return __global_rg
    # provided to support legacy RandomState generator
    # of Numpy. It is not the best thing to do however
    # and recommend the user to keep using integer seeds
    # and proper Numpy Generator API.
    if isinstance(seed, RandomState):
        mt19937 = MT19937()
        mt19937.state = seed.get_state()
        return Generator(mt19937)

    if isinstance(seed, Generator):
        return seed
    else:
        return default_rng(seed)

    




def check_vector(array, allow_reshape=True, allow_timespans=True, caller=None):
    msg = "."
    if caller is not None:
        if hasattr(caller, "name"):
            msg = f" in {caller.name}."
        else:
            msg = f"in {caller}."

    if not isinstance(array, np.ndarray):
        # maybe a single number, make it an array
        if isinstance(array, numbers.Number):
            array = np.asarray(array)
        else:
            msg = (
                f"Data type '{type(array)}' not understood. All vectors "
                f"should be Numpy arrays" + msg
            )
            raise TypeError(msg)

    if not (np.issubdtype(array.dtype, np.number)):
        msg = f"Impossible to operate on non-numerical data, in array: {array}" + msg
        raise TypeError(msg)

    if allow_reshape:
        array = np.atleast_2d(array)

    if not allow_timespans:
        if array.shape[0] > 1:
            msg = (
                f"Impossible to operate on multiple timesteps. Data should"
                f" have shape (1, n) but is {array.shape}" + msg
            )
            raise ValueError(msg)

    # TODO: choose axis to expand and/or np.atleast_2d

    return array