389 lines
14 KiB
Python
Executable File
389 lines
14 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# vim: fdm=indent ts=4
|
|
|
|
"""
|
|
Show how the Power changes when incorporating the
|
|
various clock offsets by plotting on a grid.
|
|
"""
|
|
|
|
import matplotlib.pyplot as plt
|
|
from mpl_toolkits.mplot3d import Axes3D # required for projection='3d' on old matplotliblib versions
|
|
import numpy as np
|
|
from os import path
|
|
import joblib
|
|
|
|
from earsim import REvent
|
|
from atmocal import AtmoCal
|
|
import aa_generate_beacon as beacon
|
|
import lib
|
|
from lib import rit
|
|
|
|
try:
|
|
from tqdm import tqdm
|
|
except:
|
|
tqdm = lambda x: x
|
|
|
|
def save_overlapping_traces_figure(test_location, ev, N_plot = 30, wx=200, title_extra=None, fname_distinguish='', fig_dir=None, location_in_plot_text_loc=(0.02, 0.95), location_text=None, **fig_kwargs):
|
|
P, t_, a_, a_sum, t_sum = rit.pow_and_time(test_location, ev, dt=1)
|
|
|
|
fig, axs = plt.subplots(**fig_kwargs)
|
|
#axs.set_title(
|
|
#"Antenna traces" + ("\n" if title_extra is not None else '') +
|
|
# title_extra )
|
|
axs.set_xlabel("Time [ns]")
|
|
axs.set_ylabel("Amplitude [$\\mu V/m$]")
|
|
if False and location_in_plot_text_loc:
|
|
text_loc = location_in_plot_text_loc
|
|
|
|
if not location_text:
|
|
shower_plane_location = rit.location_to_shower_plane(test_location, ev=ev)
|
|
location_text = '(' + ', '.join(['{:.1g}'.format(x) for x in shower_plane_location]) + ')'
|
|
|
|
axs.text(*text_loc, location_text, ha='left', transform=axs.transAxes)
|
|
|
|
a_max = [ np.amax(ant.E_AxB) for ant in ev.antennas ]
|
|
power_sort_idx = np.argsort(a_max)
|
|
|
|
if N_plot is None:
|
|
N_plot = len(a_max)
|
|
|
|
for i, idx in enumerate(reversed(power_sort_idx)):
|
|
if i >= N_plot:
|
|
break
|
|
|
|
alpha = max(0.4, 1/N_plot)
|
|
|
|
axs.plot(t_[idx], a_[idx], color='r', alpha=alpha, lw=2)
|
|
|
|
axs.set_ylim([None, max(a_max)*1.2])
|
|
if True:
|
|
axs.set_xlim([-28000, -27300])
|
|
|
|
if fig_dir:
|
|
if fname_distinguish:
|
|
fname_distinguish = "." + fname_distinguish
|
|
|
|
fig.tight_layout()
|
|
|
|
fig.savefig(path.join(fig_dir, path.basename(__file__) + f'{fname_distinguish}.trace_overlap.{case}.pdf'))
|
|
#fig.savefig(path.join(fig_dir, path.basename(__file__) + f'{fname_distinguish}.trace_overlap.{case}.png'), transparent=True)
|
|
|
|
# Take center between t_low and t_high
|
|
if True:
|
|
orig_xlims = axs.get_xlim()
|
|
|
|
if not True: # t_high and t_low from strongest signal
|
|
t_low = np.min(t_[power_sort_idx[-1]])
|
|
t_high = np.max(t_[power_sort_idx[-1]])
|
|
else: # take t_high and t_low from plotted signals
|
|
a = [np.min(t_[idx]) for idx in power_sort_idx[-N_plot:]]
|
|
t_low = np.nanmin(a)
|
|
b = [np.max(t_[idx]) for idx in power_sort_idx[-N_plot:]]
|
|
t_high = np.nanmax(b)
|
|
|
|
if False:
|
|
axs.plot(a, [0]*N_plot, 'gx', ms=10)
|
|
axs.plot(b, [0]*N_plot, 'b+', ms=10)
|
|
|
|
center_x = (t_high - t_low)/2 + t_low
|
|
|
|
low_xlim = max(orig_xlims[0], center_x - wx)
|
|
high_xlim = min(orig_xlims[1], center_x + wx)
|
|
|
|
axs.set_xlim(low_xlim, high_xlim)
|
|
fig.savefig(path.join(fig_dir, path.basename(__file__) + f'{fname_distinguish}.trace_overlap.zoomed.{case}.pdf'))
|
|
#fig.savefig(path.join(fig_dir, path.basename(__file__) + f'{fname_distinguish}.trace_overlap.zoomed.{case}.png'), transparent=True)
|
|
|
|
return fig
|
|
|
|
if __name__ == "__main__":
|
|
valid_cases = ['no_offset', 'repair_none', 'repair_phases', 'repair_full']
|
|
|
|
import sys
|
|
import os
|
|
import matplotlib
|
|
if os.name == 'posix' and "DISPLAY" not in os.environ:
|
|
matplotlib.use('Agg')
|
|
|
|
atm = AtmoCal()
|
|
|
|
from scriptlib import MyArgumentParser
|
|
parser = MyArgumentParser()
|
|
parser.add_argument('--input-fname', type=str, default=None, help='Path to mysim.sry, either directory or path. If empty it takes DATA_DIR and appends mysim.sry. (Default: %(default)s)')
|
|
parser.add_argument('-X', type=int, default=400, help='Atmospheric depth to start the initial grid.')
|
|
|
|
group = parser.add_argument_group('figures')
|
|
for case in valid_cases:
|
|
group.add_argument('--'+case.replace('_','-'), dest='figures', action='append_const', const=case)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not args.input_fname:
|
|
args.input_fname = args.data_dir
|
|
|
|
if path.isdir(args.input_fname):
|
|
args.input_fname = path.join(args.input_fname, "mysim.sry")
|
|
|
|
wanted_cases = args.figures
|
|
|
|
if not wanted_cases or 'all' in wanted_cases:
|
|
wanted_cases = valid_cases
|
|
|
|
figsize = (6, 4.5)
|
|
figsize_map = (6, 4.8)
|
|
if True:
|
|
from matplotlib import rcParams
|
|
#rcParams["text.usetex"] = True
|
|
rcParams["font.family"] = "serif"
|
|
rcParams["font.size"] = "15"
|
|
rcParams["grid.linestyle"] = 'dotted'
|
|
rcParams["figure.figsize"] = figsize
|
|
figsize = rcParams['figure.figsize']
|
|
|
|
fig_dir = args.fig_dir
|
|
show_plots = args.show_plots
|
|
|
|
remove_beacon_from_traces = True
|
|
apply_signal_window_from_max = True
|
|
|
|
####
|
|
fname_dir = args.data_dir
|
|
antennas_fname = path.join(fname_dir, beacon.antennas_fname)
|
|
tx_fname = path.join(fname_dir, beacon.tx_fname)
|
|
beacon_snr_fname = path.join(fname_dir, beacon.beacon_snr_fname)
|
|
|
|
# create fig_dir
|
|
if fig_dir:
|
|
os.makedirs(fig_dir, exist_ok=True)
|
|
|
|
# Read in antennas from file
|
|
_, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)
|
|
_, __, txdata = beacon.read_tx_file(tx_fname)
|
|
# Read original REvent
|
|
ev = REvent(args.input_fname)
|
|
bak_ants = ev.antennas
|
|
# .. patch in our antennas
|
|
ev.antennas = antennas
|
|
# Read in snr info
|
|
beacon_snrs = beacon.read_snr_file(beacon_snr_fname)
|
|
snr_str = f"$\\langle SNR \\rangle$ = {beacon_snrs['mean']: .1g}"
|
|
|
|
##
|
|
## Setup grid
|
|
##
|
|
X = args.X
|
|
zgr = 0 #not exact
|
|
dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),X,zgr+ev.core[2])
|
|
scale2d = dXref*np.tan(np.deg2rad(2.))
|
|
scale4d = dXref*np.tan(np.deg2rad(4.))
|
|
scale02d = dXref*np.tan(np.deg2rad(0.2))
|
|
|
|
Nx, Ny = 31, 31
|
|
#Nx, Ny = 15, 15
|
|
scales = {
|
|
'scale2d': scale2d,
|
|
'scale4d': scale4d,
|
|
'scale02d': scale02d,
|
|
}
|
|
|
|
N_plot = None
|
|
trace_zoom_wx = 100
|
|
|
|
plot_titling = {
|
|
'no_offset': "no clock offsets",
|
|
'repair_none': "randomised clock offsets",
|
|
'repair_phases': "phase clock offsets",
|
|
'repair_full': "phase + period clock offsets"
|
|
}
|
|
|
|
power_on_grid_sc_kwargs = dict(
|
|
s=90, # 250 for (10,8)
|
|
cmap='inferno'
|
|
)
|
|
|
|
# For now only implement using one freq_name
|
|
freq_names = ev.antennas[0].beacon_info.keys()
|
|
if len(freq_names) > 1:
|
|
raise NotImplementedError
|
|
|
|
freq_name = next(iter(freq_names))
|
|
|
|
# Pre remove the beacon from the traces
|
|
#
|
|
# We need to remove it here, so we do not shoot ourselves in
|
|
# the foot when changing to the various clock offsets.
|
|
#
|
|
# Note that the bandpass filter is applied only after E_AxB is
|
|
# reconstructed so we have to manipulate the original traces.
|
|
if remove_beacon_from_traces:
|
|
tx_amps = txdata['amplitudes']
|
|
tx_amps_sum = np.sum(tx_amps)
|
|
|
|
for i, ant in enumerate(ev.antennas):
|
|
beacon_phase = ant.beacon_info[freq_name]['beacon_phase']
|
|
f = ant.beacon_info[freq_name]['freq']
|
|
ampl_AxB = ant.beacon_info[freq_name]['amplitude']
|
|
calc_beacon = lib.sine_beacon(f, ev.antennas[i].t, amplitude=ampl_AxB, phase=beacon_phase)
|
|
|
|
# Split up contribution to the various polarisations
|
|
for j, amp in enumerate(tx_amps):
|
|
if j == 0:
|
|
ev.antennas[i].Ex -= amp*(1/tx_amps_sum)*calc_beacon
|
|
elif j == 1:
|
|
ev.antennas[i].Ey -= amp*(1/tx_amps_sum)*calc_beacon
|
|
elif j == 2:
|
|
ev.antennas[i].Ez -= amp*(1/tx_amps_sum)*calc_beacon
|
|
|
|
# Subtract the beacon from E_AxB
|
|
ev.antennas[i].E_AxB -= calc_beacon
|
|
|
|
# Slice the traces to a small part around the peak
|
|
if apply_signal_window_from_max:
|
|
N_pre, N_post = 250, 250 # TODO: make this configurable
|
|
|
|
for i, ant in enumerate(ev.antennas):
|
|
# Get max idx from all the traces
|
|
# and select the strongest
|
|
max_idx = []
|
|
maxs = []
|
|
|
|
for trace in [ant.Ex, ant.Ey, ant.Ez]:
|
|
idx = np.argmax(np.abs(trace))
|
|
max_idx.append(idx)
|
|
maxs.append( np.abs(trace[idx]) )
|
|
|
|
idx = np.argmax(maxs)
|
|
max_idx = max_idx[idx]
|
|
|
|
|
|
low_idx = max(0, max_idx-N_pre)
|
|
high_idx = min(len(ant.t), max_idx+N_post)
|
|
|
|
ev.antennas[i].t = ant.t[low_idx:high_idx]
|
|
ev.antennas[i].t_AxB = ant.t_AxB[low_idx:high_idx]
|
|
|
|
ev.antennas[i].Ex = ant.Ex[low_idx:high_idx]
|
|
ev.antennas[i].Ey = ant.Ey[low_idx:high_idx]
|
|
ev.antennas[i].Ez = ant.Ez[low_idx:high_idx]
|
|
ev.antennas[i].E_AxB = ant.E_AxB[low_idx:high_idx]
|
|
|
|
## Apply polarisation and bandpass filter
|
|
rit.set_pol_and_bp(ev)
|
|
|
|
# backup antenna times
|
|
backup_antenna_t = [ ant.t for ant in ev.antennas ]
|
|
backup_antenna_t_AxB = [ ant.t_AxB for ant in ev.antennas ]
|
|
|
|
fig = save_overlapping_traces_figure([0,0,0], ev, N_plot=1, wx=trace_zoom_wx, title_extra = plot_titling[case], fname_distinguish=f'single', fig_dir=fig_dir, figsize=figsize)
|
|
plt.close(fig)
|
|
|
|
no_offset_maximum_power = None # set in below loop
|
|
with joblib.parallel_backend("loky"):
|
|
# always force no_offset first to determine the reference vmax in the colorbar
|
|
for case in ['no_offset'] + [c for c in wanted_cases if c != 'no_offset']:
|
|
print(f"Starting {case} figure")
|
|
# Repair clock offsets with the measured offsets
|
|
transl_modes = {'no_offset':'orig', 'repair_phases':'phases', 'repair_full':'full'}
|
|
|
|
if case in transl_modes:
|
|
transl_mode = transl_modes[case]
|
|
measured_offsets = beacon.read_antenna_clock_repair_offsets(antennas, mode=transl_mode, freq_name=freq_name)
|
|
else:
|
|
print(f"Warning: unknown repair case requested '{case}', defaulting to none.")
|
|
measured_offsets = [0]*len(ev.antennas)
|
|
|
|
for i, ant in enumerate(ev.antennas):
|
|
total_clock_offset = measured_offsets[i]
|
|
|
|
ev.antennas[i].t = backup_antenna_t[i] + total_clock_offset
|
|
ev.antennas[i].t_AxB = backup_antenna_t_AxB[i] + total_clock_offset
|
|
|
|
if i == 0:
|
|
# Specifically compare the times
|
|
print("backup time, time with measured_offset, true clock offset, measured clock offset")
|
|
print(bak_ants[i].t[0], ev.antennas[i].t[0], ev.antennas[i].attrs['clock_offset'], measured_offsets[i])
|
|
|
|
#
|
|
# determine the maximum power for the true timing
|
|
#
|
|
if case == 'no_offset':
|
|
xx, yy, p, maxp_loc = rit.shower_plane_slice(ev, X=X, Nx=1, Ny=1, wx=0, wy=0, zgr=zgr)
|
|
|
|
no_offset_maximum_power = p
|
|
if True and np.isclose(X, 400):
|
|
no_offset_maximum_power = 200
|
|
|
|
|
|
# we forced no_offset for determing the reference vmax
|
|
# simply continue with the rest
|
|
if case not in wanted_cases:
|
|
continue
|
|
|
|
#
|
|
# Plot overlapping traces at 0,0,0
|
|
#
|
|
loc = (0,0,0)
|
|
fig = save_overlapping_traces_figure(loc, ev, N_plot=N_plot, wx=trace_zoom_wx, title_extra = plot_titling[case], fname_distinguish=f'{case}.0', fig_dir=fig_dir, figsize=figsize)
|
|
plt.close(fig)
|
|
|
|
#
|
|
# Plot overlapping traces at simulation shower axis
|
|
#
|
|
dX = atm.distance_to_slant_depth(np.deg2rad(ev.zenith), X, zgr)
|
|
loc = (0)*ev.uAxB + (0)*ev.uAxAxB + dX * ev.uA
|
|
loc[0:1] = 0
|
|
|
|
fig = save_overlapping_traces_figure(loc, ev, N_plot=N_plot, wx=trace_zoom_wx, title_extra = plot_titling[case], fname_distinguish=f'{case}.axis.X{X}', location_text=f"on simulation axis, $X={X}$", fig_dir=fig_dir, figsize=figsize)
|
|
plt.close(fig)
|
|
|
|
if not True:
|
|
continue;
|
|
|
|
#
|
|
# Measure power on grid
|
|
# and plot overlapping traces at position with highest power
|
|
#
|
|
for scalename, scale in tqdm(scales.items()):
|
|
wx, wy = scale, scale
|
|
print(f"Starting grid measurement for figure {case} with {scalename}")
|
|
xx, yy, p, maxp_loc = rit.shower_plane_slice(ev, X=X, Nx=Nx, Ny=Nx, wx=wx, wy=wy, zgr=zgr)
|
|
fig, axs = rit.slice_figure(ev, X, xx, yy, p, mode='sp', scatter_kwargs={**dict(
|
|
vmax=no_offset_maximum_power,
|
|
vmin=0,
|
|
), **power_on_grid_sc_kwargs},
|
|
figsize=figsize_map)
|
|
|
|
suptitle = fig._suptitle.get_text()
|
|
fig.suptitle("")
|
|
#axs.set_title("Shower plane slice\n" + plot_titling[case] + "\n" + suptitle)
|
|
axs.set_aspect('equal', 'datalim')
|
|
#axs.legend(title=snr_str,loc='upper right')
|
|
|
|
axs.set_xlim(1.1*min(xx)/1e3, 1.1*max(xx)/1e3)
|
|
axs.set_ylim(1.1*min(yy)/1e3, 1.1*max(yy)/1e3)
|
|
|
|
if fig_dir:
|
|
fig.tight_layout()
|
|
fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.X{X}.{case}.{scalename}.pdf'))
|
|
plt.close(fig)
|
|
|
|
#
|
|
# Plot overlapping traces at highest power of each scale
|
|
#
|
|
fig = save_overlapping_traces_figure(maxp_loc, ev, N_plot=N_plot, wx=trace_zoom_wx, title_extra = plot_titling[case] + ', ' + scalename + ' best', fname_distinguish=scalename+f'.best.X{X}', fig_dir=fig_dir, figsize=figsize)
|
|
|
|
#
|
|
# and plot overlapping traces at two other locations
|
|
#
|
|
if False:
|
|
for dist in [ 0.5, 5, 10, 50, 100]:
|
|
# only add distance horizontally
|
|
location = maxp_loc + np.sqrt(dist*1e3)*np.array([1,1,0])
|
|
|
|
fig = save_overlapping_traces_figure(location, ev, N_plot=N_plot, wx=wx, title_extra = plot_titling[case] + ', ' + scalename + f', x + {dist}km', fname_distinguish=f'{scalename}.x{dist}.X{X}', fig_dir=fig_dir, figsize=figsize)
|
|
plt.close(fig)
|
|
|
|
|
|
if args.show_plots:
|
|
plt.show()
|