Skip to content
Permalink
master
Go to file
 
 
Cannot retrieve contributors at this time
478 lines (427 sloc) 12.4 KB
/*
* PRNG support.
*/
#ifdef _MSC_VER
#define HAVE_PTHREAD_ATFORK 0
#else
#define HAVE_PTHREAD_ATFORK 1
#include <pthread.h>
#endif
/* Magic Mersenne Twister constants */
#define MT_N 624
#define MT_M 397
#define MT_MATRIX_A 0x9908b0dfU
#define MT_UPPER_MASK 0x80000000U
#define MT_LOWER_MASK 0x7fffffffU
/*
* Note this structure is accessed in numba.targets.randomimpl,
* any changes here should be reflected there too.
*/
typedef struct {
int index;
/* unsigned int is sufficient on modern machines as we only need 32 bits */
unsigned int mt[MT_N];
int has_gauss;
double gauss;
int is_initialized;
} rnd_state_t;
/* Some code portions below from CPython's _randommodule.c, some others
from Numpy's and Jean-Sebastien Roy's randomkit.c. */
NUMBA_EXPORT_FUNC(void)
numba_rnd_shuffle(rnd_state_t *state)
{
int i;
unsigned int y;
for (i = 0; i < MT_N - MT_M; i++) {
y = (state->mt[i] & MT_UPPER_MASK) | (state->mt[i+1] & MT_LOWER_MASK);
state->mt[i] = state->mt[i+MT_M] ^ (y >> 1) ^
(-(int) (y & 1) & MT_MATRIX_A);
}
for (; i < MT_N - 1; i++) {
y = (state->mt[i] & MT_UPPER_MASK) | (state->mt[i+1] & MT_LOWER_MASK);
state->mt[i] = state->mt[i+(MT_M-MT_N)] ^ (y >> 1) ^
(-(int) (y & 1) & MT_MATRIX_A);
}
y = (state->mt[MT_N - 1] & MT_UPPER_MASK) | (state->mt[0] & MT_LOWER_MASK);
state->mt[MT_N - 1] = state->mt[MT_M - 1] ^ (y >> 1) ^
(-(int) (y & 1) & MT_MATRIX_A);
}
/* Initialize mt[] with an integer seed */
NUMBA_EXPORT_FUNC(void)
numba_rnd_init(rnd_state_t *state, unsigned int seed)
{
unsigned int pos;
seed &= 0xffffffffU;
/* Knuth's PRNG as used in the Mersenne Twister reference implementation */
for (pos = 0; pos < MT_N; pos++) {
state->mt[pos] = seed;
seed = (1812433253U * (seed ^ (seed >> 30)) + pos + 1) & 0xffffffffU;
}
state->index = MT_N;
state->has_gauss = 0;
state->gauss = 0.0;
state->is_initialized = 1;
}
/* Perturb mt[] with a key array */
static void
rnd_init_by_array(rnd_state_t *state, unsigned int init_key[], size_t key_length)
{
size_t i, j, k;
unsigned int *mt = state->mt;
numba_rnd_init(state, 19650218U);
i = 1; j = 0;
k = (MT_N > key_length ? MT_N : key_length);
for (; k; k--) {
mt[i] = (mt[i] ^ ((mt[i-1] ^ (mt[i-1] >> 30)) * 1664525U))
+ init_key[j] + (unsigned int) j; /* non linear */
mt[i] &= 0xffffffffU;
i++; j++;
if (i >= MT_N) { mt[0] = mt[MT_N - 1]; i = 1; }
if (j >= key_length) j = 0;
}
for (k = MT_N - 1; k; k--) {
mt[i] = (mt[i] ^ ((mt[i-1] ^ (mt[i-1] >> 30)) * 1566083941U))
- (unsigned int) i; /* non linear */
mt[i] &= 0xffffffffU;
i++;
if (i >= MT_N) { mt[0] = mt[MT_N - 1]; i=1; }
}
mt[0] = 0x80000000U; /* MSB is 1; ensuring non-zero initial array */
state->index = MT_N;
state->has_gauss = 0;
state->gauss = 0.0;
state->is_initialized = 1;
}
/*
* Management of thread-local random state.
*/
static int rnd_globally_initialized;
#ifdef _MSC_VER
#define THREAD_LOCAL(ty) __declspec(thread) ty
#else
/* Non-standard C99 extension that's understood by gcc and clang */
#define THREAD_LOCAL(ty) __thread ty
#endif
static THREAD_LOCAL(rnd_state_t) numba_py_random_state;
static THREAD_LOCAL(rnd_state_t) numba_np_random_state;
/* Seed the state with random bytes */
static int
rnd_seed_with_bytes(rnd_state_t *state, Py_buffer *buf)
{
unsigned int *keys;
unsigned char *bytes;
size_t i, nkeys;
nkeys = buf->len / sizeof(unsigned int);
keys = (unsigned int *) PyMem_Malloc(nkeys * sizeof(unsigned int));
if (keys == NULL) {
PyBuffer_Release(buf);
return -1;
}
bytes = (unsigned char *) buf->buf;
/* Convert input bytes to int32 keys, without violating alignment
* constraints.
*/
for (i = 0; i < nkeys; i++, bytes += 4) {
keys[i] = (bytes[3] << 24) + (bytes[2] << 16) +
(bytes[1] << 8) + (bytes[0] << 0);
}
PyBuffer_Release(buf);
rnd_init_by_array(state, keys, nkeys);
PyMem_Free(keys);
return 0;
}
#if HAVE_PTHREAD_ATFORK
/* After a fork(), the child should reseed its random states.
* Since only the main thread survives in the child, it's enough to mark
* the current thread-local states as uninitialized.
*/
static void
rnd_atfork_child(void)
{
numba_py_random_state.is_initialized = 0;
numba_np_random_state.is_initialized = 0;
}
#endif
/* Global initialization routine. It must be called as early as possible.
*/
NUMBA_EXPORT_FUNC(void)
numba_rnd_ensure_global_init(void)
{
if (!rnd_globally_initialized) {
#if HAVE_PTHREAD_ATFORK
pthread_atfork(NULL, NULL, rnd_atfork_child);
#endif
numba_py_random_state.is_initialized = 0;
numba_np_random_state.is_initialized = 0;
rnd_globally_initialized = 1;
}
}
/* First-time init a random state */
static void
rnd_implicit_init(rnd_state_t *state)
{
/* Initialize with random bytes. The easiest way to get good-quality
* cross-platform random bytes is still to call os.urandom()
* using the Python interpreter...
*/
PyObject *module, *bufobj;
Py_buffer buf;
PyGILState_STATE gilstate = PyGILState_Ensure();
module = PyImport_ImportModuleNoBlock("os");
if (module == NULL)
goto error;
/* Read as many bytes as necessary to get the full entropy
* exploitable by the MT generator.
*/
bufobj = PyObject_CallMethod(module, "urandom", "i",
(int) (MT_N * sizeof(unsigned int)));
Py_DECREF(module);
if (bufobj == NULL)
goto error;
if (PyObject_GetBuffer(bufobj, &buf, PyBUF_SIMPLE))
goto error;
Py_DECREF(bufobj);
if (rnd_seed_with_bytes(state, &buf))
goto error;
/* state->is_initialized is set now */
PyGILState_Release(gilstate);
return;
error:
/* In normal conditions, os.urandom() and PyMem_Malloc() shouldn't fail,
* and we don't want the caller to deal with errors, so just bail out.
*/
if (PyErr_Occurred())
PyErr_Print();
Py_FatalError(NULL);
}
/* Functions returning the thread-local random state pointer.
* The LLVM JIT doesn't support thread-local variables so we rely
* on the C compiler instead.
*/
NUMBA_EXPORT_FUNC(rnd_state_t *)
numba_get_py_random_state(void)
{
rnd_state_t *state = &numba_py_random_state;
if (!state->is_initialized)
rnd_implicit_init(state);
return state;
}
NUMBA_EXPORT_FUNC(rnd_state_t *)
numba_get_np_random_state(void)
{
rnd_state_t *state = &numba_np_random_state;
if (!state->is_initialized)
rnd_implicit_init(state);
return state;
}
/*
* Python-exposed helpers for state management and testing.
*/
static int
rnd_state_converter(PyObject *obj, rnd_state_t **state)
{
*state = (rnd_state_t *) PyLong_AsVoidPtr(obj);
return (*state != NULL || !PyErr_Occurred());
}
NUMBA_EXPORT_FUNC(PyObject *)
_numba_rnd_get_py_state_ptr(PyObject *self)
{
return PyLong_FromVoidPtr(numba_get_py_random_state());
}
NUMBA_EXPORT_FUNC(PyObject *)
_numba_rnd_get_np_state_ptr(PyObject *self)
{
return PyLong_FromVoidPtr(numba_get_np_random_state());
}
NUMBA_EXPORT_FUNC(PyObject *)
_numba_rnd_shuffle(PyObject *self, PyObject *arg)
{
rnd_state_t *state;
if (!rnd_state_converter(arg, &state))
return NULL;
numba_rnd_shuffle(state);
Py_RETURN_NONE;
}
NUMBA_EXPORT_FUNC(PyObject *)
_numba_rnd_set_state(PyObject *self, PyObject *args)
{
int i, index;
rnd_state_t *state;
PyObject *tuplearg, *intlist;
if (!PyArg_ParseTuple(args, "O&O!:rnd_set_state",
rnd_state_converter, &state,
&PyTuple_Type, &tuplearg))
return NULL;
if (!PyArg_ParseTuple(tuplearg, "iO!", &index, &PyList_Type, &intlist))
return NULL;
if (PyList_GET_SIZE(intlist) != MT_N) {
PyErr_SetString(PyExc_ValueError, "list object has wrong size");
return NULL;
}
state->index = index;
for (i = 0; i < MT_N; i++) {
PyObject *v = PyList_GET_ITEM(intlist, i);
unsigned long x = PyLong_AsUnsignedLong(v);
if (x == (unsigned long) -1 && PyErr_Occurred())
return NULL;
state->mt[i] = (unsigned int) x;
}
state->has_gauss = 0;
state->gauss = 0.0;
state->is_initialized = 1;
Py_RETURN_NONE;
}
NUMBA_EXPORT_FUNC(PyObject *)
_numba_rnd_get_state(PyObject *self, PyObject *arg)
{
PyObject *intlist;
int i;
rnd_state_t *state;
if (!rnd_state_converter(arg, &state))
return NULL;
intlist = PyList_New(MT_N);
if (intlist == NULL)
return NULL;
for (i = 0; i < MT_N; i++) {
PyObject *v = PyLong_FromUnsignedLong(state->mt[i]);
if (v == NULL) {
Py_DECREF(intlist);
return NULL;
}
PyList_SET_ITEM(intlist, i, v);
}
return Py_BuildValue("iN", state->index, intlist);
}
NUMBA_EXPORT_FUNC(PyObject *)
_numba_rnd_seed(PyObject *self, PyObject *args)
{
unsigned int seed;
rnd_state_t *state;
if (!PyArg_ParseTuple(args, "O&I:rnd_seed",
rnd_state_converter, &state, &seed)) {
/* rnd_seed_*(bytes-like object) */
Py_buffer buf;
PyErr_Clear();
if (!PyArg_ParseTuple(args, "O&s*:rnd_seed",
rnd_state_converter, &state, &buf))
return NULL;
if (rnd_seed_with_bytes(state, &buf))
return NULL;
else
Py_RETURN_NONE;
}
else {
/* rnd_seed_*(int32) */
numba_rnd_init(state, seed);
Py_RETURN_NONE;
}
}
/*
* Random distribution helpers.
* Most code straight from Numpy's distributions.c.
*/
#ifndef M_PI
#define M_PI 3.14159265358979323846264338328
#endif
NUMBA_EXPORT_FUNC(unsigned int)
get_next_int32(rnd_state_t *state)
{
unsigned int y;
if (state->index == MT_N) {
numba_rnd_shuffle(state);
state->index = 0;
}
y = state->mt[state->index++];
/* Tempering */
y ^= (y >> 11);
y ^= (y << 7) & 0x9d2c5680U;
y ^= (y << 15) & 0xefc60000U;
y ^= (y >> 18);
return y;
}
NUMBA_EXPORT_FUNC(double)
get_next_double(rnd_state_t *state)
{
double a = get_next_int32(state) >> 5;
double b = get_next_int32(state) >> 6;
return (a * 67108864.0 + b) / 9007199254740992.0;
}
NUMBA_EXPORT_FUNC(double)
loggam(double x)
{
double x0, x2, xp, gl, gl0;
long k, n;
static double a[10] = {8.333333333333333e-02,-2.777777777777778e-03,
7.936507936507937e-04,-5.952380952380952e-04,
8.417508417508418e-04,-1.917526917526918e-03,
6.410256410256410e-03,-2.955065359477124e-02,
1.796443723688307e-01,-1.39243221690590e+00};
x0 = x;
n = 0;
if ((x == 1.0) || (x == 2.0))
{
return 0.0;
}
else if (x <= 7.0)
{
n = (long)(7 - x);
x0 = x + n;
}
x2 = 1.0/(x0*x0);
xp = 2*M_PI;
gl0 = a[9];
for (k=8; k>=0; k--)
{
gl0 *= x2;
gl0 += a[k];
}
gl = gl0/x0 + 0.5*log(xp) + (x0-0.5)*log(x0) - x0;
if (x <= 7.0)
{
for (k=1; k<=n; k++)
{
gl -= log(x0-1.0);
x0 -= 1.0;
}
}
return gl;
}
NUMBA_EXPORT_FUNC(int64_t)
numba_poisson_ptrs(rnd_state_t *state, double lam)
{
/* This method is invoked only if the parameter lambda of this
* distribution is big enough ( >= 10 ). The algorithm used is
* described in "Hörmann, W. 1992. 'The Transformed Rejection
* Method for Generating Poisson Random Variables'.
* The implementation comes straight from Numpy.
*/
int64_t k;
double U, V, slam, loglam, a, b, invalpha, vr, us;
slam = sqrt(lam);
loglam = log(lam);
b = 0.931 + 2.53*slam;
a = -0.059 + 0.02483*b;
invalpha = 1.1239 + 1.1328/(b-3.4);
vr = 0.9277 - 3.6224/(b-2);
while (1)
{
U = get_next_double(state) - 0.5;
V = get_next_double(state);
us = 0.5 - fabs(U);
k = (int64_t) floor((2*a/us + b)*U + lam + 0.43);
if ((us >= 0.07) && (V <= vr))
{
return k;
}
if ((k < 0) ||
((us < 0.013) && (V > us)))
{
continue;
}
if ((log(V) + log(invalpha) - log(a/(us*us)+b)) <=
(-lam + (double) k*loglam - loggam((double) k+1)))
{
return k;
}
}
}
You can’t perform that action at this time.