Rit: squashed commits

This commit is contained in:
Eric Teunis de Boone 2023-11-03 17:18:57 +01:00
parent 42236b03a8
commit 8e5a8b608d

View file

@ -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)