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

fixes and extensions

parent 6203d28e
......@@ -454,7 +454,7 @@ template<typename T> class Baselines
size_t offset(uint32_t index) const
{ return (index>>channelbits)*nchan + (index&channelmask); }
pyarr_c<complex<T>> ms2vis(const pyarr_c<complex<T>> &ms_,
template<typename T2> pyarr_c<T2> ms2vis(const pyarr_c<T2> &ms_,
const pyarr_c<uint32_t> &idx_) const
{
myassert(idx_.ndim()==1, "idx array must be 1D");
......@@ -465,14 +465,14 @@ template<typename T> class Baselines
auto idx = idx_.data();
auto ms = ms_.data();
auto res=makearray<complex<T>>({nvis});
auto res=makearray<T2>({nvis});
auto vis = res.mutable_data();
for (size_t i=0; i<nvis; ++i)
vis[i] = ms[offset(idx[i])];
return res;
}
pyarr_c<complex<T>> vis2ms(const pyarr_c<complex<T>> &vis_,
template<typename T2> pyarr_c<T2> vis2ms(const pyarr_c<T2> &vis_,
const pyarr_c<uint32_t> &idx_) const
{
myassert(idx_.ndim()==1, "idx array must be 1D");
......@@ -482,17 +482,17 @@ template<typename T> class Baselines
auto idx = idx_.data();
auto vis = vis_.data();
auto res = makearray<complex<T>>({nrows, nchan});
auto res = makearray<T2>({nrows, nchan});
auto ms = res.mutable_data();
for (size_t i=0; i<nrows*nchan; ++i)
ms[i] = complex<T>(0);
ms[i] = T2(0);
for (size_t i=0; i<nvis; ++i)
ms[offset(idx[i])] = vis[i];
return res;
}
pyarr_c<complex<T>> add_vis_to_ms(const pyarr_c<complex<T>> &vis_,
const pyarr_c<uint32_t> &idx_, pyarr_c<complex<T>> &ms_) const
template<typename T2> pyarr_c<T2> add_vis_to_ms(const pyarr_c<T2> &vis_,
const pyarr_c<uint32_t> &idx_, pyarr_c<T2> &ms_) const
{
myassert(idx_.ndim()==1, "idx array must be 1D");
myassert(vis_.ndim()==1, "vis array must be 1D");
......@@ -550,6 +550,8 @@ template<typename T> class GridderConfig
size_t Nu() const { return nu; }
size_t Nv() const { return nv; }
size_t W() const { return w; }
T Ucorr() const { return ucorr; }
T Vcorr() const { return vcorr; }
size_t coord2peano(const UVW<T> &coord) const
{
double u=fmodulo(coord.u*ucorr, T(1))*nu,
......@@ -821,6 +823,7 @@ template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
auto res = makearray<complex<T>>({nu, nv});
auto grid = res.mutable_data();
for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
T ucorr = gconf.Ucorr(), vcorr=gconf.Vcorr();
#pragma omp parallel
{
......@@ -834,7 +837,7 @@ template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
for (size_t ipart=0; ipart<nvis; ++ipart)
{
UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
hlp.prep_write(coord.u, coord.v);
hlp.prep_write(coord.u*ucorr, coord.v*vcorr);
auto * RESTRICT ptr = hlp.p0;
int w = hlp.w;
auto v(vis[ipart]*emb);
......@@ -866,6 +869,7 @@ template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
auto res = makearray<complex<T>>({nvis});
auto vis = res.mutable_data();
T ucorr = gconf.Ucorr(), vcorr=gconf.Vcorr();
// Loop over sampling points
#pragma omp parallel
......@@ -879,7 +883,7 @@ template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
for (size_t ipart=0; ipart<nvis; ++ipart)
{
UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
hlp.prep_read(coord.u, coord.v);
hlp.prep_read(coord.u*ucorr, coord.v*vcorr);
complex<T> r = 0;
auto * RESTRICT ptr = hlp.p0;
int w = hlp.w;
......@@ -986,9 +990,10 @@ PYBIND11_MODULE(nifty_gridder, m)
py::class_<Baselines<double>> (m, "Baselines", Baselines_DS)
.def(py::init<const pyarr_c<double> &, const pyarr_c<double> &,
const pyarr_c<bool> &>(), "coord"_a, "scaling"_a, "flags"_a)
.def ("ms2vis",&Baselines<double>::ms2vis, BL_ms2vis_DS, "ms"_a, "idx"_a)
.def ("vis2ms",&Baselines<double>::vis2ms, BL_vis2ms_DS, "vis"_a, "idx"_a)
.def ("add_vis_to_ms",&Baselines<double>::add_vis_to_ms, BL_add_vis_to_ms_DS,
.def ("ms2vis",&Baselines<double>::ms2vis<complex<double>>, BL_ms2vis_DS, "ms"_a, "idx"_a)
.def ("ms2vis_f32",&Baselines<double>::ms2vis<float>, "ms"_a, "idx"_a)
.def ("vis2ms",&Baselines<double>::vis2ms<complex<double>>, BL_vis2ms_DS, "vis"_a, "idx"_a)
.def ("add_vis_to_ms",&Baselines<double>::add_vis_to_ms<complex<double>>, BL_add_vis_to_ms_DS,
"vis"_a, "idx"_a, "ms"_a.noconvert());
py::class_<GridderConfig<double>> (m, "GridderConfig")
.def(py::init<size_t, size_t, double, double, double>(),"nxdirty"_a,
......@@ -1005,9 +1010,9 @@ PYBIND11_MODULE(nifty_gridder, m)
py::class_<Baselines<float>> (m, "Baselines_f", Baselines_DS)
.def(py::init<const pyarr_c<float> &, const pyarr_c<float> &,
const pyarr_c<bool> &>(), "coord"_a, "scaling"_a, "flags"_a)
.def ("ms2vis",&Baselines<float>::ms2vis, BL_ms2vis_DS, "ms"_a, "idx"_a)
.def ("vis2ms",&Baselines<float>::vis2ms, BL_vis2ms_DS, "vis"_a, "idx"_a)
.def ("add_vis_to_ms",&Baselines<float>::add_vis_to_ms, BL_add_vis_to_ms_DS,
.def ("ms2vis",&Baselines<float>::ms2vis<complex<float>>, BL_ms2vis_DS, "ms"_a, "idx"_a)
.def ("vis2ms",&Baselines<float>::vis2ms<complex<float>>, BL_vis2ms_DS, "vis"_a, "idx"_a)
.def ("add_vis_to_ms",&Baselines<float>::add_vis_to_ms<complex<float>>, BL_add_vis_to_ms_DS,
"vis"_a, "idx"_a, "ms"_a.noconvert());
py::class_<GridderConfig<float>> (m, "GridderConfig_f")
.def(py::init<size_t, size_t, float, float, float>(),"nxdirty"_a,
......
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