test.py 12.7 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

pmp = pytest.mark.parametrize


Philipp Arras's avatar
Philipp Arras committed
9
10
11
12
13
def _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow):
    pixsize = np.pi/180/60/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
Philipp Arras's avatar
Philipp Arras committed
14
                            pixsize_x=0.368*pixsize,
Philipp Arras's avatar
Philipp Arras committed
15
16
17
18
19
20
21
22
23
24
                            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


25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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
Philipp Arras's avatar
Philipp Arras committed
43
44


Philipp Arras's avatar
Philipp Arras committed
45
46
47
48
def _l2error(a, b):
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))


Martin Reinecke's avatar
Martin Reinecke committed
49
50
51
52
53
54
@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):
55
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
56
    bl, gconf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
57
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
Philipp Arras's avatar
Philipp Arras committed
58
    vis = bl.ms2vis(ms, idx)
Martin Reinecke's avatar
Martin Reinecke committed
59
60
    grid = ng.vis2grid_c(bl, gconf, idx, vis)
    dirty = gconf.grid2dirty_c(grid)
Martin Reinecke's avatar
Martin Reinecke committed
61
    dirty2 = np.random.rand(*dirty.shape)-0.5
62
63
    ms2 = bl.vis2ms(ng.grid2vis_c(bl, gconf, idx, gconf.dirty2grid_c(dirty2)),
                    idx)
Martin Reinecke's avatar
Martin Reinecke committed
64
    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty.real, dirty2.real))
65
66


67
68
69
70
71
72
73
74
@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)
75
    vis = np.random.rand(*idx.shape)-0.5 + 1j*(np.random.rand(*idx.shape)-0.5)
76
77
78
    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)
79
    assert_allclose(np.vdot(vis, vis2), np.vdot(dirty2, dirty),
Philipp Arras's avatar
Fixup    
Philipp Arras committed
80
                    rtol=epsilon)
81
82


83
84
85
86
87
88
@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):
89
    np.random.seed(42)
90
91
92
93
94
95
96
97
98
99
100
101
102
103
    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)
Simon Perkins's avatar
Simon Perkins committed
104
    weights = np.random.rand(nrow, nchan)
105
106
107
    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)
Martin Reinecke's avatar
Martin Reinecke committed
108

109
110
111
112
113
114
    # ----------------------------------
    # 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)
Simon Perkins's avatar
Simon Perkins committed
115

116
117
118
119
120
    # 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

Simon Perkins's avatar
Simon Perkins committed
121
    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
122
123
124
    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)
Simon Perkins's avatar
Simon Perkins committed
125
126
127
128
129
130

    # 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

131
132
133
134
135
136
137
138
139
140
141
    # ----------------------------------
    # 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
142

Martin Reinecke's avatar
Martin Reinecke committed
143
144

def test_pickling():
145
    np.random.seed(42)
Martin Reinecke's avatar
Martin Reinecke committed
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
    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())
167
168


169
170
171
172
173
@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))
Philipp Arras's avatar
Philipp Arras committed
174
175
@pmp("weight", (True, False))
def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
176
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
177
178
179
180
    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)
Philipp Arras's avatar
Philipp Arras committed
181
    wgt = np.random.rand(*vis.shape) if weight else None
Philipp Arras's avatar
Philipp Arras committed
182
183
184
    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)
185
186
187
188
189
190
191
192
193
194
    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))
Philipp Arras's avatar
Philipp Arras committed
195
196
@pmp("weight", (True, False))
def test_correlations(nxdirty, nydirty, nrow, nchan, epsilon, du, dv, weight):
197
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
198
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
199
    if conf.Supp() <= abs(du) or conf.Supp() <= dv:
200
        return
Philipp Arras's avatar
Philipp Arras committed
201
202
    wgt = np.random.rand(*idx.shape) if weight else None
    y0 = ng.get_correlations(bl, conf, idx, du, dv, wgt)
Philipp Arras's avatar
Philipp Arras committed
203
204
205
206
207
    pts = ((nxdirty, nydirty),
           (nxdirty + 1, nydirty),
           (nxdirty, nydirty//2),
           (0, 0),
           (2*nxdirty-1, 2*nydirty-1))
208
209
    for pp in pts:
        grid = np.zeros((2*nxdirty, 2*nydirty))
Philipp Arras's avatar
Philipp Arras committed
210
        grid[pp] = 1
Philipp Arras's avatar
Philipp Arras committed
211
        y1 = ng.apply_holo(bl, conf, idx, grid, wgt)
Philipp Arras's avatar
Philipp Arras committed
212
213
214
        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)
Philipp Arras's avatar
Philipp Arras committed
215
216


Philipp Arras's avatar
Philipp Arras committed
217
218
219
220
221
222
223
224
@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
Philipp Arras's avatar
Fixups    
Philipp Arras committed
225
    w = conf.Supp()
Philipp Arras's avatar
Philipp Arras committed
226
227
228
229
230
231
232
233
234
    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)


Philipp Arras's avatar
Philipp Arras committed
235
236
237
238
239
240
241
242
243
244
@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)
Philipp Arras's avatar
Fixups    
Philipp Arras committed
245
246
    res0 = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, idx, vis))
    res1 = _dft(uvw, vis, conf, False)
Philipp Arras's avatar
Philipp Arras committed
247
    assert_(_l2error(res0, res1) < epsilon)
248
249


Philipp Arras's avatar
Philipp Arras committed
250
@pmp('nxdirty', [16, 64])
251
@pmp('nydirty', [64])
Philipp Arras's avatar
Philipp Arras committed
252
253
254
@pmp("nrow", (10, 100, 1000))
@pmp("nchan", (1, 10))
@pmp("fov", (1,))
255
def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
Philipp Arras's avatar
Philipp Arras committed
256
    epsilon = 1e-7
257
    np.random.seed(40)
Philipp Arras's avatar
Philipp Arras committed
258
    pixsize = fov*np.pi/180/nxdirty
259
260
261
262
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
Martin Reinecke's avatar
Martin Reinecke committed
263
                            pixsize_y=0.5*pixsize)
264
265
266
267
268
269
270
271
272
    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)
273
    res0 = np.zeros((nxdirty, nydirty), dtype=np.complex128)
274
    mi, ma = np.min(uvw[:, 2]), np.max(uvw[:, 2])
Philipp Arras's avatar
Philipp Arras committed
275
276
277
    nplanes = 10000
    ws = mi + np.arange(nplanes)*(ma-mi)/(nplanes-1)
    for ii in range(len(ws)-1):
Philipp Arras's avatar
Philipp Arras committed
278
279
280
281
        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:
282
            continue
Philipp Arras's avatar
Philipp Arras committed
283
        dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
284
285
286
        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
Philipp Arras's avatar
Philipp Arras committed
287

Martin Reinecke's avatar
Martin Reinecke committed
288
289
290
291
292
293

@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100, 1000))
@pmp("nchan", (1,))
@pmp("fov", (50,))
294
def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
Martin Reinecke's avatar
Martin Reinecke committed
295
296
297
298
299
300
301
302
303
304
305
    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)
306
    uvw[:, 2] = uvw[0, 2]  # use exactly one random w
Martin Reinecke's avatar
Martin Reinecke committed
307
308
309
310
311
312
    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)
313
    w = np.min(uvw[:, 2])
Martin Reinecke's avatar
Martin Reinecke committed
314
315
    iind = ng.getIndices(bl, conf, flags)
    dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
316
317
318
    res0 = conf.apply_wscreen(dd, w, adjoint=False)
    res1 = _dft(uvw, vis, conf, True)
    assert_(_l2error(res0, res1) < epsilon)
Martin Reinecke's avatar
Martin Reinecke committed
319

320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344

@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)
345
    res1 = _dft(uvw, vis, conf, True)
346
    assert_(_l2error(res0, res1) < epsilon)