Commit 5e415e5d authored by Martin Reinecke's avatar Martin Reinecke

make index type more flexible

parent 42892ca2
...@@ -296,11 +296,6 @@ void legendre_prep(int n, vector<double> &x, vector<double> &w, size_t nthreads) ...@@ -296,11 +296,6 @@ void legendre_prep(int n, vector<double> &x, vector<double> &w, size_t nthreads)
// Start of real gridder functionality // Start of real gridder functionality
// //
struct RowChan
{
size_t row, chan;
};
template<typename T> void complex2hartley template<typename T> void complex2hartley
(const const_mav<complex<T>, 2> &grid, const mav<T,2> &grid2, size_t nthreads) (const const_mav<complex<T>, 2> &grid, const mav<T,2> &grid2, size_t nthreads)
{ {
...@@ -431,6 +426,13 @@ vector<double> correction_factors(size_t n, size_t nval, size_t supp, ...@@ -431,6 +426,13 @@ vector<double> correction_factors(size_t n, size_t nval, size_t supp,
return res; return res;
} }
using idx_t = uint32_t;
struct RowChan
{
idx_t row, chan;
};
struct UVW struct UVW
{ {
double u, v, w; double u, v, w;
...@@ -452,8 +454,8 @@ class Baselines ...@@ -452,8 +454,8 @@ class Baselines
protected: protected:
vector<UVW> coord; vector<UVW> coord;
vector<double> f_over_c; vector<double> f_over_c;
size_t nrows, nchan; idx_t nrows, nchan;
size_t shift, mask; idx_t shift, mask;
public: public:
template<typename T> Baselines(const const_mav<T,2> &coord_, template<typename T> Baselines(const const_mav<T,2> &coord_,
...@@ -461,12 +463,16 @@ class Baselines ...@@ -461,12 +463,16 @@ class Baselines
{ {
constexpr double speedOfLight = 299792458.; constexpr double speedOfLight = 299792458.;
myassert(coord_.shape(1)==3, "dimension mismatch"); myassert(coord_.shape(1)==3, "dimension mismatch");
auto hugeval = size_t(~(idx_t(0)));
myassert(coord_.shape(0)<hugeval, "too many entries in MS");
myassert(coord_.shape(1)<hugeval, "too many entries in MS");
myassert(coord_.size()<hugeval, "too many entries in MS");
nrows = coord_.shape(0); nrows = coord_.shape(0);
nchan = freq.shape(0); nchan = freq.shape(0);
shift=0; shift=0;
while((size_t(1)<<shift)<nchan) ++shift; while((idx_t(1)<<shift)<nchan) ++shift;
mask=(size_t(1)<<shift)-1; mask=(idx_t(1)<<shift)-1;
myassert(nrows*nchan<(size_t(1)<<32), "too many entries in MS"); myassert(nrows*nchan<hugeval, "too many entries in MS");
f_over_c.resize(nchan); f_over_c.resize(nchan);
for (size_t i=0; i<nchan; ++i) for (size_t i=0; i<nchan; ++i)
{ {
...@@ -478,19 +484,19 @@ class Baselines ...@@ -478,19 +484,19 @@ class Baselines
coord[i] = UVW(coord_(i,0), coord_(i,1), coord_(i,2)); coord[i] = UVW(coord_(i,0), coord_(i,1), coord_(i,2));
} }
RowChan getRowChan(uint32_t index) const RowChan getRowChan(idx_t index) const
{ return RowChan{index>>shift, index&mask}; } { return RowChan{index>>shift, index&mask}; }
UVW effectiveCoord(const RowChan &rc) const UVW effectiveCoord(const RowChan &rc) const
{ return coord[rc.row]*f_over_c[rc.chan]; } { return coord[rc.row]*f_over_c[rc.chan]; }
UVW effectiveCoord(uint32_t index) const UVW effectiveCoord(idx_t index) const
{ return effectiveCoord(getRowChan(index)); } { return effectiveCoord(getRowChan(index)); }
size_t Nrows() const { return nrows; } size_t Nrows() const { return nrows; }
size_t Nchannels() const { return nchan; } size_t Nchannels() const { return nchan; }
uint32_t getIdx(size_t irow, size_t ichan) const idx_t getIdx(idx_t irow, idx_t ichan) const
{ return ichan+(irow<<shift); } { return ichan+(irow<<shift); }
template<typename T> void effectiveUVW(const mav<const uint32_t,1> &idx, template<typename T> void effectiveUVW(const mav<const idx_t,1> &idx,
mav<T,2> &res) const mav<T,2> &res) const
{ {
size_t nvis = idx.shape(0); size_t nvis = idx.shape(0);
...@@ -505,7 +511,7 @@ class Baselines ...@@ -505,7 +511,7 @@ class Baselines
} }
template<typename T> void ms2vis(const mav<const T,2> &ms, template<typename T> void ms2vis(const mav<const T,2> &ms,
const mav<const uint32_t,1> &idx, mav<T,1> &vis, size_t nthreads) const const mav<const idx_t,1> &idx, mav<T,1> &vis, size_t nthreads) const
{ {
checkShape(ms.shape(), {nrows, nchan}); checkShape(ms.shape(), {nrows, nchan});
size_t nvis = idx.shape(0); size_t nvis = idx.shape(0);
...@@ -519,7 +525,7 @@ class Baselines ...@@ -519,7 +525,7 @@ class Baselines
} }
template<typename T> void vis2ms(const mav<const T,1> &vis, template<typename T> void vis2ms(const mav<const T,1> &vis,
const mav<const uint32_t,1> &idx, mav<T,2> &ms, size_t nthreads) const const mav<const idx_t,1> &idx, mav<T,2> &ms, size_t nthreads) const
{ {
size_t nvis = vis.shape(0); size_t nvis = vis.shape(0);
checkShape(idx.shape(), {nvis}); checkShape(idx.shape(), {nvis});
...@@ -872,10 +878,10 @@ template<class T, class Serv> class SubServ ...@@ -872,10 +878,10 @@ template<class T, class Serv> class SubServ
{ {
private: private:
const Serv &srv; const Serv &srv;
const_mav<uint32_t,1> subidx; const_mav<idx_t,1> subidx;
public: public:
SubServ(const Serv &orig, const const_mav<uint32_t,1> &subidx_) SubServ(const Serv &orig, const const_mav<idx_t,1> &subidx_)
: srv(orig), subidx(subidx_){} : srv(orig), subidx(subidx_){}
size_t Nvis() const { return subidx.size(); } size_t Nvis() const { return subidx.size(); }
const Baselines &getBaselines() const { return srv.getBaselines(); } const Baselines &getBaselines() const { return srv.getBaselines(); }
...@@ -883,7 +889,7 @@ template<class T, class Serv> class SubServ ...@@ -883,7 +889,7 @@ template<class T, class Serv> class SubServ
{ return srv.getCoord(subidx[i]); } { return srv.getCoord(subidx[i]); }
complex<T> getVis(size_t i) const complex<T> getVis(size_t i) const
{ return srv.getVis(subidx[i]); } { return srv.getVis(subidx[i]); }
uint32_t getIdx(size_t i) const { return srv.getIdx(subidx[i]); } idx_t getIdx(size_t i) const { return srv.getIdx(subidx[i]); }
void setVis (size_t i, const complex<T> &v) const void setVis (size_t i, const complex<T> &v) const
{ srv.setVis(subidx[i], v); } { srv.setVis(subidx[i], v); }
void addVis (size_t i, const complex<T> &v) const void addVis (size_t i, const complex<T> &v) const
...@@ -894,7 +900,7 @@ template<class T, class T2> class VisServ ...@@ -894,7 +900,7 @@ template<class T, class T2> class VisServ
{ {
private: private:
const Baselines &baselines; const Baselines &baselines;
const_mav<uint32_t,1> idx; const_mav<idx_t,1> idx;
T2 vis; T2 vis;
const_mav<T,1> wgt; const_mav<T,1> wgt;
size_t nvis; size_t nvis;
...@@ -902,14 +908,14 @@ template<class T, class T2> class VisServ ...@@ -902,14 +908,14 @@ template<class T, class T2> class VisServ
public: public:
VisServ(const Baselines &baselines_, VisServ(const Baselines &baselines_,
const const_mav<uint32_t,1> &idx_, T2 vis_, const const_mav<T,1> &wgt_) const const_mav<idx_t,1> &idx_, T2 vis_, const const_mav<T,1> &wgt_)
: baselines(baselines_), idx(idx_), vis(vis_), wgt(wgt_), : baselines(baselines_), idx(idx_), vis(vis_), wgt(wgt_),
nvis(vis.shape(0)), have_wgt(wgt.size()!=0) nvis(vis.shape(0)), have_wgt(wgt.size()!=0)
{ {
checkShape(idx.shape(), {nvis}); checkShape(idx.shape(), {nvis});
if (have_wgt) checkShape(wgt.shape(), {nvis}); if (have_wgt) checkShape(wgt.shape(), {nvis});
} }
SubServ<T, VisServ> getSubserv(const const_mav<uint32_t,1> &subidx) const SubServ<T, VisServ> getSubserv(const const_mav<idx_t,1> &subidx) const
{ return SubServ<T, VisServ>(*this, subidx); } { return SubServ<T, VisServ>(*this, subidx); }
size_t Nvis() const { return nvis; } size_t Nvis() const { return nvis; }
const Baselines &getBaselines() const { return baselines; } const Baselines &getBaselines() const { return baselines; }
...@@ -917,7 +923,7 @@ template<class T, class T2> class VisServ ...@@ -917,7 +923,7 @@ template<class T, class T2> class VisServ
{ return baselines.effectiveCoord(idx[i]); } { return baselines.effectiveCoord(idx[i]); }
complex<T> getVis(size_t i) const complex<T> getVis(size_t i) const
{ return have_wgt ? vis(i)*wgt(i) : vis(i); } { return have_wgt ? vis(i)*wgt(i) : vis(i); }
uint32_t getIdx(size_t i) const { return idx[i]; } idx_t getIdx(size_t i) const { return idx[i]; }
void setVis (size_t i, const complex<T> &v) const void setVis (size_t i, const complex<T> &v) const
{ vis(i) = have_wgt ? v*wgt(i) : v; } { vis(i) = have_wgt ? v*wgt(i) : v; }
void addVis (size_t i, const complex<T> &v) const void addVis (size_t i, const complex<T> &v) const
...@@ -925,14 +931,14 @@ template<class T, class T2> class VisServ ...@@ -925,14 +931,14 @@ template<class T, class T2> class VisServ
}; };
template<class T, class T2> VisServ<T, T2> makeVisServ template<class T, class T2> VisServ<T, T2> makeVisServ
(const Baselines &baselines, (const Baselines &baselines,
const const_mav<uint32_t,1> &idx, const T2 &vis, const const_mav<T,1> &wgt) const const_mav<idx_t,1> &idx, const T2 &vis, const const_mav<T,1> &wgt)
{ return VisServ<T, T2>(baselines, idx, vis, wgt); } { return VisServ<T, T2>(baselines, idx, vis, wgt); }
template<class T, class T2> class MsServ template<class T, class T2> class MsServ
{ {
private: private:
const Baselines &baselines; const Baselines &baselines;
const_mav<uint32_t,1> idx; const_mav<idx_t,1> idx;
T2 ms; T2 ms;
const_mav<T,2> wgt; const_mav<T,2> wgt;
size_t nvis; size_t nvis;
...@@ -940,14 +946,14 @@ template<class T, class T2> class MsServ ...@@ -940,14 +946,14 @@ template<class T, class T2> class MsServ
public: public:
MsServ(const Baselines &baselines_, MsServ(const Baselines &baselines_,
const const_mav<uint32_t,1> &idx_, T2 ms_, const const_mav<T,2> &wgt_) const const_mav<idx_t,1> &idx_, T2 ms_, const const_mav<T,2> &wgt_)
: baselines(baselines_), idx(idx_), ms(ms_), wgt(wgt_), : baselines(baselines_), idx(idx_), ms(ms_), wgt(wgt_),
nvis(idx.shape(0)), have_wgt(wgt.size()!=0) nvis(idx.shape(0)), have_wgt(wgt.size()!=0)
{ {
checkShape(ms.shape(), {baselines.Nrows(), baselines.Nchannels()}); checkShape(ms.shape(), {baselines.Nrows(), baselines.Nchannels()});
if (have_wgt) checkShape(wgt.shape(), ms.shape()); if (have_wgt) checkShape(wgt.shape(), ms.shape());
} }
SubServ<T, MsServ> getSubserv(const const_mav<uint32_t,1> &subidx) const SubServ<T, MsServ> getSubserv(const const_mav<idx_t,1> &subidx) const
{ return SubServ<T, MsServ>(*this, subidx); } { return SubServ<T, MsServ>(*this, subidx); }
size_t Nvis() const { return nvis; } size_t Nvis() const { return nvis; }
const Baselines &getBaselines() const { return baselines; } const Baselines &getBaselines() const { return baselines; }
...@@ -959,7 +965,7 @@ template<class T, class T2> class MsServ ...@@ -959,7 +965,7 @@ template<class T, class T2> class MsServ
return have_wgt ? ms(rc.row, rc.chan)*wgt(rc.row, rc.chan) return have_wgt ? ms(rc.row, rc.chan)*wgt(rc.row, rc.chan)
: ms(rc.row, rc.chan); : ms(rc.row, rc.chan);
} }
uint32_t getIdx(size_t i) const { return idx[i]; } idx_t getIdx(size_t i) const { return idx[i]; }
void setVis (size_t i, const complex<T> &v) const void setVis (size_t i, const complex<T> &v) const
{ {
auto rc = baselines.getRowChan(idx(i)); auto rc = baselines.getRowChan(idx(i));
...@@ -973,7 +979,7 @@ template<class T, class T2> class MsServ ...@@ -973,7 +979,7 @@ template<class T, class T2> class MsServ
}; };
template<class T, class T2> MsServ<T, T2> makeMsServ template<class T, class T2> MsServ<T, T2> makeMsServ
(const Baselines &baselines, (const Baselines &baselines,
const const_mav<uint32_t,1> &idx, const T2 &ms, const const_mav<T,2> &wgt) const const_mav<idx_t,1> &idx, const T2 &ms, const const_mav<T,2> &wgt)
{ return MsServ<T, T2>(baselines, idx, ms, wgt); } { return MsServ<T, T2>(baselines, idx, ms, wgt); }
template<typename T, typename Serv> void x2grid_c template<typename T, typename Serv> void x2grid_c
...@@ -1027,13 +1033,13 @@ template<typename T, typename Serv> void x2grid_c ...@@ -1027,13 +1033,13 @@ template<typename T, typename Serv> void x2grid_c
template<typename T> void vis2grid_c template<typename T> void vis2grid_c
(const Baselines &baselines, const GridderConfig &gconf, (const Baselines &baselines, const GridderConfig &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,1> &vis, const const_mav<idx_t,1> &idx, const const_mav<complex<T>,1> &vis,
mav<complex<T>,2> &grid, const const_mav<T,1> &wgt, double w0=-1, double dw=-1) mav<complex<T>,2> &grid, const const_mav<T,1> &wgt, double w0=-1, double dw=-1)
{ x2grid_c(gconf, makeVisServ(baselines, idx, vis, wgt), grid, w0, dw); } { x2grid_c(gconf, makeVisServ(baselines, idx, vis, wgt), grid, w0, dw); }
template<typename T> void ms2grid_c template<typename T> void ms2grid_c
(const Baselines &baselines, const GridderConfig &gconf, (const Baselines &baselines, const GridderConfig &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &ms, const const_mav<idx_t,1> &idx, const const_mav<complex<T>,2> &ms,
mav<complex<T>,2> &grid, const const_mav<T,2> &wgt) mav<complex<T>,2> &grid, const const_mav<T,2> &wgt)
{ x2grid_c(gconf, makeMsServ(baselines, idx, ms, wgt), grid); } { x2grid_c(gconf, makeMsServ(baselines, idx, ms, wgt), grid); }
...@@ -1087,19 +1093,19 @@ template<typename T, typename Serv> void grid2x_c ...@@ -1087,19 +1093,19 @@ template<typename T, typename Serv> void grid2x_c
} }
template<typename T> void grid2vis_c template<typename T> void grid2vis_c
(const Baselines &baselines, const GridderConfig &gconf, (const Baselines &baselines, const GridderConfig &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &grid, const const_mav<idx_t,1> &idx, const const_mav<complex<T>,2> &grid,
mav<complex<T>,1> &vis, const const_mav<T,1> &wgt, double w0=-1, double dw=-1) mav<complex<T>,1> &vis, const const_mav<T,1> &wgt, double w0=-1, double dw=-1)
{ grid2x_c(gconf, grid, makeVisServ(baselines, idx, vis, wgt), w0, dw); } { grid2x_c(gconf, grid, makeVisServ(baselines, idx, vis, wgt), w0, dw); }
template<typename T> void grid2ms_c template<typename T> void grid2ms_c
(const Baselines &baselines, const GridderConfig &gconf, (const Baselines &baselines, const GridderConfig &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &grid, const const_mav<idx_t,1> &idx, const const_mav<complex<T>,2> &grid,
mav<complex<T>,2> &ms, const const_mav<T,2> &wgt) mav<complex<T>,2> &ms, const const_mav<T,2> &wgt)
{ grid2x_c(gconf, grid, makeMsServ(baselines, idx, ms, wgt)); } { grid2x_c(gconf, grid, makeMsServ(baselines, idx, ms, wgt)); }
template<typename T> void apply_holo template<typename T> void apply_holo
(const Baselines &baselines, const GridderConfig &gconf, (const Baselines &baselines, const GridderConfig &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &grid, const const_mav<idx_t,1> &idx, const const_mav<complex<T>,2> &grid,
mav<complex<T>,2> &ogrid, const const_mav<T,1> &wgt) mav<complex<T>,2> &ogrid, const const_mav<T,1> &wgt)
{ {
checkShape(grid.shape(), {gconf.Nu(), gconf.Nv()}); checkShape(grid.shape(), {gconf.Nu(), gconf.Nv()});
...@@ -1169,7 +1175,7 @@ template<typename T> void apply_holo ...@@ -1169,7 +1175,7 @@ template<typename T> void apply_holo
template<typename T> void get_correlations template<typename T> void get_correlations
(const Baselines &baselines, const GridderConfig &gconf, (const Baselines &baselines, const GridderConfig &gconf,
const const_mav<uint32_t,1> &idx, int du, int dv, const const_mav<idx_t,1> &idx, int du, int dv,
mav<T,2> &ogrid, const const_mav<T,1> &wgt) mav<T,2> &ogrid, const const_mav<T,1> &wgt)
{ {
size_t nvis = idx.shape(0); size_t nvis = idx.shape(0);
...@@ -1314,7 +1320,7 @@ template<typename T> void update_idx(vector<T> &v, vector<T> &vold, ...@@ -1314,7 +1320,7 @@ template<typename T> void update_idx(vector<T> &v, vector<T> &vold,
template<typename Serv> void wstack_common( template<typename Serv> void wstack_common(
const GridderConfig &gconf, const Serv &srv, double &wmin, const GridderConfig &gconf, const Serv &srv, double &wmin,
double &dw, size_t &nplanes, vector<size_t> &nvis_plane, double &dw, size_t &nplanes, vector<size_t> &nvis_plane,
vector<vector<uint32_t>> &minplane, size_t verbosity) vector<vector<idx_t>> &minplane, size_t verbosity)
{ {
size_t nvis = srv.Nvis(); size_t nvis = srv.Nvis();
size_t nthreads = gconf.Nthreads(); size_t nthreads = gconf.Nthreads();
...@@ -1369,7 +1375,7 @@ template<typename Serv> void wstack_common( ...@@ -1369,7 +1375,7 @@ template<typename Serv> void wstack_common(
for (size_t j=0; j<nplanes; ++j) for (size_t j=0; j<nplanes; ++j)
{ {
size_t l=0; size_t l=0;
for (size_t i=0; i<my_num_threads(); ++i) for (int i=0; i<my_num_threads(); ++i)
l+=tcnt[i*nplanes+j]; l+=tcnt[i*nplanes+j];
minplane[j].resize(l); minplane[j].resize(l);
} }
...@@ -1382,7 +1388,7 @@ template<typename Serv> void wstack_common( ...@@ -1382,7 +1388,7 @@ template<typename Serv> void wstack_common(
for (size_t ipart=0; ipart<nvis; ++ipart) for (size_t ipart=0; ipart<nvis; ++ipart)
{ {
int plane0 = max(0,int(1+(abs(srv.getCoord(ipart).w)-(0.5*supp*dw)-wmin)/dw)); int plane0 = max(0,int(1+(abs(srv.getCoord(ipart).w)-(0.5*supp*dw)-wmin)/dw));
minplane[plane0][myofs[plane0]++]=uint32_t(ipart); minplane[plane0][myofs[plane0]++]=idx_t(ipart);
} }
} }
} }
...@@ -1401,11 +1407,11 @@ template<typename T, typename Serv> void x2dirty( ...@@ -1401,11 +1407,11 @@ template<typename T, typename Serv> void x2dirty(
double dw; double dw;
size_t nplanes; size_t nplanes;
vector<size_t> nvis_plane; vector<size_t> nvis_plane;
vector<vector<uint32_t>> minplane; vector<vector<idx_t>> minplane;
if (verbosity>0) cout << "Gridding using improved w-stacking" << endl; if (verbosity>0) cout << "Gridding using improved w-stacking" << endl;
wstack_common(gconf, srv, wmin, dw, nplanes, nvis_plane, minplane, verbosity); wstack_common(gconf, srv, wmin, dw, nplanes, nvis_plane, minplane, verbosity);
dirty.fill(0); dirty.fill(0);
vector<uint32_t> subidx, subidx_old; vector<idx_t> subidx, subidx_old;
tmpStorage<complex<T>,2> grid_({gconf.Nu(),gconf.Nv()}); tmpStorage<complex<T>,2> grid_({gconf.Nu(),gconf.Nv()});
auto grid=grid_.getMav(); auto grid=grid_.getMav();
tmpStorage<complex<T>,2> tdirty_(dirty.shape()); tmpStorage<complex<T>,2> tdirty_(dirty.shape());
...@@ -1418,8 +1424,8 @@ template<typename T, typename Serv> void x2dirty( ...@@ -1418,8 +1424,8 @@ template<typename T, typename Serv> void x2dirty(
<< " visibilities" << endl; << " visibilities" << endl;
double wcur = wmin+iw*dw; double wcur = wmin+iw*dw;
update_idx(subidx, subidx_old, minplane[iw], iw>=supp ? minplane[iw-supp] : vector<uint32_t>()); update_idx(subidx, subidx_old, minplane[iw], iw>=supp ? minplane[iw-supp] : vector<idx_t>());
auto subidx2 = const_mav<uint32_t, 1>(subidx.data(), {subidx.size()}); auto subidx2 = const_mav<idx_t, 1>(subidx.data(), {subidx.size()});
auto subsrv = srv.getSubserv(subidx2); auto subsrv = srv.getSubserv(subidx2);
grid.fill(0); grid.fill(0);
x2grid_c(gconf, subsrv, grid, wcur, dw); x2grid_c(gconf, subsrv, grid, wcur, dw);
...@@ -1451,7 +1457,7 @@ template<typename T, typename Serv> void x2dirty( ...@@ -1451,7 +1457,7 @@ template<typename T, typename Serv> void x2dirty(
} }
} }
template<typename T> void vis2dirty(const Baselines &baselines, template<typename T> void vis2dirty(const Baselines &baselines,
const GridderConfig &gconf, const const_mav<uint32_t,1> &idx, const GridderConfig &gconf, const const_mav<idx_t,1> &idx,
const const_mav<complex<T>,1> &vis, const const_mav<T,1> &wgt, const const_mav<complex<T>,1> &vis, const const_mav<T,1> &wgt,
mav<T,2> &dirty, bool do_wstacking, size_t verbosity) mav<T,2> &dirty, bool do_wstacking, size_t verbosity)
{ {
...@@ -1473,7 +1479,7 @@ template<typename T, typename Serv> void dirty2x( ...@@ -1473,7 +1479,7 @@ template<typename T, typename Serv> void dirty2x(
double dw; double dw;
size_t nplanes; size_t nplanes;
vector<size_t> nvis_plane; vector<size_t> nvis_plane;
vector<vector<uint32_t>> minplane; vector<vector<idx_t>> minplane;
if (verbosity>0) cout << "Degridding using improved w-stacking" << endl; if (verbosity>0) cout << "Degridding using improved w-stacking" << endl;
wstack_common(gconf, srv, wmin, dw, nplanes, nvis_plane, minplane, verbosity); wstack_common(gconf, srv, wmin, dw, nplanes, nvis_plane, minplane, verbosity);
...@@ -1485,7 +1491,7 @@ template<typename T, typename Serv> void dirty2x( ...@@ -1485,7 +1491,7 @@ template<typename T, typename Serv> void dirty2x(
// correct for w gridding // correct for w gridding
apply_wcorr(gconf, tdirty, ES_Kernel(supp, nthreads), dw); apply_wcorr(gconf, tdirty, ES_Kernel(supp, nthreads), dw);
vector<uint32_t> subidx, subidx_old; vector<idx_t> subidx, subidx_old;
tmpStorage<complex<T>,2> grid_({gconf.Nu(),gconf.Nv()}); tmpStorage<complex<T>,2> grid_({gconf.Nu(),gconf.Nv()});
auto grid=grid_.getMav(); auto grid=grid_.getMav();
tmpStorage<complex<T>,2> tdirty2_({nx_dirty,ny_dirty}); tmpStorage<complex<T>,2> tdirty2_({nx_dirty,ny_dirty});
...@@ -1505,8 +1511,8 @@ template<typename T, typename Serv> void dirty2x( ...@@ -1505,8 +1511,8 @@ template<typename T, typename Serv> void dirty2x(
gconf.apply_wscreen(cmav(tdirty2), tdirty2, wcur, false); gconf.apply_wscreen(cmav(tdirty2), tdirty2, wcur, false);
gconf.dirty2grid_c(cmav(tdirty2), grid); gconf.dirty2grid_c(cmav(tdirty2), grid);
update_idx(subidx, subidx_old, minplane[iw], iw>=supp ? minplane[iw-supp] : vector<uint32_t>()); update_idx(subidx, subidx_old, minplane[iw], iw>=supp ? minplane[iw-supp] : vector<idx_t>());
auto subidx2 = const_mav<uint32_t, 1>(subidx.data(), {subidx.size()}); auto subidx2 = const_mav<idx_t, 1>(subidx.data(), {subidx.size()});
auto subsrv = srv.getSubserv(subidx2); auto subsrv = srv.getSubserv(subidx2);
grid2x_c(gconf, cmav(grid), subsrv, wcur, dw); grid2x_c(gconf, cmav(grid), subsrv, wcur, dw);
} }
...@@ -1528,7 +1534,7 @@ template<typename T, typename Serv> void dirty2x( ...@@ -1528,7 +1534,7 @@ template<typename T, typename Serv> void dirty2x(
} }
} }
template<typename T> void dirty2vis(const Baselines &baselines, template<typename T> void dirty2vis(const Baselines &baselines,
const GridderConfig &gconf, const const_mav<uint32_t,1> &idx, const GridderConfig &gconf, const const_mav<idx_t,1> &idx,
const const_mav<T,2> &dirty, const const_mav<T,1> &wgt, const const_mav<T,2> &dirty, const const_mav<T,1> &wgt,
mav<complex<T>,1> &vis, bool do_wstacking, size_t verbosity) mav<complex<T>,1> &vis, bool do_wstacking, size_t verbosity)
{ {
...@@ -1550,11 +1556,11 @@ size_t getIdxSize(const Baselines &baselines, ...@@ -1550,11 +1556,11 @@ size_t getIdxSize(const Baselines &baselines,
checkShape(flags.shape(), {nrow, nchan}); checkShape(flags.shape(), {nrow, nchan});
size_t res=0; size_t res=0;
#pragma omp parallel for num_threads(nthreads) reduction(+:res) #pragma omp parallel for num_threads(nthreads) reduction(+:res)
for (size_t irow=0; irow<nrow; ++irow) for (idx_t irow=0; irow<nrow; ++irow)
for (int ichan=chbegin; ichan<chend; ++ichan) for (int ichan=chbegin; ichan<chend; ++ichan)
if (!flags(irow,ichan)) if (!flags(irow,ichan))
{ {
auto w = abs(baselines.effectiveCoord(RowChan{irow,size_t(ichan)}).w); auto w = abs(baselines.effectiveCoord(RowChan{irow,idx_t(ichan)}).w);
if ((w>=wmin) && (w<wmax)) if ((w>=wmin) && (w<wmax))
++res; ++res;
} }
...@@ -1563,7 +1569,7 @@ size_t getIdxSize(const Baselines &baselines, ...@@ -1563,7 +1569,7 @@ size_t getIdxSize(const Baselines &baselines,
void fillIdx(const Baselines &baselines, void fillIdx(const Baselines &baselines,
const GridderConfig &gconf, const const_mav<bool,2> &flags, int chbegin, const GridderConfig &gconf, const const_mav<bool,2> &flags, int chbegin,
int chend, double wmin, double wmax, mav<uint32_t,1> &res) int chend, double wmin, double wmax, mav<idx_t,1> &res)
{ {
size_t nrow=baselines.Nrows(), size_t nrow=baselines.Nrows(),
nchan=baselines.Nchannels(), nchan=baselines.Nchannels(),
...@@ -1577,13 +1583,13 @@ void fillIdx(const Baselines &baselines, ...@@ -1577,13 +1583,13 @@ void fillIdx(const Baselines &baselines,
constexpr int side=1<<logsquare; constexpr int side=1<<logsquare;
size_t nbu = (gconf.Nu()+1+side-1) >> logsquare, size_t nbu = (gconf.Nu()+1+side-1) >> logsquare,
nbv = (gconf.Nv()+1+side-1) >> logsquare; nbv = (gconf.Nv()+1+side-1) >> logsquare;
vector<uint32_t> acc(nbu*nbv+1, 0); vector<idx_t> acc(nbu*nbv+1, 0);
vector<uint32_t> tmp(nrow*(chend-chbegin)); vector<idx_t> tmp(nrow*(chend-chbegin));
for (size_t irow=0, idx=0; irow<nrow; ++irow) for (idx_t irow=0, idx=0; irow<nrow; ++irow)
for (int ichan=chbegin; ichan<chend; ++ichan) for (int ichan=chbegin; ichan<chend; ++ichan)
if (!flags(irow, ichan)) if (!flags(irow, ichan))
{ {
auto uvw = baselines.effectiveCoord(RowChan{irow,size_t(ichan)}); auto uvw = baselines.effectiveCoord(RowChan{irow,idx_t(ichan)});
if (uvw.w<0) uvw.Flip(); if (uvw.w<0) uvw.Flip();
if ((uvw.w>=wmin) && (uvw.w<wmax)) if ((uvw.w>=wmin) && (uvw.w<wmax))
{ {
...@@ -1600,17 +1606,17 @@ void fillIdx(const Baselines &baselines, ...@@ -1600,17 +1606,17 @@ void fillIdx(const Baselines &baselines,
for (size_t i=1; i<acc.size(); ++i) for (size_t i=1; i<acc.size(); ++i)
acc[i] += acc[i-1]; acc[i] += acc[i-1];
myassert(res.shape(0)==acc.back(), "array size mismatch"); myassert(res.shape(0)==acc.back(), "array size mismatch");
for (size_t irow=0, idx=0; irow<nrow; ++irow) for (idx_t irow=0, idx=0; irow<nrow; ++irow)
for (int ichan=chbegin; ichan<chend; ++ichan) for (int ichan=chbegin; ichan<chend; ++ichan)
if (!flags(irow, ichan)) if (!flags(irow, ichan))
{ {
auto w = abs(baselines.effectiveCoord(RowChan{irow,size_t(ichan)}).w); auto w = abs(baselines.effectiveCoord(RowChan{irow,idx_t(ichan)}).w);
if ((w>=wmin) && (w<wmax)) if ((w>=wmin) && (w<wmax))
res[acc[tmp[idx++]]++] = baselines.getIdx(irow, ichan); res[acc[tmp[idx++]]++] = baselines.getIdx(irow, ichan);
} }
} }
template<typename T> vector<uint32_t> getWgtIndices(const Baselines &baselines, template<typename T> vector<idx_t> getWgtIndices(const Baselines &baselines,
const GridderConfig &gconf, const const_mav<T,2> &wgt) const GridderConfig &gconf, const const_mav<T,2> &wgt)
{ {