Skip to content

GitLab

N
nifty_gridder
Commit 5e415e5d authored by Martin Reinecke's avatar Martin Reinecke
Browse files
Options

make index type more flexible

parent 42892ca2
Show whitespace changes
Inline Side-by-side
Showing with 72 additions and 66 deletions
gridder_cxx.h
View file @ 5e415e5d
......@@ -296,11 +296,6 @@ void legendre_prep(int n, vector<double> &x, vector<double> &w, size_t nthreads)
// Start of real gridder functionality
//
struct RowChan
{
size_t row, chan;
};
template<typename T> void complex2hartley
(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,
return res;
}
using idx_t = uint32_t;
struct RowChan
{
idx_t row, chan;
};
struct UVW
{
double u, v, w;
......@@ -452,8 +454,8 @@ class Baselines
protected:
vector<UVW> coord;
vector<double> f_over_c;
size_t nrows, nchan;
size_t shift, mask;
idx_t nrows, nchan;
idx_t shift, mask;
public:
template<typename T> Baselines(const const_mav<T,2> &coord_,
......@@ -461,12 +463,16 @@ class Baselines
{
constexpr double speedOfLight = 299792458.;
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);
nchan = freq.shape(0);
shift=0;
while((size_t(1)<<shift)<nchan) ++shift;
mask=(size_t(1)<<shift)-1;
myassert(nrows*nchan<(size_t(1)<<32), "too many entries in MS");
while((idx_t(1)<<shift)<nchan) ++shift;
mask=(idx_t(1)<<shift)-1;
myassert(nrows*nchan<hugeval, "too many entries in MS");
f_over_c.resize(nchan);
for (size_t i=0; i<nchan; ++i)
{
......@@ -478,19 +484,19 @@ class Baselines
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}; }
UVW effectiveCoord(const RowChan &rc) const
{ 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)); }
size_t Nrows() const { return nrows; }
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); }
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
{
size_t nvis = idx.shape(0);
......@@ -505,7 +511,7 @@ class Baselines
}
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});
size_t nvis = idx.shape(0);
......@@ -519,7 +525,7 @@ class Baselines
}
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);
checkShape(idx.shape(), {nvis});
......@@ -872,10 +878,10 @@ template<class T, class Serv> class SubServ
{
private:
const Serv &srv;
const_mav<uint32_t,1> subidx;
const_mav<idx_t,1> subidx;
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_){}
size_t Nvis() const { return subidx.size(); }
const Baselines &getBaselines() const { return srv.getBaselines(); }
......@@ -883,7 +889,7 @@ template<class T, class Serv> class SubServ
{ return srv.getCoord(subidx[i]); }
complex<T> getVis(size_t i) const
{ 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
{ srv.setVis(subidx[i], v); }
void addVis (size_t i, const complex<T> &v) const
......@@ -894,7 +900,7 @@ template<class T, class T2> class VisServ
{
private:
const Baselines &baselines;
const_mav<uint32_t,1> idx;
const_mav<idx_t,1> idx;
T2 vis;
const_mav<T,1> wgt;
size_t nvis;
......@@ -902,14 +908,14 @@ template<class T, class T2> class VisServ
public:
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_),
nvis(vis.shape(0)), have_wgt(wgt.size()!=0)
{
checkShape(idx.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); }
size_t Nvis() const { return nvis; }
const Baselines &getBaselines() const { return baselines; }
......@@ -917,7 +923,7 @@ template<class T, class T2> class VisServ
{ return baselines.effectiveCoord(idx[i]); }
complex<T> getVis(size_t i) const
{ 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
{ vis(i) = have_wgt ? v*wgt(i) : v; }
void addVis (size_t i, const complex<T> &v) const
......@@ -925,14 +931,14 @@ template<class T, class T2> class VisServ
};
template<class T, class T2> VisServ<T, T2> makeVisServ
(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); }
template<class T, class T2> class MsServ
{
private:
const Baselines &baselines;
const_mav<uint32_t,1> idx;
const_mav<idx_t,1> idx;
T2 ms;
const_mav<T,2> wgt;
size_t nvis;
......@@ -940,14 +946,14 @@ template<class T, class T2> class MsServ
public:
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_),
nvis(idx.shape(0)), have_wgt(wgt.size()!=0)
{
checkShape(ms.shape(), {baselines.Nrows(), baselines.Nchannels()});
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); }
size_t Nvis() const { return nvis; }
const Baselines &getBaselines() const { return baselines; }
......@@ -959,7 +965,7 @@ template<class T, class T2> class MsServ
return have_wgt ? ms(rc.row, rc.chan)*wgt(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
{
auto rc = baselines.getRowChan(idx(i));
......@@ -973,7 +979,7 @@ template<class T, class T2> class MsServ
};
template<class T, class T2> MsServ<T, T2> makeMsServ
(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); }
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
(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)
{ x2grid_c(gconf, makeVisServ(baselines, idx, vis, wgt), grid, w0, dw); }
template<typename T> void ms2grid_c
(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)
{ x2grid_c(gconf, makeMsServ(baselines, idx, ms, wgt), grid); }
......@@ -1087,19 +1093,19 @@ template<typename T, typename Serv> void grid2x_c
}
template<typename T> void grid2vis_c
(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)
{ grid2x_c(gconf, grid, makeVisServ(baselines, idx, vis, wgt), w0, dw); }
template<typename T> void grid2ms_c
(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)
{ grid2x_c(gconf, grid, makeMsServ(baselines, idx, ms, wgt)); }
template<typename T> void apply_holo
(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)
{
checkShape(grid.shape(), {gconf.Nu(), gconf.Nv()});
......@@ -1169,7 +1175,7 @@ template<typename T> void apply_holo
template<typename T> void get_correlations
(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)
{
size_t nvis = idx.shape(0);
......@@ -1314,7 +1320,7 @@ template<typename T> void update_idx(vector<T> &v, vector<T> &vold,
template<typename Serv> void wstack_common(
const GridderConfig &gconf, const Serv &srv, double &wmin,
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 nthreads = gconf.Nthreads();
......@@ -1369,7 +1375,7 @@ template<typename Serv> void wstack_common(
for (size_t j=0; j<nplanes; ++j)
{
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];
minplane[j].resize(l);
}
......@@ -1382,7 +1388,7 @@ template<typename Serv> void wstack_common(
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));
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(
double dw;
size_t nplanes;
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;
wstack_common(gconf, srv, wmin, dw, nplanes, nvis_plane, minplane, verbosity);
dirty.fill(0);
vector<uint32_t> subidx, subidx_old;
vector<idx_t> subidx, subidx_old;
tmpStorage<complex<T>,2> grid_({gconf.Nu(),gconf.Nv()});
auto grid=grid_.getMav();
tmpStorage<complex<T>,2> tdirty_(dirty.shape());
......@@ -1418,8 +1424,8 @@ template<typename T, typename Serv> void x2dirty(
<< " visibilities" << endl;
double wcur = wmin+iw*dw;
update_idx(subidx, subidx_old, minplane[iw], iw>=supp ? minplane[iw-supp] : vector<uint32_t>());
auto subidx2 = const_mav<uint32_t, 1>(subidx.data(), {subidx.size()});
update_idx(subidx, subidx_old, minplane[iw], iw>=supp ? minplane[iw-supp] : vector<idx_t>());
auto subidx2 = const_mav<idx_t, 1>(subidx.data(), {subidx.size()});
auto subsrv = srv.getSubserv(subidx2);
grid.fill(0);
x2grid_c(gconf, subsrv, grid, wcur, dw);
......@@ -1451,7 +1457,7 @@ 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 GridderConfig &gconf, const const_mav<idx_t,1> &idx,
const const_mav<complex<T>,1> &vis, const const_mav<T,1> &wgt,
mav<T,2> &dirty, bool do_wstacking, size_t verbosity)
{
......@@ -1473,7 +1479,7 @@ template<typename T, typename Serv> void dirty2x(
double dw;
size_t nplanes;
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;
wstack_common(gconf, srv, wmin, dw, nplanes, nvis_plane, minplane, verbosity);
......@@ -1485,7 +1491,7 @@ template<typename T, typename Serv> void dirty2x(
// correct for w gridding
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()});
auto grid=grid_.getMav();
tmpStorage<complex<T>,2> tdirty2_({nx_dirty,ny_dirty});
......@@ -1505,8 +1511,8 @@ template<typename T, typename Serv> void dirty2x(
gconf.apply_wscreen(cmav(tdirty2), tdirty2, wcur, false);
gconf.dirty2grid_c(cmav(tdirty2), grid);
update_idx(subidx, subidx_old, minplane[iw], iw>=supp ? minplane[iw-supp] : vector<uint32_t>());
auto subidx2 = const_mav<uint32_t, 1>(subidx.data(), {subidx.size()});
update_idx(subidx, subidx_old, minplane[iw], iw>=supp ? minplane[iw-supp] : vector<idx_t>());
auto subidx2 = const_mav<idx_t, 1>(subidx.data(), {subidx.size()});
auto subsrv = srv.getSubserv(subidx2);
grid2x_c(gconf, cmav(grid), subsrv, wcur, dw);
}
......@@ -1528,7 +1534,7 @@ 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 GridderConfig &gconf, const const_mav<idx_t,1> &idx,
const const_mav<T,2> &dirty, const const_mav<T,1> &wgt,
mav<complex<T>,1> &vis, bool do_wstacking, size_t verbosity)
{
......@@ -1550,11 +1556,11 @@ size_t getIdxSize(const Baselines &baselines,
checkShape(flags.shape(), {nrow, nchan});
size_t res=0;
#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)
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))
++res;
}
......@@ -1563,7 +1569,7 @@ size_t getIdxSize(const Baselines &baselines,
void fillIdx(const Baselines &baselines,
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(),
nchan=baselines.Nchannels(),
......@@ -1577,13 +1583,13 @@ void fillIdx(const Baselines &baselines,
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*(chend-chbegin));
for (size_t irow=0, idx=0; irow<nrow; ++irow)
vector<idx_t> acc(nbu*nbv+1, 0);
vector<idx_t> tmp(nrow*(chend-chbegin));
for (idx_t irow=0, idx=0; irow<nrow; ++irow)
for (int ichan=chbegin; ichan<chend; ++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>=wmin) && (uvw.w<wmax))
{
......@@ -1600,17 +1606,17 @@ void fillIdx(const Baselines &baselines,
for (size_t i=1; i<acc.size(); ++i)
acc[i] += acc[i-1];
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)
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))
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)
{
size_t nrow=baselines.Nrows(),
......@@ -1621,14 +1627,14 @@ template<typename T> vector<uint32_t> getWgtIndices(const Baselines &baselines,
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);
vector<idx_t> acc(nbu*nbv+1, 0);
vector<idx_t> tmp(nrow*nchan);
for (size_t irow=0, idx=0; irow<nrow; ++irow)
for (size_t ichan=0; ichan<nchan; ++ichan)
for (idx_t irow=0, idx=0; irow<nrow; ++irow)
for (idx_t ichan=0; ichan<nchan; ++ichan)
if ((!have_wgt) || (wgt(irow,ichan)!=0))
{
auto uvw = baselines.effectiveCoord(RowChan{irow,size_t(ichan)});
auto uvw = baselines.effectiveCoord(RowChan{irow,idx_t(ichan)});
if (uvw.w<0) uvw.Flip();
double u, v;
int iu0, iv0;
......@@ -1642,7 +1648,7 @@ template<typename T> vector<uint32_t> getWgtIndices(const Baselines &baselines,
for (size_t i=1; i<acc.size(); ++i)
acc[i] += acc[i-1];
vector<uint32_t> res(acc.back());
vector<idx_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))
......@@ -1658,7 +1664,7 @@ template<typename T> void ms2dirty(const const_mav<double,2> &uvw,
Baselines baselines(uvw, freq);
GridderConfig gconf(dirty.shape(0), dirty.shape(1), epsilon, pixsize_x, pixsize_y, nthreads);
auto idx = getWgtIndices(baselines, gconf, wgt);
auto idx2 = const_mav<uint32_t,1>(idx.data(),{idx.size()});
auto idx2 = const_mav<idx_t,1>(idx.data(),{idx.size()});
x2dirty(gconf, makeMsServ(baselines,idx2,ms,wgt), dirty, do_wstacking, verbosity);
}
......@@ -1670,7 +1676,7 @@ template<typename T> void dirty2ms(const const_mav<double,2> &uvw,
Baselines baselines(uvw, freq);
GridderConfig gconf(dirty.shape(0), dirty.shape(1), epsilon, pixsize_x, pixsize_y, nthreads);
auto idx = getWgtIndices(baselines, gconf, wgt);
auto idx2 = const_mav<uint32_t,1>(idx.data(),{idx.size()});
auto idx2 = const_mav<idx_t,1>(idx.data(),{idx.size()});
ms.fill(0);
dirty2x(gconf, dirty, makeMsServ(baselines,idx2,ms,wgt), do_wstacking, verbosity);
}
......
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