test.py 16.7 KB
Newer Older
Philipp Arras's avatar
Philipp Arras committed
1
2
from itertools import product

Martin Reinecke's avatar
Martin Reinecke committed
3
4
5
import nifty_gridder as ng
import numpy as np
import pytest
6
from numpy.testing import assert_, assert_allclose, assert_array_almost_equal
Martin Reinecke's avatar
Martin Reinecke committed
7
8
9

pmp = pytest.mark.parametrize

Martin Reinecke's avatar
Martin Reinecke committed
10

Philipp Arras's avatar
Philipp Arras committed
11
12
13
14
15
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
16
                            pixsize_x=0.368*pixsize,
Philipp Arras's avatar
Philipp Arras committed
17
                            pixsize_y=pixsize)
Martin Reinecke's avatar
Martin Reinecke committed
18
    speedoflight, f0 = 299792458., 1e9
Philipp Arras's avatar
Philipp Arras committed
19
20
21
22
23
24
25
26
    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


Philipp Arras's avatar
Philipp Arras committed
27
def _l2error(a, b):
Martin Reinecke's avatar
Martin Reinecke committed
28
29
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.maximum(np.sum(np.abs(a)**2),
                                                     np.sum(np.abs(b)**2)))
30
31


Martin Reinecke's avatar
Martin Reinecke committed
32
33
def explicit_gridder(uvw, freq, ms, wgt, nxdirty, nydirty, xpixsize, ypixsize,
                     apply_w):
34
35
    speedoflight = 299792458.
    x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]],
Martin Reinecke's avatar
Martin Reinecke committed
36
                       indexing='ij')
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
65
66
67
68
69
70
71
72
    x *= xpixsize
    y *= ypixsize
    res = np.zeros((nxdirty, nydirty))
    eps = x**2+y**2
    if apply_w:
        nm1 = -eps/(np.sqrt(1.-eps)+1.)
        n = nm1+1
    else:
        nm1 = 0.
        n = 1.
    for row in range(ms.shape[0]):
        for chan in range(ms.shape[1]):
            phase = (freq[chan]/speedoflight *
                     (x*uvw[row, 0] + y*uvw[row, 1] - uvw[row, 2]*nm1))
            if wgt is None:
                res += (ms[row, chan]*np.exp(2j*np.pi*phase)).real
            else:
                res += (ms[row, chan]*wgt[row, chan]*np.exp(2j*np.pi*phase)).real
    return res/n


@pmp("nxdirty", (30, 128))
@pmp("nydirty", (128, 250))
@pmp("nrow", (2, 27))
@pmp("nchan", (1, 5))
@pmp("epsilon", (1e-3, 1e-5, 1e-7, 5e-13))
@pmp("singleprec", (True, False))
@pmp("wstacking", (True, False))
@pmp("use_wgt", (True, False))
@pmp("nthreads", (1, 2))
def test_adjointness_ms2dirty(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking, use_wgt, nthreads):
    if singleprec and epsilon < 5e-6:
        return
    np.random.seed(42)
    pixsizex = np.pi/180/60/nxdirty*0.2398
    pixsizey = np.pi/180/60/nxdirty
Martin Reinecke's avatar
Martin Reinecke committed
73
    speedoflight, f0 = 299792458., 1e9
74
75
76
77
78
79
80
81
82
83
84
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsizey*f0/speedoflight)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    wgt = np.random.rand(nrow, nchan) if use_wgt else None
    dirty = np.random.rand(nxdirty, nydirty)-0.5
    if singleprec:
        ms = ms.astype("c8")
        dirty = dirty.astype("f4")
        if wgt is not None:
            wgt = wgt.astype("f4")
    dirty2 = ng.ms2dirty(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex,
Martin Reinecke's avatar
Martin Reinecke committed
85
                         pixsizey, epsilon, wstacking, nthreads, 0).astype("f8")
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
    ms2 = ng.dirty2ms(uvw, freq, dirty, wgt, pixsizex, pixsizey, epsilon,
                      wstacking, nthreads, 0).astype("c16")
    tol = 5e-5 if singleprec else 5e-13
    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty2, dirty), rtol=tol)


@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (2, 27))
@pmp("nchan", (1, 5))
@pmp("epsilon", (1e-2, 1e-3, 5e-5, 1e-7, 5e-13))
@pmp("singleprec", (True, False))
@pmp("wstacking", (True, False))
@pmp("use_wgt", (False, True))
@pmp("nthreads", (1, 2))
Martin Reinecke's avatar
Martin Reinecke committed
101
@pmp("fov", (1., 20.))
102
103
104
105
106
107
def test_ms2dirty_against_wdft(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking, use_wgt, fov, nthreads):
    if singleprec and epsilon < 5e-5:
        return
    np.random.seed(40)
    pixsizex = fov*np.pi/180/nxdirty
    pixsizey = fov*np.pi/180/nydirty*1.1
Martin Reinecke's avatar
Martin Reinecke committed
108
    speedoflight, f0 = 299792458., 1e9
109
110
111
112
113
114
115
116
117
118
119
120
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsizex*f0/speedoflight)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    wgt = np.random.rand(nrow, nchan) if use_wgt else None
    if singleprec:
        ms = ms.astype("c8")
        if wgt is not None:
            wgt = wgt.astype("f4")
    dirty = ng.ms2dirty(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex,
                        pixsizey, epsilon, wstacking, nthreads, 0).astype("f8")
    ref = explicit_gridder(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex, pixsizey, wstacking)
    assert_allclose(_l2error(dirty, ref), 0, atol=epsilon)
Philipp Arras's avatar
Philipp Arras committed
121
122


Martin Reinecke's avatar
Martin Reinecke committed
123
124
@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
125
126
@pmp("nrow", (1, 100))
@pmp("nchan", (1, 10))
Martin Reinecke's avatar
Martin Reinecke committed
127
128
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
129
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
130
    bl, gconf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
131
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
Philipp Arras's avatar
Philipp Arras committed
132
    vis = bl.ms2vis(ms, idx)
Martin Reinecke's avatar
Martin Reinecke committed
133
134
    grid = ng.vis2grid_c(bl, gconf, idx, vis)
    dirty = gconf.grid2dirty_c(grid)
Martin Reinecke's avatar
Martin Reinecke committed
135
136
    dirty2 = (np.random.rand(*dirty.shape)-0.5 +
              1j*(np.random.rand(*dirty.shape)-0.5))
137
138
    ms2 = bl.vis2ms(ng.grid2vis_c(bl, gconf, idx, gconf.dirty2grid_c(dirty2)),
                    idx)
Martin Reinecke's avatar
Martin Reinecke committed
139
    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty, dirty2).real)
140
141


142
143
@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
144
145
@pmp("nrow", (1, 127))
@pmp("nchan", (1, 17))
146
147
148
149
@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)
150
    vis = np.random.rand(*idx.shape)-0.5 + 1j*(np.random.rand(*idx.shape)-0.5)
151
    dirty = np.random.rand(conf.Nxdirty(), conf.Nydirty())-0.5
Martin Reinecke's avatar
Martin Reinecke committed
152
153
    dirty2 = ng.vis2dirty(bl, conf, idx, vis, do_wstacking=True)
    vis2 = ng.dirty2vis(bl, conf, idx, dirty, do_wstacking=True)
154
    assert_allclose(np.vdot(vis, vis2).real, np.vdot(dirty2, dirty),
Martin Reinecke's avatar
Martin Reinecke committed
155
                    rtol=2e-13)
156
157


158
159
@pmp("nxdirty", (128,))
@pmp("nydirty", (128,))
160
161
@pmp("nrow", (1000,))
@pmp("nchan", (10,))
162
163
@pmp("epsilon", (2e-13,))
def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
164
    np.random.seed(42)
165
166
167
168
169
    gconf = ng.GridderConfig(1024, 1024, 2e-13, 2.0, 2.0)
    freq = np.linspace(.856e9, 2*.856e9, nchan)
    f0 = freq[freq.shape[0]//2]

    pixsize = np.pi/180/60/nxdirty  # assume 1 arcmin FOV
Martin Reinecke's avatar
Martin Reinecke committed
170
    speedoflight = 299792458.
171
172
173
174
175
176
177
    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
178
    weights = np.random.rand(nrow, nchan)
179
180
181
    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
182

183
184
185
186
187
188
    # ----------------------------------
    # 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
189

190
191
192
193
194
    # 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
195
    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
196
197
198
    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
199
200
201
202
203
204

    # 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

205
206
207
208
209
210
211
212
213
214
215
    # ----------------------------------
    # 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
216

Martin Reinecke's avatar
Martin Reinecke committed
217
218

def test_pickling():
219
    np.random.seed(42)
Martin Reinecke's avatar
Martin Reinecke committed
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
    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())
241
242


243
244
@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
245
246
@pmp("nrow", (1, 39))
@pmp("nchan", (1, 70))
247
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
Philipp Arras's avatar
Philipp Arras committed
248
249
@pmp("weight", (True, False))
def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
250
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
251
252
253
254
    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
255
    wgt = np.random.rand(*vis.shape) if weight else None
Philipp Arras's avatar
Philipp Arras committed
256
257
258
    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)
259
260
261
262
263
    assert_allclose(y0, y1)


@pmp("nxdirty", (300,))
@pmp("nydirty", (250,))
264
@pmp("nrow", (10, 100))
265
266
267
268
@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
269
270
@pmp("weight", (True, False))
def test_correlations(nxdirty, nydirty, nrow, nchan, epsilon, du, dv, weight):
271
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
272
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
273
    if conf.Supp() <= abs(du) or conf.Supp() <= dv:
274
        return
Philipp Arras's avatar
Philipp Arras committed
275
276
    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
277
278
279
280
281
    pts = ((nxdirty, nydirty),
           (nxdirty + 1, nydirty),
           (nxdirty, nydirty//2),
           (0, 0),
           (2*nxdirty-1, 2*nydirty-1))
282
    for pp in pts:
Martin Reinecke's avatar
Martin Reinecke committed
283
        grid = np.zeros((2*nxdirty, 2*nydirty))+0j
Philipp Arras's avatar
Philipp Arras committed
284
        grid[pp] = 1
Philipp Arras's avatar
Philipp Arras committed
285
        y1 = ng.apply_holo(bl, conf, idx, grid, wgt)
Philipp Arras's avatar
Philipp Arras committed
286
287
288
        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
289
290


291
@pmp("nrow", (1, 100))
Philipp Arras's avatar
Philipp Arras committed
292
293
294
295
296
297
298
@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
299
    w = conf.Supp()
Philipp Arras's avatar
Philipp Arras committed
300
301
302
303
304
305
306
307
308
    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
309
@pmp('nxdirty', [16, 64])
310
@pmp('nydirty', [64])
311
@pmp("nrow", (10, 100))
Philipp Arras's avatar
Philipp Arras committed
312
313
@pmp("nchan", (1, 10))
@pmp("fov", (1,))
314
def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
Philipp Arras's avatar
Philipp Arras committed
315
    epsilon = 1e-7
316
    np.random.seed(40)
Philipp Arras's avatar
Philipp Arras committed
317
    pixsize = fov*np.pi/180/nxdirty
318
319
320
321
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
Martin Reinecke's avatar
Martin Reinecke committed
322
                            pixsize_y=0.5*pixsize)
Martin Reinecke's avatar
Martin Reinecke committed
323
    speedoflight, f0 = 299792458., 1e9
324
325
326
327
328
329
330
    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)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
331
    res0 = np.zeros((nxdirty, nydirty), dtype=np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
332
333
    mi = np.min(np.abs(uvw[:, 2]))*freq[0]/speedoflight
    ma = np.max(np.abs(uvw[:, 2]))*freq[-1]/speedoflight
Philipp Arras's avatar
Philipp Arras committed
334
335
336
    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
337
338
339
340
        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:
341
            continue
Philipp Arras's avatar
Philipp Arras committed
342
        dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
343
        res0 += conf.apply_wscreen(dd, 0.5*(ws[ii+1]+ws[ii]), adjoint=True).real
Martin Reinecke's avatar
Martin Reinecke committed
344
345
    res1 = explicit_gridder(uvw, freq, ms, None, nxdirty, nydirty, pixsize,
                            0.5*pixsize, True)
346
    assert_(_l2error(res0, res1) < 1e-4)  # Very inaccurate
Philipp Arras's avatar
Philipp Arras committed
347

Martin Reinecke's avatar
Martin Reinecke committed
348
349
350

@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
351
@pmp("nrow", (10, 100))
Martin Reinecke's avatar
Martin Reinecke committed
352
353
@pmp("nchan", (1,))
@pmp("fov", (50,))
354
def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
Martin Reinecke's avatar
Martin Reinecke committed
355
356
357
358
359
360
361
362
    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)
Martin Reinecke's avatar
Martin Reinecke committed
363
    speedoflight, f0 = 299792458., 1e9
Martin Reinecke's avatar
Martin Reinecke committed
364
365
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
366
    uvw[:, 2] = uvw[0, 2]  # use exactly one random w
Martin Reinecke's avatar
Martin Reinecke committed
367
368
369
370
371
    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)
Martin Reinecke's avatar
Martin Reinecke committed
372
    w = np.min(uvw[:, 2])*freq[0]/speedoflight
Martin Reinecke's avatar
Martin Reinecke committed
373
374
    iind = ng.getIndices(bl, conf, flags)
    dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
375
    res0 = conf.apply_wscreen(dd, w, adjoint=True).real
Martin Reinecke's avatar
Martin Reinecke committed
376
377
    res1 = explicit_gridder(uvw, freq, ms, None, nxdirty, nydirty, pixsize,
                            0.2*pixsize, True)
378
    assert_(_l2error(res0, res1) < epsilon)
Martin Reinecke's avatar
Martin Reinecke committed
379

380
381
382

@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
383
@pmp("nrow", (10, 100))
384
385
386
387
388
389
390
391
392
393
394
@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)
Martin Reinecke's avatar
Martin Reinecke committed
395
    speedoflight, f0 = 299792458., 1e9
396
397
398
399
400
401
402
    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)
Martin Reinecke's avatar
fix    
Martin Reinecke committed
403
    res0 = ng.vis2dirty(bl, conf, idx, vis, do_wstacking=True)
Martin Reinecke's avatar
Martin Reinecke committed
404
405
    res1 = explicit_gridder(uvw, freq, ms, None, nxdirty, nydirty, pixsize,
                            0.5*pixsize, True)
406
    assert_(_l2error(res0, res1) < epsilon)
407

Philipp Arras's avatar
Philipp Arras committed
408

Philipp Arras's avatar
Philipp Arras committed
409
410
@pmp('nxdirty', (32, 16))
@pmp('nydirty', (32, 64))
411
@pmp("nrow", (100, 127))
Philipp Arras's avatar
Philipp Arras committed
412
413
@pmp("nchan", (1, 5))
@pmp('epsilon', (1e-3, 1e-6))
414
def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
Philipp Arras's avatar
Philipp Arras committed
415
416
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
Philipp Arras's avatar
Philipp Arras committed
417
    nx, ny = 2*nxdirty, 2*nydirty
Philipp Arras's avatar
Philipp Arras committed
418
    grid = np.random.randn(nx, ny) + 1.0j*np.random.randn(nx, ny)
Martin Reinecke's avatar
merge    
Martin Reinecke committed
419
    W = conf.Supp()
Philipp Arras's avatar
Philipp Arras committed
420
    res0 = np.zeros_like(grid)
421
422
423
424

    # Precompute 2*W**2-2*W+1 images (naive: 4*W**2-4*W+1)
    d0, d1 = [], []
    for du, dv in product(range(-W + 1, W), range(1, W)):
Philipp Arras's avatar
Cleanup    
Philipp Arras committed
425
        d0.append(ng.get_correlations(bl, conf, idx, du, dv))
426
    for du in range(1, W):
Philipp Arras's avatar
Cleanup    
Philipp Arras committed
427
428
        d1.append(ng.get_correlations(bl, conf, idx, du, 0))
    d2 = ng.get_correlations(bl, conf, idx, 0, 0)
429
430
431
432
433
434
435
436
437
438

    # Apply
    for ii, (du, dv) in enumerate(product(range(-W + 1, W), range(1, W))):
        tmp = np.roll(grid, -du, axis=0)
        res0 += d0[ii]*np.roll(tmp, -dv, axis=1)
        res0 += np.roll(d0[ii]*tmp, dv, axis=1)
    for ii, du in enumerate(range(1, W)):
        res0 += d1[ii]*np.roll(grid, -du, axis=0)
        res0 += np.roll(d1[ii]*grid, du, axis=0)
    res0 += d2*grid
Philipp Arras's avatar
Philipp Arras committed
439
    res1 = ng.apply_holo(bl, conf, idx, grid)
Martin Reinecke's avatar
Martin Reinecke committed
440
    assert_allclose(res0, res1, rtol=1e-13, atol=1e-13)