Correct random generators for MPI (#737)

* correct non-reentrant gaussian generator

atintegrators/TestRandomPass.c

@@ -7,6 +7,9 @@
 #include "atelem.c"
 #include "atlalib.c"
 #include "atrandom.c"
+#ifdef MPI
+#include <mpi.h>
+#endif
 
 struct elem 
 {
@@ -19,22 +22,26 @@ static void RandomPass(double *r_in,
         int num_particles)
 {	
     double common_val = atrandn_r(common_rng, 0.0, 0.001);
+#ifdef MPI
+    int rank;
+    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+#else
+    int rank = 0;
+#endif /* MPI */
 
     for (int c = 0; c<num_particles; c++) {	/*Loop over particles  */
         double *r6 = r_in+c*6;
-        double thread_val = atrandn_r(thread_rng, 0.0, 0.001);
-        r6[0] = thread_val;
+        r6[0] = atrandn_r(thread_rng, 0.0, 0.001);
         r6[2] = common_val;
-        r6[4] = 0.0;
-        r6[5] = 0.0;
+        r6[4] = 0.01*rank;
+        r6[5] = 0.01*c;
     }
 
     #pragma omp parallel for if (num_particles > OMP_PARTICLE_THRESHOLD) default(none)                      \
     shared(r_in, num_particles, common_val, thread_rng)
     for (int c = 0; c<num_particles; c++) {	/*Loop over particles  */
         double *r6 = r_in+c*6;
-        double thread_val = atrandn_r(thread_rng, 0.0, 0.001);
-        r6[1] = thread_val;
+        r6[1] = atrandn_r(thread_rng, 0.0, 0.001);;
         r6[3] = common_val;
     }
 }

atintegrators/atrandom.c

@@ -27,19 +27,21 @@
 #define AT_RNG_STATE 0x853c49e6748fea9bULL
 #define AT_RNG_INC 0xda3e39cb94b95bdbULL
 
-#define COMMON_PCG32_INITIALIZER   { AT_RNG_STATE, AT_RNG_INC }
-#define THREAD_PCG32_INITIALIZER   { AT_RNG_STATE, 1ULL }
+#define COMMON_PCG32_INITIALIZER   { AT_RNG_STATE, AT_RNG_INC, 0.0, false }
+#define THREAD_PCG32_INITIALIZER   { AT_RNG_STATE, 1ULL, 0.0, false }
 
 struct pcg_state_setseq_64 {    // Internals are *Private*.
     uint64_t state;             // RNG state.  All values are possible.
     uint64_t inc;               // Controls which RNG sequence (stream) is
                                 // selected. Must *always* be odd.
+    double spare;               // spare value for normal distribution
+    bool hasSpare;
 };
 typedef struct pcg_state_setseq_64 pcg32_random_t;
 
 // If you *must* statically initialize it, here's one.
 
-#define PCG32_INITIALIZER   { AT_RNG_STATE, AT_RNG_INC }
+#define PCG32_INITIALIZER   { AT_RNG_STATE, AT_RNG_INC, 0.0, false }
 
 static pcg32_random_t pcg32_global = PCG32_INITIALIZER;
 
@@ -99,24 +101,22 @@ static double atrandn_r(pcg32_random_t* rng, double mean, double stdDev)
 {
     /* Marsaglia polar method: https://en.wikipedia.org/wiki/Marsaglia_polar_method */
 
-	static bool hasSpare = false;
-	static double spare;
 	double u, v, s;
 
-	if (hasSpare) {
-		hasSpare = false;
-		return mean + stdDev * spare;
+	if (rng->hasSpare) {
+		rng->hasSpare = false;
+		return mean + stdDev * rng->spare;
 	}
 
-	hasSpare = true;	
+	rng->hasSpare = true;
 	do {
 		u = 2.0 * atrandd_r(rng) - 1.0;
 		v = 2.0 * atrandd_r(rng) - 1.0;
 		s = u * u + v * v;
 	}
 	while ((s >= 1.0) || (s == 0.0));
 	s = sqrt(-2.0 * log(s) / s);
-	spare = v * s;
+	rng->spare = v * s;
 	return mean + stdDev * u * s;
 }
 

pyat/at.c

@@ -9,6 +9,9 @@
 #include <string.h>
 #include <omp.h>
 #endif /*_OPENMP*/
+#ifdef MPI
+#include <mpi.h>
+#endif /* MPI */
 #include "attypes.h"
 #include <stdbool.h> 
 #include <math.h>
@@ -694,10 +697,16 @@ static PyObject *at_elempass(PyObject *self, PyObject *args, PyObject *kwargs)
 static PyObject *reset_rng(PyObject *self, PyObject *args, PyObject *kwargs)
 {
     static char *kwlist[] = {"rank", "seed", NULL};
-    uint64_t rank = 0;
     uint64_t seed = AT_RNG_STATE;
+#ifdef MPI
+    int irank;
+    MPI_Comm_rank(MPI_COMM_WORLD, &irank);
+    uint64_t rank = irank;
+#else
+    uint64_t rank = 0;
+#endif /* MPI */
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|K$K", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|$KK", kwlist,
         &rank, &seed)) {
         return NULL;
     }
@@ -756,11 +765,13 @@ static PyMethodDef AtMethods[] = {
               ":meta private:"
             )},
     {"reset_rng",  (PyCFunction)reset_rng, METH_VARARGS | METH_KEYWORDS,
-    PyDoc_STR("reset_rng(rank=0, seed=None)\n\n"
+    PyDoc_STR("reset_rng(*, rank=0, seed=None)\n\n"
               "Reset the *common* and *thread* random generators.\n\n"
+              "The seed is applied unchanged to the \"common\" generator, and modified in a\n"
+              "thread-specific way to the \"thread\" generator\n\n"
               "Parameters:\n"
               "    rank (int):    thread identifier (for MPI and python multiprocessing)\n"
-              "    seed (int):    single seed for both generators\n"
+              "    seed (int):    single seed for both generators. Default: initial seed\n"
             )},
     {"common_rng",  (PyCFunction)common_rng, METH_NOARGS,
     PyDoc_STR("common_rng()\n\n"

pyat/at/tracking/__init__.py

@@ -1,11 +1,10 @@
 """
 Tracking functions
 """
-from ..lattice import DConstant
 from .atpass import reset_rng, common_rng, thread_rng
 from .track import *
 from .particles import *
 from .utils import *
 from .deprecated import *
 # initialise the C random generators
-reset_rng(DConstant.rank)
+reset_rng()

pyat/at/tracking/atpass.pyi

@@ -20,6 +20,6 @@ def elempass(element: Element, r_in,
              particle: Optional[Particle] = None,
              ): ...
 
-def reset_rng(rank: int = 0, seed: Optional[int] = None) -> None: ...
+def reset_rng(*, rank: int = 0, seed: Optional[int] = None) -> None: ...
 def common_rng() -> float: ...
 def thread_rng() -> float: ...

pyat/at/tracking/track.py

@@ -23,14 +23,14 @@
 
 def _atpass_fork(seed, rank, rin, **kwargs):
     """Single forked job"""
-    reset_rng(rank, seed=seed)
+    reset_rng(rank=rank, seed=seed)
     result = _atpass(_globring, rin, **kwargs)
     return rin, result
 
 
 def _atpass_spawn(ring, seed, rank, rin, **kwargs):
     """Single spawned job"""
-    reset_rng(rank, seed=seed)
+    reset_rng(rank=rank, seed=seed)
     result = _atpass(ring, rin, **kwargs)
     return rin, result