m-thesis-introduction/simulations/airshower_beacon_simulation/dc_grid_power_time_fixes.py

259 lines
9.2 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
if __name__ == "__main__":
valid_cases = ['no_offset', 'repair_none', 'repair_phases', 'repair_all']
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()
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()
wanted_cases = args.figures
if not wanted_cases or 'all' in wanted_cases:
wanted_cases = valid_cases
fname = "ZH_airshower/mysim.sry"
figsize = (12,8)
fig_dir = args.fig_dir
show_plots = args.show_plots
remove_beacon_from_traces = True
apply_signal_window_from_max = True
####
fname_dir = path.dirname(fname)
antennas_fname = path.join(fname_dir, beacon.antennas_fname)
pickle_fname = path.join(fname_dir, 'res.pkl')
tx_fname = path.join(fname_dir, beacon.tx_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(fname)
bak_ants = ev.antennas
# .. patch in our antennas
ev.antennas = antennas
##
## Setup grid
##
X = 400
zgr = 0 #not exact
dXref = atm.distance_to_slant_depth(np.deg2rad(ev.zenith),750,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 = 21, 21
scales = {
'scale2d': scale2d,
'scale4d': scale4d,
'scale02d': scale02d,
}
plot_titling = {
'no_offset': "no clock offset",
'repair_none': "unrepaired clock offset",
'repair_phases': "phase resolved clock offsets repaired",
'repair_all': "final measured clock offsets repaired"
}
# 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
#
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):
max_idx = np.argmax(ant.E_AxB)
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]
# 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 ]
## Apply polarisation and bandpass filter
rit.set_pol_and_bp(ev)
with joblib.parallel_backend("loky"):
for case in wanted_cases:
print(f"Starting {case} figure")
# Repair clock offsets with the measured offsets
transl_modes = {'no_offset':'orig', 'repair_phases':'phases', 'repair_all':'all'}
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:
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(bak_ants[i].t[0], ev.antennas[i].t[0], ev.antennas[i].t[0], ev.antennas[i].attrs['clock_offset'], measured_offsets[i])
#
# Plot overlapping traces at 0,0,0
#
if True:
P, t_, a_, a_sum, t_sum = rit.pow_and_time([0,0,0], ev, dt=1)
fig, axs = plt.subplots(figsize=figsize)
axs.set_title("Antenna traces" + "\n" + plot_titling[case])
axs.set_xlabel("Time [ns]")
axs.set_ylabel("Amplitude [$\\mu V/m$]")
a_max = [ np.amax(ant.E_AxB) for ant in ev.antennas ]
power_sort_idx = np.argsort(a_max)
N_plot = 30
for i, idx in enumerate(reversed(power_sort_idx)):
if i > N_plot:
break
alpha = max(0.4, 1/len(a_))
axs.plot(t_[idx], a_[idx], color='r', alpha=alpha, lw=2)
if fig_dir:
fig.tight_layout()
fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.trace_overlap.{case}.pdf'))
fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.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:]]
axs.plot(a, [0]*N_plot, 'gx', ms=10)
t_low = np.nanmin(a)
b = [np.max(t_[idx]) for idx in power_sort_idx[-N_plot:]]
axs.plot(b, [0]*N_plot, 'b+', ms=10)
t_high = np.nanmax(b)
center_x = (t_high - t_low)/2 + t_low
wx = 200
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'.trace_overlap.zoomed.{case}.pdf'))
fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.trace_overlap.zoomed.{case}.png'), transparent=True)
if True:
continue
# Measure power on grid
for scalename, scale in scales.items():
wx, wy = scale, scale
print(f"Starting grid measurement for figure {case} with {scalename}")
xx, yy, p, maxp = rit.shower_plane_slice(ev, X=X, Nx=Nx, Ny=Nx, wx=wx, wy=wy)
fig, axs = rit.slice_figure(ev, X, xx, yy, p, mode='sp')
suptitle = fig._suptitle.get_text()
fig.suptitle("")
axs.set_title(suptitle +"\n" +plot_titling[case])
#axs.set_aspect('equal', 'datalim')
if fig_dir:
fig.tight_layout()
fig.savefig(path.join(fig_dir, path.basename(__file__) + f'.X{X}.{case}.{scalename}.pdf'))
if args.show_plots:
plt.show()