2022-11-24 17:54:48 +01:00
#!/usr/bin/env python3
# vim: fdm=indent ts=4
import h5py
from itertools import combinations , zip_longest
import matplotlib . pyplot as plt
2022-12-22 16:33:19 +01:00
from matplotlib . colors import Normalize
import matplotlib as mpl
2022-11-24 17:54:48 +01:00
import numpy as np
2023-03-27 16:56:19 +02:00
import json
2022-11-24 17:54:48 +01:00
import aa_generate_beacon as beacon
import lib
2023-02-07 17:58:45 +01:00
from lib import figlib
2022-11-24 17:54:48 +01:00
if __name__ == " __main__ " :
from os import path
import sys
2022-12-21 17:24:44 +01:00
import os
import matplotlib
if os . name == ' posix ' and " DISPLAY " not in os . environ :
matplotlib . use ( ' Agg ' )
2023-01-12 14:31:21 +01:00
from scriptlib import MyArgumentParser
parser = MyArgumentParser ( )
args = parser . parse_args ( )
2023-02-02 08:57:03 +01:00
figsize = ( 12 , 8 )
2022-11-24 17:54:48 +01:00
2023-01-12 14:31:21 +01:00
show_plots = args . show_plots
2022-11-24 17:54:48 +01:00
ref_ant_id = None # leave None for all baselines
####
2023-02-02 17:55:37 +01:00
fname_dir = args . data_dir
2022-11-24 17:54:48 +01:00
antennas_fname = path . join ( fname_dir , beacon . antennas_fname )
2022-12-21 17:24:44 +01:00
time_diffs_fname = ' time_diffs.hdf5 ' if False else antennas_fname
2023-01-12 14:49:54 +01:00
fig_dir = args . fig_dir # set None to disable saving
2023-04-12 22:42:59 +02:00
snr_fname = path . join ( fname_dir , beacon . snr_fname )
2022-11-24 17:54:48 +01:00
2023-01-19 16:24:05 +01:00
basenames , time_diffs , f_beacons , clock_phase_diffs , k_periods = beacon . read_baseline_time_diffs_hdf5 ( time_diffs_fname )
2022-11-24 17:54:48 +01:00
f_beacon , tx , antennas = beacon . read_beacon_hdf5 ( antennas_fname )
2022-12-21 17:24:44 +01:00
# TODO: allow multiple frequencies
if ( f_beacon != f_beacons ) . any ( ) :
raise NotImplementedError
N_base = len ( basenames )
N_ant = len ( antennas )
2022-12-22 16:55:04 +01:00
# reshape time_diffs into N_ant x N_ant matrix
2023-01-19 17:13:29 +01:00
clock_phase_matrix = np . full ( ( N_ant , N_ant ) , np . nan , dtype = float )
2022-12-21 17:24:44 +01:00
2022-12-22 16:55:04 +01:00
## set i=i terms to 0
for i in range ( N_ant ) :
2023-01-19 17:13:29 +01:00
clock_phase_matrix [ i , i ] = 0
2022-12-22 16:55:04 +01:00
## fill matrix
2022-12-21 17:24:44 +01:00
name2idx = lambda name : int ( name ) - 1
for i , b in enumerate ( basenames ) :
idx = ( name2idx ( b [ 0 ] ) , name2idx ( b [ 1 ] ) )
2023-04-12 22:42:59 +02:00
#if idx[1] < idx[0]:
# idx = (idx[1], idx[0])
2023-01-19 17:13:29 +01:00
clock_phase_matrix [ ( idx [ 0 ] , idx [ 1 ] ) ] = lib . phase_mod ( clock_phase_diffs [ i ] )
clock_phase_matrix [ ( idx [ 1 ] , idx [ 0 ] ) ] = lib . phase_mod ( - 1 * clock_phase_diffs [ i ] )
2022-12-21 17:24:44 +01:00
2022-12-22 16:33:19 +01:00
mat_kwargs = dict (
norm = Normalize ( vmin = - np . pi , vmax = + np . pi ) ,
cmap = mpl . cm . get_cmap ( ' Spectral_r ' )
)
2022-12-21 17:24:44 +01:00
2022-12-22 16:33:19 +01:00
# Show Matrix as figure
if True :
2023-02-01 14:13:26 +01:00
fig , ax = plt . subplots ( figsize = figsize )
2023-01-24 15:24:38 +01:00
ax . set_title ( " Measured clock phase differences Baseline i,j " )
2022-12-22 16:33:19 +01:00
ax . set_ylabel ( " Antenna no. i " )
ax . set_xlabel ( " Antenna no. j " )
2023-01-19 17:13:29 +01:00
im = ax . imshow ( clock_phase_matrix , interpolation = ' none ' , * * mat_kwargs )
2022-12-22 16:33:19 +01:00
fig . colorbar ( im , ax = ax )
if fig_dir :
2023-01-04 11:05:20 +01:00
fig . savefig ( path . join ( fig_dir , path . basename ( __file__ ) + f " .matrix.original.pdf " ) )
2022-12-22 16:33:19 +01:00
2022-12-22 17:59:23 +01:00
plt . close ( fig )
2022-12-22 16:33:19 +01:00
# Modify the matrix to let each column represent multiple
# measurements of the same baseline (j,0) phase difference
if True :
# for each row j subtract the 0,j element from the whole row
# and apply phase_mod
2023-01-19 17:13:29 +01:00
first_row = - 1 * ( clock_phase_matrix [ 0 , : ] * np . ones_like ( clock_phase_matrix ) ) . T
2022-12-22 16:33:19 +01:00
# Show subtraction Matrix as figure
if True :
2023-02-01 14:13:26 +01:00
fig , ax = plt . subplots ( figsize = figsize )
2023-01-24 15:24:38 +01:00
ax . set_title ( " Clock Phase Subtraction matrix i,j " )
2022-12-22 16:33:19 +01:00
ax . set_ylabel ( " Antenna no. i " )
ax . set_xlabel ( " Antenna no. j " )
im = ax . imshow ( first_row , interpolation = ' none ' , * * mat_kwargs )
fig . colorbar ( im , ax = ax )
if fig_dir :
2023-01-04 11:05:20 +01:00
fig . savefig ( path . join ( fig_dir , path . basename ( __file__ ) + f " .matrix.first_row.pdf " ) )
2022-12-22 16:33:19 +01:00
2022-12-22 17:59:23 +01:00
plt . close ( fig )
2023-01-19 17:13:29 +01:00
clock_phase_matrix = clock_phase_matrix - first_row
clock_phase_matrix = lib . phase_mod ( clock_phase_matrix )
2022-12-21 17:24:44 +01:00
# Except for the first row, these are all separate measurements
# Condense into phase offset per antenna
if True : # do not use the first row
2023-01-19 17:13:29 +01:00
my_mask = np . arange ( 1 , len ( clock_phase_matrix ) , dtype = int )
2022-12-21 17:24:44 +01:00
else :
2023-01-19 17:13:29 +01:00
my_mask = np . arange ( 0 , len ( clock_phase_matrix ) , dtype = int )
2022-12-21 17:24:44 +01:00
2023-01-19 17:13:29 +01:00
mean_clock_phase = np . nanmean ( clock_phase_matrix [ my_mask ] , axis = 0 )
std_clock_phase = np . nanstd ( clock_phase_matrix [ my_mask ] , axis = 0 )
2022-12-21 17:24:44 +01:00
2023-03-27 16:56:19 +02:00
# Remove the mean from the matrix
if False :
clock_phase_matrix = clock_phase_matrix - np . mean ( mean_clock_phase )
mean_clock_phase = np . nanmean ( clock_phase_matrix [ my_mask ] , axis = 0 )
std_clock_phase = np . nanstd ( clock_phase_matrix [ my_mask ] , axis = 0 )
2022-12-22 16:33:19 +01:00
# Show resulting matrix as figure
if True :
2023-02-01 14:13:26 +01:00
fig , axs = plt . subplots ( 2 , 1 , sharex = True , figsize = figsize )
2023-01-24 15:24:38 +01:00
axs [ 0 ] . set_title ( " Modified clock phase differences Baseline 0,j " )
2022-12-22 16:33:19 +01:00
axs [ 0 ] . set_ylabel ( " Antenna no. 0 " )
axs [ - 1 ] . set_xlabel ( " Antenna no. j " )
2023-01-19 17:13:29 +01:00
im = axs [ 0 ] . imshow ( clock_phase_matrix , interpolation = ' none ' , * * mat_kwargs )
2022-12-22 16:33:19 +01:00
fig . colorbar ( im , ax = axs )
axs [ 0 ] . set_aspect ( ' auto ' )
2023-01-19 17:13:29 +01:00
colours = [ mat_kwargs [ ' cmap ' ] ( mat_kwargs [ ' norm ' ] ( x ) ) for x in mean_clock_phase ]
2022-12-22 16:33:19 +01:00
axs [ 1 ] . set_ylabel ( " Mean Baseline Phase (0, j)[rad] " )
2023-01-19 17:13:29 +01:00
axs [ 1 ] . errorbar ( np . arange ( N_ant ) , mean_clock_phase , yerr = std_clock_phase , ls = ' none ' )
axs [ 1 ] . scatter ( np . arange ( N_ant ) , mean_clock_phase , c = colours , s = 4 )
2022-12-22 16:33:19 +01:00
if fig_dir :
2023-01-04 11:05:20 +01:00
fig . savefig ( path . join ( fig_dir , path . basename ( __file__ ) + f " .matrix.modified.pdf " ) )
2022-12-22 16:33:19 +01:00
2022-12-22 17:59:23 +01:00
plt . close ( fig )
2022-12-21 17:24:44 +01:00
# write into antenna hdf5
with h5py . File ( antennas_fname , ' a ' ) as fp :
group = fp [ ' antennas ' ]
freq_name = None
for i , ant in enumerate ( antennas ) :
h5ant = group [ ant . name ]
h5beacon_info = h5ant [ ' beacon_info ' ]
# find out freq_name
if freq_name is None :
freq_name = [ k for k in h5beacon_info . keys ( ) if np . isclose ( h5beacon_info [ k ] . attrs [ ' freq ' ] , f_beacon ) ] [ 0 ]
h5attrs = h5beacon_info [ freq_name ] . attrs
idx = name2idx ( ant . name )
2023-01-19 17:13:29 +01:00
h5attrs [ ' clock_phase_mean ' ] = mean_clock_phase [ idx ]
h5attrs [ ' clock_phase_std ' ] = std_clock_phase [ idx ]
2022-12-21 17:24:44 +01:00
##############################
# Compare actual time shifts #
##############################
2023-04-12 22:42:59 +02:00
snrs = beacon . read_snr_file ( snr_fname )
snr_str = f " $ \\ langle SNR \\ rangle$ = { snrs [ ' mean ' ] : .1e } "
2022-12-22 18:00:34 +01:00
actual_antenna_time_shifts = { a . name : a . attrs [ ' clock_offset ' ] for a in sorted ( antennas , key = lambda a : int ( a . name ) ) }
2022-11-24 17:54:48 +01:00
2022-12-21 17:24:44 +01:00
if True :
2022-12-22 18:00:34 +01:00
actual_antenna_phase_shifts = [ - 1 * lib . phase_mod ( 2 * np . pi * f_beacon * v ) for k , v in actual_antenna_time_shifts . items ( ) ]
antenna_names = [ int ( k ) - 1 for k , v in actual_antenna_time_shifts . items ( ) ]
2022-11-24 17:54:48 +01:00
2022-12-22 18:12:04 +01:00
# Make sure to shift all antennas by a global phase
2023-01-19 17:13:29 +01:00
global_phase_shift = actual_antenna_phase_shifts [ 0 ] - mean_clock_phase [ 0 ]
2022-12-22 18:12:04 +01:00
actual_antenna_phase_shifts = lib . phase_mod ( actual_antenna_phase_shifts - global_phase_shift )
2023-03-27 16:56:19 +02:00
fit_info = { }
2022-12-22 12:04:22 +01:00
for i in range ( 2 ) :
plot_residuals = i == 1
2023-02-07 17:58:45 +01:00
true_phases = actual_antenna_phase_shifts
measured_phases = mean_clock_phase
2022-12-21 17:24:44 +01:00
2023-02-07 17:58:45 +01:00
hist_kwargs = { }
if plot_residuals :
measured_phases = lib . phase_mod ( measured_phases - actual_antenna_phase_shifts )
2023-03-27 16:56:19 +02:00
fig , _fit_info = figlib . phase_comparison_figure (
2023-02-07 17:58:45 +01:00
measured_phases ,
true_phases ,
plot_residuals = plot_residuals ,
f_beacon = f_beacon ,
figsize = figsize ,
hist_kwargs = hist_kwargs ,
2023-02-13 10:40:26 +01:00
fit_gaussian = plot_residuals ,
2023-04-12 22:42:59 +02:00
fit_ricianphase = plot_residuals ,
2023-03-27 16:56:19 +02:00
return_fit_info = True ,
2023-02-07 17:58:45 +01:00
)
2022-11-24 17:54:48 +01:00
2023-02-07 17:58:45 +01:00
axs = fig . get_axes ( )
2023-04-12 22:42:59 +02:00
axs [ 0 ] . legend ( title = snr_str )
2022-12-22 12:04:22 +01:00
if plot_residuals :
2023-04-12 22:42:59 +02:00
axs [ 0 ] . set_title ( " Difference between Measured and Actual phases (minus global phase) \n for Antenna $i$ " )
2023-02-03 17:56:24 +01:00
axs [ - 1 ] . set_xlabel ( " Antenna Mean Phase Residual $ \\ Delta_ \\ varphi$ " )
2022-12-22 12:04:22 +01:00
else :
2023-04-12 22:42:59 +02:00
axs [ 0 ] . set_title ( " Comparison Measured and Actual phases (minus global phase) \n for Antenna $i$ " )
2023-02-03 17:56:24 +01:00
axs [ - 1 ] . set_xlabel ( " Antenna Mean Phase $ \\ varphi$ " )
2022-12-22 12:04:22 +01:00
i = 1
2022-12-22 12:06:34 +01:00
axs [ i ] . set_ylabel ( " Antenna no. " )
2023-02-07 17:58:45 +01:00
#axs[i].errorbar(mean_clock_phase, np.arange(N_ant), yerr=std_clock_phase, marker='4', alpha=0.7, ls='none', color=colors[0], label='Measured')
2022-12-22 12:04:22 +01:00
fig . tight_layout ( )
2022-12-22 12:06:34 +01:00
2022-12-22 12:04:22 +01:00
if fig_dir :
extra_name = " measured "
if plot_residuals :
extra_name = " residuals "
2023-01-04 11:05:20 +01:00
fig . savefig ( path . join ( fig_dir , path . basename ( __file__ ) + f " .phase. { extra_name } .pdf " ) )
2022-12-21 17:24:44 +01:00
2023-03-27 16:56:19 +02:00
# Save fit_info to data file
if fname_dir and plot_residuals :
with open ( path . join ( fname_dir , ' phase_time_residuals.json ' ) , ' w ' ) as fp :
json . dump (
{
' mean ' : np . mean ( measured_phases ) ,
' std ' : np . std ( measured_phases ) ,
' values ' : measured_phases . tolist ( ) ,
} ,
fp )
2022-12-21 17:24:44 +01:00
##########################
##########################
##########################
2022-12-22 18:00:34 +01:00
actual_baseline_time_shifts = [ ]
2022-12-21 17:24:44 +01:00
for i , b in enumerate ( basenames ) :
2022-12-22 18:00:34 +01:00
actual_baseline_time_shift = actual_antenna_time_shifts [ b [ 0 ] ] - actual_antenna_time_shifts [ b [ 1 ] ]
2022-12-21 17:24:44 +01:00
2022-12-22 18:00:34 +01:00
actual_baseline_time_shifts . append ( actual_baseline_time_shift )
2022-12-21 17:24:44 +01:00
2023-01-19 17:13:29 +01:00
# unpack mean_clock_phase back into a list of time diffs
2022-12-22 18:00:34 +01:00
measured_baseline_time_diffs = [ ]
2022-12-21 17:24:44 +01:00
for i , b in enumerate ( basenames ) :
2023-01-19 17:13:29 +01:00
phase0 , phase1 = mean_clock_phase [ name2idx ( b [ 0 ] ) ] , mean_clock_phase [ name2idx ( b [ 1 ] ) ]
2022-12-22 18:00:34 +01:00
measured_baseline_time_diffs . append ( lib . phase_mod ( phase1 - phase0 ) / ( 2 * np . pi * f_beacon ) )
2022-12-21 17:24:44 +01:00
# Make a plot
if True :
2022-12-22 18:00:34 +01:00
for i in range ( 2 ) :
2023-02-01 14:13:26 +01:00
fig , ax = plt . subplots ( figsize = figsize )
2022-12-22 18:00:34 +01:00
ax . set_title ( " Baseline Time difference reconstruction " + ( ' ' if i == 0 else ' (wrapped time) ' ) )
ax . set_xlabel ( " Baseline no. " )
ax . set_ylabel ( " Time $ \\ Delta t$ [ns] " )
if True :
forward = lambda x : x / ( 2 * np . pi * f_beacon )
inverse = lambda x : 2 * np . pi * x * f_beacon
secax = ax . secondary_yaxis ( ' right ' , functions = ( inverse , forward ) )
secax . set_ylabel ( ' Phase $ \\ Delta \\ varphi$ [rad] ' )
if True : # indicate single beacon period span
ax . plot ( ( - 1 , - 1 ) , ( - 1 / ( 2 * f_beacon ) , 1 / ( 2 * f_beacon ) ) , marker = ' 3 ' , ms = 10 , label = ' 1/f_beacon ' )
if i == 0 :
2023-01-11 16:52:05 +01:00
ax . plot ( np . arange ( N_base ) , actual_baseline_time_shifts , ls = ' none ' , marker = ' h ' , alpha = 0.8 , label = ' actual time shifts ' )
2022-12-22 18:00:34 +01:00
else :
2023-01-11 16:52:05 +01:00
ax . plot ( np . arange ( N_base ) , ( actual_baseline_time_shifts + 1 / ( 2 * f_beacon ) ) % ( 1 / f_beacon ) - 1 / ( 2 * f_beacon ) , ls = ' none ' , marker = ' h ' , label = ' actual time shifts ' , alpha = 0.8 )
ax . plot ( np . arange ( N_base ) , measured_baseline_time_diffs , ls = ' none ' , alpha = 0.8 , marker = ' x ' , label = ' calculated ' )
2022-12-22 18:00:34 +01:00
2023-04-12 22:42:59 +02:00
ax . legend ( title = snr_str )
2022-12-22 18:00:34 +01:00
if fig_dir :
extra_name = ' '
if i == 1 :
extra_name = ' .wrapped '
2023-01-11 16:52:05 +01:00
old_lims = ( ax . get_xlim ( ) , ax . get_ylim ( ) )
for j in range ( 2 ) :
if j == 1 :
ax . set_xlim ( - 5 , 50 )
extra_name + = ' .n50 '
fig . savefig ( path . join ( fig_dir , path . basename ( __file__ ) + f " .time_comparison { extra_name } .pdf " ) )
ax . set_xlim ( * old_lims [ 0 ] )
ax . set_ylim ( * old_lims [ 1 ] )
2022-12-21 17:24:44 +01:00
if show_plots :
plt . show ( )