Support hoisted allocation in real functions

nifty_gridder.cc

@@ -203,13 +203,13 @@ template<typename T> pyarr_c<T> provideCArray(py::object &in,
   }
 
 template<typename T> pyarr_c<T> complex2hartley
-  (const pyarr_c<complex<T>> &grid_)
+  (const pyarr_c<complex<T>> &grid_, py::object &grid_in)
   {
   checkArray(grid_, "grid", {0,0});
   size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
   auto grid = grid_.data();
 
-  auto res = makeArray<T>({nu,nv});
+  auto res = provideCArray<T>(grid_in, {nu, nv});
   auto grid2 = res.mutable_data();
   {
   py::gil_scoped_release release;
@@ -222,8 +222,8 @@ template<typename T> pyarr_c<T> complex2hartley
       size_t xv = (v==0) ? 0 : nv-v;
       size_t i1 = u*nv+v;
       size_t i2 = xu*nv+xv;
-      grid2[i1] = T(0.5)*(grid[i1].real()+grid[i1].imag()+
-                          grid[i2].real()-grid[i2].imag());
+      grid2[i1] += T(0.5)*(grid[i1].real()+grid[i1].imag()+
+                           grid[i2].real()-grid[i2].imag());
       }
     }
   }
@@ -715,8 +715,8 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c(
 
 template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
   const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
-  const pyarr<complex<T>> &vis_)
-  { return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_, None)); }
+  const pyarr<complex<T>> &vis_, py::object &grid_in)
+  { return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_, None), grid_in); }
 
 template<typename T> pyarr_c<complex<T>> ms2grid_c(
   const Baselines<T> &baselines, const GridderConfig<T> &gconf,
@@ -773,8 +773,8 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c(
 
 template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines,
   const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
-  const pyarr<complex<T>> &ms_)
-  { return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_, None)); }
+  const pyarr<complex<T>> &ms_, py::object &grid_in)
+  { return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_, None), grid_in); }
 
 template<typename T> pyarr_c<complex<T>> ms2grid_c_wgt(
   const Baselines<T> &baselines, const GridderConfig<T> &gconf,
@@ -833,8 +833,9 @@ template<typename T> pyarr_c<complex<T>> ms2grid_c_wgt(
 
 template<typename T> pyarr_c<T> ms2grid_wgt(const Baselines<T> &baselines,
   const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
-  const pyarr<complex<T>> &ms_, const pyarr<T> &wgt_)
-  { return complex2hartley(ms2grid_c_wgt(baselines, gconf, idx_, ms_, wgt_, None)); }
+  const pyarr<complex<T>> &ms_, const pyarr<T> &wgt_,
+  py::object &grid_in)
+  { return complex2hartley(ms2grid_c_wgt(baselines, gconf, idx_, ms_, wgt_, None), grid_in); }
 
 template<typename T> pyarr_c<complex<T>> grid2vis_c(
   const Baselines<T> &baselines, const GridderConfig<T> &gconf,
@@ -1343,10 +1344,11 @@ PYBIND11_MODULE(nifty_gridder, m)
     "gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
     "wmin"_a=-1e30, "wmax"_a=1e30);
   m.def("vis2grid",&vis2grid<double>, vis2grid_DS, "baselines"_a, "gconf"_a,
-    "idx"_a, "vis"_a);
-  m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a);
+    "idx"_a, "vis"_a, "grid_in"_a=None);
+  m.def("ms2grid",&ms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a,
+    "grid_in"_a=None);
   m.def("ms2grid_wgt",&ms2grid_wgt<double>, "baselines"_a, "gconf"_a, "idx"_a,
-    "ms"_a, "wgt"_a);
+    "ms"_a, "wgt"_a, "grid_in"_a=None);
   m.def("grid2vis",&grid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a,
     "idx"_a, "grid"_a);
   m.def("grid2ms",&grid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,

test.py

@@ -60,16 +60,53 @@ def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
     idx = ng.getIndices(baselines, gconf, flags)
     ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
     vis = baselines.ms2vis(ms, idx)
-    user_grid = np.zeros((gconf.Nu(), gconf.Nv()),
-                         dtype=np.complex128)
-    grid = ng.vis2grid_c(baselines, gconf, idx, vis, grid_in=None)
-    grid2 = ng.vis2grid_c(baselines, gconf, idx, vis, grid_in=user_grid)
 
-    # Almost same visibilities
+    # -----------------------------
+    # Test real grid case (ms2grid)
+    # -----------------------------
+    real_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.float64)
+    grid = ng.ms2grid(baselines, gconf, idx, ms, grid_in=None)
+    grid2 = ng.ms2grid(baselines, gconf, idx, ms, grid_in=real_grid)
+
+    assert_array_almost_equal(grid, grid2)
+    assert grid.dtype == real_grid.dtype == grid2.dtype
+    assert id(grid2) == id(real_grid)  # Same grid object
+
+    # ------------------------------
+    # Test real grid case (vis2grid)
+    # ------------------------------
+    real_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.float64)
+    grid = ng.vis2grid(baselines, gconf, idx, vis, grid_in=None)
+    grid2 = ng.vis2grid(baselines, gconf, idx, vis, grid_in=real_grid)
+
+    # Almost same result
     assert_array_almost_equal(grid, grid2)
+    assert grid.dtype == real_grid.dtype == grid2.dtype
+    assert id(grid2) == id(real_grid)  # Same grid object
+
+    # ----------------------------------
+    # Test complex grid case (ms2grid_c)
+    # ----------------------------------
+    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
+    grid = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=None)
+    grid2 = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=cplx_grid)
 
-    # Same grid object
-    assert id(grid2) == id(user_grid)
+    # Almost same result
+    assert_array_almost_equal(grid, grid2)
+    assert grid.dtype == cplx_grid.dtype == grid2.dtype
+    assert id(grid2) == id(cplx_grid)  # Same grid object
+
+    # ----------------------------------
+    # Test complex grid case (vis2grid_c)
+    # ----------------------------------
+    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
+    grid = ng.vis2grid_c(baselines, gconf, idx, vis, grid_in=None)
+    grid2 = ng.vis2grid_c(baselines, gconf, idx, vis, grid_in=cplx_grid)
+
+    # Almost same result
+    assert_array_almost_equal(grid, grid2)
+    assert grid.dtype == cplx_grid.dtype == grid2.dtype
+    assert id(grid2) == id(cplx_grid)  # Same grid object
 
 
 def test_pickling():