test.py

import nifty_gridder as ng
import numpy as np
import pytest
from numpy.testing import assert_, assert_allclose, assert_array_almost_equal

pmp = pytest.mark.parametrize


def _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow):
    pixsize = np.pi/180/60/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=0.368*pixsize,
                            pixsize_y=pixsize)
    speedoflight, f0 = 3e8, 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
    baselines = ng.Baselines(coord=uvw, freq=freq)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
    idx = ng.getIndices(baselines, conf, flags)
    return baselines, conf, idx


def _dft(uvw, vis, conf, apply_w):
    shp = (conf.Nxdirty(), conf.Nydirty())
    x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in shp],
                       indexing='ij')
    x *= conf.Pixsize_x()
    y *= conf.Pixsize_y()
    dirty = np.zeros(shp, dtype=np.complex128)
    if apply_w:
        n = np.sqrt(1-x**2-y**2)
        nm1 = n-1
    for ii, vv in enumerate(vis):
        phase = x*uvw[ii, 0] + y*uvw[ii, 1]
        if apply_w:
            phase += uvw[ii, 2]*nm1
        dirty += (vv*np.exp(2j*np.pi*phase))
    if apply_w:
        return dirty/n
    return dirty


def _l2error(a, b):
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))


@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 10, 10000))
@pmp("nchan", (1, 10, 100))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
    np.random.seed(42)
    bl, gconf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    grid = ng.vis2grid_c(bl, gconf, idx, vis)
    dirty = gconf.grid2dirty_c(grid)
    dirty2 = np.random.rand(*dirty.shape)-0.5
    ms2 = bl.vis2ms(ng.grid2vis_c(bl, gconf, idx, gconf.dirty2grid_c(dirty2)),
                    idx)
    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty.real, dirty2.real))


@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 10, 10000))
@pmp("nchan", (1, 10, 100))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
def test_adjointness_wgridding(nxdirty, nydirty, nrow, nchan, epsilon):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    vis = np.random.rand(*idx.shape)-0.5 + 1j*(np.random.rand(*idx.shape)-0.5)
    dirty = np.random.rand(conf.Nxdirty(), conf.Nydirty())-0.5
    dirty2 = ng.vis2dirty_wstack(bl, conf, idx, vis)
    vis2 = ng.dirty2vis_wstack(bl, conf, idx, dirty)
    assert_allclose(np.vdot(vis, vis2), np.vdot(dirty2, dirty),
                    rtol=epsilon)


@pmp("nxdirty", (128,))
@pmp("nydirty", (128,))
@pmp("nrow", (10000,))
@pmp("nchan", (100,))
@pmp("epsilon", (2e-13,))
def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
    np.random.seed(42)
    gconf = ng.GridderConfig(1024, 1024, 2e-13, 2.0, 2.0)
    freq = np.linspace(.856e9, 2*.856e9, nchan)
    f0 = freq[freq.shape[0]//2]
    speedoflight = 3e8

    pixsize = np.pi/180/60/nxdirty  # assume 1 arcmin FOV
    speedoflight = 3e8
    uvw = (np.random.rand(nrow, 3)-0.5) / (pixsize*f0/speedoflight)
    baselines = ng.Baselines(coord=uvw, freq=freq)
    gconf = ng.GridderConfig(nxdirty=nxdirty, nydirty=nydirty,
                             epsilon=epsilon, pixsize_x=pixsize,
                             pixsize_y=pixsize)

    flags = np.zeros((nrow, nchan), dtype=np.bool)
    weights = np.random.rand(nrow, nchan)
    idx = ng.getIndices(baselines, gconf, flags)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = baselines.ms2vis(ms, idx)

    # ----------------------------------
    # Test complex grid case (ms2grid_c)
    # ----------------------------------
    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
    grid = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=None)
    grid2 = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=cplx_grid)

    # Almost same result
    assert_array_almost_equal(grid, grid2)
    assert grid.dtype == cplx_grid.dtype == grid2.dtype
    assert id(grid2) == id(cplx_grid)  # Same grid object

    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
    grid = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=None, wgt=weights)
    grid2 = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=cplx_grid,
                         wgt=weights)

    # Almost same result
    assert_array_almost_equal(grid, grid2)
    assert grid.dtype == cplx_grid.dtype == grid2.dtype
    assert id(grid2) == id(cplx_grid)  # Same grid object

    # ----------------------------------
    # Test complex grid case (vis2grid_c)
    # ----------------------------------
    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
    grid = ng.vis2grid_c(baselines, gconf, idx, vis, grid_in=None)
    grid2 = ng.vis2grid_c(baselines, gconf, idx, vis, grid_in=cplx_grid)

    # Almost same result
    assert_array_almost_equal(grid, grid2)
    assert grid.dtype == cplx_grid.dtype == grid2.dtype
    assert id(grid2) == id(cplx_grid)  # Same grid object


def test_pickling():
    np.random.seed(42)
    try:
        import cPickle as pickle
    except ImportError:
        import pickle

    import nifty_gridder as ng

    # Have to use cPickle and pickle protocol 2 for this to work
    # unfortunately
    pickle_protocol = 2

    gc = ng.GridderConfig(1024, 1024, 2e-13, 2.0, 2.0)
    pickled = pickle.dumps(gc, pickle_protocol)
    gc2 = pickle.loads(pickled)

    assert_(gc is not gc2)
    assert_(gc.Nxdirty() == gc2.Nxdirty())
    assert_(gc.Nydirty() == gc2.Nydirty())
    assert_(gc.Epsilon() == gc2.Epsilon())
    assert_(gc.Pixsize_x() == gc2.Pixsize_x())
    assert_(gc.Pixsize_y() == gc2.Pixsize_y())


@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 10, 10000))
@pmp("nchan", (1, 10, 100))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
@pmp("weight", (True, False))
def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    grid = ng.vis2grid_c(bl, conf, idx, vis)
    wgt = np.random.rand(*vis.shape) if weight else None
    y0 = ng.apply_holo(bl, conf, idx, grid, wgt)
    y1 = ng.vis2grid_c(bl, conf, idx, ng.grid2vis_c(bl, conf, idx, grid, wgt),
                       wgt=wgt)
    assert_allclose(y0, y1)


@pmp("nxdirty", (300,))
@pmp("nydirty", (250,))
@pmp("nrow", (10, 10000))
@pmp("nchan", (1, 10))
@pmp("epsilon", (1e-2, 1e-7))
@pmp("du", (-2, 0, 1, 4))
@pmp("dv", (-1, 0, 1))
@pmp("weight", (True, False))
def test_correlations(nxdirty, nydirty, nrow, nchan, epsilon, du, dv, weight):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    if conf.Supp() <= abs(du) or conf.Supp() <= dv:
        return
    wgt = np.random.rand(*idx.shape) if weight else None
    y0 = ng.get_correlations(bl, conf, idx, du, dv, wgt)
    pts = ((nxdirty, nydirty),
           (nxdirty + 1, nydirty),
           (nxdirty, nydirty//2),
           (0, 0),
           (2*nxdirty-1, 2*nydirty-1))
    for pp in pts:
        grid = np.zeros((2*nxdirty, 2*nydirty))
        grid[pp] = 1
        y1 = ng.apply_holo(bl, conf, idx, grid, wgt)
        ind = (pp[0]+du) % (2*nxdirty), (pp[1]+dv) % (2*nydirty)
        assert_allclose(y0[pp], y1[ind].real)
        assert_allclose(np.zeros_like(y1), y1.imag)


@pmp("nrow", (1, 1000))
@pmp("nchan", (1, 10))
@pmp("epsilon", (1e-2, 1e-7, 2e-12))
@pmp("weight", (True, False))
def test_no_correlation(nrow, nchan, epsilon, weight):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(128, 128, epsilon, nchan, nrow)
    wgt = np.random.rand(*idx.shape) if weight else None
    w = conf.Supp()
    for uu in range(-2*w, 2*w):
        for vv in range(-2*w, 2*w):
            if abs(uu) >= w or abs(vv) >= w:
                with pytest.raises(RuntimeError):
                    ng.get_correlations(bl, conf, idx, uu, vv, wgt)
            else:
                ng.get_correlations(bl, conf, idx, uu, vv, wgt)


@pmp('epsilon', [1e-2, 1e-4, 1e-7, 1e-10, 1e-11, 1e-12, 2e-13])
@pmp('nxdirty', [12, 128])
@pmp('nydirty', [4, 12, 128])
@pmp("nrow", (10, 100))
@pmp("nchan", (1, 10))
def test_against_dft(nxdirty, nydirty, epsilon, nchan, nrow):
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    uvw = bl.effectiveuvw(idx)
    res0 = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, idx, vis))
    res1 = _dft(uvw, vis, conf, False)
    assert_(_l2error(res0, res1) < epsilon)


@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100, 1000))
@pmp("nchan", (1, 10))
@pmp("fov", (1,))
def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
    epsilon = 1e-7
    np.random.seed(40)
    pixsize = fov*np.pi/180/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
                            pixsize_y=0.5*pixsize)
    speedoflight, f0 = 3e8, 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
    bl = ng.Baselines(coord=uvw, freq=freq)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
    idx = ng.getIndices(bl, conf, flags)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    uvw = bl.effectiveuvw(idx)
    res0 = np.zeros((nxdirty, nydirty), dtype=np.complex128)
    mi, ma = np.min(uvw[:, 2]), np.max(uvw[:, 2])
    nplanes = 10000
    ws = mi + np.arange(nplanes)*(ma-mi)/(nplanes-1)
    for ii in range(len(ws)-1):
        wmax = ws[ii+1] if ii != nplanes-2 else ws[ii+1] + abs(ws[ii+1])
        wmax = ws[ii+1] if ii != nplanes-2 else np.inf
        iind = ng.getIndices(bl, conf, flags, wmin=ws[ii], wmax=wmax)
        if len(iind) == 0:
            continue
        dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
        res0 += conf.apply_wscreen(dd, 0.5*(ws[ii+1]+ws[ii]), adjoint=False)
    res1 = _dft(uvw, vis, conf, True)
    assert_(_l2error(res0, res1) < 1e-4)  # Very inaccurate


@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100, 1000))
@pmp("nchan", (1,))
@pmp("fov", (50,))
def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
    epsilon = 1e-10
    np.random.seed(40)
    pixsize = fov*np.pi/180/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
                            pixsize_y=0.2*pixsize)
    speedoflight, f0 = 3e8, 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
    uvw[:, 2] = uvw[0, 2]  # use exactly one random w
    bl = ng.Baselines(coord=uvw, freq=freq)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
    idx = ng.getIndices(bl, conf, flags)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    uvw = bl.effectiveuvw(idx)
    w = np.min(uvw[:, 2])
    iind = ng.getIndices(bl, conf, flags)
    dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
    res0 = conf.apply_wscreen(dd, w, adjoint=False)
    res1 = _dft(uvw, vis, conf, True)
    assert_(_l2error(res0, res1) < epsilon)


@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100, 1000))
@pmp("nchan", (1, 10))
@pmp("fov", (1, 10))
@pmp("epsilon", (1e-2, 1e-7, 2e-12))
def test_wgridder_against_wdft(nxdirty, nydirty, nchan, nrow, fov, epsilon):
    np.random.seed(40)
    pixsize = fov*np.pi/180/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
                            pixsize_y=0.5*pixsize)
    speedoflight, f0 = 3e8, 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
    bl = ng.Baselines(coord=uvw, freq=freq)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
    idx = ng.getIndices(bl, conf, flags)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    uvw = bl.effectiveuvw(idx)
    res0 = ng.vis2dirty_wstack(bl, conf, idx, vis)
    res1 = _dft(uvw, vis, conf, True)
    assert_(_l2error(res0, res1) < epsilon)