test.py 14.2 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
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
39
        dirty += (vv*np.exp(2j*np.pi*phase)).real
40
41
42
    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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
    assert_allclose(np.vdot(vis, vis2).real, np.vdot(dirty2, dirty),
                    rtol=epsilon)


@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 10, 10000))
@pmp("nchan", (1, 10, 100))
@pmp("epsilon", (1e-2, 1e-3, 1e-5, 1e-6, 1e-7, 2e-13))
@pmp("singleprec", (True, False))
@pmp("wstacking", (True, False))
def test_adjointness_wgridding_highlevel(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking):
    np.random.seed(42)
    pixsizex = np.pi/180/60/nxdirty*0.2398
    pixsizey = np.pi/180/60/nxdirty
    speedoflight, f0 = 3e8, 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsizey*f0/speedoflight)
    vis = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    dirty = np.random.rand(nxdirty, nydirty)-0.5
    if wstacking:
        f1 = ng.full_gridding
        f2 = ng.full_degridding
    else:
        f1 = ng.gridding
        f2 = ng.degridding
    if singleprec:
        vis = vis.astype("c8")
        uvw = uvw.astype("f4")
        dirty = dirty.astype("f4")
        freq = freq.astype("f4")
        if wstacking:
            f1 = ng.full_gridding_f
            f2 = ng.full_degridding_f
        else:
            f1 = ng.gridding_f
            f2 = ng.degridding_f
    dirty2 = f1(uvw, freq, vis, None, nxdirty, nydirty, pixsizex, pixsizey, epsilon, 1, 0)
    vis2 = f2(uvw, freq, dirty, None, pixsizex, pixsizey, epsilon, 1, 0)
118
    assert_allclose(np.vdot(vis, vis2).real, np.vdot(dirty2, dirty),
Philipp Arras's avatar
Fixup    
Philipp Arras committed
119
                    rtol=epsilon)
120
121


122
123
124
125
126
127
@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):
128
    np.random.seed(42)
129
130
131
132
133
134
135
136
137
138
139
140
141
142
    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
143
    weights = np.random.rand(nrow, nchan)
144
145
146
    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
147

148
149
150
151
152
153
    # ----------------------------------
    # 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
154

155
156
157
158
159
    # 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
160
    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
161
162
163
    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
164
165
166
167
168
169

    # 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

170
171
172
173
174
175
176
177
178
179
180
    # ----------------------------------
    # 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
181

Martin Reinecke's avatar
Martin Reinecke committed
182
183

def test_pickling():
184
    np.random.seed(42)
Martin Reinecke's avatar
Martin Reinecke committed
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
    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())
206
207


208
209
210
211
212
@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
213
214
@pmp("weight", (True, False))
def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
215
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
216
217
218
219
    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
220
    wgt = np.random.rand(*vis.shape) if weight else None
Philipp Arras's avatar
Philipp Arras committed
221
222
223
    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)
224
225
226
227
228
229
230
231
232
233
    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
234
235
@pmp("weight", (True, False))
def test_correlations(nxdirty, nydirty, nrow, nchan, epsilon, du, dv, weight):
236
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
237
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
238
    if conf.Supp() <= abs(du) or conf.Supp() <= dv:
239
        return
Philipp Arras's avatar
Philipp Arras committed
240
241
    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
242
243
244
245
246
    pts = ((nxdirty, nydirty),
           (nxdirty + 1, nydirty),
           (nxdirty, nydirty//2),
           (0, 0),
           (2*nxdirty-1, 2*nydirty-1))
247
248
    for pp in pts:
        grid = np.zeros((2*nxdirty, 2*nydirty))
Philipp Arras's avatar
Philipp Arras committed
249
        grid[pp] = 1
Philipp Arras's avatar
Philipp Arras committed
250
        y1 = ng.apply_holo(bl, conf, idx, grid, wgt)
Philipp Arras's avatar
Philipp Arras committed
251
252
253
        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
254
255


Philipp Arras's avatar
Philipp Arras committed
256
257
258
259
260
261
262
263
@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
264
    w = conf.Supp()
Philipp Arras's avatar
Philipp Arras committed
265
266
267
268
269
270
271
272
273
    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
274
@pmp('epsilon', [1e-2, 1e-4, 1e-7, 1e-10, 1e-11, 1e-12, 2e-13])
Martin Reinecke's avatar
Martin Reinecke committed
275
276
@pmp('nxdirty', [2, 12, 128])
@pmp('nydirty', [2, 4, 12, 128])
Philipp Arras's avatar
Philipp Arras committed
277
278
279
@pmp("nrow", (10, 100))
@pmp("nchan", (1, 10))
def test_against_dft(nxdirty, nydirty, epsilon, nchan, nrow):
Martin Reinecke's avatar
Martin Reinecke committed
280
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
281
282
283
284
    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)
285
    res0 = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, idx, vis)).real
Philipp Arras's avatar
Fixups    
Philipp Arras committed
286
    res1 = _dft(uvw, vis, conf, False)
Philipp Arras's avatar
Philipp Arras committed
287
    assert_(_l2error(res0, res1) < epsilon)
288
289


Philipp Arras's avatar
Philipp Arras committed
290
@pmp('nxdirty', [16, 64])
291
@pmp('nydirty', [64])
Philipp Arras's avatar
Philipp Arras committed
292
293
294
@pmp("nrow", (10, 100, 1000))
@pmp("nchan", (1, 10))
@pmp("fov", (1,))
295
def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
Philipp Arras's avatar
Philipp Arras committed
296
    epsilon = 1e-7
297
    np.random.seed(40)
Philipp Arras's avatar
Philipp Arras committed
298
    pixsize = fov*np.pi/180/nxdirty
299
300
301
302
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
Martin Reinecke's avatar
Martin Reinecke committed
303
                            pixsize_y=0.5*pixsize)
304
305
306
307
308
309
310
311
312
    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)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
313
314
    res0 = np.zeros((nxdirty, nydirty), dtype=np.float64)
    mi, ma = np.min(np.abs(uvw[:, 2])), np.max(np.abs(uvw[:, 2]))
Philipp Arras's avatar
Philipp Arras committed
315
316
317
    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
318
319
320
321
        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:
322
            continue
Philipp Arras's avatar
Philipp Arras committed
323
        dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
324
        res0 += conf.apply_wscreen(dd, 0.5*(ws[ii+1]+ws[ii]), adjoint=False).real
325
326
    res1 = _dft(uvw, vis, conf, True)
    assert_(_l2error(res0, res1) < 1e-4)  # Very inaccurate
Philipp Arras's avatar
Philipp Arras committed
327

Martin Reinecke's avatar
Martin Reinecke committed
328
329
330
331
332
333

@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100, 1000))
@pmp("nchan", (1,))
@pmp("fov", (50,))
334
def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
Martin Reinecke's avatar
Martin Reinecke committed
335
336
337
338
339
340
341
342
343
344
345
    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)
346
    uvw[:, 2] = uvw[0, 2]  # use exactly one random w
Martin Reinecke's avatar
Martin Reinecke committed
347
348
349
350
351
352
    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)
353
    w = np.min(uvw[:, 2])
Martin Reinecke's avatar
Martin Reinecke committed
354
355
    iind = ng.getIndices(bl, conf, flags)
    dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
356
    res0 = conf.apply_wscreen(dd, w, adjoint=False).real
357
358
    res1 = _dft(uvw, vis, conf, True)
    assert_(_l2error(res0, res1) < epsilon)
Martin Reinecke's avatar
Martin Reinecke committed
359

360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384

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