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
