add high-level routines

gridder_cxx.h

@@ -949,7 +949,7 @@ template<typename T> void get_correlations
 
 
 template<typename T, typename T2> void apply_wcorr(const GridderConfig<T> &gconf,
-  mav<T2,2> &dirty, const EC_Kernel_with_correction<T> &kernel, double dw)
+  const mav<T2,2> &dirty, const EC_Kernel_with_correction<T> &kernel, double dw)
   {
   auto nx_dirty=gconf.Nxdirty();
   auto ny_dirty=gconf.Nydirty();
@@ -983,7 +983,7 @@ template<typename T, typename T2> void apply_wcorr(const GridderConfig<T> &gconf
   }
 
 template<typename T, typename Serv> void x2dirty_wstack(const Baselines<T> &baselines,
-  const GridderConfig<T> &gconf, const Serv &srv, mav<complex<T>,2> &dirty)
+  const GridderConfig<T> &gconf, const Serv &srv, const mav<complex<T>,2> &dirty)
   {
   auto nx_dirty=gconf.Nxdirty();
   auto ny_dirty=gconf.Nydirty();
@@ -1180,6 +1180,74 @@ template<typename T> void dirty2vis_wstack(const Baselines<T> &baselines,
   dirty2x_wstack(baselines, gconf, dirty, VisServ<T,mav<complex<T>,1>,const_mav<T,1>>(baselines, idx, vis, const_mav<T,1>(nullptr,{0})));
   }
 
+template<typename T> vector<uint32_t> getWgtIndices(const Baselines<T> &baselines,
+  const GridderConfig<T> &gconf, const const_mav<T,2> &wgt)
+  {
+  size_t nrow=baselines.Nrows(),
+         nchan=baselines.Nchannels(),
+         nsafe=gconf.Nsafe();
+  bool have_wgt=wgt.size()!=0;
+  if (have_wgt) checkShape(wgt.shape(),{nrow,nchan});
+  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);
+
+  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))
+        {
+        auto uvw = baselines.effectiveCoord(RowChan{irow,size_t(ichan)});
+        T u, v;
+        int iu0, iv0;
+        gconf.getpix(uvw.u, uvw.v, u, v, iu0, iv0);
+        iu0 = (iu0+nsafe)>>logsquare;
+        iv0 = (iv0+nsafe)>>logsquare;
+        ++acc[nbv*iu0 + iv0 + 1];
+        tmp[idx++] = nbv*iu0 + iv0;
+        }
+
+  for (size_t i=1; i<acc.size(); ++i)
+    acc[i] += acc[i-1];
+
+  vector<uint32_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))
+        res[acc[tmp[idx++]]++] = irow*nchan+ichan;
+  return res;
+  }
+
+template<typename T> void full_gridding(const const_mav<T,2> &uvw,
+  const const_mav<T,1> &freq, const const_mav<complex<T>,2> &ms,
+  const const_mav<T,2> &wgt, T pixsize_x, T pixsize_y, double epsilon,
+  const mav<complex<T>,2> &dirty)
+  {
+  Baselines<T> baselines(uvw, freq);
+  GridderConfig<T> gconf(dirty.shape(0), dirty.shape(1), epsilon, pixsize_x, pixsize_y);
+  auto idx = getWgtIndices(baselines, gconf, wgt);
+  auto idx2 = const_mav<uint32_t,1>(idx.data(),{idx.size()});
+  auto srv = MsServ<T,const_mav<complex<T>,2>,const_mav<T,2>>(baselines,idx2,ms,wgt);
+  x2dirty_wstack(baselines, gconf, srv, dirty);
+  }
+
+template<typename T> void full_degridding(const const_mav<T,2> &uvw,
+  const const_mav<T,1> &freq, const const_mav<complex<T>,2> &dirty,
+  const const_mav<T,2> &wgt, T pixsize_x, T pixsize_y, double epsilon,
+  const mav<complex<T>,2> &ms)
+  {
+  Baselines<T> baselines(uvw, freq);
+  GridderConfig<T> gconf(dirty.shape(0), dirty.shape(1), epsilon, pixsize_x, pixsize_y);
+  auto idx = getWgtIndices(baselines, gconf, wgt);
+  auto idx2 = const_mav<uint32_t,1>(idx.data(),{idx.size()});
+  for (size_t i=0; i<ms.shape(0); ++i)
+    for (size_t j=0; j<ms.shape(1); ++j)
+      ms(i,j) = 0;
+  auto srv = MsServ<T,mav<complex<T>,2>,const_mav<T,2>>(baselines,idx2,ms,wgt);
+  dirty2x_wstack(baselines, gconf, dirty, srv);
+  }
+
 } // namespace gridder
 
 #endif

nifty_gridder.cc

@@ -614,7 +614,7 @@ np.array((nvis,), dtype=np.uint32)
     the compressed indices for all entries which match the selected criteria
     and are not flagged.
 )""";
-template<typename T> pyarr<uint32_t> getIndices(const PyBaselines<T> &baselines,
+template<typename T> pyarr<uint32_t> PygetIndices(const PyBaselines<T> &baselines,
   const PyGridderConfig<T> &gconf, const pyarr<bool> &flags_, int chbegin,
   int chend, T wmin, T wmax)
   {
@@ -704,6 +704,36 @@ template<typename T> pyarr<complex<T>> Pydirty2vis_wstack(const PyBaselines<T> &
   return vis;
   }
 
+template<typename T> pyarr<complex<T>> Pyfull_gridding(const pyarr<T> &uvw,
+  const pyarr<T> &freq, const pyarr<complex<T>> &ms, const py::object &wgt_,
+  size_t npix_x, size_t npix_y, T pixsize_x, T pixsize_y, double epsilon)
+  {
+  auto uvw2 = make_const_mav<2>(uvw);
+  auto freq2 = make_const_mav<1>(freq);
+  auto ms2 = make_const_mav<2>(ms);
+  auto wgt=providePotentialArray<T>(wgt_,{ms2.shape(0),ms2.shape(1)});
+  auto wgt2 = make_const_mav<2>(wgt);
+  auto dirty=makeArray<complex<T>>({npix_x,npix_y});
+  auto dirty2 =make_mav<2>(dirty);
+  full_gridding(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,dirty2);
+  return dirty;
+  }
+
+template<typename T> pyarr<complex<T>> Pyfull_degridding(const pyarr<T> &uvw,
+  const pyarr<T> &freq, const pyarr<complex<T>> &dirty, const py::object &wgt_,
+  T pixsize_x, T pixsize_y, double epsilon)
+  {
+  auto uvw2 = make_const_mav<2>(uvw);
+  auto freq2 = make_const_mav<1>(freq);
+  auto dirty2 = make_const_mav<2>(dirty);
+  auto wgt=providePotentialArray<T>(wgt_,{uvw2.shape(0),freq2.shape(0)});
+  auto wgt2 = make_const_mav<2>(wgt);
+  auto ms=makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)});
+  auto ms2 =make_mav<2>(ms);
+  full_degridding(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,ms2);
+  return ms;
+  }
+
 } // unnamed namespace
 
 PYBIND11_MODULE(nifty_gridder, m)
@@ -756,7 +786,7 @@ PYBIND11_MODULE(nifty_gridder, m)
                                        t[4].cast<double>());
 
         }));
-  m.def("getIndices", getIndices<double>, getIndices_DS, "baselines"_a,
+  m.def("getIndices", PygetIndices<double>, getIndices_DS, "baselines"_a,
     "gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
     "wmin"_a=-1e30, "wmax"_a=1e30);
   m.def("vis2grid_c",&Pyvis2grid_c<double>, vis2grid_c_DS, "baselines"_a,
@@ -777,4 +807,9 @@ PYBIND11_MODULE(nifty_gridder, m)
     "idx"_a, "vis"_a);
   m.def("dirty2vis_wstack",&Pydirty2vis_wstack<double>, "baselines"_a, "gconf"_a,
     "idx"_a, "dirty"_a);
+
+  m.def("full_gridding",&Pyfull_gridding<double>,"uvw"_a,"freq"_a,"vis"_a,
+    "wgt"_a,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a);
+  m.def("full_degridding",&Pyfull_degridding<double>,"uvw"_a,"freq"_a,"dirty"_a,
+    "wgt"_a,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a);
   }