From 3c99ac711850715431515de1600e531dc2e3b268 Mon Sep 17 00:00:00 2001 From: Eric Teunis de Boone Date: Wed, 11 Jan 2023 19:43:48 +0100 Subject: [PATCH] ZH: rit attempt multiprocessing in get_axis_points --- .../airshower_beacon_simulation/lib/rit.py | 30 +++++++++++++++---- 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/simulations/airshower_beacon_simulation/lib/rit.py b/simulations/airshower_beacon_simulation/lib/rit.py index b68493e..3486a97 100644 --- a/simulations/airshower_beacon_simulation/lib/rit.py +++ b/simulations/airshower_beacon_simulation/lib/rit.py @@ -10,6 +10,11 @@ from scipy.optimize import curve_fit,minimize import pandas as pd import os +try: + from joblib import Parallel, delayed +except: + Parallel = None + plt.rcParams.update({'font.size': 16}) atm = AtmoCal() @@ -147,16 +152,15 @@ def dist_to_line_sum(param,data,weights): return dsum/len(data) def get_axis_points(e,savefig=True,path="",zgr=0,Xlow=300, Xhigh=1000, N_X=15): - axis_points = [] - max_vals = [] Xsteps = np.linspace(Xlow, Xhigh, N_X) zgr=zgr+e.core[2] #not exact dXref = atm.distance_to_slant_depth(np.deg2rad(e.zenith),750,zgr) scale2d = dXref*np.tan(np.deg2rad(2.)) scale4d = dXref*np.tan(np.deg2rad(4.)) scale0_2d=dXref*np.tan(np.deg2rad(0.2)) - for X in Xsteps: - print(X) + + def loop_func(X): + print("Starting", X) x,y,p,loc_max = shower_plane_slice(e,X,21,21,scale2d,scale2d) if savefig: fig,axs = slice_figure(e,X,x,y,p,'sp') @@ -175,8 +179,22 @@ def get_axis_points(e,savefig=True,path="",zgr=0,Xlow=300, Xhigh=1000, N_X=15): fig,axs = slice_figure(e,X,x,y,p,'sp') fig.savefig(path+'X%d_b.pdf'%(X)) plt.close(fig) - max_vals.append(np.max(p)) - axis_points.append(loc_max) + + print("Finished", X) + return np.max(p), loc_max + + if Parallel: + n_jobs= None # if None change with `with parallel_backend` + res = Parallel(n_jobs=n_jobs)(delayed(loop_func)(X) for X in Xsteps) + + max_vals = [ item[0] for item in res ] + axis_points = [ item[1] for item in res ] + + else: + for X in Xsteps: + _max_val, _loc_max = loop_func(X) + max_vals.append(_max_val) + axis_points.append(_loc_max) return Xsteps,axis_points,max_vals