numpy implementing a custom RNG

Question

I'm trying to get numpy to use my own implementation of an RNG in for consistency reasons. My understanding, based on the little documentation I could find, from the numpy docs here and here is that I need to provide a custom BitGenerator class that implements the random_raw method and then initialise using np.random.Generator, so I tried this:

import numpy as np

class TestBitGenerator(np.random.BitGenerator):
  def __init__(self):
    super().__init__(0)
    self.counter = 0
  def random_raw(self, size=None):
    self.counter += 1
    if size is None:
      return self.counter
    return np.full(n, self.counter)  

mc_adapter = TestBitGenerator() 
npgen = np.random.Generator(mc_adapter)

print(npgen.random())

which results in a segfault:

$ python bitgen.py 
Segmentation fault (core dumped)

I assume I'm missing something (from TestBitGenerator?) here, can anyone point me in the right direction?

I tried not subclassing np.random.BitGenerator, and got a different error: object has no attribute 'capsule'

I using numpy 1.19.2 and python 3.8.2

Answer 1

Actually you can do it in pure python if you wrap it using the library RandomGen (this was the incubator for the current np.random.Generator). So there is a UserBitGenerator that allow you to use only python: https://bashtage.github.io/randomgen/bit_generators/userbitgenerator.html

Sad that this did not make it in numpy (or if it is I did not find it yet...).

Answer 2

If you are using pybind11, it's not too hard, although there's no good documentation explaining any of this, so it took me a while to figure it out. Here is what I came up with, so hopefully this will save a few people some consternation.

1. Define your custom bit generator in C++. You'll need to define three functions: next_unint32, next_uint64, and next_double, each returning the next random value of the given type. (Typically one of these will be primary, and you'll define the other two based on it.) They all need to take one argument, given as void*, which you'll want to recast to whatever you actually have for a state object. Say a pointer to an instance of some class CustomRNG or whatever.
2. Write a function that takes a pointer to your state object and a py::capsule (I use namespace py = pybind11, so spell that out if you don't do this.) This capsule wraps a bitgen_t struct, and writes the appropriate values to its elements. Something like the following:

#include <numpy/random/bitgen.h>  // For bitgen_t struct

// For reference:
/*
typedef struct bitgen {
  void *state;
  uint64_t (*next_uint64)(void *st);
  uint32_t (*next_uint32)(void *st);
  double (*next_double)(void *st);
  uint64_t (*next_raw)(void *st);
} bitgen_t;
*/

void SetupBitGen(CustomRNG* rng, py::capsule capsule)
{
    bitgen_t* bg(capsule);
    bg->state = rng;
    bg->next_uint64 = next_uint64;
    bg->next_uint32 = next_uint32;
    bg->next_double = next_double;
    bg->next_raw = next_uint64;
};

3. Have pybind11 wrap this for you so you can call it from python. I'm also assuming that your CustomRNG type is already wrapped and accessible in python.

my_module.def("setup_bitgen", &SetupBitGen)

Or you could make this function a method of your CustomRNG class:

py::class_<CustomRNG>
    [...]
    .def("setup_bitgen", &SetupBitGen);

4. In Python make a class derived from numpy.random.BitGenerator, which takes an instance of your CustomRNG as an argument. For the __init__, first call super().__init__() to make the things numpy expects to be there. Then call lib.setup_bitgen(self.rng, self.capsule) (or self.rng.setup_bitgen(self.capsule) if you went the method route) to update the elements of the capsule.

That's it. Now you can make your own BitGenerator object using this class, and pass that as an argument of numpy.random.Generator.

This is the method I used in GalSim, so you can see a worked example of it here.