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.
- 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.
- 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;
};
- 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);
- 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.
