Skip to content

Commit 530dc2d

Browse files
scpmwscpmw
committed
Completed conversion to CompressedVisibility
All tests run. Will rename soon. It gets confusing as to whether a CompressedVisibility is compressed or not!
1 parent 75589d4 commit 530dc2d

24 files changed

Lines changed: 467 additions & 710 deletions

‎arl/data/data_models.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,10 @@ def w(self):
410410
def time(self):
411411
return self.data['time']
412412

413+
@property
414+
def integration_time(self):
415+
return self.data['integration_time']
416+
413417
@property
414418
def frequency(self):
415419
return self.data['frequency']
@@ -419,9 +423,9 @@ def channel_width(self):
419423
return self.data['channel_width']
420424

421425
@property
422-
def integration_time(self):
423-
return self.data['integration_time']
424-
426+
def polarisation(self):
427+
return self.data['polarisation']
428+
425429
@property
426430
def antenna1(self):
427431
return self.data['antenna1']

‎arl/fourier_transforms/convolutional_gridding.py‎

Lines changed: 104 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -298,7 +298,7 @@ def frac_coord(npixel, kernel_oversampling, p):
298298
:param kernel_oversampling: Fractional values to round to
299299
:param p: Coordinate in range [-.5,.5[
300300
"""
301-
assert numpy.array(p >= -0.5).all() and numpy.array(p < 0.5).all(), "uv overflows grid uv= %s" % str(p)
301+
assert numpy.array(p >= -0.5).all() and numpy.array(p < 0.5).all(), "Cellsize is too large: uv overflows grid uv= %s" % str(p)
302302
x = npixel // 2 + p * npixel
303303
flx = numpy.floor(x + 0.5 / kernel_oversampling)
304304
fracx = numpy.around((x - flx) * kernel_oversampling)
@@ -318,48 +318,69 @@ def frac_coords(shape, kernel_oversampling, xycoords):
318318
return x, xf, y, yf
319319

320320

321-
def fixed_kernel_degrid(kernel, vshape, uvgrid, uv, frequency, polarisation, uvscale, vpolmap = lambda pol: pol,
322-
vchanmap = lambda frequency: 0):
321+
def fixed_kernel_degrid(kernels, vshape, uvgrid, vuvwmap, vfrequencymap, vpolarisationmap):
323322
"""Convolutional degridding with frequency and polarisation independent
324323
325324
Takes into account fractional `uv` coordinate values where the GCF
326325
is oversampled
327326
328-
:param kernel: Oversampled convolution kernel
327+
:param kernel: list of oversampled convolution kernel
329328
:param vshape: Shape of visibility
330329
:param uvgrid: The uv plane to de-grid from
331-
:param uv: fractional uv coordinates in range[-0.5,0.5[
332-
:param uvscale: scaling for each channel
330+
:param uvw: fractional uvw coordinates in range[-0.5,0.5[
331+
:param vuvwmap: function to map uvw to grid fractions
332+
:param vfrequencymap: function to map frequency to image channels
333+
:param vpolarisationwmap: function to map polarisation to image polarisation
333334
:returns: Array of visibilities.
334335
"""
335-
kernel_oversampling, _, gh, gw = kernel.shape
336+
kernels = numpy.conjugate(list(kernels))
337+
kernel_oversampling, _, gh, gw = kernels[0].shape
336338
assert gh % 2 == 0, "Convolution kernel must have even number of pixels"
337339
assert gw % 2 == 0, "Convolution kernel must have even number of pixels"
338340
inchan, inpol, ny, nx = uvgrid.shape
339-
nvis, vnchan, vnpol = vshape
340-
assert vnpol == inpol, "Number of polarizations must be the same"
341+
nvis = vshape[0]
341342
vis = numpy.zeros([nvis], dtype='complex')
342343
wt = numpy.zeros([nvis])
343-
344-
ichan = vchanmap(frequency)
345-
ipol = vpolmap(polarisation)
346-
# Calculate grid point and fractional part for scaled u and v
347-
y, yf = frac_coord(uvgrid.shape[2], kernel_oversampling, uvscale[1] * uv[..., 1])
348-
x, xf = frac_coord(uvgrid.shape[3], kernel_oversampling, uvscale[0] * uv[..., 0])
344+
345+
# uvw -> fraction of grid ampping
346+
# uvw -> fraction of grid mapping
347+
y, yf = frac_coord(ny, kernel_oversampling, vuvwmap[:,1])
348+
x, xf = frac_coord(nx, kernel_oversampling, vuvwmap[:,0])
349+
349350
# Now calculate slices for the footprint of each sample
350351
slicey = []
351352
slicex = []
352353
for xx in x:
353354
slicex.append(slice((xx - gw // 2), (xx + (gw + 1) // 2)))
354355
for yy in y:
355356
slicey.append(slice((yy - gh // 2), (yy + (gh + 1) // 2)))
356-
coords = slicex, xf, slicey, yf
357+
357358
# Now we can degrid the points and accumulate weight as well
358-
vis[...] = [
359-
numpy.sum(uvgrid[ichan, ipol, sly, slx] * kernel[yf, xf, :, :])
360-
for slx, xf, sly, yf in zip(*coords)
361-
]
362-
wt[...] = [numpy.sum(kernel[yf, xf, :, :].real) for x, xf, y, yf in zip(*coords)]
359+
if len(kernels) > 1:
360+
361+
coords = list(vfrequencymap), list(vpolarisationmap), slicex, xf, slicey, yf
362+
vis[...] = [
363+
numpy.sum(uvgrid[ic, ip, sly, slx] * kernel[yf, xf, :, :])
364+
for kernel, ic, ip, slx, xf, sly, yf in zip(kernels, *coords)
365+
]
366+
367+
wt[...] = [
368+
numpy.sum(kernel[yf, xf, :, :].real)
369+
for kernel, ic, ip, slx, xf, sly, yf in zip(kernels, *coords)
370+
]
371+
else:
372+
kernel = kernels[0]
373+
374+
coords = list(vfrequencymap), list(vpolarisationmap), slicex, xf, slicey, yf
375+
vis[...] = [
376+
numpy.sum(uvgrid[ic, ip, sly, slx] * kernel[yf, xf, :, :])
377+
for ic, ip, slx, xf, sly, yf in zip(*coords)
378+
]
379+
wt[...] = [
380+
numpy.sum(kernel[yf, xf, :, :].real)
381+
for ic, ip, slx, xf, sly, yf in zip(*coords)
382+
]
383+
363384
vis[numpy.where(wt > 0)] = vis[numpy.where(wt > 0)] / wt[numpy.where(wt > 0)]
364385
vis[numpy.where(wt < 0)] = 0.0
365386
return numpy.array(vis)
@@ -381,7 +402,7 @@ def gridder(uvgrid, vis, xs, ys, kernel=numpy.ones((1,1)), kernel_ixs=None):
381402
"""
382403

383404
if kernel_ixs is None:
384-
kernel_ixs = numpy.zeroes((len(vis),0))
405+
kernel_ixs = numpy.zeros((len(vis),0))
385406
else:
386407
kernel_ixs = numpy.array(kernel_ixs)
387408
if len(kernel_ixs.shape) == 1:
@@ -392,134 +413,107 @@ def gridder(uvgrid, vis, xs, ys, kernel=numpy.ones((1,1)), kernel_ixs=None):
392413
uvgrid[y:y+gh, x:x+gw] += kernel[tuple(kern_ix)] * v
393414

394415

395-
def fixed_kernel_grid(kernel, uvgrid, uv, frequency, polarisation, uvscale, vis, visweights,
396-
vpolmap = lambda pol: pol, vchanmap = lambda frequency: 0):
416+
def fixed_kernel_grid(kernels, uvgrid, vis, visweights, vuvwmap, vfrequencymap,
417+
vpolarisationmap):
397418
"""Grid after convolving with frequency and polarisation independent gcf
398419
399420
Takes into account fractional `uv` coordinate values where the GCF
400421
is oversampled
401422
402-
:param kernel: Oversampled convolution kernel
423+
:param kernel: List of versampled convolution kernels
403424
:param uvgrid: Grid to add to
404-
:param uv: UVW positions
405-
:param uvscale: Scaling for each axis (u,v) for each channel
425+
:param uvw: UVW positions
406426
:param vis: CompressedVisibility values
407427
:param vis: CompressedVisibility weights
428+
:param vuvwmap: map uvw to grid fractions
429+
:param vfrequencymap: map frequency to image channels
430+
:param vpolarisationwmap: map polarisation to image polarisation
408431
:returns: uv grid[nchan, npol, ny, nx], sumwt[nchan, npol]
409432
"""
410433

411-
kernel_oversampling, _, gh, gw = kernel.shape
434+
kernels = list(kernels)
435+
kernel_oversampling, _, gh, gw = kernels[0].shape
412436
assert gh % 2 == 0, "Convolution kernel must have even number of pixels"
413437
assert gw % 2 == 0, "Convolution kernel must have even number of pixels"
414438
inchan, inpol, ny, nx = uvgrid.shape
415-
nvis, vnchan, vnpol = vis.shape
416-
assert vnpol == inpol, "Number of polarizations must be the same"
417-
wtgrid = numpy.zeros(uvgrid.shape, dtype='float')
439+
440+
# Construct output grids (in uv space)
418441
sumwt = numpy.zeros([inchan, inpol])
419-
ichan = vchanmap(frequency)
420-
ipol = vpolmap(polarisation)
421-
y, yf = frac_coord(ny, kernel_oversampling, uvscale[1] * uv[..., 1])
422-
x, xf = frac_coord(nx, kernel_oversampling, uvscale[0] * uv[..., 0])
442+
443+
# uvw -> fraction of grid mapping
444+
# vuvwmap is a list of generators, one per axis so we need to
445+
# select a generator and then instantiate it by making it into
446+
# a list, and finally use the values. This defers filling out
447+
# the arrays until they are needed.
448+
y, yf = frac_coord(ny, kernel_oversampling, vuvwmap[:,1])
449+
x, xf = frac_coord(nx, kernel_oversampling, vuvwmap[:,0])
450+
451+
# Construct all the slices that we will need: trade off memory to
452+
# simplify loop
423453
slicey = []
424454
slicex = []
425455
for xx in x:
426456
slicex.append(slice((xx - gw // 2), (xx + (gw + 1) // 2)))
427457
for yy in y:
428458
slicey.append(slice((yy - gh // 2), (yy + (gh + 1) // 2)))
429-
coords = slicex, xf, slicey, yf
459+
460+
# Now we can loop over all rows
430461
wts = visweights[...]
431462
viswt = vis[...] * visweights[...]
432-
coords = ichan, ipol, slicex, xf, slicey, yf
433-
for v, vwt, ic, ip, slx, xf, sly, yf in zip(viswt, wts, *coords):
434-
uvgrid[ic, ip, sly, slx] += kernel[yf, xf, :, :] * v
435-
wtgrid[ic, ip, sly, slx] += kernel[yf, xf, :, :].real * vwt
436-
sumwt[ichan, ipol] += numpy.sum(wtgrid[ichan, ipol, ...])
437-
438-
return uvgrid, sumwt
439-
463+
# There are two cases:row independent and row dependent kernel. We try to make these as
464+
# similar as possible
465+
if len(kernels) > 1:
466+
467+
coords = list(vfrequencymap), list(vpolarisationmap), slicex, xf, slicey, yf
468+
for v, vwt, kernel, ic, ip, slx, xf, sly, yf in zip(viswt, wts, kernels, *coords):
469+
uvgrid[ic, ip, sly, slx] += kernel[yf, xf, :, :] * v
470+
sumwt[ic, ip] += numpy.sum(kernel[yf, xf, :, :].real * vwt)
471+
else:
472+
kernel = kernels[0]
440473

441-
def box_grid(kernel, uvgrid, uv, frequency, polarisation, uvscale, vis, visweights, vpolmap = lambda pol: pol,
442-
vchanmap = lambda chan: chan):
443-
"""Grid with a box function
474+
coords = list(vfrequencymap), list(vpolarisationmap), slicex, xf, slicey, yf
475+
for v, vwt, ic, ip, slx, xf, sly, yf in zip(viswt, wts, *coords):
476+
uvgrid[ic, ip, sly, slx] += kernel[yf, xf, :, :] * v
477+
sumwt[ic, ip] += numpy.sum(kernel[yf, xf, :, :].real * vwt)
444478

445-
Takes into account fractional `uv` coordinate values where the GCF
446-
is oversampled
447-
448-
:param kernel: Oversampled convolution kernel
449-
:param uvgrid: Grid to add to
450-
:param uv: UVW positions
451-
:param uvscale: Scaling for each axis (u,v) for each channel
452-
:param vis: CompressedVisibility values
453-
:param vis: CompressedVisibility weights
454-
:param chanmap: function to map visibility channels to image channels
455-
"""
456-
457-
inchan, inpol, ny, nx = uvgrid.shape
458-
nvis, vnchan, vnpol = vis.shape
459-
assert vnpol == inpol, "Number of polarizations must be the same"
460-
wtgrid = numpy.zeros(uvgrid.shape, dtype='float')
461-
sumwt = numpy.zeros([inchan, inpol])
462-
ichan = vchanmap(frequency)
463-
ipol = vpolmap(polarisation)
464-
y, _ = frac_coord(ny, 1, uvscale[1] * uv[..., 1])
465-
x, _ = frac_coord(nx, 1, uvscale[0] * uv[..., 0])
466-
coords = frequency, polarisation, x, y
467-
wts = visweights[...]
468-
viswt = vis[...] * visweights[...]
469-
for v, vwt, ic, ip, x, y in zip(viswt, wts, *coords):
470-
uvgrid[ic, ip, y, x] += v
471-
wtgrid[ic, ip, y, x] += vwt
472-
sumwt = numpy.sum(wtgrid, axis=(2,3))
473-
474479
return uvgrid, sumwt
475480

476481

477-
def weight_gridding(shape, uv, uvscale, visweights, weighting='uniform',
478-
vpolmap = lambda pol: pol, vchanmap = lambda chan: chan):
482+
def weight_gridding(shape, visweights, vuvwmap, vfrequencymap, vpolarisationmap, weighting='uniform'):
479483
"""Reweight data using one of a number of algorithms
480484
481485
Takes into account fractional `uv` coordinate values where the GCF
482486
is oversampled
483487
484488
:param shape:
485-
:param uv: UVW positions
486-
:param uvscale: Scaling for each axis (u,v) for each channel
489+
:param uvw: UVW positions
487490
:param vis: CompressedVisibility values
488-
:param visweights: CompressedVisibility weights
489-
:param weighting: Weighting algorithm (natural|uniform) (uniform)
490-
:param vpolmap = lambda pol: pol, vchanmap: function to map visibility channels to image channels
491+
:param vis: CompressedVisibility weights
492+
:param vuvwmap: map uvw to grid fractions
493+
:param vfrequencymap: map frequency to image channels
494+
:param vpolarisationwmap: map polarisation to image polarisation
495+
:returns: visweights, density, desnitygrid
491496
"""
492-
densitygrid = numpy.zeros(shape)
497+
wtgrid = numpy.zeros(shape)
493498
density = numpy.zeros_like(visweights)
494499
if weighting == 'uniform':
495500
log.info("weight_gridding: Performing uniform weighting")
496501
inchan, inpol, ny, nx = shape
497-
nvis, vnchan, vnpol = visweights.shape
498-
# Add all visibility points to a float grid
499-
for vchan in range(vnchan):
500-
ichan = vpolmap = lambda pol: pol, vchanmap(vchan)
501-
y, _ = frac_coord(ny, 1, uvscale[1] * uv[..., 1])
502-
x, _ = frac_coord(nx, 1, uvscale[0] * uv[..., 0])
503-
coords = x, y
504-
for pol in range(vnpol):
505-
wts = visweights[..., vchan, pol]
506-
for wt, x, y, in zip(wts, *coords):
507-
densitygrid[ichan, pol, y, x] += wt
508-
# for i in range(visweights[..., chan, pol].shape[0]):
509-
# densitygrid[chan, pol, y[i], x[i]] += wts[i]
510-
502+
503+
# uvw -> fraction of grid mapping
504+
# uvw -> fraction of grid mapping
505+
y, yf = frac_coord(ny, 1.0, vuvwmap[:, 1])
506+
x, xf = frac_coord(nx, 1.0, vuvwmap[:, 0])
507+
wts = visweights[...]
508+
coords = list(vfrequencymap), list(vpolarisationmap), x, y
509+
for vwt, ic, ip, x, y in zip(wts, *coords):
510+
wtgrid[ic, ip, y, x] += vwt
511+
511512
# Normalise each visibility weight to sum to one in a grid cell
512513
newvisweights = numpy.zeros_like(visweights)
513-
for vchan in range(vnchan):
514-
ichan = vpolmap = lambda pol: pol, vchanmap(vchan)
515-
coords = frac_coord(shape[3], 1, uvscale[1] * uv[..., 1]), \
516-
frac_coord(shape[2], 1, uvscale[0] * uv[..., 0])
517-
for pol in range(vnpol):
518-
wts = visweights[..., vchan, pol]
519-
for wt, x, y in zip(wts, *coords):
520-
density[..., vchan, pol] += densitygrid[ichan, pol, y, x]
521-
newvisweights[..., vchan, pol][densitygrid[ichan, pol, y, x] > 0.0] = \
522-
wt / densitygrid[ichan, pol, y, x][densitygrid[ichan, pol, y, x] > 0.0]
523-
return newvisweights, density, densitygrid
514+
for wt, ic, ip, x, y in zip(wts, *coords):
515+
density[...] += wtgrid[ic, ip, y, x]
516+
newvisweights[...] = wt / wtgrid[ic, ip, y, x][wtgrid[ic, ip, y, x] > 0.0]
517+
return newvisweights, density, wtgrid
524518
else:
525-
return visweights, density, densitygrid
519+
return visweights, density, wtgrid

0 commit comments

Comments
 (0)