def convert(inputs, input_schema, output_schema):
(kernel, block, in_shape, out_shape,
dtype) = _generate_kernel(inputs, input_schema, output_schema)
# Flatten non-schema input dimensions,
# from inspection of the cupy reshape code,
# this incurs a copy when inputs is non-contiguous
nsrc = reduce(mul, inputs.shape[:-len(in_shape)], 1)
nelems = reduce(mul, in_shape, 1)
rinputs = inputs.reshape(nsrc, nelems)
assert rinputs.flags.c_contiguous
grid = grids((nsrc, 1, 1), block)
outputs = cp.empty(shape=rinputs.shape, dtype=dtype)
try:
kernel(grid, block, (rinputs, outputs))
except CompileException:
log.exception(format_code(kernel.code))
raise
shape = inputs.shape[:-len(in_shape)] + out_shape
outputs = outputs.reshape(shape)
assert outputs.flags.c_contiguous
return outputs
def _generate_kernel(inputs, input_schema, output_schema):
mapping, in_shape, out_shape, out_dtype = stokes_convert_setup(
inputs, input_schema, output_schema)
# Flatten input and output shapes
# Check that number elements are the same
in_elems = reduce(mul, in_shape, 1)
out_elems = reduce(mul, out_shape, 1)
if in_elems != out_elems:
raise ValueError("Number of input_schema elements %s "
"and output schema elements %s "
"must match for CUDA kernel." % (in_shape, out_shape))
# Infer the output data type
if out_dtype == "real":
if np.iscomplexobj(inputs):
out_dtype = inputs.real.dtype
else:
out_dtype = inputs.dtype
elif out_dtype == "complex":
if np.iscomplexobj(inputs):
out_dtype = inputs.dtype
else:
out_dtype = np.result_type(inputs.dtype, np.complex64)
else:
raise ValueError("Invalid setup dtype %s" % out_dtype)
cuda_out_dtype = cuda_type(out_dtype)
assign_exprs = []
# Render the assignment expression for each element
for (c1, c1i), (c2, c2i), outi, template_fn in mapping:
# Flattened indices
flat_outi = np.ravel_multi_index(outi, out_shape)
render = jinja_env.from_string(template_fn).render
kwargs = {
c1: "in[%d]" % np.ravel_multi_index(c1i, in_shape),
c2: "in[%d]" % np.ravel_multi_index(c2i, in_shape),
"out_type": cuda_out_dtype
}
expr_str = render(**kwargs)
assign_exprs.append("out[%d] = %s;" % (flat_outi, expr_str))
# Now render the main template
render = jinja_env.get_template(_TEMPLATE_PATH).render
name = "stokes_convert"
code = render(kernel_name=name,
input_type=cuda_type(inputs.dtype),
output_type=cuda_type(out_dtype),
assign_exprs=assign_exprs,
elements=in_elems).encode("utf-8")
# cuda block, flatten non-schema dims into a single source dim
blockdimx = 512
block = (blockdimx, 1, 1)
return (cp.RawKernel(code, name), block, in_shape, out_shape, out_dtype)
def degrid(grids, uvw, weights, ref_wave,
convolution_filter, w_bins, cell_size,
dtype=np.complex64):
"""
Convolutional W-stacking degridder (continuum)
Variable numbers of correlations are supported.
* :code:`(ny, nx, corr_1, corr_2)` ``grid`` will result in a
:code:`(row, chan, corr_1, corr_2)` ``vis``
* :code:`(ny, nx, corr_1)` ``grid`` will result in a
:code:`(row, chan, corr_1)` ``vis``
Parameters
----------
grids : list of np.ndarray
list of visibility grids of length :code:`nw`.
of shape :code:`(ny, nx, corr_1, corr_2)`
uvw : np.ndarray
float64 array of UVW coordinates of shape :code:`(row, 3)`
weights : np.ndarray
float32 or float64 array of weights. Set this to
``np.ones_like(vis, dtype=np.float32)`` as default.
ref_wave : np.ndarray
float64 array of wavelengths of shape :code:`(chan,)`
convolution_filter : :class:`~africanus.filters.ConvolutionFilter`
Convolution Filter
w_bins : :class:`numpy.ndarray`
W coordinate bins of shape :code:`(nw + 1,)`
cell_size : float
Cell size in arcseconds.
dtype : :class:`numpy.dtype`, optional
Numpy type of the resulting array. Defaults to
:class:`numpy.complex64`.
Returns
-------
np.ndarray
:code:`(row, chan, corr_1, corr_2)` complex ndarray of visibilities
"""
corrs = grids[0].shape[2:]
nrow = uvw.shape[0]
nchan = ref_wave.shape[0]
# Flatten the correlation dimensions
flat_corrs = reduce(mul, corrs)
# Create output visibilities
vis = np.empty((nrow, nchan, flat_corrs), dtype=dtype)
grids = [g.reshape(g.shape[0:2] + (flat_corrs,)) for g in grids]
numba_degrid(grids, uvw, weights, ref_wave, convolution_filter,
w_bins, cell_size, vis)
return vis.reshape((nrow, nchan) + corrs)
def degrid(grid, uvw, weights, ref_wave,
convolution_filter, cell_size, dtype=np.complex64):
"""
Convolutional degridder (continuum)
Variable numbers of correlations are supported.
* :code:`(ny, nx, corr_1, corr_2)` ``grid`` will result in a
:code:`(row, chan, corr_1, corr_2)` ``vis``
* :code:`(ny, nx, corr_1)` ``grid`` will result in a
:code:`(row, chan, corr_1)` ``vis``
Parameters
----------
grid : np.ndarray
float or complex grid of visibilities
of shape :code:`(ny, nx, corr_1, corr_2)`
uvw : np.ndarray
float64 array of UVW coordinates of shape :code:`(row, 3)`
in wavelengths.
weights : np.ndarray
float32 or float64 array of weights. Set this to
``np.ones_like(vis, dtype=np.float32)`` as default.
ref_wave : np.ndarray
float64 array of wavelengths of shape :code:`(chan,)`
convolution_filter : :class:`~africanus.filters.ConvolutionFilter`
Convolution Filter
cell_size : float
Cell size in arcseconds.
dtype : :class:`numpy.dtype`
Data type of the visibilities
Returns
-------
np.ndarray
:code:`(row, chan, corr_1, corr_2)` complex ndarray of visibilities
"""
nrow = uvw.shape[0]
nchan = ref_wave.shape[0]
corrs = flat_corrs = grid.shape[2:]
# Flatten if necessary
if len(corrs) > 1:
flat_corrs = (reduce(mul, corrs),)
grid = grid.reshape(grid.shape[:2] + flat_corrs)
weights = weights.reshape(weights.shape[:2] + flat_corrs)
vis = np.zeros((nrow, nchan) + flat_corrs, dtype=dtype)
vis = numba_degrid(grid, uvw, weights, ref_wave,
convolution_filter, cell_size, vis)
return vis.reshape(weights.shape[:2] + corrs)
def stream_reduction(time_index, antenna1, antenna2, dde1_jones, source_coh,
dde2_jones, predict_check_tup, out_dtype, streams):
"""
Reduces source coherencies + ddes over the source dimension in
``N`` parallel streams.
This is accomplished by calling predict_vis on on ddes and source
coherencies to produce visibilities which are passed into
the `base_vis` argument of ``predict_vis`` for the next chunk.
"""
# Unique name and token for this operation
token = tokenize(time_index, antenna1, antenna2, dde1_jones, source_coh,
dde2_jones, streams)
name = 'stream-coherency-reduction-' + token
# Number of dim blocks
blocks = _extract_blocks(time_index, dde1_jones, source_coh, dde2_jones)
(src_blocks, row_blocks, _,
chan_blocks), corr_blocks = blocks[:4], blocks[4:]
# Total number of other dimension blocks
nblocks = reduce(mul, (row_blocks, chan_blocks) + corr_blocks, 1)
# Create the compressed mapping
layers = CoherencyStreamReduction(time_index, antenna1, antenna2,
dde1_jones, source_coh, dde2_jones, name,
streams)
# Create the graph
extra_deps = [
a for a in (dde1_jones, source_coh, dde2_jones) if a is not None
]
deps = [time_index, antenna1, antenna2] + extra_deps
graph = HighLevelGraph.from_collections(name, layers, deps)
chunks = ((1, ) * src_blocks, (1, ) * nblocks)
# This should never be directly computed, reported chunks
# and dtype don't match the actual data. We create it
# because it makes chaining HighLevelGraphs easier
stream_reduction = da.Array(graph, name, chunks, dtype=np.int8)
name = "coherency-reduction-" + tokenize(stream_reduction)
layers = CoherencyFinalReduction(name, layers)
graph = HighLevelGraph.from_collections(name, layers, [stream_reduction])
chunks = _extract_chunks(time_index, dde1_jones, source_coh, dde2_jones)
return da.Array(graph, name, chunks[1:], dtype=out_dtype)
def feed_rotation(parallactic_angles, feed_type='linear'):
if feed_type == 'linear':
poltype = 0
elif feed_type == 'circular':
poltype = 1
else:
raise ValueError("Invalid feed_type '%s'" % feed_type)
if parallactic_angles.dtype == np.float32:
dtype = np.complex64
elif parallactic_angles.dtype == np.float64:
dtype = np.complex128
else:
raise ValueError("parallactic_angles has "
"none-floating point type %s" %
parallactic_angles.dtype)
# Create result array with flattened parangles
shape = (reduce(mul, parallactic_angles.shape), ) + (2, 2)
result = np.empty(shape, dtype=dtype)
return _nb_feed_rotation(parallactic_angles, poltype, result)
def grid(vis, uvw, flags, weights, ref_wave,
convolution_filter,
cell_size,
nx=1024, ny=1024,
grid=None):
"""
Convolutional gridder which grids visibilities ``vis``
at the specified ``uvw`` coordinates and
``ref_wave`` reference wavelengths using
the specified ``convolution_filter``.
Variable numbers of correlations are supported.
* :code:`(row, chan, corr_1, corr_2)` ``vis`` will result in a
:code:`(ny, nx, corr_1, corr_2)` ``grid``.
* :code:`(row, chan, corr_1)` ``vis`` will result in a
:code:`(ny, nx, corr_1)` ``grid``.
Parameters
----------
vis : np.ndarray
complex visibility array of shape :code:`(row, chan, corr_1, corr_2)`
uvw : np.ndarray
float64 array of UVW coordinates of shape :code:`(row, 3)`
in wavelengths.
weights : np.ndarray
float32 or float64 array of weights. Set this to
``np.ones_like(vis, dtype=np.float32)`` as default.
flags : np.ndarray
flagged array of shape :code:`(row, chan, corr_1, corr_2)`.
Any positive quantity will indicate that the corresponding
visibility should be flagged.
Set to ``np.zeros_like(vis, dtype=np.bool)`` as default.
ref_wave : np.ndarray
float64 array of wavelengths of shape :code:`(chan,)`
convolution_filter : :class:`~africanus.filters.ConvolutionFilter`
Convolution filter
cell_size : float
Cell size in arcseconds.
nx : integer, optional
Size of the grid's X dimension
ny : integer, optional
Size of the grid's Y dimension
grid : np.ndarray, optional
complex64/complex128 array of shape :code:`(ny, nx, corr_1, corr_2)`
If supplied, this array will be used as the gridding target,
and ``nx`` and ``ny`` will be derived from this grid's
dimensions.
Returns
-------
np.ndarray
:code:`(ny, nx, corr_1, corr_2)` complex ndarray of
gridded visibilities. The number of correlations may vary,
depending on the shape of vis.
"""
# Flatten the correlation dimensions
corrs = vis.shape[2:]
flat_corrs = (reduce(mul, corrs),)
# Create grid of flatten correlations or reshape
if grid is None:
grid = np.zeros((ny, nx) + flat_corrs, dtype=vis.dtype)
else:
ny, nx = grid.shape[0:2]
grid = grid.reshape((ny, nx) + flat_corrs)
return numba_grid(vis, uvw, flags, weights, ref_wave,
convolution_filter, cell_size, grid)
def grid(vis, uvw, flags, weights, ref_wave,
convolution_filter, w_bins,
cell_size,
nx=1024, ny=1024,
grids=None):
"""
Convolutional W-stacking gridder.
This function grids visibilities ``vis`` onto multiple
grids, each associated with a W-layer defined by ``w_bins``.
The W coordinate of the ``uvw`` array is used to bin the visibility
into the appropriate grid.
Variable numbers of correlations are supported.
* :code:`(row, chan, corr_1, corr_2)` ``vis`` will result in a
:code:`(ny, nx, corr_1, corr_2)` ``grid``.
* :code:`(row, chan, corr_1)` ``vis`` will result in a
:code:`(ny, nx, corr_1)` ``grid``.
Parameters
----------
vis : :class:`numpy.ndarray`
complex visibility array of shape :code:`(row, chan, corr_1, corr_2)`
uvw : :class:`numpy.ndarray`
float64 array of UVW coordinates of shape :code:`(row, 3)`
weights : :class:`numpy.ndarray`
float32 or float64 array of weights. Set this to
``np.ones_like(vis, dtype=np.float32)`` as default.
flags : np.ndarray
flagged array of shape :code:`(row, chan, corr_1, corr_2)`.
Any positive quantity will indicate that the corresponding
visibility should be flagged.
Set to ``np.zeros_like(vis, dtype=np.bool)`` as default.
ref_wave : np.ndarray
float64 array of wavelengths of shape :code:`(chan,)`
convolution_filter : :class:`~africanus.filters.ConvolutionFilter`
Convolution filter
w_bins : :class:`numpy.ndarray`
W coordinate bins of shape :code:`(nw + 1,)`
cell_size : float
Cell size in arcseconds.
nx : integer, optional
Size of the grid's X dimension
ny : integer, optional
Size of the grid's Y dimension
grids : list of np.ndarray, optional
list of complex arrays of length :code:`nw`,
each with shape :code:`(ny, nx, corr_1, corr_2)`.
If supplied, this array will be used as the gridding target,
and ``nx`` and ``ny`` will be derived from the grid's
dimensions.
Returns
-------
list of np.ndarray
list of complex arrays of gridded visibilities, of length :code:`nw`,
each with shape :code:`(ny, nx, corr_1, corr_2)`.
The number of correlations may vary,
depending on the shape of vis.
"""
corrs = vis.shape[2:]
# Create grid of flatten correlations or reshape
if grids is None:
nw = w_bins.shape[0] - 1
grids = [np.zeros((ny, nx) + corrs, dtype=vis.dtype)
for _ in range(nw)]
elif not isinstance(grids, list):
grids = [grids]
# Flatten the correlation dimensions
flat_corrs = (reduce(mul, corrs),)
grids = [g.reshape(g.shape[0:2] + flat_corrs) for g in grids]
return numba_grid(vis, uvw, flags, weights, ref_wave,
convolution_filter, w_bins, cell_size, grids)
def __len__(self):
(source, row, _, chan), corrs = self.blocks[:4], self.blocks[4:]
return reduce(mul, (source, row, chan) + corrs, 1)
def __len__(self):
# Extract dimension blocks
(source, row, _, chan), corr = self.blocks[:4], self.blocks[4:]
return reduce(mul, (source, row, chan) + corr, 1)
def _generate_kernel(time_index, antenna1, antenna2, dde1_jones, source_coh,
dde2_jones, die1_jones, base_vis, die2_jones, corrs,
out_ndim):
tup = predict_checks(time_index, antenna1, antenna2, dde1_jones,
source_coh, dde2_jones, die1_jones, base_vis,
die2_jones)
(have_ddes1, have_coh, have_ddes2, have_dies1, have_bvis, have_dies2) = tup
# Check types
if time_index.dtype != np.int32:
raise TypeError("time_index.dtype != np.int32 '%s'" % time_index.dtype)
if antenna1.dtype != np.int32:
raise TypeError("antenna1.dtype != np.int32 '%s'" % antenna1.dtype)
if antenna2.dtype != np.int32:
raise TypeError("antenna2.dtype != np.int32 '%s'" % antenna2.dtype)
# Create template
render = jinja_env.get_template(_TEMPLATE_PATH).render
name = "predict_vis"
# Complex output type
out_dtype = np.result_type(dde1_jones, source_coh, dde2_jones, die1_jones,
base_vis, die2_jones)
ncorrs = reduce(mul, corrs, 1)
# corrs x channels, rows
blockdimx = 32
blockdimy = 24 if out_dtype == np.complex128 else 32
block = (blockdimx, blockdimy, 1)
code = render(kernel_name=name,
blockdimx=blockdimx,
blockdimy=blockdimy,
have_dde1=have_ddes1,
dde1_type=cuda_type(dde1_jones) if have_ddes1 else "int",
dde1_ndim=dde1_jones.ndim if have_ddes1 else 1,
have_dde2=have_ddes2,
dde2_type=cuda_type(dde2_jones) if have_ddes2 else "int",
dde2_ndim=dde2_jones.ndim if have_ddes2 else 1,
have_coh=have_coh,
coh_type=cuda_type(source_coh) if have_coh else "int",
coh_ndim=source_coh.ndim if have_coh else 1,
have_die1=have_dies1,
die1_type=cuda_type(die1_jones) if have_dies1 else "int",
die1_ndim=die1_jones.ndim if have_dies1 else 1,
have_base_vis=have_bvis,
base_vis_type=cuda_type(base_vis) if have_bvis else "int",
base_vis_ndim=base_vis.ndim if have_bvis else 1,
have_die2=have_dies2,
die2_type=cuda_type(die2_jones) if have_dies2 else "int",
die2_ndim=die2_jones.ndim if have_dies2 else 1,
out_type=cuda_type(out_dtype),
corrs=ncorrs,
out_ndim=out_ndim,
warp_size=32).encode('utf-8')
return cp.RawKernel(code, name), block, out_dtype
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):
""" Cupy implementation of the feed_rotation kernel. """
have_ddes = dde1_jones is not None and dde2_jones is not None
have_dies = die1_jones is not None and die2_jones is not None
have_coh = source_coh is not None
have_bvis = base_vis is not None
# Infer the output shape
if have_ddes:
row = time_index.shape[0]
chan = dde1_jones.shape[3]
corrs = dde1_jones.shape[4:]
elif have_coh:
row = time_index.shape[0]
chan = source_coh.shape[2]
corrs = source_coh.shape[3:]
elif have_dies:
row = time_index.shape[0]
chan = die1_jones.shape[2]
corrs = die1_jones.shape[3:]
elif have_bvis:
row = time_index.shape[0]
chan = base_vis.shape[1]
corrs = base_vis.shape[2:]
else:
raise ValueError("Insufficient inputs supplied for determining "
"the output shape")
ncorrs = len(corrs)
# Flatten correlations
if ncorrs == 2:
flat_corrs = reduce(mul, corrs, 1)
if have_ddes:
dde_shape = dde1_jones.shape[:-ncorrs] + (flat_corrs, )
dde1_jones = dde1_jones.reshape(dde_shape)
dde2_jones = dde2_jones.reshape(dde_shape)
if have_coh:
coh_shape = source_coh.shape[:-ncorrs] + (flat_corrs, )
source_coh = source_coh.reshape(coh_shape)
if have_dies:
die_shape = die1_jones.shape[:-ncorrs] + (flat_corrs, )
die1_jones = die1_jones.reshape(die_shape)
die2_jones = die2_jones.reshape(die_shape)
if have_bvis:
bvis_shape = base_vis.shape[:-ncorrs] + (flat_corrs, )
base_vis = base_vis.reshape(bvis_shape)
elif ncorrs == 1:
flat_corrs = corrs[0]
else:
raise ValueError("Invalid correlation setup %s" % (corrs, ))
out_shape = (row, chan) + (flat_corrs, )
kernel, block, out_dtype = _generate_kernel(time_index, antenna1, antenna2,
dde1_jones, source_coh,
dde2_jones, die1_jones,
base_vis, die2_jones, corrs,
len(out_shape))
grid = grids((chan * flat_corrs, row, 1), block)
out = cp.empty(shape=out_shape, dtype=out_dtype)
# Normalise the time index
# TODO(sjperkins)
# Normalise the time index with a device-wide reduction
norm_time_index = time_index - time_index.min()
args = (norm_time_index, antenna1, antenna2, dde1_jones, source_coh,
dde2_jones, die1_jones, base_vis, die2_jones, out)
try:
kernel(grid, block, tuple(a for a in args if a is not None))
except CompileException:
log.exception(format_code(kernel.code))
raise
return out.reshape((row, chan) + corrs)
