cleanup

nifty_gridder.cc

@@ -34,6 +34,10 @@ namespace py = pybind11;
 
 namespace {
 
+//
+// basic utilities
+//
+
 void myassert(bool cond, const char *msg)
   {
   if (cond) return;
@@ -69,6 +73,18 @@ template<typename I> inline int bits_needed (I arg)
   if (arg==1) return 0;
   return ilog2(arg-1)+1;
   }
+
+/*! Returns the remainder of the division \a v1/v2.
+    The result is non-negative.
+    \a v1 can be positive or negative; \a v2 must be positive. */
+template<typename T> inline T fmodulo (T v1, T v2)
+  {
+  if (v1>=0)
+    return (v1<v2) ? v1 : fmod(v1,v2);
+  T tmp=fmod(v1,v2)+v2;
+  return (tmp==v2) ? T(0) : tmp;
+  }
+
 //
 // Utilities for converting between Cartesian coordinates and Peano index
 //
@@ -182,8 +198,7 @@ template<typename It, typename T2, typename Comp>
   T2 num=end-begin;
   idx.resize(num);
   for (T2 i=0; i<num; ++i) idx[i] = i;
-//  sort (idx.begin(),idx.end(),IdxComp__<It,Comp>(begin,comp));
-  stable_sort (idx.begin(),idx.end(),IdxComp__<It,Comp>(begin,comp));
+  sort (idx.begin(),idx.end(),IdxComp__<It,Comp>(begin,comp));
   }
 
 /*! Performs an indirect sort on the supplied iterator range and returns in
@@ -197,17 +212,6 @@ template<typename It, typename T2> inline void buildIndex (It begin, It end,
   buildIndex(begin,end,idx,less<T>());
   }
 
-/*! Returns the remainder of the division \a v1/v2.
-    The result is non-negative.
-    \a v1 can be positive or negative; \a v2 must be positive. */
-template<typename T> inline T fmodulo (T v1, T v2)
-  {
-  if (v1>=0)
-    return (v1<v2) ? v1 : fmod(v1,v2);
-  T tmp=fmod(v1,v2)+v2;
-  return (tmp==v2) ? T(0) : tmp;
-  }
-
 //
 // Utilities for Gauss-Legendre quadrature
 //
@@ -277,41 +281,8 @@ template<typename T>
   using pyarr = py::array_t<T>;
 template<typename T>
   using pyarr_c = py::array_t<T, py::array::c_style | py::array::forcecast>;
-using a_u32_c = pyarr_c<uint32_t>;
-using a_d_c = pyarr_c<double>;
-using a_f_c = pyarr_c<float>;
-using a_cd_c = pyarr_c<complex<double>>;
-
-a_u32_c peanoindex(const a_d_c &uv_, int nu, int nv)
-  {
-  myassert(uv_.ndim()==2, "uv array must be 2D");
-  myassert(uv_.shape(1)==2, "uv.shape[1] must be 2");
-  int nvis = uv_.shape(0);
-  auto uv = uv_.data();
-
-  int npmax = max(nu, nv);
-  int nbits = 0;
-  for (int istart = npmax-1; istart!=0; istart>>=1, ++nbits);
-  vector<int> ipeano(nvis);
-  for (int i=0; i<nvis; ++i)
-    {
-    auto u = fmodulo(uv[2*i], 1.)*nu;
-    auto iu = min(nu-1, int(u));
-    auto v = fmodulo(uv[2*i+1], 1.)*nv;
-    auto iv = min(nv-1, int(v));
-    ipeano[i] = morton2peano2D_32(coord2morton2D_32(iu,iv),nbits);
-    }
-  vector<int> newind;
-  buildIndex(ipeano.begin(), ipeano.end(), newind);
-  int odim[] = {nvis};
-  a_u32_c res(odim);
-  auto iout = res.mutable_data();
-  for (int i=0; i<nvis; ++i)
-    iout[i] = newind[i];
-  return res;
-  }
 
-int get_w(double epsilon)
+size_t get_w(double epsilon)
   {
   static const vector<double> maxmaperr { 1e8, 0.32, 0.021, 6.2e-4,
     1.08e-5, 1.25e-7, 8.25e-10, 5.70e-12, 1.22e-13, 2.48e-15, 4.82e-17,
@@ -324,14 +295,14 @@ int get_w(double epsilon)
   throw runtime_error("requested epsilon too small - minimum is 2e-13");
   }
 
-a_d_c complex2hartley (const a_cd_c &grid_)
+pyarr_c<double> complex2hartley (const pyarr_c<complex<double>> &grid_)
   {
   myassert(grid_.ndim()==2, "grid array must be 2D");
   int nu = grid_.shape(0), nv = grid_.shape(1);
   auto grid = grid_.data();
 
   int odim[] = {nu,nv};
-  a_d_c res(odim);
+  pyarr_c<double> res(odim);
   auto grid2 = res.mutable_data();
   for (int u=0; u<nu; ++u)
     {
@@ -348,15 +319,14 @@ a_d_c complex2hartley (const a_cd_c &grid_)
   return res;
   }
 
-
-a_cd_c hartley2complex (const a_d_c &grid_)
+pyarr_c<complex<double>> hartley2complex (const pyarr_c<double> &grid_)
   {
   myassert(grid_.ndim()==2, "grid array must be 2D");
   int nu = grid_.shape(0), nv = grid_.shape(1);
   auto grid = grid_.data();
 
   int odim[] = {nu,nv};
-  a_cd_c res(odim);
+  pyarr_c<complex<double>> res(odim);
   auto grid2 = res.mutable_data();
   for (int u=0; u<nu; ++u)
     {
@@ -377,7 +347,7 @@ a_cd_c hartley2complex (const a_d_c &grid_)
 
 /* Compute correction factors for the ES gridding kernel
    This implementation follows eqs. (3.8) to (3.10) of Barnett et al. 2018 */
-a_d_c correction_factors (size_t n, size_t nval, size_t w)
+pyarr_c<double> correction_factors (size_t n, size_t nval, size_t w)
   {
   constexpr double pi = 3.141592653589793238462643383279502884197;
   auto beta = 2.3*w;
@@ -389,7 +359,7 @@ a_d_c correction_factors (size_t n, size_t nval, size_t w)
   for (auto &v:psi)
     v = exp(beta*(sqrt(1-v*v)-1.));
   int odim[] = {int(nval)};
-  a_d_c res(odim);
+  pyarr_c<double> res(odim);
   auto val = res.mutable_data();
   for (size_t k=0; k<nval; ++k)
     {
@@ -516,7 +486,7 @@ class GridderConfig
       vector<size_t> newind;
       buildIndex(peano.begin(), peano.end(), newind);
       peano=vector<size_t>(); // deallocate
-      int odim[] = {idx.shape(0)};
+      int odim[] = {int(idx.shape(0))};
       pyarr_c<uint32_t> res(odim);
       auto iout = res.mutable_data();
       for (int i=0; i<idx.shape(0); ++i)
@@ -674,16 +644,16 @@ pyarr_c<double> ms2grid(const Baselines &baselines, const GridderConfig &gconf,
   myassert(idx_.ndim()==1, "idx array must be 1D");
   myassert(data_.ndim()==2, "data must be 2D");
   auto data=data_.data();
-  myassert(data_.shape(0)==baselines.Nrows(), "bad data dimension");
-  myassert(data_.shape(1)==baselines.Nchannels(), "bad data dimension");
+  myassert(data_.shape(0)==int(baselines.Nrows()), "bad data dimension");
+  myassert(data_.shape(1)==int(baselines.Nchannels()), "bad data dimension");
   int nvis = idx_.shape(0);
   auto idx = idx_.data();
 
   size_t nu=gconf.Nu(), nv=gconf.Nv();
-  int odim[] = {nu, nv};
+  int odim[] = {int(nu), int(nv)};
   pyarr_c<complex<double>> res(odim);
   auto grid = res.mutable_data();
-  for (int i=0; i<nu*nv; ++i) grid[i] = 0.;
+  for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
 
 #pragma omp parallel
 {
@@ -721,13 +691,13 @@ pyarr_c<complex<double>> grid2ms(const Baselines &baselines,
   auto grid_ = hartley2complex(grid0_);
   auto grid = grid_.data();
   myassert(grid_.ndim()==2, "data must be 2D");
-  myassert(grid_.shape(0)==nu, "bad grid dimension");
-  myassert(grid_.shape(1)==nv, "bad grid dimension");
+  myassert(grid_.shape(0)==int(nu), "bad grid dimension");
+  myassert(grid_.shape(1)==int(nv), "bad grid dimension");
   int nvis = idx_.shape(0);
   auto idx = idx_.data();
 
   int odim[] = {nvis};
-  a_cd_c res(odim);
+  pyarr_c<complex<double>> res(odim);
   auto vis = res.mutable_data();
 
   // Loop over sampling points
@@ -766,7 +736,7 @@ PYBIND11_MODULE(nifty_gridder, m)
   using namespace pybind11::literals;
 
   py::class_<Baselines> (m, "Baselines")
-    .def(py::init<a_d_c, a_d_c>(), "coord"_a, "scaling"_a)
+    .def(py::init<pyarr_c<double>, pyarr_c<double>>(), "coord"_a, "scaling"_a)
     .def("getIndices", &Baselines::getIndices);
   py::class_<GridderConfig> (m, "GridderConfig")
     .def(py::init<size_t, size_t, double, double, double>(),"nxdirty"_a,