There is a maintenance of MPCDF Gitlab on Thursday, April 22st 2020, 9:00 am CEST - Expect some service interruptions during this time

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)