codekingpro/portable-devtools

1from ._generator import Generator
2from ._mt19937 import MT19937
3from ._pcg64 import PCG64, PCG64DXSM
4from ._philox import Philox
5from ._sfc64 import SFC64
6from .bit_generator import BitGenerator
7from .mtrand import RandomState
8
9BitGenerators = {'MT19937': MT19937,
10                 'PCG64': PCG64,
11                 'PCG64DXSM': PCG64DXSM,
12                 'Philox': Philox,
13                 'SFC64': SFC64,
14                 }
15
16
17def __bit_generator_ctor(bit_generator: str | type[BitGenerator] = 'MT19937'):
18    """
19    Pickling helper function that returns a bit generator object
20
21    Parameters
22    ----------
23    bit_generator : type[BitGenerator] or str
24        BitGenerator class or string containing the name of the BitGenerator
25
26    Returns
27    -------
28    BitGenerator
29        BitGenerator instance
30    """
31    if isinstance(bit_generator, type):
32        bit_gen_class = bit_generator
33    elif bit_generator in BitGenerators:
34        bit_gen_class = BitGenerators[bit_generator]
35    else:
36        raise ValueError(
37            str(bit_generator) + ' is not a known BitGenerator module.'
38        )
39
40    return bit_gen_class()
41
42
43def __generator_ctor(bit_generator_name="MT19937",
44                     bit_generator_ctor=__bit_generator_ctor):
45    """
46    Pickling helper function that returns a Generator object
47
48    Parameters
49    ----------
50    bit_generator_name : str or BitGenerator
51        String containing the core BitGenerator's name or a
52        BitGenerator instance
53    bit_generator_ctor : callable, optional
54        Callable function that takes bit_generator_name as its only argument
55        and returns an instantized bit generator.
56
57    Returns
58    -------
59    rg : Generator
60        Generator using the named core BitGenerator
61    """
62    if isinstance(bit_generator_name, BitGenerator):
63        return Generator(bit_generator_name)
64    # Legacy path that uses a bit generator name and ctor
65    return Generator(bit_generator_ctor(bit_generator_name))
66
67
68def __randomstate_ctor(bit_generator_name="MT19937",
69                       bit_generator_ctor=__bit_generator_ctor):
70    """
71    Pickling helper function that returns a legacy RandomState-like object
72
73    Parameters
74    ----------
75    bit_generator_name : str
76        String containing the core BitGenerator's name
77    bit_generator_ctor : callable, optional
78        Callable function that takes bit_generator_name as its only argument
79        and returns an instantized bit generator.
80
81    Returns
82    -------
83    rs : RandomState
84        Legacy RandomState using the named core BitGenerator
85    """
86    if isinstance(bit_generator_name, BitGenerator):
87        return RandomState(bit_generator_name)
88    return RandomState(bit_generator_ctor(bit_generator_name))
89