From 77a3ecbffd03f36b5ebab2ad7ad2a2001c7381e7 Mon Sep 17 00:00:00 2001
From: Martin Reinecke <martin@mpa-garching.mpg.de>
Date: Wed, 16 Oct 2019 15:54:30 +0200
Subject: [PATCH] add some routines for uncropped dirty images

---
 gridder_cxx.h    | 99 +++++++++++++++++++++++++++++++++++++++++-------
 nifty_gridder.cc | 63 ++++++++++++++++++++++++++++++
 2 files changed, 149 insertions(+), 13 deletions(-)

diff --git a/gridder_cxx.h b/gridder_cxx.h
index 5650324..a4705b5 100644
--- a/gridder_cxx.h
+++ b/gridder_cxx.h
@@ -546,8 +546,8 @@ class GridderConfig
     double eps, psx, psy;
     size_t supp, nsafe, nu, nv;
     double beta;
-    vector<double> cfu, cfv;
     size_t nthreads;
+    vector<double> cfu, cfv;
     double ushift, vshift;
     int maxiu0, maxiv0;
 
@@ -563,6 +563,17 @@ class GridderConfig
       return complex<double>(cos(phase)*xn, sin(phase)*xn);
       }
 
+    vector<double> taper(size_t nd, size_t n) const
+      {
+      auto tmp = correction_factors(nd, n/2+1, supp, nthreads);
+      vector<double> res(n);
+      res[n/2]=tmp[0];
+      res[0]=tmp[n/2];
+      for (size_t i=1; i<n/2; ++i)
+        res[n/2-i] = res[n/2+i] = tmp[i];
+      return res;
+      }
+
   public:
     GridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
       double pixsize_x, double pixsize_y, size_t nthreads_)
@@ -571,7 +582,7 @@ class GridderConfig
         supp(ES_Kernel::get_supp(epsilon)), nsafe((supp+1)/2),
         nu(max(2*nsafe,2*nx_dirty)), nv(max(2*nsafe,2*ny_dirty)),
         beta(ES_Kernel::get_beta(supp)*supp),
-        cfu(nx_dirty), cfv(ny_dirty), nthreads(nthreads_),
+        nthreads(nthreads_), cfu(taper(nu, nx_dirty)), cfv(taper(nv, ny_dirty)),
         ushift(supp*(-0.5)+1+nu), vshift(supp*(-0.5)+1+nv),
         maxiu0((nu+nsafe)-supp), maxiv0((nv+nsafe)-supp)
       {
@@ -580,17 +591,6 @@ class GridderConfig
       myassert(epsilon>0, "epsilon must be positive");
       myassert(pixsize_x>0, "pixsize_x must be positive");
       myassert(pixsize_y>0, "pixsize_y must be positive");
-
-      auto tmp = correction_factors(nu, nx_dirty/2+1, supp, nthreads);
-      cfu[nx_dirty/2]=tmp[0];
-      cfu[0]=tmp[nx_dirty/2];
-      for (size_t i=1; i<nx_dirty/2; ++i)
-        cfu[nx_dirty/2-i] = cfu[nx_dirty/2+i] = tmp[i];
-      tmp = correction_factors(nv, ny_dirty/2+1, supp, nthreads);
-      cfv[ny_dirty/2]=tmp[0];
-      cfv[0]=tmp[ny_dirty/2];
-      for (size_t i=1; i<ny_dirty/2; ++i)
-        cfv[ny_dirty/2-i] = cfv[ny_dirty/2+i] = tmp[i];
       }
     size_t Nxdirty() const { return nx_dirty; }
     size_t Nydirty() const { return ny_dirty; }
@@ -617,6 +617,21 @@ class GridderConfig
           for (size_t j=0; j<ny_dirty; ++j)
             img2(i,j) = img(i,j)*cfu[i]*cfv[j];
       }
+    template<typename T> void apply_taper_full(const mav<const T,2> &img, mav<T,2> &img2, bool divide) const
+      {
+      checkShape(img.shape(), {nu, nv});
+      checkShape(img2.shape(), {nu, nv});
+      auto cfu2 = taper(nu, nu);
+      auto cfv2 = taper(nv, nv);
+      if (divide)
+        for (size_t i=0; i<nu; ++i)
+          for (size_t j=0; j<nv; ++j)
+            img2(i,j) = img(i,j)/(cfu2[i]*cfv2[j]);
+      else
+        for (size_t i=0; i<nu; ++i)
+          for (size_t j=0; j<nv; ++j)
+            img2(i,j) = img(i,j)*cfu2[i]*cfv2[j];
+      }
 
     template<typename T> void grid2dirty_post(const mav<T,2> &tmav, const mav<T,2> &dirty) const
       {
@@ -634,6 +649,24 @@ class GridderConfig
           }
       }
 
+    template<typename T> void grid2dirty_post_full(const mav<T,2> &tmav, const mav<T,2> &dirty) const
+      {
+      checkShape(dirty.shape(), {nu, nv});
+      auto cfu2 = taper(nu, nu);
+      auto cfv2 = taper(nv, nv);
+      for (size_t i=0; i<nu; ++i)
+        for (size_t j=0; j<nv; ++j)
+          {
+          size_t i2 = nu-nu/2+i;
+          if (i2>=nu) i2-=nu;
+          size_t j2 = nv-nv/2+j;
+          if (j2>=nv) j2-=nv;
+          // FIXME: for some reason g++ warns about double-to-float conversion
+          // here, even though there is an explicit cast...
+          dirty(i,j) = tmav(i2,j2)*T(cfu2[i]*cfv2[j]);
+          }
+      }
+
     template<typename T> void grid2dirty(const const_mav<T,2> &grid, const mav<T,2> &dirty) const
       {
       checkShape(grid.shape(), {nu,nv});
@@ -655,6 +688,18 @@ class GridderConfig
       grid2dirty_post(tmp, dirty);
       }
 
+    template<typename T> void grid2dirty_c_full(const mav<const complex<T>,2> &grid, mav<complex<T>,2> &dirty) const
+      {
+      checkShape(grid.shape(), {nu,nv});
+      tmpStorage<complex<T>,2> tmpdat({nu,nv});
+      auto tmp = tmpdat.getMav();
+      constexpr auto sc = ptrdiff_t(sizeof(complex<T>));
+      pocketfft::c2c({nu,nv},{grid.stride(0)*sc,grid.stride(1)*sc},
+        {tmp.stride(0)*sc, tmp.stride(1)*sc}, {0,1}, pocketfft::BACKWARD,
+        grid.data(), tmp.data(), T(1), nthreads);
+      grid2dirty_post(tmp, dirty);
+      }
+
     template<typename T> void dirty2grid_pre(const const_mav<T,2> &dirty, mav<T,2> &grid) const
       {
       checkShape(dirty.shape(), {nx_dirty, ny_dirty});
@@ -672,6 +717,24 @@ class GridderConfig
           grid(i2,j2) = dirty(i,j)*T(cfu[i]*cfv[j]);
           }
       }
+    template<typename T> void dirty2grid_pre_full(const const_mav<T,2> &dirty, mav<T,2> &grid) const
+      {
+      checkShape(dirty.shape(), {nu, nv});
+      checkShape(grid.shape(), {nu, nv});
+      auto cfu2(taper(nu, nu));
+      auto cfv2 = taper(nv, nv);
+      for (size_t i=0; i<nu; ++i)
+        for (size_t j=0; j<nv; ++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;
+          // FIXME: for some reason g++ warns about double-to-float conversion
+          // here, even though there is an explicit cast...
+          grid(i2,j2) = dirty(i,j)*T(cfu2[i]*cfv2[j]);
+          }
+      }
 
     template<typename T> void dirty2grid(const const_mav<T,2> &dirty, mav<T,2> &grid) const
       {
@@ -689,6 +752,16 @@ class GridderConfig
         grid.data(), grid.data(), T(1), nthreads);
       }
 
+    template<typename T> void dirty2grid_c_full(const const_mav<complex<T>,2> &dirty,
+      mav<complex<T>,2> &grid) const
+      {
+      dirty2grid_pre_full(dirty, grid);
+      constexpr auto sc = ptrdiff_t(sizeof(complex<T>));
+      pocketfft::stride_t strides{grid.stride(0)*sc,grid.stride(1)*sc};
+      pocketfft::c2c({nu,nv}, strides, strides, {0,1}, pocketfft::FORWARD,
+        grid.data(), grid.data(), T(1), nthreads);
+      }
+
     void getpix(double u_in, double v_in, double &u, double &v, int &iu0, int &iv0) const
       {
       u=fmod1(u_in*psx)*nu;
diff --git a/nifty_gridder.cc b/nifty_gridder.cc
index c84f07b..a147169 100644
--- a/nifty_gridder.cc
+++ b/nifty_gridder.cc
@@ -335,6 +335,25 @@ class PyGridderConfig: public GridderConfig
         return apply_taper2<double>(img, divide);
       myfail("type matching failed: 'img' has neither type 'f4' nor 'f8'");
       }
+    template<typename T> pyarr<T> apply_taper_full2(const pyarr<T> &img, bool divide) const
+      {
+      auto res = makeArray<T>({nu, nv});
+      auto img2 = make_const_mav<2>(img);
+      auto res2 = make_mav<2>(res);
+      {
+      py::gil_scoped_release release;
+      GridderConfig::apply_taper_full(img2, res2, divide);
+      }
+      return res;
+      }
+    py::array apply_taper_full(const py::array &img, bool divide) const
+      {
+      if (isPytype<float>(img))
+        return apply_taper_full2<float>(img, divide);
+      if (isPytype<double>(img))
+        return apply_taper_full2<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});
@@ -375,6 +394,27 @@ class PyGridderConfig: public GridderConfig
       myfail("type matching failed: 'grid' has neither type 'c8' nor 'c16'");
       }
 
+    template<typename T> pyarr<complex<T>> grid2dirty_c_full2
+      (const pyarr<complex<T>> &grid) const
+      {
+      auto res = makeArray<complex<T>>({nu, nv});
+      auto grid2=make_const_mav<2>(grid);
+      auto res2=make_mav<2>(res);
+      {
+      py::gil_scoped_release release;
+      GridderConfig::grid2dirty_c_full(grid2,res2);
+      }
+      return res;
+      }
+    py::array grid2dirty_c_full(const py::array &grid) const
+      {
+      if (isPytype<complex<float>>(grid))
+        return grid2dirty_c_full2<float>(grid);
+      if (isPytype<complex<double>>(grid))
+        return grid2dirty_c_full2<double>(grid);
+      myfail("type matching failed: 'grid' has neither type 'c8' nor 'c16'");
+      }
+
     template<typename T> pyarr<T> dirty2grid2(const pyarr<T> &dirty) const
       {
       auto dirty2 = make_const_mav<2>(dirty);
@@ -413,6 +453,25 @@ class PyGridderConfig: public GridderConfig
         return dirty2grid_c2<double>(dirty);
       myfail("type matching failed: 'dirty' has neither type 'c8' nor 'c16'");
       }
+    template<typename T> pyarr<complex<T>> dirty2grid_c_full2(const pyarr<complex<T>> &dirty) const
+      {
+      auto dirty2 = make_const_mav<2>(dirty);
+      auto grid = makeArray<complex<T>>({nu, nv});
+      auto grid2=make_mav<2>(grid);
+      {
+      py::gil_scoped_release release;
+      GridderConfig::dirty2grid_c_full(dirty2, grid2);
+      }
+      return grid;
+      }
+    py::array dirty2grid_c_full(const py::array &dirty) const
+      {
+      if (isPytype<complex<float>>(dirty))
+        return dirty2grid_c_full2<float>(dirty);
+      if (isPytype<complex<double>>(dirty))
+        return dirty2grid_c_full2<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
       {
@@ -1056,12 +1115,16 @@ PYBIND11_MODULE(nifty_gridder, m)
     .def("Supp", &PyGridderConfig::Supp)
     .def("apply_taper", &PyGridderConfig::apply_taper, apply_taper_DS,
       "img"_a, "divide"_a=false)
+    .def("apply_taper_full", &PyGridderConfig::apply_taper_full,
+      "img"_a, "divide"_a=false)
      .def("grid2dirty", &PyGridderConfig::grid2dirty,
         grid2dirty_DS, "grid"_a)
     .def("grid2dirty_c", &PyGridderConfig::grid2dirty_c, "grid"_a)
+    .def("grid2dirty_c_full", &PyGridderConfig::grid2dirty_c_full, "grid"_a)
     .def("dirty2grid", &PyGridderConfig::dirty2grid,
        dirty2grid_DS, "dirty"_a)
     .def("dirty2grid_c", &PyGridderConfig::dirty2grid_c, "dirty"_a)
+    .def("dirty2grid_c_full", &PyGridderConfig::dirty2grid_c_full, "dirty"_a)
     .def("apply_wscreen", &PyGridderConfig::apply_wscreen,
       apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)
 
-- 
GitLab