test.py 13.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
10

pmp = pytest.mark.parametrize


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
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
                            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


def _wscreen(npix, dst, w):
    dc = (np.linspace(start=-npix/2, stop=npix/2 - 1, num=npix)*dst)**2
    ls = np.broadcast_to(dc, (dc.shape[0],)*2)
    theta = np.sqrt(ls + ls.T)
    n = np.cos(theta)
    wscreen = np.exp(2*np.pi*1j*w*(n - 1))/n
    return wscreen


Philipp Arras's avatar
Philipp Arras committed
36
37
38
39
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
40
41
42
43
44
45
@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):
46
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
47
    bl, gconf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
48
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
Philipp Arras's avatar
Philipp Arras committed
49
50
    vis = bl.ms2vis(ms, idx)
    grid = ng.vis2grid(bl, gconf, idx, vis)
Martin Reinecke's avatar
Martin Reinecke committed
51
52
    dirty = gconf.grid2dirty(grid)
    dirty2 = np.random.rand(*dirty.shape)-0.5
Philipp Arras's avatar
Philipp Arras committed
53
    ms2 = bl.vis2ms(ng.grid2vis(bl, gconf, idx, gconf.dirty2grid(dirty2)), idx)
54
55
56
57
58
59
60
61
62
    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):
63
    np.random.seed(42)
64
65
66
67
68
69
70
71
72
73
74
75
76
77
    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
78
    weights = np.random.rand(nrow, nchan)
79
80
81
    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
82

83
84
85
86
87
88
89
90
91
92
93
    # -----------------------------
    # Test real grid case (ms2grid)
    # -----------------------------
    real_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.float64)
    grid = ng.ms2grid(baselines, gconf, idx, ms, grid_in=None)
    grid2 = ng.ms2grid(baselines, gconf, idx, ms, grid_in=real_grid)

    assert_array_almost_equal(grid, grid2)
    assert grid.dtype == real_grid.dtype == grid2.dtype
    assert id(grid2) == id(real_grid)  # Same grid object

Simon Perkins's avatar
Simon Perkins committed
94
    real_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
95
96
97
    grid = ng.ms2grid(baselines, gconf, idx, ms, grid_in=None, wgt=weights)
    grid2 = ng.ms2grid(baselines, gconf, idx, ms, grid_in=real_grid,
                       wgt=weights)
Simon Perkins's avatar
Simon Perkins committed
98
99
100
101
102

    assert_array_almost_equal(grid, grid2)
    assert grid.dtype == real_grid.dtype == grid2.dtype
    assert id(grid2) == id(real_grid)  # Same grid object

103
104
105
106
107
108
109
110
    # ------------------------------
    # Test real grid case (vis2grid)
    # ------------------------------
    real_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.float64)
    grid = ng.vis2grid(baselines, gconf, idx, vis, grid_in=None)
    grid2 = ng.vis2grid(baselines, gconf, idx, vis, grid_in=real_grid)

    # Almost same result
111
    assert_array_almost_equal(grid, grid2)
112
113
114
115
116
117
118
119
120
    assert grid.dtype == real_grid.dtype == grid2.dtype
    assert id(grid2) == id(real_grid)  # Same grid object

    # ----------------------------------
    # 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
121

122
123
124
125
126
    # 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
127
    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
128
129
130
    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
131
132
133
134
135
136

    # 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

137
138
139
140
141
142
143
144
145
146
147
    # ----------------------------------
    # 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
148

Martin Reinecke's avatar
Martin Reinecke committed
149
150

def test_pickling():
151
    np.random.seed(42)
Martin Reinecke's avatar
Martin Reinecke committed
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
    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())
173
174


Philipp Arras's avatar
Philipp Arras committed
175
@pmp('nx', [4, 18, 54])
Philipp Arras's avatar
Philipp Arras committed
176
@pmp('fov', [0.1, 1, 5])  # deg
Philipp Arras's avatar
Philipp Arras committed
177
@pmp('w', [0, 10, 8489])
Philipp Arras's avatar
Philipp Arras committed
178
def test_wstacking(nx, fov, w):
179
    np.random.seed(42)
Philipp Arras's avatar
Philipp Arras committed
180
    dx = fov*np.pi/180/nx
Philipp Arras's avatar
Philipp Arras committed
181
182
183
184
185
    ny, dy = nx, dx
    conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy)
    x, y = conf.Nu(), conf.Nv()
    fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
    res0 = (_wscreen(nx, dy, w).conjugate()*conf.grid2dirty_c(fld)).real
Martin Reinecke's avatar
test1    
Martin Reinecke committed
186
    res1 = conf.apply_wscreen(conf.grid2dirty_c(fld), w, True).real
Philipp Arras's avatar
Philipp Arras committed
187
188
189
190
    np.testing.assert_allclose(res0, res1)

    fld = np.random.randn(nx, ny)
    res0 = conf.dirty2grid_c(_wscreen(nx, dx, w)*fld.real)
Martin Reinecke's avatar
test1    
Martin Reinecke committed
191
    res1 = conf.dirty2grid_c(conf.apply_wscreen(fld, w, False))
Philipp Arras's avatar
Philipp Arras committed
192
    np.testing.assert_allclose(res0, res1)
193
194
195
196
197
198
199


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


Philipp Arras's avatar
Philipp Arras committed
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
@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
    w = conf.W()
    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
261
262
263
264
265
@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))
Philipp Arras's avatar
Philipp Arras committed
266
267
@pmp("complex_mode", (False, True))
def test_against_dft(nxdirty, nydirty, epsilon, nchan, nrow, complex_mode):
Philipp Arras's avatar
Philipp Arras committed
268
269
270
271
    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)

Philipp Arras's avatar
Philipp Arras committed
272
273
274
275
    if complex_mode:
        res0 = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, idx, vis))
    else:
        res0 = conf.grid2dirty(ng.vis2grid(bl, conf, idx, vis))
Philipp Arras's avatar
Philipp Arras committed
276
277
278
279
280

    x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]],
                       indexing='ij')
    x *= conf.Pixsize_x()
    y *= conf.Pixsize_y()
Philipp Arras's avatar
Philipp Arras committed
281
282
    dt = np.complex128 if complex_mode else np.float64
    res1 = np.zeros_like(res0, dtype=dt)
Philipp Arras's avatar
Philipp Arras committed
283
284
285
    uvw = bl.effectiveuvw(idx)
    for ii in idx:
        phase = x*uvw[ii, 0] + y*uvw[ii, 1]
Philipp Arras's avatar
Philipp Arras committed
286
287
288
289
290
        foo = vis[ii]*np.exp(2j*np.pi*phase)
        if complex_mode:
            res1 += foo
        else:
            res1 += foo.real
Philipp Arras's avatar
Philipp Arras committed
291
    assert_(_l2error(res0, res1) < epsilon)
292
293


Philipp Arras's avatar
Philipp Arras committed
294
@pmp('nxdirty', [16, 64])
295
@pmp('nydirty', [64])
Philipp Arras's avatar
Philipp Arras committed
296
297
298
299
300
@pmp("nrow", (10, 100, 1000))
@pmp("nchan", (1, 10))
@pmp("fov", (1,))
def test_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
    epsilon = 1e-7
301
    np.random.seed(40)
Philipp Arras's avatar
Philipp Arras committed
302
    pixsize = fov*np.pi/180/nxdirty
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
                            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)
    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 = np.zeros((nxdirty, nydirty))
    mi, ma = np.min(uvw[:, 2]), np.max(uvw[:, 2])
Philipp Arras's avatar
Philipp Arras committed
320
321
322
323
324
325
326
    nplanes = 10000
    ws = mi + np.arange(nplanes)*(ma-mi)/(nplanes-1)
    for ii in range(len(ws)-1):
        wkp1 = ws[ii+1]
        if ii == nplanes-2:
            wkp1 += abs(wkp1)
        jj = ng.getIndices(bl, conf, flags, wmin=ws[ii], wmax=wkp1)
327
328
329
        if len(jj) == 0:
            continue
        dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, jj, bl.ms2vis(ms, jj)))
Philipp Arras's avatar
Philipp Arras committed
330
331
        wforplane = 0.5*(ws[ii+1]+ws[ii])
        res0 += conf.apply_wscreen(dd, wforplane, adjoint=False).real
332
333
334
335
336
337
338

    # Compute dft with w term
    x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]],
                       indexing='ij')
    x *= conf.Pixsize_x()
    y *= conf.Pixsize_y()
    res1 = np.zeros_like(res0)
Philipp Arras's avatar
Philipp Arras committed
339
340
341
342
343

    eps = np.sqrt(x**2+y**2)
    s = np.sin(eps)
    nm1 = -s*s/(1+np.cos(eps))
    n = nm1+1
344
    for ii in idx:
Philipp Arras's avatar
Philipp Arras committed
345
        phase = x*uvw[ii, 0] + y*uvw[ii, 1] + uvw[ii, 2]*nm1
346
347
        res1 += (vis[ii]*np.exp(2j*np.pi*phase)).real
    res1 /= n
Philipp Arras's avatar
Philipp Arras committed
348
    assert_(_l2error(res0, res1) < 1e-4)
349

Philipp Arras's avatar
Philipp Arras committed
350

Philipp Arras's avatar
Philipp Arras committed
351
352
353
354
355
@pmp('nxdirty', (32, 16))
@pmp('nydirty', (32, 64))
@pmp("nrow", (100, 1000))
@pmp("nchan", (1, 5))
@pmp('epsilon', (1e-3, 1e-6))
356
def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
Philipp Arras's avatar
Philipp Arras committed
357
358
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
Philipp Arras's avatar
Philipp Arras committed
359
    nx, ny = 2*nxdirty, 2*nydirty
Philipp Arras's avatar
Philipp Arras committed
360
    grid = np.random.randn(nx, ny) + 1.0j*np.random.randn(nx, ny)
361
    W = conf.W()
Philipp Arras's avatar
Philipp Arras committed
362
    res0 = np.zeros_like(grid)
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381

    # 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)):
        print(du)
        d0.append(ng.get_correlations(bl, conf, idx, du=du, dv=dv))
    for du in range(1, W):
        d1.append(ng.get_correlations(bl, conf, idx, du=du, dv=0))
    d2 = ng.get_correlations(bl, conf, idx, du=0, dv=0)

    # 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
382
    res1 = ng.apply_holo(bl, conf, idx, grid)
Philipp Arras's avatar
Philipp Arras committed
383
    assert_allclose(res0, res1, rtol=1e-13)