2.5 Space Charge in Particle Tracking

docs/source/usage/parameters.rst

@@ -903,6 +903,7 @@ See there ``nslice`` option on lattice elements for slicing.
     When running in envelope mode (when ``algo.track = "envelope"``), this model currently assumes that ``<xy> = <yt> = <tx> = 0``.
 
   * ``"Gauss3D"``: Calculate 3D space charge forces as if the beam was a Gaussian distribution.
+
   * ``"Gauss2p5D"``: Calculate 2.5D space charge forces as if the beam was a transverse Gaussian distribution.
 
     These models are supported only in particle tracking mode (when ``algo.track = "particles"``).
@@ -923,6 +924,17 @@ See there ``nslice`` option on lattice elements for slicing.
 
       Number of bins for longitudinal line density deposition.
 
+  * ``"2p5D"``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, while the transverse space charge kicks are weighted by the
+longitudinal line density determined by charge deposition (2.5D model).  Longitudinal space charge kicks are determined by the derivative of the line charge density.
+
+    This model is supported only in particle tracking mode (when ``algo.track = "particles"``).
+
+    This model supports the following sub-option:
+
+    * ``algo.space_charge.charge_z_bins`` (``int``, default: ``129``)
+
+      Number of bins for longitudinal line density deposition.
+
 * ``amr.n_cell`` (3 integers) optional (default: 1 `blocking_factor <https://amrex-codes.github.io/amrex/docs_html/GridCreation.html>`__ per MPI process)
 
   The number of grid points along each direction (on the **coarsest level**)

docs/source/usage/python.rst

@@ -75,6 +75,7 @@ Collective Effects & Overall Simulation Parameters
         When running in envelope mode (when ``algo.track = "envelope"``), this model currently assumes that ``<xy> = <yt> = <tx> = 0``.
 
       * ``"Gauss3D"``: Calculate 3D space charge forces as if the beam was a Gaussian distribution.
+
       * ``"Gauss2p5D"``: Calculate 2.5D space charge forces as if the beam was a transverse Gaussian distribution.
 
         These models are supported only in particle tracking mode (when ``algo.track = "particles"``).
@@ -93,6 +94,15 @@ Collective Effects & Overall Simulation Parameters
 
       Number of bins for longitudinal charge density deposition (default: ``129``).
 
+      * ``"2p5D"``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, while the transverse space charge kicks are weighted by the
+        longitudinal line density determined by charge deposition (2.5D model).  Longitudinal space charge kicks are determined by the derivative of the line charge density.
+
+        These models are supported only in particle tracking mode (when ``algo.track = "particles"``).
+
+   .. py:property:: space_charge_z_bins
+
+      Number of bins for longitudinal charge density deposition (default: ``129``).
+
    .. py:property:: poisson_solver
 
       The numerical solver to solve the Poisson equation when calculating space charge effects.

src/initialization/Algorithms.H

@@ -23,7 +23,8 @@ namespace impactx
         True_3D,  /**< 3D beam distribution */
         Gauss3D,  /**< Assume a 3D Gaussian beam distribution */
         Gauss2p5D,  /**< Assume a transverse 2D Gaussian beam distribution */
-        True_2D   /**< Averaged 2D transverse beam distribution with a current along s */
+        True_2D,   /**< Averaged 2D transverse beam distribution */
+        True_2p5D   /**< Averaged 2D transverse beam distribution with a current along s */
     );
 
     /** Return the currently active space charge algorithm */

src/initialization/Algorithms.cpp

@@ -74,6 +74,10 @@ namespace impactx
         {
             return SpaceChargeAlgo::True_2D;
         }
+        else if (space_charge == "2p5D")
+        {
+            return SpaceChargeAlgo::True_2p5D;
+        }
         else
         {
             throw std::runtime_error("algo.space_charge = " + space_charge + " is not a valid option");

src/initialization/AmrCoreData.cpp

@@ -150,7 +150,7 @@ namespace impactx::initialization
         int const num_components_phi = 1;
         amrex::IntVect num_guards_phi{num_guards_rho + 1}; // todo: I think this just depends on max(MLMG, force calc)
         amrex::BoxArray phi_ba = cba;
-        if (space_charge == SpaceChargeAlgo::True_2D) {
+        if (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D) {
             num_guards_phi[2] = 0;
             amrex::BoxList bl(amrex::IndexType{rho_nodal_flag});
             bl.reserve(cba.size());
@@ -169,7 +169,7 @@ namespace impactx::initialization
 
         // space charge force
         amrex::IntVect num_guards_force(num_guards_rho);
-        if (space_charge == SpaceChargeAlgo::True_2D) { num_guards_force[2] = 0; }
+        if (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D) { num_guards_force[2] = 0; }
         std::unordered_map<std::string, amrex::MultiFab> f_comp;
         for (std::string const comp : {"x", "y", "z"})
         {

src/initialization/InitDistribution.cpp

@@ -591,7 +591,7 @@ namespace impactx
 
             amrex::ParticleReal intensity = 0.0; // bunch charge (C) for 3D model, beam current (A) for 2D model
 
-            if (space_charge == SpaceChargeAlgo::True_3D)
+            if (space_charge == SpaceChargeAlgo::True_3D || space_charge == SpaceChargeAlgo::True_2p5D)
             {
                 pp_dist.get("charge", intensity);
                 amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist, intensity);

src/particles/spacecharge/ForceFromSelfFields.cpp

@@ -56,7 +56,7 @@ namespace impactx::particles::spacecharge
                 auto scf_arr_y = space_charge_field[lev]["y"][mfi].array();
                 auto scf_arr_z = space_charge_field[lev]["z"][mfi].array();
 
-                if (space_charge == SpaceChargeAlgo::True_2D) {
+                if (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D) {
                     AMREX_ALWAYS_ASSERT_WITH_MESSAGE(bx.size()[2] == 1,
                                                      "2D space charge requires exactly 1 slice in z");
                     amrex::ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int ) noexcept {

src/particles/spacecharge/GatherAndPush.cpp

@@ -10,13 +10,15 @@
 #include "GatherAndPush.H"
 
 #include "initialization/Algorithms.H"
+#include "particles/wakefields/ChargeBinning.H"
 
 #include <ablastr/particles/NodalFieldGather.H>
 
 #include <AMReX_BLProfiler.H>
 #include <AMReX_REAL.H>       // for Real
 #include <AMReX_SPACE.H>      // for AMREX_D_DECL
-
+#include <AMReX_GpuContainers.H>
+#include <AMReX_ParmParse.H>
 
 namespace impactx::particles::spacecharge
 {
@@ -87,8 +89,6 @@
                     auto prob_lo_2D = gm.ProbLoArray();
                     prob_lo_2D[2] = 0.0_rt;
 
-                    // TODO: add in z-dependent scaling by current?
-
                     amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
                         // access SoA Real data
                         amrex::ParticleReal & AMREX_RESTRICT x = part_x[i];
@@ -111,11 +111,93 @@
                         px += field_interp[0] * push_consts * dr[2] / (beta * c0_SI);
                         py += field_interp[1] * push_consts * dr[2] / (beta * c0_SI);
                         pz += 0.0_rt;
-                        //pz += field_interp[2] * push_consts;  // TODO: non-zero in 2.5D, but we will add a toggle to turn it off there, too
 
                         // push position is done in the lattice elements
                     });
                 }
+                if (space_charge == SpaceChargeAlgo::True_2p5D) {
+                    // flatten 3rd dimension
+                    auto prob_lo_2D = gm.ProbLoArray();
+                    prob_lo_2D[2] = 0.0_rt;
+
+                    // Calculate z-dependent scaling by current
+                    int tp5d_bins = 129;
+                    amrex::ParmParse pp_algo("algo.space_charge");
+                    pp_algo.queryAddWithParser("gauss_charge_z_bins", tp5d_bins);
+
+                    // Set parameters for charge deposition
+                    bool const is_unity_particle_weight = false;
+                    bool const GetNumberDensity = true;
+
+                    // Measure beam size, extract the min, max of particle positions
+                    [[maybe_unused]] auto const [x_min, y_min, t_min, x_max, y_max, t_max] =
+                    pc.MinAndMaxPositions();
+
+                    int const num_bins = tp5d_bins;  // Set resolution
+                    amrex::Real const bin_min = t_min;
+                    amrex::Real const bin_max = t_max;
+                    amrex::Real const bin_size = (bin_max - bin_min) / (num_bins - 1);  // number of evaluation points
+                    // Allocate memory for the charge profile
+                    amrex::Gpu::DeviceVector<amrex::Real> charge_distribution(num_bins + 1, 0.0);
+                    // Call charge deposition function
+                    impactx::particles::wakefields::DepositCharge1D(pc, charge_distribution, bin_min, bin_size, is_unity_particle_weight);
+
+                    // Sum up all partial charge histograms to each MPI process to calculate
+                    // the global charge slope.
+                    amrex::ParallelAllReduce::Sum(
+                        charge_distribution.data(),
+                        charge_distribution.size(),
+                        amrex::ParallelDescriptor::Communicator()
+                    );
+
+                    // Call charge density derivative function
+                    amrex::Gpu::DeviceVector<amrex::Real> slopes(charge_distribution.size() - 1, 0.0);
+                    impactx::particles::wakefields::DerivativeCharge1D(charge_distribution, slopes, bin_size,GetNumberDensity);
+
+                    amrex::Real const * const beam_profile_slope = slopes.data();
+                    amrex::Real const * const beam_profile = charge_distribution.data();
+
+                    // group together constants for the momentum push
+                    amrex::ParticleReal const chargesign = charge / std::abs(charge);
+
+                    amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
+                        // access SoA Real data
+                        amrex::ParticleReal & AMREX_RESTRICT x = part_x[i];
+                        amrex::ParticleReal & AMREX_RESTRICT y = part_y[i];
+                        amrex::ParticleReal z = 0.0_prt;  // flatten 3rd dimension
+                        amrex::ParticleReal & AMREX_RESTRICT px = part_px[i];
+                        amrex::ParticleReal & AMREX_RESTRICT py = part_py[i];
+                        amrex::ParticleReal & AMREX_RESTRICT pz = part_pz[i];
+
+                        // force gather
+                        amrex::GpuArray<amrex::Real, 3> const field_interp =
+                            ablastr::particles::doGatherVectorFieldNodal<2>(
+                                x, y, z,
+                                scf_arr_x, scf_arr_y, scf_arr_z,
+                                invdr,
+                                prob_lo_2D
+                            );
+
+                       // Update momentae with the 2.5D SC force
+                       int const idx = static_cast<int>((z - bin_min) / bin_size);  // Find index position along z
+                       #if (defined(AMREX_DEBUG) || defined(DEBUG)) && !defined(AMREX_USE_GPU)
+                       if (idx < 0 || idx >= num_bins)
+                       {
+                            std::cerr << "Warning: Index out of range for 2.5D SC: " << idx << std::endl;
+                       }
+                       #endif
+                       [[maybe_unused]] amrex::ParticleReal const Fxy = beam_profile[idx] * chargesign;
+                       [[maybe_unused]] amrex::ParticleReal const Fz = beam_profile_slope[idx] * charge;
+
+                       // push momentum
+                       px += field_interp[0] * Fxy * push_consts;
+                       py += field_interp[1] * Fxy * push_consts;
+                       pz += 0.0_rt;
+                       //pz -= (eintz + pz_push_const) * Fz * push_consts;
+
+                    // push position is done in the lattice elements
+                    });
+                }
                 if (space_charge == SpaceChargeAlgo::True_3D) {
                     amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) {
                         // access SoA Real data

src/particles/spacecharge/HandleSpacecharge.cpp

@@ -43,7 +43,7 @@ namespace impactx::particles::spacecharge
         // turn off if less than 2 particles
         if (amr_data->track_particles.m_particle_container->TotalNumberOfParticles(true, false) < 2) { return; }
 
-        if (space_charge != SpaceChargeAlgo::True_2D)
+        if (space_charge != SpaceChargeAlgo::True_2D && space_charge != SpaceChargeAlgo::True_2p5D)
         {
             // transform from x',y',t to x,y,z
             transformation::CoordinateTransformation(
@@ -60,7 +60,7 @@ namespace impactx::particles::spacecharge
         {
             Gauss2p5dPush(*amr_data->track_particles.m_particle_container, slice_ds);
         }
-        else if (space_charge == SpaceChargeAlgo::True_3D || space_charge == SpaceChargeAlgo::True_2D)
+        else if (space_charge == SpaceChargeAlgo::True_3D || space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D)
         {
             // Note: The following operations assume that
             // the particles are in x, y, z coordinates.
@@ -79,8 +79,6 @@ namespace impactx::particles::spacecharge
                 amr_data->refRatio()
             );
 
-            // TODO for 2.5D: deposit charge/current in 1D array
-
             // poisson solve in x,y,z
             spacecharge::PoissonSolve(
                 *amr_data->track_particles.m_particle_container,
@@ -107,7 +105,7 @@ namespace impactx::particles::spacecharge
             );
         }
 
-        if (space_charge != SpaceChargeAlgo::True_2D)
+        if (space_charge != SpaceChargeAlgo::True_2D && space_charge != SpaceChargeAlgo::True_2p5D)
         {
             // transform from x,y,z to x',y',t
             transformation::CoordinateTransformation(

src/particles/spacecharge/PoissonSolve.cpp

@@ -70,6 +70,9 @@ namespace impactx::particles::spacecharge
         if (space_charge == SpaceChargeAlgo::True_2D && poisson_solver != "fft") {
             throw std::runtime_error("algo.poisson_solver must be fft for SpaceChargeAlgo::True_2D");
         }
+        if (space_charge == SpaceChargeAlgo::True_2p5D && poisson_solver != "fft") {
+            throw std::runtime_error("algo.poisson_solver must be fft for SpaceChargeAlgo::True_2p5D");
+        }
 
         // MLMG options
         amrex::Real mlmg_relative_tolerance = 1.e-7; // relative TODO: make smaller for SP
@@ -84,7 +87,7 @@ namespace impactx::particles::spacecharge
 
         // flatten rho to 2D
         std::unordered_map<int, std::pair<amrex::MultiFab, amrex::MultiFab>> rho_2d;  // pair: local & unique boxes
-        if (space_charge == SpaceChargeAlgo::True_2D) {
+        if (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D) {
             auto geom_3d = pc.GetParGDB()->Geom();
             amrex::Box domain_3d = geom_3d[0].Domain();  // whole simulation index space (level 0)
             rho_2d = flatten_charge_to_2D(rho, domain_3d);
@@ -98,7 +101,7 @@ namespace impactx::particles::spacecharge
 
         // create phi_2d and sort rho/phi pointers
         for (int lev = 0; lev <= finest_level; ++lev) {
-            if (space_charge == SpaceChargeAlgo::True_2D) {
+            if (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D) {
                 int nz = pc.GetParGDB()->Geom(lev).Domain().length(2);
                 if (nz == 1) {
                     sorted_phi.emplace_back(&phi[lev]);
@@ -119,7 +122,7 @@ namespace impactx::particles::spacecharge
             }
         }
 
-        const bool is_igf_2d = space_charge == SpaceChargeAlgo::True_2D;
+        const bool is_igf_2d = (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D);
         const bool do_single_precision_comms = false;
         const bool eb_enabled = false;
         ablastr::fields::computePhi(
@@ -154,7 +157,7 @@ namespace impactx::particles::spacecharge
         }
 
         // We may need to copy phi from phi_2d
-        if (space_charge == SpaceChargeAlgo::True_2D) {
+        if (space_charge == SpaceChargeAlgo::True_2D || space_charge == SpaceChargeAlgo::True_2p5D) {
             for (int lev=0; lev<=finest_level; lev++) {
                 if (&(phi[lev]) != sorted_phi[lev]) {
                     phi[lev].ParallelCopy(*sorted_phi[lev]);

src/python/ImpactX.cpp

@@ -285,14 +285,14 @@ void init_ImpactX (py::module& m)
                 else
                 {
                     std::string const space_charge = std::get<std::string>(space_charge_v);
-                    if (space_charge != "false" && space_charge != "off" && space_charge != "2D" && space_charge != "3D" && space_charge != "Gauss3D" && space_charge != "Gauss2p5D" ) {
-                        throw std::runtime_error("Space charge model must be 2D, 3D, Gauss3D or Gauss2p5D but is: " + space_charge);
+                    if (space_charge != "false" && space_charge != "off" && space_charge != "2D" && space_charge != "3D" && space_charge != "Gauss3D" && space_charge != "Gauss2p5D"  && space_charge != "2p5D") {
+                        throw std::runtime_error("Space charge model must be 2D, 3D, Gauss3D, Gauss2p5D, or 2p5D but is: " + space_charge);
                     }
                     amrex::ParmParse pp_algo("algo");
                     pp_algo.add("space_charge", space_charge);
                 }
             },
-            "The model to be used when calculating space charge effects. Either off, 2D, or 3D."
+            "The model to be used when calculating space charge effects. Either off, 2D, 3D, Gauss3D, Gauss2p5D, or 2p5D."
         )
         .def_property("space_charge_gauss_nint",
             [](ImpactX & /* ix */) {
@@ -320,7 +320,7 @@ void init_ImpactX (py::module& m)
                 amrex::ParmParse pp_algo("algo.space_charge");
                 pp_algo.add("gauss_charge_z_bins", gauss_charge_z_bins);
             },
-            "Number of steps for computing the integrals (default: ``129``)."
+            "Number of longitudinal bins for computing the linear charge density (default: ``129``)."
         )
         .def_property("space_charge_gauss_taylor_delta",
             [](ImpactX & /* ix */) {