Allow setting the seed of random generators

atintegrators/TestRandomPass.c

@@ -1,4 +1,4 @@
-/* RandomPass.c
+ /* RandomPass.c
    Accelerator Toolbox
 
    Test of random generators
@@ -21,14 +21,15 @@ static void RandomPass(double *r_in,
         pcg32_random_t* thread_rng,
         int num_particles)
 {	
-    double common_val = atrandn_r(common_rng, 0.0, 0.001);
+    double common_val;
 #ifdef MPI
     int rank;
     MPI_Comm_rank(MPI_COMM_WORLD, &rank);
 #else
     int rank = 0;
 #endif /* MPI */
 
+    common_val = atrandn_r(common_rng, 0.0, 0.001);
     for (int c = 0; c<num_particles; c++) {	/*Loop over particles  */
         double *r6 = r_in+c*6;
         r6[0] = atrandn_r(thread_rng, 0.0, 0.001);
@@ -37,6 +38,7 @@ static void RandomPass(double *r_in,
         r6[5] = 0.01*c;
     }
 
+    common_val = atrandn_r(common_rng, 0.0, 0.001);
     #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  */

atmat/attrack/atpass.c

@@ -52,6 +52,7 @@ typedef double mxDouble;
 #define RINGPROPERTIES prhs[9]
 #define TURN prhs[10]
 #define KEEPCOUNTER prhs[11]
+#define SEED prhs[12]
 
 #define LIMIT_AMPLITUDE		1	/*  if any of the phase space variables (except the sixth N.C.) 
 									exceeds this limit it is marked as lost */
@@ -249,6 +250,7 @@ static mxDouble *passhook(mxArray *mxPassArg[], mxArray *mxElem, mxArray *func)
 @param[in]      [9] RINGPROPERTIES
 @param[in]     [10] TURN
 @param[in]     [11] KEEPCOUNTER
+@param[in]     [12] SEED
 */
 void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
 {
@@ -281,6 +283,7 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
     int num_particles = mxGetN(INITCONDITIONS);
     int counter = (nrhs >= 11) ? (int)mxGetScalar(TURN) : 0;
     int keep_counter = (nrhs >= 12) ? (int)mxGetScalar(KEEPCOUNTER) : 0;
+    int seed = (nrhs >= 13) ? (int)mxGetScalar(SEED) : -1;
     int np6 = num_particles*6;
     int ihist, lhist;
     mxDouble *histbuf = NULL;
@@ -306,7 +309,11 @@ void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
     param.nbunch = 1;
     param.num_turns = num_turns;
     param.bdiff = NULL;
-
+
+    if (seed >= 0) {
+        pcg32_srandom_r(&common_state, seed, AT_RNG_INC);
+        pcg32_srandom_r(&thread_state, seed, 0);
+    }
     if (keep_counter)
         param.nturn = last_turn;
     else

atmat/attrack/linepass.m

@@ -58,6 +58,12 @@
 % ROUT=LINEPASS(...,'reuse') is kept for compatibilty with previous
 % versions. It has no effect.
 %
+% ROUT=LINEPASS(...,'seed',SEED)  The random generators are reset to start 
+%   with SEED.
+%
+% ROUT=LINEPASS(...,'omp_num_threads',NTHREADS)  Number of OpenMP threads.
+%   By default, OpenMP chooses the number of threads.
+%
 % Rfin=LINEPASS(...,PREFUNC)
 % Rfin=LINEPASS(...,PREFUNC,POSTFUNC)
 % Rfin=LINEPASS(...,cell(0),POSTFUNC)
@@ -75,6 +81,7 @@
 [dummy,args]=getflag(args,'reuse');	%#ok<ASGLU> % Kept for compatibility and ignored
 [nhist,args]=getoption(args,'nhist',1);
 [omp_num_threads,args]=getoption(args,'omp_num_threads');
+[seed,args]=getoption(args,'seed',-1);
 funcargs=cellfun(@(arg) isa(arg,'function_handle'), args);
 refpts=getargs(args(~funcargs),length(line)+1);
 [prefunc,postfunc]=getargs(args(funcargs),cell(0),cell(0));
@@ -88,7 +95,7 @@
 
 try
     [Rout,lossinfo] = atpass(line,Rin,newlattice,1,refpts, ...
-        prefunc,postfunc,nhist,omp_num_threads,props,0);
+        prefunc,postfunc,nhist,omp_num_threads,props,0,0,seed);
 
     if nargout>1
         if nargout>2, varargout{2}=lossinfo; end

atmat/attrack/ringpass.m

@@ -48,7 +48,10 @@
 % ROUT=RINGPASS(...,'reuse') is kept for compatibilty with previous
 % versions. It has no effect.
 %
-% ROUT=RINGPASS(...,'turn',turn)    Initial turn number. Default 0.
+% ROUT=RINGPASS(...,'seed',SEED)  The random generators are reset to start 
+%   with SEED.
+%
+% ROUT=RINGPASS(...,'turn',TURN)    Initial turn number. Default 0.
 %   The turn number is necessary to compute the absolute path length used
 %   by RFCavityPass. Ignored if KeepCounter is set.
 %
@@ -60,6 +63,9 @@
 % results), one must use one of the 'turn' option or 'KeepCounter' flag to
 % ensure the continuity of the turn number.
 %
+% ROUT=RINGPASS(...,'omp_num_threads',NTHREADS)  Number of OpenMP threads.
+%   By default, OpenMP chooses the number of threads.
+%
 % ROUT=RINGPASS(...,'Silent') does not output the particle coordinates at
 %    each turn but only at the end of the tracking
 %
@@ -83,6 +89,7 @@
 [turn,args]=getoption(args,'turn',0);
 [keep_counter,args]=getflag(args,'KeepCounter');
 [omp_num_threads,args]=getoption(args,'omp_num_threads');
+[seed,args]=getoption(args,'seed',-1);
 funcargs=cellfun(@(arg) isa(arg,'function_handle'), args);
 nturns=getargs(args(~funcargs),1);
 [prefunc,postfunc]=getargs(args(funcargs),cell(0),cell(0));
@@ -99,7 +106,7 @@
 
 try
     [Rout,lossinfo] = atpass(ring,Rin,newlattice,nturns,refpts, ...
-        prefunc,postfunc,nhist,omp_num_threads,props,turn,double(keep_counter));
+        prefunc,postfunc,nhist,omp_num_threads,props,turn,double(keep_counter),seed);
 
     if nargout>1
         if nargout>3, varargout{3}=lossinfo; end

pyat/at/tracking/atpass.pyi

@@ -10,9 +10,9 @@ _defref = np.array([], dtype=np.uint32)
 
 def atpass(
     line: list[Element],
-    r_in: np.ndarray[np.float64],
+    r_in: np.ndarray,
     nturns: int,
-    refpts: np.ndarray[np.uint32] = _defref,
+    refpts: np.ndarray = _defref,
     turn: int | None = None,
     energy: float | None = None,
     particle: Particle | None = None,
@@ -25,13 +25,13 @@ def atpass(
 ): ...
 def elempass(
     element: Element,
-    r_in: np.ndarray[np.float64],
+    r_in: np.ndarray,
     energy: float | None = None,
     particle: Particle | None = None,
 ): ...
 def diffusion_matrix(
     element: Element,
-    r_in: np.ndarray[np.float64],
+    r_in: np.ndarray,
     energy: float | None = None,
     particle: Particle | None = None,
 ): ...