Coverage for PyFHD/plotting/calibration.py: 6%
95 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-07-01 10:58 +0800
1import numpy as np
2import matplotlib.pyplot as plt
3import matplotlib.ticker as ticker
4import matplotlib
5from pathlib import Path
7matplotlib.use("pdf")
10def plot_cals(obs: dict, cal: dict, pyfhd_config: dict):
11 """
12 Plot the calibration solutions, the calibration residuals, and the raw calibration solutions
13 in a grid of 128 per page.
15 Parameters
16 ----------
17 obs : dict
18 Observation metadata dictionary
19 cal : dict
20 Calibration dictionary
21 pyfhd_config : dict
22 Run option dictionary
23 """
25 # Plotting only unflagged frequencies
26 freq_i_use = np.where(obs["baseline_info"]["freq_use"])[0]
27 freq_arr_use = obs["baseline_info"]["freq"][freq_i_use] / 1e6
29 # Plotting only unflagged tiles
30 tile_i_use = np.where(obs["baseline_info"]["tile_use"])[0]
32 # Plotting amplitude and phase for the gain, gain_residual, and their sum (raw solutions)
33 obs_id = pyfhd_config["obs_id"]
34 cal_plot_dir = Path(pyfhd_config["output_dir"], "plots", "calibration")
35 cal_plot_dir.mkdir(parents=True, exist_ok=True)
36 n_types = 3 # plotting three types of amp and phase gains
37 save_path_roots = [
38 Path(cal_plot_dir, f"{obs_id}_cal_amp"),
39 Path(cal_plot_dir, f"{obs_id}_cal_phase"),
40 Path(cal_plot_dir, f"{obs_id}_cal_residual_amp"),
41 Path(cal_plot_dir, f"{obs_id}_cal_residual_phase"),
42 Path(cal_plot_dir, f"{obs_id}_cal_raw_amp"),
43 Path(cal_plot_dir, f"{obs_id}_cal_raw_phase"),
44 ]
46 # Calibration solutions are referenced mutliple times, put them in variables for speed
47 cal_sol_amp = np.abs(cal["gain"][:, freq_i_use, :])
48 cal_sol_phase = np.arctan2(
49 (cal["gain"][:, freq_i_use, :]).imag, (cal["gain"][:, freq_i_use, :]).real
50 )
51 cal_raw_amp = np.abs(
52 cal["gain"][:, freq_i_use, :] + cal["gain_residual"][:, freq_i_use, :]
53 )
54 cal_raw_phase = np.arctan2(
55 (cal["gain"][:, freq_i_use, :] + cal["gain_residual"][:, freq_i_use, :]).imag,
56 (cal["gain"][:, freq_i_use, :] + cal["gain_residual"][:, freq_i_use, :]).real,
57 )
59 # NOTE: The residuals that are plotted are the differences in raw amp/phase and solution amp/phase,
60 # *not* the amp/phase of the raw and solution difference
61 cal_res_amp = cal_raw_amp - cal_sol_amp
62 cal_res_phase = np.unwrap(cal_raw_phase, axis=1) - np.unwrap(cal_sol_phase, axis=1)
64 # Find the min/max amplitude and phase for plotting, only using unflagged tiles
65 amp_minmax = np.zeros((6))
66 phase_minmax = np.zeros((6))
67 amp_minmax[0:2] = [np.nanmin(cal_sol_amp), np.nanmax(cal_sol_amp)]
68 amp_minmax[2:4] = [np.nanmin(cal_res_amp), np.nanmax(cal_res_amp)]
69 amp_minmax[4:6] = [np.nanmin(cal_raw_amp), np.nanmax(cal_raw_amp)]
70 phase_minmax[0:2] = [-np.pi - 0.1, np.pi + 0.1]
71 phase_minmax[2:4] = [np.nanmin(cal_res_phase), np.nanmax(cal_res_phase)]
72 phase_minmax[4:6] = [-np.pi - 0.1, np.pi + 0.1]
74 # Create patches for the legend
75 patches = [
76 plt.plot(
77 [],
78 [],
79 color=(["blue", "red"])[pol_i],
80 linestyle="solid",
81 label="{:s}".format(obs["pol_names"][pol_i]),
82 )[0]
83 for pol_i in range(cal["n_pol"])
84 ]
86 # Arrange calibration subplots in rows, cols with a maximum per page.
87 # Multiple pages will be required beyond 128 stations
88 n_rows = 8
89 n_cols = 16
90 n_pages = obs["n_tile"] // (n_cols * n_rows)
92 # Create dictionary for the subplot grid to create reliable, static plots
93 gridspec_kw = {
94 "left": 0.05,
95 "right": 0.98,
96 "top": 0.92,
97 "bottom": 0.05,
98 "hspace": 0.35,
99 }
101 # Create calibration solution plots for each page of 128 maximum stations for the specified types
102 for type_i in range(n_types):
103 for page_i in range(n_pages):
105 # Calculate number of stations for this page
106 if obs["n_tile"] // (n_cols * n_rows * (page_i + 1)) == 0:
107 page_tiles = obs["n_tile"] - (n_cols * n_rows * (page_i + 1))
108 else:
109 page_tiles = n_cols * n_rows
111 # Define the subplots, add extra rows on the figsize to account for labels
112 fig_amp, ax_amp = plt.subplots(
113 n_rows, n_cols, figsize=(n_cols, n_rows + 2), gridspec_kw=gridspec_kw
114 )
115 fig_phase, ax_phase = plt.subplots(
116 n_rows, n_cols, figsize=(n_cols, n_rows + 2), gridspec_kw=gridspec_kw
117 )
119 # Cycle through each row and col on the page and set tick and labels
120 for row_i in range(n_rows):
121 for col_i in range(n_cols):
123 # Set fontsize maximum number of ticks on xaxis to 2
124 ax_amp[row_i, col_i].xaxis.set_major_locator(ticker.MaxNLocator(2))
125 ax_phase[row_i, col_i].xaxis.set_major_locator(
126 ticker.MaxNLocator(2)
127 )
128 ax_amp[row_i, col_i].tick_params(labelsize=8, direction="in")
129 ax_phase[row_i, col_i].tick_params(labelsize=8, direction="in")
131 # Only set tick labels on the leftmost column and the bottommost row
132 if col_i > 0:
133 ax_amp[row_i, col_i].set_yticks([])
134 ax_phase[row_i, col_i].set_yticks([])
135 if row_i < 7:
136 ax_amp[row_i, col_i].set_xticks([])
137 ax_phase[row_i, col_i].set_xticks([])
139 # Set the y-axis min and max given the plot type
140 ax_amp[row_i, col_i].set_ylim(
141 amp_minmax[2 * type_i], amp_minmax[2 * type_i + 1]
142 )
143 ax_phase[row_i, col_i].set_ylim(
144 phase_minmax[2 * type_i], phase_minmax[2 * type_i + 1]
145 )
147 # Plot each tile on the specified subplot
148 tile_start = page_i * (n_cols * n_rows)
149 for tile_i in range(tile_start, page_tiles + tile_start):
150 col = (tile_i - page_i * (n_cols * n_rows)) % n_cols
151 row = (tile_i - page_i * (n_cols * n_rows)) // n_cols
153 ax_amp[row, col].set_title(
154 str(obs["baseline_info"]["tile_names"][(page_i + 1) * tile_i]),
155 fontsize=9,
156 )
157 ax_phase[row, col].set_title(
158 str(obs["baseline_info"]["tile_names"][(page_i + 1) * tile_i]),
159 fontsize=9,
160 )
162 for pol_i in range(cal["n_pol"]):
163 if pol_i == 0:
164 colour = "b-"
165 elif pol_i == 1:
166 colour = "r-"
168 # Select current plotting type
169 if type_i == 0:
170 # Gain and phase fit solutions
171 amp = cal_sol_amp[pol_i, :, tile_i].squeeze()
172 phase = cal_sol_phase[pol_i, :, tile_i].squeeze()
173 if type_i == 1:
174 # Gain and phase residuals (per-frequency solutions minus fit solutions)
175 amp = cal_res_amp[pol_i, :, tile_i].squeeze()
176 phase = cal_res_phase[pol_i, :, tile_i].squeeze()
177 if type_i == 2:
178 # Gain and phase per-frequency solutions
179 amp = cal_raw_amp[pol_i, :, tile_i].squeeze()
180 phase = cal_raw_phase[pol_i, :, tile_i].squeeze()
182 # Leave blanks for flagged or undefined tiles
183 if np.isnan([amp, phase]).all() or tile_i not in tile_i_use:
184 ax_amp[row, col].set_axis_off()
185 ax_phase[row, col].set_axis_off()
186 else:
187 ax_amp[row, col].plot(freq_arr_use, amp, colour, lw=1)
188 ax_phase[row, col].plot(freq_arr_use, phase, colour, lw=1)
190 fig_amp.suptitle(pyfhd_config["obs_id"], fontsize=14)
191 fig_phase.suptitle(pyfhd_config["obs_id"], fontsize=14)
193 fig_amp.supylabel("Amplitude")
194 fig_amp.supxlabel("Frequency (MHz)")
196 fig_phase.supylabel("Phase (radians)")
197 fig_phase.supxlabel("Frequency (MHz)")
199 fig_amp.legend(
200 handles=patches,
201 loc="outside upper right",
202 ncol=2,
203 frameon=False,
204 fontsize=12,
205 handlelength=1,
206 )
207 fig_phase.legend(
208 handles=patches,
209 loc="outside upper right",
210 ncol=2,
211 frameon=False,
212 fontsize=12,
213 handlelength=1,
214 )
216 # Page naming conventions
217 if n_pages > 1:
218 page_name = "_page" + str(page_i)
219 else:
220 page_name = ""
222 # Save amplitude and phase plot given the type
223 fig_amp.savefig(
224 f"{save_path_roots[2 * type_i]}{page_name}.png",
225 bbox_inches="tight",
226 dpi=200,
227 )
228 fig_phase.savefig(
229 f"{save_path_roots[2 * type_i + 1]}{page_name}.png",
230 bbox_inches="tight",
231 dpi=200,
232 )
234 plt.close(fig_amp)
235 plt.close(fig_phase)