diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 5db4dc058..b0ca2284b 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -882,9 +882,23 @@ See there ``nslice`` option on lattice elements for slicing. * ``"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). + Ref.: J. Qiang, "Two-and-a-half dimensional symplectic space-charge solver", LBNL Report Number: LBNL-2001674 (2025). (This reference describes both 3D and 2.5D models.) + This model supports the following sub-option: + + * ``algo.space_charge.gauss_nint`` (``int``, default: ``101``) + + Number of steps for computing the integrals (Eqs. 45-47 in the above paper). + + * ``algo.space_charge.gauss_taylor_delta`` (``float``, default: ``0.01``) + + Initial integral region to avoid divergence of integrand at 0. + + * ``algo.space_charge.gauss_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 9851eb80f..df6b57ad5 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -79,8 +79,21 @@ Collective Effects & Overall Simulation Parameters * ``"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). + Ref.: J. Qiang, "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:: space_charge_gauss_nint + + Number of steps for computing the integrals (default: ``101``). + + .. py:property:: space_charge_gauss_taylor_delta + + Initial integral region to avoid integrand divergence at 0 (default: ``0.01``). + + .. py:property:: space_charge_gauss_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/examples/CMakeLists.txt b/examples/CMakeLists.txt index 5fa986f0f..7253bffa8 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1403,6 +1403,21 @@ add_impactx_test(FODO.Gauss3d.sc.py OFF # no plot script yet ) +# FODO cell with Gaussian 2.5D space charge using particle tracking ############# +# +add_impactx_test(FODO.Gauss2p5d.sc + examples/fodo_space_charge/input_fodo_Gauss2p5D_sc.in + ON # ImpactX MPI-parallel + examples/fodo_space_charge/analysis_fodo_Gauss2p5D_sc.py + OFF # no plot script yet +) +add_impactx_test(FODO.Gauss2p5d.sc.py + examples/fodo_space_charge/run_fodo_Gauss2p5D_sc.py + ON # ImpactX MPI-parallel + examples/fodo_space_charge/analysis_fodo_Gauss2p5D_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 9e6562649..0fbe64a3d 100644 --- a/examples/fodo_space_charge/README.rst +++ b/examples/fodo_space_charge/README.rst @@ -53,7 +53,7 @@ We run the following script to analyze correctness: :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_envelope_sc.py``. -.. _examples-fodo-envelope-sc-gaussian: +.. _examples-fodo-Gaussian-sc: FODO Cell with 3D Gaussian Space Charge Using Particle Tracking =============================================================== @@ -104,3 +104,7 @@ We run the following script to analyze correctness: .. literalinclude:: analysis_fodo_Gauss3D_sc.py :language: python3 :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py``. + +.. _examples-fodo-2p5dGaussian-sc: + +FODO Cell with 2.5D Gaussian Space Charge Using Particle Tracking diff --git a/examples/fodo_space_charge/analysis_fodo_Gauss2p5D_sc.py b/examples/fodo_space_charge/analysis_fodo_Gauss2p5D_sc.py new file mode 100755 index 000000000..0e915b02f --- /dev/null +++ b/examples/fodo_space_charge/analysis_fodo_Gauss2p5D_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 = 3.0 * 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.51e-05, 7.51e-05, 9.99e-4, 1.98e-09, 1.98e-09, 1.97e-06], + rtol=rtol, + atol=atol, +) + + +print("") +print("Final Beam:") +sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf = get_moments(final) +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 = 3.0 * 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], + [ + 9.22e-05, + 8.55e-05, + 0.000996, + 2.04e-09, + 2.01e-09, + 1.97e-06, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py b/examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py index f7e292529..8cca07f17 100755 --- a/examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py +++ b/examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py @@ -53,12 +53,12 @@ def get_moments(beam): ) atol = 0.0 # ignored -rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +rtol = 3.5 * 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], + [7.51e-05, 7.51e-05, 9.99e-4, 1.98e-09, 1.98e-09, 1.97e-06], rtol=rtol, atol=atol, ) @@ -66,25 +66,25 @@ def get_moments(beam): print("") print("Final Beam:") -sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf = get_moments(initial) +sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf = get_moments(final) 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 +rtol = 3.5 * 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, + 9.21e-05, + 8.54e-05, + 0.000996, + 2.04e-09, + 2.01e-09, + 1.97e-06, ], rtol=rtol, atol=atol, diff --git a/examples/fodo_space_charge/input_fodo_Gauss2p5D_sc.in b/examples/fodo_space_charge/input_fodo_Gauss2p5D_sc.in new file mode 100644 index 000000000..b49a8804a --- /dev/null +++ b/examples/fodo_space_charge/input_fodo_Gauss2p5D_sc.in @@ -0,0 +1,60 @@ +############################################################################### +# 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 = Gauss2p5D +algo.space_charge.gauss_nint = 101 +algo.space_charge.gauss_charge_z_bins = 129 +algo.space_charge.gauss_taylor_delta = 0.01 + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/fodo_space_charge/input_fodo_Gauss3D_sc.in b/examples/fodo_space_charge/input_fodo_Gauss3D_sc.in index c3be8b254..97ce7d59d 100644 --- a/examples/fodo_space_charge/input_fodo_Gauss3D_sc.in +++ b/examples/fodo_space_charge/input_fodo_Gauss3D_sc.in @@ -50,7 +50,7 @@ drift3.ds = 0.25 ############################################################################### algo.particle_shape = 2 algo.space_charge = Gauss3D - +algo.space_charge.gauss_nint = 101 ############################################################################### # Diagnostics diff --git a/examples/fodo_space_charge/run_fodo_Gauss2p5D_sc.py b/examples/fodo_space_charge/run_fodo_Gauss2p5D_sc.py new file mode 100755 index 000000000..e5ead31c9 --- /dev/null +++ b/examples/fodo_space_charge/run_fodo_Gauss2p5D_sc.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Ji Qiang +# 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 = "Gauss2p5D" +sim.space_charge_gauss_nint = 101 +sim.space_charge_gauss_charge_z_bins = 129 +sim.space_charge_gauss_taylor_delta = 0.01 +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/examples/fodo_space_charge/run_fodo_Gauss3D_sc.py b/examples/fodo_space_charge/run_fodo_Gauss3D_sc.py index d76b0904f..e1f370c00 100755 --- a/examples/fodo_space_charge/run_fodo_Gauss3D_sc.py +++ b/examples/fodo_space_charge/run_fodo_Gauss3D_sc.py @@ -13,6 +13,7 @@ # set numerical parameters and IO control sim.particle_shape = 2 # B-spline order sim.space_charge = "Gauss3D" +sim.space_charge_gauss_nint = 101 sim.slice_step_diagnostics = True # domain decomposition & space charge mesh diff --git a/src/initialization/Algorithms.H b/src/initialization/Algorithms.H index bdb685790..3a8efdd89 100644 --- a/src/initialization/Algorithms.H +++ b/src/initialization/Algorithms.H @@ -22,6 +22,7 @@ namespace impactx False, /**< Disabled: no space charge calculation */ 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 */ ); diff --git a/src/initialization/Algorithms.cpp b/src/initialization/Algorithms.cpp index 699ad2e27..cbc780638 100644 --- a/src/initialization/Algorithms.cpp +++ b/src/initialization/Algorithms.cpp @@ -66,6 +66,10 @@ namespace impactx { return SpaceChargeAlgo::Gauss3D; } + else if (space_charge == "Gauss2p5D") + { + return SpaceChargeAlgo::Gauss2p5D; + } else if (space_charge == "2D") { return SpaceChargeAlgo::True_2D; diff --git a/src/particles/spacecharge/CMakeLists.txt b/src/particles/spacecharge/CMakeLists.txt index 669c2f0d0..46b36595d 100644 --- a/src/particles/spacecharge/CMakeLists.txt +++ b/src/particles/spacecharge/CMakeLists.txt @@ -2,6 +2,7 @@ target_sources(lib PRIVATE ForceFromSelfFields.cpp Gauss3dPush.cpp + Gauss2p5dPush.cpp GatherAndPush.cpp HandleSpacecharge.cpp PoissonSolve.cpp diff --git a/src/particles/spacecharge/Gauss2p5dPush.H b/src/particles/spacecharge/Gauss2p5dPush.H new file mode 100644 index 000000000..7a1fa7d79 --- /dev/null +++ b/src/particles/spacecharge/Gauss2p5dPush.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_GAUSS2p5D_PUSH_H +#define IMPACT_GAUSS2p5D_PUSH_H + +#include "particles/ImpactXParticleContainer.H" + +#include +#include + + +namespace impactx::particles::spacecharge +{ + /**------------------------ + * + * This calculates the space charge fields from a 3D Gaussian distribution + * with respect to particle position from the following paper. + * "Two-and-a-half dimensional symplectic space-charge slover" + * 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 Gauss2p5dPush ( + ImpactXParticleContainer & pc, + amrex::ParticleReal slice_ds + ); + +} // namespace impactx::particles::spacecharge + +#endif // IMPACT_GAUSS2p5D_PUSH_H diff --git a/src/particles/spacecharge/Gauss2p5dPush.cpp b/src/particles/spacecharge/Gauss2p5dPush.cpp new file mode 100644 index 000000000..fab90567f --- /dev/null +++ b/src/particles/spacecharge/Gauss2p5dPush.cpp @@ -0,0 +1,263 @@ +/* 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 "Gauss2p5dPush.H" + +#include // for Real +#include +#include +#include +#include // for AMREX_D_DECL +#include "diagnostics/ReducedBeamCharacteristics.H" +#include "particles/wakefields/ChargeBinning.H" + + +namespace impactx::particles::spacecharge +{ + + // compute integrals Eqs 25, 31,32 used in the space-charge fields from a 2D transverse Gaussian distribution (including 2A in Eq. 25). + // Input particle locations (x,y) and RMS sizes (sigx,sigy) and return the integrals for SC fields. + // + AMREX_GPU_DEVICE + void potInt(amrex::ParticleReal delta, int nint, amrex::ParticleReal xin, amrex::ParticleReal yin, amrex::ParticleReal sigx, amrex::ParticleReal sigy, + amrex::ParticleReal& pintex, amrex::ParticleReal& pintey, amrex::ParticleReal& pintez) + { + using namespace amrex::literals; + + // enforce odd grid points + if (nint % 2 == 0) nint += 1; + + amrex::ParticleReal xmin = delta; + amrex::ParticleReal xmax = 1.0; + amrex::ParticleReal h = (xmax - xmin) / (nint - 1); + + amrex::ParticleReal xp = xin / sigx; + amrex::ParticleReal yp = yin / sigy; + amrex::ParticleReal asp = sigx / sigy; + amrex::ParticleReal xp2 = xp*xp; + amrex::ParticleReal yp2 = yp*yp; + + amrex::ParticleReal sum0, sum0ex, sum0ey; + sum0 = sum0ex = sum0ey = 0.0; + + // Trapezoidal rule approximation integral from 0 to delta + amrex::ParticleReal x = xmin; + amrex::ParticleReal t2 = asp*asp + (1_prt - asp*asp)*x*x; + amrex::ParticleReal sqrt2 = std::sqrt(t2); + amrex::ParticleReal exparg = -(xp*xp + yp*yp / t2) / 2_prt; + amrex::ParticleReal f1 = x*exparg; + + sum0 = x*f1/sqrt2/2_prt; + f1 = x*std::exp(x*x*exparg); + sum0ex = x*f1 / sqrt2 / 2_prt; + sum0ey = x*f1 / (t2*sqrt2) / 2_prt; + + // Simpson's rule for two end points + // Ez + f1 = std::exp(x*x*exparg); + amrex::ParticleReal f2 = x*sqrt2; + amrex::ParticleReal fx = (f1 - 1_prt) / f2; + sum0 = sum0 - fx*h; + + amrex::ParticleReal f1x = x*f1; + f2 = sqrt2; + fx = f1x / f2; + sum0ex = sum0ex - fx*h; + + f2 = t2*sqrt2; + fx = f1x / f2; + sum0ey = sum0ey - fx*h; + + // Ez at xmax + x = xmax; + t2 = asp*asp + (1_prt - asp*asp)*x*x; + sqrt2 = std::sqrt(t2); + exparg = -(xp*xp + yp*yp / t2) / 2_prt; + f1 = std::exp(x*x*exparg); + f2 = x*sqrt2; + fx = (f1 - 1_prt) / f2; + sum0 = sum0 - fx*h; + f1x = x*f1; + f2 = sqrt2; + fx = f1x / f2; + sum0ex = sum0ex - fx*h; + f2 = t2*sqrt2; + fx = f1x / f2; + sum0ey = sum0ey - fx*h; + + amrex::ParticleReal fy = 0.0; + amrex::ParticleReal fz = 0.0; + amrex::ParticleReal t2sqrt = 0.0; + amrex::ParticleReal f2y = 0.0; + amrex::ParticleReal asp2 = asp*asp; + amrex::ParticleReal asp2m = 1_prt - asp*asp; + for (int i = 0; i < nint; ++i) + { + x = xmin + i*h; + t2 = asp2 + asp2m * x*x; + t2sqrt = std::sqrt(t2); + f2y = t2*t2sqrt; + f2 = x*t2sqrt; + f1 = std::exp(-x*x*(xp2 + yp2 / t2) / 2_prt); + f1x = x*f1; + + // Ez + fz = (f1 - 1_prt) / f2; + if (i%2 == 0) + sum0 += 2_prt * fz * h; + else + sum0 += 4_prt * fz * h; + + // Ex + fx = f1x / t2sqrt; + if (i%2 == 0) + sum0ex += 2_prt * fx * h; + else + sum0ex += 4_prt * fx * h; + + // Ey + fy = f1x / f2y; + if (i%2 == 0) + sum0ey += 2_prt * fy * h; + else + sum0ey += 4_prt * fy * h; + } + + pintez = 2 * asp * sum0 / 3_prt; + pintex = 2 * asp * xp * sum0ex / 3_prt / sigx; + pintey = 2 * asp * yp * sum0ey / 3_prt / sigy; + } + + //advanced particle momentae using 2.5D SC fields with a transverse Gaussian distribution. + void Gauss2p5dPush (ImpactXParticleContainer & pc, amrex::ParticleReal const slice_ds) + { + BL_PROFILE("impactx::spacecharge::Gauss2p5dPush"); + + using namespace amrex::literals; + + amrex::ParticleReal const charge = pc.GetRefParticle().charge; + + std::unordered_map const rbc = diagnostics::reduced_beam_characteristics(pc); + // get RMS sigma of beam + // Standard deviation of the particle positions (speed of light times time delay for t) (unit: meter) + amrex::ParticleReal const sigx = rbc.at("sig_x"); + amrex::ParticleReal const sigy = rbc.at("sig_y"); + + // 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; + + int nint = 101; + amrex::Real delta = 0.01; + amrex::ParmParse pp_algo("algo.space_charge"); + pp_algo.queryAddWithParser("gauss_nint", nint); + pp_algo.queryAddWithParser("gauss_taylor_delta", delta); + + int tp5d_bins = 129; + pp_algo.queryAddWithParser("gauss_charge_z_bins", tp5d_bins); + + // 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(); + + // Set parameters for charge deposition + bool const is_unity_particle_weight = false; // Only true if w = 1 + bool const GetNumberDensity = true; + + 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 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); + + // 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); + + 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 push_consts = dt * charge * inv_gamma2 / pz_ref_SI; + amrex::ParticleReal const chargesign = charge/std::abs(charge); + + // 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 AMREX_RESTRICT part_x = soa_real[RealSoA::x].dataPtr(); + amrex::ParticleReal* const AMREX_RESTRICT part_y = soa_real[RealSoA::y].dataPtr(); + amrex::ParticleReal* 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 & AMREX_RESTRICT x = part_x[i]; + amrex::ParticleReal & AMREX_RESTRICT y = part_y[i]; + amrex::ParticleReal & 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; +// amrex::Real const delta=0.01; + + amrex::ParticleReal eintx, einty, eintz; + potInt(delta,nint,x,y,sigx,sigy,eintx,einty,eintz); + + // 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 Gaussian SC: " << idx << std::endl; + } +#endif + amrex::ParticleReal const Fxy = beam_profile[idx]*chargesign; + amrex::ParticleReal const Fz = beam_profile_slope[idx]*charge; + + // push momentum + using ablastr::constant::math::pi; + px += eintx*Fxy*rfpiepslon * push_consts; + py += einty*Fxy*rfpiepslon * push_consts; + pz -= (eintz-std::log(2.0)+0.577216-2.0*std::log((sigx+sigy)/2.0))*Fz*rfpiepslon * 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/Gauss3dPush.cpp b/src/particles/spacecharge/Gauss3dPush.cpp index 14fc36949..ad6bbb97c 100644 --- a/src/particles/spacecharge/Gauss3dPush.cpp +++ b/src/particles/spacecharge/Gauss3dPush.cpp @@ -4,11 +4,11 @@ * * This file is part of ImpactX. * - * Authors: Ji Qiang + * Authors: Ji Qiang, Axel Huebl * License: BSD-3-Clause-LBNL */ #include "Gauss3dPush.H" - +#include #include #include // for Real #include "diagnostics/ReducedBeamCharacteristics.H" @@ -134,6 +134,10 @@ namespace impactx::particles::spacecharge amrex::ParticleReal const dt = slice_ds / pc.GetRefParticle().beta() / c0_SI; + int nint = 101; + amrex::ParmParse pp_algo("algo.space_charge"); + pp_algo.queryAddWithParser("gauss_nint", nint); + // group together constants for the momentum using ablastr::constant::math::pi; amrex::ParticleReal const asp = sigx/sigy; @@ -175,7 +179,6 @@ namespace impactx::particles::spacecharge amrex::ParticleReal & AMREX_RESTRICT pz = part_pz[i]; // field integrals from a 3D Gaussian bunch - int const nint = 101; // 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); diff --git a/src/particles/spacecharge/HandleSpacecharge.cpp b/src/particles/spacecharge/HandleSpacecharge.cpp index 89a19e25f..d8c009592 100644 --- a/src/particles/spacecharge/HandleSpacecharge.cpp +++ b/src/particles/spacecharge/HandleSpacecharge.cpp @@ -15,6 +15,7 @@ #include "particles/spacecharge/ForceFromSelfFields.H" #include "particles/spacecharge/GatherAndPush.H" #include "particles/spacecharge/Gauss3dPush.H" +#include "particles/spacecharge/Gauss2p5dPush.H" #include "particles/spacecharge/PoissonSolve.H" #include "particles/transformation/CoordinateTransformation.H" @@ -52,6 +53,10 @@ namespace impactx::particles::spacecharge { Gauss3dPush(*amr_data->track_particles.m_particle_container, slice_ds); } + else if (space_charge == SpaceChargeAlgo::Gauss2p5D) + { + Gauss2p5dPush(*amr_data->track_particles.m_particle_container, slice_ds); + } else if (space_charge == SpaceChargeAlgo::True_3D) { // Note: The following operations assume that diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 7ed9e27c2..aa8e66556 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -271,7 +271,7 @@ 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") { + 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 or Gauss3D but is: " + space_charge); } amrex::ParmParse pp_algo("algo"); @@ -280,6 +280,50 @@ void init_ImpactX (py::module& m) }, "The model to be used when calculating space charge effects. Either off, 2D, or 3D." ) + .def_property("space_charge_gauss_nint", + [](ImpactX & /* ix */) { + return detail::get_or_throw("algo.space_charge", "gauss_nint"); + }, + [](ImpactX & /* ix */, int const gauss_nint) { + if (gauss_nint < 1) { + throw std::runtime_error("space_charge_gauss_nint must be strictly positive"); + } + + amrex::ParmParse pp_algo("algo.space_charge"); + pp_algo.add("gauss_nint", gauss_nint); + }, + "Number of steps for computing the integrals (default: ``101``)." + ) + .def_property("space_charge_gauss_charge_z_bins", + [](ImpactX & /* ix */) { + return detail::get_or_throw("algo.space_charge", "gauss_charge_z_bins"); + }, + [](ImpactX & /* ix */, int const gauss_charge_z_bins) { + if (gauss_charge_z_bins < 1) { + throw std::runtime_error("space_charge_gauss_charge_z_bins must be strictly positive"); + } + + 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``)." + ) + .def_property("space_charge_gauss_taylor_delta", + [](ImpactX & /* ix */) { + return detail::get_or_throw("algo.space_charge", "gauss_taylor_delta"); + }, + [](ImpactX & /* ix */, amrex::Real const gauss_taylor_delta) { + if (gauss_taylor_delta < 0) { + throw std::runtime_error("space_charge_gauss_taylor_delta must be strictly positive"); + } + if (gauss_taylor_delta > 0.05) { + throw std::runtime_error("space_charge_gauss_taylor_delta must be less than 0.05"); + } + amrex::ParmParse pp_algo("algo.space_charge"); + pp_algo.add("gauss_taylor_delta", gauss_taylor_delta); + }, + "Initial region for computing the integrals (default: ``0.01``)." + ) .def_property("poisson_solver", [](ImpactX & /* ix */) { return detail::get_or_throw("algo", "poisson_solver");