Rit: squashed commits

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)
+    x = xs if xs is not None else np.linspace(-wx,wx,Nx)
-    y = np.linspace(-wy,wy,Ny)
+    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, axs = plt.subplots(1,figsize=figsize)
-    fig.suptitle(r'E = %.1f EeV, $\theta$ = %.1f$^\circ$, $\phi$ = %.1f$^\circ$ X = %.f'%(e.energy,e.zenith,e.azimuth,X))
+    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.plot(xc/1e3,yc/1e3,'r+', **crosshair_kwargs)
-        axs.set_xlabel('x (km)')
+        axs.set_xlabel('x [km]')
-        axs.set_ylabel('y (km)')
+        axs.set_ylabel('y [km]')
     elif mode == "sp":
-        axs.plot(0,0,'r+',ms=30)
+        axs.plot(0,0, 'r+', **crosshair_kwargs)
-        axs.set_xlabel('-v x B (km)')
+        axs.set_xlabel('-v x B [km]')
-        axs.set_ylabel(' vxvxB (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)