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

add wgt argument to vis2dirty/dirty2vis

parent 9f85bcd9
......@@ -1346,9 +1346,10 @@ template<typename T, typename Serv> void x2dirty(
}
template<typename T> void vis2dirty(const Baselines &baselines,
const GridderConfig &gconf, const const_mav<uint32_t,1> &idx,
const const_mav<complex<T>,1> &vis, mav<T,2> &dirty, bool do_wstacking)
const const_mav<complex<T>,1> &vis, const const_mav<T,1> &wgt,
mav<T,2> &dirty, bool do_wstacking)
{
x2dirty(gconf, makeVisServ(baselines, idx, vis, nullmav<T,1>()), dirty, do_wstacking);
x2dirty(gconf, makeVisServ(baselines, idx, vis, wgt), dirty, do_wstacking);
}
template<typename T, typename Serv> void dirty2x(
......@@ -1450,9 +1451,10 @@ template<typename T, typename Serv> void dirty2x(
}
template<typename T> void dirty2vis(const Baselines &baselines,
const GridderConfig &gconf, const const_mav<uint32_t,1> &idx,
const const_mav<T,2> &dirty, mav<complex<T>,1> &vis, bool do_wstacking)
const const_mav<T,2> &dirty, const const_mav<T,1> &wgt,
mav<complex<T>,1> &vis, bool do_wstacking)
{
dirty2x(gconf, dirty, makeVisServ(baselines, idx, vis, nullmav<T,1>()), do_wstacking);
dirty2x(gconf, dirty, makeVisServ(baselines, idx, vis, wgt), do_wstacking);
}
......
......@@ -698,7 +698,7 @@ pyarr<uint32_t> PygetIndices(const PyBaselines &baselines,
template<typename T> pyarr<T> vis2dirty2(const PyBaselines &baselines,
const PyGridderConfig &gconf, const py::array &idx_,
const py::array &vis_, bool do_wstacking)
const py::array &vis_, const py::object &wgt_, bool do_wstacking)
{
auto idx = getPyarr<uint32_t>(idx_, "idx");
auto idx2 = make_const_mav<1>(idx);
......@@ -706,49 +706,53 @@ template<typename T> pyarr<T> vis2dirty2(const PyBaselines &baselines,
auto dirty2 = make_mav<2>(dirty);
auto vis = getPyarr<complex<T>>(vis_, "vis");
auto vis2 = make_const_mav<1>(vis);
auto wgt = providePotentialArray<T>(wgt_, "wgt", {idx2.shape(0)});
auto wgt2 = make_const_mav<1>(wgt);
{
py::gil_scoped_release release;
vis2dirty<T>(baselines, gconf, idx2, vis2, dirty2, do_wstacking);
vis2dirty<T>(baselines, gconf, idx2, vis2, wgt2, dirty2, do_wstacking);
}
return dirty;
}
py::array Pyvis2dirty(const PyBaselines &baselines,
const PyGridderConfig &gconf, const py::array &idx,
const py::array &vis, bool do_wstacking)
const py::array &vis, const py::object &wgt, bool do_wstacking)
{
if (isPytype<complex<float>>(vis))
return vis2dirty2<float>(baselines, gconf, idx, vis, do_wstacking);
return vis2dirty2<float>(baselines, gconf, idx, vis, wgt, do_wstacking);
if (isPytype<complex<double>>(vis))
return vis2dirty2<double>(baselines, gconf, idx, vis, do_wstacking);
return vis2dirty2<double>(baselines, gconf, idx, vis, wgt, do_wstacking);
myfail("type matching failed: 'vis' has neither type 'c8' nor 'c16'");
}
template<typename T> pyarr<complex<T>> dirty2vis2(const PyBaselines &baselines,
const PyGridderConfig &gconf, const pyarr<uint32_t> &idx_,
const pyarr<T> &dirty_, bool do_wstacking)
const pyarr<T> &dirty_, const py::object &wgt_, bool do_wstacking)
{
auto idx = getPyarr<uint32_t>(idx_, "idx");
auto idx2 = make_const_mav<1>(idx);
auto dirty = getPyarr<T>(dirty_, "dirty");
auto dirty2 = make_const_mav<2>(dirty_);
auto wgt = providePotentialArray<T>(wgt_, "wgt", {idx2.shape(0)});
auto wgt2 = make_const_mav<1>(wgt);
auto vis = makeArray<complex<T>>({idx2.shape(0)});
auto vis2 = make_mav<1>(vis);
vis2.fill(0);
{
py::gil_scoped_release release;
vis2.fill(0);
dirty2vis<T>(baselines, gconf, idx2, dirty2, vis2, do_wstacking);
dirty2vis<T>(baselines, gconf, idx2, dirty2, wgt2, vis2, do_wstacking);
}
return vis;
}
py::array Pydirty2vis(const PyBaselines &baselines,
const PyGridderConfig &gconf, const py::array &idx,
const py::array &dirty, bool do_wstacking)
const PyGridderConfig &gconf, const py::array &idx, const py::array &dirty,
const py::object &wgt, bool do_wstacking)
{
if (isPytype<float>(dirty))
return dirty2vis2<float>(baselines, gconf, idx, dirty, do_wstacking);
return dirty2vis2<float>(baselines, gconf, idx, dirty, wgt, do_wstacking);
if (isPytype<double>(dirty))
return dirty2vis2<double>(baselines, gconf, idx, dirty, do_wstacking);
return dirty2vis2<double>(baselines, gconf, idx, dirty, wgt, do_wstacking);
myfail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
}
......@@ -906,9 +910,9 @@ PYBIND11_MODULE(nifty_gridder, m)
m.def("get_correlations", &Pyget_correlations<double>, "baselines"_a, "gconf"_a,
"idx"_a, "du"_a, "dv"_a, "wgt"_a=None);
m.def("vis2dirty",&Pyvis2dirty, "baselines"_a, "gconf"_a,
"idx"_a, "vis"_a, "do_wstacking"_a=false);
"idx"_a, "vis"_a, "wgt"_a=None, "do_wstacking"_a=false);
m.def("dirty2vis",&Pydirty2vis, "baselines"_a, "gconf"_a,
"idx"_a, "dirty"_a, "do_wstacking"_a=false);
"idx"_a, "dirty"_a, "wgt"_a=None, "do_wstacking"_a=false);
m.def("ms2dirty",&Pyms2dirty,"uvw"_a,"freq"_a,"ms"_a,
"wgt"_a=None,"npix_x"_a,"npix_y"_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