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

cleanup and tweaks

parent fb5509a1
......@@ -665,51 +665,6 @@ template<typename T> class ReadHelper: public Helper<T>
}
};
template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<complex<T>> &vis_)
{
myassert(idx_.ndim()==1, "idx array must be 1D");
myassert(vis_.ndim()==1, "vis must be 1D");
auto vis=vis_.data();
myassert(vis_.shape(0)==idx_.shape(0), "bad vis dimension");
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.data();
size_t nu=gconf.Nu(), nv=gconf.Nv();
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 beta = gconf.Beta();
size_t w = gconf.W();
#pragma omp parallel
{
WriteHelper<T> hlp(gconf, grid);
T emb = exp(-2*beta);
const T * RESTRICT ku = hlp.kernel.data();
const T * RESTRICT kv = hlp.kernel.data()+w;
// Loop over sampling points
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
hlp.prep_write(coord.u, coord.v);
auto * RESTRICT ptr = hlp.p0;
auto v(vis[ipart]*emb);
for (size_t cu=0; cu<w; ++cu)
{
complex<T> tmp(v*ku[cu]);
for (size_t cv=0; cv<w; ++cv)
ptr[cv] += tmp*kv[cv];
ptr+=hlp.sv;
}
}
} // end of parallel region
return complex2hartley(res);
}
template<typename T> pyarr_c<complex<T>> vis2grid_c(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<complex<T>> &vis_)
......@@ -755,53 +710,10 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c(const Baselines<T> &baseline
return res;
}
template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<T> &grid0_)
{
size_t nu=gconf.Nu(), nv=gconf.Nv();
myassert(idx_.ndim()==1, "idx array must be 1D");
auto grid_ = hartley2complex(grid0_);
auto grid = grid_.data();
myassert(grid_.ndim()==2, "data must be 2D");
myassert(grid_.shape(0)==int(nu), "bad grid dimension");
myassert(grid_.shape(1)==int(nv), "bad grid dimension");
size_t nvis = size_t(idx_.shape(0));
auto idx = idx_.data();
auto res = makearray<complex<T>>({nvis});
auto vis = res.mutable_data();
T beta = gconf.Beta();
size_t w = gconf.W();
// Loop over sampling points
#pragma omp parallel
{
ReadHelper<T> hlp(gconf, grid);
T emb = exp(-2*beta);
const T * RESTRICT ku = hlp.kernel.data();
const T * RESTRICT kv = hlp.kernel.data()+w;
#pragma omp for schedule(guided,100)
for (size_t ipart=0; ipart<nvis; ++ipart)
{
UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
hlp.prep_read(coord.u, coord.v);
complex<T> r = 0;
auto * RESTRICT ptr = hlp.p0;
for (size_t cu=0; cu<w; ++cu)
{
complex<T> tmp(0);
for (size_t cv=0; cv<w; ++cv)
tmp += ptr[cv] * kv[cv];
r += tmp*ku[cu];
ptr += hlp.sv;
}
vis[ipart] = r*emb;
}
}
return res;
}
const pyarr_c<complex<T>> &vis_)
{ return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_)); }
template<typename T> pyarr_c<complex<T>> grid2vis_c(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
......@@ -850,6 +762,11 @@ template<typename T> pyarr_c<complex<T>> grid2vis_c(const Baselines<T> &baseline
return res;
}
template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
const pyarr_c<T> &grid_)
{ return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_)); }
template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const pyarr_c<bool> &flags_, int chbegin,
int chend, T wmin, T wmax)
......@@ -866,10 +783,11 @@ template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &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> bincnt(nbu*nbv, 0);
for (size_t irow=0; irow<nrow; ++irow)
for (int ichan=chbegin; ichan<chend; ++ichan)
if (!flags[irow*nchan + ichan])
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)
for (int ichan=chbegin; ichan<chend; ++ichan, ++idx)
if (!flags[irow*nchan+ichan])
{
auto uvw = baselines.effectiveCoord(irow, ichan);
if ((uvw.w>=wmin) && (uvw.w<wmax))
......@@ -879,31 +797,19 @@ template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
gconf.getpix(uvw.u, uvw.v, u, v, iu0, iv0);
iu0 = (iu0+nsafe)>>logsquare;
iv0 = (iv0+nsafe)>>logsquare;
++bincnt[nbv*iu0 + iv0];
++acc[nbv*iu0 + iv0 + 1];
tmp[idx] = nbv*iu0 + iv0;
}
}
vector<uint32_t> acc(bincnt.size()+1);
acc[0] = 0;
for (size_t i=0; i<bincnt.size(); ++i)
acc[i+1] = acc[i] + bincnt[i];
for (size_t i=1; i<acc.size(); ++i)
acc[i] += acc[i-1];
auto res = makearray<uint32_t>({acc.back()});
auto iout = res.mutable_data();
for (size_t irow=0; irow<nrow; ++irow)
for (int ichan=chbegin; ichan<chend; ++ichan)
if (!flags[irow*nchan + ichan])
{
auto uvw = baselines.effectiveCoord(irow, ichan);
if ((uvw.w>=wmin) && (uvw.w<wmax))
{
T u, v;
int iu0, iv0;
gconf.getpix(uvw.u, uvw.v, u, v, iu0, iv0);
iu0 = (iu0+nsafe)>>logsquare;
iv0 = (iv0+nsafe)>>logsquare;
iout[acc[nbv*iu0 + iv0]] = irow*nchan + ichan;
++acc[nbv*iu0 + iv0];
}
}
for (size_t irow=0, idx=0; irow<nrow; ++irow)
for (int ichan=chbegin; ichan<chend; ++ichan, ++idx)
if (!flags[irow*nchan+ichan])
iout[acc[tmp[idx]]++] = irow*nchan+ichan;
return res;
}
......
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