cleanup and tweaks

nifty_gridder.cc

@@ -665,51 +665,6 @@ template<typename T> class ReadHelper: public Helper<T>
       }
   };
 
-template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
-  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
-  const pyarr_c<complex<T>> &vis_)
-  {
-  myassert(idx_.ndim()==1, "idx array must be 1D");
-  myassert(vis_.ndim()==1, "vis must be 1D");
-  auto vis=vis_.data();
-  myassert(vis_.shape(0)==idx_.shape(0), "bad vis dimension");
-  size_t nvis = size_t(idx_.shape(0));
-  auto idx = idx_.data();
-
-  size_t nu=gconf.Nu(), nv=gconf.Nv();
-  auto res = makearray<complex<T>>({nu, nv});
-  auto grid = res.mutable_data();
-  for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
-  T beta = gconf.Beta();
-  size_t w = gconf.W();
-
-#pragma omp parallel
-{
-  WriteHelper<T> hlp(gconf, grid);
-  T emb = exp(-2*beta);
-  const T * RESTRICT ku = hlp.kernel.data();
-  const T * RESTRICT kv = hlp.kernel.data()+w;
-
-  // Loop over sampling points
-#pragma omp for schedule(guided,100)
-  for (size_t ipart=0; ipart<nvis; ++ipart)
-    {
-    UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
-    hlp.prep_write(coord.u, coord.v);
-    auto * RESTRICT ptr = hlp.p0;
-    auto v(vis[ipart]*emb);
-    for (size_t cu=0; cu<w; ++cu)
-      {
-      complex<T> tmp(v*ku[cu]);
-      for (size_t cv=0; cv<w; ++cv)
-        ptr[cv] += tmp*kv[cv];
-      ptr+=hlp.sv;
-      }
-    }
-} // end of parallel region
-  return complex2hartley(res);
-  }
-
 template<typename T> pyarr_c<complex<T>> vis2grid_c(const Baselines<T> &baselines,
   const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
   const pyarr_c<complex<T>> &vis_)
@@ -755,53 +710,10 @@ template<typename T> pyarr_c<complex<T>> vis2grid_c(const Baselines<T> &baseline
   return res;
   }
 
-template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
+template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
   const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
-  const pyarr_c<T> &grid0_)
-  {
-  size_t nu=gconf.Nu(), nv=gconf.Nv();
-  myassert(idx_.ndim()==1, "idx array must be 1D");
-  auto grid_ = hartley2complex(grid0_);
-  auto grid = grid_.data();
-  myassert(grid_.ndim()==2, "data must be 2D");
-  myassert(grid_.shape(0)==int(nu), "bad grid dimension");
-  myassert(grid_.shape(1)==int(nv), "bad grid dimension");
-  size_t nvis = size_t(idx_.shape(0));
-  auto idx = idx_.data();
-
-  auto res = makearray<complex<T>>({nvis});
-  auto vis = res.mutable_data();
-  T beta = gconf.Beta();
-  size_t w = gconf.W();
-
-  // Loop over sampling points
-#pragma omp parallel
-{
-  ReadHelper<T> hlp(gconf, grid);
-  T emb = exp(-2*beta);
-  const T * RESTRICT ku = hlp.kernel.data();
-  const T * RESTRICT kv = hlp.kernel.data()+w;
-
-#pragma omp for schedule(guided,100)
-  for (size_t ipart=0; ipart<nvis; ++ipart)
-    {
-    UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
-    hlp.prep_read(coord.u, coord.v);
-    complex<T> r = 0;
-    auto * RESTRICT ptr = hlp.p0;
-    for (size_t cu=0; cu<w; ++cu)
-      {
-      complex<T> tmp(0);
-      for (size_t cv=0; cv<w; ++cv)
-        tmp += ptr[cv] * kv[cv];
-      r += tmp*ku[cu];
-      ptr += hlp.sv;
-      }
-    vis[ipart] = r*emb;
-    }
-}
-  return res;
-  }
+  const pyarr_c<complex<T>> &vis_)
+  { return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_)); }
 
 template<typename T> pyarr_c<complex<T>> grid2vis_c(const Baselines<T> &baselines,
   const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
@@ -850,6 +762,11 @@ template<typename T> pyarr_c<complex<T>> grid2vis_c(const Baselines<T> &baseline
   return res;
   }
 
+template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
+  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
+  const pyarr_c<T> &grid_)
+  { return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_)); }
+
 template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
   const GridderConfig<T> &gconf, const pyarr_c<bool> &flags_, int chbegin,
   int chend, T wmin, T wmax)
@@ -866,10 +783,11 @@ template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
   constexpr int side=1<<logsquare;
   size_t nbu = (gconf.Nu()+1+side-1) >> logsquare,
          nbv = (gconf.Nv()+1+side-1) >> logsquare;
-  vector<uint32_t> bincnt(nbu*nbv, 0);
-  for (size_t irow=0; irow<nrow; ++irow)
-    for (int ichan=chbegin; ichan<chend; ++ichan)
-      if (!flags[irow*nchan + ichan])
+  vector<uint32_t> acc(nbu*nbv+1, 0);
+  vector<uint32_t> tmp(nrow*(chend-chbegin));
+  for (size_t irow=0, idx=0; irow<nrow; ++irow)
+    for (int ichan=chbegin; ichan<chend; ++ichan, ++idx)
+      if (!flags[irow*nchan+ichan])
         {
         auto uvw = baselines.effectiveCoord(irow, ichan);
         if ((uvw.w>=wmin) && (uvw.w<wmax))
@@ -879,31 +797,19 @@ template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
           gconf.getpix(uvw.u, uvw.v, u, v, iu0, iv0);
           iu0 = (iu0+nsafe)>>logsquare;
           iv0 = (iv0+nsafe)>>logsquare;
-          ++bincnt[nbv*iu0 + iv0];
+          ++acc[nbv*iu0 + iv0 + 1];
+          tmp[idx] = nbv*iu0 + iv0;
           }
         }
-  vector<uint32_t> acc(bincnt.size()+1);
-  acc[0] = 0;
-  for (size_t i=0; i<bincnt.size(); ++i)
-    acc[i+1] = acc[i] + bincnt[i];
+
+  for (size_t i=1; i<acc.size(); ++i)
+    acc[i] += acc[i-1];
   auto res = makearray<uint32_t>({acc.back()});
   auto iout = res.mutable_data();
-  for (size_t irow=0; irow<nrow; ++irow)
-    for (int ichan=chbegin; ichan<chend; ++ichan)
-      if (!flags[irow*nchan + ichan])
-        {
-        auto uvw = baselines.effectiveCoord(irow, ichan);
-        if ((uvw.w>=wmin) && (uvw.w<wmax))
-          {
-          T u, v;
-          int iu0, iv0;
-          gconf.getpix(uvw.u, uvw.v, u, v, iu0, iv0);
-          iu0 = (iu0+nsafe)>>logsquare;
-          iv0 = (iv0+nsafe)>>logsquare;
-          iout[acc[nbv*iu0 + iv0]] = irow*nchan + ichan;
-          ++acc[nbv*iu0 + iv0];
-          }
-        }
+  for (size_t irow=0, idx=0; irow<nrow; ++irow)
+    for (int ichan=chbegin; ichan<chend; ++ichan, ++idx)
+      if (!flags[irow*nchan+ichan])
+        iout[acc[tmp[idx]]++] = irow*nchan+ichan;
   return res;
   }