def __init__(self,
template,
command_queue,
image_shape,
border,
allocator=None):
if image_shape[0] != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
if border >= 0.5:
raise ValueError('Border must be less than half the image size')
super().__init__(command_queue, allocator)
border_pixels = round(border * min(image_shape[1], image_shape[2]))
self._num_tiles_x = accel.divup(image_shape[2] - 2 * border_pixels,
template.tilex)
self._num_tiles_y = accel.divup(image_shape[1] - 2 * border_pixels,
template.tiley)
self.template = template
self.border_pixels = border_pixels
self.slots['dirty'] = accel.IOSlot([
accel.Dimension(template.num_polarizations, exact=True),
image_shape[1], image_shape[2]
], template.dtype)
self.slots['rank'] = accel.IOSlot([
accel.Dimension(self._num_tiles_y, exact=True),
accel.Dimension(self._num_tiles_x, exact=True)
], np.uint32)
self.kernel = template.program.get_kernel('compute_rank')
def __init__(self,
template,
command_queue,
loop_gain,
image_shape,
psf_shape,
allocator=None):
super().__init__(command_queue, allocator)
pol_dim = accel.Dimension(template.num_polarizations, exact=True)
if image_shape[0] != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
if psf_shape[0] != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
image_dims = [
pol_dim,
accel.Dimension(image_shape[1]),
accel.Dimension(image_shape[2])
]
psf_dims = [
pol_dim,
accel.Dimension(psf_shape[1]),
accel.Dimension(psf_shape[2])
]
self.slots['dirty'] = accel.IOSlot(image_dims, template.dtype)
self.slots['model'] = accel.IOSlot(image_dims, template.dtype)
self.slots['psf'] = accel.IOSlot(psf_dims, template.dtype)
self.slots['peak_pixel'] = accel.IOSlot([pol_dim], template.dtype)
self.loop_gain = loop_gain
self.template = template
self.kernel = template.program.get_kernel('subtract_psf')
def __init__(self,
template,
command_queue,
image_shape,
border,
allocator=None):
if image_shape[0] != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
if border >= 0.5:
raise ValueError('Border must be less than half the image size')
super().__init__(command_queue, allocator)
border_pixels = round(border * min(image_shape[1], image_shape[2]))
num_tiles_x = accel.divup(image_shape[2] - 2 * border_pixels,
template.tilex)
num_tiles_y = accel.divup(image_shape[1] - 2 * border_pixels,
template.tiley)
image_width = accel.Dimension(image_shape[2])
image_height = accel.Dimension(image_shape[1])
image_pols = accel.Dimension(template.num_polarizations, exact=True)
tiles_width = accel.Dimension(num_tiles_x)
tiles_height = accel.Dimension(num_tiles_y)
self.template = template
self.border_pixels = border_pixels
self.slots['dirty'] = accel.IOSlot(
[image_pols, image_height, image_width], template.dtype)
self.slots['tile_max'] = accel.IOSlot([tiles_height, tiles_width],
template.dtype)
self.slots['tile_pos'] = accel.IOSlot(
[tiles_height, tiles_width,
accel.Dimension(2, exact=True)], np.int32)
self.kernel = template.program.get_kernel('update_tiles')
def __init__(self, template, command_queue, shape, allocator=None):
super().__init__(command_queue, allocator)
self.template = template
if len(shape) != 3:
raise ValueError('Wrong number of dimensions in shape')
if shape[0] != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
self.slots['src'] = accel.IOSlot(shape, template.dtype)
self.slots['dest'] = accel.IOSlot(shape, template.dtype)
self.kernel = template.program.get_kernel('add_image')
def __init__(self, template, command_queue, grid_shape, allocator=None):
super().__init__(command_queue, allocator)
self.template = template
self.slots['grid'] = accel.IOSlot(
(grid_shape[0], accel.Dimension(grid_shape[1],
self.template.wgs_y),
accel.Dimension(grid_shape[2], self.template.wgs_x)), np.float32)
self.slots['sums'] = accel.IOSlot((2, ), np.float32)
self._kernel = template.program.get_kernel('mean_weight')
self._sums_host = accel.HostArray((2, ),
np.float32,
context=self.template.context)
def __init__(self, template, command_queue, shape, threshold, replacement, allocator=None):
super().__init__(command_queue, allocator)
self.template = template
if len(shape) != 3:
raise ValueError('Wrong number of dimensions in shape')
if shape[0] != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
y_dim = accel.Dimension(shape[1])
x_dim = accel.Dimension(shape[2])
self.slots['data'] = accel.IOSlot((shape[0], y_dim, x_dim), template.dtype)
self.slots['beam_power'] = accel.IOSlot((y_dim, x_dim), template.dtype)
self.kernel = template.program.get_kernel('apply_primary_beam')
self.threshold = threshold
self.replacement = replacement
def __init__(self, template, command_queue, grid_shape, allocator=None):
super().__init__(command_queue, allocator)
self.template = template
if grid_shape[0] != self.template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
self.a = 1.0
self.b = 0.0
self.slots['grid'] = accel.IOSlot(
(grid_shape[0], accel.Dimension(grid_shape[1],
self.template.wgs_y),
accel.Dimension(grid_shape[2], self.template.wgs_x)), np.float32)
self.slots['sums'] = accel.IOSlot((3, ), np.float32)
self._kernel = self.template.program.get_kernel('density_weights')
self._sums_host = accel.HostArray((3, ),
np.float32,
context=self.template.context)
def __init__(self, template, command_queue, shape, allocator=None):
if shape[0] != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
super().__init__(command_queue, allocator)
max_wg_x = accel.divup(shape[2], template.wgsx)
max_wg_y = accel.divup(shape[1], template.wgsy)
polarizations = accel.Dimension(template.num_polarizations, exact=True)
self.slots['psf'] = accel.IOSlot([polarizations, shape[1], shape[2]],
template.dtype)
self.slots['bound'] = accel.IOSlot([
accel.Dimension(max_wg_y * max_wg_x),
accel.Dimension(2, exact=True)
], np.int32)
self._bound_host = accel.HostArray((max_wg_y * max_wg_x, 2),
np.int32,
context=command_queue.context)
self.template = template
self.kernel = template.program.get_kernel('psf_patch')
def __init__(self, template, command_queue, shape, lm_scale, lm_bias, kernel_name,
allocator=None):
if len(shape) != 3 or shape[-1] != shape[-2]:
raise ValueError('shape must be square, not {}'.format(shape))
if shape[-1] % 2 != 0:
raise ValueError('image size must be even, not {}'.format(shape[-1]))
super().__init__(command_queue, allocator)
self.template = template
complex_dtype = katsdpimager.types.real_to_complex(template.real_dtype)
dims = [accel.Dimension(x) for x in shape]
self.slots['layer'] = accel.IOSlot(dims[-2:], complex_dtype)
self.slots['image'] = accel.IOSlot(dims, template.real_dtype)
self.slots['kernel1d'] = accel.IOSlot((shape[-1],), template.real_dtype)
self.kernel = template.program.get_kernel(kernel_name)
self.kernel_name = kernel_name
self.lm_scale = lm_scale
self.lm_bias = lm_bias
self.w = 0
self.polarization = 0
def __init__(self, template, command_queue, image_shape, allocator=None):
if len(image_shape) != 2:
raise ValueError('image_shape must be 2D')
super().__init__(command_queue, allocator=allocator)
self.template = template
self.image_shape = image_shape
output_shape = (image_shape[0], image_shape[1] // 2 + 1)
complex_type = katsdpimager.types.real_to_complex(template.dtype)
self.slots['data'] = accel.IOSlot(output_shape, complex_type)
self.beam = None
self._kernel = template.program.get_kernel('fourier_beam')
def __init__(self,
template,
command_queue,
image_shape,
tile_shape,
allocator=None):
if image_shape[0] != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
super().__init__(command_queue, allocator)
self.template = template
image_dims = [
accel.Dimension(image_shape[0], exact=True),
accel.Dimension(image_shape[1]),
accel.Dimension(image_shape[2])
]
tile_dims = [
accel.Dimension(tile_shape[0]),
accel.Dimension(tile_shape[1])
]
pair = accel.Dimension(2, exact=True)
self.slots['dirty'] = accel.IOSlot(image_dims, template.dtype)
self.slots['tile_max'] = accel.IOSlot(tile_dims, template.dtype)
self.slots['tile_pos'] = accel.IOSlot(tile_dims + [pair], np.int32)
self.slots['peak_value'] = accel.IOSlot([1], template.dtype)
self.slots['peak_pos'] = accel.IOSlot([2], np.int32)
self.slots['peak_pixel'] = accel.IOSlot([template.num_polarizations],
template.dtype)
self.kernel = template.program.get_kernel('find_peak')
def __init__(self,
template,
command_queue,
grid_shape,
max_vis,
allocator=None):
super().__init__(command_queue, allocator)
self.template = template
if grid_shape[0] != self.template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
if grid_shape[1] % 2 or grid_shape[2] % 2:
raise ValueError('Odd-sized grid not currently supported')
self.max_vis = max_vis
self.slots['grid'] = accel.IOSlot(grid_shape, np.float32)
self.slots['uv'] = accel.IOSlot(
(max_vis, accel.Dimension(4, exact=True)), np.int16)
self.slots['weights'] = accel.IOSlot(
(max_vis,
accel.Dimension(self.template.num_polarizations, exact=True)),
np.float32)
self._num_vis = 0
self._kernel = self.template.program.get_kernel('grid_weights')
def __init__(self, template, command_queue, image_parameters, grid_parameters,
max_vis, max_sources, allocator=None):
if len(image_parameters.fixed.polarizations) != template.num_polarizations:
raise ValueError('Mismatch in number of polarizations')
super().__init__(command_queue, template.num_polarizations, max_vis,
allocator)
self.template = template
pol_dim = accel.Dimension(template.num_polarizations, exact=True)
sources_dim = max(1, max_sources) # Cannot allocate 0-byte buffer
self.slots['lmn'] = accel.IOSlot((sources_dim, accel.Dimension(3, exact=True)), np.float32)
self.slots['flux'] = accel.IOSlot((sources_dim, pol_dim), np.float32)
self.slots['weights'] = accel.IOSlot(
(max_vis, accel.Dimension(template.num_polarizations, exact=True)), np.float32)
self._kernel = self.template.program.get_kernel('predict')
self._num_sources = 0
self.max_sources = max_sources
self.image_parameters = image_parameters
self.grid_parameters = grid_parameters
self._w = 0.0
self._host_lmn = accel.HostArray(
(max_sources, 3), np.float32, context=command_queue.context)
self._host_flux = accel.HostArray(
(max_sources, template.num_polarizations), np.float32, context=command_queue.context)
self._transfer_event = None
def __init__(self, template, command_queue, shape_grid,
lm_scale, lm_bias, fft_plan, allocator=None):
polarizations = shape_grid[0]
shape_image = (polarizations,) + tuple(fft_plan.shape)
self._fft = fft_plan.instantiate(command_queue, fft.FftMode.FORWARD, allocator)
self._image_to_layer = template.image_to_layer.instantiate(
command_queue, shape_image, lm_scale, lm_bias, allocator)
operations = [
('image_to_layer', self._image_to_layer),
('fft', self._fft),
]
compounds = {
'layer': ['fft:src', 'fft:dest', 'image_to_layer:layer'],
'image': ['image_to_layer:image'],
'kernel1d': ['image_to_layer:kernel1d']
}
super().__init__(command_queue, operations, compounds, allocator=allocator)
self.slots['grid'] = accel.IOSlot(shape_grid, fft_plan.dtype_dest)
def __init__(self, template, command_queue, shape_grid,
lm_scale, lm_bias, fft_plan, allocator=None):
self._ifft = fft_plan.instantiate(command_queue, fft.FftMode.INVERSE, allocator)
polarizations = shape_grid[0]
shape_image = (polarizations,) + tuple(fft_plan.shape)
self._layer_to_image = template.layer_to_image.instantiate(
command_queue, shape_image, lm_scale, lm_bias, allocator)
operations = [
('ifft', self._ifft),
('layer_to_image', self._layer_to_image)
]
compounds = {
'layer': ['ifft:src', 'ifft:dest', 'layer_to_image:layer'],
'image': ['layer_to_image:image'],
'kernel1d': ['layer_to_image:kernel1d']
}
super().__init__(command_queue, operations, compounds, allocator=allocator)
self.slots['grid'] = accel.IOSlot(shape_grid, fft_plan.dtype_src)
