def chan_average(chan_meta, chan_freq=None, chan_width=None): have_chan_freq = not is_numba_type_none(chan_freq) have_chan_width = not is_numba_type_none(chan_width) chan_freq_output = chan_output_factory(have_chan_freq) chan_width_output = chan_output_factory(have_chan_width) chan_freq_normaliser = normaliser_factory(have_chan_freq) chan_freq_adder = add_factory(have_chan_freq) chan_width_adder = add_factory(have_chan_width) def impl(chan_meta, chan_freq=None, chan_width=None): chan_map, out_chans = chan_meta chan_freq_avg = chan_freq_output(out_chans, chan_freq) chan_width_avg = chan_width_output(out_chans, chan_width) counts = np.zeros(out_chans, dtype=np.uint32) for in_chan, out_chan in enumerate(chan_map): counts[out_chan] += 1 chan_freq_adder(chan_freq_avg, out_chan, chan_freq, in_chan) chan_width_adder(chan_width_avg, out_chan, chan_width, in_chan) for out_chan in range(out_chans): chan_freq_normaliser(chan_freq_avg, out_chan, counts[out_chan]) return ChannelAverageOutput(chan_freq_avg, chan_width_avg) return impl
def merge_flags(flag_row, flag): have_flag_row = not is_numba_type_none(flag_row) have_flag = not is_numba_type_none(flag) if have_flag_row and have_flag: def impl(flag_row, flag): """ Check flag_row and flag agree """ for r in range(flag.shape[0]): all_flagged = True for f in range(flag.shape[1]): for c in range(flag.shape[2]): if flag[r, f, c] == 0: all_flagged = False break if not all_flagged: break if (flag_row[r] != 0) != all_flagged: raise ValueError("flag_row and flag arrays mismatch") return flag_row elif have_flag_row and not have_flag: def impl(flag_row, flag): """ Return flag_row """ return flag_row elif not have_flag_row and have_flag: def impl(flag_row, flag): """ Construct flag_row from flag """ new_flag_row = np.empty(flag.shape[0], dtype=flag.dtype) for r in range(flag.shape[0]): all_flagged = True for f in range(flag.shape[1]): for c in range(flag.shape[2]): if flag[r, f, c] == 0: all_flagged = False break if not all_flagged: break new_flag_row[r] = (1 if all_flagged else 0) return new_flag_row else: def impl(flag_row, flag): return None return impl
def chan_average(chan_meta, chan_freq=None, chan_width=None, effective_bw=None, resolution=None): have_chan_freq = not is_numba_type_none(chan_freq) have_chan_width = not is_numba_type_none(chan_width) have_effective_bw = not is_numba_type_none(effective_bw) have_resolution = not is_numba_type_none(resolution) chan_freq_output = chan_output_factory(have_chan_freq) chan_width_output = chan_output_factory(have_chan_width) effective_bw_output = chan_output_factory(have_effective_bw) resolution_output = chan_output_factory(have_resolution) chan_freq_normaliser = normaliser_factory(have_chan_freq) chan_freq_adder = add_factory(have_chan_freq) chan_width_adder = add_factory(have_chan_width) effective_bw_adder = add_factory(have_effective_bw) resolution_adder = add_factory(have_resolution) def impl(chan_meta, chan_freq=None, chan_width=None, effective_bw=None, resolution=None): chan_map, out_chans = chan_meta chan_freq_avg = chan_freq_output(out_chans, chan_freq) chan_width_avg = chan_width_output(out_chans, chan_width) effective_bw_avg = effective_bw_output(out_chans, effective_bw) resolution_avg = resolution_output(out_chans, resolution) counts = np.zeros(out_chans, dtype=np.uint32) for in_chan, out_chan in enumerate(chan_map): counts[out_chan] += 1 chan_freq_adder(chan_freq_avg, out_chan, chan_freq, in_chan) chan_width_adder(chan_width_avg, out_chan, chan_width, in_chan) effective_bw_adder(effective_bw_avg, out_chan, effective_bw, in_chan) resolution_adder(resolution_avg, out_chan, resolution, in_chan) for out_chan in range(out_chans): chan_freq_normaliser(chan_freq_avg, out_chan, counts[out_chan]) return ChannelAverageOutput(chan_freq_avg, chan_width_avg, effective_bw_avg, resolution_avg) return impl
def radec_to_lm(radec, phase_centre=None): dtype = radec.dtype if is_numba_type_none(phase_centre): _maybe_create_phase_centre = _create_phase_centre else: _maybe_create_phase_centre = _return_phase_centre def _radec_to_lm_impl(radec, phase_centre=None): sources, components = radec.shape if radec.ndim != 2 or components != 2: raise ValueError("radec must have shape (source, 2)") lm = np.empty(shape=(sources, 2), dtype=dtype) pc_ra, pc_dec = _maybe_create_phase_centre(phase_centre, dtype) sin_pc_dec = np.sin(pc_dec) cos_pc_dec = np.cos(pc_dec) for s in range(sources): da = radec[s, 0] - pc_ra sin_ra_delta = np.sin(da) cos_ra_delta = np.cos(da) sin_dec = np.sin(radec[s, 1]) cos_dec = np.cos(radec[s, 1]) lm[s, 0] = cos_dec * sin_ra_delta lm[s, 1] = sin_dec * cos_pc_dec - cos_dec * sin_pc_dec * cos_ra_delta return lm return _radec_to_lm_impl
def lm_to_radec(lm, phase_centre=None): dtype = lm.dtype if is_numba_type_none(phase_centre): _maybe_create_phase_centre = _create_phase_centre else: _maybe_create_phase_centre = _return_phase_centre def _lm_to_radec_impl(lm, phase_centre=None): if lm.ndim != 2 or lm.shape[1] != 2: raise ValueError("lm must have shape (source, 2)") radec = np.empty(shape=(lm.shape[0], 2), dtype=dtype) pc_ra, pc_dec = _maybe_create_phase_centre(phase_centre, dtype) sin_pc_dec = np.sin(pc_dec) cos_pc_dec = np.cos(pc_dec) for s in range(radec.shape[0]): l, m = lm[s] n = np.sqrt(1.0 - l**2 - m**2) radec[s, 1] = np.arcsin(m * cos_pc_dec + n * sin_pc_dec) radec[s, 0] = pc_ra + np.arctan(l / (n * cos_pc_dec - m * sin_pc_dec)) return radec return _lm_to_radec_impl
def lmn_to_radec(lmn, phase_centre=None): dtype = lmn.dtype if is_numba_type_none(phase_centre): _maybe_create_phase_centre = _create_phase_centre else: _maybe_create_phase_centre = _return_phase_centre @wraps(lmn_to_radec) def _lmn_to_radec_impl(lmn, phase_centre=None): if lmn.ndim != 2 or lmn.shape[1] != 3: raise ValueError("lmn must have shape (source, 3)") radec = np.empty(shape=(lmn.shape[0], 2), dtype=dtype) pc_ra, pc_dec = _maybe_create_phase_centre(phase_centre, dtype) sin_pc_dec = np.sin(pc_dec) cos_pc_dec = np.cos(pc_dec) for s in range(radec.shape[0]): l, m, n = lmn[s] radec[s, 1] = np.arcsin(m*cos_pc_dec + n*sin_pc_dec) radec[s, 0] = pc_ra + np.arctan(l / (n*cos_pc_dec - m*sin_pc_dec)) return radec return _lmn_to_radec_impl
def vis_to_im(vis, uvw, lm, frequency, flags, dtype=None): # Infer output dtype if none provided if is_numba_type_none(dtype): # Support both real and complex visibilities... if isinstance(vis.dtype, numba.types.scalars.Complex): vis_comp_dtype = np.dtype(vis.dtype.underlying_float.name) else: vis_comp_dtype = np.dtype(vis.dtype.name) out_dtype = np.result_type( vis_comp_dtype, *(np.dtype(a.dtype.name) for a in (uvw, lm, frequency))) else: if isinstance(dtype, numba.types.scalars.Complex): raise TypeError("dtype must be complex") out_dtype = dtype.dtype assert np.shape(vis) == np.shape(flags) def _vis_to_im_impl(vis, uvw, lm, frequency, flags, dtype=None): nrows = uvw.shape[0] nsrc = lm.shape[0] nchan = frequency.shape[0] ncorr = vis.shape[-1] im_of_vis = np.zeros((nsrc, nchan, ncorr), dtype=out_dtype) # For each source for s in range(nsrc): l, m = lm[s] n = np.sqrt(1.0 - l**2 - m**2) - 1.0 # For each uvw coordinate for r in range(nrows): u, v, w = uvw[r] # e^(-2*pi*(l*u + m*v + n*w)/c) real_phase = -minus_two_pi_over_c * (l * u + m * v + n * w) # Multiple in frequency for each channel for nu in range(nchan): p = real_phase * frequency[nu] # do not compute if any of the correlations # are flagged (complicates uncertainties) if np.any(flags[r, nu]): continue for c in range(ncorr): # elide the call to exp since result is real im_of_vis[s, nu, c] += (np.cos(p) * vis[r, nu, c].real - np.sin(p) * vis[r, nu, c].imag) return im_of_vis return _vis_to_im_impl
def _is_chan_flagged(flag, r, f, c): if is_numba_type_none(flag): def impl(flag, r, f, c): return True else: def impl(flag, r, f, c): return flag[r, f, c] return impl
def _flags_match(flag_row, ri, out_flag_row, ro): if is_numba_type_none(flag_row): def impl(flag_row, ri, out_flag_row, ro): return True else: def impl(flag_row, ri, out_flag_row, ro): return flag_row[ri] == out_flag_row[ro] return impl
def im_to_vis(image, uvw, lm, frequency, convention='fourier', dtype=None): # Infer complex output dtype if none provided if is_numba_type_none(dtype): out_dtype = np.result_type( np.complex64, *(np.dtype(a.dtype.name) for a in (image, uvw, lm, frequency))) else: out_dtype = dtype.dtype def impl(image, uvw, lm, frequency, convention='fourier', dtype=None): if convention == 'fourier': constant = minus_two_pi_over_c elif convention == 'casa': constant = two_pi_over_c else: raise ValueError("convention not in ('fourier', 'casa')") nrows = uvw.shape[0] nsrc = lm.shape[0] nchan = frequency.shape[0] ncorr = image.shape[-1] vis_of_im = np.zeros((nrows, nchan, ncorr), dtype=out_dtype) # For each uvw coordinate for r in range(nrows): u, v, w = uvw[r] # For each source for s in range(nsrc): l, m = lm[s] n = np.sqrt(1.0 - l**2 - m**2) - 1.0 # e^(-2*pi*(l*u + m*v + n*w)/c) real_phase = constant * (l * u + m * v + n * w) # Multiple in frequency for each channel for nu in range(nchan): p = real_phase * frequency[nu] * 1.0j for c in range(ncorr): if image[s, nu, c]: vis_of_im[r, nu, c] += np.exp(p) * image[s, nu, c] return vis_of_im return impl
def shape_or_invalid_shape(array, ndim): """ Return array shape tuple or (-1,)*ndim if the array is None """ try: ndim_lit = getattr(ndim, "literal_value") except AttributeError: raise ValueError("ndim must be a integer literal") if is_numba_type_none(array): tup = (-1, ) * ndim_lit def impl(array, ndim): return tup else: def impl(array, ndim): return array.shape return impl
def _get_jones_types(name, numba_ndarray_type, corr_1_dims, corr_2_dims): """ Determine which of the following three cases are valid: 1. The array is not present (None) and therefore no Jones Matrices 2. single (1,) or (2,) dual correlation 3. (2, 2) full correlation Parameters ---------- name: str Array name numba_ndarray_type: numba.type Array numba type corr_1_dims: int Number of `numba_ndarray_type` dimensions, including correlations (first option) corr_2_dims: int Number of `numba_ndarray_type` dimensions, including correlations (second option) Returns ------- int Enumeration describing the Jones Matrix Type - 0 -- Not Present - 1 -- (1,) or (2,) - 2 -- (2, 2) """ if is_numba_type_none(numba_ndarray_type): return JONES_NOT_PRESENT if numba_ndarray_type.ndim == corr_1_dims: return JONES_1_OR_2 elif numba_ndarray_type.ndim == corr_2_dims: return JONES_2X2 else: raise ValueError("%s.ndim not in (%d, %d)" % (name, corr_1_dims, corr_2_dims))
def radec_to_lmn(radec, phase_centre=None): dtype = radec.dtype if is_numba_type_none(phase_centre): _maybe_create_phase_centre = _create_phase_centre else: _maybe_create_phase_centre = _return_phase_centre @wraps(radec_to_lmn) def _radec_to_lmn_impl(radec, phase_centre=None): sources, components = radec.shape if radec.ndim != 2 or components != 2: raise ValueError("radec must have shape (source, 2)") lmn = np.empty(shape=(sources, 3), dtype=dtype) pc_ra, pc_dec = _maybe_create_phase_centre(phase_centre, dtype) sin_pc_dec = np.sin(pc_dec) cos_pc_dec = np.cos(pc_dec) for s in range(sources): ra_delta = radec[s, 0] - pc_ra sin_ra_delta = np.sin(ra_delta) cos_ra_delta = np.cos(ra_delta) sin_dec = np.sin(radec[s, 1]) cos_dec = np.cos(radec[s, 1]) lmn[s, 0] = l = cos_dec*sin_ra_delta # noqa lmn[s, 1] = m = (sin_dec*cos_pc_dec - cos_dec*sin_pc_dec*cos_ra_delta) lmn[s, 2] = np.sqrt(1.0 - l**2 - m**2) return lmn return _radec_to_lmn_impl
def predict_vis(time_index, antenna1, antenna2, dde1_jones=None, source_coh=None, dde2_jones=None, die1_jones=None, base_vis=None, die2_jones=None): tup = predict_checks(time_index, antenna1, antenna2, dde1_jones, source_coh, dde2_jones, die1_jones, base_vis, die2_jones, lambda x: not is_numba_type_none(x)) (have_ddes1, have_coh, have_ddes2, have_dies1, have_bvis, have_dies2) = tup # Infer the output dtype dtype_arrays = (dde1_jones, source_coh, dde2_jones, die1_jones, base_vis, die2_jones) out_dtype = np.result_type(*(np.dtype(a.dtype.name) for a in dtype_arrays if not is_numba_type_none(a))) jones_types = [ _get_jones_types("dde1_jones", dde1_jones, 5, 6), _get_jones_types("source_coh", source_coh, 4, 5), _get_jones_types("dde2_jones", dde2_jones, 5, 6), _get_jones_types("die1_jones", die1_jones, 4, 5), _get_jones_types("base_vis", base_vis, 3, 4), _get_jones_types("die2_jones", die2_jones, 4, 5) ] ptypes = [t for t in jones_types if t != JONES_NOT_PRESENT] if not all(ptypes[0] == p for p in ptypes[1:]): raise ValueError("Jones Matrix Correlations were mismatched") try: jones_type = ptypes[0] except IndexError: raise ValueError("No Jones Matrices were supplied") have_ddes = have_ddes1 and have_ddes2 have_dies = have_dies1 and have_dies2 # Create functions that we will use inside our predict function out_fn = output_factory(have_ddes, have_coh, have_dies, have_bvis, out_dtype) sum_coh_fn = sum_coherencies_factory(have_ddes, have_coh, jones_type) apply_dies_fn = apply_dies_factory(have_dies, have_bvis, jones_type) add_coh_fn = add_coh_factory(have_bvis) @wraps(predict_vis) def _predict_vis_fn(time_index, antenna1, antenna2, dde1_jones=None, source_coh=None, dde2_jones=None, die1_jones=None, base_vis=None, die2_jones=None): # Get the output shape out = out_fn(time_index, dde1_jones, source_coh, dde2_jones, die1_jones, base_vis, die2_jones) # Minimum time index, used to normalise within function tmin = time_index.min() # Sum coherencies if any sum_coh_fn(time_index, antenna1, antenna2, dde1_jones, source_coh, dde2_jones, tmin, out) # Add base visibilities to the output, if any add_coh_fn(base_vis, out) # Apply direction independent effects, if any apply_dies_fn(time_index, antenna1, antenna2, die1_jones, die2_jones, tmin, out) return out return _predict_vis_fn
def time_and_channel(time, interval, antenna1, antenna2, time_centroid=None, exposure=None, flag_row=None, uvw=None, weight=None, sigma=None, chan_freq=None, chan_width=None, vis=None, flag=None, weight_spectrum=None, sigma_spectrum=None, time_bin_secs=1.0, chan_bin_size=1): valid_types = (types.misc.Omitted, types.scalars.Float, types.scalars.Integer) if not isinstance(time_bin_secs, valid_types): raise TypeError("time_bin_secs must be a scalar float") valid_types = (types.misc.Omitted, types.scalars.Integer) if not isinstance(chan_bin_size, valid_types): raise TypeError("chan_bin_size must be a scalar integer") have_vis = not is_numba_type_none(vis) have_flag = not is_numba_type_none(flag) have_weight_spectrum = not is_numba_type_none(weight_spectrum) have_sigma_spectrum = not is_numba_type_none(sigma_spectrum) chan_corrs = chan_corr_factory(have_vis, have_flag, have_weight_spectrum, have_sigma_spectrum) def impl(time, interval, antenna1, antenna2, time_centroid=None, exposure=None, flag_row=None, uvw=None, weight=None, sigma=None, chan_freq=None, chan_width=None, vis=None, flag=None, weight_spectrum=None, sigma_spectrum=None, time_bin_secs=1.0, chan_bin_size=1): # Get the number of channels + correlations nchan, ncorr = chan_corrs(vis, flag, weight_spectrum, sigma_spectrum) # Merge flag_row and flag arrays flag_row = merge_flags(flag_row, flag) # Generate row mapping metadata row_meta = row_mapper(time, interval, antenna1, antenna2, flag_row=flag_row, time_bin_secs=time_bin_secs) # Generate channel mapping metadata chan_meta = channel_mapper(nchan, chan_bin_size) # Average row data row_data = row_average(row_meta, antenna1, antenna2, flag_row=flag_row, time_centroid=time_centroid, exposure=exposure, uvw=uvw, weight=weight, sigma=sigma) # Average channel data chan_data = chan_average(chan_meta, chan_freq=chan_freq, chan_width=chan_width) # Average row and channel data row_chan_data = row_chan_average(row_meta, chan_meta, flag_row=flag_row, weight=weight, vis=vis, flag=flag, weight_spectrum=weight_spectrum, sigma_spectrum=sigma_spectrum) # Have to explicitly write it out because numba tuples # are highly constrained types return AverageOutput(row_meta.time, row_meta.interval, row_meta.flag_row, row_data.antenna1, row_data.antenna2, row_data.time_centroid, row_data.exposure, row_data.uvw, row_data.weight, row_data.sigma, chan_data.chan_freq, chan_data.chan_width, row_chan_data.vis, row_chan_data.flag, row_chan_data.weight_spectrum, row_chan_data.sigma_spectrum) return impl
def bda_mapper(time, interval, ant1, ant2, uvw, chan_width, chan_freq, max_uvw_dist, flag_row=None, max_fov=3.0, decorrelation=0.98, time_bin_secs=None, min_nchan=1): have_time_bin_secs = not is_numba_type_none(time_bin_secs) Omitted = numba.types.misc.Omitted decorr_type = (numba.typeof(decorrelation.value) if isinstance( decorrelation, Omitted) else decorrelation) fov_type = (numba.typeof(max_fov.value) if isinstance(max_fov, Omitted) else max_fov) # If time_bin_secs is None, # then we set it to the max of the time dtype # lower down time_bin_secs_type = time_bin_secs if have_time_bin_secs else time.dtype spec = [('tbin', numba.uintp), ('bin_count', numba.uintp), ('bin_flag_count', numba.uintp), ('time_sum', time.dtype), ('interval_sum', interval.dtype), ('rs', numba.uintp), ('re', numba.uintp), ('bin_half_Δψ', uvw.dtype), ('max_lm', fov_type), ('n_max', fov_type), ('decorrelation', decorr_type), ('time_bin_secs', time_bin_secs_type), ('max_chan_freq', chan_freq.dtype), ('max_uvw_dist', max_uvw_dist)] JitBinner = jitclass(spec)(Binner) def impl(time, interval, ant1, ant2, uvw, chan_width, chan_freq, max_uvw_dist, flag_row=None, max_fov=3.0, decorrelation=0.98, time_bin_secs=None, min_nchan=1): # 𝞓 𝝿 𝞇 𝞍 𝝼 if decorrelation < 0.0 or decorrelation > 1.0: raise ValueError("0.0 <= decorrelation <= 1.0 must hold") if max_fov <= 0.0 or max_fov > 90.0: raise ValueError("0.0 < max_fov <= 90.0 must hold") max_lm = np.deg2rad(max_fov) ubl, _, bl_inv, _ = unique_baselines(ant1, ant2) utime, _, time_inv, _ = unique_time(time) nrow = time.shape[0] ntime = utime.shape[0] nbl = ubl.shape[0] nchan = chan_width.shape[0] nchan_factors = factors(nchan) bandwidth = chan_width.sum() if min_nchan is None: min_nchan = 1 else: s = np.searchsorted(nchan_factors, min_nchan, side='left') min_nchan = max(min_nchan, nchan_factors[s]) if nchan == 0: raise ValueError("zero channels") # Create the row lookup row_lookup = np.full((nbl, ntime), -1, dtype=np.int32) bin_lookup = np.full((nbl, ntime), -1, dtype=np.int32) bin_chan_width = np.full((nbl, ntime), 0.0, dtype=chan_width.dtype) sentinel = np.finfo(time.dtype).max time_lookup = np.full((nbl, ntime), sentinel, dtype=time.dtype) interval_lookup = np.full((nbl, ntime), sentinel, dtype=interval.dtype) # Is the entire bin flagged? bin_flagged = np.zeros((nbl, ntime), dtype=np.bool_) bin_chan_map = np.empty((nbl, ntime, nchan), dtype=np.int32) out_rows = 0 nr_of_time_bins = 0 out_row_chans = 0 def update_lookups(finalised, bl): """ Closure which updates lookups for a baseline, given a binner's finalisation data """ # NOTE(sjperkins) Why do scalars need this, but not arrays? nonlocal out_rows nonlocal out_row_chans nonlocal min_nchan tbin = finalised.tbin time_lookup[bl, tbin] = finalised.time interval_lookup[bl, tbin] = finalised.interval bin_flagged[bl, tbin] = finalised.flag nchan = max(finalised.nchan, min_nchan) bin_nchan = chan_width.shape[0] // nchan bin_chan_width[bl, tbin] = bandwidth / finalised.nchan # Construct the channel map for c in range(chan_width.shape[0]): bin_chan_map[bl, tbin, c] = c // bin_nchan out_rows += 1 out_row_chans += nchan for r in range(nrow): t = time_inv[r] bl = bl_inv[r] if row_lookup[bl, t] != -1: raise ValueError("Duplicate (TIME, ANTENNA1, ANTENNA2)") row_lookup[bl, t] = r # If we don't have time_bin_secs # set it to the maximum floating point value, # effectively ignoring this limit if not have_time_bin_secs: time_bin_secs = np.finfo(time.dtype).max # This derived from Synthesis & Imaging II (18-31) # Converts decrease in amplitude into change in phase dphi = np.sqrt(6. * (1. - decorrelation)) binner = JitBinner(0, 0, max_lm, dphi, time_bin_secs, chan_freq.max()) for bl in range(nbl): # Reset the binner for this baseline binner.reset() # Auto-correlated baseline auto_corr = ubl[bl, 0] == ubl[bl, 1] for t in range(ntime): # Lookup row, continue if non-existent r = row_lookup[bl, t] if r == -1: continue # Start a new bin if binner.empty: binner.start_bin(r, time, interval, flag_row) # Try add the row to the bin # If this fails, finalise the current bin and start a new one elif not binner.add_row(r, auto_corr, time, interval, uvw, flag_row): f = binner.finalise_bin(auto_corr, uvw, nchan_factors, chan_width, chan_freq) update_lookups(f, bl) # Post-finalisation, the bin is empty, start a new bin binner.start_bin(r, time, interval, flag_row) # Record the time bin associated with this row bin_lookup[bl, t] = binner.tbin # Finalise any remaining data in the bin if not binner.empty: f = binner.finalise_bin(auto_corr, uvw, nchan_factors, chan_width, chan_freq) update_lookups(f, bl) nr_of_time_bins += binner.tbin # Mark remaining bins as unoccupied and unflagged for tbin in range(binner.tbin, ntime): time_lookup[bl, tbin] = sentinel bin_flagged[bl, tbin] = False assert out_rows == nr_of_time_bins # Flatten the time lookup and argsort it flat_time = time_lookup.ravel() argsort = np.argsort(flat_time, kind='mergesort') inv_argsort = np.empty_like(argsort) # Generate lookup from flattened (bl, time) to output row for i, a in enumerate(argsort): inv_argsort[a] = i # Generate row offsets fbin_chan_map = bin_chan_map.reshape((-1, nchan)) offsets = np.zeros(out_rows + 1, dtype=np.uint32) decorr_chan_width = np.empty(out_rows, dtype=chan_width.dtype) # NOTE(sjperkins) # This: out_rows > 0 # does not work here for some strange (numba reason?) if offsets.shape[0] > 0: offsets[0] = 0 for r in range(1, out_rows + 1): prev_bin_chans = fbin_chan_map[argsort[r - 1]].max() + 1 offsets[r] = offsets[r - 1] + prev_bin_chans # Construct the final row map row_chan_map = np.full((nrow, nchan), -1, dtype=np.int32) time_ret = np.full(out_row_chans, -1, dtype=time.dtype) int_ret = np.full(out_row_chans, -1, dtype=interval.dtype) chan_width_ret = np.full(out_row_chans, -1, dtype=chan_width.dtype) # Construct output flag row, if necessary out_flag_row = (None if flag_row is None else np.empty( out_row_chans, dtype=flag_row.dtype)) # foreach input row for in_row in range(time.shape[0]): # Lookup baseline and time bl = bl_inv[in_row] t = time_inv[in_row] # lookup time bin and output row in inv_argsort tbin = bin_lookup[bl, t] bin_time = time_lookup[bl, tbin] bin_interval = interval_lookup[bl, tbin] flagged = bin_flagged[bl, tbin] out_row = inv_argsort[bl * ntime + tbin] decorr_chan_width[out_row] = bin_chan_width[bl, tbin] # Should never happen, but check if out_row >= out_rows: raise RowMapperError("out_row >= out_rows") # Handle output row flagging if flag_row is not None and flag_row[in_row] == 0 and flagged: raise RowMapperError("Unflagged input row " "contributing to " "flagged output row. " "This should never happen!") # Set up the row channel map, populate # time, interval and chan_width for c in range(nchan): out_offset = offsets[out_row] + bin_chan_map[bl, tbin, c] # Should never happen, but check if out_offset >= out_row_chans: raise RowMapperError("out_offset >= out_row_chans") # Set the output row for this input row and channel row_chan_map[in_row, c] = out_offset # Broadcast the time and interval to the output row time_ret[out_offset] = bin_time int_ret[out_offset] = bin_interval # Add channel contribution for each row chan_width_ret[out_offset] += chan_width[c] if flag_row is not None: out_flag_row[out_offset] = 1 if flagged else 0 return RowMapOutput(row_chan_map, offsets, decorr_chan_width, time_ret, int_ret, chan_width_ret, out_flag_row) return impl
def row_chan_average(row_meta, chan_meta, flag_row=None, weight=None, vis=None, flag=None, weight_spectrum=None, sigma_spectrum=None): have_flag_row = not is_numba_type_none(flag_row) have_vis = not is_numba_type_none(vis) have_flag = not is_numba_type_none(flag) have_weight = not is_numba_type_none(weight) have_weight_spectrum = not is_numba_type_none(weight_spectrum) have_sigma_spectrum = not is_numba_type_none(sigma_spectrum) flags_match = matching_flag_factory(have_flag_row) is_chan_flagged = is_chan_flagged_factory(have_flag) vis_factory = chan_output_factory(have_vis) weight_sum_factory = weight_sum_output_factory(have_vis) flag_factory = chan_output_factory(have_flag) weight_factory = chan_output_factory(have_weight_spectrum) sigma_factory = chan_output_factory(have_sigma_spectrum) vis_adder = vis_add_factory(have_vis, have_weight, have_weight_spectrum) weight_adder = chan_add_factory(have_weight_spectrum) sigma_adder = sigma_spectrum_add_factory(have_sigma_spectrum, have_weight, have_weight_spectrum) vis_normaliser = vis_normaliser_factory(have_vis) sigma_normaliser = sigma_spectrum_normaliser_factory(have_sigma_spectrum) weight_normaliser = weight_spectrum_normaliser_factory( have_weight_spectrum) set_flagged = set_flagged_factory(have_flag) dummy_chan_freq = None dummy_chan_width = None def impl(row_meta, chan_meta, flag_row=None, weight=None, vis=None, flag=None, weight_spectrum=None, sigma_spectrum=None): out_rows = row_meta.time.shape[0] nchan, ncorrs = chan_corrs(vis, flag, weight_spectrum, sigma_spectrum, dummy_chan_freq, dummy_chan_width) chan_map, out_chans = chan_meta out_shape = (out_rows, out_chans, ncorrs) vis_avg = vis_factory(out_shape, vis) vis_weight_sum = weight_sum_factory(out_shape, vis) weight_spectrum_avg = weight_factory(out_shape, weight_spectrum) sigma_spectrum_avg = sigma_factory(out_shape, sigma_spectrum) sigma_spectrum_weight_sum = sigma_factory(out_shape, sigma_spectrum) flagged_vis_avg = vis_factory(out_shape, vis) flagged_vis_weight_sum = weight_sum_factory(out_shape, vis) flagged_weight_spectrum_avg = weight_factory(out_shape, weight_spectrum) flagged_sigma_spectrum_avg = sigma_factory(out_shape, sigma_spectrum) flagged_sigma_spectrum_weight_sum = sigma_factory( out_shape, sigma_spectrum) flag_avg = flag_factory(out_shape, flag) counts = np.zeros(out_shape, dtype=np.uint32) flag_counts = np.zeros(out_shape, dtype=np.uint32) # Iterate over input rows, accumulating into output rows for in_row, out_row in enumerate(row_meta.map): # TIME_CENTROID/EXPOSURE case applies here, # must have flagged input and output OR unflagged input and output if not flags_match(flag_row, in_row, row_meta.flag_row, out_row): continue for in_chan, out_chan in enumerate(chan_map): for corr in range(ncorrs): if is_chan_flagged(flag, in_row, in_chan, corr): # Increment flagged averages and counts flag_counts[out_row, out_chan, corr] += 1 vis_adder(flagged_vis_avg, flagged_vis_weight_sum, vis, weight, weight_spectrum, out_row, out_chan, in_row, in_chan, corr) weight_adder(flagged_weight_spectrum_avg, weight_spectrum, out_row, out_chan, in_row, in_chan, corr) sigma_adder(flagged_sigma_spectrum_avg, flagged_sigma_spectrum_weight_sum, sigma_spectrum, weight, weight_spectrum, out_row, out_chan, in_row, in_chan, corr) else: # Increment unflagged averages and counts counts[out_row, out_chan, corr] += 1 vis_adder(vis_avg, vis_weight_sum, vis, weight, weight_spectrum, out_row, out_chan, in_row, in_chan, corr) weight_adder(weight_spectrum_avg, weight_spectrum, out_row, out_chan, in_row, in_chan, corr) sigma_adder(sigma_spectrum_avg, sigma_spectrum_weight_sum, sigma_spectrum, weight, weight_spectrum, out_row, out_chan, in_row, in_chan, corr) for r in range(out_rows): for f in range(out_chans): for c in range(ncorrs): if counts[r, f, c] > 0: # We have some unflagged samples and # only these are used as averaged output vis_normaliser(vis_avg, vis_avg, r, f, c, vis_weight_sum) sigma_normaliser(sigma_spectrum_avg, sigma_spectrum_avg, r, f, c, sigma_spectrum_weight_sum) elif flag_counts[r, f, c] > 0: # We only have flagged samples and # these are used as averaged output vis_normaliser(vis_avg, flagged_vis_avg, r, f, c, flagged_vis_weight_sum) sigma_normaliser(sigma_spectrum_avg, flagged_sigma_spectrum_avg, r, f, c, flagged_sigma_spectrum_weight_sum) weight_normaliser(weight_spectrum_avg, flagged_weight_spectrum_avg, r, f, c) # Flag the output bin set_flagged(flag_avg, r, f, c) else: raise RowChannelAverageException("Zero-filled bin") return RowChanAverageOutput(vis_avg, flag_avg, weight_spectrum_avg, sigma_spectrum_avg) return impl
def row_average(meta, ant1, ant2, flag_row=None, time_centroid=None, exposure=None, uvw=None, weight=None, sigma=None): have_flag_row = not is_numba_type_none(flag_row) have_uvw = not is_numba_type_none(uvw) have_time_centroid = not is_numba_type_none(time_centroid) have_exposure = not is_numba_type_none(exposure) have_weight = not is_numba_type_none(weight) have_sigma = not is_numba_type_none(sigma) flags_match = matching_flag_factory(have_flag_row) uvw_factory = output_factory(have_uvw) time_centroid_factory = output_factory(have_time_centroid) exposure_factory = output_factory(have_exposure) weight_factory = output_factory(have_weight) sigma_factory = output_factory(have_sigma) time_centroid_adder = add_factory(have_time_centroid) exposure_adder = add_factory(have_exposure) uvw_adder = comp_add_factory(have_uvw) weight_adder = comp_add_factory(have_weight) sigma_adder = sigma_add_factory(have_sigma, have_weight) uvw_normaliser = normaliser_factory(have_uvw) sigma_normaliser = sigma_normaliser_factory(have_sigma) time_centroid_normaliser = normaliser_factory(have_time_centroid) def impl(meta, ant1, ant2, flag_row=None, time_centroid=None, exposure=None, uvw=None, weight=None, sigma=None): out_rows = meta.time.shape[0] counts = np.zeros(out_rows, dtype=np.uint32) # These outputs are always present ant1_avg = np.empty(out_rows, ant1.dtype) ant2_avg = np.empty(out_rows, ant2.dtype) # Possibly present outputs for possibly present inputs uvw_avg = uvw_factory(out_rows, uvw) time_centroid_avg = time_centroid_factory(out_rows, time_centroid) exposure_avg = exposure_factory(out_rows, exposure) weight_avg = weight_factory(out_rows, weight) sigma_avg = sigma_factory(out_rows, sigma) sigma_weight_sum = sigma_factory(out_rows, sigma) # Iterate over input rows, accumulating into output rows for in_row, out_row in enumerate(meta.map): # Input and output flags must match in order for the # current row to contribute to these columns if flags_match(flag_row, in_row, meta.flag_row, out_row): uvw_adder(uvw_avg, out_row, uvw, in_row) weight_adder(weight_avg, out_row, weight, in_row) sigma_adder(sigma_avg, sigma_weight_sum, out_row, sigma, weight, in_row) time_centroid_adder(time_centroid_avg, out_row, time_centroid, in_row) exposure_adder(exposure_avg, out_row, exposure, in_row) counts[out_row] += 1 # Here we can simply assign because input_row baselines # should always match output row baselines ant1_avg[out_row] = ant1[in_row] ant2_avg[out_row] = ant2[in_row] # Normalise for out_row in range(out_rows): count = counts[out_row] if count > 0: uvw_normaliser(uvw_avg, out_row, count) time_centroid_normaliser(time_centroid_avg, out_row, count) sigma_normaliser(sigma_avg, out_row, sigma_weight_sum) return RowAverageOutput(ant1_avg, ant2_avg, time_centroid_avg, exposure_avg, uvw_avg, weight_avg, sigma_avg) return impl
def row_average(meta, ant1, ant2, flag_row=None, time_centroid=None, exposure=None, uvw=None, weight=None, sigma=None): have_flag_row = not is_numba_type_none(flag_row) flags_match = matching_flag_factory(have_flag_row) def impl(meta, ant1, ant2, flag_row=None, time_centroid=None, exposure=None, uvw=None, weight=None, sigma=None): out_rows = meta.time.shape[0] counts = np.zeros(out_rows, dtype=np.uint32) # These outputs are always present ant1_avg = np.empty(out_rows, ant1.dtype) ant2_avg = np.empty(out_rows, ant2.dtype) # Possibly present outputs for possibly present inputs uvw_avg = (None if uvw is None else np.zeros( (out_rows, ) + uvw.shape[1:], dtype=uvw.dtype)) time_centroid_avg = (None if time_centroid is None else np.zeros( (out_rows, ) + time_centroid.shape[1:], dtype=time_centroid.dtype)) exposure_avg = (None if exposure is None else np.zeros( (out_rows, ) + exposure.shape[1:], dtype=exposure.dtype)) weight_avg = (None if weight is None else np.zeros( (out_rows, ) + weight.shape[1:], dtype=weight.dtype)) sigma_avg = (None if sigma is None else np.zeros( (out_rows, ) + sigma.shape[1:], dtype=sigma.dtype)) sigma_weight_sum = (None if sigma is None else np.zeros( (out_rows, ) + sigma.shape[1:], dtype=sigma.dtype)) # Iterate over input rows, accumulating into output rows for in_row, out_row in enumerate(meta.map): # Input and output flags must match in order for the # current row to contribute to these columns if flags_match(flag_row, in_row, meta.flag_row, out_row): if uvw is not None: uvw_avg[out_row, 0] += uvw[in_row, 0] uvw_avg[out_row, 1] += uvw[in_row, 1] uvw_avg[out_row, 2] += uvw[in_row, 2] if time_centroid is not None: time_centroid_avg[out_row] += time_centroid[in_row] if exposure is not None: exposure_avg[out_row] += exposure[in_row] if weight is not None: for co in range(weight.shape[1]): weight_avg[out_row, co] += weight[in_row, co] if sigma is not None: for co in range(sigma.shape[1]): sva = sigma[in_row, co]**2 # Use provided weights if weight is not None: wt = weight[in_row, co] sva *= wt**2 sigma_weight_sum[out_row, co] += wt # Natural weights else: sigma_weight_sum[out_row, co] += 1.0 # Assign sigma_avg[out_row, co] += sva counts[out_row] += 1 # Here we can simply assign because input_row baselines # should always match output row baselines ant1_avg[out_row] = ant1[in_row] ant2_avg[out_row] = ant2[in_row] # Normalise for out_row in range(out_rows): count = counts[out_row] if count > 0: # Normalise uvw if uvw is not None: uvw_avg[out_row, 0] /= count uvw_avg[out_row, 1] /= count uvw_avg[out_row, 2] /= count # Normalise time centroid if time_centroid is not None: time_centroid_avg[out_row] /= count # Normalise sigma if sigma is not None: for co in range(sigma.shape[1]): ssva = sigma_avg[out_row, co] wt = sigma_weight_sum[out_row, co] if wt != 0.0: ssva /= (wt**2) sigma_avg[out_row, co] = np.sqrt(ssva) return RowAverageOutput(ant1_avg, ant2_avg, time_centroid_avg, exposure_avg, uvw_avg, weight_avg, sigma_avg) return impl
def row_mapper(time, interval, antenna1, antenna2, flag_row=None, time_bin_secs=1): """ Generates a mapping from a high resolution row index to a low resolution row index in support of time and channel averaging code. The `time` and `interval` columns are also respectively averaged and summed in the process of creating the mapping and a `flag_row` column is returned if one is provided. In order to average a chunk of row data, it is necessary to group each row (or sample) by baseline and then average the time samples present in each baseline in bins of `time_bin_secs`. Flagged data is handled as follows: 1. It does not contribute to a bin at all if there are other unflagged samples in the bin. 2. It is the only contribution to a bin if all samples in the bin are flagged. The algorithm is robust in the presence of missing time and baseline data. The algorithm works as follows: 1. `time`, `interval`, `antenna1` and `antenna2` are used to construct a `row_lookup` array of shape `(ubl, utime)` mapping a baseline and time to a row of input data. 2. For each baseline, `time_bin_secs` times are averaged together into two separate `time_lookup` arrays of shape `(ubl, utime)`. The first contains the average of unflagged samples, while the second contains the average of flagged samples. If the bin contains some unflagged samples, the unflagged average is used as the bin average, whereas if all samples are flagged the flagged average is used. Not all bins may be filled for a baseline if data is missing -- these bins are assigned a sentinel value set to the maximum floating point value. A secondary `bin_lookup` array of shape `(ubl, utime)` is constructed mapping a time in the `row_lookup` array to a time bin in `time_lookup`. 3. The `time_lookup` array is flattened and argsorted with a stable merge sort. As missing values are set to the maximum floating point value, this moves valid data to the front and missing data to the back. This has the effect of lexicographically sorts the data in an ascending `(time, bl)` order 4. Input rows are then mapped via the `row_lookup`, `bin_lookup` and argsorted `time_lookup` arrays to an output row. .. code-block:: python ret = row_mapper(time, interval, ant1, ant2, flag_row, time_bin_secs=3) # Only add a bin's contribution if both input and output # are (a) flagged or (b) unflagged sel = flag_row == ret.flag_row[ret.map] sel_map = ret.map[sel] # Recompute time average using row map time = np.zeros_like(ret.time) ant1_avg = np.empty(time.shape, ant1.dtype) ant2_avg = np.empty(time.shape, ant2.dtype) counts = np.empty(time.shape, np.uint32) # Add time and 1 at map indices to time and counts np.add.at(time, sel_map, time[sel]) np.add.at(counts, sel_map, 1) # Normalise time /= count np.testing.assert_array_equal(time, ret.time) # We assign baselines because each input baseline # is mapped to the same output baseline ant1_avg[sel_map] = ant1[sel] ant2_avg[sel_map] = ant2[sel] Parameters ---------- time : :class:`numpy.ndarray` Time values of shape :code:`(row,)`. interval : :class:`numpy.ndarray` Exposure times of shape :code:`(row,)`. antenna1 : :class:`numpy.ndarray` Antenna 1 values of shape :code:`(row,)`. antenna2 : :class:`numpy.ndarray` Antenna 2 values of shape :code:`(row,)`. flag_row : :class:`numpy.ndarray`, optional Positive values indicate that a row is flagged, while zero implies unflagged. Has shape :code:`(row,)`. time_bin_secs : int, optional Number of timesteps to average into each bin Returns ------- map : :class:`numpy.ndarray` Mapping from `np.arange(row)` to output row indices of shape :code:`(row,)` time : :class:`numpy.ndarray` Averaged time values of shape :code:`(out_row,)` interval : :class:`numpy.ndarray` Summed interval values of shape :code:`(out_row,)` flag_row : :class:`numpy.ndarray` or None Output flag rows of shape :code:`(out_row,)`. None if no input flag_row was supplied. Raises ------ RowMapperError Raised if an illegal condition occurs """ have_flag_row = not is_numba_type_none(flag_row) is_flagged_fn = is_flagged_factory(have_flag_row) output_flag_row = output_factory(have_flag_row) set_flag_row = set_flag_row_factory(have_flag_row) def impl(time, interval, antenna1, antenna2, flag_row=None, time_bin_secs=1): ubl, _, bl_inv, _ = unique_baselines(antenna1, antenna2) utime, _, time_inv, _ = unique_time(time) nbl = ubl.shape[0] ntime = utime.shape[0] sentinel = np.finfo(time.dtype).max out_rows = numba.uint32(0) scratch = np.full(3 * nbl * ntime, -1, dtype=np.int32) row_lookup = scratch[:nbl * ntime].reshape(nbl, ntime) bin_lookup = scratch[nbl * ntime:2 * nbl * ntime].reshape(nbl, ntime) inv_argsort = scratch[2 * nbl * ntime:] time_lookup = np.zeros((nbl, ntime), dtype=time.dtype) interval_lookup = np.zeros((nbl, ntime), dtype=interval.dtype) bin_flagged = np.zeros((nbl, ntime), dtype=np.bool_) # Create a mapping from the full bl x time resolution back # to the original input rows for r in range(time.shape[0]): bl = bl_inv[r] t = time_inv[r] row_lookup[bl, t] = r # Average times over each baseline and construct the # bin_lookup and time_lookup arrays for bl in range(ubl.shape[0]): tbin = numba.int32(0) bin_count = numba.int32(0) bin_flag_count = numba.int32(0) bin_low = time.dtype.type(0) for t in range(utime.shape[0]): # Lookup input row r = row_lookup[bl, t] # Ignore if not present if r == -1: continue # At this point, we decide whether to contribute to # the current bin, or create a new one. We don't add # the current sample to the current bin if # high - low >= time_bin_secs half_int = interval[r] * 0.5 # We're starting a new bin anyway, # just set the lower bin value if bin_count == 0: bin_low = time[r] - half_int # If we exceed the seconds in the bin, # normalise the time and start a new bin elif time[r] + half_int - bin_low > time_bin_secs: # Normalise and flag the bin # if total counts match flagged counts if bin_count > 0: time_lookup[bl, tbin] /= bin_count bin_flagged[bl, tbin] = bin_count == bin_flag_count # There was nothing in the bin else: time_lookup[bl, tbin] = sentinel bin_flagged[bl, tbin] = False tbin += 1 bin_count = 0 bin_flag_count = 0 # Record the output bin associated with the row bin_lookup[bl, t] = tbin # Time + Interval take unflagged + unflagged # samples into account (nominal value) time_lookup[bl, tbin] += time[r] interval_lookup[bl, tbin] += interval[r] bin_count += 1 # Record flags if is_flagged_fn(flag_row, r): bin_flag_count += 1 # Normalise the last bin if it has entries in it if bin_count > 0: time_lookup[bl, tbin] /= bin_count bin_flagged[bl, tbin] = bin_count == bin_flag_count tbin += 1 # Add this baseline's number of bins to the output rows out_rows += tbin # Set any remaining bins to sentinel value and unflagged for b in range(tbin, ntime): time_lookup[bl, b] = sentinel bin_flagged[bl, b] = False # Flatten the time lookup and argsort it flat_time = time_lookup.ravel() flat_int = interval_lookup.ravel() argsort = np.argsort(flat_time, kind='mergesort') # Generate lookup from flattened (bl, time) to output row for i, a in enumerate(argsort): inv_argsort[a] = i # Construct the final row map row_map = np.empty((time.shape[0]), dtype=np.uint32) # Construct output flag row, if necessary out_flag_row = output_flag_row(out_rows, flag_row) # foreach input row for in_row in range(time.shape[0]): # Lookup baseline and time bl = bl_inv[in_row] t = time_inv[in_row] # lookup time bin and output row tbin = bin_lookup[bl, t] # lookup output row in inv_argsort out_row = inv_argsort[bl * ntime + tbin] if out_row >= out_rows: raise RowMapperError("out_row >= out_rows") # Handle output row flagging set_flag_row(flag_row, in_row, out_flag_row, out_row, bin_flagged[bl, tbin]) row_map[in_row] = out_row time_ret = flat_time[argsort[:out_rows]] int_ret = flat_int[argsort[:out_rows]] return RowMapOutput(row_map, time_ret, int_ret, out_flag_row) return impl
def _shape_or_invalid_shape(array, ndim): """ Return array shape tuple or (-1,)*ndim if the array is None """ import numba.core.types as nbtypes from numba.extending import SentryLiteralArgs SentryLiteralArgs(['ndim']).for_function(_shape_or_invalid_shape).bind( array, ndim) try: ndim_lit = getattr(ndim, "literal_value") except AttributeError: raise ValueError("ndim must be a integer literal") if is_numba_type_none(array): tup = (-1, ) * ndim_lit def impl(array, ndim): return tup return impl elif isinstance(array, nbtypes.Array): def impl(array, ndim): return array.shape return impl elif (isinstance(array, nbtypes.UniTuple) and isinstance(array.dtype, nbtypes.Array)): if len(array) == 1: def impl(array, ndim): return array[0].shape else: def impl(array, ndim): shape = array[0].shape for a in array[1:]: if a.shape != shape: raise ValueError("Array shapes in Tuple don't match") return shape return impl elif isinstance(array, nbtypes.Tuple): if not all(isinstance(a, nbtypes.Array) for a in array.types): raise ValueError("Must be Tuple of Arrays") if not all(array.types[0].ndim == a.ndim for a in array.types[1:]): raise ValueError("Array ndims in Tuple don't match") if len(array) == 1: def impl(array, ndim): return array[0].shape else: def impl(array, ndim): shape = array[0].shape for a in array[1:]: if a.shape != shape: raise ValueError("Array shapes in Tuple don't match") return shape return impl
def row_average(meta, ant1, ant2, flag_row=None, time_centroid=None, exposure=None, uvw=None, weight=None, sigma=None): have_flag_row = not is_numba_type_none(flag_row) have_time_centroid = not is_numba_type_none(time_centroid) have_exposure = not is_numba_type_none(exposure) have_uvw = not is_numba_type_none(uvw) have_weight = not is_numba_type_none(weight) have_sigma = not is_numba_type_none(sigma) def impl(meta, ant1, ant2, flag_row=None, time_centroid=None, exposure=None, uvw=None, weight=None, sigma=None): out_rows = meta.time.shape[0] counts = np.zeros(out_rows, dtype=np.uint32) # These outputs are always present ant1_avg = np.empty(out_rows, ant1.dtype) ant2_avg = np.empty(out_rows, ant2.dtype) # Possibly present outputs for possibly present inputs uvw_avg = ( None if not have_uvw else np.zeros((out_rows,) + uvw.shape[1:], dtype=uvw.dtype)) time_centroid_avg = ( None if not have_time_centroid else np.zeros((out_rows,) + time_centroid.shape[1:], dtype=time_centroid.dtype)) exposure_avg = ( None if not have_exposure else np.zeros((out_rows,) + exposure.shape[1:], dtype=exposure.dtype)) weight_avg = ( None if not have_weight else np.zeros((out_rows,) + weight.shape[1:], dtype=weight.dtype)) sigma_avg = ( None if not have_sigma else np.zeros((out_rows,) + sigma.shape[1:], dtype=sigma.dtype)) sigma_weight_sum = ( None if not have_sigma else np.zeros((out_rows,) + sigma.shape[1:], dtype=sigma.dtype)) # Average each array, if present # The output is a flattened row-channel array # where the values for each row are repeated along the channel # Individual runs in this output are described by meta.offset # Thus, we only compute the sum in the first position for ri in range(meta.map.shape[0]): ro = meta.map[ri, 0] # Here we can simply assign because input_row baselines # should always match output row baselines ant1_avg[ro] = ant1[ri] ant2_avg[ro] = ant2[ri] # Input and output flags must match in order for the # current row to contribute to these columns if have_flag_row and flag_row[ri] != meta.flag_row[ro]: continue counts[ro] += 1 if have_uvw: uvw_avg[ro, 0] += uvw[ri, 0] uvw_avg[ro, 1] += uvw[ri, 1] uvw_avg[ro, 2] += uvw[ri, 2] if have_time_centroid: time_centroid_avg[ro] += time_centroid[ri] if have_exposure: exposure_avg[ro] += exposure[ri] if have_weight: for co in range(weight.shape[1]): weight_avg[ro, co] += weight[ri, co] if have_sigma: for co in range(sigma.shape[1]): # Use weights if present else natural weights wt = weight[ri, co] if have_weight else 1.0 # Aggregate sigma_avg[ro, co] += sigma[ri, co]**2 * wt**2 sigma_weight_sum[ro, co] += wt # Normalise and copy for o in range(len(meta.offsets) - 1): bro = meta.offsets[o] nchan = meta.offsets[o + 1] - bro count = counts[bro] for c in range(1, nchan): ant1_avg[bro + c] = ant1_avg[bro] ant2_avg[bro + c] = ant2_avg[bro] if have_uvw: # Normalise channel 0 value if count > 0: uvw_avg[bro, 0] /= count uvw_avg[bro, 1] /= count uvw_avg[bro, 2] /= count # Copy to other channels for c in range(1, nchan): uvw_avg[bro + c, 0] += uvw_avg[bro, 0] uvw_avg[bro + c, 1] += uvw_avg[bro, 1] uvw_avg[bro + c, 2] += uvw_avg[bro, 2] if have_time_centroid: # Normalise channel 0 value if count > 0: time_centroid_avg[bro] /= count # Copy to other channels for c in range(1, nchan): time_centroid_avg[bro + c] = time_centroid_avg[bro] if have_exposure: # Copy to other channels for c in range(1, nchan): exposure_avg[bro + c] = exposure_avg[bro] if have_weight: # Copy to other channels for c in range(1, nchan): for co in range(weight.shape[1]): weight_avg[bro + c, co] = weight_avg[bro, co] if have_sigma: # Normalise channel 0 values for co in range(sigma.shape[1]): sswsum = sigma_weight_sum[bro, co] if sswsum != 0.0: ssva = sigma_avg[bro, co] sigma_avg[bro, co] = np.sqrt(ssva / (sswsum**2)) # Copy values to other channels for c in range(1, nchan): for co in range(sigma.shape[1]): sigma_avg[bro + c, co] = sigma_avg[bro, co] return RowAverageOutput(ant1_avg, ant2_avg, time_centroid_avg, exposure_avg, uvw_avg, weight_avg, sigma_avg) return impl
def row_chan_average(meta, flag_row=None, weight=None, visibilities=None, flag=None, weight_spectrum=None, sigma_spectrum=None): have_vis = not is_numba_type_none(visibilities) have_flag = not is_numba_type_none(flag) have_flag_row = not is_numba_type_none(flag_row) have_flags = have_flag_row or have_flag have_weight = not is_numba_type_none(weight) have_weight_spectrum = not is_numba_type_none(weight_spectrum) have_sigma_spectrum = not is_numba_type_none(sigma_spectrum) def impl(meta, flag_row=None, weight=None, visibilities=None, flag=None, weight_spectrum=None, sigma_spectrum=None): out_rows = meta.time.shape[0] nchan, ncorrs = chan_corrs(visibilities, flag, weight_spectrum, sigma_spectrum, None, None, None, None) out_shape = (out_rows, ncorrs) if not have_flag: flag_avg = None else: flag_avg = np.zeros(out_shape, np.bool_) # If either flag_row or flag is present, we need to ensure that # effective averaging takes place. if have_flags: flags_match = np.zeros(meta.map.shape + (ncorrs,), dtype=np.bool_) flag_counts = np.zeros(out_shape, dtype=np.uint32) else: flags_match = None flag_counts = None counts = np.zeros(out_shape, dtype=np.uint32) # Determine output bin counts both unflagged and flagged for ri in range(meta.map.shape[0]): row_flagged = have_flag_row and flag_row[ri] != 0 for fi in range(meta.map.shape[1]): ro = meta.map[ri, fi] for co in range(ncorrs): flagged = (row_flagged or (have_flag and flag[ri, fi, co] != 0)) if have_flags and flagged: flag_counts[ro, co] += 1 else: counts[ro, co] += 1 # ------ # Flags # ------ # Determine whether input samples should contribute to an output bin # and, if flags are parent, whether the output bin is flagged # This follows from the definition of an effective average: # # * bad or flagged values should be excluded # when calculating the average # # Note that if a bin is completely flagged we still compute an average, # to which all relevant input samples contribute. for ri in range(meta.map.shape[0]): row_flagged = have_flag_row and flag_row[ri] != 0 for fi in range(meta.map.shape[1]): ro = meta.map[ri, fi] for co in range(ncorrs): if counts[ro, co] > 0: # Output bin should only contain unflagged samples out_flag = False if have_flag: # Set output flags flag_avg[ro, co] = False elif have_flags and flag_counts[ro, co] > 0: # Output bin is completely flagged out_flag = True if have_flag: # Set output flags flag_avg[ro, co] = True else: raise RowChannelAverageException("Zero-filled bin") # We should only add a sample to an output bin # if the input flag matches the output flag. # This is because flagged samples don't contribute # to a bin with some unflagged samples while # unflagged samples never contribute to a # completely flagged bin if have_flags: in_flag = (row_flagged or (have_flag and flag[ri, fi, co] != 0)) flags_match[ri, fi, co] = in_flag == out_flag # ------------- # Visibilities # ------------- if not have_vis: vis_avg = None else: vis_avg, vis_weight_sum = vis_output_arrays( visibilities, out_shape) # Aggregate for ri in range(meta.map.shape[0]): for fi in range(meta.map.shape[1]): ro = meta.map[ri, fi] for co in range(ncorrs): if have_flags and not flags_match[ri, fi, co]: continue wt = (weight_spectrum[ri, fi, co] if have_weight_spectrum else weight[ri, co] if have_weight else 1.0) average_visibilities(visibilities, vis_avg, vis_weight_sum, wt, ri, fi, ro, co) # Normalise for ro in range(out_rows): for co in range(ncorrs): normalise_visibilities(vis_avg, vis_weight_sum, ro, co) # ---------------- # Weight Spectrum # ---------------- if not have_weight_spectrum: weight_spectrum_avg = None else: weight_spectrum_avg = np.zeros(out_shape, weight_spectrum.dtype) # Aggregate for ri in range(meta.map.shape[0]): for fi in range(meta.map.shape[1]): ro = meta.map[ri, fi] for co in range(ncorrs): if have_flags and not flags_match[ri, fi, co]: continue weight_spectrum_avg[ro, co] += ( weight_spectrum[ri, fi, co]) # --------------- # Sigma Spectrum # --------------- if not have_sigma_spectrum: sigma_spectrum_avg = None else: sigma_spectrum_avg = np.zeros(out_shape, sigma_spectrum.dtype) sigma_spectrum_weight_sum = np.zeros_like(sigma_spectrum_avg) # Aggregate for ri in range(meta.map.shape[0]): for fi in range(meta.map.shape[1]): ro = meta.map[ri, fi] for co in range(ncorrs): if have_flags and not flags_match[ri, fi, co]: continue wt = (weight_spectrum[ri, fi, co] if have_weight_spectrum else weight[ri, co] if have_weight else 1.0) ssv = sigma_spectrum[ri, fi, co]**2 * wt**2 sigma_spectrum_avg[ro, co] += ssv sigma_spectrum_weight_sum[ro, co] += wt # Normalise for ro in range(out_rows): for co in range(ncorrs): if sigma_spectrum_weight_sum[ro, co] != 0.0: ssv = sigma_spectrum_avg[ro, co] sswsum = sigma_spectrum_weight_sum[ro, co] sigma_spectrum_avg[ro, co] = np.sqrt(ssv / sswsum**2) return RowChanAverageOutput(vis_avg, flag_avg, weight_spectrum_avg, sigma_spectrum_avg) return impl