Writing a BitGenerator¶

import numpy as np


# The rotation
def rotr_64(value, rot):
    value = np.uint64(value)
    rot = np.uint64(rot)
    return int((value >> rot) | (value << ((-rot) & np.uint(63))))


class PythonPCG64:
    # A 128 bit multiplier
    PCG_DEFAULT_MULTIPLIER = (2549297995355413924 << 64) + 4865540595714422341
    MODULUS = 2**128

    def __init__(self, state, inc):
        """Directly set the state and increment, no seed support"""
        self.state = state
        self.inc = inc
        self._has_uint32 = False
        self._uinteger = 0
        self._next_32 = self._next_64 = None

    def random_raw(self):
        """Generate the next "raw" value, which is 64 bits"""
        state = self.state * self.PCG_DEFAULT_MULTIPLIER + self.inc
        state = state % self.MODULUS
        self.state = state
        return rotr_64((state >> 64) ^ (state & 0xFFFFFFFFFFFFFFFF), state >> 122)

    @property
    def next_64(self):
        """
        Return a callable that accepts a single input. The input is usually
        a void pointer that is cast to a struct that contains the PRNGs
        state. When wiring a bit generator in Python, it is simpler to use
        a closure than to wrap the state in an array, pass it's address as a
        ctypes void pointer, and then to get the pointer in the function.
        """

        def _next_64(void_p):
            return self.random_raw()

        self._next_64 = _next_64
        return _next_64

    @property
    def next_32(self):
        """
        Return a callable that accepts a single input. This is identical to
        ``next_64`` except that it return a 32-bit unsigned int. Here we save
        half of the raw 64 bit output for subsequent calls.
        """

        def _next_32(void_p):
            if self._has_uint32:
                self._has_uint32 = False
                return self._uinteger
            next_value = self.random_raw()
            self._has_uint32 = True
            self._uinteger = next_value >> 32
            return next_value & 0xFFFFFFFF

        self._next_32 = _next_32
        return _next_32

    @property
    def state_getter(self):
        def f():
            return {"state": self.state, "inc": self.inc}

        return f

    @property
    def state_setter(self):
        def f(value):
            self.state = value["state"]
            self.inc = value["inc"]

        return f

import ctypes

import numpy as np
from numba import cfunc, types, carray, jit

from randomgen.wrapper import UserBitGenerator


rotate64_sig = types.uint64(types.uint64, types.int_)


@jit(signature_or_function=rotate64_sig, inline="always")
def rotate64(x, k):
    return (x << k) | (x >> (64 - k))


jsf_next_sig = types.uint64(types.uint64[:])


@jit(signature_or_function=jsf_next_sig, inline="always")
def jsf_next(state):
    """
    Update the state in place

    This is a literal translation of the C code where the value p, q,
    and r are fixed.
    """
    # Default values
    p = 7
    q = 13
    r = 37
    # Update
    e = state[0] - rotate64(state[1], p)
    state[0] = state[1] ^ rotate64(state[2], q)
    state[1] = state[2] + (rotate64(state[3], r) if r else state[3])
    state[2] = state[3] + e
    state[3] = e + state[0]
    return state[3]


class NumbaJSF:
    def __init__(self, seed):
        if not isinstance(seed, (int, np.integer)) or not (0 <= state < 2**64):
            raise ValueError("seed must be a valid uint64")
        # state[0:4] is the JSF state
        # state[4] contains both the has_uint flag in bit 0
        #   uinteger in bits 32...63
        self._state = np.zeros(5, dtype=np.uint64)
        self._next_raw = None
        self._next_64 = None
        self._next_32 = None
        self._next_double = None
        self.seed(seed)

    def seed(self, value):
        self._state[0] = 0xF1EA5EED
        self._state[1] = value
        self._state[2] = value
        self._state[3] = value
        for i in range(20):
            jsf_next(self._state)

    @property
    def state_address(self):
        """Get the location in memory of the state NumPy array."""
        return self._state.ctypes.data_as(ctypes.c_void_p)

    @property
    def next_64(self):
        """Same as raw since a 64 bit generator"""

        # Ensure a reference is held
        self._next_64 = self.next_raw

        return self.next_raw

    @property
    def next_32(self):
        """A CFunc generating the next 32 bits"""
        sig = types.uint32(types.CPointer(types.uint64))

        @cfunc(sig)
        def next_32(st):
            # Get the NumPy uint64 array
            bit_gen_state = carray(st, (5,), dtype=np.uint64)
            # We use the first bit to indicate that 32 bits are stored in 32...63
            if bit_gen_state[4] & np.uint64(0x1):
                # Get the upper 32 bits
                out = bit_gen_state[4] >> np.uint64(32)
                # Clear the stored value
                bit_gen_state[4] = 0
                return out
            # If no bits available, genrate a new value
            z = jsf_next(bit_gen_state)
            # Store the new value always with 1 in bit 0
            bit_gen_state[4] = z | np.uint64(0x1)
            # Return the lower 32 (0...31)
            return z & 0xFFFFFFFF

        # Ensure a reference is held
        self._next_32 = next_32

        return next_32

    @property
    def next_double(self):
        """A CFunc that generates the next ouble"""
        sig = types.double(types.CPointer(types.uint64))

        @cfunc(sig)
        def next_double(st):
            # Get the state
            bit_gen_state = carray(st, (5,), dtype=np.uint64)
            # Return the next value / 2**53
            return (
                np.uint64(jsf_next(bit_gen_state)) >> np.uint64(11)
            ) / 9007199254740992.0

        # Ensure a reference is held
        self._next_double = next_double

        return next_double

    @property
    def next_raw(self):
        sig = types.uint64(types.CPointer(types.uint64))

        @cfunc(sig)
        def next_64(st):
            # Get the NumPy array containing the state
            bit_gen_state = carray(st, (5,), dtype=np.uint64)
            # Return the next value
            return jsf_next(bit_gen_state)

        # Ensure a reference is held
        self._next_64 = next_64

        return next_64

    @property
    def state_getter(self):
        """A function that returns the state. This is Python and is not decorated"""

        def f() -> dict:
            return {
                "bit_gen": type(self).__name__,
                "state": self._state[:4],
                "has_uint": self._state[4] & np.uint64(0x1),
                "uinteger": self._state[4] >> np.uint64(32),
            }

        return f

    @property
    def state_setter(self):
        """A function that sets the state. This is Python and is not decorated"""

        def f(value: dict):
            name = value.get("bit_gen", None)
            if name != type(self).__name__:
                raise ValueError(f"state must be from a {type(self).__name__}")
            self._state[:4] = np.uint64(value["state"])
            temp = np.uint64(value["uinteger"]) << np.uint64(32)
            temp |= np.uint64(value["has_uint"]) & np.uint64(0x1)
            self._state[4] = temp

        return f