add experimental get_diagonal() method

nifty_gridder.cc

@@ -681,19 +681,19 @@ template<typename T> class GridderConfig
       }
   };
 
-template<typename T> class Helper
+template<typename T, typename T2=complex<T>> class Helper
   {
   private:
     const GridderConfig<T> &gconf;
     int nu, nv, nsafe, w;
     T beta;
-    const complex<T> *grid_r;
-    complex<T> *grid_w;
+    const T2 *grid_r;
+    T2 *grid_w;
     int su, sv;
     int iu0, iv0; // start index of the current visibility
     int bu0, bv0; // start index of the current buffer
 
-    vector<complex<T>> rbuf, wbuf;
+    vector<T2> rbuf, wbuf;
 
     void dump() const
       {
@@ -733,12 +733,11 @@ template<typename T> class Helper
       }
 
   public:
-    const complex<T> *p0r;
-    complex<T> *p0w;
+    const T2 *p0r;
+    T2 *p0w;
     vector<T> kernel;
 
-    Helper(const GridderConfig<T> &gconf_, const complex<T> *grid_r_,
-      complex<T> *grid_w_)
+    Helper(const GridderConfig<T> &gconf_, const T2 *grid_r_, T2 *grid_w_)
       : gconf(gconf_), nu(gconf.Nu()), nv(gconf.Nv()), nsafe(gconf.Nsafe()),
         w(gconf.W()), beta(gconf.Beta()), grid_r(grid_r_), grid_w(grid_w_),
         su(2*nsafe+(1<<logsquare)), sv(2*nsafe+(1<<logsquare)),
@@ -1274,6 +1273,53 @@ template<typename T> pyarr_c<complex<T>> apply_holo(
   return res;
   }
 
+template<typename T> pyarr_c<T> get_diagonal(
+  const Baselines<T> &baselines, const GridderConfig<T> &gconf,
+  const pyarr<uint32_t> &idx_, size_t dx, size_t dy)
+  {
+  size_t nu=gconf.Nu(), nv=gconf.Nv();
+  size_t w = gconf.W();
+  if (dx>=w) throw runtime_error("dx must be smaller than W");
+  if (dy>=w) throw runtime_error("dy must be smaller than W");
+  checkArray(idx_, "idx", {0});
+  size_t nvis = size_t(idx_.shape(0));
+  auto idx = idx_.template unchecked<1>();
+
+  auto res = makeArray<T>({nu, nv});
+  auto ogrid = res.mutable_data();
+  {
+  py::gil_scoped_release release;
+  T beta = gconf.Beta();
+  for (size_t i=0; i<nu*nv; ++i) ogrid[i] = 0.;
+
+  // Loop over sampling points
+#pragma omp parallel num_threads(nthreads)
+{
+  Helper<T,T> hlp(gconf, nullptr, ogrid);
+  T emb = exp(-2*beta);
+  int jump = hlp.lineJump();
+  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(coord.u, coord.v);
+    auto * RESTRICT wptr = hlp.p0w;
+    for (size_t cu=0; cu<w-dx; ++cu)
+      {
+      auto f1=ku[cu]*ku[cu+dx]*emb*emb;
+      for (size_t cv=0; cv<w-dy; ++cv)
+        wptr[cv] += f1*kv[cv]*kv[cv+dy];
+      wptr += jump;
+      }
+    }
+}
+  }
+  return res;
+  }
+
 constexpr auto getIndices_DS = R"""(
 Selects a subset of entries from a `Baselines` object.
 
@@ -1438,6 +1484,8 @@ PYBIND11_MODULE(nifty_gridder, m)
     "idx"_a, "grid"_a, "wgt"_a, "ms_in"_a=None);
   m.def("apply_holo",&apply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
     "grid"_a);
+  m.def("get_diagonal", &get_diagonal<double>, "baselines"_a, "gconf"_a,
+    "idx"_a, "dx"_a, "dy"_a);
   m.def("set_nthreads",&set_nthreads, set_nthreads_DS, "nthreads"_a);
   m.def("get_nthreads",&get_nthreads, get_nthreads_DS);
   }