Merge pull request #318 from bashtage/reenable-pickle

ENH: Allow pickle with NumPy Generator

Author: bashtage

README.md

@@ -98,10 +98,10 @@ The RNGs include:
 ## Status
 
 * Builds and passes all tests on:
-  * Linux 32/64 bit, Python 2.7, 3.5, 3.6, 3.7
-  * Linux (ARM/ARM64), Python 3.7
-  * OSX 64-bit, Python 2.7, 3.5, 3.6, 3.7
-  * Windows 32/64 bit, Python 2.7, 3.5, 3.6, 3.7
+  * Linux 32/64 bit, Python 3.7, 3.8, 3.9, 3.10
+  * Linux (ARM/ARM64), Python 3.8
+  * OSX 64-bit, Python 3.9
+  * Windows 32/64 bit, Python 3.7, 3.8, 3.9, 3.10
   * FreeBSD 64-bit
 
 ## Version

ci/azure/azure_template_posix.yml

@@ -26,12 +26,14 @@ jobs:
           python.version: '3.8'
           coverage: true
           NUMPY: 1.17.0
+        python37_latest:
+          python.version: '3.7'
+        python38_latest:
+          python.version: '3.8'
         python39_latest:
           python.version: '3.9'
         python310_latest:
           python.version: '3.10'
-        python36_latest:
-          python.version: '3.6'
         python38_mid_conda:
           python.version: '3.8'
           use.conda: true

ci/azure/azure_template_windows.yml

@@ -16,10 +16,14 @@ jobs:
     vmImage: ${{ parameters.vmImage }}
   strategy:
     matrix:
+      python37_win_latest:
+        python.version: '3.7'
       python38_win_latest:
         python.version: '3.8'
       python39_win_latest:
         python.version: '3.9'
+      python310_win_latest:
+        python.version: '3.10'
     maxParallel: 10
 
   steps:

doc/source/change-log.rst

@@ -13,6 +13,16 @@ Change Log
   You should be using :class:`numpy.random.Generator` or
   :class:`numpy.random.RandomState` which are maintained. 
 
+v1.23.1
+=======
+- Registered the bit generators included in ``randomgen`` with NumPy
+  so that NumPy :class:`~numpy.random.Generator` instances can be pickled
+  and unpickled when using a ``randomstate`` bit generator.
+- Changed the canonical name of the bit generators to be their fully qualified
+  name. For example, :class:`~randomgen.pcg64.PCG64` is not named ``"randomgen.pcg64.PCG64"``
+  instead of ``"PCG64"``.  This was done to avoid ambiguity with NumPy's supplied
+  bit generators with the same name.
+
 v1.23.0
 =======
 - Removed ``Generator`` and ``RandomState``.

randomgen/__init__.py

@@ -2,6 +2,7 @@
 import sys
 from typing import List, Union
 
+from randomgen._register import BitGenerators
 from randomgen.aes import AESCounter
 from randomgen.chacha import ChaCha
 from randomgen.dsfmt import DSFMT
@@ -37,6 +38,7 @@
 
 __all__ = [
     "AESCounter",
+    "BitGenerators",
     "ChaCha",
     "DSFMT",
     "EFIIX64",

randomgen/_pickle.py

@@ -6,13 +6,12 @@
 from randomgen.common import BitGenerator
 from randomgen.dsfmt import DSFMT
 from randomgen.efiix64 import EFIIX64
-from randomgen.generator import ExtendedGenerator, Generator
+from randomgen.generator import ExtendedGenerator
 from randomgen.hc128 import HC128
 from randomgen.jsf import JSF
 from randomgen.lxm import LXM
 from randomgen.mt64 import MT64
 from randomgen.mt19937 import MT19937
-from randomgen.mtrand import RandomState
 from randomgen.pcg32 import PCG32
 from randomgen.pcg64 import PCG64, PCG64DXSM, LCG128Mix
 from randomgen.philox import Philox
@@ -54,6 +53,10 @@
     "RDRAND": RDRAND,
 }
 
+# Assign the fully qualified name for future proofness
+for value in list(BitGenerators.values()):
+    BitGenerators[f"{value.__module__}.{value.__name__}"] = value
+
 
 def _get_bitgenerator(bit_generator_name: str) -> Type[BitGenerator]:
     """
@@ -75,29 +78,6 @@ def _decode(name: Union[str, bytes]) -> str:
     return name.decode("ascii")
 
 
-def __generator_ctor(bit_generator_name: Union[bytes, str] = "MT19937") -> Generator:
-    """
-    Pickling helper function that returns a Generator object
-
-    Parameters
-    ----------
-    bit_generator_name: str
-        String containing the core BitGenerator
-
-    Returns
-    -------
-    rg: Generator
-        Generator using the named core BitGenerator
-    """
-    bit_generator_name = _decode(bit_generator_name)
-    assert isinstance(bit_generator_name, str)
-    bit_generator = _get_bitgenerator(bit_generator_name)
-    with warnings.catch_warnings():
-        warnings.filterwarnings("ignore", category=FutureWarning)
-        bit_gen = bit_generator()
-    return Generator(bit_gen)
-
-
 def __extended_generator_ctor(
     bit_generator_name: Union[str, bytes] = "MT19937"
 ) -> ExtendedGenerator:
@@ -146,28 +126,3 @@ def __bit_generator_ctor(
         warnings.filterwarnings("ignore", category=FutureWarning)
         bit_gen = bit_generator()
     return bit_gen
-
-
-def __randomstate_ctor(
-    bit_generator_name: Union[str, bytes] = "MT19937"
-) -> RandomState:
-    """
-    Pickling helper function that returns a legacy RandomState-like object
-
-    Parameters
-    ----------
-    bit_generator_name: str
-        String containing the core BitGenerator
-
-    Returns
-    -------
-    rs: RandomState
-        Legacy RandomState using the named core BitGenerator
-    """
-    bit_generator_name = _decode(bit_generator_name)
-    assert isinstance(bit_generator_name, str)
-    bit_generator = _get_bitgenerator(bit_generator_name)
-    with warnings.catch_warnings():
-        warnings.filterwarnings("ignore", category=FutureWarning)
-        bit_gen = bit_generator()
-    return RandomState(bit_gen)

randomgen/_register.py

@@ -0,0 +1,62 @@
+from numpy.random._pickle import BitGenerators
+
+from randomgen.aes import AESCounter
+from randomgen.chacha import ChaCha
+from randomgen.dsfmt import DSFMT
+from randomgen.efiix64 import EFIIX64
+from randomgen.hc128 import HC128
+from randomgen.jsf import JSF
+from randomgen.lxm import LXM
+from randomgen.mt64 import MT64
+from randomgen.mt19937 import MT19937
+from randomgen.pcg32 import PCG32
+from randomgen.pcg64 import PCG64, PCG64DXSM, LCG128Mix
+from randomgen.philox import Philox
+from randomgen.rdrand import RDRAND
+from randomgen.romu import Romu
+from randomgen.sfc import SFC64
+from randomgen.sfmt import SFMT
+from randomgen.speck128 import SPECK128
+from randomgen.threefry import ThreeFry
+from randomgen.wrapper import UserBitGenerator
+from randomgen.xoroshiro128 import Xoroshiro128
+from randomgen.xorshift1024 import Xorshift1024
+from randomgen.xoshiro256 import Xoshiro256
+from randomgen.xoshiro512 import Xoshiro512
+
+bit_generators = [
+    AESCounter,
+    ChaCha,
+    DSFMT,
+    EFIIX64,
+    HC128,
+    JSF,
+    LXM,
+    MT19937,
+    MT64,
+    PCG32,
+    PCG64,
+    PCG64DXSM,
+    LCG128Mix,
+    Philox,
+    RDRAND,
+    Romu,
+    SFC64,
+    SFMT,
+    SPECK128,
+    ThreeFry,
+    UserBitGenerator,
+    Xoroshiro128,
+    Xorshift1024,
+    Xoshiro256,
+    Xoshiro512,
+]
+
+for bitgen in bit_generators:
+    key = f"{bitgen.__name__}"
+    if key not in BitGenerators:
+        BitGenerators[key] = bitgen
+    full_key = f"{bitgen.__module__}.{bitgen.__name__}"
+    BitGenerators[full_key] = bitgen
+
+__all__ = ["BitGenerators"]

randomgen/aes.pyx

@@ -277,7 +277,7 @@ cdef class AESCounter(BitGenerator):
         for i in range(16 * 4):
             state[i] = self.rng_state.state[i]
         offset = self.rng_state.offset
-        return {"bit_generator": type(self).__name__,
+        return {"bit_generator": fully_qualified_name(self),
                 "s": {"state": state, "seed": seed, "counter": counter,
                       "offset": offset},
                 "has_uint32": self.rng_state.has_uint32,
@@ -290,7 +290,7 @@ cdef class AESCounter(BitGenerator):
         if not isinstance(value, dict):
             raise TypeError("state must be a dict")
         bitgen = value.get("bit_generator", "")
-        if bitgen != type(self).__name__:
+        if bitgen not in (type(self).__name__, fully_qualified_name(self)):
             raise ValueError("state must be for a {0} "
                              "PRNG".format(type(self).__name__))
         state =value["s"]["state"]

randomgen/chacha.pyx

@@ -279,7 +279,7 @@ cdef class ChaCha(BitGenerator):
         for i in range(2):
             ctr[i] = self.rng_state.ctr[i]
 
-        return {"bit_generator": type(self).__name__,
+        return {"bit_generator": fully_qualified_name(self),
                 "state": {"block": block, "keysetup": keysetup, "ctr": ctr,
                           "rounds": self.rng_state.rounds}}
 
@@ -288,7 +288,7 @@ cdef class ChaCha(BitGenerator):
         if not isinstance(value, dict):
             raise TypeError("state must be a dict")
         bitgen = value.get("bit_generator", "")
-        if bitgen != type(self).__name__:
+        if bitgen not in (type(self).__name__, fully_qualified_name(self)):
             raise ValueError("state must be for a {0} "
                              "PRNG".format(type(self).__name__))
 

randomgen/common.pxd

@@ -148,3 +148,6 @@ cdef inline void compute_complex(double *rv_r, double *rv_i, double loc_r,
 
     rv_i[0] = loc_i + scale_i * (rho * rv_r[0] + scale_c * rv_i[0])
     rv_r[0] = loc_r + scale_r * rv_r[0]
+
+
+cdef object fully_qualified_name(instance)