PractRand driver for testing interleaved PCG32 streams

pcg_streams.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import json
import secrets
import sys

import numpy as np
from numpy.random import PCG32

# The value from the pcg-c sources is the actual increment, not the half-size
# stream ID that we actually need to pass in.
DEFAULT_PCG32_INC = 2891336453 // 2

MASK63 = 0x7FFFFFFFFFFFFFFF


def gen_interleaved_bytes(bitgens, n_per_gen=1024, output_dtype=np.uint32):
    while True:
        draws = [g.random_raw(n_per_gen).astype(output_dtype) for g in bitgens]
        interleaved = np.column_stack(draws).ravel()
        bytes_chunk = bytes(interleaved.data)
        yield bytes_chunk


def bitgen_from_state(state):
    cls = getattr(np.random, state['bit_generator'])
    bitgen = cls()
    bitgen.state = state
    return bitgen


def hash63(x):
    """ 63-bit hash.

    Derived from https://gist.github.com/degski/6e2069d6035ae04d5d6f64981c995ec2
    """
    MUL = 0xD6E8FEB86659FD93

    # Make sure that we don't map 0 to 0.
    x ^= DEFAULT_PCG32_INC

    x ^= x >> 32
    x = (x * MUL) & MASK63
    x ^= x >> 32
    x = (x * MUL) & MASK63
    x ^= x >> 32
    assert x <= MASK63
    return x


def murmur63(x):
    """ Adaptation of a 63-bit variant of the Murmur3 finalizer.

    http://zimbry.blogspot.com/2011/09/better-bit-mixing-improving-on.html
    """
    # Make sure that we don't map 0 to 0.
    x = x ^ DEFAULT_PCG32_INC

    x ^= x >> 33
    x = (x * 0xff51afd7ed558ccd) & MASK63
    x ^= x >> 33
    x = (x * 0xc4ceb9fe1a85ec53) & MASK63
    x ^= x >> 33
    assert x <= MASK63
    return x


def same_seed_incr_streams(n_streams=2, seed=None, inc0=DEFAULT_PCG32_INC):
    if seed is None:
        seed = secrets.randbits(64)
    bitgens = [PCG32(seed, inc=inc)
               for inc in range(inc0, inc0 + n_streams)]
    return bitgens


def same_seed_hashed_streams(n_streams=2, seed=None, inc0=0):
    if seed is None:
        seed = secrets.randbits(64)
    bitgens = [PCG32(seed, inc=murmur63(inc))
               for inc in range(inc0, inc0 + n_streams)]
    return bitgens


def jumped_state(n_streams=2, seed=None, inc=DEFAULT_PCG32_INC):
    if seed is None:
        seed = secrets.randbits(64)
    bg = PCG32(seed, inc=inc)
    bitgens = [bg]
    for i in range(n_streams - 1):
        bg = bg.jumped()
        bitgens.append(bg)
    return bitgens


def inverted_punched_state(seed=None, inc=DEFAULT_PCG32_INC):
    if seed is None:
        seed = secrets.randbits(64)
    bitgen1 = PCG32(seed, inc=inc)
    state1 = bitgen1.state
    state2 = state1.copy()
    state2['state']['state'] = -state1['state']['state'] % (1<<64)
    state2['state']['inc'] = -state1['state']['inc'] % (1<<64)
    bitgen2 = PCG32()
    bitgen2.state = state2
    return [bitgen1, bitgen2]


def inverted_with_init(seed=None, inc1=DEFAULT_PCG32_INC):
    if seed is None:
        seed = secrets.randbits(64)
    bitgen1 = PCG32(seed, inc=inc1)
    inc2 = ((-(inc1 * 2 + 1) % (1<<64)) % (1<<64)) // 2
    bitgen2 = PCG32(seed, inc=inc2)
    return [bitgen1, bitgen2]


def dump_states(bitgens, file=sys.stderr):
    text = json.dumps([g.state for g in bitgens], indent=4)
    print(text, file=file)


def from_json(filename):
    with open(filename) as f:
        states = json.load(f)
    bitgens = [bitgen_from_state(s) for s in states]
    return bitgens


def main():
    import argparse
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-j', '--json',
                        help='Load BitGenerators from JSON file.')
    parser.add_argument('--jumped', action='store_true',
                        help='Use jumped() to get new streams.')
    parser.add_argument('-s', '--seed', type=int,
                        help='Set a single seed')
    parser.add_argument('-i', '--inc', type=int, default=DEFAULT_PCG32_INC,
                        help='The starting increment')
    parser.add_argument('-n', '--n-streams', type=int, default=2,
                        help='The number of streams to interleave')

    args = parser.parse_args()
    if args.json is not None:
        bitgens = from_json(args.json)
    elif args.jumped:
        bitgens = jumped_state(n_streams=args.n_streams, seed=args.seed,
                               inc=args.inc)
    else:
        # bitgens = same_seed_incr_streams(n_streams=args.n_streams,
        #                                  seed=args.seed, inc0=args.inc)
        bitgens = same_seed_hashed_streams(n_streams=args.n_streams,
                                         seed=args.seed, inc0=args.inc)
        # bitgens = inverted_punched_state(seed=args.seed, inc=args.inc)
        # bitgens = inverted_with_init(seed=args.seed, inc1=args.inc)
    dump_states(bitgens)
    for chunk in gen_interleaved_bytes(bitgens):
        sys.stdout.buffer.write(chunk)


if __name__ == '__main__':
    main()