Adding options to switch kernel at runtime

Author: DiamonDinoia

include/finufft.fh

@@ -20,6 +20,7 @@ c     DEPRECATED: spread_kerpad kept for ABI compatibility, ignored by library
       real*8 upsampfac
       integer spread_thread, maxbatchsize, spread_nthr_atomic
       integer spread_max_sp_size
+      integer spread_kernel
       integer fftw_lock_fun, fftw_unlock_fun, fftw_lock_data
 
       end type

include/finufft/spreadinterp.h

@@ -21,9 +21,9 @@ namespace finufft {
 namespace spreadinterp {
 
 template<typename T>
-FINUFFT_EXPORT_TEST int setup_spreader(finufft_spread_opts &opts, T eps, double upsampfac,
-                                       int kerevalmeth, int debug, int showwarn,
-                                       int spreadinterponly, int dim);
+FINUFFT_EXPORT_TEST int setup_spreader(
+    finufft_spread_opts &opts, T eps, double upsampfac, int kerevalmeth, int debug,
+    int showwarn, int spreadinterponly, int dim, int kernel_type = 0);
 
 int spreadcheck(UBIGINT N1, UBIGINT N2, UBIGINT N3, const finufft_spread_opts &opts);
 template<typename T>

include/finufft_common/kernel.h

@@ -6,6 +6,7 @@
 #include <vector>
 
 #include <finufft_common/constants.h>
+#include <finufft_common/utils.h>
 
 namespace finufft::kernel {
 
@@ -56,14 +57,25 @@ template<class T, class F> std::vector<T> fit_monomials(F &&f, int n, T a, T b)
   return c;
 }
 
-template<typename T> T evaluate_kernel(T x, T beta, T c) {
-  /* ES ("exp sqrt") kernel evaluation at single real argument:
-      phi(x) = exp(beta.(sqrt(1 - (2x/n_s)^2) - 1)),    for |x| < nspread/2
-     related to an asymptotic approximation to the Kaiser--Bessel, itself an
-     approximation to prolate spheroidal wavefunction (PSWF) of order 0.
-     This is the "reference implementation", used by eg finufft/onedim_* 2/17/17.
-     Rescaled so max is 1, Barnett 7/21/24
+template<typename T> T evaluate_kernel(T x, T beta, T c, int kernel_type = 0) {
+  /* Kernel evaluation at single real argument.
+     kernel_type == 0 : ES ("exp sqrt") kernel (default)
+       phi_ES(x) = exp(beta*(sqrt(1 - c*x^2) - 1))
+     kernel_type == 1 : Kaiser--Bessel (KB) kernel
+       phi_KB(x) = I_0(beta*sqrt(1 - c*x^2)) / I_0(beta)
+     Note: `std::cyl_bessel_i` from <cmath> is used for I_0.
+     Rescaled so max is 1.
   */
+  if (kernel_type == 1) {
+    // Kaiser--Bessel (normalized by I0(beta)). Use std::cyl_bessel_i from <cmath>.
+    const T inner        = std::sqrt(T(1) - c * x * x);
+    const T arg          = beta * inner;
+    const double i0_arg  = ::finufft::common::cyl_bessel_i(0, static_cast<double>(arg));
+    const double i0_beta = ::finufft::common::cyl_bessel_i(0, static_cast<double>(beta));
+    return static_cast<T>(i0_arg / i0_beta);
+  }
+
+  // default to ES
   return std::exp(beta * (std::sqrt(T(1) - c * x * x) - T(1)));
 }
 

include/finufft_common/utils.h

@@ -13,6 +13,9 @@ namespace common {
 FINUFFT_EXPORT_TEST void gaussquad(int n, double *xgl, double *wgl);
 std::tuple<double, double> leg_eval(int n, double x);
 
+// Series implementation of the modified Bessel function of the first kind I_nu(x)
+double cyl_bessel_i(double nu, double x) noexcept;
+
 // helper to generate the integer sequence in range [Start, End]
 template<int Offset, typename Seq> struct offset_seq;
 

include/finufft_mod.f90

@@ -21,6 +21,7 @@ module finufft_mod
    real(kind=C_DOUBLE) :: upsampfac
    integer(kind=C_INT) :: spread_thread, maxbatchsize
    integer(kind=C_INT) :: spread_nthr_atomic, spread_max_sp_size
+   integer(kind=C_INT) :: spread_kernel
    integer(kind=C_SIZE_T) :: fftw_lock_fun, fftw_unlock_fun, fftw_lock_data
    !  really, last should be type(C_PTR) :: etc, but fails to print nicely
 

include/finufft_opts.h

@@ -32,6 +32,8 @@ typedef struct finufft_opts { // defaults see finufft_core.cpp:finufft_default_o
   int spread_nthr_atomic; // if >=0, threads above which spreader OMP critical goes
                           // atomic
   int spread_max_sp_size; // if >0, overrides spreader (dir=1) max subproblem size
+  int spread_kernel;      // (dev only) 0:DEFAULT, (do not change), there is no guarantee
+                          // what non-zero values do and behaviour can change anytime
                           // sphinx tag (don't remove): @opts_end
 
   // User can provide their own FFTW planner lock functions for thread safety

include/finufft_spread_opts.h

@@ -26,6 +26,8 @@ typedef struct finufft_spread_opts {
   double ES_beta;
   double ES_halfwidth;
   double ES_c;
+  // Kernel selector: 0 = ES (default), 1 = Kaiser--Bessel (KB)
+  int kernel_type; // default 0
 } finufft_spread_opts;
 
 #endif // FINUFFT_SPREAD_OPTS_H

matlab/finufft.mw

@@ -103,6 +103,9 @@ $     }
 $     else if (strcmp(fname[ifield],"spread_max_sp_size") == 0) {
 $       oc->spread_max_sp_size = (int)round(*mxGetPr(mxGetFieldByNumber(om,idx,ifield)));
 $     }
+$     else if (strcmp(fname[ifield],"spread_kernel") == 0) {
+$       oc->spread_kernel = (int)round(*mxGetPr(mxGetFieldByNumber(om,idx,ifield)));
+$     }
 $     else if (strcmp(fname[ifield],"spreadinterponly") == 0) {
 $       oc->spreadinterponly = (int)round(*mxGetPr(mxGetFieldByNumber(om,idx,ifield)));
 $     }

python/finufft/finufft/_finufft.py

@@ -86,6 +86,7 @@ class FinufftOpts(ctypes.Structure):
                       ('maxbatchsize', c_int),
                       ('spread_nthr_atomic', c_int),
                       ('spread_max_sp_size', c_int),
+                      ('spread_kernel', c_int),
                       ('fftw_lock_fun', c_void_p),
                       ('fftw_unlock_fun', c_void_p),
                       ('fftw_lock_data', c_void_p)]

src/finufft_core.cpp

@@ -108,7 +108,7 @@ static int setup_spreader_for_nufft(finufft_spread_opts &spopts, T eps,
   // this calls spreadinterp.cpp...
   int ier = setup_spreader(spopts, eps, opts.upsampfac, opts.spread_kerevalmeth,
                            opts.spread_debug, opts.showwarn, opts.spreadinterponly,
-                           dim);
+                           dim, opts.spread_kernel);
   // override various spread opts from their defaults...
   spopts.debug    = opts.spread_debug;
   spopts.sort     = opts.spread_sort; // could make dim or CPU choices here?
@@ -532,6 +532,7 @@ template<typename TF> void FINUFFT_PLAN_T<TF>::precompute_horner_coeffs() {
   // Precompute kernel parameters once
   const TF beta    = TF(this->spopts.ES_beta);
   const TF c_param = TF(this->spopts.ES_c);
+  const int kernel_type = this->spopts.kernel_type;
 
   nc = MIN_NC;
 
@@ -549,9 +550,9 @@ template<typename TF> void FINUFFT_PLAN_T<TF>::precompute_horner_coeffs() {
     // original: 0.5 * (x - nspread + 2*j + 1)
     const TF shift = TF(2 * j + 1 - nspread);
 
-    const auto kernel = [shift, beta, c_param](TF x) -> TF {
+    const auto kernel = [shift, beta, c_param, kernel_type](TF x) -> TF {
       const TF t = TF(0.5) * (x + shift);
-      return evaluate_kernel(t, beta, c_param);
+      return evaluate_kernel(t, beta, c_param, kernel_type);
     };
 
     const auto coeffs = fit_monomials(kernel, static_cast<int>(max_degree), a, b);
@@ -715,6 +716,7 @@ void finufft_default_opts_t(finufft_opts *o)
   o->maxbatchsize       = 0;
   o->spread_nthr_atomic = -1;
   o->spread_max_sp_size = 0;
+  o->spread_kernel      = 0;
   o->fftw_lock_fun      = nullptr;
   o->fftw_unlock_fun    = nullptr;
   o->fftw_lock_data     = nullptr;