add dirty2vis_wstack (untested)

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

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);
   }