import os
os.environ['HDF5_USE_FILE_LOCKING'] = 'FALSE'
import h5py
import hdf5plugin
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
import matplotlib.patches as mpatches
import matplotlib.gridspec as gridspec
import numpy as np
from pyuvdata import UVCal, UVData
import sys
import glob
import uvtools as uvt
from astropy.time import Time
from astropy.coordinates import EarthLocation, AltAz, Angle
from astropy.coordinates import SkyCoord as sc
import pandas
import warnings 
import copy
from hera_notebook_templates import utils
import hera_qm
from hera_mc import cm_hookup
import importlib
from scipy import stats
from IPython.display import display, HTML
#warnings.filterwarnings('ignore')

%matplotlib inline
%config InlineBackend.figure_format = 'retina'
display(HTML("<style>.container { width:100% !important; }</style>"))

#get data location
data_path = os.environ['DATA_PATH']
print(f'DATA_PATH = "{data_path}"')
statuses = os.environ['APRIORI_STATUSES']
print(f'APRIORI_STATUSES = {statuses}')
JD = os.environ['JULIANDATE']
print(f'JULIANDATE = {JD}')
utc = Time(JD, format='jd').datetime
print(f'Date = {utc.month}-{utc.day}-{utc.year}')

DATA_PATH = "/mnt/sn1/2460373"
APRIORI_STATUSES = dish_maintenance,dish_ok,RF_maintenance,RF_ok,digital_ok,digital_maintenance,calibration_maintenance,calibration_triage,calibration_ok
JULIANDATE = 2460373
Date = 3-3-2024

# Load in data
HHfiles, difffiles, HHautos, diffautos, uvdx, uvdy = utils.load_data(data_path,JD)
    
uvd = UVData()
unread = True
readInd=0
while unread and readInd<len(HHautos):
    try:
        uvd.read(HHautos[readInd])
        unread = False
    except:
        readInd += 1
        continue
use_ants = utils.get_use_ants(uvd,statuses,JD)
print(f'This day contains {len(use_ants)} antennas of the given status category.')
uvd.read(HHautos[::10], skip_bad_files=True, antenna_nums = use_ants)
lsts = uvd.lst_array
uvdx.select(antenna_nums=use_ants)
uvdy.select(antenna_nums=use_ants)

183 sum files found between JDs 2460373.18986 and 2460373.23639
201 diff files found between JDs 2460373.18986 and 2460373.23639
182 sum auto files found between JDs 2460373.18986 and 2460373.23639
182 diff auto files found between JDs 2460373.18986 and 2460373.23639
############################################################
######### DIFFERENT NUMBER OF SUM AND DIFF FILES ###########
############################################################

This day contains 252 antennas of the given status category.

Sky Coverage Map

Map of the sky (made using the Haslam 408MHz map). The RA/DEC range covered by this night of observation is shaded based on a 12 degree FWHM of the beam. Horizontal dashed lines represent the stripe that HERA can observe, while the shaded region is what was observed on this night. Vertical lines represent the beginning and ending LSTs of this observation. Selected sources are labelled, sources included are those in the GLEAM 4Jy catalog with a flux >10.9 Jy. Note that the map is clipped at the northern horizon.

sources = utils.gather_source_list()
utils.plot_sky_map(uvd,dec_pad=55,ra_pad=55,clip=False,sources=sources)

LST Coverage

Shows the LSTs (in hours) and JDs for which data is collected. Green represents data, red means no data.

utils.plot_lst_coverage(uvd)

Autocorrelations for a single file

This plot shows autocorrelations for one timestamp of each antenna that is active and each polarization. For each node, antennas are ordered by SNAP number, and within that by SNAP input number. The antenna number label color corresponds to the a priori status of that antenna.

### plot autos
utils.plot_autos(uvdx, uvdy)

Waterfalls of Autocorrelation Amplitudes for each Antenna and Each polarization

These plots show autocorrelation waterfalls of each antenna that is active and whose status qualifies for this notebook. For each node, antennas are ordered by SNAP number, and within that by SNAP input number. The antenna number label color corresponds to the a priori status of that antenna.

utils.plot_wfs(uvd, pol = 0)

utils.plot_wfs(uvd, pol = 1)

Correlation Metrics

The first plot shows the correlation metric (described below) for a set of baseline types, as calculated at several times throughout the night. It is expected that longer baselines (darker color) will exhibit lower values than the short baselines.

The matrices show the phase correlation between antennas. Using the even and odd visibilities, each pixel is calculated as (even/abs(even)) * (conj(odd)/abs(odd)), and then averaged across time and frequency. If the phases are noise-like, this value will average down to zero. If the antennas are well correlated, the phases should not be noise-like, and this value should average to 1. The lines denoting node boundaries are intended to help confirm that inter-node correlations are functioning - if they aren't, this plot will appear block-diagonal.

This metric has shown to be LST locked - when comparing to other nights, be sure to compare for the same LST. It is expected that some LSTs will look much better or worse than others.

Note: Within each node, the order of antennas is determined by snap, and within that by snap input number.

badAnts = []
badAnts = utils.plotNodeAveragedSummary(uvd,HHfiles,JD,use_ants,mat_pols=['xx','yy','xy','yx'])

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_18651/276258673.py in ?()
      1 badAnts = []
----> 2 badAnts = utils.plotNodeAveragedSummary(uvd,HHfiles,JD,use_ants,mat_pols=['xx','yy','xy','yx'])

~/mambaforge/envs/RTP/lib/python3.12/site-packages/hera_notebook_templates/utils.py in ?(uv, HHfiles, jd, use_ants, pols, mat_pols, baseline_groups, removeBadAnts, plotRatios, plotSummary)
    807     if baseline_groups == []:
    808         baseline_groups = [(14,0,'14m E-W'),(14,-11,'14m NW-SE'),(14,11,'14m SW-NE'),(29,0,'29m E-W'),(29,22,'29m SW-NE'),
    809                        (44,0,'44m E-W'),(58.5,0,'58m E-W'),(73,0,'73m E-W'),(87.6,0,'88m E-W'),
    810                       (102.3,0,'102m E-W')]
--> 811     nodeMedians,lsts,badAnts=get_correlation_baseline_evolutions(uv,HHfiles,jd,use_ants,pols=pols,mat_pols=mat_pols,
    812                                                                 bl_type=baseline_groups,removeBadAnts=removeBadAnts,
    813                                                                 plotRatios=plotRatios)
    814     pols = mat_pols

~/mambaforge/envs/RTP/lib/python3.12/site-packages/hera_notebook_templates/utils.py in ?(uv, HHfiles, jd, use_ants, badThresh, pols, bl_type, removeBadAnts, plotMatrix, mat_pols, plotRatios)
   1473         LSTs that metric was calculated for, spaced 1 hour apart.
   1474     bad_antennas: List
   1475         Antenna numbers flagged as bad based on badThresh parameter.
   1476     """
-> 1477     files, lsts, inds = get_hourly_files(uv, HHfiles, jd)
   1478     if use_ants == 'auto':
   1479         use_ants = uv.get_ants()
   1480     if plotRatios is True:

~/mambaforge/envs/RTP/lib/python3.12/site-packages/hera_notebook_templates/utils.py in ?(uv, HHfiles, jd)
   1385     for i,file in enumerate(HHfiles):
   1386         try:
   1387             dat = UVData()
   1388             dat.read(file, read_data=False)
-> 1389         except KeyError:
   1390             continue
   1391         jd = dat.time_array[0]
   1392         t = Time(jd,format='jd',location=loc)

~/mambaforge/envs/RTP/lib/python3.12/site-packages/pyuvdata/uvdata/uvdata.py in ?(self, filename, axis, file_type, read_data, skip_bad_files, background_lsts, astrometry_library, ignore_name, use_future_array_shapes, allow_rephase, make_multi_phase, antenna_nums, antenna_names, ant_str, bls, catalog_names, frequencies, freq_chans, times, time_range, lsts, lst_range, polarizations, blt_inds, phase_center_ids, keep_all_metadata, run_check, check_extra, run_check_acceptability, strict_uvw_antpos_check, check_autos, fix_autos, phase_type, projected, correct_lat_lon, calc_lst, fix_old_proj, fix_use_ant_pos, use_model, data_column, pol_order, ignore_single_chan, raise_error, read_weights, allow_flex_pol, multidim_index, remove_flex_pol, blt_order, blts_are_rectangular, time_axis_faster_than_bls, data_array_dtype, use_aoflagger_flags, remove_dig_gains, remove_coarse_band, correct_cable_len, correct_van_vleck, cheby_approx, flag_small_auto_ants, propagate_coarse_flags, flag_init, edge_width, start_flag, end_flag, flag_dc_offset, remove_flagged_ants, phase_to_pointing_center, nsample_array_dtype, corrchunk, receivers, sidebands, mir_select_where, apply_tsys, apply_flags, apply_dedoppler, pseudo_cont, rechunk, compass_soln, swarm_only, codes_check, recompute_nbls)
  12656                     astrometry_library=astrometry_library,
  12657                 )
  12658 
  12659             elif file_type == "uvh5":
> 12660                 self.read_uvh5(
  12661                     filename,
  12662                     antenna_nums=antenna_nums,
  12663                     antenna_names=antenna_names,

~/mambaforge/envs/RTP/lib/python3.12/site-packages/pyuvdata/uvdata/uvdata.py in ?(self, filename, **kwargs)
  11467                 "Use the generic `uvdata.read` function instead."
  11468             )
  11469 
  11470         uvh5_obj = uvh5.UVH5()
> 11471         uvh5_obj.read_uvh5(filename, **kwargs)
  11472         self._convert_from_filetype(uvh5_obj)
  11473         del uvh5_obj

~/mambaforge/envs/RTP/lib/python3.12/site-packages/pyuvdata/uvdata/uvh5.py in ?(self, filename, antenna_nums, antenna_names, ant_str, bls, frequencies, freq_chans, times, time_range, lsts, lst_range, polarizations, blt_inds, phase_center_ids, catalog_names, keep_all_metadata, read_data, data_array_dtype, multidim_index, remove_flex_pol, background_lsts, run_check, check_extra, run_check_acceptability, strict_uvw_antpos_check, fix_old_proj, fix_use_ant_pos, check_autos, fix_autos, use_future_array_shapes, blt_order, blts_are_rectangular, time_axis_faster_than_bls, recompute_nbls, astrometry_library)
   1541                     self.use_current_array_shapes()
   1542 
   1543         # check if object has all required UVParameters set
   1544         if run_check:
-> 1545             self.check(
   1546                 check_extra=check_extra,
   1547                 run_check_acceptability=run_check_acceptability,
   1548                 strict_uvw_antpos_check=strict_uvw_antpos_check,

~/mambaforge/envs/RTP/lib/python3.12/site-packages/pyuvdata/uvdata/uvdata.py in ?(self, check_extra, run_check_acceptability, check_freq_spacing, strict_uvw_antpos_check, allow_flip_conj, check_autos, fix_autos, lst_tol)
   3299 
   3300         # Check internal consistency of numbers which don't explicitly correspond
   3301         # to the shape of another array.
   3302         if self.Nants_data != self._calc_nants_data():
-> 3303             raise ValueError(
   3304                 "Nants_data must be equal to the number of unique "
   3305                 "values in ant_1_array and ant_2_array"
   3306             )

ValueError: Nants_data must be equal to the number of unique values in ant_1_array and ant_2_array

Visibility amplitude spectra for a set of redundant baselines, labeled by inter vs. intranode baselines. The red and blue should exhibit the same bandpass shape - if the red are consistently different from the blue, this indicates an issue with internode correlations.

Note: antennas that were identified as bad by the correlation matrix have been removed from this plot.

utils.plotVisibilitySpectra(HHfiles[len(HHfiles)//2+1], JD, use_ants, badAnts=[])

<Figure size 640x480 with 0 Axes>

Even and Odd File Checks

A waterfall showing the ratio between the even and odd visibilities. The purpose of this is to highlight xengine failures, which will cause this value to fall to zero or go to infinity. If things are working properly, this value should be stable at 1. The boundaries between different x-engines are shown by the vertical white lines.

if len(HHautos) == len(diffautos):
    uvd_diff = UVData()
    uvd_diff.read(diffautos[::10], skip_bad_files=True, antenna_nums=use_ants)
    rat = utils.plotEvenOddWaterfalls(uvd,uvd_diff)
else:
    uvd_diff = UVData()
    use_diffs = [f for f in diffautos if '%s/zen.%s.%s.sum.autos.uvh5' % (data_path,f.split('.')[1],f.split('.')[2]) in HHautos[::10]]
    uvd_diff.read(use_diffs, skip_bad_files=True, antenna_nums = use_ants)
    uvd_sum = uvd.select(times=np.unique(uvd_diff.time_array),inplace=False)
    rat = utils.plotEvenOddWaterfalls(uvd_sum,uvd_diff)

Crossed Antenna Check

These are differences between different panels of the correlation matrices shown above (see panel titles for specifics). Antennas showing as consistently blue are ones which are correlating stronger in the cross pols than in the auto pols, indicating that the antenna polarizations are likely crossed.

crossedAnts = utils.plotNodeAveragedSummary(uvd,HHfiles,JD,use_ants,mat_pols=['xx','yy','xy','yx'],plotRatios=True,
                                       plotSummary=False)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_18651/3454784400.py in ?()
----> 1 crossedAnts = utils.plotNodeAveragedSummary(uvd,HHfiles,JD,use_ants,mat_pols=['xx','yy','xy','yx'],plotRatios=True,
      2                                        plotSummary=False)

~/mambaforge/envs/RTP/lib/python3.12/site-packages/hera_notebook_templates/utils.py in ?(uv, HHfiles, jd, use_ants, pols, mat_pols, baseline_groups, removeBadAnts, plotRatios, plotSummary)
    807     if baseline_groups == []:
    808         baseline_groups = [(14,0,'14m E-W'),(14,-11,'14m NW-SE'),(14,11,'14m SW-NE'),(29,0,'29m E-W'),(29,22,'29m SW-NE'),
    809                        (44,0,'44m E-W'),(58.5,0,'58m E-W'),(73,0,'73m E-W'),(87.6,0,'88m E-W'),
    810                       (102.3,0,'102m E-W')]
--> 811     nodeMedians,lsts,badAnts=get_correlation_baseline_evolutions(uv,HHfiles,jd,use_ants,pols=pols,mat_pols=mat_pols,
    812                                                                 bl_type=baseline_groups,removeBadAnts=removeBadAnts,
    813                                                                 plotRatios=plotRatios)
    814     pols = mat_pols

~/mambaforge/envs/RTP/lib/python3.12/site-packages/hera_notebook_templates/utils.py in ?(uv, HHfiles, jd, use_ants, badThresh, pols, bl_type, removeBadAnts, plotMatrix, mat_pols, plotRatios)
   1473         LSTs that metric was calculated for, spaced 1 hour apart.
   1474     bad_antennas: List
   1475         Antenna numbers flagged as bad based on badThresh parameter.
   1476     """
-> 1477     files, lsts, inds = get_hourly_files(uv, HHfiles, jd)
   1478     if use_ants == 'auto':
   1479         use_ants = uv.get_ants()
   1480     if plotRatios is True:

~/mambaforge/envs/RTP/lib/python3.12/site-packages/hera_notebook_templates/utils.py in ?(uv, HHfiles, jd)
   1385     for i,file in enumerate(HHfiles):
   1386         try:
   1387             dat = UVData()
   1388             dat.read(file, read_data=False)
-> 1389         except KeyError:
   1390             continue
   1391         jd = dat.time_array[0]
   1392         t = Time(jd,format='jd',location=loc)

~/mambaforge/envs/RTP/lib/python3.12/site-packages/pyuvdata/uvdata/uvdata.py in ?(self, filename, axis, file_type, read_data, skip_bad_files, background_lsts, astrometry_library, ignore_name, use_future_array_shapes, allow_rephase, make_multi_phase, antenna_nums, antenna_names, ant_str, bls, catalog_names, frequencies, freq_chans, times, time_range, lsts, lst_range, polarizations, blt_inds, phase_center_ids, keep_all_metadata, run_check, check_extra, run_check_acceptability, strict_uvw_antpos_check, check_autos, fix_autos, phase_type, projected, correct_lat_lon, calc_lst, fix_old_proj, fix_use_ant_pos, use_model, data_column, pol_order, ignore_single_chan, raise_error, read_weights, allow_flex_pol, multidim_index, remove_flex_pol, blt_order, blts_are_rectangular, time_axis_faster_than_bls, data_array_dtype, use_aoflagger_flags, remove_dig_gains, remove_coarse_band, correct_cable_len, correct_van_vleck, cheby_approx, flag_small_auto_ants, propagate_coarse_flags, flag_init, edge_width, start_flag, end_flag, flag_dc_offset, remove_flagged_ants, phase_to_pointing_center, nsample_array_dtype, corrchunk, receivers, sidebands, mir_select_where, apply_tsys, apply_flags, apply_dedoppler, pseudo_cont, rechunk, compass_soln, swarm_only, codes_check, recompute_nbls)
  12656                     astrometry_library=astrometry_library,
  12657                 )
  12658 
  12659             elif file_type == "uvh5":
> 12660                 self.read_uvh5(
  12661                     filename,
  12662                     antenna_nums=antenna_nums,
  12663                     antenna_names=antenna_names,

~/mambaforge/envs/RTP/lib/python3.12/site-packages/pyuvdata/uvdata/uvdata.py in ?(self, filename, **kwargs)
  11467                 "Use the generic `uvdata.read` function instead."
  11468             )
  11469 
  11470         uvh5_obj = uvh5.UVH5()
> 11471         uvh5_obj.read_uvh5(filename, **kwargs)
  11472         self._convert_from_filetype(uvh5_obj)
  11473         del uvh5_obj

~/mambaforge/envs/RTP/lib/python3.12/site-packages/pyuvdata/uvdata/uvh5.py in ?(self, filename, antenna_nums, antenna_names, ant_str, bls, frequencies, freq_chans, times, time_range, lsts, lst_range, polarizations, blt_inds, phase_center_ids, catalog_names, keep_all_metadata, read_data, data_array_dtype, multidim_index, remove_flex_pol, background_lsts, run_check, check_extra, run_check_acceptability, strict_uvw_antpos_check, fix_old_proj, fix_use_ant_pos, check_autos, fix_autos, use_future_array_shapes, blt_order, blts_are_rectangular, time_axis_faster_than_bls, recompute_nbls, astrometry_library)
   1541                     self.use_current_array_shapes()
   1542 
   1543         # check if object has all required UVParameters set
   1544         if run_check:
-> 1545             self.check(
   1546                 check_extra=check_extra,
   1547                 run_check_acceptability=run_check_acceptability,
   1548                 strict_uvw_antpos_check=strict_uvw_antpos_check,

~/mambaforge/envs/RTP/lib/python3.12/site-packages/pyuvdata/uvdata/uvdata.py in ?(self, check_extra, run_check_acceptability, check_freq_spacing, strict_uvw_antpos_check, allow_flip_conj, check_autos, fix_autos, lst_tol)
   3299 
   3300         # Check internal consistency of numbers which don't explicitly correspond
   3301         # to the shape of another array.
   3302         if self.Nants_data != self._calc_nants_data():
-> 3303             raise ValueError(
   3304                 "Nants_data must be equal to the number of unique "
   3305                 "values in ant_1_array and ant_2_array"
   3306             )

ValueError: Nants_data must be equal to the number of unique values in ant_1_array and ant_2_array

Antenna Positions

Antennas outlined in black here have been identified by the correlation matrix as bad antennas. Antennas with a colorful outline correspond to their status as identified by ant_metrics (see above plot). Faded antennas are those not meeting the apriori status requirement for this notebook run. Gold stars are node box locations.

uvd1 = UVData()
uvd1.read(HHfiles[readInd], skip_bad_files=True)
utils.plot_antenna_positions(uvd1, badAnts=badAnts,use_ants=use_ants)

Observer Inspection Plots

Antennas of status digital_OK or better that are flagged as bad by any of the above metrics are plotted here so observers can inspect their failures in more detail. Additionally, a 'good' template has been used to identify outliers. The upper line plots are averages over the whole observation, and the lower line plots are slices of a single time in the middle of the observation. These plots are recommended diagnostics for demoting antennas to lower statuses or reporting issues. If the plots below look OK, check other plots in notebook to hunt why the antenna was flagged. NOTE: The colorbar/power scales in these plots are NOT locked between antennas OR polarizations so that the detail will be visible on all plots. Be sure to check for reasonable power levels, as this may be the reason the antenna was flagged for inspection.

d, tempAnts = utils.flag_by_template(uvd,HHautos,JD,use_ants=use_ants,pols=['XX','YY'],plotMap=False)

inspectAnts = utils.plot_inspect_ants(uvd,JD,badAnts=badAnts,use_ants=use_ants,
                                     tempAnts=tempAnts,crossedAnts=crossedAnts)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[15], line 2
      1 inspectAnts = utils.plot_inspect_ants(uvd,JD,badAnts=badAnts,use_ants=use_ants,
----> 2                                      tempAnts=tempAnts,crossedAnts=crossedAnts)

NameError: name 'crossedAnts' is not defined

Mean-Subtracted Waterfalls

Here the mean value in each frequency bin has been subtracted out. This effectively subtracts out the bandpass shape, making time variations more visible.

utils.plot_wfs(uvd,0,mean_sub=True,jd=JD)
utils.plot_wfs(uvd,1,mean_sub=True,jd=JD)