test.py 3.3 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
import nifty_gridder as ng
import numpy as np
import pytest
4
from numpy.testing import assert_, assert_allclose, assert_array_almost_equal
Martin Reinecke's avatar
Martin Reinecke committed
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

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)
20
21
    uvw = (np.random.rand(nrow, 3)-0.5) / (pixsize*f0/speedoflight)
    uvw[:, 2] = 0.
Martin Reinecke's avatar
Martin Reinecke committed
22
23
    baselines = ng.Baselines(coord=uvw, freq=freq)
    gconf = ng.GridderConfig(nxdirty=nxdirty, nydirty=nydirty,
24
25
26
                             epsilon=epsilon, pixsize_x=pixsize,
                             pixsize_y=pixsize)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
Martin Reinecke's avatar
Martin Reinecke committed
27
    idx = ng.getIndices(baselines, gconf, flags)
28
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
Martin Reinecke's avatar
Martin Reinecke committed
29
30
31
32
33
34
    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,
35
                           gconf.dirty2grid(dirty2)), idx)
Martin Reinecke's avatar
Martin Reinecke committed
36

37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    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)
65
66
    grid = ng.vis2grid_c(baselines, gconf, idx, vis, user_grid=None)
    grid2 = ng.vis2grid_c(baselines, gconf, idx, vis, user_grid=user_grid)
Martin Reinecke's avatar
Martin Reinecke committed
67

Simon Perkins's avatar
Simon Perkins committed
68
    # Almost same visibilities
69
    assert_array_almost_equal(grid, grid2)
Simon Perkins's avatar
Simon Perkins committed
70
71

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

Martin Reinecke's avatar
Martin Reinecke committed
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

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())
97
98