test1

nifty_gridder.cc

@@ -659,41 +659,6 @@ template<typename T> class GridderConfig
       return res;
       }
 
-    pyarr_c<T> grid2dirty_wstacking(const pyarr_c<complex<T>> &grid, double w) const
-      {
-      checkArray(grid, "grid", {nu, nv});
-      auto tmp = makeArray<complex<T>>({nu, nv});
-      auto ptmp = tmp.mutable_data();
-      pocketfft::c2c({nu,nv},{grid.strides(0),grid.strides(1)},
-        {tmp.strides(0), tmp.strides(1)}, {0,1}, pocketfft::BACKWARD,
-        grid.data(), tmp.mutable_data(), T(1), nthreads);
-      auto res = makeArray<T>({nx_dirty, ny_dirty});
-      auto pout = res.mutable_data();
-      {
-        py::gil_scoped_release release;
-#pragma omp parallel num_threads(nthreads)
-{
-#pragma omp for schedule(static)
-        for (size_t i=0; i<nx_dirty; ++i)
-          {
-            double fx = (i-nx_dirty/2.)*Pixsize_x();
-            fx *= fx;
-            for (size_t j=0; j<ny_dirty; ++j)
-              {
-                size_t i2 = nu-nx_dirty/2+i;
-                if (i2>=nu) i2-=nu;
-                size_t j2 = nv-ny_dirty/2+j;
-                if (j2>=nv) j2-=nv;
-                double fy = (j-ny_dirty/2.)*Pixsize_y();
-                auto ws = wscreen(fx, fy*fy, w, true);
-                pout[ny_dirty*i + j] = std::real(ptmp[nv*i2+j2]*cfu[i]*cfv[j]*ws);
-              }
-          }
- }
-      }
-      return res;
-      }
-
     pyarr_c<T> dirty2grid(const pyarr_c<T> &dirty) const
       {
       checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
@@ -742,50 +707,14 @@ template<typename T> class GridderConfig
       }
       return tmp;
       }
-    inline complex<T> wscreen(double x, double y, double w, bool adjoint) const
-    {
+    complex<T> wscreen(double x, double y, double w, bool adjoint) const
+      {
       constexpr double pi = 3.141592653589793238462643383279502884197;
-      constexpr complex<double> imagunit {0,1};
-      double n = cos(sqrt(x+y));
-      complex<T> phase = 2*pi*w*imagunit*(n-1);
+      double n = cos(sqrt(x+y)), xn = 1./n;
+      double phase = 2*pi*w*(n-1);
       if (adjoint) phase *= -1;
-      return exp(phase)/n;
-    }
-    pyarr_c<complex<T>> dirty2grid_wstacking(const pyarr_c<T> &dirty, double w) const
-    {
-      checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
-      auto pdirty = dirty.data();
-      auto tmp = makeArray<complex<T>>({nu, nv});
-      auto ptmp = tmp.mutable_data();
-      pocketfft::stride_t strides{tmp.strides(0),tmp.strides(1)};
-      {
-        py::gil_scoped_release release;
-        for (size_t i=0; i<nu*nv; ++i)
-          ptmp[i] = 0.;
-#pragma omp parallel num_threads(nthreads)
-        {
-#pragma omp for schedule(static)
-        for (size_t i=0; i<nx_dirty; ++i)
-          {
-            double fx = (i-nx_dirty/2.)*Pixsize_x();
-            fx *= fx;
-            for (size_t j=0; j<ny_dirty; ++j)
-              {
-                size_t i2 = nu-nx_dirty/2+i;
-                if (i2>=nu) i2-=nu;
-                size_t j2 = nv-ny_dirty/2+j;
-                if (j2>=nv) j2-=nv;
-                double fy = (j-ny_dirty/2.)*Pixsize_y();
-                auto ws = wscreen(fx, fy*fy, w, false);
-                ptmp[nv*i2+j2] = std::real(pdirty[ny_dirty*i + j])*ws*cfu[i]*cfv[j];
-              }
-          }
- }
-        pocketfft::c2c({nu,nv}, strides, strides, {0,1}, pocketfft::FORWARD,
-                       ptmp, ptmp, T(1), nthreads);
+      return complex<T>(cos(phase)*xn, sin(phase)*xn);
       }
-      return tmp;
-    }
     inline void getpix(T u_in, T v_in, T &u, T &v, int &iu0, int &iv0) const
       {
       u=fmodulo(u_in*psx, T(1))*nu,
@@ -795,6 +724,40 @@ template<typename T> class GridderConfig
       iv0 = int(v-w*0.5 + 1 + nv) - nv;
       if (iv0+w>nv+nsafe) iv0 = nv+nsafe-w;
       }
+    pyarr_c<complex<T>> apply_wscreen(const pyarr_c<complex<T>> &dirty_, double w, bool adjoint) const
+      {
+      checkArray(dirty_, "dirty", {nx_dirty, ny_dirty});
+      auto dirty = dirty_.data();
+      auto res_ = makeArray<complex<T>>({nx_dirty, ny_dirty});
+      auto res = res_.mutable_data();
+      double x0 = -0.5*nx_dirty*psx,
+             y0 = -0.5*ny_dirty*psy;
+      {
+      py::gil_scoped_release release;
+#pragma omp parallel num_threads(nthreads)
+{
+#pragma omp for schedule(static)
+      for (size_t i=0; i<=nx_dirty/2; ++i)
+        {
+        double fx = x0+i*psx;
+        fx *= fx;
+        for (size_t j=0; j<=ny_dirty/2; ++j)
+          {
+          double fy = y0+j*psy;
+          auto ws = wscreen(fx, fy*fy, w, adjoint);
+          res[ny_dirty*i+j] = dirty[ny_dirty*i+j]*ws;
+          if (i>0)
+            res[ny_dirty*(nx_dirty-i)+j] = dirty[ny_dirty*(nx_dirty-i)+j]*ws;
+          if (j>0)
+            res[ny_dirty*i+(ny_dirty-j)] = dirty[ny_dirty*i+(ny_dirty-j)]*ws;
+          if ((i>0)&&(j>0))
+            res[ny_dirty*(nx_dirty-i)+(ny_dirty-j)] = dirty[ny_dirty*(nx_dirty-i)+(ny_dirty-j)]*ws;
+          }
+        }
+}
+      }
+    return res_;
+    }
   };
 
 template<typename T, typename T2=complex<T>> class Helper
@@ -1560,11 +1523,10 @@ PYBIND11_MODULE(nifty_gridder, m)
     .def("grid2dirty", &GridderConfig<double>::grid2dirty,
       grid2dirty_DS, "grid"_a)
     .def("grid2dirty_c", &GridderConfig<double>::grid2dirty_c, "grid"_a)
-    .def("grid2dirty_wstacking", &GridderConfig<double>::grid2dirty_wstacking, "grid"_a, "w"_a)
-    .def("dirty2grid_wstacking", &GridderConfig<double>::dirty2grid_wstacking, "grid"_a, "w"_a)
     .def("dirty2grid", &GridderConfig<double>::dirty2grid,
       dirty2grid_DS, "dirty"_a)
     .def("dirty2grid_c", &GridderConfig<double>::dirty2grid_c, "dirty"_a)
+    .def("apply_wscreen", &GridderConfig<double>::apply_wscreen, "dirty"_a, "w"_a, "adjoint"_a)
 
     // pickle support
     .def(py::pickle(

profile.py

@@ -28,15 +28,14 @@ if __name__ == '__main__':
     w = 1000.2
     ny, dy = nx, dx
     conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy)
-    x, y = conf.Nu(), conf.Nv()
+    x, y = conf.Nxdirty(), conf.Nydirty()
 
     fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
-    func0 = lambda x: (_wscreen(nx, dy, w).conjugate()*conf.grid2dirty_c(x)
-                       ).real
-    func1 = lambda x: conf.grid2dirty_wstacking(x, w)
+    func0 = lambda x: _wscreen(nx, dy, w).conjugate()*x
+    func1 = lambda x: conf.apply_wscreen(x, w, True)
     print(time_op(func0, fld), time_op(func1, fld))
 
     fld = np.random.randn(nx, ny)
-    func0 = lambda x: conf.dirty2grid_c(_wscreen(nx, dx, w)*x.real)
-    func1 = lambda x: conf.dirty2grid_wstacking(x, w)
+    func0 = lambda x: _wscreen(nx, dx, w)*x.real
+    func1 = lambda x: conf.apply_wscreen(x, w, False)
     print(time_op(func0, fld), time_op(func1, fld))

test.py

@@ -173,10 +173,10 @@ def test_wstacking(nx, dx, w):
     x, y = conf.Nu(), conf.Nv()
     fld = np.random.randn(x, y) + 1j*np.random.randn(x, y)
     res0 = (_wscreen(nx, dy, w).conjugate()*conf.grid2dirty_c(fld)).real
-    res1 = conf.grid2dirty_wstacking(fld, w)
+    res1 = conf.apply_wscreen(conf.grid2dirty_c(fld), w, True).real
     np.testing.assert_allclose(res0, res1)
 
     fld = np.random.randn(nx, ny)
     res0 = conf.dirty2grid_c(_wscreen(nx, dx, w)*fld.real)
-    res1 = conf.dirty2grid_wstacking(fld, w)
+    res1 = conf.dirty2grid_c(conf.apply_wscreen(fld, w, False))
     np.testing.assert_allclose(res0, res1)