From c08bac64061b177bfd2307a6f83ad6f83642a26a Mon Sep 17 00:00:00 2001 From: Benjamin Hugo Date: Wed, 23 Sep 2020 10:52:12 +0200 Subject: [PATCH 01/12] Revert "Revert perley polyhedron gridder (#220)" This reverts commit ae8ae8e6410133e99e1abed80de7582d9dc66b19. --- HISTORY.rst | 2 + .../gridding/perleypolyhedron/__init__.py | 0 africanus/gridding/perleypolyhedron/dask.py | 319 ++++++++ .../gridding/perleypolyhedron/degridder.py | 265 +++++++ .../gridding/perleypolyhedron/gridder.py | 116 +++ .../gridding/perleypolyhedron/kernels.py | 178 +++++ .../perleypolyhedron/policies/__init__.py | 0 .../policies/baseline_transform_policies.py | 95 +++ .../policies/convolution_policies.py | 284 +++++++ .../policies/phase_transform_policies.py | 74 ++ .../policies/stokes_conversion_policies.py | 262 ++++++ .../perleypolyhedron/tests/__init__.py | 0 .../perleypolyhedron/tests/test_daskintrf.py | 367 +++++++++ .../perleypolyhedron/tests/test_ppgridder.py | 743 ++++++++++++++++++ africanus/util/numba.py | 2 +- setup.py | 4 +- 16 files changed, 2709 insertions(+), 2 deletions(-) create mode 100644 africanus/gridding/perleypolyhedron/__init__.py create mode 100644 africanus/gridding/perleypolyhedron/dask.py create mode 100644 africanus/gridding/perleypolyhedron/degridder.py create mode 100644 africanus/gridding/perleypolyhedron/gridder.py create mode 100644 africanus/gridding/perleypolyhedron/kernels.py create mode 100644 africanus/gridding/perleypolyhedron/policies/__init__.py create mode 100644 africanus/gridding/perleypolyhedron/policies/baseline_transform_policies.py create mode 100644 africanus/gridding/perleypolyhedron/policies/convolution_policies.py create mode 100644 africanus/gridding/perleypolyhedron/policies/phase_transform_policies.py create mode 100644 africanus/gridding/perleypolyhedron/policies/stokes_conversion_policies.py create mode 100644 africanus/gridding/perleypolyhedron/tests/__init__.py create mode 100644 africanus/gridding/perleypolyhedron/tests/test_daskintrf.py create mode 100644 africanus/gridding/perleypolyhedron/tests/test_ppgridder.py diff --git a/HISTORY.rst b/HISTORY.rst index 50706ac31..368439da4 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -6,6 +6,8 @@ History ------------------ * Add wrapper for ducc0.wgridder (:pr:204`) * Correct Irregular Grid nesting in BeamAxes (:pr:`217`) +* Add Perley Polyhedron Faceting Gridder/Degridder (:pr:`202`, :pr:`215`) + 0.2.5 (2020-07-01) ------------------ diff --git a/africanus/gridding/perleypolyhedron/__init__.py b/africanus/gridding/perleypolyhedron/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/africanus/gridding/perleypolyhedron/dask.py b/africanus/gridding/perleypolyhedron/dask.py new file mode 100644 index 000000000..ac4899af1 --- /dev/null +++ b/africanus/gridding/perleypolyhedron/dask.py @@ -0,0 +1,319 @@ +import numpy as np + +try: + import dask.array as da +except ImportError as e: + opt_import_err = e +else: + opt_import_err = None + +from africanus.gridding.perleypolyhedron.gridder import ( + gridder as np_gridder) +from africanus.gridding.perleypolyhedron.degridder import ( + degridder as np_degridder) +from africanus.gridding.perleypolyhedron.degridder import ( + degridder_serial as np_degridder_serial) +from africanus.gridding.perleypolyhedron.policies import ( + stokes_conversion_policies) +from africanus.util.requirements import requires_optional + + +def __degrid(uvw, + gridstack, + lambdas, + chanmap, + image_centres, + phase_centre, + cell=None, + convolution_kernel=None, + convolution_kernel_width=None, + convolution_kernel_oversampling=None, + baseline_transform_policy=None, + phase_transform_policy=None, + stokes_conversion_policy=None, + convolution_policy=None, + vis_dtype=np.complex128, + rowparallel=False): + image_centres = image_centres[0][0] + if image_centres.ndim != 2: + raise ValueError( + "Image centres for DASK wrapper expects list of image centres, " + "one per facet in radec radians" + ) + if image_centres.shape[1] != 2: + raise ValueError("Image centre must be a list of tuples") + uvw = uvw[0] + if uvw.ndim != 2 or uvw.shape[1] != 3: + raise ValueError("UVW array must be nrow x 3") + gridstack = gridstack[0][0][0][0] + if gridstack.ndim != 4: + raise ValueError("Gridstack must be nfacet x nband x ny x nx") + lambdas = lambdas + chanmap = chanmap + if chanmap.size != lambdas.size: + raise ValueError( + "Chanmap and corresponding lambdas must match in shape") + nchan = lambdas.size + nrow = uvw.shape[0] + ncorr = stokes_conversion_policies.ncorr_outpy( + policy_type=stokes_conversion_policy)() + vis = np.zeros((nrow, nchan, ncorr), dtype=vis_dtype) + degridcall = np_degridder_serial if not rowparallel else np_degridder + for fi, f in enumerate(image_centres): + # add contributions from all facets + vis[:, :, :] += \ + degridcall(uvw, + gridstack[fi, :, :, :], + lambdas, + chanmap, + cell, + image_centres, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + vis_dtype=vis_dtype) + return vis + + +@requires_optional("dask", opt_import_err) +def degridder(uvw, + gridstack, + lambdas, + chanmap, + cell, + image_centres, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + vis_dtype=np.complex128, + rowparallel=False): + """ + 2D Convolutional degridder, discrete to contiguous + @uvw: value coordinates, (nrow, 3) + @gridstack: complex gridded data, (nband, npix, npix) + @lambdas: wavelengths of data channels + @chanmap: MFS band mapping per channel + @cell: cell_size in degrees + @image_centre: new phase centre of image (radians, ra, dec) + @phase_centre: original phase centre of data (radians, ra, dec) + @convolution_kernel: packed kernel as generated by kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled points in kernel + @baseline_transform_policy: any accepted policy in + .policies.baseline_transform_policies, + can be used to tilt image planes for + polyhedron faceting + @phase_transform_policy: any accepted policy in + .policies.phase_transform_policies, + can be used to facet at provided + facet @image_centre + @stokes_conversion_policy: any accepted correlation to stokes + conversion policy in + .policies.stokes_conversion_policies + @convolution_policy: any accepted convolution policy in + .policies.convolution_policies + @vis_dtype: accumulation vis dtype (default complex 128) + @rowparallel: adds additional threading per row per chunk. This may be + necessary for cases where there are few facets and few chunks + to get optimal performance. Requires TBB to be installed + from your distribution package management system. See numba + documentation + http://numba.pydata.org/numba-doc/0.46.0/user/threading-layer.html + Must set: + from numba import config # have to be openmp to support + nested parallelism config.THREADING_LAYER = 'threadsafe' + before calling this function + """ + if image_centres.ndim != 2: + raise ValueError( + "Image centres for DASK wrapper expects list of " + "image centres, one per facet in radec radians" + ) + if image_centres.shape[1] != 2: + raise ValueError("Image centre must be a list of tuples") + if gridstack.ndim != 4 or gridstack.shape[0] != image_centres.shape[0]: + raise ValueError( + "Grid stack must be nfacet x nband x yy x xx and match number " + "of image centres" + ) + vis = da.blockwise( + __degrid, ("row", "chan", "corr"), + uvw, ("row", "uvw"), + gridstack, ("nfacet", "nband", "y", "x"), + lambdas, ("chan", ), + chanmap, ("chan", ), + image_centres, ("nfacet", "coord"), + convolution_kernel=convolution_kernel, + convolution_kernel_width=convolution_kernel_width, + convolution_kernel_oversampling=convolution_kernel_oversampling, + baseline_transform_policy=baseline_transform_policy, + phase_transform_policy=phase_transform_policy, + stokes_conversion_policy=stokes_conversion_policy, + convolution_policy=convolution_policy, + cell=cell, + phase_centre=phase_centre, + vis_dtype=vis_dtype, + # goes to one set of grids per row chunk + adjust_chunks={"row": 1}, + new_axes={ + "corr": + stokes_conversion_policies.ncorr_outpy( + policy_type=stokes_conversion_policy)() + }, + dtype=vis_dtype, + meta=np.empty( + (0, 0, 0), + dtype=vis_dtype) # row, chan, correlation product as per MSv2 spec + ) + return vis + + +def __grid(uvw, + vis, + image_centres, + lambdas=None, + chanmap=None, + convolution_kernel=None, + convolution_kernel_width=None, + convolution_kernel_oversampling=None, + baseline_transform_policy=None, + phase_transform_policy=None, + stokes_conversion_policy=None, + convolution_policy=None, + npix=None, + cell=None, + phase_centre=None, + grid_dtype=np.complex128, + do_normalize=False): + image_centres = image_centres[0] + if image_centres.ndim != 2: + raise ValueError( + "Image centres for DASK wrapper expects list of image " + "centres, one per facet in radec radians" + ) + if image_centres.shape[1] != 2: + raise ValueError("Image centre must be a list of tuples") + + uvw = uvw[0] + vis = vis[0][0] + lambdas = lambdas[0] + chanmap = chanmap[0] + grid_stack = np.zeros( + (1, image_centres.shape[0], 1, np.max(chanmap) + 1, npix, npix), + dtype=grid_dtype) + for fi, f in enumerate(image_centres): + grid_stack[0, fi, 0, :, :, :] = \ + np_gridder(uvw, vis, lambdas, chanmap, npix, cell, f, phase_centre, + convolution_kernel, convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, phase_transform_policy, + stokes_conversion_policy, + convolution_policy, grid_dtype, do_normalize) + return grid_stack + + +@requires_optional("dask", opt_import_err) +def gridder(uvw, + vis, + lambdas, + chanmap, + npix, + cell, + image_centres, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + grid_dtype=np.complex128, + do_normalize=False): + """ + 2D Convolutional gridder, contiguous to discrete + @uvw: value coordinates, (nrow, 3) + @vis: complex data, (nrow, nchan, ncorr) + @lambdas: wavelengths of data channels + @chanmap: MFS band mapping + @npix: number of pixels per axis + @cell: cell_size in degrees + @image_centres: new phase centres of images (nfacet, (radians ra, dec)) + @phase_centre: original phase centre of data (radians, ra, dec) + @convolution_kernel: packed kernel as generated by kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled points in kernel + @baseline_transform_policy: any accepted policy in + .policies.baseline_transform_policies, + can be used to tilt image planes for + polyhedron faceting + @phase_transform_policy: any accepted policy in + .policies.phase_transform_policies, + can be used to facet at provided + facet @image_centre + @stokes_conversion_policy: any accepted correlation to stokes + conversion policy in + .policies.stokes_conversion_policies + @convolution_policy: any accepted convolution policy in + .policies.convolution_policies + @grid_dtype: accumulation grid dtype (default complex 128) + @do_normalize: normalize grid by convolution weights + """ + if len(vis.chunks) != 3 or lambdas.chunks[0] != vis.chunks[1]: + raise ValueError( + "Visibility frequency chunking does not match " + "lambda frequency chunking" + ) + if len(vis.chunks) != 3 or chanmap.chunks[0] != vis.chunks[1]: + raise ValueError( + "Visibility frequency chunking does not match chanmap " + "frequency chunking" + ) + if len(vis.chunks) != 3 or len( + uvw.chunks) != 2 or vis.chunks[0] != uvw.chunks[0]: + raise ValueError( + "Visibility row chunking does not match uvw row chunking") + grids = da.blockwise( + __grid, ("row", "nfacet", "nstokes", "nband", "y", "x"), + uvw, ("row", "uvw"), + vis, ("row", "chan", "corr"), + image_centres, ("nfacet", "coord"), + lambdas, ("chan", ), + chanmap, ("chan", ), + convolution_kernel=convolution_kernel, + convolution_kernel_width=convolution_kernel_width, + convolution_kernel_oversampling=convolution_kernel_oversampling, + baseline_transform_policy=baseline_transform_policy, + phase_transform_policy=phase_transform_policy, + stokes_conversion_policy=stokes_conversion_policy, + convolution_policy=convolution_policy, + npix=npix, + cell=cell, + phase_centre=phase_centre, + grid_dtype=grid_dtype, + do_normalize=do_normalize, + # goes to one set of grids per row chunk + adjust_chunks={"row": 1}, + new_axes={ + "nband": np.max(chanmap) + 1, + # for now will need to be modified if + # multi-stokes cubes are supported + "nstokes": 1, + "y": npix, + "x": npix + }, + dtype=grid_dtype, + meta=np.empty((0, 0, 0, 0, 0, 0), dtype=grid_dtype)) + + # Parallel reduction over row dimension + return grids.mean(axis=0, split_every=2) diff --git a/africanus/gridding/perleypolyhedron/degridder.py b/africanus/gridding/perleypolyhedron/degridder.py new file mode 100644 index 000000000..bdf3a0d30 --- /dev/null +++ b/africanus/gridding/perleypolyhedron/degridder.py @@ -0,0 +1,265 @@ +import numpy as np +from numba import literally, prange +from africanus.util.numba import jit + +from africanus.gridding.perleypolyhedron.policies import ( + baseline_transform_policies as btp) +from africanus.gridding.perleypolyhedron.policies import ( + phase_transform_policies as ptp) +from africanus.gridding.perleypolyhedron.policies import ( + convolution_policies as cp) +from africanus.gridding.perleypolyhedron.policies import ( + stokes_conversion_policies as scp) + + +@jit(nopython=True, nogil=True, fastmath=True, inline="always") +def degridder_row_kernel(uvw, + gridstack, + lambdas, + chanmap, + cell, + image_centre, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + vis_dtype=np.complex128, + nband=0, + nrow=0, + npix=0, + nvischan=0, + ncorr=0, + vis=None, + scale_factor=0, + r=0): + ra0, dec0 = phase_centre + ra, dec = image_centre + btp.policy(uvw[r, :], ra, dec, ra0, dec0, + literally(baseline_transform_policy)) + for c in range(nvischan): + scaled_u = uvw[r, 0] * scale_factor / lambdas[c] + scaled_v = uvw[r, 1] * scale_factor / lambdas[c] + scaled_w = uvw[r, 2] * scale_factor / lambdas[c] + grid = gridstack[chanmap[c], :, :] + cp.policy(scaled_u, + scaled_v, + scaled_w, + npix, + grid, + vis, + r, + c, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + stokes_conversion_policy, + policy_type=literally(convolution_policy)) + ptp.policy(vis[r, :, :], + uvw[r, :], + lambdas, + ra0, + dec0, + ra, + dec, + policy_type=literally(phase_transform_policy), + phasesign=-1.0) + + +@jit(nopython=True, nogil=True, fastmath=True, parallel=True) +def degridder(uvw, + gridstack, + lambdas, + chanmap, + cell, + image_centre, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + vis_dtype=np.complex128): + """ + 2D Convolutional degridder, discrete to contiguous + @uvw: value coordinates, (nrow, 3) + @gridstack: complex gridded data, (nband, npix, npix) + @lambdas: wavelengths of data channels + @chanmap: MFS band mapping per channel + @cell: cell_size in degrees + @image_centres: new phase centre of image (radians, ra, dec) + @phase_centre: original phase centre of data (radians, ra, dec) + @convolution_kernel: packed kernel as generated by kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled points in kernel + @baseline_transform_policy: any accepted policy in + .policies.baseline_transform_policies, + can be used to tilt image planes for + polyhedron faceting + @phase_transform_policy: any accepted policy in + .policies.phase_transform_policies, + can be used to facet at provided + facet @image_centre + @stokes_conversion_policy: any accepted correlation to + stokes conversion policy in + .policies.stokes_conversion_policies + @convolution_policy: any accepted convolution policy in + .policies.convolution_policies + @vis_dtype: accumulation vis dtype (default complex 128) + """ + + if chanmap.size != lambdas.size: + raise ValueError( + "Chanmap and corresponding lambdas must match in shape") + chanmap = chanmap.ravel() + lambdas = lambdas.ravel() + nband = np.max(chanmap) + 1 + nrow = uvw.shape[0] + npix = gridstack.shape[1] + if gridstack.shape[1] != gridstack.shape[2]: + raise ValueError("Grid must be square") + nvischan = lambdas.size + ncorr = scp.ncorr_out(policy_type=literally(stokes_conversion_policy)) + if gridstack.shape[0] < nband: + raise ValueError( + "Not enough channel bands in grid stack to match mfs band mapping") + if uvw.shape[1] != 3: + raise ValueError("UVW array must be array of tripples") + if uvw.shape[0] != nrow: + raise ValueError( + "UVW array must have same number of rows as vis array") + if nvischan != lambdas.size: + raise ValueError("Chanmap must correspond to visibility channels") + + vis = np.zeros((nrow, nvischan, ncorr), dtype=vis_dtype) + + # scale the FOV using the simularity theorem + scale_factor = npix * cell / 3600.0 * np.pi / 180.0 + for r in prange(nrow): + degridder_row_kernel(uvw, + gridstack, + lambdas, + chanmap, + cell, + image_centre, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + vis_dtype=vis_dtype, + nband=nband, + nrow=nrow, + npix=npix, + nvischan=nvischan, + ncorr=ncorr, + vis=vis, + scale_factor=scale_factor, + r=r) + return vis + + +@jit(nopython=True, nogil=True, fastmath=True, parallel=False) +def degridder_serial(uvw, + gridstack, + lambdas, + chanmap, + cell, + image_centre, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + vis_dtype=np.complex128): + """ + 2D Convolutional degridder, discrete to contiguous + @uvw: value coordinates, (nrow, 3) + @gridstack: complex gridded data, (nband, npix, npix) + @lambdas: wavelengths of data channels + @chanmap: MFS band mapping per channel + @cell: cell_size in degrees + @image_centres: new phase centre of image (radians, ra, dec) + @phase_centre: original phase centre of data (radians, ra, dec) + @convolution_kernel: packed kernel as generated by kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled points in kernel + @baseline_transform_policy: any accepted policy in + .policies.baseline_transform_policies, + can be used to tilt image planes for + polyhedron faceting + @phase_transform_policy: any accepted policy in + .policies.phase_transform_policies, + can be used to facet at provided + facet @image_centre + @stokes_conversion_policy: any accepted correlation to stokes + conversion policy in + .policies.stokes_conversion_policies + @convolution_policy: any accepted convolution policy in + .policies.convolution_policies + @vis_dtype: accumulation vis dtype (default complex 128) + """ + + if chanmap.size != lambdas.size: + raise ValueError( + "Chanmap and corresponding lambdas must match in shape") + chanmap = chanmap.ravel() + lambdas = lambdas.ravel() + nband = np.max(chanmap) + 1 + nrow = uvw.shape[0] + npix = gridstack.shape[1] + if gridstack.shape[1] != gridstack.shape[2]: + raise ValueError("Grid must be square") + nvischan = lambdas.size + ncorr = scp.ncorr_out(policy_type=literally(stokes_conversion_policy)) + if gridstack.shape[0] < nband: + raise ValueError( + "Not enough channel bands in grid stack to match mfs band mapping") + if uvw.shape[1] != 3: + raise ValueError("UVW array must be array of tripples") + if uvw.shape[0] != nrow: + raise ValueError( + "UVW array must have same number of rows as vis array") + if nvischan != lambdas.size: + raise ValueError("Chanmap must correspond to visibility channels") + + vis = np.zeros((nrow, nvischan, ncorr), dtype=vis_dtype) + + # scale the FOV using the simularity theorem + scale_factor = npix * cell / 3600.0 * np.pi / 180.0 + for r in range(nrow): + degridder_row_kernel(uvw, + gridstack, + lambdas, + chanmap, + cell, + image_centre, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + vis_dtype=vis_dtype, + nband=nband, + nrow=nrow, + npix=npix, + nvischan=nvischan, + ncorr=ncorr, + vis=vis, + scale_factor=scale_factor, + r=r) + return vis diff --git a/africanus/gridding/perleypolyhedron/gridder.py b/africanus/gridding/perleypolyhedron/gridder.py new file mode 100644 index 000000000..a28946414 --- /dev/null +++ b/africanus/gridding/perleypolyhedron/gridder.py @@ -0,0 +1,116 @@ +import numpy as np +from numba import literally + +from africanus.util.numba import jit +from africanus.gridding.perleypolyhedron.policies import ( + baseline_transform_policies as btp) +from africanus.gridding.perleypolyhedron.policies import ( + phase_transform_policies as ptp) +from africanus.gridding.perleypolyhedron.policies import ( + convolution_policies as cp) + + +@jit(nopython=True, nogil=True, fastmath=True, parallel=False) +def gridder(uvw, + vis, + lambdas, + chanmap, + npix, + cell, + image_centre, + phase_centre, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + baseline_transform_policy, + phase_transform_policy, + stokes_conversion_policy, + convolution_policy, + grid_dtype=np.complex128, + do_normalize=False): + """ + 2D Convolutional gridder, contiguous to discrete + @uvw: value coordinates, (nrow, 3) + @vis: complex data, (nrow, nchan, ncorr) + @lambdas: wavelengths of data channels + @chanmap: MFS band mapping + @npix: number of pixels per axis + @cell: cell_size in degrees + @image_centre: new phase centre of image (radians, ra, dec) + @phase_centre: original phase centre of data (radians, ra, dec) + @convolution_kernel: packed kernel as generated by kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled points in kernel + @baseline_transform_policy: any accepted policy in + .policies.baseline_transform_policies, + can be used to tilt image planes for + polyhedron faceting + @phase_transform_policy: any accepted policy in + .policies.phase_transform_policies, + can be used to facet at provided + facet @image_centre + @stokes_conversion_policy: any accepted correlation to stokes + conversion policy in + .policies.stokes_conversion_policies + @convolution_policy: any accepted convolution policy in + .policies.convolution_policies + @grid_dtype: accumulation grid dtype (default complex 128) + @do_normalize: normalize grid by convolution weights + """ + if chanmap.size != lambdas.size: + raise ValueError( + "Chanmap and corresponding lambdas must match in shape") + chanmap = chanmap.ravel() + lambdas = lambdas.ravel() + nband = np.max(chanmap) + 1 + nrow, nvischan, ncorr = vis.shape + if uvw.shape[1] != 3: + raise ValueError("UVW array must be array of tripples") + if uvw.shape[0] != nrow: + raise ValueError( + "UVW array must have same number of rows as vis array") + if nvischan != lambdas.size: + raise ValueError("Chanmap must correspond to visibility channels") + + gridstack = np.zeros((nband, npix, npix), dtype=grid_dtype) + + # scale the FOV using the simularity theorem + scale_factor = npix * cell / 3600.0 * np.pi / 180.0 + wt_ch = np.zeros(nband, dtype=np.float64) + for r in range(nrow): + ra0, dec0 = phase_centre + ra, dec = image_centre + ptp.policy(vis[r, :, :], + uvw[r, :], + lambdas, + ra0, + dec0, + ra, + dec, + policy_type=literally(phase_transform_policy), + phasesign=1.0) + btp.policy(uvw[r, :], ra0, dec0, ra, dec, + literally(baseline_transform_policy)) + for c in range(nvischan): + scaled_u = uvw[r, 0] * scale_factor / lambdas[c] + scaled_v = uvw[r, 1] * scale_factor / lambdas[c] + scaled_w = uvw[r, 2] * scale_factor / lambdas[c] + grid = gridstack[chanmap[c], :, :] + wt_ch[chanmap[c]] += cp.policy( + scaled_u, + scaled_v, + scaled_w, + npix, + grid, + vis, + r, + c, + convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + stokes_conversion_policy, + policy_type=literally(convolution_policy)) + if do_normalize: + for c in range(nband): + gridstack[c, :, :] /= wt_ch[c] + 1.0e-8 + return gridstack diff --git a/africanus/gridding/perleypolyhedron/kernels.py b/africanus/gridding/perleypolyhedron/kernels.py new file mode 100644 index 000000000..9fee2e304 --- /dev/null +++ b/africanus/gridding/perleypolyhedron/kernels.py @@ -0,0 +1,178 @@ +import numpy as np +from numba import prange + +try: + from scipy.special import jn +except ImportError as e: + # https://stackoverflow.com/a/29268974/1611416, pep 3110 and 344 + scipy_import_error = e +else: + scipy_import_error = None + +from africanus.util.numba import jit, register_jitable +from africanus.util.requirements import requires_optional + + +@register_jitable +def uspace(W, oversample): + """ + Generates a kernel sampling of the form + 0 1 2 3 4 + |...|...|...|...|... + |<----->| + half-support (W / 2) + | x W + . x (oversample - 1) x W + where W is odd + plus padding by 1 unit on either side + """ + # must be odd so that the taps can be centred at the origin + assert W % 2 == 1 + taps = np.arange( + oversample * + (W + 2)) / float(oversample) - (W + 2) // 2 + # (|+.) * W centred at 0 + return taps + + +def sinc(W, oversample=5, a=1.0): + """ + Basic oversampled sinc window + """ + u = uspace(W, oversample) + res = np.sinc(u * a) + return res / np.sum(res) + + +@requires_optional('scipy', scipy_import_error) +def kbsinc(W, b=None, oversample=5, order=15): + """ + Modified keiser bessel windowed sinc (Jackson et al., + IEEE transactions on medical imaging, 1991) + with a modification of higher order bessels as default, as + this improves the kernel at low number of taps. + """ + autocoeffs = np.polyfit( + [1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0], + [1.9980, 2.3934, 3.3800, 4.2054, 4.9107, 5.7567, 6.6291, 7.4302], 1) + if b is None: + b = np.poly1d(autocoeffs)((W + 2)) + + u = uspace(W, oversample) + wnd = jn(order, b * np.sqrt(1 - (2 * u / + ((W + 2) + 1))**2)) * 1 / ((W + 2) + 1) + res = sinc(W, oversample=oversample) * wnd * np.sum(wnd) + return res / np.sum(res) + + +def hanningsinc(W, a=0.5, oversample=5): + """ + Basic hanning windowed sinc + """ + autocoeffs = np.polyfit([1.5, 2.0, 2.5, 3.0, 3.5], + [0.7600, 0.7146, 0.6185, 0.5534, 0.5185], 3) + if a is None: + a = np.poly1d(autocoeffs)((W + 2)) + u = uspace(W, oversample) + wnd = a + (1 - a) * np.cos(2 * np.pi / ((W + 2) + 1) * u) + res = sinc(W, oversample=oversample) * wnd + return res / np.sum(res) + + +def pack_kernel(K, W, oversample=5): + """ + Repacks kernel to be cache coherent + Expects sampled kernel of the form + |...|...|...|...|... + |<----->| + half-support (W / 2) + | x W + . x (oversample - 1) x W + """ + pkern = np.empty(oversample * (W + 2), dtype=K.dtype) + for t in range(oversample): + pkern[t * (W + 2):(t + 1) * (W + 2)] = K[t::oversample] + return pkern + + +def unpack_kernel(K, W, oversample=5): + """ + Unpacks kernel to original sampling form (non-cache coherent) + Produces sampled kernel of the form + |...|...|...|...|... + |<----->| + half-support (W / 2) + | x W + . x (oversample - 1) x W + """ + upkern = np.empty(oversample * (W + 2), dtype=K.dtype) + for t in range(oversample): + upkern[t::oversample] = K[t * (W + 2):(t + 1) * (W + 2)] + return upkern + + +def compute_detaper(npix, K, W, oversample=5): + """ + Computes detapering function of a oversampled kernel + using a memory intensive FFT and the simularity theorem + Assumes a 2D square kernel to be passed as argument K + """ + pk = np.zeros((npix * oversample, npix * oversample)) + pk[npix * oversample // 2 - K.shape[0] // 2:npix * oversample // 2 - + K.shape[0] // 2 + K.shape[0], + npix * oversample // 2 - K.shape[1] // 2:npix * oversample // 2 - + K.shape[1] // 2 + K.shape[1]] = K + fpk = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(pk))) + fk = fpk[npix * oversample // 2 - npix // 2:npix * oversample // 2 - + npix // 2 + npix, npix * oversample // 2 - + npix // 2:npix * oversample // 2 - npix // 2 + npix] + return np.abs(fk) + + +@jit(nopython=True, nogil=True, fastmath=True, cache=True) +def compute_detaper_dft(npix, K, W, oversample=5): + """ + Computes detapering function of a oversampled kernel + using a memory non-intensive DFT sampled on a grid the + size of the square image + Assumes a 2D square kernel to be passed as argument K + """ + pk = np.zeros((npix, npix), dtype=np.complex128) + ksample = uspace(W, oversample=oversample) + rK = K.ravel() + + for p in prange(npix * npix): + ll = p % npix + mm = p // npix + llN = (ll - npix // 2) / float(npix) + mmN = (mm - npix // 2) / float(npix) + for x in range(K.size): + xx = ksample[x % K.shape[1]] + yy = ksample[x // K.shape[1]] + pk[mm, ll] += rK[x] * np.exp(-2.0j * np.pi * + (llN * xx + mmN * yy)) + return np.abs(pk) + + +@jit(nopython=True, nogil=True, fastmath=True, cache=True) +def compute_detaper_dft_seperable(npix, K, W, oversample=5): + """ + Computes detapering function of a oversampled seperable kernel + using a memory non-intensive DFT sampled on a grid the size + of the square image + Assumes a 1D kernel to be passed as argument K + + The outer product of K with itself can be evalated as + F(outer(K,K))[l,m] = F(K)[l].F(K)[m] + """ + pkX = np.zeros((npix), dtype=np.complex128) + ksample = uspace(W, oversample=oversample) + rK = K.ravel() + + for ll in range(npix): + llN = (ll - npix // 2) / float(npix) + for x in range(K.size): + xx = ksample[x] + pkX[ll] += rK[x] * np.exp(-2.0j * np.pi * (llN * xx)) + + return np.abs(np.outer(pkX, pkX)) diff --git a/africanus/gridding/perleypolyhedron/policies/__init__.py b/africanus/gridding/perleypolyhedron/policies/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/africanus/gridding/perleypolyhedron/policies/baseline_transform_policies.py b/africanus/gridding/perleypolyhedron/policies/baseline_transform_policies.py new file mode 100644 index 000000000..e09b9cc29 --- /dev/null +++ b/africanus/gridding/perleypolyhedron/policies/baseline_transform_policies.py @@ -0,0 +1,95 @@ +from africanus.util.numba import jit, overload +from numpy import cos, sin + + +def uvw_norotate(uvw, ra0, dec0, ra, dec, policy_type): + pass + + +def uvw_rotate(uvw, ra0, dec0, ra, dec, policy_type): + ''' + Compute the following 3x3 coordinate transformation matrix: + Z_rot(facet_new_rotation) * \\ + T(new_phase_centre_ra,new_phase_centre_dec) * \\ + transpose(T(old_phase_centre_ra, + old_phase_centre_dec)) * \\ + transpose(Z_rot(facet_old_rotation)) + where: + | cRA -sRA 0 | + T (RA,D) = | -sDsRA -sDcRA cD | + | cDsRA cDcRA sD | + This is the similar to the one in Thompson, A. R.; Moran, J. M.; + and Swenson, G. W., Jr. Interferometry and Synthesis + in Radio Astronomy, New York: Wiley, ch. 4, but in a + lefthanded system. + We're not transforming between a coordinate system with w pointing + towards the pole and one with w pointing towards the reference + centre here, so the last rotation matrix is ignored! + This transformation will let the image be tangent to the celestial + sphere at the new delay centre + ''' + d_ra = ra - ra0 + c_d_ra = cos(d_ra) + s_d_ra = sin(d_ra) + c_new_dec = cos(dec) + c_old_dec = cos(dec0) + s_new_dec = sin(dec) + s_old_dec = sin(dec0) + mat_11 = c_d_ra + mat_12 = s_old_dec * s_d_ra + mat_13 = -c_old_dec * s_d_ra + mat_21 = -s_new_dec * s_d_ra + mat_22 = s_new_dec * s_old_dec * c_d_ra + c_new_dec * c_old_dec + mat_23 = -c_old_dec * s_new_dec * c_d_ra + c_new_dec * s_old_dec + mat_31 = c_new_dec * s_d_ra + mat_32 = -c_new_dec * s_old_dec * c_d_ra + s_new_dec * c_old_dec + mat_33 = c_new_dec * c_old_dec * c_d_ra + s_new_dec * s_old_dec + uvw[0] = mat_11 * uvw[0] + mat_12 * uvw[1] + mat_13 * uvw[3] + uvw[1] = mat_21 * uvw[0] + mat_22 * uvw[1] + mat_23 * uvw[3] + uvw[2] = mat_31 * uvw[0] + mat_32 * uvw[1] + mat_33 * uvw[3] + + +@jit(nopython=True, nogil=True, fastmath=True, parallel=False) +def uvw_planarwapprox(uvw, ra0, dec0, ra, dec, policy_type): + ''' + Implements the coordinate uv transform associated with taking a planar + approximation to w(n-1) as described in Kogan & Greisen's AIPS Memo 113 + This is essentially equivalent to rotating the facet to be tangent to + the celestial sphere as Perley suggested to limit error, but it instead + takes w into account in a linear approximation to the phase error near + the facet centre. This keeps the facets parallel to the original facet + plane. Of course this 2D taylor expansion of the first order is only + valid over a small field of view, but that true of normal tilted + faceting as well. Only a convolution can get rid of the (n-1) + factor in the ME. + ''' + d_ra = ra - ra0 + n_dec = dec + o_dec = dec0 + c_d_ra = cos(d_ra) + s_d_ra = sin(d_ra) + c_new_dec = cos(n_dec) + c_old_dec = cos(o_dec) + s_new_dec = sin(n_dec) + s_old_dec = sin(o_dec) + li0 = c_new_dec * s_d_ra + mi0 = s_new_dec * c_old_dec - c_new_dec * s_old_dec * c_d_ra + ni0 = s_new_dec * s_old_dec + c_new_dec * c_old_dec * c_d_ra + uvw[0] = uvw[0] - uvw[2] * li0 / ni0 + uvw[1] = uvw[1] - uvw[2] * mi0 / ni0 + + +def policy(uvw, ra0, dec0, ra, dec, policy_type): + pass + + +@overload(policy, inline="always") +def policy_impl(uvw, ra0, dec0, ra, dec, policy_type): + if policy_type.literal_value == "None": + return uvw_norotate + elif policy_type.literal_value == "rotate": + return uvw_rotate + elif policy_type.literal_value == "wlinapprox": + return uvw_planarwapprox + else: + raise ValueError("Invalid baseline transform policy type") diff --git a/africanus/gridding/perleypolyhedron/policies/convolution_policies.py b/africanus/gridding/perleypolyhedron/policies/convolution_policies.py new file mode 100644 index 000000000..3bff1054c --- /dev/null +++ b/africanus/gridding/perleypolyhedron/policies/convolution_policies.py @@ -0,0 +1,284 @@ +from africanus.util.numba import overload +from numba import literally +import numpy as np +from . import stokes_conversion_policies as scp + + +def convolve_1d_axisymmetric_unpacked_scatter( + scaled_u, scaled_v, scaled_w, npix, grid, vis, r, c, + convolution_kernel, convolution_kernel_width, + convolution_kernel_oversampling, stokes_conversion_policy, + policy_type): + ''' + Convolution policy for a 1D axisymmetric unpacked + AA kernel (gridding kernel) + @scaled_u: simularity theorem and lambda scaled u + @scaled_v: simularity theorem and lambda scaled v + @scaled_w: simularity theorem and lambda scaled w + @npix: number of pixels per axis + @grid: 2d grid + @r: current row in the visibility array + @c: current channel in the visibility array + @convolution_kernel: packed kernel as generated by + kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled + points in kernel + @stokes_conversion_policy: any accepted correlation to stokes + conversion policy in + .policies.stokes_conversion_policies + ''' + offset_u = scaled_u + npix // 2 + offset_v = scaled_v + npix // 2 + disc_u = int(np.round(offset_u)) + disc_v = int(np.round(offset_v)) + frac_u = int((-offset_u + disc_u) * convolution_kernel_oversampling) + frac_v = int((-offset_v + disc_v) * convolution_kernel_oversampling) + cw = 0.0 + for tv in range(convolution_kernel_width): + conv_v = convolution_kernel[(tv + 1) * convolution_kernel_oversampling + + frac_v] + grid_v_lookup = disc_v + tv - convolution_kernel_width // 2 + for tu in range(convolution_kernel_width): + conv_u = convolution_kernel[(tu + 1) * + convolution_kernel_oversampling + + frac_u] + grid_u_lookup = disc_u + tu - convolution_kernel_width // 2 + if (grid_v_lookup >= 0 and grid_v_lookup < npix + and grid_u_lookup >= 0 and grid_u_lookup < npix): + grid[grid_v_lookup, grid_u_lookup] += \ + conv_v * conv_u * \ + scp.corr2stokes(vis[r, c, :], + literally(stokes_conversion_policy)) + cw += conv_v * conv_u + return cw + + +def convolve_1d_axisymmetric_packed_scatter( + scaled_u, scaled_v, scaled_w, npix, grid, vis, r, c, + convolution_kernel, convolution_kernel_width, + convolution_kernel_oversampling, stokes_conversion_policy, + policy_type): + ''' + Convolution policy for a 1D axisymmetric packed AA kernel (gridding kernel) + @scaled_u: simularity theorem and lambda scaled u + @scaled_v: simularity theorem and lambda scaled v + @scaled_w: simularity theorem and lambda scaled w + @npix: number of pixels per axis + @grid: 2d grid + @r: current row in the visibility array + @c: current channel in the visibility array + @convolution_kernel: packed kernel as generated by kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled points in kernel + @stokes_conversion_policy: any accepted correlation to stokes + conversion policy in + .policies.stokes_conversion_policies + ''' + offset_u = scaled_u + npix // 2 + offset_v = scaled_v + npix // 2 + disc_u = int(np.round(offset_u)) + disc_v = int(np.round(offset_v)) + frac_u = int((-offset_u + disc_u) * convolution_kernel_oversampling) + frac_v = int((-offset_v + disc_v) * convolution_kernel_oversampling) + frac_offset_u = 0 if frac_u < 0 else +1 + frac_offset_v = 0 if frac_v < 0 else +1 + # Two cases here: + # <--> padding before and + # <--> after + # <---> half support + # |...|...|...|...|... + # 0123456789ABCDEFGHIJ + # repacked into + # 048CG159DH26AEI37BFJ + # if fraction is negative we index 0 + frac * (support+2) + # else we index 1 + frac * (support+2) + # where frac wraps around to negative indexing + cw = 0.0 + for tv in range(convolution_kernel_width): + conv_v = convolution_kernel[tv + frac_offset_v + frac_v * + (convolution_kernel_width + 2)] + grid_v_lookup = disc_v + tv - convolution_kernel_width // 2 + for tu in range(convolution_kernel_width): + conv_u = convolution_kernel[tu + frac_offset_u + frac_u * + (convolution_kernel_width + 2)] + grid_u_lookup = disc_u + tu - convolution_kernel_width // 2 + if (grid_v_lookup >= 0 and grid_v_lookup < npix + and grid_u_lookup >= 0 and grid_u_lookup < npix): + grid[grid_v_lookup, grid_u_lookup] += \ + conv_v * conv_u * \ + scp.corr2stokes(vis[r, c, :], + literally(stokes_conversion_policy)) + cw += conv_v * conv_u + return cw + + +def convolve_nn_scatter(scaled_u, scaled_v, scaled_w, npix, grid, vis, r, c, + convolution_kernel, convolution_kernel_width, + convolution_kernel_oversampling, + stokes_conversion_policy, policy_type): + ''' + Convolution policy for a nn scatter kernel (gridding kernel) + @scaled_u: simularity theorem and lambda scaled u + @scaled_v: simularity theorem and lambda scaled v + @scaled_w: simularity theorem and lambda scaled w + @npix: number of pixels per axis + @grid: 2d grid + @r: current row in the visibility array + @c: current channel in the visibility array + @convolution_kernel: packed kernel as generated by + kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled + points in kernel + @stokes_conversion_policy: any accepted correlation to stokes + conversion policy in + .policies.stokes_conversion_policies + ''' + offset_u = scaled_u + npix // 2 + offset_v = scaled_v + npix // 2 + disc_u = int(np.round(offset_u)) + disc_v = int(np.round(offset_v)) + cw = 1.0 + grid[disc_v, disc_u] += \ + scp.corr2stokes(vis[r, c, :], + literally(stokes_conversion_policy)) + return cw + + +def convolve_1d_axisymmetric_packed_gather(scaled_u, scaled_v, scaled_w, npix, + grid, vis, r, c, convolution_kernel, + convolution_kernel_width, + convolution_kernel_oversampling, + stokes_conversion_policy, + policy_type): + ''' + Convolution policy for a 1D axisymmetric packed AA + kernel (degridding kernel) + @scaled_u: simularity theorem and lambda scaled u + @scaled_v: simularity theorem and lambda scaled v + @scaled_w: simularity theorem and lambda scaled w + @npix: number of pixels per axis + @grid: 2d grid + @r: current row in the visibility array + @c: current channel in the visibility array + @convolution_kernel: packed kernel as generated by + kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled points in kernel + @stokes_conversion_policy: any accepted correlation to + stokes conversion policy in + .policies.stokes_conversion_policies + ''' + offset_u = scaled_u + npix // 2 + offset_v = scaled_v + npix // 2 + disc_u = int(np.round(offset_u)) + disc_v = int(np.round(offset_v)) + frac_u = int((-offset_u + disc_u) * convolution_kernel_oversampling) + frac_v = int((-offset_v + disc_v) * convolution_kernel_oversampling) + frac_offset_u = 0 if frac_u < 0 else +1 + frac_offset_v = 0 if frac_v < 0 else +1 + # Two cases here: + # <--> padding before and + # <--> after + # <---> half support + # |...|...|...|...|... + # 0123456789ABCDEFGHIJ + # repacked into + # 048CG159DH26AEI37BFJ + # if fraction is negative we index 0 + frac * (support+2) + # else we index 1 + frac * (support+2) + # where frac wraps around to negative indexing + cw = 0 + for tv in range(convolution_kernel_width): + conv_v = convolution_kernel[tv + frac_offset_v + frac_v * + (convolution_kernel_width + 2)] + grid_v_lookup = disc_v + tv - convolution_kernel_width // 2 + for tu in range(convolution_kernel_width): + conv_u = convolution_kernel[tu + frac_offset_u + frac_u * + (convolution_kernel_width + 2)] + grid_u_lookup = disc_u + tu - convolution_kernel_width // 2 + if (grid_v_lookup >= 0 and grid_v_lookup < npix + and grid_u_lookup >= 0 and grid_u_lookup < npix): + scp.stokes2corr( + grid[disc_v + tv - convolution_kernel_width // 2, disc_u + + tu - convolution_kernel_width // 2] * conv_v * conv_u, + vis[r, c, :], + policy_type=literally(stokes_conversion_policy)) + cw += conv_v * conv_u + vis[r, c, :] /= cw + 1.0e-8 + + +def convolve_1d_axisymmetric_unpacked_gather( + scaled_u, scaled_v, scaled_w, npix, grid, vis, r, c, + convolution_kernel, convolution_kernel_width, + convolution_kernel_oversampling, stokes_conversion_policy, + policy_type): + ''' + Convolution policy for a 1D axisymmetric unpacked + AA kernel (degridding kernel) + @scaled_u: simularity theorem and lambda scaled u + @scaled_v: simularity theorem and lambda scaled v + @scaled_w: simularity theorem and lambda scaled w + @npix: number of pixels per axis + @grid: 2d grid + @r: current row in the visibility array + @c: current channel in the visibility array + @convolution_kernel: packed kernel as generated by kernels package + @convolution_kernel_width: number of taps in kernel + @convolution_kernel_oversampling: number of oversampled points + in kernel + @stokes_conversion_policy: any accepted correlation to stokes + conversion policy in + .policies.stokes_conversion_policies + ''' + offset_u = scaled_u + npix // 2 + offset_v = scaled_v + npix // 2 + disc_u = int(np.round(offset_u)) + disc_v = int(np.round(offset_v)) + frac_u = int((-offset_u + disc_u) * convolution_kernel_oversampling) + frac_v = int((-offset_v + disc_v) * convolution_kernel_oversampling) + cw = 0 + for tv in range(convolution_kernel_width): + conv_v = convolution_kernel[(tv + 1) * convolution_kernel_oversampling + + frac_v] + grid_v_lookup = disc_v + tv - convolution_kernel_width // 2 + for tu in range(convolution_kernel_width): + conv_u = convolution_kernel[(tu + 1) * + convolution_kernel_oversampling + + frac_u] + grid_u_lookup = disc_u + tu - convolution_kernel_width // 2 + if (grid_v_lookup >= 0 and grid_v_lookup < npix + and grid_u_lookup >= 0 and grid_u_lookup < npix): + scp.stokes2corr( + grid[grid_v_lookup, grid_u_lookup] * conv_v * conv_u, + vis[r, c, :], + policy_type=literally(stokes_conversion_policy)) + cw += conv_v * conv_u + vis[r, c, :] /= cw + 1.0e-8 + + +def policy(scaled_u, scaled_v, scaled_w, npix, grid, vis, r, c, + convolution_kernel, convolution_kernel_width, + convolution_kernel_oversampling, stokes_conversion_policy, + policy_type): + pass + + +@overload(policy, inline="always") +def policy_impl(scaled_u, scaled_v, scaled_w, npix, grid, vis, r, c, + convolution_kernel, convolution_kernel_width, + convolution_kernel_oversampling, stokes_conversion_policy, + policy_type): + if policy_type.literal_value == "conv_1d_axisymmetric_packed_scatter": + return convolve_1d_axisymmetric_packed_scatter + elif policy_type.literal_value == "conv_nn_scatter": + return convolve_nn_scatter + elif policy_type.literal_value == "conv_1d_axisymmetric_unpacked_scatter": + return convolve_1d_axisymmetric_unpacked_scatter + elif policy_type.literal_value == "conv_1d_axisymmetric_packed_gather": + return convolve_1d_axisymmetric_packed_gather + elif policy_type.literal_value == "conv_1d_axisymmetric_unpacked_gather": + return convolve_1d_axisymmetric_unpacked_gather + else: + raise ValueError("Invalid convolution policy type") diff --git a/africanus/gridding/perleypolyhedron/policies/phase_transform_policies.py b/africanus/gridding/perleypolyhedron/policies/phase_transform_policies.py new file mode 100644 index 000000000..8afc3aec3 --- /dev/null +++ b/africanus/gridding/perleypolyhedron/policies/phase_transform_policies.py @@ -0,0 +1,74 @@ +from africanus.util.numba import overload +from numpy import pi, cos, sin, sqrt + + +def phase_norotate(vis, + uvw, + lambdas, + ra0, + dec0, + ra, + dec, + policy_type, + phasesign=1.0): + pass + + +def phase_rotate(vis, + uvw, + lambdas, + ra0, + dec0, + ra, + dec, + policy_type, + phasesign=1.0): + ''' + Convert ra,dec to l,m,n based on Synthesis Imaging II, Pg. 388 + The phase term (as documented in Perley & Cornwell (1992)) + calculation requires the delta l,m,n coordinates. + Through simplification l0,m0,n0 = (0,0,1) (assume dec == dec0 and + ra == ra0, and the simplification follows) + l,m,n is then calculated using the new and original phase centres + as per the relation on Pg. 388 + lambdas has the same shape as vis + ''' + d_ra = ra - ra0 + d_dec = dec + d_decp = dec0 + c_d_dec = cos(d_dec) + s_d_dec = sin(d_dec) + s_d_ra = sin(d_ra) + c_d_ra = cos(d_ra) + c_d_decp = cos(d_decp) + s_d_decp = sin(d_decp) + ll = c_d_dec * s_d_ra + mm = (s_d_dec * c_d_decp - c_d_dec * s_d_decp * c_d_ra) + nn = -(1 - sqrt(1 - ll * ll - mm * mm)) + for c in range(lambdas.size): + x = phasesign * 2 * pi * (uvw[0] * ll + uvw[1] * mm + + uvw[2] * nn) / lambdas[c] + vis[c, :] *= cos(x) + 1.0j * sin(x) + + +def policy(vis, uvw, lambdas, ra0, dec0, ra, dec, policy_type, phasesign=1.0): + pass + + +@overload(policy, inline="always") +def policy_impl(vis, + uvw, + lambdas, + ra0, + dec0, + ra, + dec, + policy_type, + phasesign=1.0): + if policy_type.literal_value == "None" or \ + policy_type.literal_value is None: + return phase_norotate + elif policy_type.literal_value == "phase_rotate": + return phase_rotate + else: + raise ValueError("Invalid baseline transform policy type") diff --git a/africanus/gridding/perleypolyhedron/policies/stokes_conversion_policies.py b/africanus/gridding/perleypolyhedron/policies/stokes_conversion_policies.py new file mode 100644 index 000000000..7cc5811ad --- /dev/null +++ b/africanus/gridding/perleypolyhedron/policies/stokes_conversion_policies.py @@ -0,0 +1,262 @@ +from africanus.util.numba import overload + + +def stokes2corr(vis_in, vis_out, policy_type): + pass + + +@overload(stokes2corr, inline="always") +def stokes2corrimpl(vis_in, vis_out, policy_type): + if policy_type.literal_value == "XXYY_FROM_I": + + def XXYY_FROM_I(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += vis_in + + return XXYY_FROM_I + elif policy_type.literal_value == "XXXYYXYY_FROM_I": + + def XXXYYXYY_FROM_I(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += 0 + vis_out[2] += 0 + vis_out[3] += vis_in + + return XXXYYXYY_FROM_I + elif policy_type.literal_value == "RRLL_FROM_I": + + def RRLL_FROM_I(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += vis_in + + return RRLL_FROM_I + elif policy_type.literal_value == "RRRLLRLL_FROM_I": + + def RRRLLRLL_FROM_I(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += 0 + vis_out[2] += 0 + vis_out[3] += vis_in + + return RRRLLRLL_FROM_I + elif policy_type.literal_value == "XXYY_FROM_Q": + + def XXYY_FROM_Q(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += -vis_in + + return XXYY_FROM_Q + elif policy_type.literal_value == "XXXYYXYY_FROM_Q": + + def XXXYYXYY_FROM_Q(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += 0 + vis_out[2] += 0 + vis_out[3] += -vis_in + + return XXXYYXYY_FROM_Q + elif policy_type.literal_value == "RLLR_FROM_Q": + + def RLLR_FROM_Q(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += vis_in + + return RLLR_FROM_Q + elif policy_type.literal_value == "RRRLLRLL_FROM_Q": + + def RRRLLRLL_FROM_Q(vis_in, vis_out, policy_type): + vis_out[0] += 0 + vis_out[1] += vis_in + vis_out[2] += vis_in + vis_out[3] += 0 + + return RRRLLRLL_FROM_Q + elif policy_type.literal_value == "XYYX_FROM_U": + + def XYYX_FROM_U(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += vis_in + + return XYYX_FROM_U + elif policy_type.literal_value == "XXXYYXYY_FROM_U": + + def XXXYYXYY_FROM_U(vis_in, vis_out, policy_type): + vis_out[0] += 0 + vis_out[1] += vis_in + vis_out[2] += vis_in + vis_out[3] += 0 + + return XXXYYXYY_FROM_U + elif policy_type.literal_value == "RLLR_FROM_U": + + def RLLR_FROM_U(vis_in, vis_out, policy_type): + vis_out[0] += 1.0j * vis_in + vis_out[1] += -1.0j * vis_in + + return RLLR_FROM_U + elif policy_type.literal_value == "RRRLLRLL_FROM_U": + + def RRRLLRLL_FROM_U(vis_in, vis_out, policy_type): + vis_out[0] += 0.0 + vis_out[1] += 1.0j * vis_in + vis_out[2] += -1.0j * vis_in + vis_out[3] += 0.0 + + return RRRLLRLL_FROM_U + elif policy_type.literal_value == "XYYX_FROM_V": + + def XYYX_FROM_V(vis_in, vis_out, policy_type): + vis_out[0] += 1.0j * vis_in + vis_out[1] += -1.0j * vis_in + + return XYYX_FROM_V + elif policy_type.literal_value == "XXXYYXYY_FROM_V": + + def XXXYYXYY_FROM_V(vis_in, vis_out, policy_type): + vis_out[0] += 0.0 + vis_out[1] += 1.0j * vis_in + vis_out[2] += -1.0j * vis_in + vis_out[3] += 0.0 + + return XXXYYXYY_FROM_V + elif policy_type.literal_value == "RRLL_FROM_V": + + def RRLL_FROM_V(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += -vis_in + + return RRLL_FROM_V + elif policy_type.literal_value == "RRRLLRLL_FROM_V": + + def RRRLLRLL_FROM_V(vis_in, vis_out, policy_type): + vis_out[0] += vis_in + vis_out[1] += 0 + vis_out[2] += 0 + vis_out[3] += -vis_in + + return RRRLLRLL_FROM_V + else: + raise ValueError("Invalid stokes conversion") + + +def corr2stokes(vis_in, policy_type): + pass + + +@overload(corr2stokes, inline="always") +def corr2stokesimpl(vis_in, policy_type): + if policy_type.literal_value == "I_FROM_XXYY": + return lambda vis_in, policy_type: (vis_in[0] + vis_in[1]) * 0.5 + elif policy_type.literal_value == "I_FROM_XXXYYXYY": + return lambda vis_in, policy_type: (vis_in[0] + vis_in[3]) * 0.5 + elif policy_type.literal_value == "I_FROM_RRLL": + return lambda vis_in, policy_type: (vis_in[0] + vis_in[1]) * 0.5 + elif policy_type.literal_value == "I_FROM_RRRLLRLL": + return lambda vis_in, policy_type: (vis_in[0] + vis_in[3]) * 0.5 + elif policy_type.literal_value == "Q_FROM_XXYY": + return lambda vis_in, policy_type: (vis_in[0] - vis_in[1]) * 0.5 + elif policy_type.literal_value == "Q_FROM_XXXYYXYY": + return lambda vis_in, policy_type: (vis_in[0] - vis_in[3]) * 0.5 + elif policy_type.literal_value == "Q_FROM_RRRLLRLL": + return lambda vis_in, policy_type: (vis_in[1] + vis_in[2]) * 0.5 + elif policy_type.literal_value == "U_FROM_XYYX": + return lambda vis_in, policy_type: (vis_in[0] + vis_in[1]) * 0.5 + elif policy_type.literal_value == "U_FROM_XXXYYXYY": + return lambda vis_in, policy_type: (vis_in[1] + vis_in[2]) * 0.5 + elif policy_type.literal_value == "U_FROM_RLLR": + return lambda vis_in, policy_type: -1.0j * (vis_in[0] - vis_in[1] + ) * 0.5 + elif policy_type.literal_value == "U_FROM_RRRLLRLL": + return lambda vis_in, policy_type: -1.0j * (vis_in[1] - vis_in[2] + ) * 0.5 + elif policy_type.literal_value == "V_FROM_RRLL": + return lambda vis_in, policy_type: (vis_in[0] - vis_in[1]) * 0.5 + elif policy_type.literal_value == "V_FROM_RRRLLRLL": + return lambda vis_in, policy_type: (vis_in[0] - vis_in[3]) * 0.5 + elif policy_type.literal_value == "V_FROM_XYYX": + return lambda vis_in, policy_type: -1.0j * (vis_in[0] - vis_in[1] + ) * 0.5 + elif policy_type.literal_value == "V_FROM_XXXYYXYY": + return lambda vis_in, policy_type: -1.0j * (vis_in[1] - vis_in[2] + ) * 0.5 + else: + raise ValueError("Invalid stokes conversion") + + +def ncorr_out(policy_type): + pass + + +@overload(ncorr_out, inline="always") +def ncorr_outimpl(policy_type): + if policy_type.literal_value == "XXYY_FROM_I": + return lambda policy_type: 2 + elif policy_type.literal_value == "XXXYYXYY_FROM_I": + return lambda policy_type: 4 + elif policy_type.literal_value == "RRLL_FROM_I": + return lambda policy_type: 2 + elif policy_type.literal_value == "RRRLLRLL_FROM_I": + return lambda policy_type: 4 + elif policy_type.literal_value == "XXYY_FROM_Q": + return lambda policy_type: 2 + elif policy_type.literal_value == "XXXYYXYY_FROM_Q": + return lambda policy_type: 4 + elif policy_type.literal_value == "RLLR_FROM_Q": + return lambda policy_type: 2 + elif policy_type.literal_value == "RRRLLRLL_FROM_Q": + return lambda policy_type: 4 + elif policy_type.literal_value == "XYYX_FROM_U": + return lambda policy_type: 2 + elif policy_type.literal_value == "XXXYYXYY_FROM_U": + return lambda policy_type: 4 + elif policy_type.literal_value == "RLLR_FROM_U": + return lambda policy_type: 2 + elif policy_type.literal_value == "RRRLLRLL_FROM_U": + return lambda policy_type: 4 + elif policy_type.literal_value == "XYYX_FROM_V": + return lambda policy_type: 2 + elif policy_type.literal_value == "XXXYYXYY_FROM_V": + return lambda policy_type: 4 + elif policy_type.literal_value == "RRLL_FROM_V": + return lambda policy_type: 2 + elif policy_type.literal_value == "RRRLLRLL_FROM_V": + return lambda policy_type: 4 + else: + raise ValueError("Invalid stokes conversion") + + +def ncorr_outpy(policy_type): + if policy_type == "XXYY_FROM_I": + return lambda: 2 + elif policy_type == "XXXYYXYY_FROM_I": + return lambda: 4 + elif policy_type == "RRLL_FROM_I": + return lambda: 2 + elif policy_type == "RRRLLRLL_FROM_I": + return lambda: 4 + elif policy_type == "XXYY_FROM_Q": + return lambda: 2 + elif policy_type == "XXXYYXYY_FROM_Q": + return lambda: 4 + elif policy_type == "RLLR_FROM_Q": + return lambda: 2 + elif policy_type == "RRRLLRLL_FROM_Q": + return lambda: 4 + elif policy_type == "XYYX_FROM_U": + return lambda: 2 + elif policy_type == "XXXYYXYY_FROM_U": + return lambda: 4 + elif policy_type == "RLLR_FROM_U": + return lambda: 2 + elif policy_type == "RRRLLRLL_FROM_U": + return lambda: 4 + elif policy_type == "XYYX_FROM_V": + return lambda: 2 + elif policy_type == "XXXYYXYY_FROM_V": + return lambda: 4 + elif policy_type == "RRLL_FROM_V": + return lambda: 2 + elif policy_type == "RRRLLRLL_FROM_V": + return lambda: 4 + else: + raise ValueError("Invalid stokes conversion") diff --git a/africanus/gridding/perleypolyhedron/tests/__init__.py b/africanus/gridding/perleypolyhedron/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py new file mode 100644 index 000000000..2d08b5499 --- /dev/null +++ b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py @@ -0,0 +1,367 @@ +# -*- coding: utf-8 -*- + +import time + +import numpy as np +import pytest +from africanus.gridding.perleypolyhedron import (kernels, + gridder, + degridder) +from africanus.gridding.perleypolyhedron import dask as dwrap +from africanus.dft.kernels import im_to_vis + + +class clock: + def __init__(self, identifier="untitled"): + self._id = identifier + self._elapsed = 0.0 + self._onenter = 0.0 + self._onexit = 0.0 + + def __enter__(self): + self._onenter = time.time() + return self + + def __exit__(self, extype, exval, tb): + self._onexit = time.time() + self._elapsed = self._onexit - self._onenter + + @property + def elapsed(self): + return self._elapsed + + def __str__(self): + res = "{0:s}: Walltime {1:.0f}m{2:.2f}s elapsed".format( + self._id, self.elapsed // 60, + self.elapsed - (self.elapsed // 60) * 60) + return res + + __repr__ = __str__ + + +def test_gridder_dask(): + da = pytest.importorskip("dask.array") + + with clock("DASK gridding") as tictoc: + # construct kernel + W = 5 + OS = 9 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) + nrow = int(1e6) + np.random.seed(0) + # simulate some ficticious baselines rotated by an hour angle + row_chunks = nrow // 10 + uvw = np.zeros((nrow, 3), dtype=np.float64) + blpos = np.random.uniform(26, 10000, size=(25, 3)) + ntime = int(nrow / 25.0) + d0 = np.pi / 4.0 + for n in range(25): + for ih0, h0 in enumerate( + np.linspace(np.deg2rad(-20), np.deg2rad(20), ntime)): + s = np.sin + c = np.cos + R = np.array([[s(h0), c(h0), 0], + [-s(d0) * c(h0), + s(d0) * s(h0), + c(d0)], + [c(d0) * c(h0), -c(d0) * s(h0), + s(d0)]]) + uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T) + uvw = da.from_array(uvw, chunks=(row_chunks, 3)) + pxacrossbeam = 5 + nchan = 128 + frequency = da.from_array(np.linspace(1.0e9, 1.4e9, nchan), + chunks=(nchan, )) + wavelength = 299792458.0 / frequency + cell = da.rad2deg( + wavelength[0] / + (max(da.max(da.absolute(uvw[:, 0])), + da.max(da.absolute(uvw[:, 1]))) * pxacrossbeam)) + npixfacet = 100 + fftpad = 1.1 + + image_centres = da.from_array(np.array([[0, d0]]), chunks=(1, 2)) + chanmap = da.from_array(np.zeros(nchan, dtype=np.int64), + chunks=(nchan, )) + detaper_facet = kernels.compute_detaper_dft_seperable( + int(npixfacet * fftpad), kernels.unpack_kernel(kern, W, OS), W, + OS) + vis_dft = da.ones(shape=(nrow, nchan, 2), + chunks=(row_chunks, nchan, 2), + dtype=np.complex64) + vis_grid_facet = dwrap.gridder( + uvw, + vis_dft, + wavelength, + chanmap, + int(npixfacet * fftpad), + cell * 3600.0, + image_centres, (0, d0), + kern, + W, + OS, + "None", + "None", + "I_FROM_XXYY", + "conv_1d_axisymmetric_packed_scatter", + do_normalize=True) + + vis_grid_facet = vis_grid_facet.compute() + + ftvisfacet = (np.fft.fftshift( + np.fft.ifft2(np.fft.ifftshift( + vis_grid_facet[0, :, :]))).reshape( + (1, int(npixfacet * fftpad), int( + npixfacet * fftpad)))).real / detaper_facet * int( + npixfacet * fftpad)**2 + ftvisfacet = ftvisfacet[:, + int(npixfacet * fftpad) // 2 - npixfacet // + 2:int(npixfacet * fftpad) // 2 - + npixfacet // 2 + npixfacet, + int(npixfacet * fftpad) // 2 - npixfacet // + 2:int(npixfacet * fftpad) // 2 - + npixfacet // 2 + npixfacet] + print(tictoc) + assert (np.abs(np.max(ftvisfacet[0, :, :]) - 1.0) < 1.0e-6) + + +def test_gridder_nondask(): + with clock("Non-DASK gridding") as tictoc: + # construct kernel + W = 5 + OS = 9 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) + nrow = int(1e6) + np.random.seed(0) + # simulate some ficticious baselines rotated by an hour angle + uvw = np.zeros((nrow, 3), dtype=np.float64) + blpos = np.random.uniform(26, 10000, size=(25, 3)) + ntime = int(nrow / 25.0) + d0 = np.pi / 4.0 + for n in range(25): + for ih0, h0 in enumerate( + np.linspace(np.deg2rad(-20), np.deg2rad(20), ntime)): + s = np.sin + c = np.cos + R = np.array([[s(h0), c(h0), 0], + [-s(d0) * c(h0), + s(d0) * s(h0), + c(d0)], + [c(d0) * c(h0), -c(d0) * s(h0), + s(d0)]]) + uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T) + pxacrossbeam = 5 + nchan = 128 + frequency = np.linspace(1.0e9, 1.4e9, nchan) + wavelength = 299792458.0 / frequency + cell = np.rad2deg( + wavelength[0] / + (max(np.max(np.absolute(uvw[:, 0])), + np.max(np.absolute(uvw[:, 1]))) * pxacrossbeam)) + npixfacet = 100 + fftpad = 1.1 + + image_centres = np.array([[0, d0]]) + chanmap = np.zeros(nchan, dtype=np.int64) + detaper_facet = kernels.compute_detaper_dft_seperable( + int(npixfacet * fftpad), kernels.unpack_kernel(kern, W, OS), W, + OS) + vis_dft = np.ones((nrow, nchan, 2), dtype=np.complex64) + vis_grid_facet = gridder.gridder( + uvw, + vis_dft, + wavelength, + chanmap, + int(npixfacet * fftpad), + cell * 3600.0, + image_centres[0, :], (0, d0), + kern, + W, + OS, + "None", + "None", + "I_FROM_XXYY", + "conv_1d_axisymmetric_packed_scatter", + do_normalize=True) + ftvisfacet = (np.fft.fftshift( + np.fft.ifft2(np.fft.ifftshift( + vis_grid_facet[0, :, :]))).reshape( + (1, int(npixfacet * fftpad), int( + npixfacet * fftpad)))).real / detaper_facet * int( + npixfacet * fftpad)**2 + ftvisfacet = ftvisfacet[:, + int(npixfacet * fftpad) // 2 - npixfacet // + 2:int(npixfacet * fftpad) // 2 - + npixfacet // 2 + npixfacet, + int(npixfacet * fftpad) // 2 - npixfacet // + 2:int(npixfacet * fftpad) // 2 - + npixfacet // 2 + npixfacet] + print(tictoc) + assert (np.abs(np.max(ftvisfacet[0, :, :]) - 1.0) < 1.0e-6) + + +def test_degrid_dft_packed_nondask(): + # construct kernel + W = 5 + OS = 3 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), + W, + oversample=OS) + nrow = int(5e4) + uvw = np.column_stack( + (5000.0 * np.cos(np.linspace(0, 2 * np.pi, nrow)), + 5000.0 * np.sin(np.linspace(0, 2 * np.pi, nrow)), np.zeros(nrow))) + + pxacrossbeam = 10 + nchan = 1024 + frequency = np.array(np.linspace(1.0e9, 1.4e9, nchan)) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npix = 512 + mod = np.zeros((1, npix, npix), dtype=np.complex64) + mod[0, npix // 2 - 5, npix // 2 - 5] = 1.0 + + ftmod = np.fft.ifftshift(np.fft.fft2(np.fft.fftshift( + mod[0, :, :]))).reshape((1, npix, npix)) + chanmap = np.zeros(nchan, dtype=np.int64) + + with clock("Non-DASK degridding") as tictoc: + degridder.degridder( + uvw, + ftmod, + wavelength, + chanmap, + cell * 3600.0, + (0, np.pi / 4.0), + (0, np.pi / 4.0), + kern, + W, + OS, + "None", # no faceting + "None", # no faceting + "XXYY_FROM_I", + "conv_1d_axisymmetric_packed_gather") + + print(tictoc) + + +def test_degrid_dft_packed_dask(): + da = pytest.importorskip("dask.array") + + # construct kernel + W = 5 + OS = 3 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), + W, + oversample=OS) + nrow = int(5e4) + nrow_chunk = nrow // 32 + uvw = np.column_stack( + (5000.0 * np.cos(np.linspace(0, 2 * np.pi, nrow)), + 5000.0 * np.sin(np.linspace(0, 2 * np.pi, nrow)), np.zeros(nrow))) + + pxacrossbeam = 10 + nchan = 1024 + frequency = np.array(np.linspace(1.0e9, 1.4e9, nchan)) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npix = 512 + mod = np.zeros((1, npix, npix), dtype=np.complex64) + mod[0, npix // 2 - 5, npix // 2 - 5] = 1.0 + + ftmod = np.fft.ifftshift(np.fft.fft2(np.fft.fftshift( + mod[0, :, :]))).reshape((1, 1, npix, npix)) + chanmap = np.zeros(nchan, dtype=np.int64) + + with clock("DASK degridding") as tictoc: + vis_degrid = dwrap.degridder( + da.from_array(uvw, chunks=(nrow_chunk, 3)), + da.from_array(ftmod, chunks=(1, 1, npix, npix)), + da.from_array(wavelength, chunks=(nchan, )), + da.from_array(chanmap, chunks=(nchan, )), + cell * 3600.0, + da.from_array(np.array([[0, np.pi / 4.0]]), chunks=(1, 2)), + (0, np.pi / 4.0), + kern, + W, + OS, + "None", # no faceting + "None", # no faceting + "XXYY_FROM_I", + "conv_1d_axisymmetric_packed_gather") + + vis_degrid = vis_degrid.compute() + + print(tictoc) + + +def test_degrid_dft_packed_dask_dft_check(): + da = pytest.importorskip("dask.array") + + # construct kernel + W = 5 + OS = 3 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), + W, + oversample=OS) + nrow = 100 + nrow_chunk = nrow // 8 + uvw = np.column_stack( + (5000.0 * np.cos(np.linspace(0, 2 * np.pi, nrow)), + 5000.0 * np.sin(np.linspace(0, 2 * np.pi, nrow)), np.zeros(nrow))) + + pxacrossbeam = 10 + nchan = 16 + frequency = np.array(np.linspace(1.0e9, 1.4e9, nchan)) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npix = 512 + mod = np.zeros((1, npix, npix), dtype=np.complex64) + mod[0, npix // 2 - 5, npix // 2 - 5] = 1.0 + + ftmod = np.fft.ifftshift(np.fft.fft2(np.fft.fftshift( + mod[0, :, :]))).reshape((1, 1, npix, npix)) + chanmap = np.zeros(nchan, dtype=np.int64) + dec, ra = np.meshgrid( + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell), + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) + radec = np.column_stack((ra.flatten(), dec.flatten())) + vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, npix * npix).T.copy(), + uvw, radec, frequency) + + vis_degrid = dwrap.degridder( + da.from_array(uvw, chunks=(nrow_chunk, 3)), + da.from_array(ftmod, chunks=(1, 1, npix, npix)), + da.from_array(wavelength, chunks=(nchan, )), + da.from_array(chanmap, chunks=(nchan, )), + cell * 3600.0, + da.from_array(np.array([[0, np.pi / 4.0]]), chunks=(1, 2)), + (0, np.pi / 4.0), + kern, + W, + OS, + "None", # no faceting + "None", # no faceting + "XXYY_FROM_I", + "conv_1d_axisymmetric_packed_gather") + + vis_degrid = vis_degrid.compute() + + assert np.percentile( + np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real), + 99.0) < 0.05 + assert np.percentile( + np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag), + 99.0) < 0.05 diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py new file mode 100644 index 000000000..bd02befab --- /dev/null +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -0,0 +1,743 @@ +# -*- coding: utf-8 -*- + +import os + +import numpy as np +import pytest + +from africanus.gridding.perleypolyhedron import (kernels, + gridder, + degridder) +from africanus.dft.kernels import im_to_vis, vis_to_im +from africanus.coordinates import radec_to_lmn + + +def test_construct_kernels(tmp_path_factory): + matplotlib = pytest.importorskip("matplotlib") + matplotlib.use("agg") + plt = pytest.importorskip("matplotlib.pyplot") + + plt.figure() + WIDTH = 5 + OVERSAMP = 101 + ll = kernels.uspace(WIDTH, OVERSAMP) + sel = ll <= (WIDTH + 2) // 2 + plt.axvline(0.5, 0, 1, ls="--", c="k") + plt.axvline(-0.5, 0, 1, ls="--", c="k") + plt.plot( + ll[sel] * OVERSAMP / 2 / np.pi, + 10 * np.log10( + np.abs( + np.fft.fftshift( + np.fft.fft( + kernels.kbsinc(WIDTH, oversample=OVERSAMP, + order=0)[sel])))), + label="kbsinc order 0") + plt.plot( + ll[sel] * OVERSAMP / 2 / np.pi, + 10 * np.log10( + np.abs( + np.fft.fftshift( + np.fft.fft( + kernels.kbsinc( + WIDTH, oversample=OVERSAMP, order=15)[sel])))), + label="kbsinc order 15") + plt.plot(ll[sel] * OVERSAMP / 2 / np.pi, + 10 * np.log10( + np.abs( + np.fft.fftshift( + np.fft.fft( + kernels.hanningsinc( + WIDTH, oversample=OVERSAMP)[sel])))), + label="hanning sinc") + plt.plot( + ll[sel] * OVERSAMP / 2 / np.pi, + 10 * np.log10( + np.abs( + np.fft.fftshift( + np.fft.fft( + kernels.sinc(WIDTH, oversample=OVERSAMP)[sel])))), + label="sinc") + plt.xlim(-10, 10) + plt.legend() + plt.ylabel("Response [dB]") + plt.xlabel("FoV") + plt.grid(True) + plt.savefig(tmp_path_factory.mktemp("plots") / "aakernels.png") + + +def test_taps(): + oversample = 14 + W = 5 + taps = kernels.uspace(W, oversample=oversample) + assert taps[oversample * ((W + 2) // 2)] == 0 + assert taps[0] == -((W + 2) // 2) + assert taps[-oversample] == (W + 2) // 2 + + +def test_packunpack(): + oversample = 4 + W = 3 + K = kernels.uspace(W, oversample=oversample) + Kp = kernels.pack_kernel(K, W, oversample=oversample) + Kup = kernels.unpack_kernel(Kp, W, oversample=oversample) + assert np.all(K == Kup) + assert np.allclose(K, [ + -2.0, -1.75, -1.5, -1.25, -1.0, -0.75, -0.5, -0.25, 0, 0.25, 0.5, + 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5, 2.75 + ]) + assert np.allclose(Kp, [ + -2.0, -1.0, 0, 1.0, 2.0, -1.75, -0.75, 0.25, 1.25, 2.25, -1.5, + -0.5, 0.5, 1.5, 2.5, -1.25, -0.25, 0.75, 1.75, 2.75 + ]) + + +def test_facetcodepath(): + # construct kernel + W = 5 + OS = 3 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), + W, + oversample=OS) + + # offset 0 + uvw = np.array([[0, 0, 0]]) + vis = np.array([[[1.0 + 0j, 1.0 + 0j]]]) + gridder.gridder(uvw, vis, np.array([1.0]), np.array([0]), 64, + 30, (0, 0), (0, 0), kern, W, OS, "rotate", + "phase_rotate", "I_FROM_XXYY", + "conv_1d_axisymmetric_packed_scatter") + + +def test_degrid_dft(tmp_path_factory): + # construct kernel + W = 5 + OS = 3 + kern = kernels.kbsinc(W, oversample=OS) + uvw = np.column_stack( + (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 1000)), + 5000.0 * np.sin(np.linspace(0, 2 * np.pi, 1000)), np.zeros(1000))) + + pxacrossbeam = 10 + frequency = np.array([1.4e9]) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npix = 512 + mod = np.zeros((1, npix, npix), dtype=np.complex64) + mod[0, npix // 2 - 5, npix // 2 - 5] = 1.0 + + ftmod = np.fft.ifftshift(np.fft.fft2(np.fft.fftshift( + mod[0, :, :]))).reshape((1, npix, npix)) + chanmap = np.array([0]) + vis_degrid = degridder.degridder( + uvw, + ftmod, + wavelength, + chanmap, + cell * 3600.0, + (0, np.pi / 4.0), + (0, np.pi / 4.0), + kern, + W, + OS, + "None", # no faceting + "None", # no faceting + "XXYY_FROM_I", + "conv_1d_axisymmetric_unpacked_gather") + + dec, ra = np.meshgrid( + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell), + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) + radec = np.column_stack((ra.flatten(), dec.flatten())) + + vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, npix * npix).T.copy(), + uvw, radec, frequency) + + try: + import matplotlib + except ImportError: + pass + else: + matplotlib.use("agg") + from matplotlib import pyplot as plt + plt.figure() + plt.plot(vis_degrid[:, 0, 0].real, label=r"$\Re(\mathtt{degrid})$") + plt.plot(vis_dft[:, 0, 0].real, label=r"$\Re(\mathtt{dft})$") + plt.plot(np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real), + label="Error") + plt.legend() + plt.xlabel("sample") + plt.ylabel("Real of predicted") + plt.savefig( + os.path.join(os.environ.get("TMPDIR", "/tmp"), + "degrid_vs_dft_re.png")) + plt.figure() + plt.plot(vis_degrid[:, 0, 0].imag, label=r"$\Im(\mathtt{degrid})$") + plt.plot(vis_dft[:, 0, 0].imag, label=r"$\Im(\mathtt{dft})$") + plt.plot(np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag), + label="Error") + plt.legend() + plt.xlabel("sample") + plt.ylabel("Imag of predicted") + plt.savefig(tmp_path_factory.mktemp("degrid_dft") / + "degrid_vs_dft_im.png") + + assert np.percentile( + np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real), + 99.0) < 0.05 + assert np.percentile( + np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag), + 99.0) < 0.05 + + +def test_degrid_dft_packed(tmp_path_factory): + # construct kernel + W = 5 + OS = 3 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), + W, + oversample=OS) + uvw = np.column_stack( + (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 1000)), + 5000.0 * np.sin(np.linspace(0, 2 * np.pi, 1000)), np.zeros(1000))) + + pxacrossbeam = 10 + frequency = np.array([1.4e9]) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npix = 512 + mod = np.zeros((1, npix, npix), dtype=np.complex64) + mod[0, npix // 2 - 5, npix // 2 - 5] = 1.0 + + ftmod = np.fft.ifftshift(np.fft.fft2(np.fft.fftshift( + mod[0, :, :]))).reshape((1, npix, npix)) + chanmap = np.array([0]) + vis_degrid = degridder.degridder( + uvw, + ftmod, + wavelength, + chanmap, + cell * 3600.0, + (0, np.pi / 4.0), + (0, np.pi / 4.0), + kern, + W, + OS, + "None", # no faceting + "None", # no faceting + "XXYY_FROM_I", + "conv_1d_axisymmetric_packed_gather") + + dec, ra = np.meshgrid( + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell), + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) + radec = np.column_stack((ra.flatten(), dec.flatten())) + + vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, npix * npix).T.copy(), + uvw, radec, frequency) + + try: + import matplotlib + except ImportError: + pass + else: + matplotlib.use("agg") + from matplotlib import pyplot as plt + plt.figure() + plt.plot(vis_degrid[:, 0, 0].real, label=r"$\Re(\mathtt{degrid})$") + plt.plot(vis_dft[:, 0, 0].real, label=r"$\Re(\mathtt{dft})$") + plt.plot(np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real), + label="Error") + plt.legend() + plt.xlabel("sample") + plt.ylabel("Real of predicted") + plt.savefig( + os.path.join(os.environ.get("TMPDIR", "/tmp"), + "degrid_vs_dft_re_packed.png")) + plt.figure() + plt.plot(vis_degrid[:, 0, 0].imag, label=r"$\Im(\mathtt{degrid})$") + plt.plot(vis_dft[:, 0, 0].imag, label=r"$\Im(\mathtt{dft})$") + plt.plot(np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag), + label="Error") + plt.legend() + plt.xlabel("sample") + plt.ylabel("Imag of predicted") + plt.savefig(tmp_path_factory.mktemp("degrid_dft_packed") / + "degrid_vs_dft_im_packed.png") + + assert np.percentile( + np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real), + 99.0) < 0.05 + assert np.percentile( + np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag), + 99.0) < 0.05 + + +def test_detaper(tmp_path_factory): + W = 5 + OS = 3 + K1D = kernels.kbsinc(W, oversample=OS) + K2D = np.outer(K1D, K1D) + detaper = kernels.compute_detaper(128, K2D, W, OS) + detaperdft = kernels.compute_detaper_dft(128, K2D, W, OS) + detaperdftsep = kernels.compute_detaper_dft_seperable(128, K1D, W, OS) + + try: + import matplotlib + except ImportError: + pass + else: + matplotlib.use("agg") + from matplotlib import pyplot as plt + plt.figure() + plt.subplot(131) + plt.title("FFT detaper") + plt.imshow(detaper) + plt.colorbar() + plt.subplot(132) + plt.title("DFT detaper") + plt.imshow(detaperdft) + plt.colorbar() + plt.subplot(133) + plt.title("ABS error") + plt.imshow(np.abs(detaper - detaperdft)) + plt.colorbar() + plt.savefig(tmp_path_factory.mktemp("detaper") / "detaper.png") + + assert (np.percentile(np.abs(detaper - detaperdft), 99.0) < 1.0e-14) + assert (np.max(np.abs(detaperdft - detaperdftsep)) < 1.0e-14) + + +def test_grid_dft(tmp_path_factory): + # construct kernel + W = 7 + OS = 9 + kern = kernels.kbsinc(W, oversample=OS) + nrow = 5000 + np.random.seed(0) + uvw = np.random.normal(scale=6000, size=(nrow, 3)) + uvw[:, 2] = 0.0 # ignore widefield effects for now + + pxacrossbeam = 10 + frequency = np.array([30.0e9]) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npix = 256 + fftpad = 1.25 + mod = np.zeros((1, npix, npix), dtype=np.complex64) + for n in [int(n) for n in np.linspace(npix // 8, 2 * npix // 5, 5)]: + mod[0, npix // 2 + n, npix // 2 + n] = 1.0 + mod[0, npix // 2 + n, npix // 2 - n] = 1.0 + mod[0, npix // 2 - n, npix // 2 - n] = 1.0 + mod[0, npix // 2 - n, npix // 2 + n] = 1.0 + mod[0, npix // 2, npix // 2 + n] = 1.0 + mod[0, npix // 2, npix // 2 - n] = 1.0 + mod[0, npix // 2 - n, npix // 2] = 1.0 + mod[0, npix // 2 + n, npix // 2] = 1.0 + + dec, ra = np.meshgrid( + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell), + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) + radec = np.column_stack((ra.flatten(), dec.flatten())) + + vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, + npix * npix).T.copy(), uvw, + radec, frequency).repeat(2).reshape(nrow, 1, 2) + chanmap = np.array([0]) + + detaper = kernels.compute_detaper(int(npix * fftpad), + np.outer(kern, kern), W, OS) + vis_grid = gridder.gridder( + uvw, + vis_dft, + wavelength, + chanmap, + int(npix * fftpad), + cell * 3600.0, + (0, np.pi / 4.0), + (0, np.pi / 4.0), + kern, + W, + OS, + "None", # no faceting + "None", # no faceting + "I_FROM_XXYY", + "conv_1d_axisymmetric_unpacked_scatter", + do_normalize=True) + + ftvis = (np.fft.fftshift( + np.fft.ifft2(np.fft.ifftshift(vis_grid[0, :, :]))).reshape( + (1, int(npix * fftpad), int( + npix * fftpad)))).real / detaper * int(npix * fftpad)**2 + ftvis = ftvis[:, + int(npix * fftpad) // 2 - + npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix, + int(npix * fftpad) // 2 - + npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix] + dftvis = vis_to_im(vis_dft, uvw, radec, frequency, + np.zeros(vis_dft.shape, + dtype=np.bool)).T.copy().reshape( + 2, 1, npix, npix) / nrow + + try: + import matplotlib + except ImportError: + pass + else: + matplotlib.use("agg") + from matplotlib import pyplot as plt + plt.figure() + plt.subplot(131) + plt.title("FFT") + plt.imshow(ftvis[0, :, :]) + plt.colorbar() + plt.subplot(132) + plt.title("DFT") + plt.imshow(dftvis[0, 0, :, :]) + plt.colorbar() + plt.subplot(133) + plt.title("ABS diff") + plt.imshow(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :])) + plt.colorbar() + plt.savefig(tmp_path_factory.mktemp("grid_dft") / + "grid_diff_dft.png") + + assert (np.percentile(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :]), + 95.0) < 0.15) + + +def test_grid_dft_packed(tmp_path_factory): + # construct kernel + W = 7 + OS = 1009 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) + nrow = 5000 + np.random.seed(0) + uvw = np.random.normal(scale=6000, size=(nrow, 3)) + uvw[:, 2] = 0.0 # ignore widefield effects for now + + pxacrossbeam = 10 + frequency = np.array([30.0e9]) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npix = 256 + fftpad = 1.25 + mod = np.zeros((1, npix, npix), dtype=np.complex64) + for n in [int(n) for n in np.linspace(npix // 8, 2 * npix // 5, 5)]: + mod[0, npix // 2 + n, npix // 2 + n] = 1.0 + mod[0, npix // 2 + n, npix // 2 - n] = 1.0 + mod[0, npix // 2 - n, npix // 2 - n] = 1.0 + mod[0, npix // 2 - n, npix // 2 + n] = 1.0 + mod[0, npix // 2, npix // 2 + n] = 1.0 + mod[0, npix // 2, npix // 2 - n] = 1.0 + mod[0, npix // 2 - n, npix // 2] = 1.0 + mod[0, npix // 2 + n, npix // 2] = 1.0 + + dec, ra = np.meshgrid( + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell), + np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) + radec = np.column_stack((ra.flatten(), dec.flatten())) + + vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, + npix * npix).T.copy(), uvw, + radec, frequency).repeat(2).reshape(nrow, 1, 2) + chanmap = np.array([0]) + detaper = kernels.compute_detaper_dft_seperable( + int(npix * fftpad), kernels.unpack_kernel(kern, W, OS), W, OS) + vis_grid = gridder.gridder( + uvw, + vis_dft, + wavelength, + chanmap, + int(npix * fftpad), + cell * 3600.0, + (0, np.pi / 4.0), + (0, np.pi / 4.0), + kern, + W, + OS, + "None", # no faceting + "None", # no faceting + "I_FROM_XXYY", + "conv_1d_axisymmetric_packed_scatter", + do_normalize=True) + + ftvis = (np.fft.fftshift( + np.fft.ifft2(np.fft.ifftshift(vis_grid[0, :, :]))).reshape( + (1, int(npix * fftpad), int( + npix * fftpad)))).real / detaper * int(npix * fftpad)**2 + ftvis = ftvis[:, + int(npix * fftpad) // 2 - + npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix, + int(npix * fftpad) // 2 - + npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix] + dftvis = vis_to_im(vis_dft, uvw, radec, frequency, + np.zeros(vis_dft.shape, + dtype=np.bool)).T.copy().reshape( + 2, 1, npix, npix) / nrow + + try: + import matplotlib + except ImportError: + pass + else: + matplotlib.use("agg") + from matplotlib import pyplot as plt + plt.figure() + plt.subplot(131) + plt.title("FFT") + plt.imshow(ftvis[0, :, :]) + plt.colorbar() + plt.subplot(132) + plt.title("DFT") + plt.imshow(dftvis[0, 0, :, :]) + plt.colorbar() + plt.subplot(133) + plt.title("ABS diff") + plt.imshow(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :])) + plt.colorbar() + plt.savefig(tmp_path_factory.mktemp("grid_dft_packed") / + "grid_diff_dft_packed.png") + + assert (np.percentile(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :]), + 95.0) < 0.15) + + +def test_wcorrection_faceting_backward(tmp_path_factory): + # construct kernel + W = 5 + OS = 9 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) + nrow = 5000 + np.random.seed(0) + # simulate some ficticious baselines rotated by an hour angle + uvw = np.zeros((nrow, 3), dtype=np.float64) + blpos = np.random.uniform(26, 10000, size=(25, 3)) + ntime = int(nrow / 25.0) + d0 = np.pi / 4.0 + for n in range(25): + for ih0, h0 in enumerate( + np.linspace(np.deg2rad(-20), np.deg2rad(20), ntime)): + s = np.sin + c = np.cos + R = np.array([[s(h0), c(h0), 0], + [-s(d0) * c(h0), + s(d0) * s(h0), + c(d0)], [c(d0) * c(h0), -c(d0) * s(h0), + s(d0)]]) + uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T) + + pxacrossbeam = 5 + frequency = np.array([1.4e9]) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npix = 2048 + npixfacet = 100 + fftpad = 1.1 + mod = np.ones((1, 1, 1), dtype=np.complex64) + deltaradec = np.array( + [[600 * np.deg2rad(cell), 600 * np.deg2rad(cell)]]) + lm = radec_to_lmn(deltaradec + np.array([[0, d0]]), + phase_centre=np.array([0, d0])) + + vis_dft = im_to_vis(mod, uvw, lm[:, 0:2], + frequency).repeat(2).reshape(nrow, 1, 2) + chanmap = np.array([0]) + + detaper = kernels.compute_detaper_dft_seperable( + int(npix * fftpad), kernels.unpack_kernel(kern, W, OS), W, OS) + vis_grid_nofacet = gridder.gridder( + uvw, + vis_dft, + wavelength, + chanmap, + int(npix * fftpad), + cell * 3600.0, + (0, d0), + (0, d0), + kern, + W, + OS, + "None", # no faceting + "None", # no faceting + "I_FROM_XXYY", + "conv_1d_axisymmetric_packed_scatter", + do_normalize=True) + ftvis = (np.fft.fftshift( + np.fft.ifft2(np.fft.ifftshift(vis_grid_nofacet[0, :, :]))).reshape( + (1, int(npix * fftpad), int( + npix * fftpad)))).real / detaper * int(npix * fftpad)**2 + ftvis = ftvis[:, + int(npix * fftpad) // 2 - + npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix, + int(npix * fftpad) // 2 - + npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix] + + detaper_facet = kernels.compute_detaper_dft_seperable( + int(npixfacet * fftpad), kernels.unpack_kernel(kern, W, OS), W, OS) + vis_grid_facet = gridder.gridder( + uvw, + vis_dft, + wavelength, + chanmap, + int(npixfacet * fftpad), + cell * 3600.0, (deltaradec + np.array([[0, d0]]))[0, :], (0, d0), + kern, + W, + OS, + "rotate", + "phase_rotate", + "I_FROM_XXYY", + "conv_1d_axisymmetric_packed_scatter", + do_normalize=True) + ftvisfacet = (np.fft.fftshift( + np.fft.ifft2(np.fft.ifftshift(vis_grid_facet[0, :, :]))).reshape( + (1, int(npixfacet * fftpad), int( + npixfacet * fftpad)))).real / detaper_facet * int( + npixfacet * fftpad)**2 + ftvisfacet = ftvisfacet[:, + int(npixfacet * fftpad) // 2 - + npixfacet // 2:int(npixfacet * fftpad) // 2 - + npixfacet // 2 + npixfacet, + int(npixfacet * fftpad) // 2 - + npixfacet // 2:int(npixfacet * fftpad) // 2 - + npixfacet // 2 + npixfacet] + + try: + import matplotlib + except ImportError: + pass + else: + matplotlib.use("agg") + from matplotlib import pyplot as plt + plot_dir = tmp_path_factory.mktemp("wcorrection_backward") + + plt.figure() + plt.subplot(121) + plt.imshow(ftvis[0, 1624 - 50:1624 + 50, 1447 - 50:1447 + 50]) + plt.colorbar() + plt.title("Offset FFT (peak={0:.1f})".format(np.max(ftvis))) + plt.subplot(122) + plt.imshow(ftvisfacet[0, :, :]) + plt.colorbar() + plt.title("Faceted FFT (peak={0:.1f})".format(np.max(ftvisfacet))) + plt.savefig(plot_dir / "facet_imaging.png") + + assert (np.abs(np.max(ftvisfacet[0, :, :]) - 1.0) < 1.0e-6) + + +def test_wcorrection_faceting_forward(tmp_path_factory): + # construct kernel + W = 5 + OS = 9 + kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) + nrow = 5000 + np.random.seed(0) + # simulate some ficticious baselines rotated by an hour angle + uvw = np.zeros((nrow, 3), dtype=np.float64) + blpos = np.random.uniform(26, 10000, size=(25, 3)) + ntime = int(nrow / 25.0) + d0 = np.pi / 4.0 + for n in range(25): + for ih0, h0 in enumerate( + np.linspace(np.deg2rad(-20), np.deg2rad(20), ntime)): + s = np.sin + c = np.cos + R = np.array([[s(h0), c(h0), 0], + [-s(d0) * c(h0), + s(d0) * s(h0), + c(d0)], [c(d0) * c(h0), -c(d0) * s(h0), + s(d0)]]) + uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T) + + pxacrossbeam = 5 + frequency = np.array([1.4e9]) + wavelength = np.array([299792458.0 / f for f in frequency]) + + cell = np.rad2deg( + wavelength[0] / + (max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + npixfacet = 100 + mod = np.ones((1, 1, 1), dtype=np.complex64) + deltaradec = np.array([[20 * np.deg2rad(cell), 20 * np.deg2rad(cell)]]) + lm = radec_to_lmn(deltaradec + np.array([[0, d0]]), + phase_centre=np.array([0, d0])) + + vis_dft = im_to_vis(mod, uvw, lm[:, 0:2], + frequency).repeat(2).reshape(nrow, 1, 2) + chanmap = np.array([0]) + ftmod = np.ones((1, npixfacet, npixfacet), + dtype=np.complex64) # point source at centre of facet + vis_degrid = degridder.degridder( + uvw, + ftmod, + wavelength, + chanmap, + cell * 3600.0, + (deltaradec + np.array([[0, d0]]))[0, :], + (0, d0), + kern, + W, + OS, + "rotate", # no faceting + "phase_rotate", # no faceting + "XXYY_FROM_I", + "conv_1d_axisymmetric_packed_gather") + + try: + import matplotlib + except ImportError: + pass + else: + matplotlib.use("agg") + from matplotlib import pyplot as plt + plot_dir = tmp_path_factory.mktemp("wcorrection_forward") + + plt.figure() + plt.plot(vis_degrid[:, 0, 0].real, + label=r"$\Re(\mathtt{degrid facet})$") + plt.plot(vis_dft[:, 0, 0].real, label=r"$\Re(\mathtt{dft})$") + plt.plot(np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real), + label="Error") + plt.legend() + plt.xlabel("sample") + plt.ylabel("Real of predicted") + plt.savefig(plot_dir / "facet_degrid_vs_dft_re_packed.png") + plt.figure() + plt.plot(vis_degrid[:, 0, 0].imag, + label=r"$\Im(\mathtt{degrid facet})$") + plt.plot(vis_dft[:, 0, 0].imag, label=r"$\Im(\mathtt{dft})$") + plt.plot(np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag), + label="Error") + plt.legend() + plt.xlabel("sample") + plt.ylabel("Imag of predicted") + plt.savefig(plot_dir / "facet_degrid_vs_dft_im_packed.png") + + assert np.percentile( + np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real), + 99.0) < 0.05 + assert np.percentile( + np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag), + 99.0) < 0.05 diff --git a/africanus/util/numba.py b/africanus/util/numba.py index e0b2d6cb8..9574e6d7f 100644 --- a/africanus/util/numba.py +++ b/africanus/util/numba.py @@ -22,9 +22,9 @@ def wrapper(*args, **kwargs): jit = _fake_decorator njit = _fake_decorator stencil = _fake_decorator + overload = _fake_decorator register_jitable = _fake_decorator - else: from numba import cfunc, jit, njit, generated_jit, stencil # noqa from numba.extending import overload, register_jitable # noqa diff --git a/setup.py b/setup.py index d6fa5c7b1..324434fce 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,9 @@ # astropy breaks with numpy 1.15.3 # https://github.com/astropy/astropy/issues/7943 'numpy >= 1.14.0, != 1.15.3', - 'numba >= 0.46.0'] + # version 0.5 causes issues with + 'numba <= 0.49.0' + ] extras_require = { 'cuda': ['cupy >= 5.0.0', 'jinja2 >= 2.10'], From 62d07931cc435c7c556b0d751aab08f14a564cba Mon Sep 17 00:00:00 2001 From: bennahugo Date: Wed, 23 Sep 2020 11:48:18 +0200 Subject: [PATCH 02/12] Workaround for upstream bug in numba/numba#6274 --- .../gridding/perleypolyhedron/degridder.py | 22 +++++++++---------- .../gridding/perleypolyhedron/gridder.py | 2 +- .../policies/convolution_policies.py | 11 +++++----- 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/degridder.py b/africanus/gridding/perleypolyhedron/degridder.py index bdf3a0d30..efcbddb09 100644 --- a/africanus/gridding/perleypolyhedron/degridder.py +++ b/africanus/gridding/perleypolyhedron/degridder.py @@ -39,7 +39,7 @@ def degridder_row_kernel(uvw, ra0, dec0 = phase_centre ra, dec = image_centre btp.policy(uvw[r, :], ra, dec, ra0, dec0, - literally(baseline_transform_policy)) + baseline_transform_policy) for c in range(nvischan): scaled_u = uvw[r, 0] * scale_factor / lambdas[c] scaled_v = uvw[r, 1] * scale_factor / lambdas[c] @@ -57,7 +57,7 @@ def degridder_row_kernel(uvw, convolution_kernel_width, convolution_kernel_oversampling, stokes_conversion_policy, - policy_type=literally(convolution_policy)) + policy_type=convolution_policy) ptp.policy(vis[r, :, :], uvw[r, :], lambdas, @@ -65,7 +65,7 @@ def degridder_row_kernel(uvw, dec0, ra, dec, - policy_type=literally(phase_transform_policy), + policy_type=phase_transform_policy, phasesign=-1.0) @@ -151,10 +151,10 @@ def degridder(uvw, convolution_kernel, convolution_kernel_width, convolution_kernel_oversampling, - baseline_transform_policy, - phase_transform_policy, - stokes_conversion_policy, - convolution_policy, + literally(baseline_transform_policy), + literally(phase_transform_policy), + literally(stokes_conversion_policy), + literally(convolution_policy), vis_dtype=vis_dtype, nband=nband, nrow=nrow, @@ -249,10 +249,10 @@ def degridder_serial(uvw, convolution_kernel, convolution_kernel_width, convolution_kernel_oversampling, - baseline_transform_policy, - phase_transform_policy, - stokes_conversion_policy, - convolution_policy, + literally(baseline_transform_policy), + literally(phase_transform_policy), + literally(stokes_conversion_policy), + literally(convolution_policy), vis_dtype=vis_dtype, nband=nband, nrow=nrow, diff --git a/africanus/gridding/perleypolyhedron/gridder.py b/africanus/gridding/perleypolyhedron/gridder.py index a28946414..ed4826d89 100644 --- a/africanus/gridding/perleypolyhedron/gridder.py +++ b/africanus/gridding/perleypolyhedron/gridder.py @@ -108,7 +108,7 @@ def gridder(uvw, convolution_kernel, convolution_kernel_width, convolution_kernel_oversampling, - stokes_conversion_policy, + literally(stokes_conversion_policy), policy_type=literally(convolution_policy)) if do_normalize: for c in range(nband): diff --git a/africanus/gridding/perleypolyhedron/policies/convolution_policies.py b/africanus/gridding/perleypolyhedron/policies/convolution_policies.py index 3bff1054c..b0308a95e 100644 --- a/africanus/gridding/perleypolyhedron/policies/convolution_policies.py +++ b/africanus/gridding/perleypolyhedron/policies/convolution_policies.py @@ -1,5 +1,4 @@ from africanus.util.numba import overload -from numba import literally import numpy as np from . import stokes_conversion_policies as scp @@ -49,7 +48,7 @@ def convolve_1d_axisymmetric_unpacked_scatter( grid[grid_v_lookup, grid_u_lookup] += \ conv_v * conv_u * \ scp.corr2stokes(vis[r, c, :], - literally(stokes_conversion_policy)) + stokes_conversion_policy) cw += conv_v * conv_u return cw @@ -108,7 +107,7 @@ def convolve_1d_axisymmetric_packed_scatter( grid[grid_v_lookup, grid_u_lookup] += \ conv_v * conv_u * \ scp.corr2stokes(vis[r, c, :], - literally(stokes_conversion_policy)) + stokes_conversion_policy) cw += conv_v * conv_u return cw @@ -142,7 +141,7 @@ def convolve_nn_scatter(scaled_u, scaled_v, scaled_w, npix, grid, vis, r, c, cw = 1.0 grid[disc_v, disc_u] += \ scp.corr2stokes(vis[r, c, :], - literally(stokes_conversion_policy)) + stokes_conversion_policy) return cw @@ -204,7 +203,7 @@ def convolve_1d_axisymmetric_packed_gather(scaled_u, scaled_v, scaled_w, npix, grid[disc_v + tv - convolution_kernel_width // 2, disc_u + tu - convolution_kernel_width // 2] * conv_v * conv_u, vis[r, c, :], - policy_type=literally(stokes_conversion_policy)) + policy_type=stokes_conversion_policy) cw += conv_v * conv_u vis[r, c, :] /= cw + 1.0e-8 @@ -253,7 +252,7 @@ def convolve_1d_axisymmetric_unpacked_gather( scp.stokes2corr( grid[grid_v_lookup, grid_u_lookup] * conv_v * conv_u, vis[r, c, :], - policy_type=literally(stokes_conversion_policy)) + policy_type=stokes_conversion_policy) cw += conv_v * conv_u vis[r, c, :] /= cw + 1.0e-8 From c3dd06f5b6c48fd1fb6c14b150a378a745b519fa Mon Sep 17 00:00:00 2001 From: bennahugo Date: Wed, 23 Sep 2020 12:48:26 +0200 Subject: [PATCH 03/12] bump numba dependency version --- setup.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 324434fce..900aaeb90 100644 --- a/setup.py +++ b/setup.py @@ -21,8 +21,7 @@ # astropy breaks with numpy 1.15.3 # https://github.com/astropy/astropy/issues/7943 'numpy >= 1.14.0, != 1.15.3', - # version 0.5 causes issues with - 'numba <= 0.49.0' + 'numba >= 0.49.0' ] extras_require = { From 44a6d698ae9f2509f8357367e3926ba6e478a7fd Mon Sep 17 00:00:00 2001 From: bennahugo Date: Wed, 23 Sep 2020 14:34:22 +0200 Subject: [PATCH 04/12] Fix issues raised in review --- HISTORY.rst | 6 ++- .../gridding/perleypolyhedron/degridder.py | 42 +++++++++---------- .../gridding/perleypolyhedron/gridder.py | 20 ++++----- .../gridding/perleypolyhedron/kernels.py | 2 +- .../perleypolyhedron/tests/test_daskintrf.py | 11 ++--- .../perleypolyhedron/tests/test_ppgridder.py | 13 +++--- 6 files changed, 49 insertions(+), 45 deletions(-) diff --git a/HISTORY.rst b/HISTORY.rst index 368439da4..bb4a14b3c 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -2,12 +2,14 @@ History ======= +0.2.7 (YYYY-MM-DD) +---------------------------- +* Add Perley Polyhedron Faceting Gridder/Degridder (:pr:`202`, :pr:`215`, :pr:`222`) + 0.2.6 (2020-08-07) ------------------ * Add wrapper for ducc0.wgridder (:pr:204`) * Correct Irregular Grid nesting in BeamAxes (:pr:`217`) -* Add Perley Polyhedron Faceting Gridder/Degridder (:pr:`202`, :pr:`215`) - 0.2.5 (2020-07-01) ------------------ diff --git a/africanus/gridding/perleypolyhedron/degridder.py b/africanus/gridding/perleypolyhedron/degridder.py index efcbddb09..d8d35cbf1 100644 --- a/africanus/gridding/perleypolyhedron/degridder.py +++ b/africanus/gridding/perleypolyhedron/degridder.py @@ -15,7 +15,7 @@ @jit(nopython=True, nogil=True, fastmath=True, inline="always") def degridder_row_kernel(uvw, gridstack, - lambdas, + wavelengths, chanmap, cell, image_centre, @@ -41,9 +41,9 @@ def degridder_row_kernel(uvw, btp.policy(uvw[r, :], ra, dec, ra0, dec0, baseline_transform_policy) for c in range(nvischan): - scaled_u = uvw[r, 0] * scale_factor / lambdas[c] - scaled_v = uvw[r, 1] * scale_factor / lambdas[c] - scaled_w = uvw[r, 2] * scale_factor / lambdas[c] + scaled_u = uvw[r, 0] * scale_factor / wavelengths[c] + scaled_v = uvw[r, 1] * scale_factor / wavelengths[c] + scaled_w = uvw[r, 2] * scale_factor / wavelengths[c] grid = gridstack[chanmap[c], :, :] cp.policy(scaled_u, scaled_v, @@ -60,7 +60,7 @@ def degridder_row_kernel(uvw, policy_type=convolution_policy) ptp.policy(vis[r, :, :], uvw[r, :], - lambdas, + wavelengths, ra0, dec0, ra, @@ -72,7 +72,7 @@ def degridder_row_kernel(uvw, @jit(nopython=True, nogil=True, fastmath=True, parallel=True) def degridder(uvw, gridstack, - lambdas, + wavelengths, chanmap, cell, image_centre, @@ -89,7 +89,7 @@ def degridder(uvw, 2D Convolutional degridder, discrete to contiguous @uvw: value coordinates, (nrow, 3) @gridstack: complex gridded data, (nband, npix, npix) - @lambdas: wavelengths of data channels + @wavelengths: wavelengths of data channels @chanmap: MFS band mapping per channel @cell: cell_size in degrees @image_centres: new phase centre of image (radians, ra, dec) @@ -113,17 +113,17 @@ def degridder(uvw, @vis_dtype: accumulation vis dtype (default complex 128) """ - if chanmap.size != lambdas.size: + if chanmap.size != wavelengths.size: raise ValueError( - "Chanmap and corresponding lambdas must match in shape") + "Chanmap and corresponding wavelengths must match in shape") chanmap = chanmap.ravel() - lambdas = lambdas.ravel() + wavelengths = wavelengths.ravel() nband = np.max(chanmap) + 1 nrow = uvw.shape[0] npix = gridstack.shape[1] if gridstack.shape[1] != gridstack.shape[2]: raise ValueError("Grid must be square") - nvischan = lambdas.size + nvischan = wavelengths.size ncorr = scp.ncorr_out(policy_type=literally(stokes_conversion_policy)) if gridstack.shape[0] < nband: raise ValueError( @@ -133,7 +133,7 @@ def degridder(uvw, if uvw.shape[0] != nrow: raise ValueError( "UVW array must have same number of rows as vis array") - if nvischan != lambdas.size: + if nvischan != wavelengths.size: raise ValueError("Chanmap must correspond to visibility channels") vis = np.zeros((nrow, nvischan, ncorr), dtype=vis_dtype) @@ -143,7 +143,7 @@ def degridder(uvw, for r in prange(nrow): degridder_row_kernel(uvw, gridstack, - lambdas, + wavelengths, chanmap, cell, image_centre, @@ -170,7 +170,7 @@ def degridder(uvw, @jit(nopython=True, nogil=True, fastmath=True, parallel=False) def degridder_serial(uvw, gridstack, - lambdas, + wavelengths, chanmap, cell, image_centre, @@ -187,7 +187,7 @@ def degridder_serial(uvw, 2D Convolutional degridder, discrete to contiguous @uvw: value coordinates, (nrow, 3) @gridstack: complex gridded data, (nband, npix, npix) - @lambdas: wavelengths of data channels + @wavelengths: wavelengths of data channels @chanmap: MFS band mapping per channel @cell: cell_size in degrees @image_centres: new phase centre of image (radians, ra, dec) @@ -211,17 +211,17 @@ def degridder_serial(uvw, @vis_dtype: accumulation vis dtype (default complex 128) """ - if chanmap.size != lambdas.size: + if chanmap.size != wavelengths.size: raise ValueError( - "Chanmap and corresponding lambdas must match in shape") + "Chanmap and corresponding wavelengths must match in shape") chanmap = chanmap.ravel() - lambdas = lambdas.ravel() + wavelengths = wavelengths.ravel() nband = np.max(chanmap) + 1 nrow = uvw.shape[0] npix = gridstack.shape[1] if gridstack.shape[1] != gridstack.shape[2]: raise ValueError("Grid must be square") - nvischan = lambdas.size + nvischan = wavelengths.size ncorr = scp.ncorr_out(policy_type=literally(stokes_conversion_policy)) if gridstack.shape[0] < nband: raise ValueError( @@ -231,7 +231,7 @@ def degridder_serial(uvw, if uvw.shape[0] != nrow: raise ValueError( "UVW array must have same number of rows as vis array") - if nvischan != lambdas.size: + if nvischan != wavelengths.size: raise ValueError("Chanmap must correspond to visibility channels") vis = np.zeros((nrow, nvischan, ncorr), dtype=vis_dtype) @@ -241,7 +241,7 @@ def degridder_serial(uvw, for r in range(nrow): degridder_row_kernel(uvw, gridstack, - lambdas, + wavelengths, chanmap, cell, image_centre, diff --git a/africanus/gridding/perleypolyhedron/gridder.py b/africanus/gridding/perleypolyhedron/gridder.py index ed4826d89..0acbeef6b 100644 --- a/africanus/gridding/perleypolyhedron/gridder.py +++ b/africanus/gridding/perleypolyhedron/gridder.py @@ -13,7 +13,7 @@ @jit(nopython=True, nogil=True, fastmath=True, parallel=False) def gridder(uvw, vis, - lambdas, + wavelengths, chanmap, npix, cell, @@ -32,7 +32,7 @@ def gridder(uvw, 2D Convolutional gridder, contiguous to discrete @uvw: value coordinates, (nrow, 3) @vis: complex data, (nrow, nchan, ncorr) - @lambdas: wavelengths of data channels + @wavelengths: wavelengths of data channels @chanmap: MFS band mapping @npix: number of pixels per axis @cell: cell_size in degrees @@ -57,11 +57,11 @@ def gridder(uvw, @grid_dtype: accumulation grid dtype (default complex 128) @do_normalize: normalize grid by convolution weights """ - if chanmap.size != lambdas.size: + if chanmap.size != wavelengths.size: raise ValueError( - "Chanmap and corresponding lambdas must match in shape") + "Chanmap and corresponding wavelengths must match in shape") chanmap = chanmap.ravel() - lambdas = lambdas.ravel() + wavelengths = wavelengths.ravel() nband = np.max(chanmap) + 1 nrow, nvischan, ncorr = vis.shape if uvw.shape[1] != 3: @@ -69,7 +69,7 @@ def gridder(uvw, if uvw.shape[0] != nrow: raise ValueError( "UVW array must have same number of rows as vis array") - if nvischan != lambdas.size: + if nvischan != wavelengths.size: raise ValueError("Chanmap must correspond to visibility channels") gridstack = np.zeros((nband, npix, npix), dtype=grid_dtype) @@ -82,7 +82,7 @@ def gridder(uvw, ra, dec = image_centre ptp.policy(vis[r, :, :], uvw[r, :], - lambdas, + wavelengths, ra0, dec0, ra, @@ -92,9 +92,9 @@ def gridder(uvw, btp.policy(uvw[r, :], ra0, dec0, ra, dec, literally(baseline_transform_policy)) for c in range(nvischan): - scaled_u = uvw[r, 0] * scale_factor / lambdas[c] - scaled_v = uvw[r, 1] * scale_factor / lambdas[c] - scaled_w = uvw[r, 2] * scale_factor / lambdas[c] + scaled_u = uvw[r, 0] * scale_factor / wavelengths[c] + scaled_v = uvw[r, 1] * scale_factor / wavelengths[c] + scaled_w = uvw[r, 2] * scale_factor / wavelengths[c] grid = gridstack[chanmap[c], :, :] wt_ch[chanmap[c]] += cp.policy( scaled_u, diff --git a/africanus/gridding/perleypolyhedron/kernels.py b/africanus/gridding/perleypolyhedron/kernels.py index 9fee2e304..66be95354 100644 --- a/africanus/gridding/perleypolyhedron/kernels.py +++ b/africanus/gridding/perleypolyhedron/kernels.py @@ -65,7 +65,7 @@ def kbsinc(W, b=None, oversample=5, order=15): return res / np.sum(res) -def hanningsinc(W, a=0.5, oversample=5): +def hanningsinc(W, a=None, oversample=5): """ Basic hanning windowed sinc """ diff --git a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py index 2d08b5499..3a6ac95ae 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py +++ b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py @@ -9,6 +9,7 @@ degridder) from africanus.gridding.perleypolyhedron import dask as dwrap from africanus.dft.kernels import im_to_vis +from africanus.constants import c as lightspeed class clock: @@ -72,7 +73,7 @@ def test_gridder_dask(): nchan = 128 frequency = da.from_array(np.linspace(1.0e9, 1.4e9, nchan), chunks=(nchan, )) - wavelength = 299792458.0 / frequency + wavelength = lightspeed / frequency cell = da.rad2deg( wavelength[0] / (max(da.max(da.absolute(uvw[:, 0])), @@ -153,7 +154,7 @@ def test_gridder_nondask(): pxacrossbeam = 5 nchan = 128 frequency = np.linspace(1.0e9, 1.4e9, nchan) - wavelength = 299792458.0 / frequency + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / (max(np.max(np.absolute(uvw[:, 0])), @@ -215,7 +216,7 @@ def test_degrid_dft_packed_nondask(): pxacrossbeam = 10 nchan = 1024 frequency = np.array(np.linspace(1.0e9, 1.4e9, nchan)) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / @@ -267,7 +268,7 @@ def test_degrid_dft_packed_dask(): pxacrossbeam = 10 nchan = 1024 frequency = np.array(np.linspace(1.0e9, 1.4e9, nchan)) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / @@ -321,7 +322,7 @@ def test_degrid_dft_packed_dask_dft_check(): pxacrossbeam = 10 nchan = 16 frequency = np.array(np.linspace(1.0e9, 1.4e9, nchan)) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index bd02befab..2cc7d2377 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -10,6 +10,7 @@ degridder) from africanus.dft.kernels import im_to_vis, vis_to_im from africanus.coordinates import radec_to_lmn +from africanus.constants import c as lightspeed def test_construct_kernels(tmp_path_factory): @@ -120,7 +121,7 @@ def test_degrid_dft(tmp_path_factory): pxacrossbeam = 10 frequency = np.array([1.4e9]) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / @@ -207,7 +208,7 @@ def test_degrid_dft_packed(tmp_path_factory): pxacrossbeam = 10 frequency = np.array([1.4e9]) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / @@ -328,7 +329,7 @@ def test_grid_dft(tmp_path_factory): pxacrossbeam = 10 frequency = np.array([30.0e9]) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / @@ -430,7 +431,7 @@ def test_grid_dft_packed(tmp_path_factory): pxacrossbeam = 10 frequency = np.array([30.0e9]) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / @@ -545,7 +546,7 @@ def test_wcorrection_faceting_backward(tmp_path_factory): pxacrossbeam = 5 frequency = np.array([1.4e9]) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / @@ -672,7 +673,7 @@ def test_wcorrection_faceting_forward(tmp_path_factory): pxacrossbeam = 5 frequency = np.array([1.4e9]) - wavelength = np.array([299792458.0 / f for f in frequency]) + wavelength = lightspeed / frequency cell = np.rad2deg( wavelength[0] / From 7069c3ecc00879e37d1421f84f37933234b26603 Mon Sep 17 00:00:00 2001 From: bennahugo Date: Wed, 23 Sep 2020 16:49:01 +0200 Subject: [PATCH 05/12] Reduce test length --- .../perleypolyhedron/tests/test_daskintrf.py | 12 +++++------ .../perleypolyhedron/tests/test_ppgridder.py | 20 +++++++++---------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py index 3a6ac95ae..bda28d600 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py +++ b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py @@ -48,7 +48,7 @@ def test_gridder_dask(): W = 5 OS = 9 kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) - nrow = int(1e6) + nrow = int(1e5) np.random.seed(0) # simulate some ficticious baselines rotated by an hour angle row_chunks = nrow // 10 @@ -70,7 +70,7 @@ def test_gridder_dask(): uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T) uvw = da.from_array(uvw, chunks=(row_chunks, 3)) pxacrossbeam = 5 - nchan = 128 + nchan = 64 frequency = da.from_array(np.linspace(1.0e9, 1.4e9, nchan), chunks=(nchan, )) wavelength = lightspeed / frequency @@ -132,7 +132,7 @@ def test_gridder_nondask(): W = 5 OS = 9 kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) - nrow = int(1e6) + nrow = int(1e5) np.random.seed(0) # simulate some ficticious baselines rotated by an hour angle uvw = np.zeros((nrow, 3), dtype=np.float64) @@ -152,7 +152,7 @@ def test_gridder_nondask(): s(d0)]]) uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T) pxacrossbeam = 5 - nchan = 128 + nchan = 64 frequency = np.linspace(1.0e9, 1.4e9, nchan) wavelength = lightspeed / frequency cell = np.rad2deg( @@ -214,7 +214,7 @@ def test_degrid_dft_packed_nondask(): 5000.0 * np.sin(np.linspace(0, 2 * np.pi, nrow)), np.zeros(nrow))) pxacrossbeam = 10 - nchan = 1024 + nchan = 64 frequency = np.array(np.linspace(1.0e9, 1.4e9, nchan)) wavelength = lightspeed / frequency @@ -266,7 +266,7 @@ def test_degrid_dft_packed_dask(): 5000.0 * np.sin(np.linspace(0, 2 * np.pi, nrow)), np.zeros(nrow))) pxacrossbeam = 10 - nchan = 1024 + nchan = 64 frequency = np.array(np.linspace(1.0e9, 1.4e9, nchan)) wavelength = lightspeed / frequency diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index 2cc7d2377..b95e7db4f 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -116,8 +116,8 @@ def test_degrid_dft(tmp_path_factory): OS = 3 kern = kernels.kbsinc(W, oversample=OS) uvw = np.column_stack( - (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 1000)), - 5000.0 * np.sin(np.linspace(0, 2 * np.pi, 1000)), np.zeros(1000))) + (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 100)), + 5000.0 * np.sin(np.linspace(0, 2 * np.pi, 100)), np.zeros(100))) pxacrossbeam = 10 frequency = np.array([1.4e9]) @@ -203,8 +203,8 @@ def test_degrid_dft_packed(tmp_path_factory): W, oversample=OS) uvw = np.column_stack( - (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 1000)), - 5000.0 * np.sin(np.linspace(0, 2 * np.pi, 1000)), np.zeros(1000))) + (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 100)), + 5000.0 * np.sin(np.linspace(0, 2 * np.pi, 100)), np.zeros(100))) pxacrossbeam = 10 frequency = np.array([1.4e9]) @@ -322,7 +322,7 @@ def test_grid_dft(tmp_path_factory): W = 7 OS = 9 kern = kernels.kbsinc(W, oversample=OS) - nrow = 5000 + nrow = 1500 np.random.seed(0) uvw = np.random.normal(scale=6000, size=(nrow, 3)) uvw[:, 2] = 0.0 # ignore widefield effects for now @@ -416,7 +416,7 @@ def test_grid_dft(tmp_path_factory): "grid_diff_dft.png") assert (np.percentile(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :]), - 95.0) < 0.15) + 85.0) < 0.20) def test_grid_dft_packed(tmp_path_factory): @@ -424,7 +424,7 @@ def test_grid_dft_packed(tmp_path_factory): W = 7 OS = 1009 kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) - nrow = 5000 + nrow = 1500 np.random.seed(0) uvw = np.random.normal(scale=6000, size=(nrow, 3)) uvw[:, 2] = 0.0 # ignore widefield effects for now @@ -517,7 +517,7 @@ def test_grid_dft_packed(tmp_path_factory): "grid_diff_dft_packed.png") assert (np.percentile(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :]), - 95.0) < 0.15) + 85.0) < 0.20) def test_wcorrection_faceting_backward(tmp_path_factory): @@ -525,7 +525,7 @@ def test_wcorrection_faceting_backward(tmp_path_factory): W = 5 OS = 9 kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) - nrow = 5000 + nrow = 500 np.random.seed(0) # simulate some ficticious baselines rotated by an hour angle uvw = np.zeros((nrow, 3), dtype=np.float64) @@ -652,7 +652,7 @@ def test_wcorrection_faceting_forward(tmp_path_factory): W = 5 OS = 9 kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) - nrow = 5000 + nrow = 500 np.random.seed(0) # simulate some ficticious baselines rotated by an hour angle uvw = np.zeros((nrow, 3), dtype=np.float64) From 460c95e4a4ef97432e73719145f76f999e260672 Mon Sep 17 00:00:00 2001 From: bennahugo Date: Wed, 23 Sep 2020 17:05:14 +0200 Subject: [PATCH 06/12] fix flake --- africanus/gridding/perleypolyhedron/tests/test_daskintrf.py | 1 - africanus/gridding/perleypolyhedron/tests/test_ppgridder.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py index 46282286d..2301c2ce2 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py +++ b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py @@ -156,7 +156,6 @@ def test_gridder_nondask(): uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T) pxacrossbeam = 5 nchan = 64 - frequency = np.linspace(1.0e9, 1.4e9, nchan) wavelength = lightspeed / frequency cell = np.rad2deg( diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index e25706757..08e26ec6c 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -206,7 +206,6 @@ def test_degrid_dft_packed(tmp_path_factory): (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 100)), 5000.0 * np.sin(np.linspace(0, 2 * np.pi, 100)), np.zeros(100))) - pxacrossbeam = 10 frequency = np.array([1.4e9]) wavelength = lightspeed / frequency @@ -421,7 +420,6 @@ def test_grid_dft(tmp_path_factory): 85.0) < 0.20) - def test_grid_dft_packed(tmp_path_factory): # construct kernel W = 7 From 187a92f8710fddf689e2974189dfa252bba91640 Mon Sep 17 00:00:00 2001 From: bennahugo Date: Wed, 23 Sep 2020 18:30:30 +0200 Subject: [PATCH 07/12] further reductions to test time --- .../perleypolyhedron/tests/test_daskintrf.py | 10 +-- .../perleypolyhedron/tests/test_ppgridder.py | 88 +++++++++++++------ 2 files changed, 64 insertions(+), 34 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py index 2301c2ce2..a2af79f8f 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py +++ b/africanus/gridding/perleypolyhedron/tests/test_daskintrf.py @@ -48,7 +48,7 @@ def test_gridder_dask(): W = 5 OS = 9 kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) - nrow = int(1e5) + nrow = int(1e3) np.random.seed(0) # simulate some ficticious baselines rotated by an hour angle @@ -134,7 +134,7 @@ def test_gridder_nondask(): W = 5 OS = 9 kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS) - nrow = int(1e5) + nrow = int(1e3) np.random.seed(0) # simulate some ficticious baselines rotated by an hour angle @@ -211,7 +211,7 @@ def test_degrid_dft_packed_nondask(): kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, oversample=OS) - nrow = int(5e4) + nrow = int(5e3) uvw = np.column_stack( (5000.0 * np.cos(np.linspace(0, 2 * np.pi, nrow)), 5000.0 * np.sin(np.linspace(0, 2 * np.pi, nrow)), np.zeros(nrow))) @@ -263,7 +263,7 @@ def test_degrid_dft_packed_dask(): kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, oversample=OS) - nrow = int(5e4) + nrow = int(5e3) nrow_chunk = nrow // 32 uvw = np.column_stack( (5000.0 * np.cos(np.linspace(0, 2 * np.pi, nrow)), @@ -325,7 +325,7 @@ def test_degrid_dft_packed_dask_dft_check(): 5000.0 * np.sin(np.linspace(0, 2 * np.pi, nrow)), np.zeros(nrow))) pxacrossbeam = 10 - nchan = 16 + nchan = 1 frequency = np.linspace(1.0e9, 1.4e9, nchan) diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index 08e26ec6c..3e4580e54 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -109,8 +109,11 @@ def test_facetcodepath(): "phase_rotate", "I_FROM_XXYY", "conv_1d_axisymmetric_packed_scatter") +@pytest.fixture(scope="module") +def global_var_degrid(): + pytest.__CACHE_DEGRID_DFT = None -def test_degrid_dft(tmp_path_factory): +def test_degrid_dft(tmp_path_factory, global_var_degrid): # construct kernel W = 5 OS = 3 @@ -155,8 +158,11 @@ def test_degrid_dft(tmp_path_factory): np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) radec = np.column_stack((ra.flatten(), dec.flatten())) - vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, npix * npix).T.copy(), - uvw, radec, frequency) + if pytest.__CACHE_DEGRID_DFT is None: + pytest.__CACHE_DEGRID_DFT = im_to_vis(mod[0, :, :].reshape( + 1, 1, npix * npix).T.copy(), + uvw, radec, frequency) + vis_dft = pytest.__CACHE_DEGRID_DFT try: import matplotlib @@ -173,9 +179,8 @@ def test_degrid_dft(tmp_path_factory): plt.legend() plt.xlabel("sample") plt.ylabel("Real of predicted") - plt.savefig( - os.path.join(os.environ.get("TMPDIR", "/tmp"), - "degrid_vs_dft_re.png")) + plt.savefig(tmp_path_factory.mktemp("degrid_dft") / + "degrid_vs_dft_re.png") plt.figure() plt.plot(vis_degrid[:, 0, 0].imag, label=r"$\Im(\mathtt{degrid})$") plt.plot(vis_dft[:, 0, 0].imag, label=r"$\Im(\mathtt{dft})$") @@ -195,7 +200,7 @@ def test_degrid_dft(tmp_path_factory): 99.0) < 0.05 -def test_degrid_dft_packed(tmp_path_factory): +def test_degrid_dft_packed(tmp_path_factory, global_var_degrid): # construct kernel W = 5 OS = 3 @@ -242,8 +247,11 @@ def test_degrid_dft_packed(tmp_path_factory): np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) radec = np.column_stack((ra.flatten(), dec.flatten())) - vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, npix * npix).T.copy(), - uvw, radec, frequency) + if pytest.__CACHE_DEGRID_DFT is None: + pytest.__CACHE_DEGRID_DFT = im_to_vis(mod[0, :, :].reshape( + 1, 1, npix * npix).T.copy(), + uvw, radec, frequency) + vis_dft = pytest.__CACHE_DEGRID_DFT try: import matplotlib @@ -260,9 +268,8 @@ def test_degrid_dft_packed(tmp_path_factory): plt.legend() plt.xlabel("sample") plt.ylabel("Real of predicted") - plt.savefig( - os.path.join(os.environ.get("TMPDIR", "/tmp"), - "degrid_vs_dft_re_packed.png")) + plt.savefig(tmp_path_factory.mktemp("degrid_dft_packed") / + "degrid_vs_dft_re_packed.png") plt.figure() plt.plot(vis_degrid[:, 0, 0].imag, label=r"$\Im(\mathtt{degrid})$") plt.plot(vis_dft[:, 0, 0].imag, label=r"$\Im(\mathtt{dft})$") @@ -317,7 +324,13 @@ def test_detaper(tmp_path_factory): assert (np.max(np.abs(detaperdft - detaperdftsep)) < 1.0e-14) -def test_grid_dft(tmp_path_factory): +@pytest.fixture(scope="module") +def global_vars_grid(): + pytest.__CACHE_GRID_MOD = None + pytest.__CACHE_GRID_DFT = None + + +def test_grid_dft(tmp_path_factory, global_vars_grid): # construct kernel W = 7 OS = 9 @@ -354,9 +367,14 @@ def test_grid_dft(tmp_path_factory): np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) radec = np.column_stack((ra.flatten(), dec.flatten())) - vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, - npix * npix).T.copy(), uvw, - radec, frequency).repeat(2).reshape(nrow, 1, 2) + if pytest.__CACHE_GRID_MOD is None: + pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :]\ + .reshape(1, 1, + npix * + npix).T.copy(), uvw, + radec, frequency)\ + .repeat(2).reshape(nrow, 1, 2) + vis_dft = pytest.__CACHE_GRID_MOD chanmap = np.array([0]) detaper = kernels.compute_detaper(int(npix * fftpad), @@ -388,11 +406,14 @@ def test_grid_dft(tmp_path_factory): npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix, int(npix * fftpad) // 2 - npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix] - dftvis = vis_to_im(vis_dft, uvw, radec, frequency, - np.zeros(vis_dft.shape, - dtype=np.bool)).T.copy().reshape( - 2, 1, npix, npix) / nrow - + if pytest.__CACHE_GRID_DFT is None: + pytest.__CACHE_GRID_DFT = \ + vis_to_im(vis_dft, uvw, radec, frequency, + np.zeros(vis_dft.shape, + dtype=np.bool))\ + .T.copy().reshape(2, 1, + npix, npix) / nrow + dftvis = pytest.__CACHE_GRID_DFT try: import matplotlib except ImportError: @@ -420,7 +441,7 @@ def test_grid_dft(tmp_path_factory): 85.0) < 0.20) -def test_grid_dft_packed(tmp_path_factory): +def test_grid_dft_packed(tmp_path_factory, global_vars_grid): # construct kernel W = 7 OS = 1009 @@ -457,9 +478,14 @@ def test_grid_dft_packed(tmp_path_factory): np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) radec = np.column_stack((ra.flatten(), dec.flatten())) - vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, - npix * npix).T.copy(), uvw, - radec, frequency).repeat(2).reshape(nrow, 1, 2) + if pytest.__CACHE_GRID_MOD is None: + pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :]\ + .reshape(1, 1, + npix * + npix).T.copy(), uvw, + radec, frequency)\ + .repeat(2).reshape(nrow, 1, 2) + vis_dft = pytest.__CACHE_GRID_MOD chanmap = np.array([0]) detaper = kernels.compute_detaper_dft_seperable( int(npix * fftpad), kernels.unpack_kernel(kern, W, OS), W, OS) @@ -490,10 +516,14 @@ def test_grid_dft_packed(tmp_path_factory): npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix, int(npix * fftpad) // 2 - npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix] - dftvis = vis_to_im(vis_dft, uvw, radec, frequency, - np.zeros(vis_dft.shape, - dtype=np.bool)).T.copy().reshape( - 2, 1, npix, npix) / nrow + if pytest.__CACHE_GRID_DFT is None: + pytest.__CACHE_GRID_DFT = \ + vis_to_im(vis_dft, uvw, radec, frequency, + np.zeros(vis_dft.shape, + dtype=np.bool))\ + .T.copy().reshape(2, 1, + npix, npix) / nrow + dftvis = pytest.__CACHE_GRID_DFT try: import matplotlib From 37882aad334f5df59e0164dff1e042bd17d215a8 Mon Sep 17 00:00:00 2001 From: bennahugo Date: Wed, 23 Sep 2020 18:43:24 +0200 Subject: [PATCH 08/12] fix flake --- .../perleypolyhedron/tests/test_ppgridder.py | 39 +++++++++---------- 1 file changed, 18 insertions(+), 21 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index 3e4580e54..125a97224 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -1,7 +1,4 @@ # -*- coding: utf-8 -*- - -import os - import numpy as np import pytest @@ -109,10 +106,12 @@ def test_facetcodepath(): "phase_rotate", "I_FROM_XXYY", "conv_1d_axisymmetric_packed_scatter") + @pytest.fixture(scope="module") def global_var_degrid(): pytest.__CACHE_DEGRID_DFT = None + def test_degrid_dft(tmp_path_factory, global_var_degrid): # construct kernel W = 5 @@ -368,12 +367,11 @@ def test_grid_dft(tmp_path_factory, global_vars_grid): radec = np.column_stack((ra.flatten(), dec.flatten())) if pytest.__CACHE_GRID_MOD is None: - pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :]\ - .reshape(1, 1, - npix * - npix).T.copy(), uvw, - radec, frequency)\ - .repeat(2).reshape(nrow, 1, 2) + pytest.__CACHE_GRID_MOD = \ + im_to_vis(mod[0, :, :].reshape(1, 1, + npix * + npix).T.copy(), + uvw, radec, frequency).repeat(2).reshape(nrow, 1, 2) vis_dft = pytest.__CACHE_GRID_MOD chanmap = np.array([0]) @@ -410,9 +408,9 @@ def test_grid_dft(tmp_path_factory, global_vars_grid): pytest.__CACHE_GRID_DFT = \ vis_to_im(vis_dft, uvw, radec, frequency, np.zeros(vis_dft.shape, - dtype=np.bool))\ - .T.copy().reshape(2, 1, - npix, npix) / nrow + dtype=np.bool)).T.copy().reshape(2, 1, + npix, + npix) / nrow dftvis = pytest.__CACHE_GRID_DFT try: import matplotlib @@ -479,12 +477,11 @@ def test_grid_dft_packed(tmp_path_factory, global_vars_grid): radec = np.column_stack((ra.flatten(), dec.flatten())) if pytest.__CACHE_GRID_MOD is None: - pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :]\ - .reshape(1, 1, - npix * - npix).T.copy(), uvw, - radec, frequency)\ - .repeat(2).reshape(nrow, 1, 2) + pytest.__CACHE_GRID_MOD = \ + im_to_vis(mod[0, :, :].reshape(1, 1, + npix * + npix).T.copy(), + uvw, radec, frequency).repeat(2).reshape(nrow, 1, 2) vis_dft = pytest.__CACHE_GRID_MOD chanmap = np.array([0]) detaper = kernels.compute_detaper_dft_seperable( @@ -520,9 +517,9 @@ def test_grid_dft_packed(tmp_path_factory, global_vars_grid): pytest.__CACHE_GRID_DFT = \ vis_to_im(vis_dft, uvw, radec, frequency, np.zeros(vis_dft.shape, - dtype=np.bool))\ - .T.copy().reshape(2, 1, - npix, npix) / nrow + dtype=np.bool)).T.copy().reshape(2, 1, + npix, + npix) / nrow dftvis = pytest.__CACHE_GRID_DFT try: From f440752df4706dcf557a249239bbae5b4fd7e032 Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Wed, 30 Sep 2020 14:14:00 +0200 Subject: [PATCH 09/12] Implement vis_dft with fixtures --- .../perleypolyhedron/tests/test_ppgridder.py | 152 ++++++++++++------ 1 file changed, 103 insertions(+), 49 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index 125a97224..975a988be 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -10,6 +10,70 @@ from africanus.constants import c as lightspeed +@pytest.fixture(scope="module", params=[512]) +def npix(request): + return request.param + + +@pytest.fixture(scope="module", params=[3]) +def oversampling(request): + return request.param + + +@pytest.fixture(scope="module", params=[5]) +def kernel_width(request): + return request.param + + +@pytest.fixture(scope="module", params=[np.array([1.4e9])]) +def frequency(request): + return request.param + + +@pytest.fixture(scope="module") +def wavelength(frequency): + return lightspeed / frequency + + +@pytest.fixture(scope="module", params=[10]) +def pxacrossbeam(request): + return request.param + + +@pytest.fixture(scope="module", params=[100]) +def uvw(request): + nrow = request.param + x = np.linspace(0, 2 * np.pi, nrow) + + return np.column_stack(( + 5000.0 * np.cos(x), + 5000.0 * np.sin(x), + np.zeros(nrow)) + ).astype(np.float64) + + +@pytest.fixture(scope="module") +def image(npix, kernel_width): + W = kernel_width + mod = np.zeros((1, npix, npix), dtype=np.complex64) + mod[0, npix // 2 - W, npix // 2 - W] = 1.0 + return mod + + +@pytest.fixture(scope="module") +def ftimage(image): + _, npix, _ = image.shape + return np.fft.ifftshift(np.fft.fft2(np.fft.fftshift( + image[0, :, :]))).reshape((1, npix, npix)) + + +@pytest.fixture(scope="module") +def cell_size(request, uvw, wavelength, pxacrossbeam): + return np.rad2deg(wavelength[0] / + (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * + pxacrossbeam)) + + def test_construct_kernels(tmp_path_factory): matplotlib = pytest.importorskip("matplotlib") matplotlib.use("agg") @@ -112,36 +176,35 @@ def global_var_degrid(): pytest.__CACHE_DEGRID_DFT = None -def test_degrid_dft(tmp_path_factory, global_var_degrid): - # construct kernel - W = 5 - OS = 3 - kern = kernels.kbsinc(W, oversample=OS) - uvw = np.column_stack( - (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 100)), - 5000.0 * np.sin(np.linspace(0, 2 * np.pi, 100)), np.zeros(100))) +@pytest.fixture(scope="module") +def vis_dft(request, image, cell_size, uvw, frequency): + _, npix, _ = image.shape + extent = np.arange(-npix // 2, npix // 2) * np.deg2rad(cell_size) + ra, dec = np.meshgrid(extent, extent) + radec = np.column_stack((ra.flatten(), dec.flatten())) + return im_to_vis(image[0, :, :].reshape(1, 1, npix * npix).T.copy(), + uvw, radec, frequency) - pxacrossbeam = 10 - frequency = np.array([1.4e9]) - wavelength = lightspeed / frequency - cell = np.rad2deg( - wavelength[0] / - (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * - pxacrossbeam)) - npix = 512 - mod = np.zeros((1, npix, npix), dtype=np.complex64) - mod[0, npix // 2 - 5, npix // 2 - 5] = 1.0 +# We can indirectly parametrize the number of rows in the uvw +# fixture like so: +# @pytest.mark.parametrize("uvw", [172], indirect=True) +def test_degrid_dft(npix, oversampling, kernel_width, + frequency, wavelength, pxacrossbeam, uvw, + image, ftimage, cell_size, vis_dft, tmp_path_factory): + OS = oversampling + W = kernel_width + + # construct kernel + kern = kernels.kbsinc(W, oversample=OS) - ftmod = np.fft.ifftshift(np.fft.fft2(np.fft.fftshift( - mod[0, :, :]))).reshape((1, npix, npix)) chanmap = np.array([0]) vis_degrid = degridder.degridder( uvw, - ftmod, + ftimage, wavelength, chanmap, - cell * 3600.0, + cell_size * 3600.0, (0, np.pi / 4.0), (0, np.pi / 4.0), kern, @@ -152,17 +215,6 @@ def test_degrid_dft(tmp_path_factory, global_var_degrid): "XXYY_FROM_I", "conv_1d_axisymmetric_unpacked_gather") - dec, ra = np.meshgrid( - np.arange(-npix // 2, npix // 2) * np.deg2rad(cell), - np.arange(-npix // 2, npix // 2) * np.deg2rad(cell)) - radec = np.column_stack((ra.flatten(), dec.flatten())) - - if pytest.__CACHE_DEGRID_DFT is None: - pytest.__CACHE_DEGRID_DFT = im_to_vis(mod[0, :, :].reshape( - 1, 1, npix * npix).T.copy(), - uvw, radec, frequency) - vis_dft = pytest.__CACHE_DEGRID_DFT - try: import matplotlib except ImportError: @@ -367,11 +419,12 @@ def test_grid_dft(tmp_path_factory, global_vars_grid): radec = np.column_stack((ra.flatten(), dec.flatten())) if pytest.__CACHE_GRID_MOD is None: - pytest.__CACHE_GRID_MOD = \ - im_to_vis(mod[0, :, :].reshape(1, 1, - npix * - npix).T.copy(), - uvw, radec, frequency).repeat(2).reshape(nrow, 1, 2) + pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :]\ + .reshape(1, 1, + npix * + npix).T.copy(), uvw, + radec, frequency)\ + .repeat(2).reshape(nrow, 1, 2) vis_dft = pytest.__CACHE_GRID_MOD chanmap = np.array([0]) @@ -408,9 +461,9 @@ def test_grid_dft(tmp_path_factory, global_vars_grid): pytest.__CACHE_GRID_DFT = \ vis_to_im(vis_dft, uvw, radec, frequency, np.zeros(vis_dft.shape, - dtype=np.bool)).T.copy().reshape(2, 1, - npix, - npix) / nrow + dtype=np.bool))\ + .T.copy().reshape(2, 1, + npix, npix) / nrow dftvis = pytest.__CACHE_GRID_DFT try: import matplotlib @@ -477,11 +530,12 @@ def test_grid_dft_packed(tmp_path_factory, global_vars_grid): radec = np.column_stack((ra.flatten(), dec.flatten())) if pytest.__CACHE_GRID_MOD is None: - pytest.__CACHE_GRID_MOD = \ - im_to_vis(mod[0, :, :].reshape(1, 1, - npix * - npix).T.copy(), - uvw, radec, frequency).repeat(2).reshape(nrow, 1, 2) + pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :]\ + .reshape(1, 1, + npix * + npix).T.copy(), uvw, + radec, frequency)\ + .repeat(2).reshape(nrow, 1, 2) vis_dft = pytest.__CACHE_GRID_MOD chanmap = np.array([0]) detaper = kernels.compute_detaper_dft_seperable( @@ -517,9 +571,9 @@ def test_grid_dft_packed(tmp_path_factory, global_vars_grid): pytest.__CACHE_GRID_DFT = \ vis_to_im(vis_dft, uvw, radec, frequency, np.zeros(vis_dft.shape, - dtype=np.bool)).T.copy().reshape(2, 1, - npix, - npix) / nrow + dtype=np.bool))\ + .T.copy().reshape(2, 1, + npix, npix) / nrow dftvis = pytest.__CACHE_GRID_DFT try: From db28958c2e8e946edd18acbad5969bb73da54233 Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Wed, 30 Sep 2020 14:28:50 +0200 Subject: [PATCH 10/12] flake8 --- .../perleypolyhedron/tests/test_ppgridder.py | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index 975a988be..c52598493 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -69,9 +69,8 @@ def ftimage(image): @pytest.fixture(scope="module") def cell_size(request, uvw, wavelength, pxacrossbeam): - return np.rad2deg(wavelength[0] / - (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) * - pxacrossbeam)) + max_uv = np.max(uvw[:, :2], axis=(0, 1), keepdims=False) + return np.rad2deg(wavelength[0] / (2 * max_uv * pxacrossbeam)) def test_construct_kernels(tmp_path_factory): @@ -183,7 +182,7 @@ def vis_dft(request, image, cell_size, uvw, frequency): ra, dec = np.meshgrid(extent, extent) radec = np.column_stack((ra.flatten(), dec.flatten())) return im_to_vis(image[0, :, :].reshape(1, 1, npix * npix).T.copy(), - uvw, radec, frequency) + uvw, radec, frequency) # We can indirectly parametrize the number of rows in the uvw @@ -419,11 +418,11 @@ def test_grid_dft(tmp_path_factory, global_vars_grid): radec = np.column_stack((ra.flatten(), dec.flatten())) if pytest.__CACHE_GRID_MOD is None: - pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :]\ + pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :] .reshape(1, 1, npix * npix).T.copy(), uvw, - radec, frequency)\ + radec, frequency)\ .repeat(2).reshape(nrow, 1, 2) vis_dft = pytest.__CACHE_GRID_MOD chanmap = np.array([0]) @@ -462,8 +461,8 @@ def test_grid_dft(tmp_path_factory, global_vars_grid): vis_to_im(vis_dft, uvw, radec, frequency, np.zeros(vis_dft.shape, dtype=np.bool))\ - .T.copy().reshape(2, 1, - npix, npix) / nrow + .T.copy().reshape(2, 1, + npix, npix) / nrow dftvis = pytest.__CACHE_GRID_DFT try: import matplotlib @@ -530,11 +529,11 @@ def test_grid_dft_packed(tmp_path_factory, global_vars_grid): radec = np.column_stack((ra.flatten(), dec.flatten())) if pytest.__CACHE_GRID_MOD is None: - pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :]\ + pytest.__CACHE_GRID_MOD = im_to_vis(mod[0, :, :] .reshape(1, 1, npix * npix).T.copy(), uvw, - radec, frequency)\ + radec, frequency)\ .repeat(2).reshape(nrow, 1, 2) vis_dft = pytest.__CACHE_GRID_MOD chanmap = np.array([0]) @@ -572,8 +571,8 @@ def test_grid_dft_packed(tmp_path_factory, global_vars_grid): vis_to_im(vis_dft, uvw, radec, frequency, np.zeros(vis_dft.shape, dtype=np.bool))\ - .T.copy().reshape(2, 1, - npix, npix) / nrow + .T.copy().reshape(2, 1, + npix, npix) / nrow dftvis = pytest.__CACHE_GRID_DFT try: From 42da370136a8c76da753ec2bf87be374b4cb052d Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Wed, 30 Sep 2020 15:01:01 +0200 Subject: [PATCH 11/12] Use full variable name --- .../gridding/perleypolyhedron/tests/test_ppgridder.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index c52598493..3259c8a78 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -191,11 +191,9 @@ def vis_dft(request, image, cell_size, uvw, frequency): def test_degrid_dft(npix, oversampling, kernel_width, frequency, wavelength, pxacrossbeam, uvw, image, ftimage, cell_size, vis_dft, tmp_path_factory): - OS = oversampling - W = kernel_width # construct kernel - kern = kernels.kbsinc(W, oversample=OS) + kern = kernels.kbsinc(kernel_width, oversample=oversampling) chanmap = np.array([0]) vis_degrid = degridder.degridder( @@ -207,8 +205,8 @@ def test_degrid_dft(npix, oversampling, kernel_width, (0, np.pi / 4.0), (0, np.pi / 4.0), kern, - W, - OS, + kernel_width, + oversampling, "None", # no faceting "None", # no faceting "XXYY_FROM_I", From 931ddaaa6e2e331e0f13914e936d11e0ac83b8fc Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Wed, 30 Sep 2020 17:41:07 +0200 Subject: [PATCH 12/12] further simplifications + naming corrections --- .../perleypolyhedron/tests/test_ppgridder.py | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py index 3259c8a78..bfb723e3b 100644 --- a/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py +++ b/africanus/gridding/perleypolyhedron/tests/test_ppgridder.py @@ -53,7 +53,7 @@ def uvw(request): @pytest.fixture(scope="module") -def image(npix, kernel_width): +def model(npix, kernel_width): W = kernel_width mod = np.zeros((1, npix, npix), dtype=np.complex64) mod[0, npix // 2 - W, npix // 2 - W] = 1.0 @@ -61,10 +61,10 @@ def image(npix, kernel_width): @pytest.fixture(scope="module") -def ftimage(image): - _, npix, _ = image.shape +def ftmodel(model): + _, npix, _ = model.shape return np.fft.ifftshift(np.fft.fft2(np.fft.fftshift( - image[0, :, :]))).reshape((1, npix, npix)) + model[0, :, :]))).reshape((1, npix, npix)) @pytest.fixture(scope="module") @@ -176,21 +176,24 @@ def global_var_degrid(): @pytest.fixture(scope="module") -def vis_dft(request, image, cell_size, uvw, frequency): - _, npix, _ = image.shape +def vis_dft(request, model, cell_size, uvw, frequency): + _, npix, _ = model.shape extent = np.arange(-npix // 2, npix // 2) * np.deg2rad(cell_size) ra, dec = np.meshgrid(extent, extent) radec = np.column_stack((ra.flatten(), dec.flatten())) - return im_to_vis(image[0, :, :].reshape(1, 1, npix * npix).T.copy(), + return im_to_vis(model[0, :, :].reshape(1, 1, npix * npix).T.copy(), uvw, radec, frequency) -# We can indirectly parametrize the number of rows in the uvw -# fixture like so: -# @pytest.mark.parametrize("uvw", [172], indirect=True) -def test_degrid_dft(npix, oversampling, kernel_width, - frequency, wavelength, pxacrossbeam, uvw, - image, ftimage, cell_size, vis_dft, tmp_path_factory): +# We can indirectly parametrize the number of pixels +# in the model and ftmodel like so +# @pytest.mark.parametrize("npix", [256], indirect=True) +# or the number of UVW rows like so +# @pytest.mark.parametrize("uvw", [128], indirect=True) +def test_degrid_dft(oversampling, kernel_width, + wavelength, uvw, + ftmodel, cell_size, + vis_dft, tmp_path_factory): # construct kernel kern = kernels.kbsinc(kernel_width, oversample=oversampling) @@ -198,7 +201,7 @@ def test_degrid_dft(npix, oversampling, kernel_width, chanmap = np.array([0]) vis_degrid = degridder.degridder( uvw, - ftimage, + ftmodel, wavelength, chanmap, cell_size * 3600.0,