sync with wsclean version

gridder_cxx.h

@@ -152,27 +152,6 @@ template<typename T, size_t ndim> const_mav<T, ndim> nullmav()
   shp.fill(0);
   return const_mav<T, ndim>(nullptr, shp);
   }
-template<typename T, size_t ndim> class tmpStorage
-  {
-  private:
-    vector<T> d;
-    mav<T,ndim> mav_;
-
-    static size_t prod(const array<size_t,ndim> &shp)
-      {
-      size_t res=1;
-      for (auto v: shp) res*=v;
-      return res;
-      }
-
-  public:
-    tmpStorage(const array<size_t,ndim> &shp)
-      : d(prod(shp)), mav_(d.data(), shp) {}
-    mav<T,ndim> &getMav() { return mav_; }
-    const_mav<T,ndim> getCmav() { return cmav(mav_); }
-    void fill(const T & val)
-      { std::fill(d.begin(), d.end(), val); }
-  };
 
 //
 // basic utilities
@@ -239,6 +218,28 @@ template<size_t ndim> void checkShape
 template<typename T> inline T fmod1 (T v)
   { return v-floor(v); }
 
+template<typename T, size_t ndim> class tmpStorage
+  {
+  private:
+    vector<T> d;
+    mav<T,ndim> mav_;
+
+    static size_t prod(const array<size_t,ndim> &shp)
+      {
+      size_t res=1;
+      for (auto v: shp) res*=v;
+      return res;
+      }
+
+  public:
+    tmpStorage(const array<size_t,ndim> &shp)
+      : d(prod(shp)), mav_(d.data(), shp) {}
+    mav<T,ndim> &getMav() { return mav_; }
+    const_mav<T,ndim> getCmav() { return cmav(mav_); }
+    void fill(const T & val)
+      { std::fill(d.begin(), d.end(), val); }
+  };
+
 //
 // Utilities for Gauss-Legendre quadrature
 //
@@ -254,8 +255,8 @@ void legendre_prep(int n, vector<double> &x, vector<double> &w, size_t nthreads)
   x.resize(m);
   w.resize(m);
 
-  double t0 = 1 - (1-1./n) / (8.*n*n);
-  double t1 = 1./(4.*n+2.);
+  const double t0 = 1 - (1-1./n) / (8.*n*n);
+  const double t1 = 1./(4.*n+2.);
 
 #pragma omp parallel num_threads(nthreads)
 {
@@ -592,15 +593,16 @@ class GridderConfig
     double ushift, vshift;
     int maxiu0, maxiv0;
 
-    complex<double> wscreen(double x, double y, double w, bool adjoint) const
+    complex<double> wscreen(double x, double y, double w, bool adjoint,
+      bool divide_by_n) const
       {
       constexpr double pi = 3.141592653589793238462643383279502884197;
       double tmp = 1-x-y;
-      if (tmp<0) return 0.;
-      double nm1 = (-x-y)/(sqrt(tmp)+1); // more accurate form of sqrt(1-x-y)
-      double n = nm1+1., xn = 1./n;
+      if (tmp<=0) return divide_by_n ? 0. : 1.; // no phase factor beyond the horizon
+      double nm1 = (-x-y)/(sqrt(tmp)+1); // more accurate form of sqrt(1-x-y)-1
       double phase = 2*pi*w*nm1;
       if (adjoint) phase *= -1;
+      double xn = divide_by_n ? 1./(nm1+1) : 1;
       return complex<double>(cos(phase)*xn, sin(phase)*xn);
       }
 
@@ -760,7 +762,7 @@ class GridderConfig
       }
 
     template<typename T> void apply_wscreen(const const_mav<complex<T>,2> &dirty,
-      mav<complex<T>,2> &dirty2, double w, bool adjoint) const
+      mav<complex<T>,2> &dirty2, double w, bool adjoint, bool divide_by_n) const
       {
       checkShape(dirty.shape(), {nx_dirty, ny_dirty});
       checkShape(dirty2.shape(), {nx_dirty, ny_dirty});
@@ -774,7 +776,7 @@ class GridderConfig
         for (size_t j=0; j<=ny_dirty/2; ++j)
           {
           double fy = y0+j*psy;
-          auto ws = complex<T>(wscreen(fx, fy*fy, w, adjoint));
+          auto ws = complex<T>(wscreen(fx, fy*fy, w, adjoint, divide_by_n));
           dirty2(i,j) = dirty(i,j)*ws; // lower left
           size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
           if ((i>0)&&(i<i2))
@@ -847,8 +849,6 @@ template<typename T, typename T2=complex<T>> class Helper
       {
       if (bu0<-nsafe) return; // nothing written into buffer yet
 
-#pragma omp critical (gridder_writing_to_grid)
-{
       int idxu = (bu0+nu)%nu;
       int idxv0 = (bv0+nv)%nv;
       for (int iu=0; iu<su; ++iu)
@@ -863,7 +863,6 @@ template<typename T, typename T2=complex<T>> class Helper
         locks[idxu].unlock();
         if (++idxu>=nu) idxu=0;
         }
-}
       }
 
     void load()
@@ -1323,8 +1322,13 @@ template<typename T> void apply_wcorr(const GridderConfig &gconf,
       double tmp = 1.-fx-fy;
       if (tmp>=0)
         {
-        auto nm1 = (-fx-fy)/(sqrt(1.-fx-fy)+1.); // accurate form of sqrt(1-x-y)
-        fct = T((nm1<=-1) ? 0. : kernel.corfac(nm1*dw));
+        auto nm1 = (-fx-fy)/(sqrt(tmp)+1.); // accurate form of sqrt(1-x-y)-1
+        fct = T(kernel.corfac(nm1*dw));
+        }
+      else // beyond the horizon, don't really know what to do here
+        {
+        double nm1 = sqrt(-tmp)-1;
+        fct = T(kernel.corfac(nm1*dw));
         }
       size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
       dirty(i,j)*=fct;
@@ -1411,9 +1415,12 @@ template<typename Serv> class WgridHelper
       if (verbosity>0) cout << "Using single-threaded mode" << endl;
 #endif
       if (verbosity>0) cout << "W range: " << wmin << " to " << wmax << endl;
+
       double x0 = -0.5*gconf.Nxdirty()*gconf.Pixsize_x(),
              y0 = -0.5*gconf.Nydirty()*gconf.Pixsize_y();
       double nmin = sqrt(max(1.-x0*x0-y0*y0,0.))-1.;
+      if (x0*x0+y0*y0>1.)
+        nmin = -sqrt(abs(1.-x0*x0-y0*y0))-1.;
       dw = 0.25/abs(nmin);
       nplanes = size_t((wmax-wmin)/dw+2);
       dw = (1.+1e-13)*(wmax-wmin)/(nplanes-1);
@@ -1426,6 +1433,15 @@ template<typename Serv> class WgridHelper
       if (verbosity>0) cout << "nplanes: " << nplanes << endl;
 
       minplane.resize(nplanes);
+#if 0
+      // extra short, but potentially inefficient version:
+      for (size_t ipart=0; ipart<nvis; ++ipart)
+        {
+        int plane0 = max(0,int(1+(abs(srv.getCoord(ipart).w)-(0.5*supp*dw)-wmin)/dw));
+        minplane[plane0].push_back(idx_t(ipart));
+        }
+#else
+      // full fledged, overdesigned OpenMP version
       vector<size_t> tcnt(max<int>(nthreads, my_max_threads())*nplanes,0);
 #pragma omp parallel num_threads(nthreads)
 {
@@ -1461,6 +1477,7 @@ template<typename Serv> class WgridHelper
         minplane[plane0][myofs[plane0]++]=idx_t(ipart);
         }
 }
+#endif
       }
 
     typename Serv::Tsub getSubserv() const
@@ -1503,7 +1520,7 @@ template<typename T, typename Serv> void x2dirty(
       grid.fill(0);
       x2grid_c(gconf, hlp.getSubserv(), grid, hlp.W(), dw);
       gconf.grid2dirty_c_overwrite(grid, tdirty);
-      gconf.apply_wscreen(cmav(tdirty), tdirty, hlp.W(), true);
+      gconf.apply_wscreen(cmav(tdirty), tdirty, hlp.W(), true, true);
 #pragma omp parallel for num_threads(nthreads)
       for (size_t i=0; i<gconf.Nxdirty(); ++i)
         for (size_t j=0; j<gconf.Nydirty(); ++j)
@@ -1568,7 +1585,7 @@ template<typename T, typename Serv> void dirty2x(
       for (size_t i=0; i<nx_dirty; ++i)
         for (size_t j=0; j<ny_dirty; ++j)
           tdirty2(i,j) = tdirty(i,j);
-      gconf.apply_wscreen(cmav(tdirty2), tdirty2, hlp.W(), false);
+      gconf.apply_wscreen(cmav(tdirty2), tdirty2, hlp.W(), false, true);
       gconf.dirty2grid_c(cmav(tdirty2), grid);
 
       grid2x_c(gconf, cmav(grid), hlp.getSubserv(), hlp.W(), dw);

nifty_gridder.cc

@@ -284,6 +284,8 @@ w : float
     the w value to use
 adjoint: bool
     if True, apply the complex conjugate of the w screen
+ivide_by_n: bool, default=True
+    if True, divide ny n.
 
 Returns
 =======
@@ -420,23 +422,25 @@ class PyGridderConfig: public GridderConfig
       myfail("type matching failed: 'dirty' has neither type 'c8' nor 'c16'");
       }
     template<typename T> pyarr<complex<T>> apply_wscreen2
-      (const pyarr<complex<T>> &dirty, double w, bool adjoint) const
+      (const pyarr<complex<T>> &dirty, double w, bool adjoint,
+       bool divide_by_n) const
       {
       auto dirty2 = make_const_mav<2>(dirty);
       auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
       auto res2 = make_mav<2>(res);
       {
       py::gil_scoped_release release;
-      GridderConfig::apply_wscreen(dirty2, res2, w, adjoint);
+      GridderConfig::apply_wscreen(dirty2, res2, w, adjoint, divide_by_n);
       }
       return res;
       }
-    py::array apply_wscreen(const py::array &dirty, double w, bool adjoint) const
+    py::array apply_wscreen(const py::array &dirty, double w, bool adjoint,
+      bool divide_by_n) const
       {
       if (isPytype<complex<float>>(dirty))
-        return apply_wscreen2<float>(dirty, w, adjoint);
+        return apply_wscreen2<float>(dirty, w, adjoint, divide_by_n);
       if (isPytype<complex<double>>(dirty))
-        return apply_wscreen2<double>(dirty, w, adjoint);
+        return apply_wscreen2<double>(dirty, w, adjoint, divide_by_n);
       myfail("type matching failed: 'dirty' has neither type 'c8' nor 'c16'");
       }
   };
@@ -1087,7 +1091,7 @@ PYBIND11_MODULE(nifty_gridder, m)
        dirty2grid_DS, "dirty"_a)
     .def("dirty2grid_c", &PyGridderConfig::dirty2grid_c, "dirty"_a)
     .def("apply_wscreen", &PyGridderConfig::apply_wscreen,
-      apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)
+      apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a, "divide_by_n"_a=true)
 
     // pickle support
     .def(py::pickle(