_pickle.py 2.7 KB

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