Merge pull request #758 from DiamonDinoia/feature/hint-nj-density

lu1and10 · web-flow · commit b4b8a10bf626 · 2025-12-09T20:44:31.000-05:00
Using density to determine upsampfact, deferring to setpts when auto
Now type3 upsampfact may differs from inner type2 upsampfact if opts.upsampfac is set to 0.
Previously type3 upsampfact is always same as the inner type2 upsampfact.
diff --git a/include/finufft/fft.h b/include/finufft/fft.h
@@ -79,6 +79,15 @@ template<> struct Finufft_FFT_plan<float> {
     uint64_t nf = 1;
     for (auto i : dims) nf *= i;
     lock();
+    // Destroy existing plans before creating new ones (handles re-planning)
+    if (plan_) {
+      fftwf_destroy_plan(plan_);
+      plan_ = nullptr;
+    }
+    if (plan_adj_) {
+      fftwf_destroy_plan(plan_adj_);
+      plan_adj_ = nullptr;
+    }
 #ifdef _OPENMP
     fftwf_plan_with_nthreads(nthreads);
 #endif
@@ -155,6 +164,15 @@ template<> struct Finufft_FFT_plan<double> {
     uint64_t nf = 1;
     for (auto i : dims) nf *= i;
     lock();
+    // Destroy existing plans before creating new ones (handles re-planning)
+    if (plan_) {
+      fftw_destroy_plan(plan_);
+      plan_ = nullptr;
+    }
+    if (plan_adj_) {
+      fftw_destroy_plan(plan_adj_);
+      plan_adj_ = nullptr;
+    }
 #ifdef _OPENMP
     fftw_plan_with_nthreads(nthreads);
 #endif
diff --git a/include/finufft/finufft_core.h b/include/finufft/finufft_core.h
@@ -84,14 +84,16 @@ template<typename TF> struct FINUFFT_PLAN_T { // the main plan class, fully C++
   int dim;  // overall dimension: 1,2 or 3
 
 private:
-  int ntrans;    // how many transforms to do at once (vector or "many" mode)
-  BIGINT nj;     // num of NU pts in type 1,2 (for type 3, num input x pts)
-  BIGINT nk;     // number of NU freq pts (type 3 only)
-  TF tol;        // relative user tolerance
-  int batchSize; // # strength vectors to group together for FFTW, etc
-  int nbatch;    // how many batches done to cover all ntrans vectors
+  int ntrans;             // how many transforms to do at once (vector or "many" mode)
+  BIGINT nj;              // num of NU pts in type 1,2 (for type 3, num input x pts)
+  BIGINT nk;              // number of NU freq pts (type 3 only)
+  TF tol;                 // relative user tolerance
+  int batchSize;          // # strength vectors to group together for FFTW, etc
+  int nbatch;             // how many batches done to cover all ntrans vectors
 
-  int nc = 0;    // number of Horner coefficients used for ES kernel (<= MAX_NC)
+  int nc             = 0; // number of Horner coefficients used for ES kernel (<= MAX_NC)
+  bool upsamp_locked = false; // true if user specified upsampfac != 0, prevents auto
+                              // update
 
 public:
   std::array<BIGINT, 3> mstu; // number of modes in x,y,z directions
@@ -148,6 +150,10 @@ template<typename TF> struct FINUFFT_PLAN_T { // the main plan class, fully C++
 
   void precompute_horner_coeffs();
 
+  // Helper to initialize spreader, phiHat (Fourier series), and FFT plan.
+  // Used by constructor (when upsampfac given) and setpts (when upsampfac deferred).
+  int initSpreadAndFFT();
+
 public:
   FINUFFT_PLAN_T(int type, int dim, const BIGINT *n_modes, int iflag, int ntrans, TF tol,
                  const finufft_opts *opts, int &ier);
diff --git a/src/finufft_core.cpp b/src/finufft_core.cpp
@@ -594,6 +594,88 @@ template<typename TF> void FINUFFT_PLAN_T<TF>::precompute_horner_coeffs() {
   }
 }
 
+// Helper to initialize spreader, phiHat (Fourier series), and FFT plan.
+// Used by constructor (when upsampfac given) and setpts (when upsampfac deferred).
+// Returns 0 on success, or an error code if set_nf_type12 or alloc fails.
+template<typename TF> int FINUFFT_PLAN_T<TF>::initSpreadAndFFT() {
+  CNTime timer{};
+  spopts.spread_direction = type;
+  constexpr TF EPSILON    = std::numeric_limits<TF>::epsilon();
+
+  if (opts.spreadinterponly) { // (unusual case of no NUFFT, just report)
+    // spreadinterp grid will simply be the user's "mode" grid...
+    for (int idim = 0; idim < dim; ++idim) nfdim[idim] = mstu[idim];
+
+    if (opts.debug) { // "long long" here is to avoid warnings with printf...
+      printf("[%s] %dd spreadinterponly(dir=%d): (ms,mt,mu)=(%lld,%lld,%lld)"
+             "\n               ntrans=%d nthr=%d batchSize=%d kernel width ns=%d",
+             __func__, dim, type, (long long)mstu[0], (long long)mstu[1],
+             (long long)mstu[2], ntrans, opts.nthreads, batchSize, spopts.nspread);
+      if (batchSize == 1) // spread_thread has no effect in this case
+        printf("\n");
+      else
+        printf(" spread_thread=%d\n", opts.spread_thread);
+    }
+
+  } else { // ..... usual NUFFT: eval Fourier series, alloc workspace .....
+
+    if (opts.showwarn) { // user warn round-off error (due to prob condition #)...
+      for (int idim = 0; idim < dim; ++idim)
+        if (EPSILON * mstu[idim] > 1.0)
+          fprintf(stderr,
+                  "%s warning: rounding err (due to cond # of prob) eps_mach*N%d = %.3g "
+                  "> 1 !\n",
+                  __func__, idim, (double)(EPSILON * mstu[idim]));
+    }
+
+    // determine fine grid sizes, sanity check, then alloc...
+    for (int idim = 0; idim < dim; ++idim) {
+      int nfier = set_nf_type12(mstu[idim], opts, spopts, &nfdim[idim]);
+      if (nfier) return nfier;                    // nf too big; we're done
+      phiHat[idim].resize(nfdim[idim] / 2 + 1);   // alloc fseries
+    }
+
+    if (opts.debug) { // "long long" here is to avoid warnings with printf...
+      printf("[%s] %dd%d: (ms,mt,mu)=(%lld,%lld,%lld) "
+             "(nf1,nf2,nf3)=(%lld,%lld,%lld)\n               ntrans=%d nthr=%d "
+             "batchSize=%d ",
+             __func__, dim, type, (long long)mstu[0], (long long)mstu[1],
+             (long long)mstu[2], (long long)nfdim[0], (long long)nfdim[1],
+             (long long)nfdim[2], ntrans, opts.nthreads, batchSize);
+      if (batchSize == 1) // spread_thread has no effect in this case
+        printf("\n");
+      else
+        printf(" spread_thread=%d\n", opts.spread_thread);
+    }
+
+    // STEP 0: get Fourier coeffs of spreading kernel along each fine grid dim
+    timer.restart();
+    for (int idim = 0; idim < dim; ++idim)
+      onedim_fseries_kernel(nfdim[idim], phiHat[idim], spopts, horner_coeffs.data(), nc);
+    if (opts.debug)
+      printf("[%s] kernel fser (ns=%d):\t\t%.3g s\n", __func__, spopts.nspread,
+             timer.elapsedsec());
+
+    if (nf() * batchSize > MAX_NF) {
+      fprintf(stderr,
+              "[%s] fwBatch would be bigger than MAX_NF, not attempting memory "
+              "allocation!\n",
+              __func__);
+      return FINUFFT_ERR_MAXNALLOC;
+    }
+
+    timer.restart(); // plan the FFTW (to act in-place on the workspace fwBatch)
+    int nthr_fft = opts.nthreads;
+    const auto ns = gridsize_for_fft(*this);
+    std::vector<TC, xsimd::aligned_allocator<TC, 64>> fwBatch(nf() * batchSize);
+    fftPlan->plan(ns, batchSize, fwBatch.data(), fftSign, opts.fftw, nthr_fft);
+    if (opts.debug)
+      printf("[%s] FFT plan (mode %d, nthr=%d):\t%.3g s\n", __func__, opts.fftw, nthr_fft,
+             timer.elapsedsec());
+  }
+  return 0;
+}
+
 // --------------- rest is the five user guru (plan) interface drivers ----------
 // (they are not namespaced since have safe names finufft{f}_* )
 using namespace finufft::utils; // AHB since already given at top, needed again?
@@ -749,106 +831,28 @@ FINUFFT_PLAN_T<TF>::FINUFFT_PLAN_T(int type_, int dim_, const BIGINT *n_modes, i
     }
   }
 
-  // heuristic to choose default upsampfac... (currently two poss)
-  if (opts.upsampfac == 0.0) { // init to auto choice
-    // Let assume density=1 as the average use case.
-    // TODO: make a decision on how to choose density properly.
-    const auto density = TF{1};
-    opts.upsampfac     = bestUpsamplingFactor<TF>(opts.nthreads, density, dim, type, tol);
-    if (opts.debug > 1)
-      printf("[%s] threads %d, density %.3g, dim %d, nufft type %d, tol %.3g: auto "
-             "upsampfac=%.2f\n",
-             __func__, opts.nthreads, density, dim, type, tol, opts.upsampfac);
-  }
-  // use opts to choose and write into plan's spread options...
-  ier = setup_spreader_for_nufft(spopts, tol, opts, dim);
-  if (ier > 1) // proceed if success or warning
-    throw int(ier);
-  precompute_horner_coeffs();
-  //  ------------------------ types 1,2: planning needed ---------------------
-  if (type == 1 || type == 2) {
-
-    int nthr_fft = nthr; // give FFT all threads (or use o.spread_thread?)
-                         // Note: batchSize not used since might be only 1.
-
-    spopts.spread_direction = type;
-    constexpr TF EPSILON    = std::numeric_limits<TF>::epsilon();
-
-    if (opts.spreadinterponly) { // (unusual case of no NUFFT, just report)
-
-      // spreadinterp grid will simply be the user's "mode" grid...
-      for (int idim = 0; idim < dim; ++idim) nfdim[idim] = mstu[idim];
-
-      if (opts.debug) { // "long long" here is to avoid warnings with printf...
-        printf("[%s] %dd spreadinterponly(dir=%d): (ms,mt,mu)=(%lld,%lld,%lld)"
-               "\n               ntrans=%d nthr=%d batchSize=%d kernel width ns=%d",
-               __func__, dim, type, (long long)mstu[0], (long long)mstu[1],
-               (long long)mstu[2], ntrans, nthr, batchSize, spopts.nspread);
-        if (batchSize == 1) // spread_thread has no effect in this case
-          printf("\n");
-        else
-          printf(" spread_thread=%d\n", opts.spread_thread);
-      }
-
-    } else { // ..... usual NUFFT: eval Fourier series, alloc workspace .....
-
-      if (opts.showwarn) { // user warn round-off error (due to prob condition #)...
-        for (int idim = 0; idim < dim; ++idim)
-          if (EPSILON * mstu[idim] > 1.0)
-            fprintf(
-                stderr,
-                "%s warning: rounding err (due to cond # of prob) eps_mach*N%d = %.3g "
-                "> 1 !\n",
-                __func__, idim, (double)(EPSILON * mstu[idim]));
-      }
-
-      // determine fine grid sizes, sanity check, then alloc...
-      for (int idim = 0; idim < dim; ++idim) {
-        int nfier = set_nf_type12(mstu[idim], opts, spopts, &nfdim[idim]);
-        if (nfier) throw nfier;                   // nf too big; we're done
-        phiHat[idim].resize(nfdim[idim] / 2 + 1); // alloc fseries
-      }
-
-      if (opts.debug) { // "long long" here is to avoid warnings with printf...
-        printf("[%s] %dd%d: (ms,mt,mu)=(%lld,%lld,%lld) "
-               "(nf1,nf2,nf3)=(%lld,%lld,%lld)\n               ntrans=%d nthr=%d "
-               "batchSize=%d ",
-               __func__, dim, type, (long long)mstu[0], (long long)mstu[1],
-               (long long)mstu[2], (long long)nfdim[0], (long long)nfdim[1],
-               (long long)nfdim[2], ntrans, nthr, batchSize);
-        if (batchSize == 1) // spread_thread has no effect in this case
-          printf("\n");
-        else
-          printf(" spread_thread=%d\n", opts.spread_thread);
-      }
-
-      // STEP 0: get Fourier coeffs of spreading kernel along each fine grid dim
-      timer.restart();
-      for (int idim = 0; idim < dim; ++idim)
-        onedim_fseries_kernel(nfdim[idim], phiHat[idim], spopts, horner_coeffs.data(),
-                              nc);
-      if (opts.debug)
-        printf("[%s] kernel fser (ns=%d):\t\t%.3g s\n", __func__, spopts.nspread,
-               timer.elapsedsec());
-
-      if (nf() * batchSize > MAX_NF) {
-        fprintf(stderr,
-                "[%s] fwBatch would be bigger than MAX_NF, not attempting memory "
-                "allocation!\n",
-                __func__);
-        throw int(FINUFFT_ERR_MAXNALLOC);
-      }
-
-      timer.restart(); // plan the FFTW (to act in-place on the workspace fwBatch)
-      const auto ns = gridsize_for_fft(*this);
-      std::vector<TC, xsimd::aligned_allocator<TC, 64>> fwBatch(nf() * batchSize);
-      fftPlan->plan(ns, batchSize, fwBatch.data(), fftSign, opts.fftw, nthr_fft);
-      if (opts.debug)
-        printf("[%s] FFT plan (mode %d, nthr=%d):\t%.3g s\n", __func__, opts.fftw,
-               nthr_fft, timer.elapsedsec());
+  // heuristic to choose default upsampfac: defer selection to setpts unless
+  // the user explicitly forced a nonzero value in opts. In that case initialize
+  // spreader/Horner internals now using the provided upsampfac.
+  if (opts.upsampfac != 0.0) {
+    upsamp_locked = true; // user explicitly set upsampfac, don't auto-update
+    ier = setup_spreader_for_nufft(spopts, tol, opts, dim);
+    if (ier > 1) // proceed if success or warning
+      throw int(ier);
+    precompute_horner_coeffs();
+
+    //  ------------------------ types 1,2: planning needed ---------------------
+    if (type == 1 || type == 2) {
+      int code = initSpreadAndFFT();
+      if (code) throw code;
     }
+  } else {
+    // If upsampfac was left as 0.0 (auto) we defer setup_spreader to setpts.
+    // However, we must still check if tol is too small to warn the user now.
+    if (tol < std::numeric_limits<TF>::epsilon()) ier = FINUFFT_WARN_EPS_TOO_SMALL;
+  }
 
-  } else { // -------------------------- type 3 (no planning) ------------
+  if (type == 3) { // -------------------------- type 3 (no planning) ------------
 
     if (opts.debug) printf("[%s] %dd%d: ntrans=%d\n", __func__, dim, type, ntrans);
     // Type 3 will call finufft_makeplan for type 2; no need to init FFTW
@@ -908,6 +912,27 @@ int FINUFFT_PLAN_T<TF>::setpts(BIGINT nj, const TF *xj, const TF *yj, const TF *
 
   if (type != 3) {          // ------------------ TYPE 1,2 SETPTS -------------------
                             // (all we can do is check and maybe bin-sort the NU pts)
+    // If upsampfac is not locked by user (auto mode), choose or update it now
+    // based on the actual density nj/N(). Re-plan if density changed significantly.
+    if (!upsamp_locked) {
+      double density = double(nj) / double(N());
+      double upsampfac = bestUpsamplingFactor<TF>(opts.nthreads, density, dim, type, tol);
+      // Re-plan if this is the first call (upsampfac==0) or if upsampfac changed
+      if (upsampfac != opts.upsampfac) {
+        opts.upsampfac = upsampfac;
+        if (opts.debug > 1)
+          printf("[setpts] selected upsampfac=%.2f (density=%.3f, nj=%lld)\n",
+                 opts.upsampfac, density, (long long)nj);
+        int code = setup_spreader_for_nufft(spopts, tol, opts, dim);
+        if (code > 1) return code;
+        precompute_horner_coeffs();
+
+        // Perform the planning steps (first call or re-plan due to density change).
+        code = initSpreadAndFFT();
+        if (code) return code;
+      }
+    }
+
     XYZ     = {xj, yj, zj}; // plan must keep pointers to user's fixed NU pts
     int ier = spreadcheck(nfdim[0], nfdim[1], nfdim[2], spopts);
     if (ier)                // no warnings allowed here
@@ -935,6 +960,21 @@ int FINUFFT_PLAN_T<TF>::setpts(BIGINT nj, const TF *xj, const TF *yj, const TF *
     this->nk = nk; // user set # targ freq pts
     STU      = {s, t, u};
 
+    // For type 3 with deferred upsampfac (not locked by user), pick and persist
+    // an upsamp now using density=1.0 so that subsequent steps (set_nhg_type3 etc.)
+    // have a concrete upsampfac to work with. This choice is persisted so inner
+    // type-2 plans will inherit it.
+    double upsampfac = bestUpsamplingFactor<TF>(opts.nthreads, 1.0, dim, type, tol);
+    if (!upsamp_locked && (upsampfac != opts.upsampfac)) {
+      opts.upsampfac = upsampfac;
+      if (opts.debug > 1)
+        printf("[setpts t3] selected upsampfac=%.2f (density=1 used; persisted)\n",
+               opts.upsampfac);
+      int sier = setup_spreader_for_nufft(spopts, tol, opts, dim);
+      if (sier > 1) return sier;
+      precompute_horner_coeffs();
+    }
+
     // pick x, s intervals & shifts & # fine grid pts (nf) in each dim...
     std::array<TF, 3> S = {0, 0, 0};
     if (opts.debug) printf("\tM=%lld N=%lld\n", (long long)nj, (long long)nk);
@@ -1038,6 +1078,8 @@ int FINUFFT_PLAN_T<TF>::setpts(BIGINT nj, const TF *xj, const TF *yj, const TF *
     t2opts.debug        = std::max(0, opts.debug - 1);    // don't print as much detail
     t2opts.spread_debug = std::max(0, opts.spread_debug - 1);
     t2opts.showwarn     = 0;                              // so don't see warnings 2x
+    if (!upsamp_locked) t2opts.upsampfac = 0.0; // if the upsampfac was auto, let inner
+                                                // t2 pick it again (from density=nj/Nf)
     // (...could vary other t2opts here?)
     // MR: temporary hack, until we have figured out the C++ interface.
     FINUFFT_PLAN_T<TF> *tmpplan;
