diff --git a/simulations/airshower_beacon_simulation/lib/__init__.py b/simulations/airshower_beacon_simulation/lib/__init__.py new file mode 100644 index 0000000..b2d03e1 --- /dev/null +++ b/simulations/airshower_beacon_simulation/lib/__init__.py @@ -0,0 +1,2 @@ +from .lib import * +from . import rit diff --git a/simulations/airshower_beacon_simulation/lib/rit.py b/simulations/airshower_beacon_simulation/lib/rit.py new file mode 100644 index 0000000..9f06cd1 --- /dev/null +++ b/simulations/airshower_beacon_simulation/lib/rit.py @@ -0,0 +1,346 @@ +# from wappy import * +from earsim import * +from atmocal import * + +import matplotlib.pyplot as plt +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 os + +plt.rcParams.update({'font.size': 16}) + +atm = AtmoCal() + +from matplotlib import cm + +def set_pol_and_bp(e,low=0.03,high=0.08): + for ant in e.antennas: + E = [np.dot(e.uAxB,[ex,ey,ez]) for ex,ey,ez in zip(ant.Ex,ant.Ey,ant.Ez)] + dt = ant.t[1] -ant.t[0] + E = block_filter(E,dt,low,high) + ant.E_AxB = E + ant.t_AxB = ant.t + + +def pow_and_time(test_loc,ev,dt=1.0): + t_ = [] + a_ = [] + t_min = 1e9 + t_max = -1e9 + for ant in ev.antennas: + #propagate to test location + aloc = [ant.x,ant.y,ant.z] + delta,dist = atm.light_travel_time(test_loc,aloc) + delta = delta*1e9 + t__ = np.subtract(ant.t_AxB,delta) + t_.append(t__) + a_.append(ant.E_AxB) + if t__[0] < t_min: + t_min = t__[0] + if t__[-1] > t_max: + t_max = t__[-1] + + t_sum = np.arange(t_min+1,t_max-1,dt) + a_sum = np.zeros(len(t_sum)) + #interpolation + for t_r,E_ in zip (t_,a_): + f = interp1d(t_r,E_,assume_sorted=True,bounds_error=False,fill_value=0.) + a_int = f(t_sum) + a_sum = np.add(a_sum,a_int) + if len(a_sum) != 0: + P = np.sum(np.square(np.absolute(np.fft.fft(a_sum)))) + else: + print("ERROR, a_sum lenght = 0", + "tmin ",t_min, + "t_max ",t_max, + "dt",dt) + P = 0 + return P,t_,a_,a_sum,t_sum + +def shower_axis_slice(e,Xb=200,Xe=1200,dX=2,zgr=1400): + zgr = zgr + e.core[2] + N = int((Xe-Xb)/dX) + Xs = np.array(np.linspace(Xb,Xe,N+1)) + ds = np.array([atm.distance_to_slant_depth(np.deg2rad(e.zenith),X,zgr) for X in Xs]) + + locs = [] + for d_ in ds: + xc = np.sin(np.deg2rad(e.zenith))*np.cos(np.deg2rad(e.azimuth))* d_ + yc = np.sin(np.deg2rad(e.zenith))*np.sin(np.deg2rad(e.azimuth))* d_ + zc = np.cos(np.deg2rad(e.zenith))* d_ + locs.append([xc,yc,zc]) + p = [] + for loc in locs: + P,t_,pulses_,wav,twav = pow_and_time(loc,e) + p.append(P) + 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=1400): + 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) + xx = [] + yy = [] + p = [] + locs = [] + for x_ in x: + for y_ in y: + 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) + xx.append(x_+xoff) + yy.append(y_+yoff) + p.append(P) + xx = np.asarray(xx) + yy = np.asarray(yy) + p = np.asanyarray(p) + return xx,yy,p,locs[np.argmax(p)] + +def slice_figure(e,X,xx,yy,p,mode='horizontal'): + 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)) + sc = axs.scatter(xx/1e3,yy/1e3,c=p,cmap='Spectral_r',alpha=0.6) + fig.colorbar(sc,ax=axs) + zgr = 1400 + 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)') + elif mode == "sp": + axs.plot(0,0,'r+',ms=30) + axs.set_xlabel('-v x B (km)') + axs.set_ylabel(' vxvxB (km)') + im = np.argmax(p) + axs.plot(xx[im]/1e3,yy[im]/1e3,'bx',ms=30) + fig.tight_layout() + return fig,axs + +def dist_to_line(xp,core,u): + xp = np.array(xp) + xp_core = xp-core + c2 = np.dot(xp_core,xp_core) + a2 = np.dot((np.dot(xp_core,u)*u),(np.dot(xp_core,u)*u)) + d = (np.abs(c2 - a2))**0.5 + return d + +def dist_to_line_sum(param,data,weights): + #distance line point: a = xp-core is D= | (a)^2-(a dot n)n | + #where ux is direction of line and x0 is a point in the line (like t = 0) + x0 = param[0] + y0 = param[1] + theta = param[2] + phi = param[3] + core = np.array([x0, y0, 0.]) + u = np.array([np.cos(phi)*np.sin(theta),np.sin(phi)*np.sin(theta),np.cos(theta)]) + dsum = 0 + for xp,w in zip(data,weights): + dsum += dist_to_line(xp,core,u)*w**2 + # print('%.2e %.2e %.2e %.2e %.2e'%(x0,y0,theta,phi,dsum)) + return dsum/len(data) + +def get_axis_points(e,savefig=True,path="",zgr=1400): + axis_points = [] + max_vals = [] + Xsteps = np.linspace(300,1000,15) + 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) + 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') + fig.savefig(path+'X%d_a.pdf'%(X)) + plt.close(fig) + im = np.argmax(p) + if np.abs(x[im]) == np.max(x) or np.abs(y[im]) == (np.max(y)): + x,y,p,loc_max = shower_plane_slice(e,X,21,21,scale4d,scale4d) + if savefig: + fig,axs = slice_figure(e,X,x,y,p,'sp') + fig.savefig(path+'X%d_c.pdf'%(X)) + plt.close(fig) + im = np.argmax(p) + x,y,p,loc_max = shower_plane_slice(e,X,21,21,scale0_2d,scale0_2d,x[im],y[im]) + if savefig: + 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) + + return Xsteps,axis_points,max_vals + +def fit_track(e,axis_points,vals,nscale=1e0): + weights = vals/np.max(vals) + data=axis_points[:] + data = [d/nscale for d in data] #km, to have more comparable step sizes + x0=[0,0,np.deg2rad(e.zenith),np.deg2rad(e.azimuth)] + res = minimize(dist_to_line_sum,args=(data,weights),x0=x0) + zen_r = np.rad2deg(res.x[2]) + azi_r = np.rad2deg(res.x[3]) + print(res,zen_r,e.zenith,azi_r,e.azimuth) + return zen_r,azi_r,[res.x[0]*nscale,res.x[1]*nscale,0] + + +def update_event(e,core,theta,phi,axp=None): + #recalculate + e.zenith = theta + e.azimuth = phi + theta = np.deg2rad(theta) + phi = np.deg2rad(phi) + e.core = e.core+core + e.uA = np.array([np.cos(phi)*np.sin(theta),np.sin(phi)*np.sin(theta),np.cos(theta)]) + e.uAxB = np.cross(e.uA,e.uB) + e.uAxB = e.uAxB/(np.dot(e.uAxB,e.uAxB))**0.5 + e.uAxAxB = np.cross(e.uA,e.uAxB) + #antenna position + for a in e.antennas: + a.x -= core[0] + a.y -= core[1] + a.z -= core[2] + if axp != None: + for ap in axp: + ap[0] -= core[0] + ap[1] -= core[1] + ap[2] -= core[2] + +def longitudinal_figure(dist,Xs,p,mode='grammage'): + fig, axs = plt.subplots(1,figsize=(6,5)) + if mode=='grammage': + axs.plot(Xs,p/np.max(p),'o-') + axs.set_xlabel('X (g/cm$^2$)') + if mode=='distance': + axs.plot(dist/1e3,p/np.max(p),'o-') + axs.set_xlabel('distance from ground (km)') + axs.grid() + fig.tight_layout() + return fig + +def time_residuals(e,tlable=True): + ds,tp,nsum,ssum,swidth,azi,x,y,sid = lateral_parameters(e,True,[0,0,0]) + fig, axs = plt.subplots(1,figsize=(6,5),sharex=True) + tp = tp-np.min(tp) + cut_outlier = ~((ds<200)&(tp > 10)) + axs.plot(ds,tp,'o') + if tlable: + for d,t,s in zip(ds,tp,sid): + plt.text(d,t,s) + # axs.text(ds,tp,sid) + axs.set_xlabel('distance (m)') + axs.set_ylabel('$\Delta t (ns)$') + axs.grid() + z = np.polyfit(ds[cut_outlier],tp[cut_outlier],3) + pfit = np.poly1d(z) + xfit = np.linspace(np.min(ds),np.max(ds),100) + yfit = pfit(xfit) + tres = tp - pfit(ds) + sigma =np.std(tres[cut_outlier]) + axs.plot(xfit,yfit,label=r'pol3 fit, $\sigma=%.2f$ (ns)'%(sigma)) + axs.legend() + fig.tight_layout() + return fig,tres + +def figure_3D(axis_points,max_vals,zen,azi,core,res = 0): + fig = plt.figure(figsize=(5,9)) + # fig, axs = plt.subplots(1,2,figsize=(12,8)) + ax = fig.add_subplot(2,1,1,projection='3d') + xp = [ap[0]/1e3 for ap in axis_points] + yp = [ap[1]/1e3 for ap in axis_points] + zp = [ap[2]/1e3 for ap in axis_points] + max_vals = np.asarray(max_vals) + ax.scatter(xp, yp, zp,c=max_vals,s=150*(max_vals/np.max(max_vals))**2,cmap='Spectral_r') + ax = fig.add_subplot(2,1,2) + core = np.array(core) + theta = np.deg2rad(zen) + phi = np.deg2rad(azi) + u = np.array([np.cos(phi)*np.sin(theta),np.sin(phi)*np.sin(theta),np.cos(theta)]) + residuals = [dist_to_line(ap,core,u) for ap in axis_points] + dist = [np.sum((ap-core)**2)**0.5 for ap in axis_points] + ax.scatter(dist,residuals,c=max_vals,cmap='Spectral_r') + ax.grid() + # ax.plot(xl,yl,zl,'-') + # ax.set_zlim(0,18) + # ax.view_init(15, 10) + fig.tight_layout() + if res != 0: + res.track_dis.append(dist) + res.track_res.append(residuals) + res.track_val.append(max_vals) + return fig + +class RITResult(): + """docstring for RITResult.""" + def __init__(self): + super(RITResult, self).__init__() + self.xmax_rit = [] + self.xmax = [] + self.profile_rit = [] + self.dX = [] + self.dl = [] + self.zenith_ini = [] + self.azimuth_ini = [] + self.core_ini = [] + self.dcore_rec = [] + self.zenith_rec = [] + self.azimuth_rec = [] + self.index = [] + self.isMC = [] + self.track_dis = [] + self.track_res =[] + self.track_val =[] + self.station_ids =[] + self.station_x =[] + self.station_y =[] + self.station_z =[] + self.station_maxE = [] + self.has_pulse = [] + +def fill_stations_propeties(e,res): + x = np.array([a.x for a in e.antennas]) + y = np.array([a.y for a in e.antennas]) + z = np.array([a.z for a in e.antennas]) + ids = [a.name for a in e.antennas] + maxE = np.array([np.max(a.E_AxB) for a in e.antennas]) + has_pulse = np.array([np.max(a.has_pulse) for a in e.antennas]) + res.station_x.append(x) + res.station_y.append(y) + res.station_z.append(z) + res.station_ids.append(ids) + res.has_pulse.append(has_pulse) + +def reconstruction(e,outfile=''): + res = RITResult() + res.isMC.append(True) + res.zenith_ini.append(e.zenith) + res.azimuth_ini.append(e.azimuth) + res.core_ini.append(e.core) + + set_pol_and_bp(e) + + #only use signal that have a signal in data + fill_stations_propeties(e,res) + Xs,axis_points,max_vals = get_axis_points(e,False) + 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) + update_event(e,core,zen,azi) + ds,Xs,locs,p = shower_axis_slice(e) + #result + res.dX.append(Xs) + res.dl.append(ds) + res.profile_rit.append(p) + res.xmax_rit.append(Xs[np.argmax(p)]) + res.azimuth_rec.append(e.azimuth) + res.zenith_rec.append(e.zenith) + res.dcore_rec.append(core) + return res