Commit b9f6029b authored by Simon Perkins's avatar Simon Perkins
Browse files

User supplied grid to vis2grid_c

parent 8e7b0e6a
......@@ -633,16 +633,21 @@ template<typename T> class Helper
template<typename T> pyarr_c<complex<T>> vis2grid_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr_c<uint32_t> &idx_, const pyarr_c<complex<T>> &vis_)
const pyarr_c<uint32_t> &idx_, const pyarr_c<complex<T>> &vis_,
pyarr_c<complex<T>> * user_grid)
{
checkArray(vis_, "vis", {0});
size_t nvis = size_t(vis_.shape(0));
size_t nu=gconf.Nu(), nv=gconf.Nv();
checkArray(idx_, "idx", {nvis});
if(user_grid)
{ checkArray(*user_grid, "user_grid", {nu, nv}); }
auto vis=vis_.data();
auto idx = idx_.data();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = makeArray<complex<T>>({nu, nv});
auto res = user_grid ? makeArray<complex<T>>({nu, nv}) : *user_grid;
auto grid = res.mutable_data();
{
py::gil_scoped_release release;
......@@ -681,7 +686,11 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c(
template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<complex<T>> &vis_)
{ return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_)); }
{
return complex2hartley(vis2grid_c(baselines, gconf,
idx_, vis_,
static_cast<pyarr_c<std::complex<T>> *>(nullptr)));
}
template<typename T> pyarr_c<complex<T>> ms2grid_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
......@@ -1119,15 +1128,15 @@ PYBIND11_MODULE(nifty_gridder, m)
"idx"_a, "vis"_a);
m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a);
m.def("grid2vis",&grid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a,
"idx"_a, "grid"_a);
"idx"_a, "grid"_a=py::none());
m.def("grid2ms",&grid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=py::none());
m.def("vis2grid_c",&vis2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"vis"_a);
"vis"_a, "user_grid"_a);
m.def("ms2grid_c",&ms2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"ms"_a);
m.def("grid2vis_c",&grid2vis_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a);
m.def("grid2ms_c",&grid2ms_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=py::none());
"grid"_a, "ms_in"_a.none(true));
}
......@@ -17,23 +17,55 @@ def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
speedoflight = 3e8
freq = f0 + np.arange(nchan)*(f0/nchan)
uvw = (np.random.rand(nrow,3)-0.5) / (pixsize*f0/speedoflight)
uvw[:,2] = 0.
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)
epsilon=epsilon, pixsize_x=pixsize,
pixsize_y=pixsize)
flags = np.zeros((nrow, nchan), dtype=np.bool)
idx = ng.getIndices(baselines, gconf, flags)
ms = np.random.rand(nrow,nchan)-0.5 + 1j*(np.random.rand(nrow,nchan)-0.5)
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)
dirty = gconf.grid2dirty(grid)
dirty2 = np.random.rand(*dirty.shape)-0.5
ms2 = baselines.vis2ms(ng.grid2vis(baselines, gconf, idx,
gconf.dirty2grid(dirty2)), idx)
gconf.dirty2grid(dirty2)), idx)
assert_allclose(np.vdot(ms,ms2).real, np.vdot(dirty, dirty2))
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):
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)
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)
user_grid = np.zeros((gconf.Nu(), gconf.Nv()),
dtype=np.complex128)
grid = ng.vis2grid_c(baselines, gconf, idx, vis, user_grid=user_grid)
# Same base array under the hood
assert grid.base is user_grid.base
def test_pickling():
......@@ -58,3 +90,5 @@ def test_pickling():
assert_(gc.Epsilon() == gc2.Epsilon())
assert_(gc.Pixsize_x() == gc2.Pixsize_x())
assert_(gc.Pixsize_y() == gc2.Pixsize_y())
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()
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