diff --git a/gridder_cxx.h b/gridder_cxx.h index fd7041e227344016f7144efb45254f9963800c3a..d492f0053f2ff0cfd403a4653322a89cf80fc292 100644 --- a/gridder_cxx.h +++ b/gridder_cxx.h @@ -1095,6 +1095,100 @@ cout << "applying correction for gridding in w direction" << endl; apply_wcorr(gconf, dirty, kernel, dw); } +template<typename T> void dirty2vis_wstack(const Baselines<T> &baselines, + const GridderConfig<T> &gconf, const const_mav<uint32_t,1> &idx, + const const_mav<complex<T>,2> &dirty, mav<complex<T>,1> &vis) + { + auto nx_dirty=gconf.Nxdirty(); + auto ny_dirty=gconf.Nydirty(); + auto nu=gconf.Nu(); + auto nv=gconf.Nv(); + auto psx=gconf.Pixsize_x(); + auto psy=gconf.Pixsize_y(); + size_t nvis = vis.shape(0); + checkShape(idx.shape(),vis.shape()); + checkShape(dirty.shape(),{nx_dirty,ny_dirty}); + + // determine w values for every visibility, and min/max w; + T wmin=T(1e38), wmax=T(-1e38); + vector<T> wval(nvis); + for (size_t ipart=0; ipart<nvis; ++ipart) + { + wval[ipart] = baselines.effectiveCoord(idx(ipart)).w; + wmin = min(wmin,wval[ipart]); + wmax = max(wmax,wval[ipart]); + } +cout << "data w range: " << wmin << " to " << wmax << endl; + + double x0 = -0.5*nx_dirty*psx, + y0 = -0.5*ny_dirty*psy; + double nmin = sqrt(max(1.-x0*x0-y0*y0,0.))-1.; + double dw = 0.25/abs(nmin); + + auto w_supp = gconf.Supp(); + EC_Kernel_with_correction<T> kernel(w_supp); + wmin -= 0.5*w_supp*dw; + wmax += 0.5*w_supp*dw; + size_t nplanes = size_t((wmax-wmin)/dw)+2; +cout << "nplanes: " << nplanes << endl; + vector<size_t> nvis_plane(nplanes,0); + vector<int> minplane(nvis); + for (size_t ipart=0; ipart<nvis; ++ipart) + { + int plane0 = int((wval[ipart]-wmin)/dw+0.5*w_supp)-int(w_supp-1); + minplane[ipart]=plane0; + for (int i=max<int>(0,plane0); i<min<int>(nplanes,plane0+w_supp); ++i) + ++nvis_plane[i]; + } + + for (size_t i=0; i<nvis; ++i) + vis(i) = 0.; + + tmpStorage<complex<T>,2> tdirty_({nx_dirty,ny_dirty}); + auto tdirty=tdirty_.getMav(); + for (size_t i=0; i<nx_dirty; ++i) + for (size_t j=0; j<ny_dirty; ++j) + tdirty(i,j) = dirty(i,j); + // correct for w gridding +cout << "applying correction for gridding in w direction" << endl; + apply_wcorr(gconf, tdirty, kernel, dw); + + tmpStorage<complex<T>,2> grid_({nu,nv}); + auto grid=grid_.getMav(); + tmpStorage<complex<T>,2> tdirty2_({nx_dirty,ny_dirty}); + auto tdirty2=tdirty2_.getMav(); + vector<complex<T>> vis_loc_; + vector<uint32_t> idx_loc_; + for (size_t iw=0; iw<nplanes; ++iw) + { +cout << "working on w plane #" << iw << " containing " << nvis_plane[iw] + << " visibilities" << endl; + if (nvis_plane[iw]==0) continue; + double wcur = wmin+iw*dw; + gconf.apply_wscreen(tdirty, tdirty2, wcur, false); + gconf.dirty2grid_c(tdirty2, grid); + size_t cnt=0; + vis_loc_.resize(nvis_plane[iw]); + auto vis_loc = mav<complex<T>,1>(vis_loc_.data(), {nvis_plane[iw]}); + idx_loc_.resize(nvis_plane[iw]); + auto idx_loc = mav<uint32_t,1>(idx_loc_.data(), {nvis_plane[iw]}); + for (size_t ipart=0; ipart<nvis; ++ipart) + if ((iw>=minplane[ipart]) && (iw<minplane[ipart]+w_supp)) + idx_loc[cnt++] = idx[ipart]; + +myassert(cnt==nvis_plane[iw],"must not happen 2"); + grid2vis_c(baselines, gconf, const_mav<uint32_t,1>(idx_loc), const_mav<complex<T>,2>(grid), vis_loc, const_mav<T,1>(nullptr,{0})); + cnt=0; + for (size_t ipart=0; ipart<nvis; ++ipart) + if ((iw>=minplane[ipart]) && (iw<minplane[ipart]+w_supp)) + { + double x=min(1.,2./(w_supp*dw)*abs(wcur-wval[ipart])); + vis[ipart] += vis_loc[cnt]*kernel(x); + ++cnt; + } + } + } + } // namespace gridder #endif diff --git a/nifty_gridder.cc b/nifty_gridder.cc index 137f768596237499b54bde3f296eae2366d46237..38be25ff746e7081cea486ddb29f6a22341b2ab8 100644 --- a/nifty_gridder.cc +++ b/nifty_gridder.cc @@ -698,6 +698,22 @@ template<typename T> pyarr<complex<T>> Pyvis2dirty_wstack(const PyBaselines<T> & return dirty; } +template<typename T> pyarr<complex<T>> Pydirty2vis_wstack(const PyBaselines<T> &baselines, + const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_, + const pyarr<complex<T>> &dirty_) + { + auto idx2=make_const_mav<1>(idx_); + auto nvis = idx2.shape(0); + auto vis = makeArray<complex<T>>({nvis}); + auto vis2=make_mav<1>(vis); + auto dirty2=make_const_mav<2>(dirty_); + { + py::gil_scoped_release release; + dirty2vis_wstack<T>(baselines, gconf, idx2, dirty2, vis2); + } + return vis; + } + } // unnamed namespace PYBIND11_MODULE(nifty_gridder, m) @@ -769,4 +785,6 @@ PYBIND11_MODULE(nifty_gridder, m) m.def("get_nthreads",&get_nthreads, get_nthreads_DS); m.def("vis2dirty_wstack",&Pyvis2dirty_wstack<double>, "baselines"_a, "gconf"_a, "idx"_a, "vis"_a); + m.def("dirty2vis_wstack",&Pydirty2vis_wstack<double>, "baselines"_a, "gconf"_a, + "idx"_a, "dirty"_a); }