diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 6e90e8951..8986100f0 100644 --- a/docs/source/usage/parameters.rst +++ b/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 `` = = = 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 `__ per MPI process) The number of grid points along each direction (on the **coarsest level**) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index d02cfdd1c..0f93b0f96 100644 --- a/docs/source/usage/python.rst +++ b/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 `` = = = 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. diff --git a/src/initialization/Algorithms.H b/src/initialization/Algorithms.H index 3a8efdd89..4d72a0704 100644 --- a/src/initialization/Algorithms.H +++ b/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 */ diff --git a/src/initialization/Algorithms.cpp b/src/initialization/Algorithms.cpp index cbc780638..1a91a61a6 100644 --- a/src/initialization/Algorithms.cpp +++ b/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"); diff --git a/src/initialization/AmrCoreData.cpp b/src/initialization/AmrCoreData.cpp index 23fdd691a..ba8ff44c0 100644 --- a/src/initialization/AmrCoreData.cpp +++ b/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 f_comp; for (std::string const comp : {"x", "y", "z"}) { diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index b5c2f40e0..12fa994af 100644 --- a/src/initialization/InitDistribution.cpp +++ b/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); diff --git a/src/particles/spacecharge/CMakeLists.txt b/src/particles/spacecharge/CMakeLists.txt index 46b36595d..8ea8e919e 100644 --- a/src/particles/spacecharge/CMakeLists.txt +++ b/src/particles/spacecharge/CMakeLists.txt @@ -4,6 +4,7 @@ target_sources(lib Gauss3dPush.cpp Gauss2p5dPush.cpp GatherAndPush.cpp + Deposit1D.cpp HandleSpacecharge.cpp PoissonSolve.cpp ) diff --git a/src/particles/spacecharge/Deposit1D.H b/src/particles/spacecharge/Deposit1D.H new file mode 100644 index 000000000..9ebe74892 --- /dev/null +++ b/src/particles/spacecharge/Deposit1D.H @@ -0,0 +1,44 @@ +/* Copyright 2022-2025 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Ji Qiang + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACT_DEPOSIT1D_PUSH_H +#define IMPACT_DEPOSIT1D_PUSH_H + +#include "particles/ImpactXParticleContainer.H" + +#include + + +namespace impactx::particles::spacecharge +{ + /** Deposit charge using longitudinal binning for the calculation of space charge + * + * This function performs particle binning for particles in the beam along the + * longitudinal direction, for the purposes of computing space charge in + * so-called 2.5D models. + * + * @param[in] pc container of the particles that deposited rho + * @param[out] beam_profile deposited charge density + * @param[out] beam_profile_slope derivative of deposited charge density + * @param[in] bin_min location of lower endpoint for binning + * @param[in] bin_max location of upper endpoint for binning + * @param[in] num_bins number of longitudinal bins + */ + void Deposit1D ( + ImpactXParticleContainer & pc, + amrex::Real * beam_profile, + amrex::Real * beam_profile_slope, + amrex::Real bin_min, + amrex::Real bin_max, + int num_bins + ); + +} // namespace impactx::particles::spacecharge + +#endif // IMPACT_DEPOSIT1D_H diff --git a/src/particles/spacecharge/Deposit1D.cpp b/src/particles/spacecharge/Deposit1D.cpp new file mode 100644 index 000000000..a6007cfd3 --- /dev/null +++ b/src/particles/spacecharge/Deposit1D.cpp @@ -0,0 +1,66 @@ +/* Copyright 2022-2025 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Ji Qiang + * License: BSD-3-Clause-LBNL + */ +#include "Deposit1D.H" + +#include "diagnostics/ReducedBeamCharacteristics.H" +#include "particles/wakefields/ChargeBinning.H" + +#include +#include +#include +#include + +#include + + +namespace impactx::particles::spacecharge +{ + + void Deposit1D ( + ImpactXParticleContainer & pc, + [[maybe_unused]] amrex::Real * beam_profile, + [[maybe_unused]] amrex::Real * beam_profile_slope, + amrex::Real bin_min, + amrex::Real bin_max, + int num_bins + ) + { + BL_PROFILE("impactx::spacecharge::Deposit1D"); + + using namespace amrex::literals; + + // Set parameters for charge deposition + bool const is_unity_particle_weight = false; + bool const GetNumberDensity = true; + + 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 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 slopes(charge_distribution.size() - 1, 0.0); + impactx::particles::wakefields::DerivativeCharge1D(charge_distribution, slopes, bin_size,GetNumberDensity); + + beam_profile = charge_distribution.data(); + beam_profile_slope = slopes.data(); + + } + +} // namespace impactx::particles::spacecharge diff --git a/src/particles/spacecharge/ForceFromSelfFields.cpp b/src/particles/spacecharge/ForceFromSelfFields.cpp index 88c02bac3..4fa708961 100644 --- a/src/particles/spacecharge/ForceFromSelfFields.cpp +++ b/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 { diff --git a/src/particles/spacecharge/GatherAndPush.cpp b/src/particles/spacecharge/GatherAndPush.cpp index c2b96923b..eb0bd4f76 100644 --- a/src/particles/spacecharge/GatherAndPush.cpp +++ b/src/particles/spacecharge/GatherAndPush.cpp @@ -10,13 +10,16 @@ #include "GatherAndPush.H" #include "initialization/Algorithms.H" +#include "particles/wakefields/ChargeBinning.H" +#include "particles/spacecharge/Deposit1D.H" #include #include #include // for Real #include // for AMREX_D_DECL - +#include +#include namespace impactx::particles::spacecharge { @@ -35,6 +38,23 @@ namespace impactx::particles::spacecharge amrex::ParticleReal const charge = pc.GetRefParticle().charge; + // Deposit 1D charge density in cases where it is required. + amrex::Real * beam_profile = nullptr; + amrex::Real * beam_profile_slope = nullptr; + int num_bins = 129; + + [[maybe_unused]] auto const [x_min, y_min, t_min, x_max, y_max, t_max] = + pc.MinAndMaxPositions(); + + amrex::Real bin_min = t_min; + amrex::Real bin_max = t_max; + amrex::Real const bin_size = (bin_max - bin_min) / (num_bins - 1); + + if (space_charge == SpaceChargeAlgo::True_2p5D) { + + Deposit1D ( pc, beam_profile, beam_profile_slope, bin_min, bin_max, num_bins); + } + // loop over refinement levels int const nLevel = pc.finestLevel(); for (int lev = 0; lev <= nLevel; ++lev) @@ -87,8 +107,6 @@ namespace impactx::particles::spacecharge 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 +129,56 @@ namespace impactx::particles::spacecharge 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; + + // 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 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((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 diff --git a/src/particles/spacecharge/HandleSpacecharge.cpp b/src/particles/spacecharge/HandleSpacecharge.cpp index b8b23a044..00480138f 100644 --- a/src/particles/spacecharge/HandleSpacecharge.cpp +++ b/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( diff --git a/src/particles/spacecharge/PoissonSolve.cpp b/src/particles/spacecharge/PoissonSolve.cpp index 4d367e3d6..62eaff25b 100644 --- a/src/particles/spacecharge/PoissonSolve.cpp +++ b/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> 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]); diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index b50876dd0..351516523 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -285,14 +285,14 @@ void init_ImpactX (py::module& m) else { std::string const space_charge = std::get(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 */) {