separate out kernel and correction factor code

nifty_gridder.cc

@@ -316,28 +316,56 @@ template<typename T> void hartley2_2D(const pyarr_c<T> &in, pyarr_c<T> &out)
   }
   }
 
+template<typename T> class EC_Kernel
+  {
+  protected:
+    size_t supp;
+    T beta, emb_;
+
+  public:
+    EC_Kernel(size_t supp_)
+      : supp(supp_), beta(2.3*supp_), emb_(exp(-beta)) {}
+    T operator()(T v) const { return exp(beta*(sqrt(T(1)-v*v)-T(1))); }
+    T val_no_emb(T v) const { return exp(beta*sqrt(T(1)-v*v)); }
+    T emb() const { return emb_; }
+  };
+
+template<typename T> class EC_Kernel_with_correction: public EC_Kernel<T>
+  {
+  protected:
+    static constexpr T pi = T(3.141592653589793238462643383279502884197L);
+    int p;
+    vector<T> x, wgt, psi;
+    using EC_Kernel<T>::supp;
+
+  public:
+    using EC_Kernel<T>::operator();
+    EC_Kernel_with_correction(size_t supp_)
+      : EC_Kernel<T>(supp_), p(int(1.5*supp_+2))
+      {
+      legendre_prep(2*p,x,wgt);
+      psi=x;
+      for (auto &v:psi)
+        v=operator()(v);
+      }
+    T corfac(T v) const
+      {
+      T tmp=0;
+      for (int i=0; i<p; ++i)
+        tmp += wgt[i]*psi[i]*cos(pi*supp*v*x[i]);
+      return T(1)/(supp*tmp);
+      }
+  };
 /* Compute correction factors for the ES gridding kernel
    This implementation follows eqs. (3.8) to (3.10) of Barnett et al. 2018 */
 vector<double> correction_factors (size_t n, size_t nval, size_t supp)
   {
-  constexpr double pi = 3.141592653589793238462643383279502884197;
-  auto beta = 2.3*supp;
-  auto p = int(1.5*supp+2);
-  double alpha = pi*supp/n;
-  vector<double> x, wgt;
-  legendre_prep(2*p,x,wgt);
-  auto psi = x;
-  for (auto &v:psi)
-    v = exp(beta*(sqrt(1-v*v)-1.));
+  EC_Kernel_with_correction<double> kernel(supp);
   vector<double> res(nval);
+  double xn = 1./n;
 #pragma omp parallel for schedule(static) num_threads(nthreads)
   for (size_t k=0; k<nval; ++k)
-    {
-    double tmp=0;
-    for (int i=0; i<p; ++i)
-      tmp += wgt[i]*psi[i]*cos(alpha*k*x[i]);
-    res[k] = 1./(supp*tmp);
-    }
+    res[k] = kernel.corfac(k*xn);
   return res;
   }
 
@@ -1498,7 +1526,7 @@ cout << "data w range: " << wmin << " to " << wmax << endl;
 
   double w_eps=1e-7; // FIXME
   auto w_supp = get_supp(w_eps);
-  auto beta=2.3*w_supp;
+  EC_Kernel_with_correction<T> kernel(w_supp);
   wmin -= 0.5*w_supp*dw;
   wmax += 0.5*w_supp*dw;
   size_t nplanes = size_t((wmax-wmin)/dw)+2;
@@ -1533,7 +1561,7 @@ cout << "working on w plane #" << iw << " containing " << nvis_plane[iw]
         {
 myassert(cnt<nvis_plane[iw],"must not happen");
         double x=min(1.,2./(w_supp*dw)*abs(wcur-wval[ipart]));
-        vis_loc[cnt] = vis[ipart]*exp(beta*(sqrt(1.-x*x)-1.));
+        vis_loc[cnt] = vis[ipart]*kernel(x);
         idx_loc[cnt] = idx[ipart];
         ++cnt;
         }
@@ -1549,14 +1577,6 @@ myassert(cnt==nvis_plane[iw],"must not happen 2");
   py::gil_scoped_release release;
 cout << "applying correction for gridding in w direction" << endl;
 
-  constexpr double pi = 3.141592653589793238462643383279502884197;
-  auto beta = 2.3*w_supp;
-  auto p = int(1.5*w_supp+2);
-  vector<double> x, wgt;
-  legendre_prep(2*p,x,wgt);
-  auto psi = x;
-  for (auto &v:psi)
-    v = exp(beta*(sqrt(1-v*v)-1.));
 #pragma omp parallel num_threads(nthreads)
 {
 #pragma omp for schedule(static)
@@ -1573,10 +1593,7 @@ cout << "applying correction for gridding in w direction" << endl;
 #else
       auto nm1 = (-fx-fy)/(sqrt(1.-fx-fy)+1.);
 #endif
-      double fct = 0.;
-      for (int ix=0; ix<p; ++ix)
-        fct += wgt[ix]*psi[ix]*cos(pi*w_supp*nm1*dw*x[ix]);
-      fct = 1./(w_supp*fct);
+      T fct = kernel.corfac(nm1*dw);
       size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
       accum[ny_dirty*i+j]*=fct;
       if ((i>0)&&(i<i2))