Commit e631d7e8 authored by Martin Reinecke's avatar Martin Reinecke

add gridding/degridding with weights

parent aa401fbb
......@@ -757,6 +757,67 @@ template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines,
const pyarr<complex<T>> &ms_)
{ return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_)); }
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_)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
checkArray(wgt_, "wgt", {nrows, nchan});
checkArray(ms_, "ms", {nrows, nchan});
checkArray(idx_, "idx", {0});
size_t nvis = size_t(idx_.shape(0));
auto ms = ms_.template unchecked<2>();
auto wgt = wgt_.template unchecked<2>();
auto idx = idx_.template unchecked<1>();
size_t nu=gconf.Nu(), nv=gconf.Nv();
auto res = makeArray<complex<T>>({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();
#pragma omp parallel
{
Helper<T> hlp(gconf, nullptr, grid);
T emb = exp(-2*beta);
int jump = hlp.lineJump();
const T * RESTRICT ku = hlp.kernel.data();
const T * RESTRICT kv = hlp.kernel.data()+w;
// Loop over sampling points
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
auto tidx = idx(ipart);
auto row = tidx/nchan;
auto chan = tidx-row*nchan;
UVW<T> coord = baselines.effectiveCoord(tidx);
hlp.prep(coord.u, coord.v);
auto * RESTRICT ptr = hlp.p0w;
auto v(ms(row,chan)*(emb*wgt(row, chan)));
for (size_t cu=0; cu<w; ++cu)
{
complex<T> tmp(v*ku[cu]);
for (size_t cv=0; cv<w; ++cv)
ptr[cv] += tmp*kv[cv];
ptr+=jump;
}
}
} // end of parallel region
}
return res;
}
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_)); }
template<typename T> pyarr_c<complex<T>> grid2vis_c(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_)
......@@ -870,6 +931,72 @@ template<typename T> pyarr_c<complex<T>> grid2ms(const Baselines<T> &baselines,
const pyarr_c<T> &grid_, py::object &ms_in)
{ return grid2ms_c(baselines, gconf, idx_, hartley2complex(grid_), ms_in); }
template<typename T> pyarr_c<complex<T>> grid2ms_c_wgt(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_,
const pyarr<T> &wgt_, py::object &ms_in)
{
auto nrows = baselines.Nrows();
auto nchan = baselines.Nchannels();
checkArray(wgt_, "wgt", {nrows, nchan});
auto wgt = wgt_.template unchecked<2>();
size_t nu=gconf.Nu(), nv=gconf.Nv();
checkArray(idx_, "idx", {0});
auto grid = grid_.data();
checkArray(grid_, "grid", {nu, nv});
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.template unchecked<1>();
auto res = provideArray<complex<T>>(ms_in, {nrows, nchan});
auto ms = res.template mutable_unchecked<2>();
{
py::gil_scoped_release release;
T beta = gconf.Beta();
size_t w = gconf.W();
// Loop over sampling points
#pragma omp parallel
{
Helper<T> hlp(gconf, grid, nullptr);
T emb = exp(-2*beta);
int jump = hlp.lineJump();
const T * RESTRICT ku = hlp.kernel.data();
const T * RESTRICT kv = hlp.kernel.data()+w;
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
auto tidx = idx(ipart);
auto row = tidx/nchan;
auto chan = tidx-row*nchan;
UVW<T> coord = baselines.effectiveCoord(tidx);
hlp.prep(coord.u, coord.v);
complex<T> r = 0;
const auto * RESTRICT ptr = hlp.p0r;
for (size_t cu=0; cu<w; ++cu)
{
complex<T> tmp(0);
for (size_t cv=0; cv<w; ++cv)
tmp += ptr[cv] * kv[cv];
r += tmp*ku[cu];
ptr += jump;
}
ms(row,chan) += r*(emb*wgt(row,chan));
}
}
}
return res;
}
template<typename T> pyarr_c<complex<T>> grid2ms_wgt(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr_c<T> &grid_, const pyarr<T> &wgt_,
py::object &ms_in)
{
return grid2ms_c_wgt(baselines, gconf, idx_, hartley2complex(grid_), wgt_,
ms_in);
}
template<typename T> pyarr_c<complex<T>> apply_holo(
const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_)
......@@ -1200,18 +1327,26 @@ PYBIND11_MODULE(nifty_gridder, m)
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);
m.def("ms2grid_wgt",&ms2grid_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
"ms"_a, "wgt"_a);
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());
m.def("grid2ms_wgt",&grid2ms_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a, "wgt"_a, "ms_in"_a=py::none());
m.def("vis2grid_c",&vis2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"vis"_a);
m.def("ms2grid_c",&ms2grid_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
"ms"_a);
m.def("ms2grid_c_wgt",&ms2grid_c_wgt<double>, "baselines"_a, "gconf"_a,
"idx"_a, "ms"_a, "wgt"_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());
m.def("grid2ms_c_wgt",&grid2ms_c_wgt<double>, "baselines"_a, "gconf"_a,
"idx"_a, "grid"_a, "wgt"_a, "ms_in"_a=py::none());
m.def("apply_holo",&apply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
"grid"_a);
}
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