Commit bddb6c4f authored by Martin Reinecke's avatar Martin Reinecke

speed up tests without making them easier

parent fd55f2d6
......@@ -7,7 +7,6 @@ from numpy.testing import assert_, assert_allclose, assert_array_almost_equal
pmp = pytest.mark.parametrize
def _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow):
pixsize = np.pi/180/60/nxdirty
conf = ng.GridderConfig(nxdirty=nxdirty,
......@@ -45,13 +44,103 @@ def _dft(uvw, vis, conf, apply_w):
def _l2error(a, b):
return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))
return np.sqrt(np.sum(np.abs(a-b)**2)/np.maximum(np.sum(np.abs(a)**2),np.sum(np.abs(b)**2)))
def explicit_gridder(uvw, freq, ms, wgt, nxdirty, nydirty, xpixsize, ypixsize, apply_w):
speedoflight = 299792458.
x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]],
indexing='ij')
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
speedoflight, f0 = 3e8, 1e9
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,
pixsizey, epsilon, wstacking, nthreads, 2).astype("f8")
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))
@pmp("fov", (1.,20.))
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
speedoflight, f0 = 3e8, 1e9
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)
@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 10, 10000))
@pmp("nchan", (1, 10, 100))
@pmp("nrow", (1, 100))
@pmp("nchan", (1, 10))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
np.random.seed(42)
......@@ -68,8 +157,8 @@ def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 10, 10000))
@pmp("nchan", (1, 10, 100))
@pmp("nrow", (1, 127))
@pmp("nchan", (1, 17))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
def test_adjointness_wgridding(nxdirty, nydirty, nrow, nchan, epsilon):
np.random.seed(42)
......@@ -82,40 +171,10 @@ def test_adjointness_wgridding(nxdirty, nydirty, nrow, nchan, epsilon):
rtol=2e-13)
@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 10, 10000))
@pmp("nchan", (1, 10, 100))
@pmp("epsilon", (1e-2, 1e-3, 5e-5, 1e-6, 1e-7, 5e-13))
@pmp("singleprec", (True, False))
@pmp("wstacking", (True, False))
def test_adjointness_wgridding_highlevel(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking):
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
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 singleprec:
vis = vis.astype("c8")
dirty = dirty.astype("f4")
dirty2 = ng.ms2dirty(uvw, freq, vis, None, nxdirty, nydirty, pixsizex,
pixsizey, epsilon, wstacking, 1, 0).astype("f8")
vis2 = ng.dirty2ms(uvw, freq, dirty, None, pixsizex, pixsizey, epsilon,
wstacking, 1, 0).astype("c16")
tol = 5e-5 if singleprec else 5e-13
assert_allclose(np.vdot(vis, vis2).real, np.vdot(dirty2, dirty),
rtol=tol)
@pmp("nxdirty", (128,))
@pmp("nydirty", (128,))
@pmp("nrow", (10000,))
@pmp("nchan", (100,))
@pmp("nrow", (1000,))
@pmp("nchan", (10,))
@pmp("epsilon", (2e-13,))
def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
np.random.seed(42)
......@@ -200,8 +259,8 @@ def test_pickling():
@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 10, 10000))
@pmp("nchan", (1, 10, 100))
@pmp("nrow", (1, 39))
@pmp("nchan", (1, 70))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
@pmp("weight", (True, False))
def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
......@@ -219,7 +278,7 @@ def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
@pmp("nxdirty", (300,))
@pmp("nydirty", (250,))
@pmp("nrow", (10, 10000))
@pmp("nrow", (10, 100))
@pmp("nchan", (1, 10))
@pmp("epsilon", (1e-2, 1e-7))
@pmp("du", (-2, 0, 1, 4))
......@@ -246,7 +305,7 @@ def test_correlations(nxdirty, nydirty, nrow, nchan, epsilon, du, dv, weight):
assert_allclose(np.zeros_like(y1), y1.imag)
@pmp("nrow", (1, 1000))
@pmp("nrow", (1, 100))
@pmp("nchan", (1, 10))
@pmp("epsilon", (1e-2, 1e-7, 2e-12))
@pmp("weight", (True, False))
......@@ -282,7 +341,7 @@ def test_against_dft(nxdirty, nydirty, epsilon, nchan, nrow):
@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100, 1000))
@pmp("nrow", (10, 100))
@pmp("nchan", (1, 10))
@pmp("fov", (1,))
def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
......@@ -321,7 +380,7 @@ def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100, 1000))
@pmp("nrow", (10, 100))
@pmp("nchan", (1,))
@pmp("fov", (50,))
def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
......@@ -353,7 +412,7 @@ def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100, 1000))
@pmp("nrow", (10, 100))
@pmp("nchan", (1, 10))
@pmp("fov", (1, 10))
@pmp("epsilon", (1e-2, 1e-7, 2e-12))
......@@ -381,7 +440,7 @@ def test_wgridder_against_wdft(nxdirty, nydirty, nchan, nrow, fov, epsilon):
@pmp('nxdirty', (32, 16))
@pmp('nydirty', (32, 64))
@pmp("nrow", (100, 1000))
@pmp("nrow", (100, 127))
@pmp("nchan", (1, 5))
@pmp('epsilon', (1e-3, 1e-6))
def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment