Commit 8dbc0f5f authored by Martin Reinecke's avatar Martin Reinecke
Browse files

polishing

parent 6575e568
...@@ -418,7 +418,7 @@ class EC_Kernel_with_correction: public EC_Kernel ...@@ -418,7 +418,7 @@ class EC_Kernel_with_correction: public EC_Kernel
}; };
/* Compute correction factors for the ES gridding kernel /* Compute correction factors for the ES gridding kernel
This implementation follows eqs. (3.8) to (3.10) of Barnett et al. 2018 */ This implementation follows eqs. (3.8) to (3.10) of Barnett et al. 2018 */
vector<double> correction_factors (size_t n, size_t nval, size_t supp, vector<double> correction_factors(size_t n, size_t nval, size_t supp,
size_t nthreads) size_t nthreads)
{ {
EC_Kernel_with_correction kernel(supp, nthreads); EC_Kernel_with_correction kernel(supp, nthreads);
...@@ -724,7 +724,6 @@ template<typename T> class GridderConfig ...@@ -724,7 +724,6 @@ template<typename T> class GridderConfig
} }
} }
} }
}; };
constexpr int logsquare=4; constexpr int logsquare=4;
...@@ -971,28 +970,23 @@ template<typename T> void vis2grid_c ...@@ -971,28 +970,23 @@ template<typename T> void vis2grid_c
(const Baselines<T> &baselines, const GridderConfig<T> &gconf, (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,1> &vis, const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,1> &vis,
mav<complex<T>,2> &grid, const const_mav<T,1> &wgt, T w0=-1, T dw=-1) mav<complex<T>,2> &grid, const const_mav<T,1> &wgt, T w0=-1, T 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<T> &baselines, const GridderConfig<T> &gconf, (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &ms, const const_mav<uint32_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);
}
template<typename T, typename Serv> void grid2x_c template<typename T, typename Serv> void grid2x_c
(const GridderConfig<T> &gconf, const const_mav<complex<T>,2> &grid, const Serv &srv, (const GridderConfig<T> &gconf, const const_mav<complex<T>,2> &grid, const Serv &srv,
T w0=-1, T dw=-1) T w0=-1, T dw=-1)
{ {
size_t nu=gconf.Nu(), nv=gconf.Nv(); checkShape(grid.shape(), {gconf.Nu(), gconf.Nv()});
checkShape(grid.shape(), {nu, nv});
myassert(grid.contiguous(), "grid is not contiguous"); myassert(grid.contiguous(), "grid is not contiguous");
size_t supp = gconf.Supp(); size_t supp = gconf.Supp();
size_t nthreads = gconf.Nthreads(); size_t nthreads = gconf.Nthreads();
bool do_w_gridding=dw>0; bool do_w_gridding = dw>0;
// Loop over sampling points // Loop over sampling points
#pragma omp parallel num_threads(nthreads) #pragma omp parallel num_threads(nthreads)
...@@ -1035,17 +1029,13 @@ template<typename T> void grid2vis_c ...@@ -1035,17 +1029,13 @@ template<typename T> void grid2vis_c
(const Baselines<T> &baselines, const GridderConfig<T> &gconf, (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &grid, const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &grid,
mav<complex<T>,1> &vis, const const_mav<T,1> &wgt, T w0=-1, T dw=-1) mav<complex<T>,1> &vis, const const_mav<T,1> &wgt, T w0=-1, T 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<T> &baselines, const GridderConfig<T> &gconf, (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &grid, const const_mav<uint32_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<T> &baselines, const GridderConfig<T> &gconf, (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
...@@ -1057,7 +1047,7 @@ template<typename T> void apply_holo ...@@ -1057,7 +1047,7 @@ template<typename T> void apply_holo
size_t nthreads = gconf.Nthreads(); size_t nthreads = gconf.Nthreads();
bool have_wgt = wgt.size()!=0; bool have_wgt = wgt.size()!=0;
if (have_wgt) checkShape(wgt.shape(), {nvis}); if (have_wgt) checkShape(wgt.shape(), {nvis});
checkShape(ogrid.shape(), ogrid.shape()); checkShape(ogrid.shape(), grid.shape());
ogrid.fill(0); ogrid.fill(0);
size_t supp = gconf.Supp(); size_t supp = gconf.Supp();
...@@ -1308,12 +1298,9 @@ template<typename T, typename Serv> void x2dirty_wstack( ...@@ -1308,12 +1298,9 @@ template<typename T, typename Serv> void x2dirty_wstack(
{ {
auto nx_dirty=gconf.Nxdirty(); auto nx_dirty=gconf.Nxdirty();
auto ny_dirty=gconf.Nydirty(); auto ny_dirty=gconf.Nydirty();
auto nu=gconf.Nu();
auto nv=gconf.Nv();
size_t nvis = srv.Nvis(); size_t nvis = srv.Nvis();
auto w_supp = gconf.Supp(); auto supp = gconf.Supp();
size_t nthreads = gconf.Nthreads(); size_t nthreads = gconf.Nthreads();
EC_Kernel_with_correction kernel(w_supp, nthreads);
T wmin; T wmin;
double dw; double dw;
size_t nplanes; size_t nplanes;
...@@ -1324,11 +1311,10 @@ template<typename T, typename Serv> void x2dirty_wstack( ...@@ -1324,11 +1311,10 @@ template<typename T, typename Serv> void x2dirty_wstack(
dirty.fill(0); dirty.fill(0);
vector<complex<T>> vis_loc_; vector<complex<T>> vis_loc_;
vector<uint32_t> idx_loc_; vector<uint32_t> idx_loc_, mapper;
vector<uint32_t> mapper; tmpStorage<complex<T>,2> grid_({gconf.Nu(),gconf.Nv()});
tmpStorage<complex<T>,2> grid_({nu,nv});
auto grid=grid_.getMav(); auto grid=grid_.getMav();
tmpStorage<complex<T>,2> tdirty_({nx_dirty,ny_dirty}); tmpStorage<complex<T>,2> tdirty_(dirty.shape());
auto tdirty=tdirty_.getMav(); auto tdirty=tdirty_.getMav();
for (size_t iw=0; iw<nplanes; ++iw) for (size_t iw=0; iw<nplanes; ++iw)
{ {
...@@ -1344,7 +1330,7 @@ template<typename T, typename Serv> void x2dirty_wstack( ...@@ -1344,7 +1330,7 @@ template<typename T, typename Serv> void x2dirty_wstack(
auto idx_loc = mav<uint32_t,1>(idx_loc_.data(), {nvis_plane[iw]}); auto idx_loc = mav<uint32_t,1>(idx_loc_.data(), {nvis_plane[iw]});
mapper.resize(nvis_plane[iw]); mapper.resize(nvis_plane[iw]);
for (size_t ipart=0; ipart<nvis; ++ipart) for (size_t ipart=0; ipart<nvis; ++ipart)
if ((int(iw)>=minplane[ipart]) && (iw<minplane[ipart]+w_supp)) if ((int(iw)>=minplane[ipart]) && (iw<minplane[ipart]+supp))
mapper[cnt++] = ipart; mapper[cnt++] = ipart;
#pragma omp parallel for num_threads(nthreads) #pragma omp parallel for num_threads(nthreads)
...@@ -1364,7 +1350,7 @@ template<typename T, typename Serv> void x2dirty_wstack( ...@@ -1364,7 +1350,7 @@ template<typename T, typename Serv> void x2dirty_wstack(
dirty(i,j) += tdirty(i,j).real(); dirty(i,j) += tdirty(i,j).real();
} }
// correct for w gridding // correct for w gridding
apply_wcorr(gconf, dirty, kernel, dw); apply_wcorr(gconf, dirty, EC_Kernel_with_correction(supp, nthreads), dw);
} }
template<typename T> void vis2dirty_wstack(const Baselines<T> &baselines, template<typename T> void vis2dirty_wstack(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const const_mav<uint32_t,1> &idx, const GridderConfig<T> &gconf, const const_mav<uint32_t,1> &idx,
...@@ -1377,12 +1363,9 @@ template<typename T, typename Serv> void x2dirty_wstack2( ...@@ -1377,12 +1363,9 @@ template<typename T, typename Serv> void x2dirty_wstack2(
{ {
auto nx_dirty=gconf.Nxdirty(); auto nx_dirty=gconf.Nxdirty();
auto ny_dirty=gconf.Nydirty(); auto ny_dirty=gconf.Nydirty();
auto nu=gconf.Nu();
auto nv=gconf.Nv();
size_t nvis = srv.Nvis(); size_t nvis = srv.Nvis();
auto w_supp = gconf.Supp(); auto supp = gconf.Supp();
size_t nthreads = gconf.Nthreads(); size_t nthreads = gconf.Nthreads();
EC_Kernel_with_correction kernel(w_supp, nthreads);
T wmin; T wmin;
double dw; double dw;
size_t nplanes; size_t nplanes;
...@@ -1393,9 +1376,9 @@ template<typename T, typename Serv> void x2dirty_wstack2( ...@@ -1393,9 +1376,9 @@ template<typename T, typename Serv> void x2dirty_wstack2(
dirty.fill(0); dirty.fill(0);
vector<uint32_t> newidx; vector<uint32_t> newidx;
tmpStorage<complex<T>,2> grid_({nu,nv}); tmpStorage<complex<T>,2> grid_({gconf.Nu(),gconf.Nv()});
auto grid=grid_.getMav(); auto grid=grid_.getMav();
tmpStorage<complex<T>,2> tdirty_({nx_dirty,ny_dirty}); tmpStorage<complex<T>,2> tdirty_(dirty.shape());
auto tdirty=tdirty_.getMav(); auto tdirty=tdirty_.getMav();
for (size_t iw=0; iw<nplanes; ++iw) for (size_t iw=0; iw<nplanes; ++iw)
{ {
...@@ -1407,7 +1390,7 @@ template<typename T, typename Serv> void x2dirty_wstack2( ...@@ -1407,7 +1390,7 @@ template<typename T, typename Serv> void x2dirty_wstack2(
size_t cnt=0; size_t cnt=0;
newidx.resize(nvis_plane[iw]); newidx.resize(nvis_plane[iw]);
for (size_t ipart=0; ipart<nvis; ++ipart) for (size_t ipart=0; ipart<nvis; ++ipart)
if ((int(iw)>=minplane[ipart]) && (iw<minplane[ipart]+w_supp)) if ((int(iw)>=minplane[ipart]) && (iw<minplane[ipart]+supp))
newidx[cnt++] = srv.getIdx(ipart); newidx[cnt++] = srv.getIdx(ipart);
grid.fill(0); grid.fill(0);
...@@ -1421,7 +1404,7 @@ template<typename T, typename Serv> void x2dirty_wstack2( ...@@ -1421,7 +1404,7 @@ template<typename T, typename Serv> void x2dirty_wstack2(
dirty(i,j) += tdirty(i,j).real(); dirty(i,j) += tdirty(i,j).real();
} }
// correct for w gridding // correct for w gridding
apply_wcorr(gconf, dirty, kernel, dw); apply_wcorr(gconf, dirty, EC_Kernel_with_correction(supp, nthreads), dw);
} }
template<typename T> void vis2dirty_wstack2(const Baselines<T> &baselines, template<typename T> void vis2dirty_wstack2(const Baselines<T> &baselines,
const GridderConfig<T> &gconf, const const_mav<uint32_t,1> &idx, const GridderConfig<T> &gconf, const const_mav<uint32_t,1> &idx,
...@@ -1436,12 +1419,9 @@ template<typename T, typename Serv> void dirty2x_wstack( ...@@ -1436,12 +1419,9 @@ template<typename T, typename Serv> void dirty2x_wstack(
{ {
auto nx_dirty=gconf.Nxdirty(); auto nx_dirty=gconf.Nxdirty();
auto ny_dirty=gconf.Nydirty(); auto ny_dirty=gconf.Nydirty();
auto nu=gconf.Nu();
auto nv=gconf.Nv();
size_t nvis = srv.Nvis(); size_t nvis = srv.Nvis();
auto w_supp = gconf.Supp(); auto supp = gconf.Supp();
size_t nthreads = gconf.Nthreads(); size_t nthreads = gconf.Nthreads();
EC_Kernel_with_correction kernel(w_supp, nthreads);
T wmin; T wmin;
double dw; double dw;
size_t nplanes; size_t nplanes;
...@@ -1456,12 +1436,11 @@ template<typename T, typename Serv> void dirty2x_wstack( ...@@ -1456,12 +1436,11 @@ template<typename T, typename Serv> void dirty2x_wstack(
for (size_t j=0; j<ny_dirty; ++j) for (size_t j=0; j<ny_dirty; ++j)
tdirty(i,j) = dirty(i,j); tdirty(i,j) = dirty(i,j);
// correct for w gridding // correct for w gridding
apply_wcorr(gconf, tdirty, kernel, dw); apply_wcorr(gconf, tdirty, EC_Kernel_with_correction(supp, nthreads), dw);
vector<complex<T>> vis_loc_; vector<complex<T>> vis_loc_;
vector<uint32_t> idx_loc_; vector<uint32_t> idx_loc_, mapper;
vector<uint32_t> mapper; tmpStorage<complex<T>,2> grid_({gconf.Nu(),gconf.Nv()});
tmpStorage<complex<T>,2> grid_({nu,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});
auto tdirty2=tdirty2_.getMav(); auto tdirty2=tdirty2_.getMav();
...@@ -1479,7 +1458,7 @@ template<typename T, typename Serv> void dirty2x_wstack( ...@@ -1479,7 +1458,7 @@ template<typename T, typename Serv> void dirty2x_wstack(
auto idx_loc = mav<uint32_t,1>(idx_loc_.data(), {nvis_plane[iw]}); auto idx_loc = mav<uint32_t,1>(idx_loc_.data(), {nvis_plane[iw]});
mapper.resize(nvis_plane[iw]); mapper.resize(nvis_plane[iw]);
for (size_t ipart=0; ipart<nvis; ++ipart) for (size_t ipart=0; ipart<nvis; ++ipart)
if ((int(iw)>=minplane[ipart]) && (iw<minplane[ipart]+w_supp)) if ((int(iw)>=minplane[ipart]) && (iw<minplane[ipart]+supp))
mapper[cnt++] = ipart; mapper[cnt++] = ipart;
#pragma omp parallel for num_threads(nthreads) #pragma omp parallel for num_threads(nthreads)
for (size_t i=0; i<nvis_plane[iw]; ++i) for (size_t i=0; i<nvis_plane[iw]; ++i)
......
...@@ -478,7 +478,10 @@ template<typename T> pyarr<T> Pyvis2grid(const PyBaselines<T> &baselines, ...@@ -478,7 +478,10 @@ template<typename T> pyarr<T> Pyvis2grid(const PyBaselines<T> &baselines,
{ {
auto tmp=Pyvis2grid_c(baselines, gconf, idx_, vis_, None, wgt_); auto tmp=Pyvis2grid_c(baselines, gconf, idx_, vis_, None, wgt_);
auto grd=provideCArray<T>(grid_in,{gconf.Nu(), gconf.Nv()}); auto grd=provideCArray<T>(grid_in,{gconf.Nu(), gconf.Nv()});
{
py::gil_scoped_release release;
gridder::detail::complex2hartley(make_const_mav<2>(tmp), make_mav<2>(grd), gconf.Nthreads()); gridder::detail::complex2hartley(make_const_mav<2>(tmp), make_mav<2>(grd), gconf.Nthreads());
}
return grd; return grd;
} }
...@@ -534,7 +537,10 @@ template<typename T> pyarr<T> Pyms2grid( ...@@ -534,7 +537,10 @@ template<typename T> pyarr<T> Pyms2grid(
auto tmp = Pyms2grid_c(baselines, gconf, idx_, ms_, None, wgt_); auto tmp = Pyms2grid_c(baselines, gconf, idx_, ms_, None, wgt_);
auto res_ = provideCArray<T>(grid_in, {gconf.Nu(), gconf.Nv()}); auto res_ = provideCArray<T>(grid_in, {gconf.Nu(), gconf.Nv()});
auto res = make_mav<2>(res_); auto res = make_mav<2>(res_);
{
py::gil_scoped_release release;
gridder::detail::complex2hartley(make_const_mav<2>(tmp), res, gconf.Nthreads()); gridder::detail::complex2hartley(make_const_mav<2>(tmp), res, gconf.Nthreads());
}
return res_; return res_;
} }
...@@ -746,7 +752,10 @@ template<typename T> pyarr<T> Pygridding(const pyarr<T> &uvw, ...@@ -746,7 +752,10 @@ template<typename T> pyarr<T> Pygridding(const pyarr<T> &uvw,
auto wgt2 = make_const_mav<2>(wgt); auto wgt2 = make_const_mav<2>(wgt);
auto dirty=makeArray<T>({npix_x,npix_y}); auto dirty=makeArray<T>({npix_x,npix_y});
auto dirty2 =make_mav<2>(dirty); auto dirty2 =make_mav<2>(dirty);
{
py::gil_scoped_release release;
gridding(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,dirty2,verbosity); gridding(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,dirty2,verbosity);
}
return dirty; return dirty;
} }
...@@ -761,7 +770,10 @@ template<typename T> pyarr<complex<T>> Pydegridding(const pyarr<T> &uvw, ...@@ -761,7 +770,10 @@ template<typename T> pyarr<complex<T>> Pydegridding(const pyarr<T> &uvw,
auto wgt2 = make_const_mav<2>(wgt); auto wgt2 = make_const_mav<2>(wgt);
auto ms=makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)}); auto ms=makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)});
auto ms2 =make_mav<2>(ms); auto ms2 =make_mav<2>(ms);
{
py::gil_scoped_release release;
degridding(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,ms2,verbosity); degridding(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,ms2,verbosity);
}
return ms; return ms;
} }
...@@ -777,7 +789,10 @@ template<typename T> pyarr<T> Pyfull_gridding(const pyarr<T> &uvw, ...@@ -777,7 +789,10 @@ template<typename T> pyarr<T> Pyfull_gridding(const pyarr<T> &uvw,
auto wgt2 = make_const_mav<2>(wgt); auto wgt2 = make_const_mav<2>(wgt);
auto dirty=makeArray<T>({npix_x,npix_y}); auto dirty=makeArray<T>({npix_x,npix_y});
auto dirty2 =make_mav<2>(dirty); auto dirty2 =make_mav<2>(dirty);
{
py::gil_scoped_release release;
full_gridding(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,dirty2,verbosity); full_gridding(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,dirty2,verbosity);
}
return dirty; return dirty;
} }
...@@ -792,7 +807,10 @@ template<typename T> pyarr<complex<T>> Pyfull_degridding(const pyarr<T> &uvw, ...@@ -792,7 +807,10 @@ template<typename T> pyarr<complex<T>> Pyfull_degridding(const pyarr<T> &uvw,
auto wgt2 = make_const_mav<2>(wgt); auto wgt2 = make_const_mav<2>(wgt);
auto ms=makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)}); auto ms=makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)});
auto ms2 =make_mav<2>(ms); auto ms2 =make_mav<2>(ms);
{
py::gil_scoped_release release;
full_degridding(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,ms2,verbosity); full_degridding(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,ms2,verbosity);
}
return ms; return ms;
} }
......
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