Merge branch 'wstacking' into 'master'

Wstacking

See merge request !7

nifty_gridder.cc

@@ -412,6 +412,24 @@ template<typename T> class Baselines
     np.array((nvis,), dtype=np.complex)
         The visibility data for the index array
     )""";
+
+    pyarr_c<T> ind2effectiveuvw(const pyarr_c<uint32_t> &idx_) const
+    {
+      checkArray(idx_, "idx", {0});
+      size_t nvis = size_t(idx_.shape(0));
+      auto idx = idx_.template unchecked<1>();
+      auto res_=makeArray<T>({nvis, 3});
+      auto res = res_.template mutable_unchecked<2>();
+      for (size_t i=0; i<nvis; i++)
+        {
+          auto uvw = effectiveCoord(idx(i));
+          res(i,0) = uvw.u;
+          res(i,1) = uvw.v;
+          res(i,2) = uvw.w;
+        }
+      return res_;
+    }
+
     template<typename T2> pyarr_c<T2> ms2vis(const pyarr<T2> &ms_,
       const pyarr_c<uint32_t> &idx_) const
       {
@@ -526,6 +544,24 @@ np.array((nxdirty, nydirty), dtype=np.float64)
     the image with the taper applied
 )""";
 
+constexpr auto apply_wscreen_DS = R"""(
+Applies the w screen to an image
+
+Parameters
+==========
+dirty: nd.array((nxdirty, nydirty), dtype=np.complex128)
+    the image
+w : float
+    the w value to use
+adjoint: bool
+    if True, apply the complex conjugate of the w screen
+
+Returns
+=======
+np.array((nxdirty, nydirty), dtype=np.complex128)
+    the image with the w screen applied
+)""";
+
 constexpr auto GridderConfig_DS = R"""(
 Class storing information related to the gridding/degridding process.
 
@@ -552,6 +588,15 @@ template<typename T> class GridderConfig
     T beta;
     vector<T> cfu, cfv;
 
+    complex<T> wscreen(double x, double y, double w, bool adjoint) const
+      {
+      constexpr double pi = 3.141592653589793238462643383279502884197;
+      double n = cos(sqrt(x+y)), xn = 1./n;
+      double phase = 2*pi*w*(n-1);
+      if (adjoint) phase *= -1;
+      return complex<T>(cos(phase)*xn, sin(phase)*xn);
+      }
+
   public:
     GridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
       double pixsize_x, double pixsize_y)
@@ -716,6 +761,43 @@ 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; // lower left
+          size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
+          if ((i>0)&&(i<i2))
+            {
+            res[ny_dirty*i2+j] = dirty[ny_dirty*i2+j]*ws; // lower right
+            if ((j>0)&&(j<j2))
+              res[ny_dirty*i2+j2] = dirty[ny_dirty*i2+j2]*ws; // upper right
+            }
+          if ((j>0)&&(j<j2))
+            res[ny_dirty*i+j2] = dirty[ny_dirty*i+j2]*ws; // upper left
+          }
+        }
+}
+      }
+    return res_;
+    }
   };
 
 template<typename T, typename T2=complex<T>> class Helper
@@ -1464,6 +1546,7 @@ PYBIND11_MODULE(nifty_gridder, m)
     .def ("Nchannels",&Baselines<double>::Nchannels)
     .def ("ms2vis",&Baselines<double>::ms2vis<complex<double>>,
       Baselines<double>::ms2vis_DS, "ms"_a, "idx"_a)
+    .def ("ind2effectiveuvw",&Baselines<double>::ind2effectiveuvw, "idx"_a)
     .def ("vis2ms",&Baselines<double>::vis2ms<complex<double>>,
       Baselines<double>::vis2ms_DS, "vis"_a, "idx"_a, "ms_in"_a=None);
   py::class_<GridderConfig<double>> (m, "GridderConfig", GridderConfig_DS)
@@ -1484,6 +1567,8 @@ PYBIND11_MODULE(nifty_gridder, m)
     .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,
+      apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)
 
     // pickle support
     .def(py::pickle(

profile.py

+import nifty_gridder as ng
+from time import time
+import numpy as np
+
+
+def _wscreen(npix, dst, w):
+    dc = (np.linspace(start=-npix/2, stop=npix/2 - 1, num=npix)*dst)**2
+    ls = np.broadcast_to(dc, (dc.shape[0],)*2)
+    theta = np.sqrt(ls + ls.T)
+    n = np.cos(theta)
+    wscreen = np.exp(2*np.pi*1j*w*(n - 1))/n
+    return wscreen
+
+
+def time_op(func, x, ntries=5):
+    t0 = time()
+    for ii in range(ntries):
+        print(ii)
+        func(x)
+    return (time() - t0)/ntries
+
+
+if __name__ == '__main__':
+    ng.set_nthreads(4)
+    ntries = 20
+    nx = 2048
+    dx = 12
+    w = 1000.2
+    ny, dy = nx, dx
+    conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy)
+    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()*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: _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

@@ -155,3 +155,28 @@ def test_pickling():
     assert_(gc.Pixsize_y() == gc2.Pixsize_y())
 
 
+def _wscreen(npix, dst, w):
+    dc = (np.linspace(start=-npix/2, stop=npix/2 - 1, num=npix)*dst)**2
+    ls = np.broadcast_to(dc, (dc.shape[0],)*2)
+    theta = np.sqrt(ls + ls.T)
+    n = np.cos(theta)
+    wscreen = np.exp(2*np.pi*1j*w*(n - 1))/n
+    return wscreen
+
+
+@pmp('nx', [4, 18, 54])
+@pmp('dx', [1., 0.13, 132])
+@pmp('w', [0, 10, 8489])
+def test_wstacking(nx, dx, w):
+    ny, dy = nx, dx
+    conf = ng.GridderConfig(nx, ny, 1e-7, dx, dy)
+    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.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_c(conf.apply_wscreen(fld, w, False))
+    np.testing.assert_allclose(res0, res1)