humble beginnings

nifty_gridder.cc

@@ -1462,6 +1462,83 @@ template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
   return res;
   }
 
+template<typename T> pyarr_c<T> vis2dirty_wstack(const Baselines<T> &baselines,
+  const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
+  const pyarr<complex<T>> &vis_)
+  {
+  auto nxdirty=gconf.Nxdirty();
+  auto nydirty=gconf.Nydirty();
+  checkArray(vis_, "vis", {0});
+  size_t nvis = size_t(vis_.shape(0));
+  checkArray(idx_, "idx", {nvis});
+  auto vis=vis_.template unchecked<1>();
+  auto idx = idx_.template unchecked<1>();
+
+  // 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;
+
+//FIXME temporary
+double psz=gconf.Pixsize_x();
+  double dw = psz;// distance between two w planes, FIXME
+cout << "delta w: " << dw << endl;
+
+  double w_eps=1e-7;
+  auto w_supp = get_supp(w_eps);
+  auto beta=2.3*w_supp;
+  wmin -= 0.5*w_supp*dw;
+  wmax += 0.5*w_supp*dw;
+  size_t nplanes = size_t((wmax-wmin)/dw)+1;
+cout << "nplanes: " << nplanes << endl;
+  double dwmax=0.5*w_supp*dw;
+  vector<size_t> nvis_plane(nplanes,0);
+  for (size_t ipart=0; ipart<nvis; ++ipart)
+    {
+    int iplane = int(wval[ipart]-wmin)/dw;
+    for (int i=max<int>(0,iplane-w_supp); i<min<int>(nplanes,iplane+w_supp+1); ++i)
+      if (abs(wval[ipart]-(wmin+i*dw))<dwmax) ++nvis_plane[i];
+    }
+cout << "nvis/plane: ";
+for (auto nv:nvis_plane) cout << nv << " ";
+cout << endl;
+  auto accum_ = makeArray<complex<T>>({nxdirty, nydirty});
+  auto accum = accum_.mutable_data();
+  for (ptrdiff_t i=0; i<accum_.size(); ++i)
+    accum[i] = 0.;
+  for (size_t iw=0; iw<nplanes; ++iw)
+    {
+    if (nvis_plane[iw]==0) continue;
+    double wcur = wmin+iw*dw;
+    size_t cnt=0;
+    pyarr<complex<T>> vis_loc_ = makeArray<complex<T>>({nvis_plane[iw]});
+    auto vis_loc = vis_loc_.mutable_data();
+    pyarr<uint32_t> idx_loc_ = makeArray<uint32_t>({nvis_plane[iw]});
+    auto idx_loc = idx_loc_.mutable_data();
+    for (size_t ipart=0; ipart<nvis; ++ipart)
+      if (abs(wcur-wval[ipart]) < dwmax)
+        {
+        double x=2./(w_supp*dw)*abs(wcur-wval[ipart]);
+        vis_loc[cnt] = vis[ipart]*exp(beta*sqrt(1.-x*x));
+        idx_loc[cnt] = idx[ipart];
+        ++cnt;
+        }
+    auto dirty_ = gconf.apply_wscreen(
+      gconf.grid2dirty_c(vis2grid_c(baselines, gconf, idx_loc_, vis_loc_, None, None)), wcur, false);
+    auto dirty = dirty_.data();
+    for (ptrdiff_t i=0; i<accum_.size(); ++i)
+      accum[i] += dirty[i];
+    }
+  // correct for w gridding
+  return accum_;
+  }
+
 } // unnamed namespace
 
 PYBIND11_MODULE(nifty_gridder, m)
@@ -1543,4 +1620,6 @@ PYBIND11_MODULE(nifty_gridder, m)
     "idx"_a, "du"_a, "dv"_a, "wgt"_a=None);
   m.def("set_nthreads",&set_nthreads, set_nthreads_DS, "nthreads"_a);
   m.def("get_nthreads",&get_nthreads, get_nthreads_DS);
+  m.def("vis2dirty_wstack",&vis2dirty_wstack<double>, "baselines"_a, "gconf"_a,
+    "idx"_a, "vis"_a);
   }