Implement particle suborbits for the implicit solver (BLAST-WarpX#5969)

dpgrote · web-flow · commit 350b568bc0d7 · 2025-09-03T22:25:20.000Z
In the implicit solver, a Picard iteration is done for the particle advance. Occasionally, there will be a few particles that fail to converge (for example reaching a bistable state). A solution is to do sub-orbits during the advance, dividing the time step into smaller steps, with the Picard iteration done on each substep. This PR implements this scheme. The PR is ready for merging accept for adding a CI test. A full test of this would require a full implicit simulation with large time steps (where the need for the sub-orbits arises). A CI test could be added at a later point when the implicit capability is fully implemented. Merge first: - [x] BLAST-WarpX#6021 - [x] BLAST-WarpX#6116 - [x] BLAST-WarpX#6120
diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst
@@ -149,6 +149,8 @@ Overall simulation parameters
 
           - ``implicit_evolve.max_particle_iterations`` (`integer`, default: 21)
           - ``implicit_evolve.particle_tolerance`` (`float`, default: 1.e-10)
+          - ``implicit_evolve.particle_suborbits`` (`bool`, default: false)
+          - ``implicit_evolve.print_unconverged_particle_details`` (`bool`, default: false)
 
         - ``implicit_evolve.use_mass_matrices_jacobian`` (`bool`, default: false).
           When `true`, the plasma current density is computed using the mass matrices during the linear stage of PS-JFNK, replacing direct particle calculations. This can enable large speed ups for simulations with many particles.
diff --git a/Source/FieldSolver/ImplicitSolvers/ImplicitOptions.H b/Source/FieldSolver/ImplicitSolvers/ImplicitOptions.H
@@ -7,6 +7,7 @@ using namespace amrex::literals;
 struct ImplicitOptions {
     int max_particle_iterations = 21;
     amrex::ParticleReal particle_tolerance = 1.0e-10_prt;
+    bool print_unconverged_particle_details = false;
     bool deposit_mass_matrices = false;
     bool linear_stage_of_jfnk = false;
     int nonlinear_iteration = 0;
diff --git a/Source/FieldSolver/ImplicitSolvers/ImplicitSolver.H b/Source/FieldSolver/ImplicitSolvers/ImplicitSolver.H
@@ -54,6 +54,8 @@ public:
 
     void CreateParticleAttributes () const;
 
+    bool DoParticleSuborbits () { return m_particle_suborbits; }
+
     /**
      * \brief Advance fields and particles by one time step using the specified implicit algorithm
      */
@@ -148,6 +150,18 @@ protected:
      */
     int m_max_particle_iterations = 21;
 
+    /**
+     * \brief whether to use suborbits for particles that fail to converge in
+     * m_max_particle_iterations iterations
+     */
+    bool m_particle_suborbits = false;
+
+    /**
+     * \brief whether the details of unconverged particles are printed out during
+     * the particle evolve loops
+     */
+    bool m_print_unconverged_particle_details = false;
+
     /**
      * \brief Whether to use mass matrices for the implicit solver
      */
diff --git a/Source/FieldSolver/ImplicitSolvers/ImplicitSolver.cpp b/Source/FieldSolver/ImplicitSolvers/ImplicitSolver.cpp
@@ -27,6 +27,10 @@ void ImplicitSolver::CreateParticleAttributes () const
         pc->AddRealComp("ux_n", comm);
         pc->AddRealComp("uy_n", comm);
         pc->AddRealComp("uz_n", comm);
+
+        if (m_particle_suborbits) {
+            pc->AddIntComp("nsuborbits", comm);
+        }
     }
 }
 
@@ -429,7 +433,9 @@ void ImplicitSolver::parseNonlinearSolverParams ( const amrex::ParmParse&  pp )
         m_nlsolver_type = NonlinearSolverType::Newton;
         m_nlsolver = std::make_unique<NewtonSolver<WarpXSolverVec,ImplicitSolver>>();
         pp.query("max_particle_iterations", m_max_particle_iterations);
+        pp.query("particle_suborbits", m_particle_suborbits);
         pp.query("particle_tolerance", m_particle_tolerance);
+        pp.query("print_unconverged_particle_details", m_print_unconverged_particle_details);
         pp.query("use_mass_matrices_jacobian", m_use_mass_matrices_jacobian);
         pp.query("use_mass_matrices_pc", m_use_mass_matrices_pc);
         if (m_use_mass_matrices_jacobian || m_use_mass_matrices_pc) {
diff --git a/Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.cpp b/Source/FieldSolver/ImplicitSolvers/ThetaImplicitEM.cpp
@@ -73,6 +73,7 @@ void ThetaImplicitEM::PrintParameters () const
     amrex::Print() << "Time-bias parameter theta:  " << m_theta << "\n";
     amrex::Print() << "max particle iterations:    " << m_max_particle_iterations << "\n";
     amrex::Print() << "particle tolerance:         " << m_particle_tolerance << "\n";
+    amrex::Print() << "use particle suborbits:     " << (m_particle_suborbits ? "true":"false") << "\n";
     if (m_nlsolver_type==NonlinearSolverType::Picard) {
         amrex::Print() << "Nonlinear solver type:      Picard\n";
     }
diff --git a/Source/FieldSolver/ImplicitSolvers/WarpXImplicitOps.cpp b/Source/FieldSolver/ImplicitSolvers/WarpXImplicitOps.cpp
@@ -169,6 +169,9 @@ WarpX::SaveParticlesAtImplicitStepStart ( )
                 amrex::ParticleReal* uy_n = pti.GetAttribs("uy_n").dataPtr();
                 amrex::ParticleReal* uz_n = pti.GetAttribs("uz_n").dataPtr();
 
+                // Check if nsuborbits is present, and if so it is set to 1
+                int *nsuborbits = (pc->HasiAttrib("nsuborbits") ? pti.GetiAttribs("nsuborbits").dataPtr() : nullptr);
+
                 const long np = pti.numParticles();
 
                 amrex::ParallelFor( np, [=] AMREX_GPU_DEVICE (long ip)
@@ -190,6 +193,10 @@ WarpX::SaveParticlesAtImplicitStepStart ( )
                     uy_n[ip] = uy[ip];
                     uz_n[ip] = uz[ip];
 
+                    if (nsuborbits) {
+                        nsuborbits[ip] = 1;
+                    }
+
                 });
 
             }
diff --git a/Source/Particles/Deposition/CurrentDeposition.H b/Source/Particles/Deposition/CurrentDeposition.H
@@ -2267,6 +2267,10 @@ void doVillasenorDepositionShapeNImplicit ([[maybe_unused]]const amrex::Particle
         np_to_deposit,
         [=] AMREX_GPU_DEVICE (long const ip) {
 
+            // Skip particles with zero weight.
+            // This should only be the case for particles that will be suborbited.
+            if (wp[ip] == 0.) { return; }
+
 #if !defined(WARPX_DIM_3D)
             constexpr amrex::ParticleReal inv_c2 = 1._prt/(PhysConst::c*PhysConst::c);
 
diff --git a/Source/Particles/Deposition/MassMatricesDeposition.H b/Source/Particles/Deposition/MassMatricesDeposition.H
@@ -1630,6 +1630,11 @@ void doVillasenorJandSigmaDeposition ( [[maybe_unused]] const amrex::ParticleRea
     amrex::ParallelFor(
         np_to_deposit,
         [=] AMREX_GPU_DEVICE (long const ip) {
+
+            // Skip particles with zero weight.
+            // This should only be the case for particles that will be suborbited.
+            if (wp[ip] == 0.) { return; }
+
             amrex::ParticleReal xp_nph, yp_nph, zp_nph;
             GetPosition(ip, xp_nph, yp_nph, zp_nph);
 
diff --git a/Source/Particles/PhysicalParticleContainer.H b/Source/Particles/PhysicalParticleContainer.H
@@ -135,8 +135,31 @@ public:
                          long np_to_push,
                          int lev, int gather_lev,
                          amrex::Real dt, ScaleFields scaleFields,
+                         long & num_unconverged_particles,
+                         amrex::Gpu::DeviceVector<long> & unconverged_indices,
+                         amrex::Gpu::DeviceVector<amrex::ParticleReal> & saved_weights,
                          DtType a_dt_type=DtType::Full);
 
+    void ImplicitPushXPSubOrbits (WarpXParIter& pti, ablastr::fields::MultiFabRegister& fields,
+                                  amrex::FArrayBox const * exfab,
+                                  amrex::FArrayBox const * eyfab,
+                                  amrex::FArrayBox const * ezfab,
+                                  amrex::FArrayBox const * bxfab,
+                                  amrex::FArrayBox const * byfab,
+                                  amrex::FArrayBox const * bzfab,
+                                  ImplicitOptions const * implicit_options,
+                                  amrex::IntVect ngEB,
+                                  amrex::MultiFab * const jx,
+                                  amrex::MultiFab * const jy,
+                                  amrex::MultiFab * const jz,
+                                  long offset,
+                                  int lev, int gather_lev,
+                                  amrex::Real dt, ScaleFields scaleFields,
+                                  bool skip_deposition,
+                                  long num_unconverged_particles,
+                                  amrex::Gpu::DeviceVector<long> & unconverged_indices,
+                                  amrex::Gpu::DeviceVector<amrex::ParticleReal> & saved_weights);
+
     void PushP (int lev, amrex::Real dt,
                         const amrex::MultiFab& Ex,
                         const amrex::MultiFab& Ey,
diff --git a/Source/Particles/PhysicalParticleContainer.cpp b/Source/Particles/PhysicalParticleContainer.cpp
@@ -547,6 +547,12 @@ PhysicalParticleContainer::Evolve (ablastr::fields::MultiFabRegister& fields,
 
                 int e_is_nodal = Ex.is_nodal() and Ey.is_nodal() and Ez.is_nodal();
 
+                // Temporary data used in the implicit advance
+                amrex::Gpu::DeviceVector<long> unconverged_indices;
+                amrex::Gpu::DeviceVector<amrex::ParticleReal> saved_weights;
+                long num_unconverged_particles = 0;
+                long num_unconverged_particles_c = 0;
+
                 //
                 // Gather and push for particles not in the buffer
                 //
@@ -559,11 +565,14 @@ PhysicalParticleContainer::Evolve (ablastr::fields::MultiFabRegister& fields,
                            Ex.nGrowVect(), e_is_nodal,
                            0, np_to_push, lev, gather_lev, dt, ScaleFields(false), a_dt_type);
                 } else if (push_type == PushType::Implicit) {
+                    long const offset = 0;
                     ImplicitPushXP(pti, exfab, eyfab, ezfab,
                                    bxfab, byfab, bzfab,
                                    implicit_options,
                                    Ex.nGrowVect(),
-                                   0, np_to_push, lev, gather_lev, dt, ScaleFields(false), a_dt_type);
+                                   offset, np_to_push, lev, gather_lev, dt, ScaleFields(false),
+                                   num_unconverged_particles, unconverged_indices, saved_weights,
+                                   a_dt_type);
                 }
 
                 if (np_gather < np)
@@ -614,7 +623,9 @@ PhysicalParticleContainer::Evolve (ablastr::fields::MultiFabRegister& fields,
                                        implicit_options,
                                        cEx.nGrowVect(),
                                        nfine_gather, np-nfine_gather,
-                                       lev, lev-1, dt, ScaleFields(false), a_dt_type);
+                                       lev, lev-1, dt, ScaleFields(false),
+                                       num_unconverged_particles_c, unconverged_indices, saved_weights,
+                                       a_dt_type);
                     }
                 }
 
@@ -662,7 +673,54 @@ PhysicalParticleContainer::Evolve (ablastr::fields::MultiFabRegister& fields,
                                        np_to_deposit, np-np_to_deposit, thread_num,
                                        lev, lev-1, dt, relative_time, push_type);
                     }
-                } // end of "if electrostatic_solver_id == ElectrostaticSolverAlgo::None"
+                } // end of "if skip_deposition"
+
+                if (push_type == PushType::Implicit) {
+                    if (num_unconverged_particles > 0) {
+                        amrex::MultiFab * jx = fields.get(current_fp_string, Direction{0}, lev);
+                        amrex::MultiFab * jy = fields.get(current_fp_string, Direction{1}, lev);
+                        amrex::MultiFab * jz = fields.get(current_fp_string, Direction{2}, lev);
+                        long const offset = 0;
+                        ImplicitPushXPSubOrbits(pti, fields, exfab, eyfab, ezfab,
+                                                bxfab, byfab, bzfab,
+                                                implicit_options,
+                                                Ex.nGrowVect(),
+                                                jx, jy, jz,
+                                                offset, lev, gather_lev, dt, ScaleFields(false), skip_deposition,
+                                                num_unconverged_particles, unconverged_indices, saved_weights);
+                    }
+                    if (num_unconverged_particles_c > 0) {
+
+                        amrex::MultiFab & cEx = *fields.get(FieldType::Efield_cax, Direction{0}, lev);
+                        amrex::MultiFab & cEy = *fields.get(FieldType::Efield_cax, Direction{1}, lev);
+                        amrex::MultiFab & cEz = *fields.get(FieldType::Efield_cax, Direction{2}, lev);
+                        amrex::MultiFab & cBx = *fields.get(FieldType::Bfield_cax, Direction{0}, lev);
+                        amrex::MultiFab & cBy = *fields.get(FieldType::Bfield_cax, Direction{1}, lev);
+                        amrex::MultiFab & cBz = *fields.get(FieldType::Bfield_cax, Direction{2}, lev);
+
+                        // Data on the grid
+                        FArrayBox const* cexfab = &cEx[pti];
+                        FArrayBox const* ceyfab = &cEy[pti];
+                        FArrayBox const* cezfab = &cEz[pti];
+                        FArrayBox const* cbxfab = &cBx[pti];
+                        FArrayBox const* cbyfab = &cBy[pti];
+                        FArrayBox const* cbzfab = &cBz[pti];
+
+                        amrex::MultiFab * cjx = fields.get(FieldType::current_buf, Direction{0}, lev);
+                        amrex::MultiFab * cjy = fields.get(FieldType::current_buf, Direction{1}, lev);
+                        amrex::MultiFab * cjz = fields.get(FieldType::current_buf, Direction{2}, lev);
+
+                        long const offset = num_unconverged_particles;
+                        ImplicitPushXPSubOrbits(pti, fields, cexfab, ceyfab, cezfab,
+                                                cbxfab, cbyfab, cbzfab,
+                                                implicit_options,
+                                                cEx.nGrowVect(),
+                                                cjx, cjy, cjz,
+                                                offset, lev, lev-1, dt, ScaleFields(false), skip_deposition,
+                                                num_unconverged_particles_c, unconverged_indices, saved_weights);
+                    }
+                }
+
             } // end of "if do_not_push"
 
             if (has_rho && ! skip_deposition && ! do_not_deposit) {
diff --git a/Source/Particles/Pusher/ImplicitPushPX.cpp b/Source/Particles/Pusher/ImplicitPushPX.cpp
diff --git a/Source/Particles/WarpXParticleContainer.H b/Source/Particles/WarpXParticleContainer.H

Original file line number	Diff line number	Diff line change
`@@ -73,6 +73,7 @@ void ThetaImplicitEM::PrintParameters () const`
`73`	`73`	`amrex::Print() << "Time-bias parameter theta: " << m_theta << "\n";`
`74`	`74`	`amrex::Print() << "max particle iterations: " << m_max_particle_iterations << "\n";`
`75`	`75`	`amrex::Print() << "particle tolerance: " << m_particle_tolerance << "\n";`
	`76`	`+ amrex::Print() << "use particle suborbits: " << (m_particle_suborbits ? "true":"false") << "\n";`
`76`	`77`	`if (m_nlsolver_type==NonlinearSolverType::Picard) {`
`77`	`78`	`amrex::Print() << "Nonlinear solver type: Picard\n";`
`78`	`79`	`}`