ZH: rit multiprocessing shower_plane_slice

simulations/airshower_beacon_simulation/lib/rit.py

@@ -85,26 +85,31 @@ def shower_axis_slice(e,Xb=200,Xe=1200,dX=2,zgr=0):
     p = np.asanyarray(p)
     return ds,Xs,locs,p
 
-def shower_plane_slice(e,X=750.,Nx=10,Ny=10,wx=1e3,wy=1e3,xoff=0,yoff=0,zgr=0):
+def shower_plane_slice(e,X=750.,Nx=10,Ny=10,wx=1e3,wy=1e3,xoff=0,yoff=0,zgr=0,n_jobs=None):
     zgr = zgr + e.core[2]
     dX = atm.distance_to_slant_depth(np.deg2rad(e.zenith),X,zgr)
     x = np.linspace(-wx,wx,Nx)
     y = np.linspace(-wy,wy,Ny)
-    xx = []
-    yy = []
-    p = []
-    locs = []
-    for x_ in x:
-        for y_ in y:
+
+    def loop_func(x_, y_, xoff=xoff, yoff=yoff):
             loc = (x_+xoff)* e.uAxB + (y_+yoff)*e.uAxAxB + dX *e.uA
-            locs.append(loc)
             P,t_,pulses_,wav,twav = pow_and_time(loc,e)
-            xx.append(x_+xoff)
-            yy.append(y_+yoff)
-            p.append(P)
+
+            return x_+xoff, y_+yoff, P, loc
+
+    res = ( delayed(loop_func)(x_, y_) for x_ in x for y_ in y)
+
+    if Parallel:
+        #if n_jobs is None change with `with parallel_backend`
+        res = Parallel(n_jobs=n_jobs)(res)
+
+    # unpack loop results
+    xx, yy, p, locs = zip(*res)
+
     xx = np.asarray(xx)
     yy = np.asarray(yy)
     p = np.asanyarray(p)
+
     return xx,yy,p,locs[np.argmax(p)]
 
 def slice_figure(e,X,xx,yy,p,mode='horizontal'):