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


@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):
    f0 = 1e9  # rough observation frequency
    pixsize = np.pi/180/60/nxdirty  # assume 1 arcmin FOV
    speedoflight = 3e8

    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5) / (pixsize*f0/speedoflight)
    uvw[:, 2] = 0.
    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)
    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)
    grid = ng.vis2grid(baselines, gconf, idx, vis)
    dirty = gconf.grid2dirty(grid)
    dirty2 = np.random.rand(*dirty.shape)-0.5

    ms2 = baselines.vis2ms(ng.grid2vis(baselines, gconf, idx,
                           gconf.dirty2grid(dirty2)), idx)

    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty, dirty2))


@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):
    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)
    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)
    user_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=user_grid)

    # Almost same visibilities
    assert_array_almost_equal(grid, grid2)

    # Same grid object
    assert id(grid2) == id(user_grid)


def test_pickling():
    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())