Coverage for PyFHD/gridding/gridding_utils.py: 76%

1import numpy as np 

2from numpy.typing import NDArray 

3import PyFHD.gridding.filters as filters 

4from PyFHD.pyfhd_tools.pyfhd_utils import rebin, histogram, array_match, meshgrid 

5from scipy.signal import convolve 

6from astropy.convolution import Box2DKernel 

7from math import pi 

8from logging import Logger 

9import h5py 

10 

11 

12def interpolate_kernel( 

13 kernel_arr: NDArray[np.complex128], 

14 x_offset: NDArray[np.integer], 

15 y_offset: NDArray[np.integer], 

16 dx0dy0: NDArray[np.float64], 

17 dx1dy0: NDArray[np.float64], 

18 dx0dy1: NDArray[np.float64], 

19 dx1dy1: NDArray[np.float64], 

20) -> NDArray[np.complex128]: 

21 """ 

22 Perform a bilinear interpolation given the 2D derivatives of the baseline center to the nearest 

23 hyperresolved pixel edge in {u,v} space. This will provide a quadratic estimate at the sample 

24 location to smooth out the dependence on hyperresolved pixel size. 

25 

26 Parameters 

27 ---------- 

28 kernel_arr : NDArray[np.complex128] 

29 The 2D kernel for bilinear interpolation 

30 x_offset : NDArray[np.integer] 

31 The nearest pixel offset in the x (u) direction 

32 y_offset : NDArray[np.integer] 

33 The nearest pixel offset in the y (v) direction 

34 dx0dy0 : NDArray[np.float64] 

35 (1 - derivative to pixel edge in x) * (1 - derivative to pixel edge in y) for selected baselines 

36 dx1dy0 : NDArray[np.float64] 

37 (derivative to pixel edge in x) * (1 - derivative to pixel edge in y) for selected baselines 

38 dx0dy1 : NDArray[np.float64] 

39 (1 - derivative to pixel edge in x) * (derivative to pixel edge in y) for selected baselines 

40 dx1dy1 : NDArray[np.float64] 

41 (derivative to pixel edge in x) * (derivative to pixel edge in y) for selected baselines 

42 

43 Returns 

44 ------- 

45 kernel: NDArray[np.complex128] 

46 The interpolated 2D kernel 

47 """ 

48 # x_offset and y_offset needed to be swapped around as IDL is column-major, while Python is row-major 

49 kernel = kernel_arr[y_offset, x_offset] * dx0dy0 

50 kernel += kernel_arr[y_offset, x_offset + 1] * dx1dy0 

51 kernel += kernel_arr[y_offset + 1, x_offset] * dx0dy1 

52 kernel += kernel_arr[y_offset + 1, x_offset + 1] * dx1dy1 

53 

54 return kernel 

55 

56 

57def conjugate_mirror( 

58 image: NDArray[np.complex128 | np.float64], 

59) -> NDArray[np.complex128 | np.float64]: 

60 """ 

61 Mirror an image about the origin and take the complex conjugate. The origin is considered to be the first 

62 row, and thus is not repeated in the conjugate mirror. 

63 

64 Parameters 

65 ---------- 

66 image: NDArray[np.complex128 | np.float64] 

67 A 2D array of real or complex numbers 

68 

69 Returns 

70 ------- 

71 conj_mirror_image: NDArray[np.complex128 | np.float64] 

72 The 2D conjugate mirror of the input without the origin 

73 """ 

74 # Flip image left to right (i.e. flips columns) & Flip image up to down (i.e. flips rows) 

75 conj_mirror_image = np.flip(image) 

76 # Shifts columns then rows by 1 

77 conj_mirror_image = np.roll(conj_mirror_image, 1, axis=(1, 0)) 

78 # If any of the array is complex, or its a complex array, get the conjugates 

79 if np.iscomplexobj(image): 

80 conj_mirror_image = np.conjugate(conj_mirror_image) 

81 return conj_mirror_image 

82 

83 

84def baseline_grid_locations( 

85 obs: dict, 

86 psf: dict, 

87 params: dict, 

88 vis_weights: NDArray[np.float64], 

89 logger: Logger, 

90 bi_use: NDArray[np.integer] | None = None, 

91 fi_use: NDArray[np.integer] | None = None, 

92 fill_model_visibilities: bool = False, 

93 interp_flag: bool = False, 

94 mask_mirror_indices: bool = False, 

95) -> dict: 

96 """ 

97 Calculate the histogram of baseline grid locations in units of pixels whilst also 

98 returning the minimum pixel number that an unflagged baseline contributes to (depending on the 

99 size of the kernel). Optionally return the 2D derivatives for bilinear interpolation and the 

100 indices of the unflagged baselines/frequencies. 

101 

102 Parameters 

103 ---------- 

104 obs : dict 

105 Observation metadata dictionary 

106 psf : dict 

107 Beam metadata dictionary 

108 params : dict 

109 Visibility metadata dictionary 

110 vis_weights : NDArray[np.float64] 

111 Weights (flags) of the visibilities 

112 logger : Logger 

113 PyFHD's logger 

114 bi_use : NDArray[np.integer] | None, optional 

115 Baseline index array for gridding, i.e even vs odd time stamps, by default None 

116 fi_use : NDArray[np.integer] | None, optional 

117 Frequency index array for gridding, i.e. gridding all frequencies for continuum images, by default None 

118 fill_model_visibilities : bool, optional 

119 Calculate baseline grid locations for all baselines, regardless of flags, by default False 

120 interp_flag : bool, optional 

121 Calculate derivatives for bilinear interpolation of the kernel to pixel locations, by default False 

122 mask_mirror_indices : bool, optional 

123 Exclude baselines mirrored along the v-axis, by default False 

124 

125 Returns 

126 ------- 

127 baselines_dict : dict 

128 Histogram of baseline grid locations, associated derivatives, and minimum contributing pixel location, 

129 arranged in a dictionary 

130 """ 

131 

132 # Set up the return dictionary 

133 baselines_dict = {} 

134 

135 # Retrieve information from the data structures 

136 n_tile = obs["n_tile"] 

137 n_freq = obs["n_freq"] 

138 dimension = obs["dimension"] 

139 elements = obs["elements"] 

140 kbinsize = obs["kpix"] 

141 min_baseline = obs["min_baseline"] 

142 max_baseline = obs["max_baseline"] 

143 b_info = obs["baseline_info"] 

144 psf_dim = psf["dim"] 

145 psf_resolution = psf["resolution"] 

146 if isinstance(psf, h5py.File): 

147 psf_dim = psf_dim[0] 

148 psf_resolution = psf_resolution[0] 

149 

150 # Frequency information of the visibilities 

151 if fill_model_visibilities: 

152 fi_use = np.arange(n_freq) 

153 elif fi_use is None: 

154 fi_use = np.nonzero(b_info["freq_use"])[0] 

155 frequency_array = b_info["freq"] 

156 frequency_array = frequency_array[fi_use] 

157 

158 # Set the bi_use_flag, if we set it here, we need to return it 

159 bi_use_flag = False 

160 # Baselines to use 

161 if bi_use is None: 

162 bi_use_flag = True 

163 # if the data is being gridded separately for the even/odd time samples 

164 # then force flagging to be consistent across even/odd sets 

165 if not fill_model_visibilities: 

166 flag_test = np.sum(np.maximum(vis_weights, 0), axis=0) 

167 bi_use = np.nonzero(flag_test)[0] 

168 else: 

169 tile_use = np.arange(n_tile) + 1 

170 bi_use = array_match( 

171 b_info["tile_a"].astype(int), 

172 tile_use, 

173 array_2=b_info["tile_b"].astype(int), 

174 ) 

175 

176 # Rather than calculating the flat indexes we want, lets just index the array 

177 # by the frequency use and baseline_use indexes 

178 rows, cols = np.meshgrid(fi_use, bi_use) 

179 vis_weights_use = vis_weights[rows, cols].T 

180 

181 # Units in pixel/Hz 

182 kx_arr = params["uu"][bi_use] / kbinsize 

183 ky_arr = params["vv"][bi_use] / kbinsize 

184 

185 if not fill_model_visibilities: 

186 # Flag baselines on their maximum and minimum extent in the full frequency range of the observation 

187 # This prevents the sudden dissapearance of baselines along frequency 

188 dist_test = np.sqrt(kx_arr**2 + ky_arr**2) * kbinsize 

189 dist_test_max = np.max(obs["baseline_info"]["freq"]) * dist_test 

190 dist_test_min = np.min(obs["baseline_info"]["freq"]) * dist_test 

191 flag_dist_baseline = np.where( 

192 (dist_test_min < min_baseline) | (dist_test_max > max_baseline) 

193 ) 

194 del (dist_test, dist_test_max, dist_test_min) 

195 

196 # Create the other half of the uv plane via negating the locations 

197 conj_i = np.where(ky_arr > 0)[0] 

198 if conj_i.size > 0: 198 ↛ 203line 198 didn't jump to line 203 because the condition on line 198 was always true

199 kx_arr[conj_i] = -kx_arr[conj_i] 

200 ky_arr[conj_i] = -ky_arr[conj_i] 

201 

202 # Center of baselines for x and y in units of pixels 

203 xcen = np.outer(frequency_array, kx_arr) 

204 ycen = np.outer(frequency_array, ky_arr) 

205 

206 # Pixel number offset per baseline for each uv-box subset 

207 x_offset = np.fix( 

208 np.floor((xcen - np.floor(xcen)) * psf_resolution) % psf_resolution 

209 ).astype(np.int64) 

210 y_offset = np.fix( 

211 np.floor((ycen - np.floor(ycen)) * psf_resolution) % psf_resolution 

212 ).astype(np.int64) 

213 

214 if interp_flag: 

215 # Derivatives from pixel edge to baseline center for use in interpolation 

216 dx_arr = (xcen - np.floor(xcen)) * psf_resolution - np.floor( 

217 (xcen - np.floor(xcen)) * psf_resolution 

218 ) 

219 dy_arr = (ycen - np.floor(ycen)) * psf_resolution - np.floor( 

220 (ycen - np.floor(ycen)) * psf_resolution 

221 ) 

222 baselines_dict["dx0dy0_arr"] = (1 - dx_arr) * (1 - dy_arr) 

223 baselines_dict["dx0dy1_arr"] = (1 - dx_arr) * dy_arr 

224 baselines_dict["dx1dy0_arr"] = dx_arr * (1 - dy_arr) 

225 baselines_dict["dx1dy1_arr"] = dx_arr * dy_arr 

226 

227 # The minimum pixel in the uv-grid (bottom left of the kernel) that each baseline contributes to 

228 xmin = (np.floor(xcen) + elements / 2 - (psf_dim / 2 - 1)).astype(np.int64) 

229 ymin = (np.floor(ycen) + dimension / 2 - (psf_dim / 2 - 1)).astype(np.int64) 

230 

231 # Set the minimum pixel value of baselines which fall outside of the uv-grid tom -1 to exclude them 

232 range_test_x_i = np.where((xmin <= 0) | (xmin + psf_dim - 1 >= elements - 1)) 

233 range_test_y_i = np.where((ymin <= 0) | (ymin + psf_dim - 1 >= dimension - 1)) 

234 if range_test_x_i[0].size > 0: 234 ↛ 237line 234 didn't jump to line 237 because the condition on line 234 was always true

235 xmin[range_test_x_i] = -1 

236 ymin[range_test_x_i] = -1 

237 if range_test_y_i[0].size > 0: 

238 xmin[range_test_y_i] = -1 

239 ymin[range_test_y_i] = -1 

240 

241 # Flag baselines which fall outside the uv plane 

242 if not fill_model_visibilities: 

243 if np.size(flag_dist_baseline) > 0: 243 ↛ 251line 243 didn't jump to line 251 because the condition on line 243 was always true

244 # If baselines fall outside the desired min/max baseline range at all during the frequency range 

245 # then set their maximum pixel value to -1 to exclude them 

246 xmin[:, flag_dist_baseline] = -1 

247 ymin[:, flag_dist_baseline] = -1 

248 del flag_dist_baseline 

249 

250 # Normally we check vis_weight_switch, but its always true here so... do this 

251 flag_i = np.where(vis_weights_use <= 0) 

252 if fill_model_visibilities: 

253 n_flag = 0 

254 else: 

255 n_flag = np.size(flag_i) 

256 if n_flag > 0: 

257 xmin[flag_i] = -1 

258 ymin[flag_i] = -1 

259 

260 if mask_mirror_indices: 

261 # Option to exclude v-axis mirrored baselines 

262 if conj_i.size > 0: 262 ↛ 267line 262 didn't jump to line 267 because the condition on line 262 was always true

263 xmin[:, conj_i] = -1 

264 ymin[:, conj_i] = -1 

265 

266 # If xmin or ymin is invalid then adjust the baselines dict as necessary 

267 if xmin.size == 0 or ymin.size == 0 or np.max([np.max(xmin), np.max(ymin)]) < 0: 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true

268 logger.warning("All data flagged or cut!") 

269 baselines_dict["bin_n"] = 0 

270 baselines_dict["n_bin_use"] = 0 

271 baselines_dict["bin_i"] = -1 

272 baselines_dict["ri"] = 0 

273 else: 

274 # Match all visibilities that map from and to exactly the same pixels and store them as a histogram in bin_n 

275 # with their respective index ri. Setting min equal to 0, excludes flagged data (data set to -1). 

276 for_hist = xmin + ymin * dimension 

277 bin_n, _, ri = histogram(for_hist, min=0) 

278 bin_i = np.nonzero(bin_n)[0] 

279 

280 # Update the baselines_dict which gets returned 

281 baselines_dict["bin_n"] = bin_n 

282 baselines_dict["bin_i"] = bin_i 

283 baselines_dict["n_bin_use"] = bin_i.size 

284 baselines_dict["ri"] = ri 

285 # Add values to baselines dict 

286 baselines_dict["xmin"] = xmin 

287 baselines_dict["ymin"] = ymin 

288 baselines_dict["x_offset"] = x_offset 

289 baselines_dict["y_offset"] = y_offset 

290 if bi_use_flag: 

291 baselines_dict["bi_use"] = bi_use 

292 

293 # Return the dictionary 

294 return baselines_dict 

295 

296 

297def dirty_image_generate( 

298 dirty_image_uv: NDArray[np.complex128], 

299 pyfhd_config: dict, 

300 logger: Logger, 

301 uniform_filter_uv: NDArray[np.float64] | None = None, 

302 mask: NDArray[np.integer] | None = None, 

303 baseline_threshold: int | float = 0, 

304 normalization: float | NDArray[np.float64] | None = None, 

305 resize: int | None = None, 

306 width_smooth: int | float | None = None, 

307 degpix: float | None = None, 

308 not_real: bool = False, 

309 pad_uv_image: int | float | None = None, 

310 weights: NDArray[np.float64] | None = None, 

311 filter: NDArray[np.float64] | None = None, 

312 beam_ptr: NDArray[np.complex128] | None = None, 

313) -> tuple[ 

314 NDArray[np.complex128], NDArray[np.float64], float | NDArray[np.float64] | None 

315]: 

316 """ 

317 Generate a projected image from an input {u,v} plane through a 2D FFT. Optionally apply padding, masking, 

318 filtering, and more. 

319 

320 Parameters 

321 ---------- 

322 dirty_image_uv : NDArray[np.complex128] 

323 A 2D {u,v} plane which generally includes the beam via a gridding kernel 

324 pyfhd_config : dict 

325 PyFHD's configuration dictionary containing all the options set for a PyFHD run 

326 logger : Logger 

327 PyFHD's logger 

328 uniform_filter_uv : NDArray[np.float64] | None, optional 

329 A 2D {u,v} gridded visibility number, by default None 

330 mask : NDArray[np.integer] | None, optional 

331 A 2D {u,v} mask to apply before image creation, by default None 

332 baseline_threshold : int | float, optional 

333 The maximum baseline length to include in units of pixels, by default 0 

334 normalization : float | NDArray[np.float64] | None, optional 

335 A value by which to normalize the image by, by default None 

336 resize : int | None, optional 

337 Increase the number of pixels by a factor and rebin the input {u,v} plane, by default None 

338 width_smooth : int | float | None, optional 

339 Smooth out the harsh baseline threshold by the given number of pixels, by default None 

340 degpix : float | None, optional 

341 Degrees per pixel, by default None 

342 not_real : bool, optional 

343 Flag to return a complex image, by default False 

344 pad_uv_image : int | float | None, optional 

345 Pad the {u,v} plane by this pixel amount to increase perceived resolution of the image, by default None 

346 weights : NDArray[np.float64] | None, optional 

347 Gridded {u,v} plane of visibility weights, necessary in some filtering schemes, by default None 

348 filter : NDArray[np.float64] | None, optional 

349 Image filter to apply, by default None 

350 beam_ptr : NDArray[np.complex128] | None, optional 

351 Weight by an additional factor of the beam for optimal weighting, by default None 

352 

353 Returns 

354 ------- 

355 dirty_image : NDArray[np.complex128] 

356 A 2D {l,m} directional-cosine image plane 

357 filter : NDArray[np.float64] 

358 The filter applied to the dirty image 

359 normalization : float | NDArray[np.float64] | None 

360 The normalization (if any) applied to the dirty image 

361 """ 

362 

363 # dimension is columns, elements is rows 

364 elements, dimension = dirty_image_uv.shape 

365 di_uv_use = dirty_image_uv 

366 # If the baseline threshold has been set 

367 if baseline_threshold is not None: 367 ↛ 400line 367 didn't jump to line 400 because the condition on line 367 was always true

368 if width_smooth is None: 

369 width_smooth = np.floor(np.sqrt(dimension * elements) / 100) 

370 rarray = np.sqrt( 

371 (meshgrid(dimension, 1) - dimension / 2) ** 2 

372 + (meshgrid(elements, 2) - elements / 2) ** 2 

373 ) 

374 # Get all the values that meet the threshold 

375 if baseline_threshold >= 0: 375 ↛ 378line 375 didn't jump to line 378 because the condition on line 375 was always true

376 cut_i = np.where(rarray.flatten() < baseline_threshold) 

377 else: 

378 cut_i = np.where(rarray.flatten() > np.abs(baseline_threshold)) 

379 # Create the mask array of ones 

380 mask_bt = np.ones((elements, dimension)) 

381 # If there are values from cut, then use all those here and replace with 0 

382 if np.size(cut_i) > 0: 

383 mask_bt_flatiter = mask_bt.flat 

384 mask_bt_flatiter[cut_i] = 0 

385 # Get the kernel width 

386 kernel_width = np.max([width_smooth, 1]) 

387 # In IDL if the kernel width is even one is added to make it odd 

388 if kernel_width % 2 == 0: 

389 kernel_width += 1 

390 # Use a box width averaging filter over the mask, use valid so we can insert it in later 

391 box_averages = convolve(mask_bt, Box2DKernel(kernel_width), mode="valid") 

392 # Since IDL SMOOTH edges by default are the edges of the array used, ignore edges (its the reason why we used a valid convolve) 

393 start = int(kernel_width // 2) 

394 end = int(mask_bt.shape[1] - (kernel_width // 2)) 

395 mask_bt[start:end, start:end] = box_averages 

396 # Apply boxed mask to the dirty image 

397 di_uv_use *= mask_bt 

398 

399 # If a mask was supplied use that too 

400 if mask is not None: 

401 di_uv_use *= mask 

402 

403 # If a filter was supplied as a numpy array (we can adjust this to support different formats) 

404 if filter is not None: 

405 if isinstance(filter, np.ndarray): 405 ↛ 443line 405 didn't jump to line 443 because the condition on line 405 was always true

406 # If the filter is already the right size, use it 

407 if np.size(filter) == np.size(di_uv_use): 

408 di_uv_use *= filter 

409 # Otherwise use a filter function 

410 else: 

411 if pyfhd_config["image_filter"] == "filter_uv_uniform": 411 ↛ 413line 411 didn't jump to line 413 because the condition on line 411 was always true

412 logger.info("Using filter_uv_uniform for dirty_image_generate") 

413 elif pyfhd_config["image_filter"] == "filter_uv_hannning": 

414 logger.warning( 

415 "filter_uv_hanning hasn't been translated yet using filter_uv_uniform for dirty_image_generate instead" 

416 ) 

417 

418 elif pyfhd_config["image_filter"] == "filter_uv_natural": 

419 logger.warning( 

420 "filter_uv_natural hasn't been translated yet using filter_uv_uniform for dirty_image_generate instead" 

421 ) 

422 

423 elif pyfhd_config["image_filter"] == "filter_uv_radial": 

424 logger.warning( 

425 "filter_uv_radial hasn't been translated yet using filter_uv_uniform for dirty_image_generate instead" 

426 ) 

427 

428 elif pyfhd_config["image_filter"] == "filter_uv_tapered_uniform": 

429 logger.warning( 

430 "filter_uv_tapered_uniform hasn't been translated yet using filter_uv_uniform for dirty_image_generate instead" 

431 ) 

432 

433 elif pyfhd_config["image_filter"] == "filter_uv_optimal": 

434 logger.warning( 

435 "filter_uv_optimal hasn't been translated yet using filter_uv_uniform for dirty_image_generate instead" 

436 ) 

437 # Since we only use filter_uniform at the moment, put the call to it here. 

438 di_uv_use, filter = filters.filter_uv_uniform( 

439 di_uv_use, vis_count=uniform_filter_uv, weights=weights 

440 ) 

441 

442 # Resize the dirty image by the factor resize 

443 if resize is not None: 

444 dimension *= resize 

445 elements *= resize 

446 # Ensure elements and dimension are integers as resize should be an int 

447 # but if it's not this makes sure dimension and elements are ints afterwards 

448 dimension = int(dimension) 

449 elements = int(elements) 

450 di_uv_real = di_uv_use.real 

451 di_uv_img = di_uv_use.imag 

452 # Use rebin to resize, apply to real and complex separately 

453 di_uv_real = rebin(di_uv_real, (elements, dimension)) 

454 di_uv_img = rebin(di_uv_img, (elements, dimension)) 

455 # Combine real and complex back together 

456 di_uv_use = di_uv_real + di_uv_img * 1j 

457 

458 # Apply padding if it was supplied 

459 if pad_uv_image is not None: 

460 # dimension_new = int(np.max([np.max([dimension, elements]) * pad_uv_image, np.max([dimension, elements])])) 

461 # di_uv1 = np.zeros((dimension_new, dimension_new), dtype = "complex") 

462 # di_uv1[dimension_new // 2 - elements // 2 : dimension_new // 2 + elements // 2, 

463 # dimension_new // 2 - dimension // 2 : dimension_new // 2 + dimension // 2] = di_uv_use 

464 di_uv1 = np.pad( 

465 di_uv_use, (np.max([dimension, elements]) // 2) * (int(pad_uv_image) - 1) 

466 ) 

467 di_uv_use = di_uv1 * (pad_uv_image**2) 

468 

469 # FFT normalization 

470 if degpix is not None: 

471 di_uv_use /= np.radians(degpix) ** 2 

472 

473 # Multivariate Fast Fourier Transform 

474 dirty_image = np.fft.fftshift( 

475 np.fft.fftn(np.fft.fftshift(di_uv_use), norm="forward") 

476 ) 

477 if not not_real: 

478 dirty_image = dirty_image.real 

479 

480 # filter_uv_optimal produces images that are weighted by one factor of the beam 

481 # Weight by an additional factor of the beam to align with FHD's convention 

482 if pyfhd_config["image_filter"] == "filter_uv_optimal" and beam_ptr is not None: 482 ↛ 483line 482 didn't jump to line 483 because the condition on line 482 was never true

483 dirty_image *= beam_ptr 

484 

485 # If we are returning complex, make sure its complex 

486 if not_real: 

487 dirty_image = dirty_image.astype(np.complex128) 

488 else: 

489 dirty_image = dirty_image.real 

490 # Normalize by the matrix given, if it was given 

491 if normalization is not None: 

492 dirty_image *= normalization 

493 

494 # Return 

495 return dirty_image, filter, normalization 

496 

497 

498def grid_beam_per_baseline( 

499 psf: dict, 

500 pyfhd_config: dict, 

501 logger: Logger, 

502 uu: NDArray[np.float64], 

503 vv: NDArray[np.float64], 

504 ww: NDArray[np.float64], 

505 l_mode: NDArray[np.float64], 

506 m_mode: NDArray[np.float64], 

507 n_tracked: NDArray[np.float64], 

508 frequency_array: NDArray[np.float64], 

509 x: NDArray[np.float64], 

510 y: NDArray[np.float64], 

511 xmin_use: int, 

512 ymin_use: int, 

513 freq_i: NDArray[np.integer], 

514 bt_index: NDArray[np.integer], 

515 polarization: int, 

516 image_bot: int, 

517 image_top: int, 

518 psf_dim3: int, 

519 box_matrix: NDArray[np.complex128], 

520 vis_n: int, 

521 beam_int: NDArray[np.complex128] | None = None, 

522 beam2_int: NDArray[np.complex128] | None = None, 

523 n_grp_use: NDArray[np.integer] | None = None, 

524 degrid_flag: bool = False, 

525): 

526 """ 

527 Calculate the contribution of each baseline to the static {u,v} grid using the corrective phases in 

528 image-space, as detailed in J. Line's thesis "PUMA and MAJICK: cross-matching and imaging techniques 

529 for a detection of the epoch of reionisation" 

530 

531 Parameters 

532 ---------- 

533 psf : dict 

534 Beam metadata dictionary 

535 pyfhd_config : dict 

536 PyFHD's configuration dictionary containing all the options set for a PyFHD run 

537 logger : Logger 

538 PyFHD's logger 

539 uu : NDArray[np.float64] 

540 1D array of the u-coordinate of selected baselines in light travel time 

541 vv : NDArray[np.float64] 

542 1D array of the v-coordinate of selected baselines in light travel time 

543 ww : NDArray[np.float64] 

544 1D array of the w-coordinate of selected baselines in light travel time 

545 l_mode : NDArray[np.float64] 

546 Directional-cosine l of pixel centers of the hyperresolved beam 

547 m_mode : NDArray[np.float64] 

548 Directional-cosine m of pixel centers of the hyperresolved beam 

549 n_tracked : NDArray[np.float64] 

550 Directional-cosine n of pixel centers of the phase-tracked hyperresolved beam 

551 frequency_array : NDArray[np.float64] 

552 Array of selected frequencies in Hz 

553 x : NDArray[np.float64] 

554 1D array of gridding extent and resolution in the x-direction in wavelengths 

555 y : NDArray[np.float64] 

556 1D array of gridding extent and resolution in the y-direction in wavelengths 

557 xmin_use : int 

558 The minimum x-pixel that each selected baseline contributes to 

559 ymin_use : int 

560 The minimum y-pixel that each selected baseline contributes to 

561 freq_i : NDArray[np.integer] 

562 The current frequency index 

563 bt_index : NDArray[np.integer] 

564 The current baseline/time index 

565 polarization : int 

566 The current polarization index 

567 image_bot : int 

568 The bottom-most pixel index that the image-space hyperresolved beam contributes to 

569 image_top : int 

570 The top-most pixel index that the image-space hyperresolved beam contributes to 

571 psf_dim3 : int 

572 The pixel area of the psf footprint on the {u,v} grid 

573 box_matrix : NDArray[np.complex128] 

574 A 2D array of the number of visibilities to grid and the area of each visibility on the static {u,v} grid 

575 vis_n : int 

576 The number of visibilities to grid 

577 beam_int : NDArray[np.complex128] | None, optional 

578 The integral of the beam sensitivity in {u,v} space, by default None 

579 beam2_int : NDArray[np.complex128] | None, optional 

580 The integral of the squared beam sensitivity in {u,v} space, by default None 

581 n_grp_use : NDArray[np.integer] | None, optional 

582 The number of baselines in the current grouping, by default None 

583 degrid_flag : bool, optional 

584 Perform degridding instead of gridding, by default False 

585 

586 Returns 

587 ------- 

588 box_matrix : NDArray[np.complex128] 

589 The kernel values on the static {u,v} grid for each visibility 

590 """ 

591 

592 psf_dim = psf["dim"] 

593 psf_resolution = psf["resolution"] 

594 if isinstance(psf, h5py.File): 594 ↛ 595line 594 didn't jump to line 595 because the condition on line 594 was never true

595 psf_dim = psf_dim[0] 

596 psf_resolution = psf_resolution[0] 

597 # Loop over all visibilities that fall within the chosen visibility box 

598 for ii in range(vis_n): 

599 # Pixel center offset phases 

600 deltau_l = l_mode * ( 

601 uu[bt_index[ii]] * frequency_array[freq_i[ii]] - x[xmin_use + psf_dim // 2] 

602 ) 

603 deltav_m = m_mode * ( 

604 vv[bt_index[ii]] * frequency_array[freq_i[ii]] - y[ymin_use + psf_dim // 2] 

605 ) 

606 # w term offset phase 

607 w_n_tracked = n_tracked * ww[bt_index[ii]] * frequency_array[freq_i[ii]] 

608 

609 # Generate a UV beam from the image space beam, offset by calculated phases 

610 psf_base_superres, _, _ = dirty_image_generate( 

611 psf["image_info"]["image_power_beam_arr"][polarization] 

612 * np.exp(2 * pi * (0 + 1j) * (-w_n_tracked + deltau_l + deltav_m)), 

613 pyfhd_config, 

614 logger, 

615 not_real=True, 

616 ) 

617 psf_base_superres = psf_base_superres[ 

618 image_bot : image_top + 1, image_bot : image_top + 1 

619 ] 

620 

621 # A quick way to sum down the image by a factor of 2 in both dimensions. 

622 # A 4x4 example where we sum down by a factor of 2 

623 # 

624 # 1 2 3 4 1 2 1 2 1 2 5 6 14 46 14 22 

625 # 5 6 7 8 --> 3 4 --> 5 6 --> 9 10 13 14 --> 22 54 --> 46 54 

626 # 9 10 11 12 9 10 

627 # 13 14 15 16 5 6 13 14 3 4 7 8 

628 # 7 8 11 12 15 16 

629 # 3 4 

630 # 9 10 7 8 

631 # 11 12 11 12 

632 # 15 16 

633 # 13 14 

634 # 15 16 

635 d = psf_base_superres.shape 

636 # Note columns and rows are swapped from IDL so nx is now rows! 

637 nx = d[0] // psf_resolution 

638 ny = d[1] // psf_resolution 

639 # The same result of IDL in numpy is np.reshape, with shape swapping rows and columns, then doing transpose of this shape 

640 psf_base_superres = np.reshape( 

641 psf_base_superres, [psf_resolution * ny, nx, psf_resolution] 

642 ) 

643 psf_base_superres = np.transpose(psf_base_superres, [1, 0, 2]) 

644 psf_base_superres = np.reshape(psf_base_superres, [ny, nx, psf_resolution**2]) 

645 psf_base_superres = np.sum(psf_base_superres, -1) 

646 psf_base_superres = np.transpose(psf_base_superres) 

647 psf_base_superres = np.reshape(psf_base_superres, psf_dim**2) 

648 start = psf_dim3 * ii 

649 end = start + psf_base_superres.size 

650 box_matrix_iter = box_matrix.flat 

651 box_matrix_iter[start:end] = psf_base_superres 

652 

653 # Subtract off a small clip, set negative indices to 0, and renomalize. 

654 # This is a modification of the look-up-table beam using a few assumptions 

655 # to make it faster/feasible to run. 

656 # Modifications: done per group of baselines that fit within the current box, 

657 # rather than individually. region_grow is not used to find a contiguous 

658 # edge around the beam to cut because it is too slow. 

659 if pyfhd_config["beam_clip_floor"]: 659 ↛ 674line 659 didn't jump to line 674 because the condition on line 659 was always true

660 psf_val_ref = np.sum(box_matrix, 1) 

661 psf_amp = np.abs(box_matrix) 

662 psf_mask_threshold_use = np.max(psf_amp) / psf["beam_mask_threshold"] 

663 psf_amp -= psf_mask_threshold_use 

664 psf_phase = np.arctan2(box_matrix.imag, box_matrix.real) 

665 psf_amp = np.maximum(psf_amp, np.zeros_like(psf_amp)) 

666 box_matrix = psf_amp * np.cos(psf_phase) + (0 + 1j) * psf_amp * np.sin( 

667 psf_phase 

668 ) 

669 ref_temp = np.sum(box_matrix, -1) 

670 box_matrix[:vis_n, :] *= np.reshape( 

671 psf_val_ref / ref_temp, (psf_val_ref.size, 1) 

672 ) 

673 

674 if ( 674 ↛ 681line 674 didn't jump to line 681 because the condition on line 674 was never true

675 degrid_flag 

676 and beam_int is not None 

677 and beam2_int is not None 

678 and n_grp_use is not None 

679 ): 

680 # Calculate the beam and beam^2 integral (degridding) 

681 beam_int_temp = np.sum(box_matrix, 0) / psf_resolution**2 

682 beam2_int_temp = np.sum(np.abs(box_matrix) ** 2, 0) / psf_resolution**2 

683 for ii in range(np.size(freq_i)): 

684 beam_int[freq_i[ii]] += beam_int_temp[ii] 

685 beam2_int[freq_i[ii]] += beam2_int_temp[ii] 

686 n_grp_use[freq_i[ii]] += 1 

687 

688 return box_matrix 

689 

690 

691def visibility_count( 

692 obs: dict, 

693 psf: dict, 

694 params: dict, 

695 vis_weights: NDArray[np.float64], 

696 logger: Logger, 

697 fi_use: NDArray[np.integer] | None = None, 

698 bi_use: NDArray[np.integer] | None = None, 

699 mask_mirror_indices: bool = False, 

700 no_conjugate: bool = False, 

701 fill_model_visibilities: bool = False, 

702) -> NDArray[np.float64]: 

703 """ 

704 Calculate the number of contributing visibilities per pixel on the static {u,v} grid 

705 

706 Parameters 

707 ---------- 

708 obs : dict 

709 Observation metadata dictionary 

710 params : dict 

711 Visibility metadata dictionary 

712 vis_weights : NDArray[np.float64] 

713 Weights (flags) of the visibilities 

714 pyfhd_config : dict 

715 PyFHD's configuration dictionary containing all the options set for a PyFHD run 

716 logger : Logger 

717 PyFHD's logger 

718 fi_use : NDArray[np.integer] | None, optional 

719 Frequency index array for gridding, i.e. gridding all frequencies for continuum images, by default None 

720 bi_use : NDArray[np.integer] | None, optional 

721 Baseline index array for gridding, i.e even vs odd time stamps, by default None 

722 mask_mirror_indices : bool, optional 

723 Exclude baselines mirrored along the v-axis, by default False 

724 no_conjugate : bool, optional 

725 Do not perform the conjugate mirror of the {u,v} plane, by default False 

726 fill_model_visibilities : bool, optional 

727 Calculate baseline grid locations for all baselines, regardless of flags, by default False 

728 

729 Returns 

730 ------- 

731 uniform_filter : NDArray[np.float64] 

732 2D array of number of contributing visibilities per pixel on the {u,v} grid 

733 """ 

734 

735 # Retrieve info from the data structures 

736 dimension = int(obs["dimension"]) 

737 elements = int(obs["elements"]) 

738 psf_dim = psf["dim"] 

739 if isinstance(psf, h5py.File): 739 ↛ 740line 739 didn't jump to line 740 because the condition on line 739 was never true

740 psf_dim = psf_dim[0] 

741 

742 baselines_dict = baseline_grid_locations( 

743 obs, 

744 psf, 

745 params, 

746 vis_weights, 

747 logger, 

748 bi_use=bi_use, 

749 fi_use=fi_use, 

750 mask_mirror_indices=mask_mirror_indices, 

751 fill_model_visibilities=fill_model_visibilities, 

752 ) 

753 # Retrieve the data we need from baselines_dict 

754 bin_n = baselines_dict["bin_n"] 

755 bin_i = baselines_dict["bin_i"] 

756 xmin = baselines_dict["xmin"] 

757 ymin = baselines_dict["ymin"] 

758 ri = baselines_dict["ri"] 

759 n_bin_use = baselines_dict["n_bin_use"] 

760 # Remove baselines_dict 

761 del baselines_dict 

762 

763 uniform_filter = np.zeros((elements, dimension)) 

764 # Get the flat iterations of the xmin and ymin arrays. 

765 xmin_iter = xmin.flat 

766 ymin_iter = ymin.flat 

767 for bi in range(n_bin_use): 

768 idx = ri[ri[bin_i[bi]]] 

769 # Should all be the same, but don't want an array 

770 xmin_use = xmin_iter[idx] 

771 ymin_use = ymin_iter[idx] 

772 # Please note that due to precision differences, the result will be different compared to IDL FHD 

773 uniform_filter[ 

774 ymin_use : ymin_use + psf_dim, xmin_use : xmin_use + psf_dim 

775 ] += bin_n[bin_i[bi]] 

776 

777 if not no_conjugate: 

778 uniform_filter = (uniform_filter + conjugate_mirror(uniform_filter)) / 2 

779 

780 return uniform_filter 

781 

782 

783def holo_mapfn_convert(map_fn, psf_dim, dimension, elements=None, norm=1, threshold=0): 

784 """ 

785 Convert pointer array holographic map function to a sparse matrix. May need to be deprecated. 

786 The mapping functions were not translated into `visibility_grid` at the time of translation as 

787 it wasn't clear what PyFHD was going to use for sparse/large arrays at the time. If you wish to 

788 implement the mapping functions, I suggest using HDF5 chunk loading for the mapping function. 

789 

790 Parameters 

791 ---------- 

792 map_fn: np.ndarray 

793 Pointer array holographic map function 

794 psf_dim: int, float 

795 The number of pixels in one direction of the psf {u,v} footprint 

796 dimension: int 

797 The number of pixels in the u-direction 

798 elements: None, optional 

799 The number of pixels in the v-direction 

800 norm: int 

801 The normalization of the holographic mapping function (i.e. the number of visibilities) 

802 threshold: int, float, optional 

803 Include values from the holographic map function which are above a provided threshold 

804 

805 Returns 

806 ------- 

807 map_fn: np.recarray 

808 Sparse array holographic map function 

809 """ 

810 # Set up all the necessary arrays and numbers 

811 if elements is None: 

812 elements = dimension 

813 psf_dim2 = 2 * psf_dim 

814 psf_n = psf_dim**2 

815 sub_xv = meshgrid(psf_dim2, 1) - psf_dim 

816 sub_yv = meshgrid(psf_dim2, 2) - psf_dim 

817 # Generate an array of shape elements x dimension 

818 n_arr = np.zeros((elements, dimension)) 

819 # Replace the values of n_arr with the number from each array in mapfn that exceeds the threshold 

820 for xi in range(dimension - 1): 

821 for yi in range(elements - 1): 

822 temp_arr = map_fn[xi, yi] 

823 n1 = np.size(np.where(abs(temp_arr) > threshold)) 

824 n_arr[xi, yi] = n1 

825 # Get the ones to use 

826 i_use = np.where(n_arr) 

827 # Get the amount we're using 

828 i_use_size = np.size(i_use) 

829 # If we aren't using any then return 0 

830 if i_use_size == 0: 

831 return 0 

832 

833 # Get the reverse indices 

834 _, _, ri = histogram(i_use, min=0) 

835 # Create zeros of the same size as what we're using 

836 sa = np.zeros(i_use_size) 

837 ija = np.zeros(i_use_size) 

838 # Fill in the sa and ija arrays 

839 for index in range(i_use_size): 

840 i = i_use[index] 

841 xi = i % dimension 

842 yi = np.floor(i / dimension) 

843 map_fn_sub = map_fn[xi, yi] 

844 j_use = np.where(np.abs(map_fn_sub) > threshold) 

845 

846 xii_arr = sub_xv[j_use] + xi 

847 yii_arr = sub_yv[j_use] + yi 

848 sa[index] = map_fn_sub[j_use] 

849 ija[index] = ri[ri[xii_arr + dimension * yii_arr]]