Commit 6acee462 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

Merge branch 'more_weights' into 'master'

Add optional `wgt` argument for many functions

See merge request !11
parents 17f8b827 ddff3959
This diff is collapsed.
......@@ -6,34 +6,46 @@ 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,
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)
uvw[:, 2] = 0.
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
@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):
f0 = 1e9 # rough observation frequency
pixsize = np.pi/180/60/nxdirty # assume 1 arcmin FOV
speedoflight = 3e8
freq = f0 + np.arange(nchan)*(f0/nchan)
uvw = (np.random.rand(nrow, 3)-0.5) / (pixsize*f0/speedoflight)
uvw[:, 2] = 0.
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)
idx = ng.getIndices(baselines, gconf, flags)
np.random.seed(42)
bl, gconf, 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 = baselines.ms2vis(ms, idx)
grid = ng.vis2grid(baselines, gconf, idx, vis)
vis = bl.ms2vis(ms, idx)
grid = ng.vis2grid(bl, gconf, idx, vis)
dirty = gconf.grid2dirty(grid)
dirty2 = np.random.rand(*dirty.shape)-0.5
ms2 = baselines.vis2ms(ng.grid2vis(baselines, gconf, idx,
gconf.dirty2grid(dirty2)), idx)
ms2 = bl.vis2ms(ng.grid2vis(bl, gconf, idx, gconf.dirty2grid(dirty2)), idx)
assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty, dirty2))
......@@ -43,6 +55,7 @@ def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
@pmp("nchan", (100,))
@pmp("epsilon", (2e-13,))
def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
np.random.seed(42)
gconf = ng.GridderConfig(1024, 1024, 2e-13, 2.0, 2.0)
freq = np.linspace(.856e9, 2*.856e9, nchan)
f0 = freq[freq.shape[0]//2]
......@@ -74,10 +87,9 @@ def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
assert id(grid2) == id(real_grid) # Same grid object
real_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.float64)
grid = ng.ms2grid_wgt(baselines, gconf, idx, ms,
weights, grid_in=None)
grid2 = ng.ms2grid_wgt(baselines, gconf, idx, ms,
weights, grid_in=real_grid)
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)
assert_array_almost_equal(grid, grid2)
assert grid.dtype == real_grid.dtype == grid2.dtype
......@@ -108,10 +120,9 @@ def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
assert id(grid2) == id(cplx_grid) # Same grid object
cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
grid = ng.ms2grid_c_wgt(baselines, gconf, idx, ms,
weights, grid_in=None)
grid2 = ng.ms2grid_c_wgt(baselines, gconf, idx, ms,
weights, grid_in=cplx_grid)
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)
# Almost same result
assert_array_almost_equal(grid, grid2)
......@@ -132,6 +143,7 @@ def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
def test_pickling():
np.random.seed(42)
try:
import cPickle as pickle
except ImportError:
......@@ -155,19 +167,11 @@ def test_pickling():
assert_(gc.Pixsize_y() == gc2.Pixsize_y())
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
@pmp('nx', [4, 18, 54])
@pmp('dx', [1., 0.13, 132])
@pmp('w', [0, 10, 8489])
def test_wstacking(nx, dx, w):
np.random.seed(42)
ny, dy = nx, dx
conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy)
x, y = conf.Nu(), conf.Nv()
......@@ -180,3 +184,51 @@ def test_wstacking(nx, dx, w):
res0 = conf.dirty2grid_c(_wscreen(nx, dx, w)*fld.real)
res1 = conf.dirty2grid_c(conf.apply_wscreen(fld, w, False))
np.testing.assert_allclose(res0, res1)
@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))
@pmp("weight", (True, False))
def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
np.random.seed(42)
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)
wgt = np.random.rand(*vis.shape) if weight else None
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)
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))
@pmp("weight", (True, False))
def test_correlations(nxdirty, nydirty, nrow, nchan, epsilon, du, dv, weight):
np.random.seed(42)
bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
if conf.W() <= abs(du) or conf.W() <= dv:
return
wgt = np.random.rand(*idx.shape) if weight else None
y0 = ng.get_correlations(bl, conf, idx, du, dv, wgt)
pts = ((nxdirty, nydirty),
(nxdirty + 1, nydirty),
(nxdirty, nydirty//2),
(0, 0),
(2*nxdirty-1, 2*nydirty-1))
for pp in pts:
grid = np.zeros((2*nxdirty, 2*nydirty))
grid[pp] = 1
y1 = ng.apply_holo(bl, conf, idx, grid, wgt)
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)
Supports Markdown
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