Coverage for PyFHD/gridding/visibility_degrid.py: 3%
176 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-07-01 10:58 +0800
1import numpy as np
2from numpy.typing import NDArray
3import warnings
4import h5py
5from PyFHD.gridding.gridding_utils import (
6 baseline_grid_locations,
7 interpolate_kernel,
8 grid_beam_per_baseline,
9)
10from PyFHD.pyfhd_tools.pyfhd_utils import (
11 l_m_n,
12 rebin,
13 weight_invert,
14 histogram,
15 deriv_coefficients,
16)
19def visibility_degrid(
20 image_uv: NDArray[np.complex128],
21 vis_weights: NDArray[np.float64],
22 obs: dict,
23 psf: dict | h5py.File,
24 params: dict,
25 polarization: int = 0,
26 fill_model_visibilities: bool = False,
27 vis_input: NDArray[np.complex128] | None = None,
28 spectral_model_uv_arr: NDArray[np.float64] | None = None,
29 beam_per_baseline: bool = False,
30 uv_grid_phase_only: bool = True,
31 conserve_memory: bool = False,
32 memory_threshold: float | int = 1e8,
33):
34 """
35 Generate visibilities from a 2D hyperresolved {u,v} plane using the Fourier transform of the beam sensitivity as the
36 kernel (or integration function). The input {u,v} plane is the slant-orthographic projection of the sky when Fourier
37 transformed with no instrumental effects. The integration kernel adds the instrumental effects to the visiblities.
39 Degridding is performed only on simulated model {u,v} planes. Simulated sources towards the edge of the sky image will be
40 distorted (smeared) due to the projection of the 2D {u,v} plane.
42 The kernel is a extremely hyperresolved look-up table, which is (optionally) interpolated even further.
43 Since the {u,v} pixels are discrete and the baseline locations are not, the kernel will integrate the pixels
44 in a unique way for each individual baseline. This code is optimized to provide the best estimate for each
45 baseline whilst maintaining speed.
47 Parameters
48 ----------
49 image_uv : NDArray[np.complex128]
50 A simulated {u,v} plane with no instrumental effects
51 vis_weight : NDArray[np.float64]
52 Weights (flags) of the visibilities
53 obs : dict
54 Observation metadata dictionary
55 psf : dict | h5py.File
56 Beam metadata dictionary
57 params : dict
58 Visibility metadata dictionary
59 polarization : int
60 Index of the current polarization, by default 0
61 fill_model_visibilities : bool, optional
62 Create all model visibilities disregarding flags, by default False
63 vis_input : NDArray[np.complex128] | None, optional
64 Extra model visibilities to add to the degridded products, by default None
65 spectral_model_uv_arr : NDArray[np.float64] | None, optional
66 Additional {u,v} planes to degrid for complicated source spectral dependencies, by default None
67 beam_per_baseline : bool, optional
68 Generate beams with corrective phases given the baseline location, by default False
69 uv_grid_phase_only : bool, optional
70 Generate beams with only {u,v} corrective phases, disregarding w phases, by default True
71 conserve_memory : bool, optional
72 Reduce memory load by running loops, by default False
74 Returns
75 -------
76 visibility_array : NDArray[np.complex128]
77 A simulated visibility array from degridding the input {u,v} plane with the instrumental kernel
78 obs : dict
79 Updated observation metadata dictionary
80 """
82 complex_flag = psf["complex_flag"][0]
83 n_spectral = obs["degrid_spectral_terms"][0]
84 interp_flag = psf["interpolate_kernel"][0]
85 if conserve_memory:
86 # memory threshold is in bytes
87 if memory_threshold < 1e6:
88 memory_threshold = 1e8
90 # If both beam and interp_flag leave a warning, prioritise beam_per_baseline
91 if beam_per_baseline and interp_flag:
92 warnings.warn(
93 "Cannot have beam per baseline and interpolation at the same time, turning off interpolation"
94 )
95 interp_flag = False
97 # For each unflagged baseline, get the minimum contributing pixel number for gridding
98 # and the 2D derivatives for bilinear interpolation
99 baselines_dict = baseline_grid_locations(
100 obs,
101 psf,
102 params,
103 vis_weights,
104 fill_model_visibilities=fill_model_visibilities,
105 interp_flag=interp_flag,
106 )
108 # Extract data from the data structures
109 bin_n = baselines_dict["bin_n"]
110 bin_i = baselines_dict["bin_i"]
111 n_bin_use = baselines_dict["n_bin_use"]
112 ri = baselines_dict["ri"]
113 xmin = baselines_dict["xmin"]
114 ymin = baselines_dict["ymin"]
115 x_offset = baselines_dict["x_offset"]
116 y_offset = baselines_dict["y_offset"]
117 if interp_flag:
118 dx0dy0_arr = baselines_dict["dx0dy0_arr"]
119 dx0dy1_arr = baselines_dict["dx0dy1_arr"]
120 dx1dy0_arr = baselines_dict["dx1dy0_arr"]
121 dx1dy1_arr = baselines_dict["dx1dy1_arr"]
122 dimension = obs["dimension"][0]
123 kbinsize = obs["kpix"][0]
124 freq_bin_i = obs["baseline_info"][0]["fbin_i"][0]
125 frequency_array = obs["baseline_info"][0]["freq"][0]
126 freq_delta = (frequency_array - obs["freq_center"][0]) / obs["freq_center"][0]
127 psf_dim = psf["dim"][0]
128 psf_resolution = psf["resolution"][0]
129 psf_dim3 = int(psf_dim**2)
130 n_baselines = obs["n_baselines"][0]
131 n_samples = obs["n_time"][0]
132 n_freq_use = frequency_array.size
133 n_freq = obs["n_freq"][0]
134 group_arr = np.squeeze(psf["id"][0][:, freq_bin_i, polarization])
135 beam_arr = psf["beam_ptr"][0]
137 if beam_per_baseline:
138 uu = params["uu"][0]
139 vv = params["vv"][0]
140 ww = params["ww"][0]
141 psf_image_dim = psf["image_info"][0]["psf_image_dim"][0]
142 psf_intermediate_res = np.min(
143 np.ceil(np.sqrt(psf_resolution) / 2) * 2, psf_resolution
144 )
145 image_bot = -(psf_dim / 2) * psf_intermediate_res + psf_image_dim / 2
146 image_top = (
147 (psf_dim * psf_resolution - 1)
148 - (psf_dim / 2) * psf_intermediate_res
149 + psf_image_dim / 2
150 )
152 l_mode, m_mode, n_tracked = l_m_n(obs, psf)
154 # w-terms have not been tested, thus they've been turned off for now
155 if uv_grid_phase_only:
156 n_tracked = np.zeros_like(n_tracked)
158 beam_int = np.zeros(n_freq)
159 beam2_int = np.zeros(n_freq)
160 n_grp_use = np.zeros(n_freq)
161 primary_beam_area = np.zeros(n_freq)
162 primary_beam_sq_area = np.zeros(n_freq)
164 x = (np.arange(dimension) - dimension / 2) * kbinsize
165 y = x.copy()
167 conj_i = np.where(params["vv"][0] > 0)
168 if conj_i[0].size > 0:
169 if beam_per_baseline:
170 uu[conj_i] = -uu[conj_i]
171 vv[conj_i] = -vv[conj_i]
172 ww[conj_i] = -ww[conj_i]
174 # Create the correct size visibility array
175 vis_dimension = n_baselines * n_samples
176 visibility_array = np.zeros((vis_dimension, n_freq), dtype=np.cdouble)
178 ind_ref = np.arange(max(bin_n))
180 if complex_flag:
181 arr_type = np.cdouble
182 else:
183 arr_type = np.double
185 if n_spectral:
186 prefactor = np.empty(n_spectral, dtype=object)
187 for s_i in range(n_spectral):
188 prefactor[s_i] = deriv_coefficients(s_i + 1, divide_factorial=True)
189 box_arr_ptr = np.empty(n_spectral, dtype=object)
191 for bi in range(n_bin_use):
192 vis_n = bin_n[bin_i[bi]]
193 """
194 Python is not inclusive of end of loop
195 while IDL is, as such n_bin_use - 1 does not do
196 the last index properly on Python, as such when we
197 have reached the last index we need to change the
198 indexation of ri to include from that point to the end.
199 """
200 if bi == n_bin_use - 1:
201 inds = ri[ri[bin_i[bi]] :]
202 else:
203 inds = ri[ri[bin_i[bi]] : ri[bin_i[bi + 1]]]
205 # if constraining memory usage, then est number of loops needed
206 if conserve_memory:
207 required_bytes = 8 * vis_n * psf_dim3
208 mem_iter = int(np.ceil(required_bytes / memory_threshold))
209 if mem_iter > 1:
210 vis_n_full = vis_n
211 inds_full = inds
212 vis_n_per_iter = int(np.ceil(vis_n_full / mem_iter))
213 else:
214 mem_iter = 1
215 # loop over chunks of visibilities to grid to conserve memory
216 for mem_i in range(mem_iter):
217 if mem_iter > 1:
218 # calculate the indices of this visibility chunk if split into multiple chunks
219 if vis_n_per_iter * (mem_i + 1) > vis_n_full:
220 max_ind = vis_n_full
221 else:
222 max_ind = vis_n_per_iter * (mem_i + 1)
223 inds = inds_full[vis_n_per_iter * mem_i : max_ind]
224 vis_n = max_ind - vis_n_per_iter * mem_i
226 x_off = x_offset.flat[inds].astype(np.int64)
227 y_off = y_offset.flat[inds].astype(np.int64)
228 # xmin and ymin should be all the same
229 xmin_use = xmin.flat[inds[0]]
230 ymin_use = ymin.flat[inds[0]]
231 freq_i = inds % n_freq_use
232 fbin = freq_bin_i[freq_i]
233 baseline_inds = (inds / n_freq_use).astype(int) % n_baselines
235 box_matrix = np.zeros((vis_n, psf_dim3), dtype=arr_type)
236 box_arr = image_uv[
237 ymin_use : ymin_use + psf_dim, xmin_use : xmin_use + psf_dim
238 ].flatten()
240 if interp_flag:
241 dx0dy0 = dx0dy0_arr.flat[inds]
242 dx0dy1 = dx0dy1_arr.flat[inds]
243 dx1dy0 = dx1dy0_arr.flat[inds]
244 dx1dy1 = dx1dy1_arr.flat[inds]
245 ind_remap_flag = False
246 bt_index = inds / n_freq_use
247 for ii in range(vis_n):
248 kernel = interpolate_kernel(
249 beam_arr[baseline_inds[ii], fbin[ii], polarization],
250 x_off[ii],
251 y_off[ii],
252 dx0dy0[ii],
253 dx1dy0[ii],
254 dx0dy1[ii],
255 dx1dy1[ii],
256 )
257 box_matrix.flat[psf_dim3 * ii : psf_dim3 * ii + kernel.size] = (
258 kernel
259 )
260 else:
261 group_id = group_arr[inds]
262 group_max = np.max(group_id) + 1
263 xyf_i = (
264 x_off + y_off * psf_resolution + fbin * psf_resolution**2
265 ) * group_max + group_id
266 xyf_si = np.sort(xyf_i)
267 xyf_i = xyf_i[xyf_si]
268 xyf_ui = np.unique(xyf_i)
269 n_xyf_bin = xyf_ui.size
271 # There might be a better selection criteria to determine which is more efficient
272 if vis_n > np.ceil(1.1 * n_xyf_bin) and not beam_per_baseline:
273 ind_remap_flag = True
274 inds = inds[xyf_si]
275 inds_use = xyf_si[xyf_ui]
276 freq_i = freq_i[inds_use]
277 x_off = x_off[inds_use]
278 y_off = y_off[inds_use]
279 fbin = fbin[inds_use]
280 baseline_inds = baseline_inds[inds]
282 if n_xyf_bin == 1:
283 ind_remap = np.arange(vis_n, dtype=int)
284 else:
285 hist_inds_u, _, ri_xyf = histogram(xyf_ui, bin_size=1, min=0)
286 ind_remap = ind_ref[ri_xyf[0 : hist_inds_u.size] - ri_xyf[0]]
288 vis_n = n_xyf_bin
289 else:
290 ind_remap_flag = False
291 bt_index = inds / n_freq_use
293 if beam_per_baseline:
294 box_matrix = grid_beam_per_baseline(
295 psf,
296 uu,
297 vv,
298 ww,
299 l_mode,
300 m_mode,
301 n_tracked,
302 frequency_array,
303 x,
304 y,
305 xmin_use,
306 ymin_use,
307 freq_i,
308 bt_index,
309 polarization,
310 fbin,
311 image_bot,
312 image_top,
313 psf_dim3,
314 box_matrix,
315 vis_n,
316 beam_int=beam_int,
317 beam2_int=beam2_int,
318 n_grp_use=n_grp_use,
319 degrid_flag=True,
320 obs=obs,
321 params=params,
322 weights=vis_weights,
323 )
324 else:
325 for ii in range(vis_n):
326 # more efficient array subscript notation
327 box_matrix[psf_dim3 * ii] = beam_arr[
328 baseline_inds[ii], fbin[ii], polarization
329 ][y_off[ii], x_off[ii]]
331 if n_spectral:
332 vis_box = np.dot(box_arr, np.transpose(box_matrix))
333 freq_term_arr = rebin(
334 np.transpose(freq_delta[freq_i]), (vis_n, psf_dim3), sample=True
335 )
336 for s_i in range(n_spectral):
337 # s_i loop is over terms of the Taylor expansion, starting from the lowest-order term
338 prefactor_use = prefactor[s_i]
339 box_matrix *= freq_term_arr
340 box_arr_ptr[s_i] = spectral_model_uv_arr[s_i][
341 ymin_use : ymin_use + psf_dim - 1,
342 xmin_use : xmin_use + psf_dim - 1,
343 ].flatten()
345 for s_i_i in range(s_i):
346 # s_i_i loop is over powers of the model x alpha^n, n=s_i_i+1
347 box_arr = prefactor_use[s_i_i] * box_arr_ptr[s_i_i]
348 vis_box += np.dot(box_arr, np.transpose(box_matrix))
349 del box_arr_ptr
350 else:
351 vis_box = np.dot(box_arr, np.transpose(box_matrix))
352 if ind_remap_flag:
353 vis_box = vis_box[ind_remap]
354 visibility_array.flat[inds] = vis_box
356 if beam_per_baseline:
357 # factor of kbinsize^2 is FFT units normalization
358 beam2_int *= weight_invert(n_grp_use) / kbinsize**2
359 beam_int *= weight_invert(n_grp_use) / kbinsize**2
360 primary_beam_area = beam2_int
361 primary_beam_area = beam_int
362 # Returning obs in this case?
363 obs["primary_beam_area"][polarization] = primary_beam_area
364 obs["primary_beam_sq_area"][polarization] = primary_beam_sq_area
366 del (x_offset, y_offset, xmin, ymin, bin_n)
367 if conj_i[0].size > 0:
368 visibility_array[conj_i, :] = np.conj(visibility_array[conj_i, :])
370 if vis_input is not None:
371 visibility_array += vis_input
373 if beam_per_baseline:
374 return visibility_array, obs
375 else:
376 return visibility_array