#!/usr/bin/env python3 # vim: fdm=indent ts=4 import h5py from itertools import combinations, zip_longest import matplotlib.pyplot as plt from matplotlib.colors import Normalize import matplotlib as mpl import numpy as np import json import aa_generate_beacon as beacon import lib from lib import figlib if __name__ == "__main__": from os import path import sys import os import matplotlib if os.name == 'posix' and "DISPLAY" not in os.environ: matplotlib.use('Agg') from scriptlib import MyArgumentParser parser = MyArgumentParser() args = parser.parse_args() figsize = (12,8) show_plots = args.show_plots ref_ant_id = None # leave None for all baselines #### fname_dir = args.data_dir antennas_fname = path.join(fname_dir, beacon.antennas_fname) time_diffs_fname = 'time_diffs.hdf5' if False else antennas_fname fig_dir = args.fig_dir # set None to disable saving snr_fname = path.join(fname_dir, beacon.snr_fname) basenames, time_diffs, f_beacons, clock_phase_diffs, k_periods = beacon.read_baseline_time_diffs_hdf5(time_diffs_fname) f_beacon, tx, antennas = beacon.read_beacon_hdf5(antennas_fname) # TODO: allow multiple frequencies if (f_beacon != f_beacons).any(): raise NotImplementedError N_base = len(basenames) N_ant = len(antennas) # reshape time_diffs into N_ant x N_ant matrix clock_phase_matrix = np.full( (N_ant, N_ant), np.nan, dtype=float) ## set i=i terms to 0 for i in range(N_ant): clock_phase_matrix[i,i] = 0 ## fill matrix name2idx = lambda name: int(name)-1 for i, b in enumerate(basenames): idx = (name2idx(b[0]), name2idx(b[1])) #if idx[1] < idx[0]: # idx = (idx[1], idx[0]) 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]) mat_kwargs = dict( norm = Normalize(vmin=-np.pi, vmax=+np.pi), cmap = mpl.cm.get_cmap('Spectral_r') ) # Show Matrix as figure if True: fig, ax = plt.subplots(figsize=figsize) ax.set_title("Measured clock phase differences Baseline i,j") ax.set_ylabel("Antenna no. i") ax.set_xlabel("Antenna no. j") im = ax.imshow(clock_phase_matrix, interpolation='none', **mat_kwargs) fig.colorbar(im, ax=ax) if fig_dir: fig.savefig(path.join(fig_dir, path.basename(__file__) + f".matrix.original.pdf")) plt.close(fig) # 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 first_row = -1*(clock_phase_matrix[0,:] * np.ones_like(clock_phase_matrix)).T # Show subtraction Matrix as figure if True: fig, ax = plt.subplots(figsize=figsize) ax.set_title("Clock Phase Subtraction matrix i,j") 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: fig.savefig(path.join(fig_dir, path.basename(__file__) + f".matrix.first_row.pdf")) plt.close(fig) clock_phase_matrix = clock_phase_matrix - first_row clock_phase_matrix = lib.phase_mod(clock_phase_matrix) # Except for the first row, these are all separate measurements # Condense into phase offset per antenna if True: # do not use the first row my_mask = np.arange(1, len(clock_phase_matrix), dtype=int) else: my_mask = np.arange(0, len(clock_phase_matrix), dtype=int) mean_clock_phase = np.nanmean(clock_phase_matrix[my_mask], axis=0) std_clock_phase = np.nanstd( clock_phase_matrix[my_mask], axis=0) # 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) # Show resulting matrix as figure if True: fig, axs = plt.subplots(2,1, sharex=True, figsize=figsize) axs[0].set_title("Modified clock phase differences Baseline 0,j") axs[0].set_ylabel("Antenna no. 0") axs[-1].set_xlabel("Antenna no. j") im = axs[0].imshow(clock_phase_matrix, interpolation='none', **mat_kwargs) fig.colorbar(im, ax=axs) axs[0].set_aspect('auto') colours = [mat_kwargs['cmap'](mat_kwargs['norm'](x)) for x in mean_clock_phase] axs[1].set_ylabel("Mean Baseline Phase (0, j)[rad]") 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) if fig_dir: fig.savefig(path.join(fig_dir, path.basename(__file__) + f".matrix.modified.pdf")) plt.close(fig) # 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) h5attrs['clock_phase_mean'] = mean_clock_phase[idx] h5attrs['clock_phase_std'] = std_clock_phase[idx] ############################## # Compare actual time shifts # ############################## snrs = beacon.read_snr_file(snr_fname) snr_str = f"$\\langle SNR \\rangle$ = {snrs['mean']: .1e}" actual_antenna_time_shifts = { a.name: a.attrs['clock_offset'] for a in sorted(antennas, key=lambda a: int(a.name)) } if True: 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() ] # Make sure to shift all antennas by a global phase global_phase_shift = actual_antenna_phase_shifts[0] - mean_clock_phase[0] actual_antenna_phase_shifts = lib.phase_mod(actual_antenna_phase_shifts - global_phase_shift ) fit_info = {} for i in range(2): plot_residuals = i == 1 true_phases = actual_antenna_phase_shifts measured_phases = mean_clock_phase hist_kwargs = {} if plot_residuals: measured_phases = lib.phase_mod(measured_phases - actual_antenna_phase_shifts) fig, _fit_info = figlib.phase_comparison_figure( measured_phases, true_phases, plot_residuals=plot_residuals, f_beacon=f_beacon, figsize=figsize, hist_kwargs=hist_kwargs, fit_gaussian=plot_residuals, fit_ricianphase=plot_residuals, return_fit_info = True, ) axs = fig.get_axes() axs[0].legend(title=snr_str) if plot_residuals: axs[0].set_title("Difference between Measured and Actual phases (minus global phase)\n for Antenna $i$") axs[-1].set_xlabel("Antenna Mean Phase Residual $\\Delta_\\varphi$") else: axs[0].set_title("Comparison Measured and Actual phases (minus global phase)\n for Antenna $i$") axs[-1].set_xlabel("Antenna Mean Phase $\\varphi$") i=1 axs[i].set_ylabel("Antenna no.") #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') fig.tight_layout() if fig_dir: extra_name = "measured" if plot_residuals: extra_name = "residuals" fig.savefig(path.join(fig_dir, path.basename(__file__) + f".phase.{extra_name}.pdf")) # 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) ########################## ########################## ########################## actual_baseline_time_shifts = [] for i,b in enumerate(basenames): actual_baseline_time_shift = actual_antenna_time_shifts[b[0]] - actual_antenna_time_shifts[b[1]] actual_baseline_time_shifts.append(actual_baseline_time_shift) # unpack mean_clock_phase back into a list of time diffs measured_baseline_time_diffs = [] for i,b in enumerate(basenames): phase0, phase1 = mean_clock_phase[name2idx(b[0])], mean_clock_phase[name2idx(b[1])] measured_baseline_time_diffs.append(lib.phase_mod(phase1 - phase0)/(2*np.pi*f_beacon)) # Make a plot if True: for i in range(2): fig, ax = plt.subplots(figsize=figsize) 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: ax.plot(np.arange(N_base), actual_baseline_time_shifts, ls='none', marker='h', alpha=0.8, label='actual time shifts') else: 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') ax.legend(title=snr_str) if fig_dir: extra_name = '' if i == 1: extra_name = '.wrapped' 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]) if show_plots: plt.show()