Coverage for PyFHD/healpix/healpix_utils.py: 58%
247 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-07-01 10:58 +0800
1import sys
2from logging import Logger
3from pathlib import Path
4import h5py
5import importlib_resources
6import numpy as np
7from numpy.typing import NDArray
8from healpy import query_disc
9from healpy.pixelfunc import ang2vec, pix2vec, vec2ang
10from PyFHD.beam_setup.beam_utils import beam_image
11from PyFHD.gridding.visibility_grid import visibility_grid
12from PyFHD.gridding.gridding_utils import dirty_image_generate
13from PyFHD.io.pyfhd_io import load, save
14from PyFHD.pyfhd_tools.pyfhd_utils import (
15 angle_difference,
16 histogram,
17 meshgrid,
18 region_grow,
19)
20from PyFHD.pyfhd_tools.unit_conv import radec_to_altaz, radec_to_pixel
21from astropy.coordinates import EarthLocation
24def healpix_cnv_apply(
25 image: NDArray[np.integer | np.floating | np.complexfloating], hpx_cnv: dict
26) -> NDArray[np.float64]:
27 """
28 healpix_cnv_apply creates a map based off the array/image and healpix convention dictionary given.
29 In FHD the healpix_cnv_apply was mainly used as a wrapper for sprsax2, as such I will put the code
30 for sprsax2 in here as PyFHD will only use sprsax2 here as we don't have the holographic mapping
31 function at this time.
33 Vectors 'sa' (interpolation) and 'ija' (index) follow the row-index sparse storage mode as described
34 in section 2.7 of Numerical Recipes in C, 2nd edition.
36 Parameters
37 ----------
38 image : NDArray[np.integer | np.floating | np.complexfloating]
39 An orthoslant image array that is to be converted to a HEALPix map.
40 hpx_cnv : dict
41 The HEALPix index/interpolation dictionary
43 Returns
44 -------
45 hpx_map: NDArray[np.float64]
46 A flattened 1D HEALPix array with the interpolated values from the image.
47 """
48 # Is B in sprsax2
49 hpx_map = np.zeros(np.size(hpx_cnv["inds"]))
50 # Is X in sprsax2 or X_use
51 image_vector = image.flatten()
52 for i in range(0, np.size(hpx_cnv["i_use"])):
53 i_use = hpx_cnv["i_use"][i]
54 # Transpose is always True when used inside healpix_cnv_apply when using
55 # sprsax2, so we can ignore the transpose keyword
56 hpx_map[hpx_cnv["ija"][i]] += hpx_cnv["sa"][i] * image_vector[i_use]
57 return hpx_map
60def healpix_cnv_generate(
61 obs: dict,
62 mask: NDArray[np.int64] | None,
63 hpx_radius: float,
64 pyfhd_config: dict,
65 logger: Logger,
66 nside: float = None,
67) -> dict:
68 """
69 Generate the HEALPix index/interpolation dictionary that is used to convert
70 between orthoslant image pixelization and the RING index HEALPix pixelization
71 scheme. Binlinear interpolation is performed to HEALPix pixels centers. All
72 angles are in degrees.
74 Parameters
75 ----------
76 obs : dict
77 Observation metadata dictionary
78 mask : np.ndarray
79 If HEALPix indices are not provided in the config file, this boolean
80 mask map of the orthoslant image can be used to mask out areas from the
81 HEALPix conversion.
82 hpx_radius : float
83 If HEALPix indices are not provided in the config file, this selects
84 a radius in degrees, centred on the RA/Dec of the observation, to be
85 converted to HEALPix pixels.
86 pyfhd_config : dict
87 PyFHD configuration settings
88 logger : Logger
89 PyFHD's Logger
90 nside : float
91 The HEALPix nside parameter, calculated based on the observation
92 metadata if none provided. Defaults to None.
94 Returns
95 -------
96 dict
97 A dictionary containing the orthoslant to HEALPix pixel mapping and
98 interpolation. Vectors 'sa' (bilinear interpolation calculated from
99 fractional offsets) and 'ija' (flattened index) follow the row-index
100 sparse storage mode as described in section 2.7 of Numerical Recipes
101 in C, 2nd edition.
102 """
104 # Have ignored the code that relates to hpx_radius being empty, in PyFHD's case
105 # we will assume that a value for hpx_radius is supplied, if you want to add it
106 # yourself you can do that here
107 hpx_inds = None
108 # Fill hpx_inds and nside with values from a file if restrict_healpix_inds has been activated
109 if pyfhd_config["restrict_healpix_inds"]: 109 ↛ 171line 109 didn't jump to line 171 because the condition on line 109 was always true
110 if pyfhd_config["healpix_inds"] is None: 110 ↛ 166line 110 didn't jump to line 166 because the condition on line 110 was always true
111 # Get the healpix indexes based off the observation, comes from observation_healpix_inds_select
112 files = np.array(
113 [
114 {
115 "name": "EoR0_high_healpix_inds.h5",
116 "ra": 0,
117 "dec": -30,
118 "freq": 182,
119 },
120 {
121 "name": "EoR0_low_healpix_inds.h5",
122 "ra": 0,
123 "dec": -30,
124 "freq": 151,
125 },
126 {
127 "name": "EoR1_high_healpix_inds.h5",
128 "ra": 60,
129 "dec": -30,
130 "freq": 182,
131 },
132 {
133 "name": "EoR1_low_healpix_inds.h5",
134 "ra": 60,
135 "dec": -30,
136 "freq": 151,
137 },
138 ]
139 )
140 ang_dist = []
141 for file in files:
142 ang_dist.append(
143 np.abs(
144 angle_difference(
145 obs["obsra"],
146 obs["obsdec"],
147 file["ra"],
148 file["dec"],
149 degree=True,
150 nearest=True,
151 )
152 )
153 )
154 ang_dist = np.array(ang_dist)
155 ang_use = np.min(ang_dist)
156 i_use = np.where(np.abs(ang_dist - ang_use) <= 1)
157 files = files[i_use]
158 freq_dist = []
159 for file in files:
160 freq_dist.append(file["freq"])
161 freq_dist = np.abs(np.array(freq_dist) - (obs["freq_center"] / 1e6))
162 min_i = np.argmin(freq_dist)
163 pyfhd_config["healpix_inds"] = importlib_resources.files(
164 "PyFHD.resources.healpix"
165 ).joinpath(files[min_i]["name"])
166 hpx_inds = load(pyfhd_config["healpix_inds"], logger=logger)
167 if type(hpx_inds) is dict: 167 ↛ 171line 167 didn't jump to line 171 because the condition on line 167 was always true
168 if "nside" in hpx_inds: 168 ↛ 170line 168 didn't jump to line 170 because the condition on line 168 was always true
169 nside = hpx_inds["nside"]
170 hpx_inds = hpx_inds["hpx_inds"]
171 if nside is None: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 pix_sky = (
173 4 * np.pi * ((180 / np.pi) ** 2) / np.prod(np.abs(obs["astr"]["cdelt"]))
174 )
175 nside = 2 ** (np.ceil(np.log(np.sqrt(pix_sky / 12)) / np.log(2)))
177 # If you wish to implement the keyword divide_pixel_area implement it here and
178 # add it as an option inside pyfhd_config
180 if hpx_inds is not None: 180 ↛ 186line 180 didn't jump to line 186 because the condition on line 180 was always true
181 pix_coords = np.vstack(pix2vec(int(nside), hpx_inds)).T
182 pix_ra, pix_dec = vec2ang(pix_coords, lonlat=True)
183 # This assume the refraction fix on FHD has been implemented
184 xv_hpx, yv_hpx = radec_to_pixel(pix_ra, pix_dec, obs["astr"])
185 else:
186 cen_coords = ang2vec(obs["obsra"], obs["obsdec"], lonlat=True)
187 hpx_inds = query_disc(nside, cen_coords, hpx_radius)
188 pix_coords = np.vstack(pix2vec(nside, hpx_inds)).T
189 pix_dec, pix_ra = vec2ang(pix_coords, lonlat=True)
190 xv_hpx, yv_hpx = radec_to_pixel(pix_ra, pix_dec, obs["astr"])
191 # slightly more restrictive boundary here ('LT' and 'GT' instead of 'LE' and 'GE')
192 pix_i_use = np.where(
193 (xv_hpx > 0)
194 & (xv_hpx < obs["dimension"] - 1)
195 & (yv_hpx > 0)
196 & (yv_hpx < obs["elements"] - 1)
197 )
198 xv_hpx = xv_hpx[pix_i_use]
199 yv_hpx = yv_hpx[pix_i_use]
200 if mask is not None:
201 hpx_mask00 = mask[np.floor(xv_hpx), np.floor(yv_hpx)]
202 hpx_mask01 = mask[np.floor(xv_hpx), np.ceil(yv_hpx)]
203 hpx_mask10 = mask[np.ceil(xv_hpx), np.floor(yv_hpx)]
204 hpx_mask11 = mask[np.ceil(xv_hpx), np.ceil(yv_hpx)]
205 hpx_mask = hpx_mask00 * hpx_mask01 * hpx_mask10 * hpx_mask11
206 pix_i_use2 = np.nonzero(hpx_mask)
207 xv_hpx = xv_hpx[pix_i_use2]
208 yv_hpx = yv_hpx[pix_i_use2]
209 pix_i_use = pix_i_use[pix_i_use2]
210 hpx_inds = hpx_inds[pix_i_use]
212 # Test for pixels past the horizon. We don't need to be precise with this, so turn off precession, etc..
213 # Get the location from instrument name
214 location = EarthLocation.of_site(obs["instrument"])
215 alt, _ = radec_to_altaz(
216 pix_ra,
217 pix_dec,
218 location.lat.value,
219 location.lon.value,
220 location.height.value,
221 obs["jd0"],
222 )
223 horizon_i = np.where(alt <= 0)
224 if np.size(horizon_i) > 0: 224 ↛ 225line 224 didn't jump to line 225 because the condition on line 224 was never true
225 logger.info("Cutting the HEALPix pixels that were below the horizon")
226 h_use = np.where(alt > 0)
227 xv_hpx = xv_hpx[h_use]
228 yv_hpx = yv_hpx[h_use]
229 hpx_inds = hpx_inds[h_use]
230 # The differences in precision through the use of vec2ang, radec_to_pixel are exposed
231 # directly causing differences in the results. We can probably assume these numbers are the results. We can probably assume these numbers are
232 # "better" compared to IDL
233 x_frac = 1 - (xv_hpx - np.floor(xv_hpx))
234 y_frac = 1 - (yv_hpx - np.floor(yv_hpx))
236 v_floor = np.floor(xv_hpx) + obs["dimension"] * np.floor(yv_hpx)
237 v_ceil = np.ceil(xv_hpx) + obs["dimension"] * np.ceil(yv_hpx)
238 # Differences in precision occur here compared to IDL, unsolvable ones as we're using
239 # built in HEALPIX and astropy functions to get these arrays. We can probably assume these
240 # numbers are "better" compared to IDL, as a result the v_floor and v_ceil arrays are one off
241 # compared to IDL, making min_bin off by one
242 min_bin = max(np.nanmin(v_floor), 0)
243 max_bin = min(np.nanmax(v_ceil), obs["dimension"] * obs["elements"] - 1)
244 h00, _, ri00 = histogram(v_floor, min=min_bin, max=max_bin)
245 h01, _, ri01 = histogram(
246 np.floor(xv_hpx) + obs["dimension"] * np.ceil(yv_hpx), min=min_bin, max=max_bin
247 )
248 h10, _, ri10 = histogram(
249 np.ceil(xv_hpx) + obs["dimension"] * np.floor(yv_hpx), min=min_bin, max=max_bin
250 )
251 h11, _, ri11 = histogram(v_ceil, min=min_bin, max=max_bin)
252 htot = h00 + h01 + h10 + h11
253 inds = np.nonzero(htot)[0]
255 n_arr = htot[inds]
256 i_use = (inds + min_bin).astype(np.int64)
257 # To make sure it's an array we can store into a HDF5 file we're
258 # creating object arrays full of Nones which will be populated by
259 # NumPy arrays, creating a "variable length" array of arrays, to which
260 # the save function in pyfhd_io can now handle through the variable_length
261 # parameter
262 sa = np.full(np.size(n_arr), None, dtype=object)
263 ija = np.full(np.size(n_arr), None, dtype=object)
265 for i in range(np.size(n_arr)):
266 ind0 = inds[i]
267 sa0 = np.zeros(n_arr[i])
268 ija0 = np.zeros(n_arr[i], dtype=np.int64)
269 hist_arr = np.array(
270 [0, h00[ind0], h01[ind0], h10[ind0], h11[ind0]], dtype=np.int64
271 )
272 bin_i = np.cumsum(hist_arr)
273 if h00[ind0] > 0:
274 bi = 0
275 inds1 = ri00[ri00[ind0] : ri00[ind0 + 1]]
276 sa0[bin_i[bi] : bin_i[bi + 1]] = x_frac[inds1] * y_frac[inds1]
277 ija0[bin_i[bi] : bin_i[bi + 1]] = inds1
278 if h01[ind0] > 0:
279 bi = 1
280 inds1 = ri01[ri01[ind0] : ri01[ind0 + 1]]
281 sa0[bin_i[bi] : bin_i[bi + 1]] = x_frac[inds1] * (1 - y_frac[inds1])
282 ija0[bin_i[bi] : bin_i[bi + 1]] = inds1
283 if h10[ind0] > 0:
284 bi = 2
285 inds1 = ri10[ri10[ind0] : ri10[ind0 + 1]]
286 sa0[bin_i[bi] : bin_i[bi + 1]] = (1 - x_frac[inds1]) * y_frac[inds1]
287 ija0[bin_i[bi] : bin_i[bi + 1]] = inds1
288 if h11[ind0] > 0:
289 bi = 3
290 inds1 = ri11[ri11[ind0] : ri11[ind0 + 1]]
291 sa0[bin_i[bi] : bin_i[bi + 1]] = (1 - x_frac[inds1]) * (1 - y_frac[inds1])
292 ija0[bin_i[bi] : bin_i[bi + 1]] = inds1
293 # Since we didn't translate pixel_area we can ignore the if statement coverring
294 # the case where pixel_area is greater than 1
295 sa[i] = sa0
296 ija[i] = ija0
298 hpx_cnv = {"nside": nside, "ija": ija, "sa": sa, "i_use": i_use, "inds": hpx_inds}
299 obs["healpix"]["nside"] = nside
300 if pyfhd_config["restrict_healpix_inds"]: 300 ↛ 303line 300 didn't jump to line 303 because the condition on line 300 was always true
301 obs["healpix"]["ind_list"] = None
302 else:
303 obs["healpix"]["ind_list"] = hpx_inds
304 obs["healpix"]["n_pix"] = np.size(hpx_inds)
305 mask_test = healpix_cnv_apply(mask, hpx_cnv)
306 mask_test_i0 = np.where(mask_test == 0)
307 obs["healpix"]["n_zero"] = np.size(mask_test_i0[0])
309 return hpx_cnv, obs
312def beam_image_cube(
313 obs: dict,
314 psf: dict | h5py.File,
315 logger: Logger,
316 freq_i_arr: NDArray[np.integer] | None = None,
317 pol_i_arr: NDArray[np.integer] | None = None,
318 n_freq_bin: int | None = None,
319 square: bool = True,
320 beam_threshold: float | None = None,
321) -> tuple[NDArray[np.complex128], NDArray[np.float64]]:
322 """
323 Calculate the beam image per unflagged frequency channel to build a cube.
324 Build a contiguous mask in x and y based off of a threshold value.
326 Parameters
327 ----------
328 obs : dict
329 Observation metadata dictionary
330 psf : dict | h5py.File
331 Beam dictionary
332 logger : Logger
333 PyFHD's Logger
334 freq_i_arr : np.ndarray | None, optional
335 Index array of the beam frequency binning in the observation
336 metadata dictionary, by default None
337 pol_i_arr : np.ndarray | None, optional
338 Index array of the polarizations in the observation metadata
339 dictionary, by default None
340 n_freq_bin : int | None, optional
341 Number of frequency channels to divide the full observation bandwidth
342 into for the beam image cube, by default None
343 square : bool, optional
344 Square the beam image, by default True
345 beam_threshold : float | None, optional
346 Fractional threshold in which to build a mask for all smaller
347 beam values, by default None
349 Returns
350 -------
351 tuple[NDArray[np.complex128], NDArray[np.float64]]
352 Beam image frequency cube and mask threshold array.
353 """
355 if beam_threshold is None: 355 ↛ 356line 355 didn't jump to line 356 because the condition on line 355 was never true
356 beam_threshold = 0.05
358 if pol_i_arr is None: 358 ↛ 361line 358 didn't jump to line 361 because the condition on line 358 was always true
359 pol_i_arr = np.arange(obs["n_pol"])
361 if n_freq_bin is not None: 361 ↛ 365line 361 didn't jump to line 365 because the condition on line 361 was always true
362 freq_i_arr = np.floor(
363 np.arange(n_freq_bin) * (obs["n_freq"] / n_freq_bin)
364 ).astype(np.int64)
365 if freq_i_arr is None: 365 ↛ 366line 365 didn't jump to line 366 because the condition on line 365 was never true
366 logger.error("beam_image_cube requires n_freq_bin or freq_i_arr to be set")
367 sys.exit(1)
368 if n_freq_bin is None and freq_i_arr is not None: 368 ↛ 369line 368 didn't jump to line 369 because the condition on line 368 was never true
369 n_freq_bin = freq_i_arr.size
371 if np.median(freq_i_arr) > obs["n_freq"]: 371 ↛ 372line 371 didn't jump to line 372 because the condition on line 371 was never true
372 freq_i_use = np.interp(
373 np.arange(obs["n_freq"]), obs["baseline_info"]["freq"], freq_i_arr
374 )
375 else:
376 freq_i_use = freq_i_arr
378 beam_arr = np.zeros([obs["n_pol"], n_freq_bin, obs["dimension"], obs["elements"]])
380 bin_arr = obs["baseline_info"]["fbin_i"][freq_i_use]
381 bin_hist, _, bri = histogram(bin_arr, min=0)
382 bin_use = np.nonzero(bin_hist)[0]
383 if np.size(bin_use) == 0: 383 ↛ 384line 383 didn't jump to line 384 because the condition on line 383 was never true
384 return beam_arr
385 bin_n = bin_hist[bin_use]
386 beam_mask = np.ones([obs["dimension"], obs["elements"]])
387 for pol_i in range(obs["n_pol"]):
388 for fb_i in range(np.size(bin_use)):
389 f_i_i = bri[bri[bin_use[fb_i]] : bri[bin_use[fb_i] + 1]]
390 f_i = freq_i_use[f_i_i[0]]
391 beam_single = beam_image(psf, obs, pol_i, freq_i=f_i, square=square)
392 beam_arr[pol_i, f_i_i[0] : f_i_i[bin_n[fb_i] - 1] + 1] = beam_single
393 b_i = int(obs["obsx"] + obs["obsy"] * obs["dimension"])
394 beam_i = region_grow(
395 beam_single,
396 b_i,
397 low=beam_threshold ** (square + 1),
398 high=np.max(beam_single),
399 )
400 beam_mask1 = np.zeros([obs["dimension"], obs["elements"]])
401 beam_mask1.flat[beam_i] = 1
402 beam_mask *= beam_mask1
403 return beam_arr, beam_mask
406def phase_shift_uv_image(obs: dict) -> NDArray[np.complex128]:
407 """
408 Phase shift the entire u-v plane to the original phase center.
410 Parameters
411 ----------
412 obs : dict
413 Observation metadata dictionary
415 Returns
416 -------
417 NDArray[np.float64]
418 A 2D u-v phase shift array that will shift the current phase centre to
419 the original phase centre.
420 """
421 # Since we only use it once in PyFHD, assume we always want to do /to_orig_phase
422 # Implement the other options if you decide to use this function elsewhere
423 ra_use = obs["orig_phasera"]
424 dec_use = obs["orig_phasedec"]
426 # Skip calculations if phased correctly
427 if ( 427 ↛ 431line 427 didn't jump to line 431 because the condition on line 427 was never true
428 obs["phasera"] == obs["orig_phasera"]
429 and obs["phasedec"] == obs["orig_phasedec"]
430 ):
431 return np.ones([obs["dimension"], obs["elements"]], dtype=np.complex128)
433 x, y = radec_to_pixel(ra_use, dec_use, obs["astr"])
435 # uv_mask is not applied in FHD examples or docs decided not to translate it, if you want it put it here
437 dx = (x - (obs["dimension"] / 2)) * (2 * np.pi / obs["dimension"])
438 dy = (y - (obs["elements"] / 2)) * (2 * np.pi / obs["dimension"])
440 xvals = meshgrid(obs["dimension"], obs["elements"], 1) - (obs["dimension"] / 2)
441 yvals = meshgrid(obs["dimension"], obs["elements"], 2) - (obs["elements"] / 2)
443 phase = xvals * dx + yvals * dy
444 rephase_vals = np.cos(phase) + np.sin(phase) * 1j
446 # Again no uv_mask therefore just return the rephase as is
448 return rephase_vals
451def vis_model_freq_split(
452 obs: dict,
453 psf: dict | h5py.File,
454 params: dict,
455 vis_weights: NDArray[np.float64],
456 vis_model_arr: NDArray[np.complex128],
457 vis_arr: NDArray[np.complex128],
458 polarization: int,
459 pyfhd_config: dict,
460 logger: Logger,
461 fft: bool = True,
462 save_uvf: bool = True,
463 uvf_name: str = "",
464 bi_use: NDArray[np.integer] = None,
465) -> dict:
466 """
467 Grid as a function of frequency, which can be split into larger channels. This
468 uvf cube can be saved as is or used to create an orthoslant image as a function
469 of frequency. Optional bi_use argument allows for the separate gridding of even/odd
470 interleaved time samples for error propagation in the power spectrum.
472 Parameters
473 ----------
474 obs : dict
475 Observation metadata dictionary
476 psf : dict | h5py.File
477 Beam dictionary
478 params : dict
479 Visibility metadata dictionary
480 vis_weights : NDArray[np.float64]
481 Visibility weights array
482 vis_model_arr : NDArray[np.complex128]
483 Model visibility array
484 vis_arr : NDArray[np.complex128]
485 Calibrated visibility array
486 polarization : int
487 The polarization index
488 pyfhd_config : dict
489 PyFHD configuration settings
490 logger : Logger
491 PyFHD's Logger
492 fft : bool, optional
493 Calculate the orthoslant image instead of the u-v plane,
494 by default True
495 save_uvf : bool, optional
496 Save the uvf cube, by default True
497 uvf_name : str, optional
498 The name after the obsid in the uvf cube filename, by default ''
499 bi_use : NDArray[np.integer], optional
500 The visibility indices to use in gridding, i.e. even/odd
501 interleaved time samples, by default None
503 Returns
504 -------
505 cube_split: dict
506 Dictionary of the gridded u-v planes or orthoslant images as a function
507 of frequency for the calibrated data, model (if present), residual data
508 (if present), sampling map, and variance map, along with the observation
509 metadata dictionary.
510 """
512 freq_bin_i2 = np.arange(obs["n_freq"]) // pyfhd_config["n_avg"]
513 nf = int(np.max(freq_bin_i2) + 1)
514 if save_uvf:
515 dirty_uv_arr = np.zeros(
516 [nf, obs["dimension"], obs["dimension"]], dtype=np.complex128
517 )
518 weights_uv_arr = np.zeros(
519 [nf, obs["dimension"], obs["dimension"]], dtype=np.complex128
520 )
521 variance_uv_arr = np.zeros(
522 [nf, obs["dimension"], obs["dimension"]], dtype=np.complex128
523 )
524 model_uv_arr = np.zeros(
525 [nf, obs["dimension"], obs["dimension"]], dtype=np.complex128
526 )
527 dirty_arr = np.zeros([nf, obs["dimension"], obs["dimension"]])
528 weights_arr = np.zeros([nf, obs["dimension"], obs["dimension"]])
529 variance_arr = np.zeros([nf, obs["dimension"], obs["dimension"]])
530 model_arr = np.zeros([nf, obs["dimension"], obs["dimension"]])
531 vis_n_arr = np.zeros([nf])
532 if pyfhd_config["rephase_weights"]:
533 rephase_use = phase_shift_uv_image(obs)
534 else:
535 rephase_use = 1
536 # No x_range and y_range is used in PyFHD, if you wish to do that add that here
538 if bi_use is None:
539 if obs["n_pol"] > 1:
540 flag_test = np.maximum(np.maximum(vis_weights[0], vis_weights[1]), 0)
541 # Double check the axis used
542 flag_test = np.sum(flag_test, axis=1)
543 bi_use = np.where(flag_test > 0)[0]
544 else:
545 flag_test = np.maximum(vis_weights[0], 0)
546 flag_test = np.sum(flag_test, axis=1)
547 bi_use = np.where(flag_test > 0)[0]
549 n_vis_use = 0
550 for fi in range(nf):
551 fi_use = np.where((freq_bin_i2 == fi) & (obs["baseline_info"]["freq_use"] > 0))[
552 0
553 ]
554 if np.size(fi_use) == 0:
555 continue
556 gridding_dict = visibility_grid(
557 vis_arr[polarization],
558 vis_weights[polarization],
559 obs,
560 psf,
561 params,
562 polarization,
563 pyfhd_config,
564 logger,
565 model=vis_model_arr[polarization],
566 fi_use=fi_use,
567 bi_use=bi_use,
568 verbose_logging=False,
569 )
570 if not gridding_dict:
571 logger.warning(
572 f"No visibilities gridded for frequency channel {fi_use} and polarization {obs['pol_names'][polarization]} ({polarization})"
573 )
574 continue
575 n_vis_use += gridding_dict["n_vis"]
576 vis_n_arr[fi] = gridding_dict["n_vis"]
578 if save_uvf:
579 dirty_uv_arr[fi] = gridding_dict["image_uv"] * gridding_dict["n_vis"]
580 weights_uv_arr[fi] = (
581 gridding_dict["weights"] * rephase_use * gridding_dict["n_vis"]
582 )
583 variance_uv_arr[fi] = (
584 gridding_dict["variance"] * rephase_use * gridding_dict["n_vis"]
585 )
586 model_uv_arr[fi] = gridding_dict["model_return"] * gridding_dict["n_vis"]
588 if fft:
589 # No x_range and y_range hence no check for it here
590 dirty_arr[fi], _, _ = dirty_image_generate(
591 gridding_dict["image_uv"],
592 pyfhd_config,
593 logger,
594 degpix=obs["degpix"],
595 )
596 dirty_arr[fi] *= gridding_dict["n_vis"]
597 weights_arr[fi], _, _ = dirty_image_generate(
598 gridding_dict["weights"] * rephase_use,
599 pyfhd_config,
600 logger,
601 degpix=obs["degpix"],
602 )
603 weights_arr[fi] *= gridding_dict["n_vis"]
604 variance_arr[fi], _, _ = dirty_image_generate(
605 gridding_dict["variance"] * rephase_use,
606 pyfhd_config,
607 logger,
608 degpix=obs["degpix"],
609 )
610 variance_arr[fi] *= gridding_dict["n_vis"]
611 model_arr[fi], _, _ = dirty_image_generate(
612 gridding_dict["model_return"],
613 pyfhd_config,
614 logger,
615 degpix=obs["degpix"],
616 )
617 model_arr[fi] *= gridding_dict["n_vis"]
618 else:
619 dirty_arr[fi] = gridding_dict["image_uv"] * gridding_dict["n_vis"]
620 weights_arr[fi] = (
621 gridding_dict["weights"] * rephase_use * gridding_dict["n_vis"]
622 )
623 variance_arr[fi] = (
624 gridding_dict["variance"] * rephase_use * gridding_dict["n_vis"]
625 )
626 model_arr[fi] = gridding_dict["model_return"] * gridding_dict["n_vis"]
627 obs["n_vis"] = n_vis_use
629 if save_uvf:
630 h5_save_dict = {
631 "dirty_uv": dirty_uv_arr,
632 "weights_uv": weights_uv_arr,
633 "variance_uv": variance_uv_arr,
634 "model_uv": model_uv_arr,
635 }
636 uvf_dir = Path(Path(pyfhd_config["output_dir"], "healpix", "uvf_grid"))
637 uvf_dir.mkdir(exist_ok=True, parents=True)
638 save(
639 Path(
640 uvf_dir,
641 f'{pyfhd_config["obs_id"]}_{uvf_name}_{obs["pol_names"][polarization]}_dirty_uv_arr_gridded_uvf.h5',
642 ),
643 h5_save_dict,
644 f'{pyfhd_config["obs_id"]}_{uvf_name}_{obs["pol_names"][polarization]}_dirty_uv_arr_gridded_uvf.h5',
645 logger=logger,
646 )
647 save(
648 Path(
649 uvf_dir,
650 f'{pyfhd_config["obs_id"]}_{uvf_name}_{obs["pol_names"][polarization]}_weights_uv_gridded_uvf.h5',
651 ),
652 h5_save_dict,
653 f'{pyfhd_config["obs_id"]}_{uvf_name}_{obs["pol_names"][polarization]}_weights_uv_gridded_uvf.h5',
654 logger=logger,
655 )
656 save(
657 Path(
658 uvf_dir,
659 f'{pyfhd_config["obs_id"]}_{uvf_name}_{obs["pol_names"][polarization]}_variance_uv_arr_gridded_uvf.h5',
660 ),
661 h5_save_dict,
662 f'{pyfhd_config["obs_id"]}_{uvf_name}_{obs["pol_names"][polarization]}_variance_uv_arr_gridded_uvf.h5',
663 logger=logger,
664 )
665 save(
666 Path(
667 uvf_dir,
668 f'{pyfhd_config["obs_id"]}_{uvf_name}_{obs["pol_names"][polarization]}_model_uv_arr_gridded_uvf.h5',
669 ),
670 h5_save_dict,
671 f'{pyfhd_config["obs_id"]}_{uvf_name}_{obs["pol_names"][polarization]}_model_uv_arr_gridded_uvf.h5',
672 logger=logger,
673 )
675 cube_split = {
676 "obs": obs,
677 "residual_arr": dirty_arr,
678 "weights_arr": weights_arr,
679 "variance_arr": variance_arr,
680 "model_arr": model_arr,
681 }
683 return cube_split