Add quasi-3D Integrated Green Functions solver

Docs/source/usage/parameters.rst

@@ -290,6 +290,19 @@ Overall simulation parameters
         In electromagnetic mode, this solver can be used to initialize the species' self fields
         (``<species_name>.initialize_self_fields=1``) provided that the field BCs are PML (``boundary.field_lo,hi = PML``).
 
+          * warpx.use_2d_slices_fft_solver (`bool`, default: 0): Select the type of Integrated Green Function solver.
+            If 0, solve Poisson equation in full 3D geometry.
+            If 1, solve Poisson equation in a quasi 3D geometry, neglecting the :math:`z` derivatives in the Laplacian of the Poisson equation.
+            In practice, in this case, the code performes many 2D Poisson solves on all :math:`(x,y)` slices, each slice at a given :math:`z`.
+            This is often a good approximation for ultra-relativistic beams propagating along the :math:`z` direction, with the relativistic solver.
+            As a consequence, this solver does not need to do an FFT along the :math:`z` direction,
+            and instead uses only transverse FFTs (along :math:`x` and :math:`y`) at each :math:`z` position (or :math:`z` "slice").
+
+          * warpx.use_distributed_3d_fft_solver (`bool`, default: 0): Choose whether the 3D FFTs performed in the
+            full 3D Integrated Green Function solver are distributed.
+            If 0, the FFTs are performed on a single MPI rank (the rest of the code is still fully parallel).
+            If 1, the FFTs are distributed using the heFFTe library. The code must be compiled with `-DWarpX_HEFFTE=ON`.
+
 * ``warpx.self_fields_required_precision`` (`float`, default: 1.e-11)
     The relative precision with which the electrostatic space-charge fields should
     be calculated. More specifically, the space-charge fields are

Examples/Tests/open_bc_poisson_solver/CMakeLists.txt

@@ -12,3 +12,15 @@ if(WarpX_FFT)
         OFF  # dependency
     )
 endif()
+
+if(WarpX_FFT)
+    add_warpx_test(
+        test_3d_open_bc_poisson_solver_sliced  # name
+        3  # dims
+        2  # nprocs
+        inputs_test_3d_open_bc_poisson_solver_sliced  # inputs
+        analysis.py  # analysis
+        diags/diag1000001  # output
+        OFF  # dependency
+    )
+endif()

Examples/Tests/open_bc_poisson_solver/analysis.py

@@ -61,7 +61,6 @@ def evaluate_E(x, y, z):
     assert np.allclose(Ex_warpx, Ex_theory, rtol=0.032, atol=0)
     assert np.allclose(Ey_warpx, Ey_theory, rtol=0.029, atol=0)
 
-
 # compare checksums
 evaluate_checksum(
     test_name=os.path.split(os.getcwd())[1],

Examples/Tests/open_bc_poisson_solver/inputs_test_3d_open_bc_poisson_solver_sliced

@@ -0,0 +1,3 @@
+FILE = inputs_test_3d_open_bc_poisson_solver
+
+warpx.use_2d_slices_fft_solver = 1

Regression/Checksum/benchmarks_json/test_3d_open_bc_poisson_solver_heffte.json

@@ -0,0 +1,20 @@
+{
+  "lev=0": {
+    "Bx": 100915933.44552392,
+    "By": 157610622.18550807,
+    "Bz": 8.739731456892512e-12,
+    "Ex": 4.725065270619816e+16,
+    "Ey": 3.025394898923131e+16,
+    "Ez": 3276573.95147766,
+    "rho": 10994013582437.193
+  },
+  "electron": {
+    "particle_momentum_x": 5.701277606055911e-19,
+    "particle_momentum_y": 3.650451663619896e-19,
+    "particle_momentum_z": 1.145432768297242e-10,
+    "particle_position_x": 17.314086912497864,
+    "particle_position_y": 0.2583691267187796,
+    "particle_position_z": 10066.329600000008,
+    "particle_weight": 19969036501.910976
+  }
+}

Regression/Checksum/benchmarks_json/test_3d_open_bc_poisson_solver_sliced.json

@@ -0,0 +1,21 @@
+{
+  "lev=0": {
+    "Bx": 100915933.44551668,
+    "By": 157610622.18551716,
+    "Bz": 2.598515299403035e-15,
+    "Ex": 4.725065270620093e+16,
+    "Ey": 3.0253948989229424e+16,
+    "Ez": 2787743.3330717986,
+    "rho": 10994013582437.193
+  },
+  "electron": {
+    "particle_momentum_x": 5.701277606056779e-19,
+    "particle_momentum_y": 3.650451663675671e-19,
+    "particle_momentum_z": 1.145432768297242e-10,
+    "particle_position_x": 17.314086912497864,
+    "particle_position_y": 0.25836912671877954,
+    "particle_position_z": 10066.329600000008,
+    "particle_weight": 19969036501.910976
+  }
+}
+

Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.H

@@ -92,7 +92,8 @@ public:
         amrex::Real required_precision,
         amrex::Real absolute_tolerance,
         int max_iters,
-        int verbosity
+        int verbosity,
+        bool is_igf_2d_slices
     ) const;
 
     /**
@@ -153,6 +154,10 @@ public:
      *  2 : convergence progress at every MLMG iteration
      */
     int self_fields_verbosity = 2;
+
+    /** Parameters for FFT Poisson solver aka IGF */
+    // 0: full 3D, 1: many 2D z-slices (quasi-3D)
+    bool is_igf_2d_slices = false;
 };
 
 #endif // WARPX_ELECTROSTATICSOLVER_H_

Source/FieldSolver/ElectrostaticSolvers/ElectrostaticSolver.cpp

@@ -38,7 +38,12 @@ void ElectrostaticSolver::ReadParameters () {
         pp_warpx, "self_fields_absolute_tolerance", self_fields_absolute_tolerance);
     utils::parser::queryWithParser(
         pp_warpx, "self_fields_max_iters", self_fields_max_iters);
-    pp_warpx.query("self_fields_verbosity", self_fields_verbosity);
+   utils::parser::queryWithParser(
+        pp_warpx, "self_fields_verbosity", self_fields_verbosity);
+
+    // FFT solver flags
+   utils::parser::queryWithParser(
+        pp_warpx, "use_2d_slices_fft_solver", is_igf_2d_slices);
 }
 
 void
@@ -121,7 +126,9 @@ ElectrostaticSolver::computePhi (
     Real const required_precision,
     Real absolute_tolerance,
     int const max_iters,
-    int const verbosity) const
+    int const verbosity,
+    bool const is_igf_2d
+) const
 {
     using ablastr::fields::Direction;
 
@@ -202,6 +209,7 @@ ElectrostaticSolver::computePhi (
         warpx.boxArray(),
         WarpX::grid_type,
         is_solver_igf_on_lev0,
+        is_igf_2d,
         EB::enabled(),
         WarpX::do_single_precision_comms,
         warpx.refRatio(),

Source/FieldSolver/ElectrostaticSolvers/LabFrameExplicitES.cpp

@@ -75,7 +75,7 @@ void LabFrameExplicitES::ComputeSpaceChargeField (
         // Use the AMREX MLMG or the FFT (IGF) solver otherwise
         computePhi(rho_fp, phi_fp, beta, self_fields_required_precision,
                    self_fields_absolute_tolerance, self_fields_max_iters,
-                   self_fields_verbosity);
+                   self_fields_verbosity, is_igf_2d_slices);
 #endif
 
     }

Source/FieldSolver/ElectrostaticSolvers/RelativisticExplicitES.cpp

@@ -130,7 +130,7 @@ void RelativisticExplicitES::AddSpaceChargeField (
     computePhi( amrex::GetVecOfPtrs(rho), amrex::GetVecOfPtrs(phi),
                 beta, pc.self_fields_required_precision,
                 pc.self_fields_absolute_tolerance, pc.self_fields_max_iters,
-                pc.self_fields_verbosity );
+                pc.self_fields_verbosity, is_igf_2d_slices);
 
     // Compute the corresponding electric and magnetic field, from the potential phi
     computeE( Efield_fp, amrex::GetVecOfPtrs(phi), beta );
@@ -168,7 +168,7 @@ void RelativisticExplicitES::AddBoundaryField (ablastr::fields::MultiLevelVector
     computePhi( amrex::GetVecOfPtrs(rho), amrex::GetVecOfPtrs(phi),
                 beta, self_fields_required_precision,
                 self_fields_absolute_tolerance, self_fields_max_iters,
-                self_fields_verbosity );
+                self_fields_verbosity, is_igf_2d_slices);
 
     // Compute the corresponding electric field, from the potential phi.
     computeE( Efield_fp, amrex::GetVecOfPtrs(phi), beta );

Source/ablastr/fields/IntegratedGreenFunctionSolver.H

@@ -15,13 +15,14 @@
 #include <AMReX_REAL.H>
 #include <AMReX_Vector.H>
 
-
 #include <array>
 #include <cmath>
 
 
 namespace ablastr::fields
 {
+    using namespace amrex::literals;
+
 
     /** @brief Implements equation 2 in https://doi.org/10.1103/PhysRevSTAB.10.129901
      *         with some modification to symmetrize the function.
@@ -30,54 +31,92 @@ namespace ablastr::fields
      * @param[in] y y-coordinate of given location
      * @param[in] z z-coordinate of given location
      *
-     * @return the integrated Green function G
+     * @return the integrated Green function G in 3D
      */
     AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     amrex::Real
-    IntegratedPotential (amrex::Real x, amrex::Real y, amrex::Real z)
+    IntegratedPotential3D (amrex::Real x, amrex::Real y, amrex::Real z)
     {
-        using namespace amrex::literals;
-
         amrex::Real const r = std::sqrt( x*x + y*y + z*z );
-        amrex::Real const G =
-            - 0.5_rt * z*z * std::atan( x*y/(z*r) )
-            - 0.5_rt * y*y * std::atan( x*z/(y*r) )
-            - 0.5_rt * x*x * std::atan( y*z/(x*r) )
-            + y*z*std::asinh( x/std::sqrt(y*y + z*z) )
-            + x*z*std::asinh( y/std::sqrt(x*x + z*z) )
-            + x*y*std::asinh( z/std::sqrt(x*x + y*y) );
+        amrex::Real const G = - 0.5_rt * z*z * std::atan( x*y/(z*r) )
+                              - 0.5_rt * y*y * std::atan( x*z/(y*r) )
+                              - 0.5_rt * x*x * std::atan( y*z/(x*r) )
+                              + y*z*std::asinh( x/std::sqrt(y*y + z*z) )
+                              + x*z*std::asinh( y/std::sqrt(x*x + z*z) )
+                              + x*y*std::asinh( z/std::sqrt(x*x + y*y) );
+        return G;
+    }
+
+
+    /** @brief Implements minus equation 58 in https://doi.org/10.1016/j.jcp.2004.01.008
+     *
+     * @param[in] x x-coordinate of given location
+     * @param[in] y y-coordinate of given location
+     *
+     * @return the integrated Green function G in 2D
+     */
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    amrex::Real
+    IntegratedPotential2D (amrex::Real x, amrex::Real y)
+    {
+        amrex::Real const G = 3_rt*x*y
+                              - x*x * std::atan(y/x)
+                              - y*y * std::atan(x/y)
+                              - x*y * std::log(x*x + y*y);
         return G;
     }
 
+
     /** @brief add
      *
      * @param[in] x x-coordinate of given location
      * @param[in] y y-coordinate of given location
      * @param[in] z z-coordinate of given location
+     * @param[in] dx cell size along x
+     * @param[in] dy cell size along y
+     * @param[in] dz cell size along z
      *
-     * @return the sum of integrated Green function G
+     * @return the sum of integrated Green function G in 3D
      */
     AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     amrex::Real
-    SumOfIntegratedPotential (amrex::Real x, amrex::Real y, amrex::Real z, amrex::Real dx, amrex::Real dy, amrex::Real dz)
+    SumOfIntegratedPotential3D (amrex::Real x, amrex::Real y, amrex::Real z, amrex::Real dx, amrex::Real dy, amrex::Real dz)
     {
-        using namespace amrex::literals;
-
+        return 1._rt/(4._rt*ablastr::constant::math::pi*ablastr::constant::SI::ep0) * (
+                  IntegratedPotential3D( x+0.5_rt*dx, y+0.5_rt*dy, z+0.5_rt*dz )
+                - IntegratedPotential3D( x-0.5_rt*dx, y+0.5_rt*dy, z+0.5_rt*dz )
+                - IntegratedPotential3D( x+0.5_rt*dx, y-0.5_rt*dy, z+0.5_rt*dz )
+                + IntegratedPotential3D( x-0.5_rt*dx, y-0.5_rt*dy, z+0.5_rt*dz )
+                - IntegratedPotential3D( x+0.5_rt*dx, y+0.5_rt*dy, z-0.5_rt*dz )
+                + IntegratedPotential3D( x-0.5_rt*dx, y+0.5_rt*dy, z-0.5_rt*dz )
+                + IntegratedPotential3D( x+0.5_rt*dx, y-0.5_rt*dy, z-0.5_rt*dz )
+                - IntegratedPotential3D( x-0.5_rt*dx, y-0.5_rt*dy, z-0.5_rt*dz )
+                );
+    }
 
-        amrex::Real const G_value = 1._rt/(4._rt*ablastr::constant::math::pi*ablastr::constant::SI::ep0) * (
-            IntegratedPotential( x+0.5_rt*dx, y+0.5_rt*dy, z+0.5_rt*dz )
-          - IntegratedPotential( x-0.5_rt*dx, y+0.5_rt*dy, z+0.5_rt*dz )
-          - IntegratedPotential( x+0.5_rt*dx, y-0.5_rt*dy, z+0.5_rt*dz )
-          + IntegratedPotential( x-0.5_rt*dx, y-0.5_rt*dy, z+0.5_rt*dz )
-          - IntegratedPotential( x+0.5_rt*dx, y+0.5_rt*dy, z-0.5_rt*dz )
-          + IntegratedPotential( x-0.5_rt*dx, y+0.5_rt*dy, z-0.5_rt*dz )
-          + IntegratedPotential( x+0.5_rt*dx, y-0.5_rt*dy, z-0.5_rt*dz )
-          - IntegratedPotential( x-0.5_rt*dx, y-0.5_rt*dy, z-0.5_rt*dz )
-        );
 
-        return G_value;
+    /** @brief add
+     *
+     * @param[in] x x-coordinate of given location
+     * @param[in] y y-coordinate of given location
+     * @param[in] dx cell size along x
+     * @param[in] dy cell size along y
+     *
+     * @return the sum of integrated Green function G in 2D
+     */
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    amrex::Real
+    SumOfIntegratedPotential2D (amrex::Real x, amrex::Real y, amrex::Real dx, amrex::Real dy)
+    {
+        return 1._rt/(4._rt*ablastr::constant::math::pi*ablastr::constant::SI::ep0) * (
+                  IntegratedPotential2D( x+0.5_rt*dx, y+0.5_rt*dy )
+                - IntegratedPotential2D( x+0.5_rt*dx, y-0.5_rt*dy )
+                - IntegratedPotential2D( x-0.5_rt*dx, y+0.5_rt*dy )
+                + IntegratedPotential2D( x-0.5_rt*dx, y-0.5_rt*dy )
+                );
     }
 
+
     /** @brief Compute the electrostatic potential using the Integrated Green Function method
      *         as in http://dx.doi.org/10.1103/PhysRevSTAB.9.044204
      *
@@ -90,7 +129,8 @@ namespace ablastr::fields
     computePhiIGF (amrex::MultiFab const & rho,
                    amrex::MultiFab & phi,
                    std::array<amrex::Real, 3> const & cell_size,
-                   amrex::BoxArray const & ba);
+                   amrex::BoxArray const & ba,
+                   bool do_2d_slices);
 
 } // namespace ablastr::fields
 

Source/ablastr/fields/IntegratedGreenFunctionSolver.cpp

@@ -34,7 +34,8 @@ void
 computePhiIGF ( amrex::MultiFab const & rho,
                 amrex::MultiFab & phi,
                 std::array<amrex::Real, 3> const & cell_size,
-                amrex::BoxArray const & ba)
+                amrex::BoxArray const & ba,
+                bool const do_2d_slices)
 {
     using namespace amrex::literals;
 
@@ -45,9 +46,6 @@ computePhiIGF ( amrex::MultiFab const & rho,
     domain.surroundingNodes(); // get nodal points, since `phi` and `rho` are nodal
     domain.grow( phi.nGrowVect() ); // include guard cells
 
-    // Do we grow the domain in the z-direction in the 2D mode?
-    bool const do_2d_fft = false;
-
     int nprocs = amrex::ParallelDescriptor::NProcs();
     {
         amrex::ParmParse pp("ablastr");
@@ -61,7 +59,7 @@ computePhiIGF ( amrex::MultiFab const & rho,
     }
     if (!obc_solver || obc_solver->Domain() != domain) {
         amrex::FFT::Info info{};
-        if (do_2d_fft) { info.setBatchMode(true); }
+        if (do_2d_slices) { info.setBatchMode(true); } // do 2D FFTs
         info.setNumProcs(nprocs);
         obc_solver = std::make_unique<amrex::FFT::OpenBCSolver<amrex::Real>>(domain, info);
     }
@@ -71,7 +69,9 @@ computePhiIGF ( amrex::MultiFab const & rho,
     amrex::Real const dy = cell_size[1];
     amrex::Real const dz = cell_size[2];
 
-    obc_solver->setGreensFunction(
+    if (!do_2d_slices){
+        // 2D sliced solver
+        obc_solver->setGreensFunction(
         [=] AMREX_GPU_DEVICE (int i, int j, int k) -> amrex::Real
         {
             int const i0 = i - lo[0];
@@ -80,9 +80,26 @@ computePhiIGF ( amrex::MultiFab const & rho,
             amrex::Real const x = i0*dx;
             amrex::Real const y = j0*dy;
             amrex::Real const z = k0*dz;
-            return SumOfIntegratedPotential(x, y, z, dx, dy, dz);
+
+            return SumOfIntegratedPotential3D(x, y, z, dx, dy, dz);
+        });
+    }else{
+        // fully 3D solver
+        obc_solver->setGreensFunction(
+        [=] AMREX_GPU_DEVICE (int i, int j, int k) -> amrex::Real
+        {
+            int const i0 = i - lo[0];
+            int const j0 = j - lo[1];
+            amrex::Real const x = i0*dx;
+            amrex::Real const y = j0*dy;
+            amrex::ignore_unused(k);
+
+            return SumOfIntegratedPotential2D(x, y, dx, dy);
         });
 
+    }
+
     obc_solver->solve(phi, rho);
-}
+} // computePhiIGF
+
 } // namespace ablastr::fields