Skip to content

Commit

Permalink
Merge branch 'master' into blchunk-udpates
Browse files Browse the repository at this point in the history
  • Loading branch information
@steven-murray
steven-murray authored Jul 31, 2024
2 parents 6015390 + 5133d40 commit 45035f8
Show file tree
Hide file tree
Showing 7 changed files with 321 additions and 106 deletions.
Binary file added hera_notebook_templates/data/SSM_autocorrelations_downsampled_sum_pol_convention.uvh5
View file
Binary file not shown.
39 changes: 22 additions & 17 deletions hera_notebook_templates/notebooks/lststack.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
"source": [
"import os\n",
"import sys\n",
"import pickle\n",
"from pathlib import Path\n",
"from functools import partial\n",
"from datetime import datetime\n",
Expand Down Expand Up @@ -146,7 +147,6 @@
"outdir: str = \".\"\n",
"bl_chunk_size: int = 0\n",
"rephase: bool = True\n",
"vis_units: str = \"Jy\"\n",
"fname_format: str = '{inpaint_mode}/zen.{kind}.{lst:7.5f}.{blchunk}.sum.uvh5'\n",
"overwrite: bool = True\n",
"write_med_mad: bool = False\n",
Expand Down Expand Up @@ -759,6 +759,7 @@
"source": [
"%%time\n",
"if do_lstcal and n_bl_chunks==1: # can't fit all the baselines in memory if not redavg'd, and need all of them at once to do lstcal\n",
" all_calibration_parameters = {}\n",
" for i, (stack, lstavg_model, auto_model) in enumerate(zip(cross_stacks, cross_lstavg, autos_lstavg)):\n",
" calibration_parameters, gains = lstbin_absolute_calibration(\n",
" stack=stack, \n",
Expand All @@ -768,8 +769,9 @@
" auto_stack=auto_stacks[i],\n",
" auto_model=auto_model['data'],\n",
" calibrate_inplace=True, # calibrate inplace\n",
" run_phase_cal=True, # run phase calibration\n",
" run_amplitude_cal=True, # run amplitude calibration\n",
" run_phase_cal=True, # run phase calibration\n",
" run_cross_pol_phase_cal=True\n",
" )\n",
" \n",
" # Write out calibration parameters and metadata\n",
Expand All @@ -779,18 +781,22 @@
" calibration_parameters[\"antpairs\"] = stack.antpairs\n",
" calibration_parameters[\"lst\"] = stackconf.lst_grid[i]\n",
" \n",
" # Get the calibration filename\n",
" cal_fname = lstbin.io.format_outfile_name(\n",
" lst=stackconf.lst_grid_edges[i],\n",
" pols=stackconf.pols,\n",
" fname_format=fname_format.replace(\"sum.uvh5\", \"lstcal.npz\"),\n",
" inpaint_mode=True,\n",
" lst_branch_cut=stackconf.properties[\"lst_branch_cut\"],\n",
" kind=\"LST\"\n",
" )\n",
" outfile = outdir / cal_fname\n",
" all_calibration_parameters[stackconf.lst_grid[i]] = calibration_parameters\n",
" \n",
" # Get the calibration filename\n",
" cal_fname = lstbin.io.format_outfile_name(\n",
" lst=stackconf.lst_grid_edges[0],\n",
" pols=stackconf.pols,\n",
" fname_format=fname_format.replace(\"sum.uvh5\", \"lstcal.pkl\"),\n",
" inpaint_mode=True,\n",
" lst_branch_cut=stackconf.properties[\"lst_branch_cut\"],\n",
" kind=\"LST\"\n",
" )\n",
" outfile = outdir / cal_fname\n",
" \n",
" np.savez(outfile, **calibration_parameters)\n",
" # Write out calibration parameters\n",
" with open(outfile, 'wb') as handle:\n",
" pickle.dump(all_calibration_parameters, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
" \n",
" \n",
" # Recompute auto and cross averages after calibration - needed for STD files \n",
Expand Down Expand Up @@ -1334,7 +1340,6 @@
" freq_max=freq_max,\n",
" lst_branch_cut=stackconf.properties[\"lst_branch_cut\"],\n",
" lsts=stackconf.lst_grid,\n",
" vis_units=vis_units\n",
")\n",
"\n",
"create_file = partial(\n",
Expand Down Expand Up @@ -2754,9 +2759,9 @@
],
"metadata": {
"kernelspec": {
"display_name": "h6c",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "h6c"
"name": "python3"
},
"language_info": {
"codemirror_mode": {
Expand All @@ -2768,7 +2773,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.6"
"version": "3.11.3"
},
"toc": {
"base_numbering": 1,
Expand Down
127 changes: 95 additions & 32 deletions notebooks/calibration_smoothing.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
"source": [
"# Calibration Smoothing\n",
"\n",
"**by Josh Dillon**, last updated March 29, 2023\n",
"**by Josh Dillon**, last updated July 19, 2023\n",
"\n",
"This notebook runs calibration smoothing to the gains coming out of [file_calibration](https://github.com/HERA-Team/hera_notebook_templates/blob/master/notebooks/file_calibration.ipynb) notebook. It removes any flags founds on by that notebook and replaces them with flags generated from [full_day_rfi](https://github.com/HERA-Team/hera_notebook_templates/blob/master/notebooks/full_day_rfi.ipynb) and [full_day_antenna_flagging](https://github.com/HERA-Team/hera_notebook_templates/blob/master/notebooks/full_day_antenna_flagging.ipynb). It also plots the results for a couple of antennas.\n",
"\n",
Expand Down Expand Up @@ -71,7 +71,7 @@
"source": [
"# get files\n",
"SUM_FILE = os.environ.get(\"SUM_FILE\", None)\n",
"# SUM_FILE = '/users/jsdillon/lustre/H6C/abscal/2459853/zen.2459853.25518.sum.uvh5'\n",
"# SUM_FILE = '/lustre/aoc/projects/hera/h6c-analysis/IDR2/2459867/zen.2459867.43004.sum.uvh5'\n",
"SUM_SUFFIX = os.environ.get(\"SUM_SUFFIX\", 'sum.uvh5')\n",
"CAL_SUFFIX = os.environ.get(\"CAL_SUFFIX\", 'sum.omni.calfits')\n",
"SMOOTH_CAL_SUFFIX = os.environ.get(\"SMOOTH_CAL_SUFFIX\", 'sum.smooth.calfits')\n",
Expand All @@ -80,9 +80,10 @@
"FREQ_SMOOTHING_SCALE = float(os.environ.get(\"FREQ_SMOOTHING_SCALE\", 10.0)) # MHz\n",
"TIME_SMOOTHING_SCALE = float(os.environ.get(\"TIME_SMOOTHING_SCALE\", 6e5)) # seconds\n",
"EIGENVAL_CUTOFF = float(os.environ.get(\"EIGENVAL_CUTOFF\", 1e-12))\n",
"PER_POL_REFANT = os.environ.get(\"PER_POL_REFANT\", \"False\").upper() == \"TRUE\"\n",
"\n",
"for setting in ['SUM_FILE', 'SUM_SUFFIX', 'CAL_SUFFIX', 'SMOOTH_CAL_SUFFIX', 'ANT_FLAG_SUFFIX',\n",
" 'RFI_FLAG_SUFFIX', 'FREQ_SMOOTHING_SCALE', 'TIME_SMOOTHING_SCALE', 'EIGENVAL_CUTOFF']:\n",
" 'RFI_FLAG_SUFFIX', 'FREQ_SMOOTHING_SCALE', 'TIME_SMOOTHING_SCALE', 'EIGENVAL_CUTOFF', 'PER_POL_REFANT']:\n",
" print(f'{setting} = {eval(setting)}')"
]
},
Expand All @@ -91,7 +92,7 @@
"id": "c6ecd16f",
"metadata": {},
"source": [
"## Load files"
"## Load files and select reference antenna(s)"
]
},
{
Expand Down Expand Up @@ -134,22 +135,41 @@
{
"cell_type": "code",
"execution_count": null,
"id": "ee831559",
"metadata": {
"scrolled": true
},
"id": "0976d9b5-b08c-4899-a866-8badc6d9173d",
"metadata": {},
"outputs": [],
"source": [
"cs = smooth_cal.CalibrationSmoother(cal_files, flag_file_list=(ant_flag_files + rfi_flag_files),\n",
" ignore_calflags=True, pick_refant=False, load_chisq=True, load_cspa=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4c2af8ac-efb8-4489-80f6-d0f30d8b62a8",
"metadata": {},
"outputs": [],
"source": [
"cs = smooth_cal.CalibrationSmoother(cal_files, flag_file_list=(ant_flag_files + rfi_flag_files), ignore_calflags=True,\n",
" pick_refant=True, propagate_refant_flags=True, load_chisq=True, load_cspa=True)\n",
"cs.refant = smooth_cal.pick_reference_antenna(cs.gain_grids, cs.flag_grids, cs.freqs, per_pol=True)\n",
"for pol in cs.refant:\n",
" print(f'Reference antenna {cs.refant[pol][0]} selected for {pol}.')"
" print(f'Reference antenna {cs.refant[pol][0]} selected for smoothing {pol} gains.')\n",
"\n",
"if not PER_POL_REFANT:\n",
" # in this case, rephase both pols separately before smoothing, but also smooth the relative polarization calibration phasor\n",
" overall_refant = smooth_cal.pick_reference_antenna({ant: cs.gain_grids[ant] for ant in cs.refant.values()}, \n",
" {ant: cs.flag_grids[ant] for ant in cs.refant.values()}, \n",
" cs.freqs, per_pol=False)\n",
" print(f'Overall reference antenna {overall_refant} selected.')\n",
" other_refant = [ant for ant in cs.refant.values() if ant != overall_refant][0]\n",
"\n",
" relative_pol_phasor = cs.gain_grids[overall_refant] * cs.gain_grids[other_refant].conj() # TODO: is this conjugation right?\n",
" relative_pol_phasor /= np.abs(relative_pol_phasor)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "482760b0",
"id": "9f459577-0ba2-4c55-a767-df9d10c7226d",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -160,7 +180,21 @@
"candidate_ants = [ant for ant, nflags in zip(antnums, flags_per_antnum) if (ant not in refant_nums) and (nflags <= np.percentile(flags_per_antnum, 25))\n",
" and not np.all(cs.flag_grids[ant, 'Jee']) and not np.all(cs.flag_grids[ant, 'Jnn'])]\n",
"ants_to_plot = [func(candidate_ants) for func in (np.min, np.max)]\n",
"abscal_gains = {(ant, pol): np.array(cs.gain_grids[(ant, pol)]) for ant in ants_to_plot for pol in ['Jee', 'Jnn']}"
"abscal_gains = {}\n",
"for pol in ['Jee', 'Jnn']:\n",
" for antnum in ants_to_plot:\n",
" refant_here = (cs.refant[pol] if PER_POL_REFANT else overall_refant)\n",
" abscal_gains[antnum, pol] = cs.gain_grids[(antnum, pol)] * np.abs(cs.gain_grids[refant_here]) / cs.gain_grids[refant_here]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ff18e5a9-23e5-41f5-ab0b-35f09ab67404",
"metadata": {},
"outputs": [],
"source": [
"cs.rephase_to_refant(propagate_refant_flags=True)"
]
},
{
Expand All @@ -171,6 +205,19 @@
"## Perform smoothing"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "178ce9a5-570e-4cbc-a7d0-978da71033a6",
"metadata": {},
"outputs": [],
"source": [
"if not PER_POL_REFANT:\n",
" # treat the relative_pol_phasor as if it were antenna -1\n",
" cs.gain_grids[(-1, other_refant[1])] = relative_pol_phasor\n",
" cs.flag_grids[(-1, other_refant[1])] = cs.flag_grids[overall_refant] | cs.flag_grids[other_refant]"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand All @@ -182,6 +229,23 @@
" method='DPSS', fit_method='lu_solve', fix_phase_flips=True, flag_phase_flip_ints=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "10ff0705-f3fd-4570-b47a-e96e360d183c",
"metadata": {},
"outputs": [],
"source": [
"if not PER_POL_REFANT:\n",
" # put back in the smoothed phasor, ensuring the amplitude is 1 and that data are flagged anywhere either polarization's refant is flagged\n",
" smoothed_relative_pol_phasor = cs.gain_grids[(-1, other_refant[-1])] / np.abs(cs.gain_grids[(-1, other_refant[-1])])\n",
" for ant in cs.gain_grids:\n",
" if ant[0] >= 0 and ant[1] == other_refant[1]:\n",
" cs.gain_grids[ant] /= smoothed_relative_pol_phasor\n",
" cs.flag_grids[ant] |= (cs.flag_grids[(-1, other_refant[1])])\n",
" cs.refant = overall_refant"
]
},
{
"cell_type": "markdown",
"id": "c40160e2",
Expand All @@ -205,9 +269,7 @@
"cell_type": "code",
"execution_count": null,
"id": "4fe13af1",
"metadata": {
"scrolled": false
},
"metadata": {},
"outputs": [],
"source": [
"def amplitude_plot(ant_to_plot):\n",
Expand Down Expand Up @@ -289,14 +351,16 @@
" warnings.simplefilter(\"ignore\") \n",
" display(HTML(f'<h2>Antenna {ant_to_plot} Phase Waterfalls</h2>'))\n",
" fig, axes = plt.subplots(4, 2, figsize=(14,14), gridspec_kw={'height_ratios': [1, 1, .4, .4]})\n",
"\n",
" \n",
" # Plot phase waterfalls for a single antenna \n",
" for ax, pol in zip(axes[0], ['Jee', 'Jnn']):\n",
" ant = (ant_to_plot, pol)\n",
" extent=[cs.freqs[0]/1e6, cs.freqs[-1]/1e6, lst_grid[-1], lst_grid[0]]\n",
" im = ax.imshow(np.where(cs.flag_grids[ant], np.nan, np.angle(cs.gain_grids[ant])), aspect='auto', cmap='inferno', \n",
" interpolation='nearest', vmin=-np.pi, vmax=np.pi, extent=extent)\n",
" ax.set_title(f'Smoothcal Gain Phase of Ant {ant[0]} / Ant {cs.refant[pol][0]}: {pol[-1]}-polarized')\n",
"\n",
" refant = (cs.refant[pol] if isinstance(cs.refant, dict) else cs.refant)\n",
" ax.set_title(f'Smoothcal Gain Phase of Ant {ant[0]}{pol[-1]} / Ant {refant[0]}{refant[1][-1]}')\n",
" ax.set_xlabel('Frequency (MHz)')\n",
" ax.set_ylabel('LST (Hours)')\n",
" ax.set_xlim([cs.freqs[0]/1e6, cs.freqs[-1]/1e6])\n",
Expand All @@ -309,7 +373,8 @@
" extent=[cs.freqs[0]/1e6, cs.freqs[-1]/1e6, lst_grid[-1], lst_grid[0]]\n",
" im = ax.imshow(np.where(cs.flag_grids[ant], np.nan, np.angle(abscal_gains[ant])), aspect='auto', cmap='inferno', \n",
" interpolation='nearest', vmin=-np.pi, vmax=np.pi, extent=extent)\n",
" ax.set_title(f'Abscal Gain Phase of Ant {ant[0]} / Ant {cs.refant[pol][0]}: {pol[-1]}-polarized')\n",
" refant = (cs.refant[pol] if isinstance(cs.refant, dict) else cs.refant)\n",
" ax.set_title(f'Abscal Gain Phase of Ant {ant[0]}{pol[-1]} / Ant {refant[0]}{refant[1][-1]}')\n",
" ax.set_xlabel('Frequency (MHz)')\n",
" ax.set_ylabel('LST (Hours)')\n",
" ax.set_xlim([cs.freqs[0]/1e6, cs.freqs[-1]/1e6])\n",
Expand All @@ -326,8 +391,9 @@
" ax.set_ylim([-np.pi, np.pi])\n",
" ax.set_xlim([cs.freqs[0]/1e6, cs.freqs[-1]/1e6]) \n",
" ax.set_xlabel('Frequency (MHz)')\n",
" ax.set_ylabel(f'Phase of g$_{{{ant[0]}}}$ / g$_{{{cs.refant[pol][0]}}}$')\n",
" ax.set_title(f'Median Infrequently-Flagged Gain Phase of Ant {ant[0]} / Ant {cs.refant[pol][0]}: {pol[-1]}-polarized')\n",
" refant = (cs.refant[pol] if isinstance(cs.refant, dict) else cs.refant)\n",
" ax.set_ylabel(f'Phase of g$_{{{ant[0]}{pol[-1]}}}$ / g$_{{{refant[0]}{refant[1][-1]}}}$')\n",
" ax.set_title(f'Median Infrequently-Flagged Gain Phase of Ant {ant[0]}{pol[-1]} / Ant {refant[0]}{refant[1][-1]}')\n",
" ax.legend(loc='upper left')\n",
"\n",
" # # Now plot median gain time series\n",
Expand All @@ -339,8 +405,9 @@
" ax.plot(lst_grid[to_plot], np.nanmean(np.where(cs.flag_grids[ant], np.nan, np.angle(cs.gain_grids[ant])), axis=1)[to_plot], 'k.', ms=2, label='Smoothed') \n",
" ax.set_ylim([-np.pi, np.pi]) \n",
" ax.set_xlabel('LST (hours)')\n",
" ax.set_ylabel(f'Phase of g$_{{{ant[0]}}}$ / g$_{{{cs.refant[pol][0]}}}$')\n",
" ax.set_title(f'Mean Infrequently-Flagged Gain Phase of Ant {ant[0]} / Ant {cs.refant[pol][0]}: {pol[-1]}-polarized')\n",
" refant = (cs.refant[pol] if isinstance(cs.refant, dict) else cs.refant)\n",
" ax.set_ylabel(f'Phase of g$_{{{ant[0]}{pol[-1]}}}$ / g$_{{{refant[0]}{refant[1][-1]}}}$')\n",
" ax.set_title(f'Mean Infrequently-Flagged Gain Phase of Ant {ant[0]}{pol[-1]} / Ant {refant[0]}{refant[1][-1]}')\n",
" ax.set_xticklabels(ax.get_xticks() % 24) \n",
" ax.legend(loc='upper left')\n",
"\n",
Expand All @@ -362,9 +429,7 @@
"cell_type": "code",
"execution_count": null,
"id": "43438042",
"metadata": {
"scrolled": false
},
"metadata": {},
"outputs": [],
"source": [
"for ant_to_plot in ants_to_plot:\n",
Expand All @@ -385,9 +450,7 @@
"cell_type": "code",
"execution_count": null,
"id": "5c6fa5d0",
"metadata": {
"scrolled": false
},
"metadata": {},
"outputs": [],
"source": [
"for ant_to_plot in ants_to_plot:\n",
Expand All @@ -413,8 +476,8 @@
" fig, axes = plt.subplots(1, 2, figsize=(14, 10), sharex=True, sharey=True)\n",
" extent = [cs.freqs[0]/1e6, cs.freqs[-1]/1e6, lst_grid[-1], lst_grid[0]]\n",
" for ax, pol in zip(axes, ['Jee', 'Jnn']):\n",
"\n",
" im = ax.imshow(np.where(cs.flag_grids[cs.refant[pol]], np.nan, cs.chisq_grids[pol]), vmin=1, vmax=5, \n",
" refant = (cs.refant[pol] if isinstance(cs.refant, dict) else cs.refant)\n",
" im = ax.imshow(np.where(cs.flag_grids[refant], np.nan, cs.chisq_grids[pol]), vmin=1, vmax=5, \n",
" aspect='auto', cmap='turbo', interpolation='none', extent=extent)\n",
" ax.set_yticklabels(ax.get_yticks() % 24)\n",
" ax.set_title(f'{pol[1:]}-Polarized $\\\\chi^2$ / DoF')\n",
Expand Down Expand Up @@ -605,7 +668,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.16"
"version": "3.10.12"
},
"toc": {
"base_numbering": 1,
Expand Down
Loading

0 comments on commit 45035f8

Please sign in to comment.