Merge branch 'master' into improve_scaling

nifty_gridder.cc

@@ -232,12 +232,12 @@ Converts from UV grid to dirty image (FFT, cropping, correction)
 
 Parameters
 ==========
-grid: np.array((nu, nv), dtype=np.float64)
+grid: np.array((nu, nv), dtype=np.float32 or np.float64)
     gridded UV data
 
 Returns
 =======
-nd.array((nxdirty, nydirty), dtype=np.float64)
+nd.array((nxdirty, nydirty), same dtype as `grid`)
     the dirty image
 )""";
 
@@ -246,12 +246,12 @@ Converts from a dirty image to a UV grid (correction, padding, FFT)
 
 Parameters
 ==========
-dirty: nd.array((nxdirty, nydirty), dtype=np.float64)
+dirty: nd.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
     the dirty image
 
 Returns
 =======
-np.array((nu, nv), dtype=np.float64)
+np.array((nu, nv), same dtype as `dirty`)
     gridded UV data
 )""";
 
@@ -260,14 +260,14 @@ Applies the taper (or its inverse) to an image
 
 Parameters
 ==========
-img: nd.array((nxdirty, nydirty), dtype=np.float64)
+img: nd.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
     the image
 divide: bool
     if True, the routine dividex by the taper, otherwise it multiplies by it
 
 Returns
 =======
-np.array((nxdirty, nydirty), dtype=np.float64)
+np.array((nxdirty, nydirty), same dtype as `img`)
     the image with the taper applied
 )""";
 
@@ -276,7 +276,7 @@ Applies the w screen to an image
 
 Parameters
 ==========
-dirty: nd.array((nxdirty, nydirty), dtype=np.complex128)
+dirty: nd.array((nxdirty, nydirty), dtype=np.complex64 or np.complex128)
     the image
 w : float
     the w value to use
@@ -285,7 +285,7 @@ adjoint: bool
 
 Returns
 =======
-np.array((nxdirty, nydirty), dtype=np.complex128)
+np.array((nxdirty, nydirty), same dtype as 'dirty')
     the image with the w screen applied
 )""";
 
@@ -316,7 +316,7 @@ class PyGridderConfig: public GridderConfig
       double pixsize_x, double pixsize_y, size_t nthreads)
       : GridderConfig(nxdirty, nydirty, epsilon, pixsize_x, pixsize_y, nthreads) {}
 
-    template<typename T> pyarr<T> apply_taper(const pyarr<T> &img, bool divide) const
+    template<typename T> pyarr<T> apply_taper2(const pyarr<T> &img, bool divide) const
       {
       auto res = makeArray<T>({nx_dirty, ny_dirty});
       auto img2 = make_const_mav<2>(img);
@@ -327,7 +327,15 @@ class PyGridderConfig: public GridderConfig
       }
       return res;
       }
-    template<typename T> pyarr<T> grid2dirty(const pyarr<T> &grid) const
+    py::array apply_taper(const py::array &img, bool divide) const
+      {
+      if (isPytype<complex<float>>(img))
+        return apply_taper2<float>(img, divide);
+      if (isPytype<complex<double>>(img))
+        return apply_taper2<double>(img, divide);
+      myfail("type matching failed: 'img' has neither type 'f4' nor 'f8'");
+      }
+    template<typename T> pyarr<T> grid2dirty2(const pyarr<T> &grid) const
       {
       auto res = makeArray<T>({nx_dirty, ny_dirty});
       auto grid2=make_const_mav<2>(grid);
@@ -338,7 +346,16 @@ class PyGridderConfig: public GridderConfig
       }
       return res;
       }
-    template<typename T> pyarr<complex<T>> grid2dirty_c(const pyarr<complex<T>> &grid) const
+    py::array grid2dirty(const py::array &grid) const
+      {
+      if (isPytype<float>(grid))
+        return grid2dirty2<float>(grid);
+      if (isPytype<double>(grid))
+        return grid2dirty2<double>(grid);
+      myfail("type matching failed: 'grid' has neither type 'f4' nor 'f8'");
+      }
+    template<typename T> pyarr<complex<T>> grid2dirty_c2
+      (const pyarr<complex<T>> &grid) const
       {
       auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
       auto grid2=make_const_mav<2>(grid);
@@ -349,8 +366,16 @@ class PyGridderConfig: public GridderConfig
       }
       return res;
       }
+    py::array grid2dirty_c(const py::array &grid) const
+      {
+      if (isPytype<complex<float>>(grid))
+        return grid2dirty_c2<float>(grid);
+      if (isPytype<complex<double>>(grid))
+        return grid2dirty_c2<double>(grid);
+      myfail("type matching failed: 'grid' has neither type 'c8' nor 'c16'");
+      }
 
-    template<typename T> pyarr<T> dirty2grid(const pyarr<T> &dirty) const
+    template<typename T> pyarr<T> dirty2grid2(const pyarr<T> &dirty) const
       {
       auto dirty2 = make_const_mav<2>(dirty);
       auto grid = makeArray<T>({nu, nv});
@@ -361,7 +386,15 @@ class PyGridderConfig: public GridderConfig
       }
       return grid;
       }
-    template<typename T> pyarr<complex<T>> dirty2grid_c(const pyarr<complex<T>> &dirty) const
+    py::array dirty2grid(const py::array &dirty) const
+      {
+      if (isPytype<float>(dirty))
+        return dirty2grid2<float>(dirty);
+      if (isPytype<double>(dirty))
+        return dirty2grid2<double>(dirty);
+      myfail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
+      }
+    template<typename T> pyarr<complex<T>> dirty2grid_c2(const pyarr<complex<T>> &dirty) const
       {
       auto dirty2 = make_const_mav<2>(dirty);
       auto grid = makeArray<complex<T>>({nu, nv});
@@ -372,7 +405,16 @@ class PyGridderConfig: public GridderConfig
       }
       return grid;
       }
-    template<typename T> pyarr<complex<T>> apply_wscreen(const pyarr<complex<T>> &dirty, double w, bool adjoint) const
+    py::array dirty2grid_c(const py::array &dirty) const
+      {
+      if (isPytype<complex<float>>(dirty))
+        return dirty2grid_c2<float>(dirty);
+      if (isPytype<complex<double>>(dirty))
+        return dirty2grid_c2<double>(dirty);
+      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
       {
       auto dirty2 = make_const_mav<2>(dirty);
       auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
@@ -383,6 +425,14 @@ class PyGridderConfig: public GridderConfig
       }
       return res;
       }
+    py::array apply_wscreen(const py::array &dirty, double w, bool adjoint) const
+      {
+      if (isPytype<complex<float>>(dirty))
+        return apply_wscreen2<float>(dirty, w, adjoint);
+      if (isPytype<complex<double>>(dirty))
+        return apply_wscreen2<double>(dirty, w, adjoint);
+      myfail("type matching failed: 'dirty' has neither type 'c8' nor 'c16'");
+      }
   };
 
 
@@ -988,15 +1038,15 @@ PYBIND11_MODULE(nifty_gridder, m)
     .def("Nu", &PyGridderConfig::Nu)
     .def("Nv", &PyGridderConfig::Nv)
     .def("Supp", &PyGridderConfig::Supp)
-    .def("apply_taper", &PyGridderConfig::apply_taper<double>, apply_taper_DS,
+    .def("apply_taper", &PyGridderConfig::apply_taper, apply_taper_DS,
       "img"_a, "divide"_a=false)
-     .def("grid2dirty", &PyGridderConfig::grid2dirty<double>,
+     .def("grid2dirty", &PyGridderConfig::grid2dirty,
         grid2dirty_DS, "grid"_a)
-    .def("grid2dirty_c", &PyGridderConfig::grid2dirty_c<double>, "grid"_a)
-    .def("dirty2grid", &PyGridderConfig::dirty2grid<double>,
+    .def("grid2dirty_c", &PyGridderConfig::grid2dirty_c, "grid"_a)
+    .def("dirty2grid", &PyGridderConfig::dirty2grid,
        dirty2grid_DS, "dirty"_a)
-    .def("dirty2grid_c", &PyGridderConfig::dirty2grid_c<double>, "dirty"_a)
-    .def("apply_wscreen", &PyGridderConfig::apply_wscreen<double>,
+    .def("dirty2grid_c", &PyGridderConfig::dirty2grid_c, "dirty"_a)
+    .def("apply_wscreen", &PyGridderConfig::apply_wscreen,
       apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)
 
     // pickle support

test.py

@@ -7,6 +7,7 @@ from numpy.testing import assert_, assert_allclose, assert_array_almost_equal
 
 pmp = pytest.mark.parametrize
 
+
 def _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow):
     pixsize = np.pi/180/60/nxdirty
     conf = ng.GridderConfig(nxdirty=nxdirty,
@@ -44,13 +45,15 @@ def _dft(uvw, vis, conf, apply_w):
 
 
 def _l2error(a, b):
-    return np.sqrt(np.sum(np.abs(a-b)**2)/np.maximum(np.sum(np.abs(a)**2),np.sum(np.abs(b)**2)))
+    return np.sqrt(np.sum(np.abs(a-b)**2)/np.maximum(np.sum(np.abs(a)**2),
+                                                     np.sum(np.abs(b)**2)))
 
 
-def explicit_gridder(uvw, freq, ms, wgt, nxdirty, nydirty, xpixsize, ypixsize, apply_w):
+def explicit_gridder(uvw, freq, ms, wgt, nxdirty, nydirty, xpixsize, ypixsize,
+                     apply_w):
     speedoflight = 299792458.
     x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]],
-        indexing='ij')
+                       indexing='ij')
     x *= xpixsize
     y *= ypixsize
     res = np.zeros((nxdirty, nydirty))
@@ -115,7 +118,7 @@ def test_adjointness_ms2dirty(nxdirty, nydirty, nrow, nchan, epsilon, singleprec
 @pmp("wstacking", (True, False))
 @pmp("use_wgt", (False, True))
 @pmp("nthreads", (1, 2))
-@pmp("fov", (1.,20.))
+@pmp("fov", (1., 20.))
 def test_ms2dirty_against_wdft(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking, use_wgt, fov, nthreads):
     if singleprec and epsilon < 5e-5:
         return
@@ -149,10 +152,11 @@ def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
     vis = bl.ms2vis(ms, idx)
     grid = ng.vis2grid_c(bl, gconf, idx, vis)
     dirty = gconf.grid2dirty_c(grid)
-    dirty2 = np.random.rand(*dirty.shape)-0.5
+    dirty2 = (np.random.rand(*dirty.shape)-0.5 +
+              1j*(np.random.rand(*dirty.shape)-0.5))
     ms2 = bl.vis2ms(ng.grid2vis_c(bl, gconf, idx, gconf.dirty2grid_c(dirty2)),
                     idx)
-    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty.real, dirty2.real))
+    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty, dirty2).real)
 
 
 @pmp("nxdirty", (128, 300))
@@ -470,4 +474,3 @@ def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
     res0 += d2*grid
     res1 = ng.apply_holo(bl, conf, idx, grid)
     assert_allclose(res0, res1, rtol=1e-13)
-