ZH: rit multiprocessing shower_plane_slice

This commit is contained in:
Eric Teunis de Boone 2023-01-13 18:48:31 +01:00
parent d1700da445
commit c742926c4d

View file

@ -85,26 +85,31 @@ def shower_axis_slice(e,Xb=200,Xe=1200,dX=2,zgr=0):
p = np.asanyarray(p) p = np.asanyarray(p)
return ds,Xs,locs,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] zgr = zgr + e.core[2]
dX = atm.distance_to_slant_depth(np.deg2rad(e.zenith),X,zgr) dX = atm.distance_to_slant_depth(np.deg2rad(e.zenith),X,zgr)
x = np.linspace(-wx,wx,Nx) x = np.linspace(-wx,wx,Nx)
y = np.linspace(-wy,wy,Ny) y = np.linspace(-wy,wy,Ny)
xx = []
yy = [] def loop_func(x_, y_, xoff=xoff, yoff=yoff):
p = []
locs = []
for x_ in x:
for y_ in y:
loc = (x_+xoff)* e.uAxB + (y_+yoff)*e.uAxAxB + dX *e.uA 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) P,t_,pulses_,wav,twav = pow_and_time(loc,e)
xx.append(x_+xoff)
yy.append(y_+yoff) return x_+xoff, y_+yoff, P, loc
p.append(P)
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) xx = np.asarray(xx)
yy = np.asarray(yy) yy = np.asarray(yy)
p = np.asanyarray(p) p = np.asanyarray(p)
return xx,yy,p,locs[np.argmax(p)] return xx,yy,p,locs[np.argmax(p)]
def slice_figure(e,X,xx,yy,p,mode='horizontal'): def slice_figure(e,X,xx,yy,p,mode='horizontal'):