Commit 53e3a866 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

add high-level routines

parent b26eebb6
......@@ -949,7 +949,7 @@ template<typename T> void get_correlations
template<typename T, typename T2> void apply_wcorr(const GridderConfig<T> &gconf,
mav<T2,2> &dirty, const EC_Kernel_with_correction<T> &kernel, double dw)
const mav<T2,2> &dirty, const EC_Kernel_with_correction<T> &kernel, double dw)
{
auto nx_dirty=gconf.Nxdirty();
auto ny_dirty=gconf.Nydirty();
......@@ -983,7 +983,7 @@ template<typename T, typename T2> void apply_wcorr(const GridderConfig<T> &gconf
}
template<typename T, typename Serv> void x2dirty_wstack(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const Serv &srv, mav<complex<T>,2> &dirty)
const GridderConfig<T> &gconf, const Serv &srv, const mav<complex<T>,2> &dirty)
{
auto nx_dirty=gconf.Nxdirty();
auto ny_dirty=gconf.Nydirty();
......@@ -1180,6 +1180,74 @@ template<typename T> void dirty2vis_wstack(const Baselines<T> &baselines,
dirty2x_wstack(baselines, gconf, dirty, VisServ<T,mav<complex<T>,1>,const_mav<T,1>>(baselines, idx, vis, const_mav<T,1>(nullptr,{0})));
}
template<typename T> vector<uint32_t> getWgtIndices(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const const_mav<T,2> &wgt)
{
size_t nrow=baselines.Nrows(),
nchan=baselines.Nchannels(),
nsafe=gconf.Nsafe();
bool have_wgt=wgt.size()!=0;
if (have_wgt) checkShape(wgt.shape(),{nrow,nchan});
constexpr int side=1<<logsquare;
size_t nbu = (gconf.Nu()+1+side-1) >> logsquare,
nbv = (gconf.Nv()+1+side-1) >> logsquare;
vector<uint32_t> acc(nbu*nbv+1, 0);
vector<uint32_t> tmp(nrow*nchan);
for (size_t irow=0, idx=0; irow<nrow; ++irow)
for (size_t ichan=0; ichan<nchan; ++ichan)
if ((!have_wgt) || (wgt(irow,ichan)!=0))
{
auto uvw = baselines.effectiveCoord(RowChan{irow,size_t(ichan)});
T u, v;
int iu0, iv0;
gconf.getpix(uvw.u, uvw.v, u, v, iu0, iv0);
iu0 = (iu0+nsafe)>>logsquare;
iv0 = (iv0+nsafe)>>logsquare;
++acc[nbv*iu0 + iv0 + 1];
tmp[idx++] = nbv*iu0 + iv0;
}
for (size_t i=1; i<acc.size(); ++i)
acc[i] += acc[i-1];
vector<uint32_t> res(acc.back());
for (size_t irow=0, idx=0; irow<nrow; ++irow)
for (size_t ichan=0; ichan<nchan; ++ichan)
if ((!have_wgt) || (wgt(irow,ichan)!=0))
res[acc[tmp[idx++]]++] = irow*nchan+ichan;
return res;
}
template<typename T> void full_gridding(const const_mav<T,2> &uvw,
const const_mav<T,1> &freq, const const_mav<complex<T>,2> &ms,
const const_mav<T,2> &wgt, T pixsize_x, T pixsize_y, double epsilon,
const mav<complex<T>,2> &dirty)
{
Baselines<T> baselines(uvw, freq);
GridderConfig<T> gconf(dirty.shape(0), dirty.shape(1), epsilon, pixsize_x, pixsize_y);
auto idx = getWgtIndices(baselines, gconf, wgt);
auto idx2 = const_mav<uint32_t,1>(idx.data(),{idx.size()});
auto srv = MsServ<T,const_mav<complex<T>,2>,const_mav<T,2>>(baselines,idx2,ms,wgt);
x2dirty_wstack(baselines, gconf, srv, dirty);
}
template<typename T> void full_degridding(const const_mav<T,2> &uvw,
const const_mav<T,1> &freq, const const_mav<complex<T>,2> &dirty,
const const_mav<T,2> &wgt, T pixsize_x, T pixsize_y, double epsilon,
const mav<complex<T>,2> &ms)
{
Baselines<T> baselines(uvw, freq);
GridderConfig<T> gconf(dirty.shape(0), dirty.shape(1), epsilon, pixsize_x, pixsize_y);
auto idx = getWgtIndices(baselines, gconf, wgt);
auto idx2 = const_mav<uint32_t,1>(idx.data(),{idx.size()});
for (size_t i=0; i<ms.shape(0); ++i)
for (size_t j=0; j<ms.shape(1); ++j)
ms(i,j) = 0;
auto srv = MsServ<T,mav<complex<T>,2>,const_mav<T,2>>(baselines,idx2,ms,wgt);
dirty2x_wstack(baselines, gconf, dirty, srv);
}
} // namespace gridder
#endif
......@@ -614,7 +614,7 @@ np.array((nvis,), dtype=np.uint32)
the compressed indices for all entries which match the selected criteria
and are not flagged.
)""";
template<typename T> pyarr<uint32_t> getIndices(const PyBaselines<T> &baselines,
template<typename T> pyarr<uint32_t> PygetIndices(const PyBaselines<T> &baselines,
const PyGridderConfig<T> &gconf, const pyarr<bool> &flags_, int chbegin,
int chend, T wmin, T wmax)
{
......@@ -704,6 +704,36 @@ template<typename T> pyarr<complex<T>> Pydirty2vis_wstack(const PyBaselines<T> &
return vis;
}
template<typename T> pyarr<complex<T>> Pyfull_gridding(const pyarr<T> &uvw,
const pyarr<T> &freq, const pyarr<complex<T>> &ms, const py::object &wgt_,
size_t npix_x, size_t npix_y, T pixsize_x, T pixsize_y, double epsilon)
{
auto uvw2 = make_const_mav<2>(uvw);
auto freq2 = make_const_mav<1>(freq);
auto ms2 = make_const_mav<2>(ms);
auto wgt=providePotentialArray<T>(wgt_,{ms2.shape(0),ms2.shape(1)});
auto wgt2 = make_const_mav<2>(wgt);
auto dirty=makeArray<complex<T>>({npix_x,npix_y});
auto dirty2 =make_mav<2>(dirty);
full_gridding(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,dirty2);
return dirty;
}
template<typename T> pyarr<complex<T>> Pyfull_degridding(const pyarr<T> &uvw,
const pyarr<T> &freq, const pyarr<complex<T>> &dirty, const py::object &wgt_,
T pixsize_x, T pixsize_y, double epsilon)
{
auto uvw2 = make_const_mav<2>(uvw);
auto freq2 = make_const_mav<1>(freq);
auto dirty2 = make_const_mav<2>(dirty);
auto wgt=providePotentialArray<T>(wgt_,{uvw2.shape(0),freq2.shape(0)});
auto wgt2 = make_const_mav<2>(wgt);
auto ms=makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)});
auto ms2 =make_mav<2>(ms);
full_degridding(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,ms2);
return ms;
}
} // unnamed namespace
PYBIND11_MODULE(nifty_gridder, m)
......@@ -756,7 +786,7 @@ PYBIND11_MODULE(nifty_gridder, m)
t[4].cast<double>());
}));
m.def("getIndices", getIndices<double>, getIndices_DS, "baselines"_a,
m.def("getIndices", PygetIndices<double>, getIndices_DS, "baselines"_a,
"gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
"wmin"_a=-1e30, "wmax"_a=1e30);
m.def("vis2grid_c",&Pyvis2grid_c<double>, vis2grid_c_DS, "baselines"_a,
......@@ -777,4 +807,9 @@ PYBIND11_MODULE(nifty_gridder, m)
"idx"_a, "vis"_a);
m.def("dirty2vis_wstack",&Pydirty2vis_wstack<double>, "baselines"_a, "gconf"_a,
"idx"_a, "dirty"_a);
m.def("full_gridding",&Pyfull_gridding<double>,"uvw"_a,"freq"_a,"vis"_a,
"wgt"_a,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a);
m.def("full_degridding",&Pyfull_degridding<double>,"uvw"_a,"freq"_a,"dirty"_a,
"wgt"_a,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_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