Add DST-2 DST-3 Poisson solver

docs/source/run/parameters.rst

@@ -341,7 +341,7 @@ The default is to use the explicit solver. **We strongly recommend to use the ex
     Which solver to use.
     Possible values: ``explicit`` and ``predictor-corrector``.
 
-* ``fields.poisson_solver`` (`string`) optional (default CPU: `FFTDirichletDirect`, GPU: `FFTDirichletFast`)
+* ``fields.poisson_solver`` (`string`) optional (default CPU: `FFTDirichletDirect`, GPU: `FFTDirichletQuick` or `FFTDirichletFast`)
     Which Poisson solver to use for ``Psi``, ``Ez`` and ``Bz``. The ``predictor-corrector`` BxBy
     solver also uses this poisson solver for ``Bx`` and ``By`` internally. Available solvers are:
 
@@ -358,6 +358,10 @@ The default is to use the explicit solver. **We strongly recommend to use the ex
         algorithm that uses FFTs of the same size as the fields.
         Preferred resolution: :math:`2^N-1`.
 
+      * ``FFTDirichletQuick`` Similar to ``FFTDirichletFast`` but uses a different type of FFT that
+        works better for even resolutions.
+        Preferred resolution: :math:`2^N`.
+
       * ``MGDirichlet`` Use the HiPACE++ multigrid solver to solve the Poisson equation with
         Dirichlet boundary conditions.
         Preferred resolution: :math:`2^N` and :math:`2^N-1`.

src/fields/Fields.H

@@ -460,8 +460,6 @@ public:
 private:
     /** Vector over levels of all required fields to compute current slice */
     amrex::Vector<amrex::MultiFab> m_slices;
-    /** Type of poisson solver to use */
-    std::string m_poisson_solver_str = "";
     /** Class to handle transverse FFT Poisson solver on 1 slice */
     amrex::Vector<std::unique_ptr<FFTPoissonSolver>> m_poisson_solver;
     /** Stores temporary values for z interpolation in Fields::Copy */

src/fields/Fields.cpp

@@ -11,6 +11,7 @@
 #include "fft_poisson_solver/FFTPoissonSolverDirichletDirect.H"
 #include "fft_poisson_solver/FFTPoissonSolverDirichletExpanded.H"
 #include "fft_poisson_solver/FFTPoissonSolverDirichletFast.H"
+#include "fft_poisson_solver/FFTPoissonSolverDirichletQuick.H"
 #include "fft_poisson_solver/MGPoissonSolverDirichlet.H"
 #include "Hipace.H"
 #include "OpenBoundary.H"
@@ -33,13 +34,6 @@ Fields::ReadParameters (const int nlev)
 
     amrex::ParmParse ppf("fields");
     DeprecatedInput("fields", "do_dirichlet_poisson", "poisson_solver", "");
-    // set default Poisson solver based on the platform
-#ifdef AMREX_USE_GPU
-    m_poisson_solver_str = "FFTDirichletFast";
-#else
-    m_poisson_solver_str = "FFTDirichletDirect";
-#endif
-    queryWithParser(ppf, "poisson_solver", m_poisson_solver_str);
     queryWithParser(ppf, "insitu_period", m_insitu_period);
     m_insitu_file_prefix = Hipace::m_output_folder + "/insitu";
     const bool set_file_prefix = queryWithParser(ppf, "insitu_file_prefix", m_insitu_file_prefix);
@@ -203,38 +197,53 @@ Fields::AllocData (
         m_slices[lev].setVal(0._rt);
     }
 
+    // set default Poisson solver based on the platform and resolution
+#ifdef AMREX_USE_GPU
+    const bool is_even = std::max(slice_ba[0].length(0), slice_ba[0].length(1)) % 2 == 0;
+    std::string poisson_solver_str = is_even ? "FFTDirichletQuick" : "FFTDirichletFast";
+#else
+    std::string poisson_solver_str = "FFTDirichletDirect";
+#endif
+    amrex::ParmParse ppf("fields");
+    queryWithParser(ppf, "poisson_solver", poisson_solver_str);
+
     // The Poisson solver operates on transverse slices only.
     // The constructor takes the BoxArray and the DistributionMap of a slice,
     // so the FFTPlans are built on a slice.
-    if (m_poisson_solver_str == "FFTDirichletDirect"){
+    if (poisson_solver_str == "FFTDirichletDirect"){
         m_poisson_solver.push_back(std::unique_ptr<FFTPoissonSolverDirichletDirect>(
             new FFTPoissonSolverDirichletDirect(getSlices(lev).boxArray(),
                                                 getSlices(lev).DistributionMap(),
                                                 geom)) );
-    } else if (m_poisson_solver_str == "FFTDirichletExpanded"){
+    } else if (poisson_solver_str == "FFTDirichletExpanded"){
         m_poisson_solver.push_back(std::unique_ptr<FFTPoissonSolverDirichletExpanded>(
             new FFTPoissonSolverDirichletExpanded(getSlices(lev).boxArray(),
                                                   getSlices(lev).DistributionMap(),
                                                   geom)) );
-    } else if (m_poisson_solver_str == "FFTDirichletFast"){
+    } else if (poisson_solver_str == "FFTDirichletFast"){
         m_poisson_solver.push_back(std::unique_ptr<FFTPoissonSolverDirichletFast>(
             new FFTPoissonSolverDirichletFast(getSlices(lev).boxArray(),
                                               getSlices(lev).DistributionMap(),
                                               geom)) );
-    } else if (m_poisson_solver_str == "FFTPeriodic") {
+    } else if (poisson_solver_str == "FFTDirichletQuick"){
+        m_poisson_solver.push_back(std::unique_ptr<FFTPoissonSolverDirichletQuick>(
+            new FFTPoissonSolverDirichletQuick(getSlices(lev).boxArray(),
+                                               getSlices(lev).DistributionMap(),
+                                               geom)) );
+    } else if (poisson_solver_str == "FFTPeriodic") {
         m_poisson_solver.push_back(std::unique_ptr<FFTPoissonSolverPeriodic>(
             new FFTPoissonSolverPeriodic(getSlices(lev).boxArray(),
                                          getSlices(lev).DistributionMap(),
                                          geom))  );
-    } else if (m_poisson_solver_str == "MGDirichlet") {
+    } else if (poisson_solver_str == "MGDirichlet") {
         m_poisson_solver.push_back(std::unique_ptr<MGPoissonSolverDirichlet>(
             new MGPoissonSolverDirichlet(getSlices(lev).boxArray(),
                                          getSlices(lev).DistributionMap(),
                                          geom))  );
     } else {
-        amrex::Abort("Unknown poisson solver '" + m_poisson_solver_str +
+        amrex::Abort("Unknown poisson solver '" + poisson_solver_str +
             "', must be 'FFTDirichletDirect', 'FFTDirichletExpanded', 'FFTDirichletFast', " +
-            "'FFTPeriodic' or 'MGDirichlet'");
+            "'FFTDirichletQuick', 'FFTPeriodic' or 'MGDirichlet'");
     }
 
     if (lev == 0 && m_insitu_period > 0) {

src/fields/fft_poisson_solver/CMakeLists.txt

@@ -12,6 +12,7 @@ target_sources(HiPACE
     FFTPoissonSolverDirichletDirect.cpp
     FFTPoissonSolverDirichletExpanded.cpp
     FFTPoissonSolverDirichletFast.cpp
+    FFTPoissonSolverDirichletQuick.cpp
     MGPoissonSolverDirichlet.cpp
 )
 

src/fields/fft_poisson_solver/FFTPoissonSolverDirichletQuick.H

@@ -0,0 +1,82 @@
+/* Copyright 2025
+ *
+ * This file is part of HiPACE++.
+ *
+ * Authors: AlexanderSinn, MaxThevenet, Severin Diederichs
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef FFT_POISSON_SOLVER_DIRICHLET_QUICK_H_
+#define FFT_POISSON_SOLVER_DIRICHLET_QUICK_H_
+
+#include "fields/fft_poisson_solver/fft/AnyFFT.H"
+#include "FFTPoissonSolver.H"
+
+#include <AMReX_MultiFab.H>
+#include <AMReX_GpuComplex.H>
+
+/**
+ * \brief This class handles functions and data to perform transverse Fourier-based Poisson solves.
+ *
+ * For a given source S, it solves equation Laplacian(F) = S and returns F.
+ * Once an instance is created, a typical use consists in:
+ * 1. Compute S directly in FFTPoissonSolver::m_stagingArea
+ * 2. Call FFTPoissonSolver::SolvePoissonEquation(mf), which will solve Poisson equation with RHS
+ *    in the staging area and return the LHS in mf.
+ */
+class FFTPoissonSolverDirichletQuick final : public FFTPoissonSolver
+{
+public:
+    /** Constructor */
+    FFTPoissonSolverDirichletQuick ( amrex::BoxArray const& a_realspace_ba,
+                                    amrex::DistributionMapping const& dm,
+                                    amrex::Geometry const& gm);
+
+    /** virtual destructor */
+    virtual ~FFTPoissonSolverDirichletQuick () override final {}
+
+    /**
+     * \param[in] realspace_ba BoxArray on which the FFT is executed.
+     * \param[in] dm DistributionMapping for the BoxArray.
+     * \param[in] gm Geometry, contains the box dimensions.
+     */
+    void define ( amrex::BoxArray const& realspace_ba,
+                  amrex::DistributionMapping const& dm,
+                  amrex::Geometry const& gm);
+
+    /**
+     * Solve Poisson equation. The source term must be stored in the staging area m_stagingArea prior to this call.
+     *
+     * \param[in] lhs_mf Destination array, where the result is stored.
+     */
+    virtual void SolvePoissonEquation (amrex::MultiFab& lhs_mf) override final;
+
+    /** Position and relative factor used to apply inhomogeneous Dirichlet boundary conditions */
+    virtual amrex::Real BoundaryOffset() override final { return 0.5; }
+    virtual amrex::Real BoundaryFactor() override final { return 2.; }
+
+private:
+    /** Real array for the FFTs */
+    amrex::Gpu::DeviceVector<amrex::Real> m_real_array;
+    /** Complex array for the FFTs */
+    amrex::Gpu::DeviceVector<amrex::GpuComplex<amrex::Real>> m_comp_array;
+    /** FFT plan in x direction */
+    AnyFFT m_x_r2cfft;
+    /** FFT plan in y direction */
+    AnyFFT m_y_r2cfft;
+    /** FFT plan in x direction */
+    AnyFFT m_x_c2rfft;
+    /** FFT plan in y direction */
+    AnyFFT m_y_c2rfft;
+    /** work area for all DST plans */
+    amrex::Gpu::DeviceVector<char> m_fft_work_area;
+    /** x prefactor for dst2_out and dst3_in */
+    amrex::Gpu::DeviceVector<amrex::GpuComplex<amrex::Real>> m_omega_x;
+    /** y prefactor for dst2_out and dst3_in */
+    amrex::Gpu::DeviceVector<amrex::GpuComplex<amrex::Real>> m_omega_y;
+    /** x prefactor for eigenvalue */
+    amrex::Gpu::DeviceVector<amrex::Real> m_eig_x;
+    /** y prefactor for eigenvalue */
+    amrex::Gpu::DeviceVector<amrex::Real> m_eig_y;
+};
+
+#endif