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)