diff --git a/airshower_beacon_simulation/lib/rit.py b/airshower_beacon_simulation/lib/rit.py index 821a4ef..5f025e0 100644 --- a/airshower_beacon_simulation/lib/rit.py +++ b/airshower_beacon_simulation/lib/rit.py @@ -7,9 +7,14 @@ from scipy.signal import hilbert from scipy import signal from scipy.interpolate import interp1d from scipy.optimize import curve_fit,minimize -import pandas as pd +#import pandas as pd import os +try: + from tqdm import tqdm +except: + tqdm = lambda x: x + try: from joblib import Parallel, delayed except: @@ -113,11 +118,11 @@ 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,n_jobs=None): +def shower_plane_slice(e,X=750.,Nx=10,Ny=10,wx=1e3,wy=1e3,xoff=0,yoff=0,zgr=0,n_jobs=None, xs=None, ys=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) + x = xs if xs is not None else np.linspace(-wx,wx,Nx) + y = ys if ys is not None else np.linspace(-wy,wy,Ny) def loop_func(x_, y_, xoff=xoff, yoff=yoff): loc = (x_+xoff)* e.uAxB + (y_+yoff)*e.uAxAxB + dX *e.uA @@ -140,7 +145,7 @@ def shower_plane_slice(e,X=750.,Nx=10,Ny=10,wx=1e3,wy=1e3,xoff=0,yoff=0,zgr=0,n_ return xx,yy,p,locs[np.argmax(p)] -def slice_figure(e,X,xx,yy,p,mode='horizontal', scatter_kwargs={}, colorbar_kwargs={'label':'Power'}): +def slice_figure(e,X,xx,yy,p,mode='horizontal', scatter_kwargs={}, colorbar_kwargs={'label':'Power'},suptitle=True, figsize=(10,8)): scatter_kwargs = { **dict( cmap='Spectral_r', @@ -149,26 +154,39 @@ def slice_figure(e,X,xx,yy,p,mode='horizontal', scatter_kwargs={}, colorbar_kwar ), **scatter_kwargs } + crosshair_kwargs = dict(ms=30, mew=3) - fig, axs = plt.subplots(1,figsize=(10,8)) - fig.suptitle(r'E = %.1f EeV, $\theta$ = %.1f$^\circ$, $\phi$ = %.1f$^\circ$ X = %.f'%(e.energy,e.zenith,e.azimuth,X)) + fig, axs = plt.subplots(1,figsize=figsize) + if suptitle: + fig.suptitle(r'E = %.1g PeV, $\theta$ = %.1f$^\circ$, $\phi$ = %.1f$^\circ$ X = %.f'%(e.energy*10**3,e.zenith,e.azimuth,X)) sc = axs.scatter(xx/1e3,yy/1e3,c=p,**scatter_kwargs) - fig.colorbar(sc,ax=axs, **colorbar_kwargs) + cbar = fig.colorbar(sc,ax=axs, **colorbar_kwargs) zgr = 0 + e.core[2] dX = atm.distance_to_slant_depth(np.deg2rad(e.zenith),X,zgr) xc = np.sin(np.deg2rad(e.zenith))*np.cos(np.deg2rad(e.azimuth))* dX yc = np.sin(np.deg2rad(e.zenith))*np.sin(np.deg2rad(e.azimuth))* dX if mode == 'horizontal': - axs.plot(xc/1e3,yc/1e3,'r+',ms=30) - axs.set_xlabel('x (km)') - axs.set_ylabel('y (km)') + axs.plot(xc/1e3,yc/1e3,'r+', **crosshair_kwargs) + axs.set_xlabel('x [km]') + axs.set_ylabel('y [km]') elif mode == "sp": - axs.plot(0,0,'r+',ms=30) - axs.set_xlabel('-v x B (km)') - axs.set_ylabel(' vxvxB (km)') + axs.plot(0,0, 'r+', **crosshair_kwargs) + axs.set_xlabel('-v x B [km]') + axs.set_ylabel(' vxvxB [km]') + + # indicate maximum power im = np.argmax(p) - axs.plot(xx[im]/1e3,yy[im]/1e3,'bx',ms=30) + im_color = 'blue' + im_norm = cbar.norm(p[im]) + if im_norm < 0.4: + im_color = '#4488FF' + if True: + cbar.add_lines(levels=[p[im]], colors=[im_color], linewidths=[crosshair_kwargs['mew']]) + #cbar.lines[-1].set_linestyles('dotted') + + axs.plot(xx[im]/1e3,yy[im]/1e3, 'x', color=im_color, **crosshair_kwargs) + #cbar.ax.axhline(cbar.norm(p[im]), color='b') fig.tight_layout() return fig,axs @@ -195,7 +213,7 @@ def dist_to_line_sum(param,data,weights): # print('%.2e %.2e %.2e %.2e %.2e'%(x0,y0,theta,phi,dsum)) return dsum/len(data) -def get_axis_points(e,savefig=True,path="",zgr=0,Xlow=300, Xhigh=1000, N_X=15, n_jobs=None): +def get_axis_points(e,savefig=True,path="",zgr=0,Xlow=300, Xhigh=1000, N_X=15, n_jobs=None, tqdm=tqdm): 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) @@ -227,7 +245,7 @@ def get_axis_points(e,savefig=True,path="",zgr=0,Xlow=300, Xhigh=1000, N_X=15, n print("Finished", X) return np.max(p), loc_max - res = (delayed(loop_func)(X) for X in Xsteps) + res = tqdm((delayed(loop_func)(X) for X in Xsteps), total=len(Xsteps)) if Parallel: #if n_jobs is None change with `with parallel_backend` @@ -376,7 +394,7 @@ def fill_stations_propeties(e,res): res.station_ids.append(ids) #res.has_pulse.append(has_pulse) -def reconstruction(e,outfile='', slice_outdir=None, Xlow=300, Xhigh=1000, N_X=15, disable_pol_and_bp=False): +def reconstruction(e,outfile='', slice_outdir=None, Xlow=300, Xhigh=1000, N_X=15, disable_pol_and_bp=False, tqdm=tqdm): res = RITResult() res.isMC.append(True) res.zenith_ini.append(e.zenith) @@ -388,7 +406,7 @@ def reconstruction(e,outfile='', slice_outdir=None, Xlow=300, Xhigh=1000, N_X=15 #only use signal that have a signal in data fill_stations_propeties(e,res) - Xs,axis_points,max_vals = get_axis_points(e,savefig=(slice_outdir is not None), path=slice_outdir, Xlow=Xlow, Xhigh=Xhigh, N_X=N_X) + Xs,axis_points,max_vals = get_axis_points(e,savefig=(slice_outdir is not None), path=slice_outdir, Xlow=Xlow, Xhigh=Xhigh, N_X=N_X, tqdm=tqdm) zen,azi,core = fit_track(e,axis_points,max_vals,1e2) fig = figure_3D(axis_points,max_vals,zen,azi,core,res) fig.savefig(outfile)