Commit c0d22569 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

Merge branch 'hoisted-grid-allocation-experiments' into 'master'

User supplied grid to vis2grid_c

See merge request !3
parents 87d0487a 8549edf4
......@@ -40,6 +40,8 @@ namespace py = pybind11;
namespace {
auto None = py::none();
//
// basic utilities
//
......@@ -169,7 +171,7 @@ void checkArray(const py::array &arr, const char *aname,
template<typename T> pyarr<T> provideArray(py::object &in,
const vector<size_t> &shape)
{
if (in.is(py::none()))
if (in.is_none())
{
auto tmp_ = makeArray<T>(shape);
size_t sz = size_t(tmp_.size());
......@@ -183,14 +185,31 @@ template<typename T> pyarr<T> provideArray(py::object &in,
return tmp_;
}
template<typename T> pyarr_c<T> provideCArray(py::object &in,
const vector<size_t> &shape)
{
if (in.is_none())
{
auto tmp_ = makeArray<T>(shape);
size_t sz = size_t(tmp_.size());
auto tmp = tmp_.mutable_data();
for (size_t i=0; i<sz; ++i)
tmp[i] = T(0);
return tmp_;
}
auto tmp_ = in.cast<pyarr_c<T>>();
checkArray(tmp_, "temporary", shape);
return tmp_;
}
template<typename T> pyarr_c<T> complex2hartley
(const pyarr_c<complex<T>> &grid_)
(const pyarr_c<complex<T>> &grid_, py::object &grid_in)
{
checkArray(grid_, "grid", {0,0});
size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
auto grid = grid_.data();
auto res = makeArray<T>({nu,nv});
auto res = provideCArray<T>(grid_in, {nu, nv});
auto grid2 = res.mutable_data();
{
py::gil_scoped_release release;
......@@ -203,8 +222,8 @@ template<typename T> pyarr_c<T> complex2hartley
size_t xv = (v==0) ? 0 : nv-v;
size_t i1 = u*nv+v;
size_t i2 = xu*nv+xv;
grid2[i1] = T(0.5)*(grid[i1].real()+grid[i1].imag()+
grid[i2].real()-grid[i2].imag());
grid2[i1] += T(0.5)*(grid[i1].real()+grid[i1].imag()+
grid[i2].real()-grid[i2].imag());
}
}
}
......@@ -384,7 +403,7 @@ template<typename T> class Baselines
auto t = idx(i);
auto row = t/nchan;
auto chan = t-row*nchan;
ms(row, chan) = vis(i);
ms(row, chan) += vis(i);
}
}
return res;
......@@ -648,7 +667,8 @@ template<typename T> class Helper
template<typename T> pyarr_c<complex<T>> vis2grid_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &vis_)
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &vis_,
py::object &grid_in)
{
checkArray(vis_, "vis", {0});
size_t nvis = size_t(vis_.shape(0));
......@@ -657,11 +677,10 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c(
auto idx = idx_.template unchecked<1>();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = makeArray<complex<T>>({nu, nv});
auto res = provideCArray<complex<T>>(grid_in, {nu, nv});
auto grid = res.mutable_data();
{
py::gil_scoped_release release;
for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
T beta = gconf.Beta();
size_t w = gconf.W();
......@@ -696,12 +715,13 @@ 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<uint32_t> &idx_,
const pyarr<complex<T>> &vis_)
{ return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_)); }
const pyarr<complex<T>> &vis_, py::object &grid_in)
{ return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_, None), grid_in); }
template<typename T> pyarr_c<complex<T>> ms2grid_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_)
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
py::object &grid_in)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
......@@ -712,11 +732,10 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
auto idx = idx_.template unchecked<1>();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = makeArray<complex<T>>({nu, nv});
auto res = provideCArray<complex<T>>(grid_in, {nu, nv});
auto grid = res.mutable_data();
{
py::gil_scoped_release release;
for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
T beta = gconf.Beta();
size_t w = gconf.W();
......@@ -754,13 +773,13 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<complex<T>> &ms_)
{ return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_)); }
const pyarr<complex<T>> &ms_, py::object &grid_in)
{ return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_, None), grid_in); }
template<typename T> pyarr_c<complex<T>> ms2grid_c_wgt(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
const pyarr<T> &wgt_)
const pyarr<T> &wgt_, py::object &grid_in)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
......@@ -773,11 +792,10 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c_wgt(
auto idx = idx_.template unchecked<1>();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = makeArray<complex<T>>({nu, nv});
auto res = provideArray<complex<T>>(grid_in, {nu, nv});
auto grid = res.mutable_data();
{
py::gil_scoped_release release;
for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
T beta = gconf.Beta();
size_t w = gconf.W();
......@@ -815,8 +833,9 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c_wgt(
template<typename T> pyarr_c<T> ms2grid_wgt(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
const pyarr<complex<T>> &ms_, const pyarr<T> &wgt_)
{ return complex2hartley(ms2grid_c_wgt(baselines, gconf, idx_, ms_, wgt_)); }
const pyarr<complex<T>> &ms_, const pyarr<T> &wgt_,
py::object &grid_in)
{ return complex2hartley(ms2grid_c_wgt(baselines, gconf, idx_, ms_, wgt_, None), grid_in); }
template<typename T> pyarr_c<complex<T>> grid2vis_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
......@@ -1284,7 +1303,7 @@ PYBIND11_MODULE(nifty_gridder, m)
.def ("ms2vis",&Baselines<double>::ms2vis<complex<double>>, BL_ms2vis_DS,
"ms"_a, "idx"_a)
.def ("vis2ms",&Baselines<double>::vis2ms<complex<double>>, BL_vis2ms_DS,
"vis"_a, "idx"_a, "ms_in"_a=py::none());
"vis"_a, "idx"_a, "ms_in"_a=None);
py::class_<GridderConfig<double>> (m, "GridderConfig", GridderConfig_DS)
.def(py::init<size_t, size_t, double, double, double>(),"nxdirty"_a,
"nydirty"_a, "epsilon"_a, "pixsize_x"_a, "pixsize_y"_a)
......@@ -1325,28 +1344,29 @@ PYBIND11_MODULE(nifty_gridder, m)
"gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
"wmin"_a=-1e30, "wmax"_a=1e30);
m.def("vis2grid",&vis2grid<double>, vis2grid_DS, "baselines"_a, "gconf"_a,
"idx"_a, "vis"_a);
m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a);
"idx"_a, "vis"_a, "grid_in"_a=None);
m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a,
"grid_in"_a=None);
m.def("ms2grid_wgt",&ms2grid_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
"ms"_a, "wgt"_a);
"ms"_a, "wgt"_a, "grid_in"_a=None);
m.def("grid2vis",&grid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a,
"idx"_a, "grid"_a);
m.def("grid2ms",&grid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "ms_in"_a=py::none());
"grid"_a, "ms_in"_a=None);
m.def("grid2ms_wgt",&grid2ms_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "wgt"_a, "ms_in"_a=py::none());
"grid"_a, "wgt"_a, "ms_in"_a=None);
m.def("vis2grid_c",&vis2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"vis"_a);
"vis"_a, "grid_in"_a);
m.def("ms2grid_c",&ms2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"ms"_a);
"ms"_a, "grid_in"_a=None);
m.def("ms2grid_c_wgt",&ms2grid_c_wgt<double>, "baselines"_a, "gconf"_a,
"idx"_a, "ms"_a, "wgt"_a);
"idx"_a, "ms"_a, "wgt"_a, "grid_in"_a=None);
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);
m.def("grid2ms_c_wgt",&grid2ms_c_wgt<double>, "baselines"_a, "gconf"_a,
"idx"_a, "grid"_a, "wgt"_a, "ms_in"_a=py::none());
"idx"_a, "grid"_a, "wgt"_a, "ms_in"_a=None);
m.def("apply_holo",&apply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a);
}
import nifty_gridder as ng
import numpy as np
import pytest
from numpy.testing import assert_, assert_allclose
from numpy.testing import assert_, assert_allclose, assert_array_almost_equal
pmp = pytest.mark.parametrize
......@@ -17,23 +17,118 @@ 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)
weights = np.random.rand(nrow, nchan)
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)
# -----------------------------
# 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
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)
assert_array_almost_equal(grid, grid2)
assert grid.dtype == real_grid.dtype == grid2.dtype
assert id(grid2) == id(real_grid) # Same grid object
# ------------------------------
# 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
assert_array_almost_equal(grid, grid2)
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)
# 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
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)
# 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
# ----------------------------------
# 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
def test_pickling():
......@@ -58,3 +153,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