diff --git a/simulations/airshower_beacon_simulation/lib/rit.py b/simulations/airshower_beacon_simulation/lib/rit.py index 68e3dc6..41fc6e5 100644 --- a/simulations/airshower_beacon_simulation/lib/rit.py +++ b/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'):