merge

nifty_gridder.cc

@@ -138,9 +138,9 @@ Class storing UVW coordinates and channel information.
 
 Parameters
 ==========
-coord: np.array((nrows, 3), dtype=np.float)
+coord: np.array((nrows, 3), dtype=np.float64)
     u, v and w coordinates for each row
-freq: np.array((nchannels,), dtype=np.float)
+freq: np.array((nchannels,), dtype=np.float64)
     frequency for each individual channel (in Hz)
 )""";
 class PyBaselines: public Baselines
@@ -305,6 +305,8 @@ pixsize_x: float
     Pixel size in x direction (radians)
 pixsize_y: float
     Pixel size in y direction (radians)
+nthreads: int
+    the number of threads to use for all calculations involving this object.
 )""";
 class PyGridderConfig: public GridderConfig
   {
@@ -714,6 +716,35 @@ template<typename T> pyarr<T> vis2dirty2(const PyBaselines &baselines,
   }
   return dirty;
   }
+
+constexpr auto vis2dirty_DS = R"""(
+Converts an array of visibilities to a dirty image.
+
+Parameters
+==========
+baselines: Baselines
+    the Baselines object
+gconf: GridderConf
+    the GridderConf object to be used
+    (used to optimize the ordering of the indices)
+idx: np.array((nvis,), dtype=np.uint32)
+    the indices for the provided visibilities
+vis: np.array(nvis,), dtype=np.complex64 or np.complex128)
+    The input visibilities
+    Its data type determines the precision in which the calculation is carried
+    out.
+wgt: np.array((nvis,), dtype=float with same precision as `vis`, optional
+    If present, visibilities are multiplied by the corresponding entries.
+do_wstacking: bool
+    if True, the full improved w-stacking algorithm is carried out, otherwise
+    the w values are assumed to be zero.
+
+Returns
+=======
+np.array((nxdirty, nydirty), dtype=float of same precision as `vis`.)
+    The dirty image
+)""";
+
 py::array Pyvis2dirty(const PyBaselines &baselines,
   const PyGridderConfig &gconf, const py::array &idx,
   const py::array &vis, const py::object &wgt, bool do_wstacking)
@@ -745,6 +776,34 @@ template<typename T> pyarr<complex<T>> dirty2vis2(const PyBaselines &baselines,
   }
   return vis;
   }
+
+constexpr auto dirty2vis_DS = R"""(
+Converts a dirty image into a 1D array of visibilities.
+
+Parameters
+==========
+baselines: Baselines
+    the Baselines object
+gconf: GridderConf
+    the GridderConf object to be used
+    (used to optimize the ordering of the indices)
+idx: np.array((nvis,), dtype=np.uint32)
+    the indices for the visibilities to be computed
+dirty: np.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
+    dirty image
+    Its data type determines the precision in which the calculation is carried
+    out.
+wgt: np.array((nvis,), same dtype as `dirty`, optional
+    If present, visibilities are multiplied by the corresponding entries.
+do_wstacking: bool
+    if True, the full improved w-stacking algorithm is carried out, otherwise
+    the w values are assumed to be zero.
+
+Returns
+=======
+np.array((nvis,), dtype=complex of same precision as `dirty`.)
+    The visibility data
+)""";
 py::array Pydirty2vis(const PyBaselines &baselines,
   const PyGridderConfig &gconf, const py::array &idx, const py::array &dirty,
   const py::object &wgt, bool do_wstacking)
@@ -779,6 +838,43 @@ template<typename T> py::array ms2dirty2(const py::array &uvw_,
   return dirty;
   }
 
+constexpr auto ms2dirty_DS = R"""(
+Converts an MS object to dirty image.
+
+Parameters
+==========
+uvw: np.array((nrows, 3), dtype=np.float64)
+    UVW coordinates from the measurement set
+freq: np.array((nchan,), dtype=np.float64)
+    channel frequencies
+ms: np.array((nrows, nchan,), dtype=np.complex64 or np.complex128)
+    the input measurement set data.
+    Its data type determines the precision in which the calculation is carried
+    out.
+wgt: np.array((nrows, nchan), float with same precision as `ms`), optional
+    If present, its values are multiplied to the output
+npix_x, npix_y: int
+    dimensions of the dirty image
+pixsize_x, pixsize_y: float
+    angular pixel size (in radians) of the dirty image
+epsilon: float
+    accuracy at which the computation should be done. Must be larger than 2e-13.
+    If `ms` has type np.complex64, it must be larger than 1e-5.
+do_wstacking: bool
+    if True, the full improved w-stacking algorithm is carried out, otherwise
+    the w values are assumed to be zero.
+nthreads: int
+    number of threads to use for the calculation
+verbosity: int
+    0: no output
+    1: some output
+    2: detailed output
+
+Returns
+=======
+np.array((nxdirty, nydirty), dtype=float of same precision as `ms`)
+    the dirty image
+)""";
 py::array Pyms2dirty(const py::array &uvw,
   const py::array &freq, const py::array &ms, const py::object &wgt,
   size_t npix_x, size_t npix_y, double pixsize_x, double pixsize_y, double epsilon,
@@ -816,6 +912,41 @@ template<typename T> py::array dirty2ms2(const py::array &uvw_,
   return ms;
   }
 
+constexpr auto dirty2ms_DS = R"""(
+Converts a dirty image to an MS object.
+
+Parameters
+==========
+uvw: np.array((nrows, 3), dtype=np.float64)
+    UVW coordinates from the measurement set
+freq: np.array((nchan,), dtype=np.float64)
+    channel frequencies
+dirty: np.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
+    dirty image
+    Its data type determines the precision in which the calculation is carried
+    out.
+wgt: np.array((nrows, nchan), same dtype as `dirty`), optional
+    If present, its values are multiplied to the output
+pixsize_x, pixsize_y: float
+    angular pixel size (in radians) of the dirty image
+epsilon: float
+    accuracy at which the computation should be done. Must be larger than 2e-13.
+    If `dirty` has type np.float32, it must be larger than 1e-5.
+do_wstacking: bool
+    if True, the full improved w-stacking algorithm is carried out, otherwise
+    the w values are assumed to be zero.
+nthreads: int
+    number of threads to use for the calculation
+verbosity: int
+    0: no output
+    1: some output
+    2: detailed output
+
+Returns
+=======
+np.array((nrows, nchan,), dtype=complex of same precision as `dirty`)
+    the measurement set data.
+)""";
 py::array Pydirty2ms(const py::array &uvw,
   const py::array &freq, const py::array &dirty, const py::object &wgt,
   double pixsize_x, double pixsize_y, double epsilon,
@@ -909,15 +1040,15 @@ PYBIND11_MODULE(nifty_gridder, m)
     "grid"_a, "wgt"_a=None);
   m.def("get_correlations", &Pyget_correlations<double>, "baselines"_a, "gconf"_a,
     "idx"_a, "du"_a, "dv"_a, "wgt"_a=None);
-  m.def("vis2dirty",&Pyvis2dirty, "baselines"_a, "gconf"_a,
+  m.def("vis2dirty",&Pyvis2dirty, vis2dirty_DS, "baselines"_a, "gconf"_a,
     "idx"_a, "vis"_a, "wgt"_a=None, "do_wstacking"_a=false);
-  m.def("dirty2vis",&Pydirty2vis, "baselines"_a, "gconf"_a,
+  m.def("dirty2vis",&Pydirty2vis, "baselines"_a, dirty2vis_DS, "gconf"_a,
     "idx"_a, "dirty"_a, "wgt"_a=None, "do_wstacking"_a=false);
 
-  m.def("ms2dirty",&Pyms2dirty,"uvw"_a,"freq"_a,"ms"_a,
-    "wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
+  m.def("ms2dirty", &Pyms2dirty, ms2dirty_DS, "uvw"_a, "freq"_a, "ms"_a,
+    "wgt"_a=None, "npix_x"_a, "npix_y"_a, "pixsize_x"_a, "pixsize_y"_a,
+    "epsilon"_a, "do_wstacking"_a=false, "nthreads"_a=1, "verbosity"_a=0);
+  m.def("dirty2ms", &Pydirty2ms, dirty2ms_DS, "uvw"_a, "freq"_a, "dirty"_a,
+    "wgt"_a=None, "pixsize_x"_a, "pixsize_y"_a, "epsilon"_a,
     "do_wstacking"_a=false, "nthreads"_a=1, "verbosity"_a=0);
-  m.def("dirty2ms",&Pydirty2ms,"uvw"_a,"freq"_a,"dirty"_a,
-    "wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a, "do_wstacking"_a=false, "nthreads"_a=1,
-    "verbosity"_a=0);
   }

test.py

@@ -7,7 +7,6 @@ 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,
@@ -45,13 +44,103 @@ def _dft(uvw, vis, conf, apply_w):
 
 
 def _l2error(a, b):
-    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**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):
+    speedoflight = 299792458.
+    x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]],
+        indexing='ij')
+    x *= xpixsize
+    y *= ypixsize
+    res = np.zeros((nxdirty, nydirty))
+    eps = x**2+y**2
+    if apply_w:
+        nm1 = -eps/(np.sqrt(1.-eps)+1.)
+        n = nm1+1
+    else:
+        nm1 = 0.
+        n = 1.
+    for row in range(ms.shape[0]):
+        for chan in range(ms.shape[1]):
+            phase = (freq[chan]/speedoflight *
+                     (x*uvw[row, 0] + y*uvw[row, 1] - uvw[row, 2]*nm1))
+            if wgt is None:
+                res += (ms[row, chan]*np.exp(2j*np.pi*phase)).real
+            else:
+                res += (ms[row, chan]*wgt[row, chan]*np.exp(2j*np.pi*phase)).real
+    return res/n
+
+
+@pmp("nxdirty", (30, 128))
+@pmp("nydirty", (128, 250))
+@pmp("nrow", (2, 27))
+@pmp("nchan", (1, 5))
+@pmp("epsilon", (1e-3, 1e-5, 1e-7, 5e-13))
+@pmp("singleprec", (True, False))
+@pmp("wstacking", (True, False))
+@pmp("use_wgt", (True, False))
+@pmp("nthreads", (1, 2))
+def test_adjointness_ms2dirty(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking, use_wgt, nthreads):
+    if singleprec and epsilon < 5e-6:
+        return
+    np.random.seed(42)
+    pixsizex = np.pi/180/60/nxdirty*0.2398
+    pixsizey = np.pi/180/60/nxdirty
+    speedoflight, f0 = 3e8, 1e9
+    freq = f0 + np.arange(nchan)*(f0/nchan)
+    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsizey*f0/speedoflight)
+    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
+    wgt = np.random.rand(nrow, nchan) if use_wgt else None
+    dirty = np.random.rand(nxdirty, nydirty)-0.5
+    if singleprec:
+        ms = ms.astype("c8")
+        dirty = dirty.astype("f4")
+        if wgt is not None:
+            wgt = wgt.astype("f4")
+    dirty2 = ng.ms2dirty(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex,
+                         pixsizey, epsilon, wstacking, nthreads, 2).astype("f8")
+    ms2 = ng.dirty2ms(uvw, freq, dirty, wgt, pixsizex, pixsizey, epsilon,
+                      wstacking, nthreads, 0).astype("c16")
+    tol = 5e-5 if singleprec else 5e-13
+    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty2, dirty), rtol=tol)
+
+
+@pmp('nxdirty', [16, 64])
+@pmp('nydirty', [64])
+@pmp("nrow", (2, 27))
+@pmp("nchan", (1, 5))
+@pmp("epsilon", (1e-2, 1e-3, 5e-5, 1e-7, 5e-13))
+@pmp("singleprec", (True, False))
+@pmp("wstacking", (True, False))
+@pmp("use_wgt", (False, True))
+@pmp("nthreads", (1, 2))
+@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
+    np.random.seed(40)
+    pixsizex = fov*np.pi/180/nxdirty
+    pixsizey = fov*np.pi/180/nydirty*1.1
+    speedoflight, f0 = 3e8, 1e9
+    freq = f0 + np.arange(nchan)*(f0/nchan)
+    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsizex*f0/speedoflight)
+    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
+    wgt = np.random.rand(nrow, nchan) if use_wgt else None
+    if singleprec:
+        ms = ms.astype("c8")
+        if wgt is not None:
+            wgt = wgt.astype("f4")
+    dirty = ng.ms2dirty(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex,
+                        pixsizey, epsilon, wstacking, nthreads, 0).astype("f8")
+    ref = explicit_gridder(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex, pixsizey, wstacking)
+    assert_allclose(_l2error(dirty, ref), 0, atol=epsilon)
 
 
 @pmp("nxdirty", (128, 300))
 @pmp("nydirty", (128, 250))
-@pmp("nrow", (1, 10, 10000))
-@pmp("nchan", (1, 10, 100))
+@pmp("nrow", (1, 100))
+@pmp("nchan", (1, 10))
 @pmp("epsilon", (1e-2, 1e-7, 2e-13))
 def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
     np.random.seed(42)
@@ -68,8 +157,8 @@ def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
 
 @pmp("nxdirty", (128, 300))
 @pmp("nydirty", (128, 250))
-@pmp("nrow", (1, 10, 10000))
-@pmp("nchan", (1, 10, 100))
+@pmp("nrow", (1, 127))
+@pmp("nchan", (1, 17))
 @pmp("epsilon", (1e-2, 1e-7, 2e-13))
 def test_adjointness_wgridding(nxdirty, nydirty, nrow, nchan, epsilon):
     np.random.seed(42)
@@ -82,40 +171,10 @@ def test_adjointness_wgridding(nxdirty, nydirty, nrow, nchan, epsilon):
                     rtol=2e-13)
 
 
-@pmp("nxdirty", (128, 300))
-@pmp("nydirty", (128, 250))
-@pmp("nrow", (1, 10, 10000))
-@pmp("nchan", (1, 10, 100))
-@pmp("epsilon", (1e-2, 1e-3, 5e-5, 1e-6, 1e-7, 5e-13))
-@pmp("singleprec", (True, False))
-@pmp("wstacking", (True, False))
-def test_adjointness_wgridding_highlevel(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking):
-    if singleprec and epsilon < 5e-6:
-        return
-    np.random.seed(42)
-    pixsizex = np.pi/180/60/nxdirty*0.2398
-    pixsizey = np.pi/180/60/nxdirty
-    speedoflight, f0 = 3e8, 1e9
-    freq = f0 + np.arange(nchan)*(f0/nchan)
-    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsizey*f0/speedoflight)
-    vis = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
-    dirty = np.random.rand(nxdirty, nydirty)-0.5
-    if singleprec:
-        vis = vis.astype("c8")
-        dirty = dirty.astype("f4")
-    dirty2 = ng.ms2dirty(uvw, freq, vis, None, nxdirty, nydirty, pixsizex,
-                         pixsizey, epsilon, wstacking, 1, 0).astype("f8")
-    vis2 = ng.dirty2ms(uvw, freq, dirty, None, pixsizex, pixsizey, epsilon,
-                       wstacking, 1, 0).astype("c16")
-    tol = 5e-5 if singleprec else 5e-13
-    assert_allclose(np.vdot(vis, vis2).real, np.vdot(dirty2, dirty),
-                    rtol=tol)
-
-
 @pmp("nxdirty", (128,))
 @pmp("nydirty", (128,))
-@pmp("nrow", (10000,))
-@pmp("nchan", (100,))
+@pmp("nrow", (1000,))
+@pmp("nchan", (10,))
 @pmp("epsilon", (2e-13,))
 def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
     np.random.seed(42)
@@ -200,8 +259,8 @@ def test_pickling():
 
 @pmp("nxdirty", (128, 300))
 @pmp("nydirty", (128, 250))
-@pmp("nrow", (1, 10, 10000))
-@pmp("nchan", (1, 10, 100))
+@pmp("nrow", (1, 39))
+@pmp("nchan", (1, 70))
 @pmp("epsilon", (1e-2, 1e-7, 2e-13))
 @pmp("weight", (True, False))
 def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
@@ -219,7 +278,7 @@ def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
 
 @pmp("nxdirty", (300,))
 @pmp("nydirty", (250,))
-@pmp("nrow", (10, 10000))
+@pmp("nrow", (10, 100))
 @pmp("nchan", (1, 10))
 @pmp("epsilon", (1e-2, 1e-7))
 @pmp("du", (-2, 0, 1, 4))
@@ -246,7 +305,7 @@ def test_correlations(nxdirty, nydirty, nrow, nchan, epsilon, du, dv, weight):
         assert_allclose(np.zeros_like(y1), y1.imag)
 
 
-@pmp("nrow", (1, 1000))
+@pmp("nrow", (1, 100))
 @pmp("nchan", (1, 10))
 @pmp("epsilon", (1e-2, 1e-7, 2e-12))
 @pmp("weight", (True, False))
@@ -282,7 +341,7 @@ def test_against_dft(nxdirty, nydirty, epsilon, nchan, nrow):
 
 @pmp('nxdirty', [16, 64])
 @pmp('nydirty', [64])
-@pmp("nrow", (10, 100, 1000))
+@pmp("nrow", (10, 100))
 @pmp("nchan", (1, 10))
 @pmp("fov", (1,))
 def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
@@ -321,7 +380,7 @@ def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
 
 @pmp('nxdirty', [16, 64])
 @pmp('nydirty', [64])
-@pmp("nrow", (10, 100, 1000))
+@pmp("nrow", (10, 100))
 @pmp("nchan", (1,))
 @pmp("fov", (50,))
 def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
@@ -353,7 +412,7 @@ def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
 
 @pmp('nxdirty', [16, 64])
 @pmp('nydirty', [64])
-@pmp("nrow", (10, 100, 1000))
+@pmp("nrow", (10, 100))
 @pmp("nchan", (1, 10))
 @pmp("fov", (1, 10))
 @pmp("epsilon", (1e-2, 1e-7, 2e-12))
@@ -381,7 +440,7 @@ def test_wgridder_against_wdft(nxdirty, nydirty, nchan, nrow, fov, epsilon):
 
 @pmp('nxdirty', (32, 16))
 @pmp('nydirty', (32, 64))
-@pmp("nrow", (100, 1000))
+@pmp("nrow", (100, 127))
 @pmp("nchan", (1, 5))
 @pmp('epsilon', (1e-3, 1e-6))
 def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):