diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 0e389a9ff..5db4dc058 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -879,6 +879,12 @@ 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. + + This model is supported only in particle tracking mode (when ``algo.track = "particles"``). + Ref.: J. Qiang et al., "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025). + (This reference describes both 3D and 2.5D models.) + * ``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 1049ec8bd..9851eb80f 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -76,6 +76,11 @@ 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. + + This model is supported only in particle tracking mode (when ``algo.track = "particles"``). + Ref.: J. Qiang et al., "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025). + (This reference describes both 3D and 2.5D models.) .. py:property:: poisson_solver The numerical solver to solve the Poisson equation when calculating space charge effects. diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 829bdafbb..64b1a2043 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1374,6 +1374,21 @@ add_impactx_test(FODO.envelope.sc.py OFF # no plot script yet ) +# FODO cell with Gaussian 3D space charge using particle tracking ############# +# +add_impactx_test(FODO.Gauss3d.sc + examples/fodo_space_charge/input_fodo_Gauss3D_sc.in + ON # ImpactX MPI-parallel + examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py + OFF # no plot script yet +) +add_impactx_test(FODO.Gauss3d.sc.py + examples/fodo_space_charge/run_fodo_Gauss3D_sc.py + ON # ImpactX MPI-parallel + examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py + OFF # no plot script yet +) + # A single bend with incoherent synchrotron radiation ###################### # # no space charge diff --git a/examples/fodo_space_charge/README.rst b/examples/fodo_space_charge/README.rst index 282aab230..9e6562649 100644 --- a/examples/fodo_space_charge/README.rst +++ b/examples/fodo_space_charge/README.rst @@ -51,3 +51,56 @@ We run the following script to analyze correctness: .. literalinclude:: analysis_fodo_envelope_sc.py :language: python3 :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_envelope_sc.py``. + + +.. _examples-fodo-envelope-sc-gaussian: + +FODO Cell with 3D Gaussian Space Charge Using Particle Tracking +=============================================================== + +This example illustrates a 1 nC electron beam with a kinetic energy of 100 MeV in a FODO cell, +with 3D Gaussian space charge included. The parameters are those described in: + +The purpose of this example is to illustrate the use of particle tracking mode with 3D space charge from a Gaussian density +distribution. + +The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell with small noticeable growth. + +In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values. + + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_fodo_Gauss3D_sc.py`` or +* ImpactX **executable** using an input file: ``impactx input_fodo_Gauss3D_sc.in`` + +For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python: Script + + .. literalinclude:: run_fodo_Gauss3D_sc.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_space_charge/run_fodo_Gauss3D_sc.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_fodo_Gauss3D_sc.in + :language: ini + :caption: You can copy this file from ``examples/fodo_space_charge/input_fodo_Gauss3D_sc.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_fodo_Gauss3D_sc.py`` + + .. literalinclude:: analysis_fodo_Gauss3D_sc.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py``. diff --git a/examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py b/examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py new file mode 100755 index 000000000..f7e292529 --- /dev/null +++ b/examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Ji Qiang +# License: BSD-3-Clause-LBNL +# + +import numpy as np +import openpmd_api as io +from scipy.stats import moment + + +def get_moments(beam): + """Calculate standard deviations of beam position & momenta + and emittance values + + Returns + ------- + sigx, sigy, sigt, emittance_x, emittance_y, emittance_t + """ + sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev. + sigpx = moment(beam["momentum_x"], moment=2) ** 0.5 + sigy = moment(beam["position_y"], moment=2) ** 0.5 + sigpy = moment(beam["momentum_y"], moment=2) ** 0.5 + sigt = moment(beam["position_t"], moment=2) ** 0.5 + sigpt = moment(beam["momentum_t"], moment=2) ** 0.5 + + epstrms = beam.cov(ddof=0) + emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5 + emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5 + emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5 + + return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t) + + +# initial/final beam +series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) +last_step = list(series.iterations)[-1] +initial = series.iterations[1].particles["beam"].to_df() +beam_final = series.iterations[last_step].particles["beam"] +final = beam_final.to_df() + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +print("Initial Beam:") +sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti = get_moments(initial) +print(f" sigx={sig_xi:e} sigy={sig_yi:e} sigt={sig_ti:e}") +print( + f" emittance_x={emittance_xi:e} emittance_y={emittance_yi:e} emittance_t={emittance_ti:e}" +) + +atol = 0.0 # ignored +rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti], + [7.515765e-05, 7.511883e-05, 9.997395e-4, 2.001510e-09, 1.999755e-09, 1.999289e-06], + rtol=rtol, + atol=atol, +) + + +print("") +print("Final Beam:") +sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf = get_moments(initial) +print(f" sigx={sig_xf:e} sigy={sig_yf:e} sigt={sig_tf:e}") +print( + f" emittance_x={emittance_xf:e} emittance_y={emittance_yf:e} emittance_t={emittance_tf:e}" +) + +atol = 0.0 # ignored +rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf], + [ + 7.51576586332169e-05, + 7.511883208451813e-05, + 0.0009997395499750136, + 2.0015106608723994e-09, + 1.999755254276969e-09, + 1.9992898444562777e-06, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/fodo_space_charge/input_fodo_Gauss3D_sc.in b/examples/fodo_space_charge/input_fodo_Gauss3D_sc.in new file mode 100644 index 000000000..c3be8b254 --- /dev/null +++ b/examples/fodo_space_charge/input_fodo_Gauss3D_sc.in @@ -0,0 +1,58 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 1.0e2 +beam.charge = 1.0e-9 +beam.particle = electron +beam.distribution = gaussian +beam.lambdaX = 3.9984884770e-5 +beam.lambdaY = beam.lambdaX +beam.lambdaT = 1.0e-3 +beam.lambdaPx = 2.6623538760e-5 +beam.lambdaPy = beam.lambdaPx +beam.lambdaPt = 2.0e-3 +beam.muxpx = -0.846574929020762 +beam.muypy = -beam.muxpx +beam.mutpt = 0.0 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift3 monitor +lattice.nslice = 25 + +monitor.type = beam_monitor +monitor.backend = h5 + +drift1.type = drift +drift1.ds = 0.25 + +quad1.type = quad +quad1.ds = 1.0 +quad1.k = 1.0 + +drift2.type = drift +drift2.ds = 0.5 + +quad2.type = quad +quad2.ds = 1.0 +quad2.k = -1.0 + +drift3.type = drift +drift3.ds = 0.25 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = Gauss3D + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/fodo_space_charge/run_fodo_Gauss3D_sc.py b/examples/fodo_space_charge/run_fodo_Gauss3D_sc.py new file mode 100755 index 000000000..d76b0904f --- /dev/null +++ b/examples/fodo_space_charge/run_fodo_Gauss3D_sc.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.particle_shape = 2 # B-spline order +sim.space_charge = "Gauss3D" +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# load a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 100 # reference energy +bunch_charge_C = 1.0e-9 # used with space charge +npart = 10000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Gaussian( + lambdaX=3.9984884770e-5, + lambdaY=3.9984884770e-5, + lambdaT=1.0e-3, + lambdaPx=2.6623538760e-5, + lambdaPy=2.6623538760e-5, + lambdaPt=2.0e-3, + muxpx=-0.846574929020762, + muypy=0.846574929020762, + mutpt=0.0, +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice) +ns = 25 # number of slices per ds in the element +fodo = [ + monitor, + elements.ChrDrift(name="drift1", ds=0.25, nslice=ns), + monitor, + elements.ChrQuad(name="quad1", ds=1.0, k=1.0, nslice=ns), + monitor, + elements.ChrDrift(name="drift2", ds=0.5, nslice=ns), + monitor, + elements.ChrQuad(name="quad2", ds=1.0, k=-1.0, nslice=ns), + monitor, + elements.ChrDrift(name="drift3", ds=0.25, nslice=ns), + monitor, +] +# assign a fodo segment +sim.lattice.extend(fodo) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/src/initialization/Algorithms.H b/src/initialization/Algorithms.H index 6a28f0903..bdb685790 100644 --- a/src/initialization/Algorithms.H +++ b/src/initialization/Algorithms.H @@ -21,7 +21,8 @@ namespace impactx AMREX_ENUM(SpaceChargeAlgo, False, /**< Disabled: no space charge calculation */ True_3D, /**< 3D beam distribution */ - True_2D /**< averaged 2D transverse beam distribution with a current along s */ + Gauss3D, /**< Assume a 3D Gaussian beam distribution */ + True_2D /**< 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 007f5bc40..699ad2e27 100644 --- a/src/initialization/Algorithms.cpp +++ b/src/initialization/Algorithms.cpp @@ -62,6 +62,10 @@ namespace impactx { return SpaceChargeAlgo::True_3D; } + else if (space_charge == "Gauss3D") + { + return SpaceChargeAlgo::Gauss3D; + } else if (space_charge == "2D") { return SpaceChargeAlgo::True_2D; diff --git a/src/particles/spacecharge/CMakeLists.txt b/src/particles/spacecharge/CMakeLists.txt index 38ad2ea66..669c2f0d0 100644 --- a/src/particles/spacecharge/CMakeLists.txt +++ b/src/particles/spacecharge/CMakeLists.txt @@ -1,6 +1,7 @@ target_sources(lib PRIVATE ForceFromSelfFields.cpp + Gauss3dPush.cpp GatherAndPush.cpp HandleSpacecharge.cpp PoissonSolve.cpp diff --git a/src/particles/spacecharge/Gauss3dPush.H b/src/particles/spacecharge/Gauss3dPush.H new file mode 100644 index 000000000..67402991b --- /dev/null +++ b/src/particles/spacecharge/Gauss3dPush.H @@ -0,0 +1,40 @@ +/* Copyright 2022-2023 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_GAUSS3D_PUSH_H +#define IMPACT_GAUSS3D_PUSH_H + +#include "particles/ImpactXParticleContainer.H" + +#include + + +namespace impactx::particles::spacecharge +{ + /** Apply a Space-Charge Kick according to a 3D Gaussian Distribution + * + * This calculates the space charge fields from a 3D Gaussian distribution + * with respect to particle position from the following paper: + * J. Qiang et al., "Two-and-a-half dimensional symplectic space-charge + * solver", LBNL Report Number: LBNL-2001674 (2025). + * The momentum of all particles is then pushed using a common + * time step given by the reference particle speed and ds slice. The + * position push is done in the lattice elements and not here. + * + * @param[inout] pc container of the particles that deposited rho + * @param[in] slice_ds segment length in meters + */ + void Gauss3dPush ( + ImpactXParticleContainer & pc, + amrex::ParticleReal slice_ds + ); + +} // namespace impactx::particles::spacecharge + +#endif // IMPACT_GAUSS3D_PUSH_H diff --git a/src/particles/spacecharge/Gauss3dPush.cpp b/src/particles/spacecharge/Gauss3dPush.cpp new file mode 100644 index 000000000..75dfd0276 --- /dev/null +++ b/src/particles/spacecharge/Gauss3dPush.cpp @@ -0,0 +1,187 @@ +/* Copyright 2022-2023 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 "Gauss3dPush.H" + +#include +#include // for Real +#include "diagnostics/ReducedBeamCharacteristics.H" + + +namespace impactx::particles::spacecharge +{ + /** Compute space-charge fields from a 3D Gaussian distribution. + * + * Compute integrals Eqs 45-47 used in the space-charge fields from a 3D Gaussian distribution. + * "A two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674, 2025. + * Input particle locations (x,y,z) and RMS sizes (sigx,sigy,sigz) and return the integrals for SC fields. + * + * @todo This function can be vectorized with AMREX_SIMD for better CPU performance. + */ + AMREX_GPU_DEVICE + void efldgauss ( + int nint, + amrex::ParticleReal const xin, + amrex::ParticleReal const yin, + amrex::ParticleReal const zin, + amrex::ParticleReal const sigx, + amrex::ParticleReal const sigy, + amrex::ParticleReal const sigz, + amrex::ParticleReal const gamma, + amrex::ParticleReal & pintex, + amrex::ParticleReal & pintey, + amrex::ParticleReal & pintez + ) + { + using namespace amrex::literals; + + // enforce an odd number of grid points + if (nint % 2 == 0) nint += 1; + + amrex::ParticleReal const xmin = 0, xmax = 1; + amrex::ParticleReal const h = (xmax - xmin) / (nint - 1); + amrex::ParticleReal const xp = xin / sigx; + amrex::ParticleReal const yp = yin / sigy; + amrex::ParticleReal const zp = zin / sigz; + amrex::ParticleReal const asp = sigx / sigy; + amrex::ParticleReal const basp = sigx / (sigz * gamma); + amrex::ParticleReal const xp2 = xp * xp; + amrex::ParticleReal const yp2 = yp * yp; + amrex::ParticleReal const zp2 = zp * zp; + + // integration results + amrex::ParticleReal sum0ex = 0, sum0ey = 0, sum0ez = 0; + + for (int i = 0; i < nint; ++i) + { + amrex::ParticleReal const x = xmin + i * h; + amrex::ParticleReal const t2 = asp * asp + (1_prt - asp * asp) * x * x; + amrex::ParticleReal const t2sqrt = std::sqrt(t2); + amrex::ParticleReal const bt2 = basp * basp + (1_prt - basp * basp) * x * x; + amrex::ParticleReal const bt2sqrt = std::sqrt(bt2); + amrex::ParticleReal const f2x = t2sqrt * bt2sqrt; + amrex::ParticleReal const f2y = t2 * t2sqrt * bt2sqrt; + amrex::ParticleReal const f2z = bt2 * t2sqrt * bt2sqrt; + amrex::ParticleReal const f1 = x * x * std::exp(-x * x * (xp2 + yp2 / t2 + zp2 / bt2) / 2_prt); + + // Simpson's rule integration + + // Ex + amrex::ParticleReal const fx = f1 / f2x; + if (i%2 == 0) + sum0ex += 2_prt * fx * h; + else + sum0ex += 4_prt * fx * h; + + // Ey + amrex::ParticleReal const fy = f1 / f2y; + if (i%2 == 0) + sum0ey += 2_prt * fy * h; + else + sum0ey += 4_prt * fy * h; + + // Ez + amrex::ParticleReal const fz = f1 / f2z; + if (i%2 == 0) + sum0ez += 2_prt * fz * h; + else + sum0ez += 4_prt * fz * h; + } + + pintex = sum0ex / 3_prt; + pintey = sum0ey / 3_prt; + pintez = sum0ez / 3_prt; + } + + void Gauss3dPush ( + ImpactXParticleContainer & pc, + amrex::ParticleReal const slice_ds + ) + { + BL_PROFILE("impactx::spacecharge::GatherAndPush"); + + using namespace amrex::literals; + + amrex::ParticleReal const charge = pc.GetRefParticle().charge; + + // get RMS sigma of beam + // Standard deviation of the particle positions (speed of light times time delay for t) (unit: meter) + std::unordered_map const rbc = diagnostics::reduced_beam_characteristics(pc); + amrex::ParticleReal const sigx = rbc.at("sig_x"); + amrex::ParticleReal const sigy = rbc.at("sig_y"); + amrex::ParticleReal const sigz = rbc.at("sig_t"); + amrex::ParticleReal const bchchg = rbc.at("charge_C"); + + // physical constants and reference quantities + amrex::ParticleReal const c0_SI = 2.99792458e8; // TODO move out + amrex::ParticleReal const mc_SI = pc.GetRefParticle().mass * c0_SI; + amrex::ParticleReal const pz_ref_SI = pc.GetRefParticle().beta_gamma() * mc_SI; + amrex::ParticleReal const gamma = pc.GetRefParticle().gamma(); + amrex::ParticleReal const inv_gamma2 = 1.0_prt / (gamma * gamma); + amrex::ParticleReal const rfpiepslon = c0_SI*c0_SI*1.0e-7; + + amrex::ParticleReal const dt = slice_ds / pc.GetRefParticle().beta() / c0_SI; + + // group together constants for the momentum + using ablastr::constant::math::pi; + amrex::ParticleReal const asp = sigx/sigy; + amrex::ParticleReal const push_consts = dt * charge * inv_gamma2 / pz_ref_SI + * 2_prt * asp * (sigx * sigx * sigz * std::sqrt(2_prt * pi)) + * bchchg * rfpiepslon; + + // loop over refinement levels + int const nLevel = pc.finestLevel(); + for (int lev = 0; lev <= nLevel; ++lev) + { + // loop over all particle boxes + using ParIt = ImpactXParticleContainer::iterator; +#ifdef AMREX_USE_OMP +#pragma omp parallel if (amrex::Gpu::notInLaunchRegion()) +#endif + for (ParIt pti(pc, lev); pti.isValid(); ++pti) { + const int np = pti.numParticles(); + + // preparing access to particle data: SoA of Reals + auto& soa_real = pti.GetStructOfArrays().GetRealData(); + amrex::ParticleReal const * const AMREX_RESTRICT part_x = soa_real[RealSoA::x].dataPtr(); + amrex::ParticleReal const * const AMREX_RESTRICT part_y = soa_real[RealSoA::y].dataPtr(); + amrex::ParticleReal const * const AMREX_RESTRICT part_z = soa_real[RealSoA::z].dataPtr(); // note: currently for a fixed t + amrex::ParticleReal* const AMREX_RESTRICT part_px = soa_real[RealSoA::px].dataPtr(); + amrex::ParticleReal* const AMREX_RESTRICT part_py = soa_real[RealSoA::py].dataPtr(); + amrex::ParticleReal* const AMREX_RESTRICT part_pz = soa_real[RealSoA::pz].dataPtr(); // note: currently for a fixed t + + // gather to each particle and push momentum + amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) { + // access SoA Real data + amrex::ParticleReal const & AMREX_RESTRICT x = part_x[i]; + amrex::ParticleReal const & AMREX_RESTRICT y = part_y[i]; + amrex::ParticleReal const & AMREX_RESTRICT z = part_z[i]; + amrex::ParticleReal & AMREX_RESTRICT px = part_px[i]; + amrex::ParticleReal & AMREX_RESTRICT py = part_py[i]; + amrex::ParticleReal & AMREX_RESTRICT pz = part_pz[i]; + + // field integrals from a 3D Gaussian bunch + int const nint = 401; // TODO: should "nint" be user-configurable? Otherwise make it constexpr in efldgauss + amrex::ParticleReal eintx, einty, eintz; + efldgauss(nint,x,y,z,sigx,sigy,sigz,gamma,eintx,einty,eintz); + + // push momentum + px += x * eintx * push_consts; + py += y * einty * push_consts; + pz += z * eintz * push_consts; + + // push position is done in the lattice elements + }); + + + } // end loop over all particle boxes + } // env mesh-refinement level loop + } + +} // namespace impactx::particles::spacecharge diff --git a/src/particles/spacecharge/HandleSpacecharge.cpp b/src/particles/spacecharge/HandleSpacecharge.cpp index 9612b236f..89a19e25f 100644 --- a/src/particles/spacecharge/HandleSpacecharge.cpp +++ b/src/particles/spacecharge/HandleSpacecharge.cpp @@ -14,13 +14,14 @@ #include "particles/ImpactXParticleContainer.H" #include "particles/spacecharge/ForceFromSelfFields.H" #include "particles/spacecharge/GatherAndPush.H" +#include "particles/spacecharge/Gauss3dPush.H" #include "particles/spacecharge/PoissonSolve.H" #include "particles/transformation/CoordinateTransformation.H" -#include #include #include +#include namespace impactx::particles::spacecharge @@ -47,45 +48,59 @@ namespace impactx::particles::spacecharge CoordSystem::t ); - // Note: The following operation assume that - // the particles are in x, y, z coordinates. - - // Resize the mesh, based on `amr_data->track_particles.m_particle_container` extent - ResizeMesh(); - - // Redistribute particles in the new mesh in x, y, z - amr_data->track_particles.m_particle_container->Redistribute(); - - // charge deposition - amr_data->track_particles.m_particle_container->DepositCharge( - amr_data->track_particles.m_rho, - amr_data->refRatio() - ); - - // poisson solve in x,y,z - spacecharge::PoissonSolve( - *amr_data->track_particles.m_particle_container, - amr_data->track_particles.m_rho, - amr_data->track_particles.m_phi, - amr_data->refRatio() - ); - - // calculate force in x,y,z - spacecharge::ForceFromSelfFields( - amr_data->track_particles.m_space_charge_field, - amr_data->track_particles.m_phi, - amr_data->Geom() - ); - - // gather and space-charge push in x,y,z , assuming the space-charge - // field is the same before/after transformation - // TODO: This is currently using linear order. - spacecharge::GatherAndPush( - *amr_data->track_particles.m_particle_container, - amr_data->track_particles.m_space_charge_field, - amr_data->Geom(), - slice_ds - ); + if (space_charge == SpaceChargeAlgo::Gauss3D) + { + Gauss3dPush(*amr_data->track_particles.m_particle_container, slice_ds); + } + else if (space_charge == SpaceChargeAlgo::True_3D) + { + // Note: The following operations assume that + // the particles are in x, y, z coordinates. + + // Resize the mesh, based on `amr_data->track_particles.m_particle_container` extent + ResizeMesh(); + + // Redistribute particles in the new mesh in x, y, z + amr_data->track_particles.m_particle_container->Redistribute(); + + // charge deposition + amr_data->track_particles.m_particle_container->DepositCharge( + amr_data->track_particles.m_rho, + amr_data->refRatio() + ); + + // poisson solve in x,y,z + spacecharge::PoissonSolve( + *amr_data->track_particles.m_particle_container, + amr_data->track_particles.m_rho, + amr_data->track_particles.m_phi, + amr_data->refRatio() + ); + + // calculate force in x,y,z + spacecharge::ForceFromSelfFields( + amr_data->track_particles.m_space_charge_field, + amr_data->track_particles.m_phi, + amr_data->Geom() + ); + + // gather and space-charge push in x,y,z , assuming the space-charge + // field is the same before/after transformation + // TODO: This is currently using linear order. + spacecharge::GatherAndPush( + *amr_data->track_particles.m_particle_container, + amr_data->track_particles.m_space_charge_field, + amr_data->Geom(), + slice_ds + ); + } + else if (space_charge == SpaceChargeAlgo::True_2D) + { + throw std::runtime_error( + "2D space charge is not implemented yet for particle tracking. " + "Please follow https://github.com/BLAST-ImpactX/impactx/pull/909 for updates." + ); + } // transform from x,y,z to x',y',t transformation::CoordinateTransformation( @@ -93,7 +108,7 @@ namespace impactx::particles::spacecharge CoordSystem::s ); - // for later: original Impact implementation as an option + // for later: original Impact implementation as an option for 3D space charge to: // Redistribute particles in x',y',t // TODO: only needed if we want to gather and push space charge // in x',y',t diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index b29dadb06..7dc097859 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -271,8 +271,8 @@ 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") { - throw std::runtime_error("Space charge model must be 2D or 3D but is: " + space_charge); + if (space_charge != "false" && space_charge != "off" && space_charge != "2D" && space_charge != "3D" && space_charge != "Gauss3D") { + throw std::runtime_error("Space charge model must be 2D, 3D or Gauss3D but is: " + space_charge); } amrex::ParmParse pp_algo("algo"); pp_algo.add("space_charge", space_charge); diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index 7475e15d7..15c1402b0 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -106,6 +106,12 @@ namespace impactx ablastr::warn_manager::WarnPriority::high ); } + if (space_charge == SpaceChargeAlgo::Gauss3D) { + throw std::runtime_error( + "Gauss3D space charge force calculation is only supported with particle tracking. " + "For envelope tracking, use: 3D" + ); + } bool csr = false; pp_algo.query("csr", csr);