diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 62b91f5e4..01606341b 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -30,6 +30,7 @@ This section allows you to **download input files** that correspond to different examples/kicker/README.rst examples/thin_dipole/README.rst examples/aperture/README.rst + examples/epac2004_benchmarks/README.rst examples/pytorch_surrogate_model/README.rst examples/apochromatic/README.rst diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index e5eb977e0..7b873b72f 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -236,6 +236,21 @@ Initial Beam Distributions * ``beam.muypy`` (``float``, dimensionless, default: ``0``) correlation Y-Py * ``beam.mutpt`` (``float``, dimensionless, default: ``0``) correlation T-Pt + * ``thermal`` for a 6D stationary thermal or bithermal distribution. + This distribution type is described, for example in: + R. D. Ryne et al, "A Test Suite of Space-Charge Problems for Code Benchmarking", in Proc. EPAC2004, Lucerne, Switzerland. + C. E. Mitchell et al, "ImpactX Modeling of Benchmark Tests for Space Charge Validation", in Proc. HB2023, Geneva, Switzerland. + With additional parameters: + + * ``beam.k`` (``float``, in inverse meters) external focusing strength + * ``beam.kT`` (``float``, dimensionless) temperature of core population + = < p_x^2 > = < p_y^2 >, where all momenta are normalized by the reference momentum + * ``beam.kT_halo`` (``float``, dimensionless, default ``kT``) temperature of halo population + * ``beam.normalize`` (``float``, dimensionless) normalizing constant for core population + * ``beam.normalize_halo`` (``float``, dimensionless) normalizing constant for halo population + * ``beam.halo`` (``float``, dimensionless) fraction of charge in halo + + .. _running-cpp-parameters-lattice: Lattice Elements diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 334c2c198..910ac1eba 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -427,6 +427,10 @@ This module provides particle beam distributions that can be used to initialize A 6D Waterbag distribution. +.. py:class:: impactx.distribution.Thermal(k, kT, kT_halo, normalize, normalize_halo, halo) + + A 6D stationary thermal or bithermal distribution. + Lattice Elements ---------------- diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 1b5bebf6d..3907f03fe 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -548,7 +548,7 @@ add_impactx_test(positron_channel.py OFF # no plot script yet ) -# Cyclotron ############################################################ +# Cyclotron ################################################################### # # w/o space charge add_impactx_test(cyclotron @@ -566,7 +566,7 @@ add_impactx_test(cyclotron.py OFF # no plot script yet ) -# Combined-function bend ######################################################## +# Combined-function bend ###################################################### # # w/o space charge add_impactx_test(cfbend @@ -584,7 +584,7 @@ add_impactx_test(cfbend.py OFF # no plot script yet ) -# Ballistic compression ######################################################## +# Ballistic compression ####################################################### # # w/o space charge add_impactx_test(compression @@ -602,7 +602,7 @@ add_impactx_test(compression.py OFF # no plot script yet ) -# Kicker test ########################################################## +# Kicker test ################################################################# # # w/o space charge add_impactx_test(kicker @@ -640,7 +640,58 @@ add_impactx_test(hvkicker_madx.py OFF # no plot script yet ) -# IOTA s-dependent nonlinear lens test ########################################################## +# FODO + RF EPAC2004 ########################################################## +# +# with space charge +add_impactx_test(fodo_rf_sc + examples/epac2004_benchmarks/input_fodo_rf_SC.in + ON # ImpactX MPI-parallel + OFF # ImpactX Python interface + examples/epac2004_benchmarks/analysis_fodo_rf_SC.py + OFF # no plot script yet +) +add_impactx_test(fodo_rf_sc.py + examples/epac2004_benchmarks/run_fodo_rf_SC.py + OFF # ImpactX MPI-parallel + ON # ImpactX Python interface + examples/epac2004_benchmarks/analysis_fodo_rf_SC.py + OFF # no plot script yet +) + +# Thermal Beam EPAC2004 ####################################################### +# +# with space charge +add_impactx_test(thermal + examples/epac2004_benchmarks/input_thermal.in + ON # ImpactX MPI-parallel + OFF # ImpactX Python interface + examples/epac2004_benchmarks/analysis_thermal.py + OFF # no plot script yet +) + +# Bithermal Beam EPAC2004 ##################################################### +# +# with space charge +add_impactx_test(bithermal + examples/epac2004_benchmarks/input_bithermal.in + ON # ImpactX MPI-parallel + OFF # ImpactX Python interface + examples/epac2004_benchmarks/analysis_bithermal.py + examples/epac2004_benchmarks/plot_bithermal.py +) + +# Python: Bithermal Beam EPAC2004 ############################################# +# +# with space charge +add_impactx_test(bithermal.py + examples/epac2004_benchmarks/run_bithermal.py + ON # ImpactX MPI-parallel + ON # ImpactX Python interface + examples/epac2004_benchmarks/analysis_bithermal.py + examples/epac2004_benchmarks/plot_bithermal.py +) + +# IOTA s-dependent nonlinear lens test ######################################## # # w/o space charge add_impactx_test(IOTA_nll @@ -658,7 +709,7 @@ add_impactx_test(IOTA_nll.py OFF # no plot script yet ) -# IOTA nonlinear lattice test ########################################################## +# IOTA nonlinear lattice test ################################################# # # w/o space charge add_impactx_test(IOTA_lattice @@ -676,7 +727,7 @@ add_impactx_test(IOTA_lattice.py OFF # no plot script yet ) -# Thin dipole ######################################################## +# Thin dipole ################################################################# # # w/o space charge add_impactx_test(thin_dipole @@ -694,7 +745,7 @@ add_impactx_test(thin_dipole.py OFF # no plot script yet ) -# Aperture collimation ############################################################ +# Aperture collimation ######################################################## # # w/o space charge add_impactx_test(aperture @@ -712,7 +763,7 @@ add_impactx_test(aperture.py OFF # no plot script yet ) -# Apochromat example ######################################################## +# Apochromat example ########################################################## # # w/o space charge add_impactx_test(apochromat diff --git a/examples/epac2004_benchmarks/README.rst b/examples/epac2004_benchmarks/README.rst new file mode 100644 index 000000000..c9583fabc --- /dev/null +++ b/examples/epac2004_benchmarks/README.rst @@ -0,0 +1,219 @@ +.. _examples-fodo-rf-sc: + +Cold Beam in a FODO Channel with RF Cavities (and Space Charge) +=============================================================== + +This example is based on the subsection of the same name in: +R. D. Ryne et al, "A Test Suite of Space-Charge Problems for Code Benchmarking", in Proc. EPAC2004, Lucerne, Switzerland. + +See additional documentation in: +C. E. Mitchell et al, "ImpactX Modeling of Benchmark Tests for Space Charge Validation", in Proc. HB2023, Geneva, Switzerland. + +A cold (zero momentum spread), uniform density, 250 MeV, 143 pC proton bunch propagates in a FODO lattice with 700 MHz RF +cavities added for longitudinal confinement. The on-axis profile of the RF electric field is given by: + +.. math:: + + E(z)=\exp(-(4z)^4)\cos(\frac{5\pi}{2}\tanh(5z)). + +The beam is matched to the 3D focusing, with space charge, using the rms envelope equations. + +The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. +This is tested using the second moments of the distribution. + +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_rf_SC.py`` or +* ImpactX **executable** using an input file: ``impactx input_fodo_rf_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_rf_SC.py + :language: python3 + :caption: You can copy this file from ``examples/epac2004_benchmarks/run_fodo_rf_SC.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_fodo_rf_SC.in + :language: ini + :caption: You can copy this file from ``examples/epac2004_benchmarks/input_fodo_rf_SC.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_fodo_rf_SC.py`` + + .. literalinclude:: analysis_fodo_rf_SC.py + :language: python3 + :caption: You can copy this file from ``examples/epac2004_benchmarks/analysis_fodo_rf_SC.py``. + + + +.. _examples-thermal-beam: + +Thermal Beam in a Constant Focusing Channel (with Space Charge) +=============================================================== + +This example is based on the subsection of the same name in: +R. D. Ryne et al, "A Test Suite of Space-Charge Problems for Code Benchmarking", in Proc. EPAC2004, Lucerne, Switzerland. + +See additional documentation in: +C. E. Mitchell et al, "ImpactX Modeling of Benchmark Tests for Space Charge Validation", in Proc. HB2023, Geneva, Switzerland. + +This example illustrates a stationary solution of the Vlasov-Poisson equations with spherical symmetry (in the beam +rest frame). The distribution represents a thermal equilibrium of the form: + +.. math:: + + f=C\exp(-H/kT), + +where :math:`C` and :math:`kT` are constants, and :math:`H` denotes the self-consistent Hamiltonian with space charge. + +In this example, a 0.1 MeV, 143 pC proton bunch with :math:`kT=36\times 10^{-6}` propagates in a constant focusing lattice +with 3D isotropic focusing. (The isotropy is exact in the beam rest frame.) + +The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. +This is tested using the second moments of the distribution. + +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_thermal.py`` or +* ImpactX **executable** using an input file: ``impactx input_thermal.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_thermal.py + :language: python3 + :caption: You can copy this file from ``examples/epac2004_benchmarks/run_thermal.py``. + + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_thermal.in + :language: ini + :caption: You can copy this file from ``examples/epac2004_benchmarks/input_thermal.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_thermal.py`` + + .. literalinclude:: analysis_thermal.py + :language: python3 + :caption: You can copy this file from ``examples/epac2004_benchmarks/analysis_thermal.py``. + + + +.. _examples-bithermal-beam: + +Bithermal Beam in a Constant Focusing Channel (with Space Charge) +================================================================= + +This example is based on the subsection of the same name in: +R. D. Ryne et al, "A Test Suite of Space-Charge Problems for Code Benchmarking", in Proc. EPAC2004, Lucerne, Switzerland. + +See additional documentation in: +C. E. Mitchell et al, "ImpactX Modeling of Benchmark Tests for Space Charge Validation", in Proc. HB2023, Geneva, Switzerland. + +This example illustrates a stationary solution of the Vlasov-Poisson equations with spherical symmetry (in the beam rest frame). +It provides a self-consistent model of a 3D bunch with a nontrivial core-halo distribution. + +The distribution represents a bithermal stationary distribution of the form: + +.. math:: + + f=c_1\exp(-H/kT_1)+c_2\exp(-H/kT_2), + +where :math:`c_j`, :math:`kT_j` :math:`(j=1,2)` are constants, and :math:`H` denotes the self-consistent Hamiltonian with space charge. + +In this example, a 0.1 MeV, 143 pC proton bunch with :math:`kT_1=36\times 10^{-6}` and :math:`kT_1=900\times 10^{-6}` propagates in a constant focusing lattice +with 3D isotropic focusing. +(The isotropy is exact in the beam rest frame.) +5% of the total charge lies in the second (halo) population. + +The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. +This is tested using the second moments of the distribution. + +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_bithermal.py`` or +* ImpactX **executable** using an input file: ``impactx input_bithermal.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_bithermal.py + :language: python3 + :caption: You can copy this file from ``examples/epac2004_benchmarks/run_bithermal.py``. + + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_bithermal.in + :language: ini + :caption: You can copy this file from ``examples/epac2004_benchmarks/input_bithermal.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_bithermal.py`` + + .. literalinclude:: analysis_bithermal.py + :language: python3 + :caption: You can copy this file from ``examples/epac2004_benchmarks/analysis_bithermal.py``. + + +Visualize +--------- + +You can run the following script to visualize the initial and final beam distribution: + +.. dropdown:: Script ``plot_bithermal.py`` + + .. literalinclude:: plot_bithermal.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/plot_bithermal.py``. + +.. figure:: https://user-images.githubusercontent.com/1353258/294003440-b16185c7-2573-48d9-8998-17e116721ab5.png + :alt: Initial and final beam distribution when running with full resolution (see inline comments in the input file/script). The bithermal distribution should stay static in this test. + + Initial and final beam distribution when running with full resolution (see inline comments in the input file/script). + The bithermal distribution should stay static in this test. diff --git a/examples/epac2004_benchmarks/analysis_bithermal.py b/examples/epac2004_benchmarks/analysis_bithermal.py new file mode 100755 index 000000000..64f54ebd6 --- /dev/null +++ b/examples/epac2004_benchmarks/analysis_bithermal.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# 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() +final = series.iterations[last_step].particles["beam"].to_df() + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +print("Initial Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 3.5 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 2.751162e-03, + 2.751725e-03, + 1.884003e-01, + 2.449966e-05, + 2.451077e-05, + 2.444195e-05, + ], + rtol=rtol, + atol=atol, +) + + +print("") +print("Final Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 6000 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 2.751162e-03, + 2.751725e-03, + 1.884003e-01, + 2.449966e-05, + 2.451077e-05, + 2.444195e-05, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/epac2004_benchmarks/analysis_fodo_rf_SC.py b/examples/epac2004_benchmarks/analysis_fodo_rf_SC.py new file mode 100755 index 000000000..2b78ab9b7 --- /dev/null +++ b/examples/epac2004_benchmarks/analysis_fodo_rf_SC.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# 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() +final = series.iterations[last_step].particles["beam"].to_df() + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +print("Initial Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt: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( + [sigx, sigy, sigt], + [9.84722273e-4, 6.96967278e-4, 4.486799242214e-03], + rtol=rtol, + atol=atol, +) + +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 4.0e-8 +print(f" atol={atol}") + +assert np.allclose( + [emittance_x, emittance_y, emittance_t], + [0.0, 0.0, 0.0], + atol=atol, +) + + +print("") +print("Final Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") + +atol = 0.0 # ignored +rtol = 3.5 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt], + [9.84722273e-4, 6.96967278e-4, 4.486799242214e-03], + rtol=rtol, + atol=atol, +) + +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 4.0e-8 +print(f" atol={atol}") + +assert np.allclose( + [emittance_x, emittance_y, emittance_t], + [0.0, 0.0, 0.0], + atol=atol, +) diff --git a/examples/epac2004_benchmarks/analysis_thermal.py b/examples/epac2004_benchmarks/analysis_thermal.py new file mode 100755 index 000000000..86a1ea1f3 --- /dev/null +++ b/examples/epac2004_benchmarks/analysis_thermal.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# 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() +final = series.iterations[last_step].particles["beam"].to_df() + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +print("Initial Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 3.5 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 2.569162e-03, + 2.569557e-03, + 1.757951e-01, + 1.540773e-05, + 1.541883e-05, + 1.538814e-05, + ], + rtol=rtol, + atol=atol, +) + + +print("") +print("Final Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 6.0 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 2.569162e-03, + 2.569557e-03, + 1.757951e-01, + 1.540773e-05, + 1.541883e-05, + 1.538814e-05, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/epac2004_benchmarks/input_bithermal.in b/examples/epac2004_benchmarks/input_bithermal.in new file mode 100644 index 000000000..2f540ea88 --- /dev/null +++ b/examples/epac2004_benchmarks/input_bithermal.in @@ -0,0 +1,48 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +#beam.npart = 100000000 #full resolution +beam.npart = 10000 +beam.units = static +beam.kin_energy = 0.1 +beam.charge = 1.4285714285714285714e-10 +beam.particle = proton +beam.distribution = thermal +beam.k = 6.283185307179586 +beam.kT = 36.0e-6 +beam.kT_halo = 900.0e-6 +beam.halo = 0.05 +beam.normalize = 0.54226 +beam.normalize_halo = 0.08195 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor constf1 monitor + +monitor.type = beam_monitor +monitor.backend = h5 + +constf1.type = constf +constf1.ds = 10.0 +constf1.kx = 6.283185307179586 +constf1.ky = 6.283185307179586 +constf1.kt = 6.283185307179586 +#constf1.nslice = 400 #full resolution +constf1.nslice = 50 + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = true + +#amr.n_cell = 128 128 128 #full resolution +amr.n_cell = 64 64 64 +geometry.prob_relative = 3.0 + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = false diff --git a/examples/epac2004_benchmarks/input_fodo_rf_SC.in b/examples/epac2004_benchmarks/input_fodo_rf_SC.in new file mode 100644 index 000000000..ade26631c --- /dev/null +++ b/examples/epac2004_benchmarks/input_fodo_rf_SC.in @@ -0,0 +1,128 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 250.0 +beam.charge = 1.42857142857142865e-10 +beam.particle = proton +beam.distribution = kurth6d +beam.sigmaX = 9.84722273e-4 +beam.sigmaY = 6.96967278e-4 +beam.sigmaT = 4.486799242214e-03 +beam.sigmaPx = 0.0 +beam.sigmaPy = 0.0 +beam.sigmaPt = 0.0 +beam.muxpx = 0.0 +beam.muypy = 0.0 +beam.mutpt = 0.0 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor fquad dr gapa1 dr dquad dr gapb1 dr fquad monitor + +monitor.type = beam_monitor +monitor.backend = h5 + +dr.type = drift +dr.ds = 0.1 +dr.nslice = 4 + +fquad.type = quad +fquad.ds = 0.15 +fquad.k = 2.4669749766168163 +fquad.nslice = 6 + +dquad.type = quad +dquad.ds = 0.3 +dquad.k = -2.4669749766168163 +dquad.nslice = 12 + +gapa1.type = rfcavity +gapa1.ds = 1.0 +gapa1.escale = 0.042631556991578 +gapa1.freq = 7.0e8 +gapa1.phase = 45.0 +gapa1.mapsteps = 100 +gapa1.nslice = 10 +gapa1.cos_coefficients = \ + 0.120864178375839 \ + -0.044057987631337 \ + -0.209107290958498 \ + -0.019831226655815 \ + 0.290428111491964 \ + 0.381974267375227 \ + 0.276801212694382 \ + 0.148265085353012 \ + 0.068569351192205 \ + 0.0290155855315696 \ + 0.011281649986680 \ + 0.004108501632832 \ + 0.0014277644197320 \ + 0.000474212125404 \ + 0.000151675768439 \ + 0.000047031436898 \ + 0.000014154595193 \ + 4.154741658e-6 \ + 1.191423909e-6 \ + 3.348293360e-7 \ + 9.203061700e-8 \ + 2.515007200e-8 \ + 6.478108000e-9 \ + 1.912531000e-9 \ + 2.925600000e-10 +gapa1.sin_coefficients = 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 \ + 0 0 0 0 0 0 0 + +gapb1.type = rfcavity +gapb1.ds = 1.0 +gapb1.escale = 0.042631556991578 +gapb1.freq = 7.0e8 +gapb1.phase = -1.0 +gapb1.mapsteps = 100 +gapb1.nslice = 10 +gapb1.cos_coefficients = \ + 0.120864178375839 \ + -0.044057987631337 \ + -0.209107290958498 \ + -0.019831226655815 \ + 0.290428111491964 \ + 0.381974267375227 \ + 0.276801212694382 \ + 0.148265085353012 \ + 0.068569351192205 \ + 0.0290155855315696 \ + 0.011281649986680 \ + 0.004108501632832 \ + 0.0014277644197320 \ + 0.000474212125404 \ + 0.000151675768439 \ + 0.000047031436898 \ + 0.000014154595193 \ + 4.154741658e-6 \ + 1.191423909e-6 \ + 3.348293360e-7 \ + 9.203061700e-8 \ + 2.515007200e-8 \ + 6.478108000e-9 \ + 1.912531000e-9 \ + 2.925600000e-10 +gapb1.sin_coefficients = 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 \ + 0 0 0 0 0 0 0 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = true + +amr.n_cell = 56 56 64 +geometry.prob_relative = 4.0 + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/epac2004_benchmarks/input_thermal.in b/examples/epac2004_benchmarks/input_thermal.in new file mode 100644 index 000000000..75f158951 --- /dev/null +++ b/examples/epac2004_benchmarks/input_thermal.in @@ -0,0 +1,45 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +#beam.npart = 100000000 #full resolution +beam.npart = 10000 +beam.units = static +beam.kin_energy = 0.1 +beam.charge = 1.4285714285714285714e-10 +beam.particle = proton +beam.distribution = thermal +beam.k = 6.283185307179586 +beam.kT = 36.0e-6 +beam.normalize = 0.41604661 + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor constf1 monitor + +monitor.type = beam_monitor +monitor.backend = h5 + +constf1.type = constf +constf1.ds = 10.0 +constf1.kx = 6.283185307179586 +constf1.ky = 6.283185307179586 +constf1.kt = 6.283185307179586 +constf1.nslice = 400 #full resolution +#constf1.nslice = 50 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = true + +#amr.n_cell = 128 128 128 #full resolution +amr.n_cell = 64 64 64 +geometry.prob_relative = 3.0 + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = false diff --git a/examples/epac2004_benchmarks/plot_bithermal.py b/examples/epac2004_benchmarks/plot_bithermal.py new file mode 100755 index 000000000..97657448e --- /dev/null +++ b/examples/epac2004_benchmarks/plot_bithermal.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + +import argparse +from math import pi + +import matplotlib.pyplot as plt +import numpy as np +import openpmd_api as io + +# options to run this script +parser = argparse.ArgumentParser(description="Plot the Bithermal benchmark.") +parser.add_argument( + "--save-png", action="store_true", help="non-interactive run: save to PNGs" +) +args = parser.parse_args() + + +# initial/final beam +series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) +last_step = list(series.iterations)[-1] +initial_beam = series.iterations[1].particles["beam"].to_df() +final_beam = series.iterations[last_step].particles["beam"].to_df() + +# Constants +w1 = 0.95 +w2 = 0.05 +bg = 0.0146003 +Min = 0.0 +Max = 0.025 +Np = 100000001 +n = 300 + + +# Function for radius calculation +def r(x, y, z): + return np.sqrt(x**2 + y**2 + z**2) + + +# Calculate radius and bin data +initial_radii = r( + bg * initial_beam["position_t"], + initial_beam["position_x"], + initial_beam["position_y"], +) +initial_hist, bin_edges = np.histogram(initial_radii, bins=n, range=(Min, Max)) +initial_bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:]) + +final_radii = r( + bg * final_beam["position_t"], final_beam["position_x"], final_beam["position_y"] +) +final_hist, _ = np.histogram(final_radii, bins=n, range=(Min, Max)) + +# dr (m) +initial_r = initial_hist / (Np * (bin_edges[1] - bin_edges[0])) +final_r = final_hist / (Np * (bin_edges[1] - bin_edges[0])) + +# Plotting +plt.figure(figsize=(10, 6)) +plt.xscale("linear") +plt.yscale("log") +plt.xlim([Min, Max]) +plt.ylim([0.1, 1.6e6]) +plt.xlabel("r (m)", fontsize=30) +plt.xticks(fontsize=25) +plt.yticks(fontsize=25) +plt.grid(True) + +# Plot the data +plt.plot( + initial_bin_centers, + initial_r / (4.0 * pi * (initial_bin_centers) ** 2), + label="Initial beam", + linewidth=2, +) +plt.plot( + initial_bin_centers, + final_r / (4.0 * pi * (initial_bin_centers) ** 2), + label="Final beam", + linewidth=2, + linestyle="dotted", +) + +# Show plot +plt.legend(fontsize=20) + +plt.tight_layout() +if args.save_png: + plt.savefig("bithermal.png") +else: + plt.show() diff --git a/examples/epac2004_benchmarks/run_bithermal.py b/examples/epac2004_benchmarks/run_bithermal.py new file mode 100755 index 000000000..871921b80 --- /dev/null +++ b/examples/epac2004_benchmarks/run_bithermal.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Marco Garten, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex.space3d as amr +from impactx import ImpactX, RefPart, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +# sim.n_cell = [128, 128, 128] # full resolution +sim.n_cell = [64, 64, 64] +sim.particle_shape = 2 # B-spline order +sim.space_charge = True +sim.dynamic_size = True +sim.prob_relative = [3.0] + +# beam diagnostics +sim.slice_step_diagnostics = False + +# domain decomposition & space charge mesh +sim.init_grids() + +# beam parameters +kin_energy_MeV = 0.1 # reference energy +bunch_charge_C = 1.4285714285714285714e-10 # used with space charge +# npart = 100000000 # full resolution +npart = 10000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Thermal( + k=6.283185307179586, + kT=36.0e-6, + kT_halo=900.0e-6, + normalize=0.54226, + normalize_halo=0.08195, + halo=0.05, +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice +sim.lattice.append(monitor) + +constf = elements.ConstF( + ds=10.0, + kx=6.283185307179586, + ky=6.283185307179586, + kt=6.283185307179586, + # nslice=400, # full resolution + nslice=50, +) + +sim.lattice.append(constf) +sim.lattice.append(monitor) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amr.finalize() diff --git a/examples/epac2004_benchmarks/run_fodo_rf_SC.py b/examples/epac2004_benchmarks/run_fodo_rf_SC.py new file mode 100755 index 000000000..7112e975b --- /dev/null +++ b/examples/epac2004_benchmarks/run_fodo_rf_SC.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Marco Garten, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex.space3d as amr +from impactx import ImpactX, RefPart, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.n_cell = [56, 56, 64] +sim.particle_shape = 2 # B-spline order +sim.space_charge = True +sim.dynamic_size = True +sim.prob_relative = [4.0] + +# beam diagnostics +sim.slice_step_diagnostics = False + +# domain decomposition & space charge mesh +sim.init_grids() + +# beam parameters +kin_energy_MeV = 250.0 # reference energy +bunch_charge_C = 1.42857142857142865e-10 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Kurth6D( + sigmaX=9.84722273e-4, + sigmaY=6.96967278e-4, + sigmaT=4.486799242214e-03, + sigmaPx=0.0, + sigmaPy=0.0, + sigmaPt=0.0, +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice +sim.lattice.append(monitor) + +# Quad elements +fquad = elements.Quad(ds=0.15, k=2.4669749766168163, nslice=6) +dquad = elements.Quad(ds=0.3, k=-2.4669749766168163, nslice=12) + +# Drift element +dr = elements.Drift(ds=0.1, nslice=4) + +# RF cavity elements +gapa1 = elements.RFCavity( + ds=1.0, + escale=0.042631556991578, + freq=7.0e8, + phase=45.0, + cos_coefficients=[ + 0.120864178375839, + -0.044057987631337, + -0.209107290958498, + -0.019831226655815, + 0.290428111491964, + 0.381974267375227, + 0.276801212694382, + 0.148265085353012, + 0.068569351192205, + 0.0290155855315696, + 0.011281649986680, + 0.004108501632832, + 0.0014277644197320, + 0.000474212125404, + 0.000151675768439, + 0.000047031436898, + 0.000014154595193, + 4.154741658e-6, + 1.191423909e-6, + 3.348293360e-7, + 9.203061700e-8, + 2.515007200e-8, + 6.478108000e-9, + 1.912531000e-9, + 2.925600000e-10, + ], + sin_coefficients=[ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ], + mapsteps=100, + nslice=4, +) + +gapb1 = elements.RFCavity( + ds=1.0, + escale=0.042631556991578, + freq=7.0e8, + phase=-1.0, + cos_coefficients=[ + 0.120864178375839, + -0.044057987631337, + -0.209107290958498, + -0.019831226655815, + 0.290428111491964, + 0.381974267375227, + 0.276801212694382, + 0.148265085353012, + 0.068569351192205, + 0.0290155855315696, + 0.011281649986680, + 0.004108501632832, + 0.0014277644197320, + 0.000474212125404, + 0.000151675768439, + 0.000047031436898, + 0.000014154595193, + 4.154741658e-6, + 1.191423909e-6, + 3.348293360e-7, + 9.203061700e-8, + 2.515007200e-8, + 6.478108000e-9, + 1.912531000e-9, + 2.925600000e-10, + ], + sin_coefficients=[ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ], + mapsteps=100, + nslice=4, +) + +lattice_no_drifts = [fquad, gapa1, dquad, gapb1, fquad] + +# set first lattice element +sim.lattice.append(lattice_no_drifts[0]) + +# intersperse all remaining elements of the lattice with a drift element +for element in lattice_no_drifts[1:]: + sim.lattice.extend([dr, element]) + +sim.lattice.append(monitor) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amr.finalize() diff --git a/examples/epac2004_benchmarks/run_thermal.py b/examples/epac2004_benchmarks/run_thermal.py new file mode 100755 index 000000000..0211077ea --- /dev/null +++ b/examples/epac2004_benchmarks/run_thermal.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Marco Garten, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex.space3d as amr +from impactx import ImpactX, RefPart, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.n_cell = [56, 56, 64] +sim.particle_shape = 2 # B-spline order +sim.space_charge = True +sim.dynamic_size = True +sim.prob_relative = [4.0] + +# beam diagnostics +sim.slice_step_diagnostics = False + +# domain decomposition & space charge mesh +sim.init_grids() + +# beam parameters +kin_energy_MeV = 0.1 # reference energy +bunch_charge_C = 1.4285714285714285714e-10 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Thermal( + k=6.283185307179586, + kT=36.0e-6, + kT_halo=36.0e-6, + normalize=0.41604661, + normalize_halo=0.0, + w_halo=0.0, +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice +sim.lattice.append(monitor) + +constf = elements.Constf( + ds=10.0, + kx=6.283185307179586, + ky=6.283185307179586, + kt=6.283185307179586, + nslice=400, +) + +# set first lattice element +sim.lattice.append(constf) +sim.lattice.append(monitor) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amr.finalize() diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 51194fd30..5e0dd0700 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -56,6 +56,7 @@ namespace impactx ); } + // init particles amrex::Vector x, y, t; amrex::Vector px, py, pt; amrex::ParticleReal ix, iy, it, ipx, ipy, ipt; @@ -74,6 +75,10 @@ namespace impactx amrex::ParticleReal(npart); std::visit([&](auto&& distribution){ + // initialize distributions + distribution.initialize(bunch_charge, ref); + + // alloc data for particle attributes x.reserve(npart_this_proc); y.reserve(npart_this_proc); t.reserve(npart_this_proc); @@ -296,6 +301,22 @@ namespace impactx add_particles(bunch_charge, triangle, npart); + } else if (distribution_type == "thermal") { + amrex::ParticleReal k, kT, kT_halo, normalize, normalize_halo; + amrex::ParticleReal halo = 0.0; + pp_dist.get("k", k); + pp_dist.get("kT", kT); + kT_halo = kT; + pp_dist.get("normalize", normalize); + normalize_halo = normalize; + pp_dist.query("kT_halo", kT_halo); + pp_dist.query("normalize_halo", normalize_halo); + pp_dist.query("halo", halo); + + distribution::KnownDistributions thermal(distribution::Thermal(k, kT, kT_halo, normalize, normalize_halo, halo)); + + add_particles(bunch_charge, thermal, npart); + } else { amrex::Abort("Unknown distribution: " + distribution_type); } diff --git a/src/particles/distribution/All.H b/src/particles/distribution/All.H index e3a833708..e3e9f8c12 100644 --- a/src/particles/distribution/All.H +++ b/src/particles/distribution/All.H @@ -16,6 +16,7 @@ #include "KVdist.H" #include "None.H" #include "Semigaussian.H" +#include "Thermal.H" #include "Triangle.H" #include "Waterbag.H" @@ -30,6 +31,7 @@ namespace impactx::distribution Kurth4D, Kurth6D, KVdist, + Thermal, Triangle, Semigaussian, Waterbag diff --git a/src/particles/distribution/Gaussian.H b/src/particles/distribution/Gaussian.H index c542d943d..83b19e1da 100644 --- a/src/particles/distribution/Gaussian.H +++ b/src/particles/distribution/Gaussian.H @@ -10,6 +10,8 @@ #ifndef IMPACTX_DISTRIBUTION_GAUSSIAN #define IMPACTX_DISTRIBUTION_GAUSSIAN +#include "particles/ReferenceParticle.H" + #include #include @@ -47,6 +49,17 @@ namespace impactx::distribution { } + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } + /** Return 1 6D particle coordinate * * @param x particle position in x diff --git a/src/particles/distribution/KVdist.H b/src/particles/distribution/KVdist.H index 233283d93..65dac699c 100644 --- a/src/particles/distribution/KVdist.H +++ b/src/particles/distribution/KVdist.H @@ -10,6 +10,8 @@ #ifndef IMPACTX_DISTRIBUTION_KVDIST #define IMPACTX_DISTRIBUTION_KVDIST +#include "particles/ReferenceParticle.H" + #include #include @@ -48,6 +50,17 @@ namespace impactx::distribution { } + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } + /** Return 1 6D particle coordinate * * @param x particle position in x diff --git a/src/particles/distribution/Kurth4D.H b/src/particles/distribution/Kurth4D.H index 68d1b5ed8..988448241 100644 --- a/src/particles/distribution/Kurth4D.H +++ b/src/particles/distribution/Kurth4D.H @@ -10,6 +10,8 @@ #ifndef IMPACTX_DISTRIBUTION_KURTH4D #define IMPACTX_DISTRIBUTION_KURTH4D +#include "particles/ReferenceParticle.H" + #include #include @@ -48,6 +50,17 @@ namespace impactx::distribution { } + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } + /** Return 1 6D particle coordinate * * @param x particle position in x diff --git a/src/particles/distribution/Kurth6D.H b/src/particles/distribution/Kurth6D.H index b411aa7cf..d1ed44ed4 100644 --- a/src/particles/distribution/Kurth6D.H +++ b/src/particles/distribution/Kurth6D.H @@ -10,6 +10,8 @@ #ifndef IMPACTX_DISTRIBUTION_KURTH6D #define IMPACTX_DISTRIBUTION_KURTH6D +#include "particles/ReferenceParticle.H" + #include #include @@ -50,6 +52,17 @@ namespace impactx::distribution { } + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } + /** Return 1 6D particle coordinate * * @param x particle position in x diff --git a/src/particles/distribution/None.H b/src/particles/distribution/None.H index 517455752..79e3657e3 100644 --- a/src/particles/distribution/None.H +++ b/src/particles/distribution/None.H @@ -10,6 +10,8 @@ #ifndef IMPACTX_DISTRIBUTION_NONE #define IMPACTX_DISTRIBUTION_NONE +#include "particles/ReferenceParticle.H" + #include #include @@ -20,9 +22,20 @@ namespace impactx::distribution { /** This distribution sets all values to zero. */ - None() - { - } + None() + { + } + + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } /** Return 1 6D particle coordinate * diff --git a/src/particles/distribution/Semigaussian.H b/src/particles/distribution/Semigaussian.H index 59cfdd2ca..9b91e568a 100644 --- a/src/particles/distribution/Semigaussian.H +++ b/src/particles/distribution/Semigaussian.H @@ -10,6 +10,8 @@ #ifndef IMPACTX_DISTRIBUTION_SEMIGAUSSIAN #define IMPACTX_DISTRIBUTION_SEMIGAUSSIAN +#include "particles/ReferenceParticle.H" + #include #include @@ -48,6 +50,17 @@ namespace impactx::distribution { } + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } + /** Return 1 6D particle coordinate * * @param x particle position in x diff --git a/src/particles/distribution/Thermal.H b/src/particles/distribution/Thermal.H new file mode 100644 index 000000000..b2cc31819 --- /dev/null +++ b/src/particles/distribution/Thermal.H @@ -0,0 +1,403 @@ +/* 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, Axel Huebl, Chad Mitchell + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_DISTRIBUTION_THERMAL +#define IMPACTX_DISTRIBUTION_THERMAL + +#include "particles/ReferenceParticle.H" + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include + + +namespace impactx +{ +namespace distribution +{ + struct ThermalData + { + static constexpr amrex::ParticleReal tolerance = 1.0e-3; ///< tolerance for matching condition + static constexpr amrex::ParticleReal rin = 1.0e-10; ///< initial r value for numerical integration + static constexpr amrex::ParticleReal rout = 10.0; ///< final r value for numerical integration + static constexpr int nsteps = 2000; ///< number of radial steps for numerical integration + + struct Rprofile + { + amrex::ParticleReal f1; ///< cumulative distribution of first population + amrex::ParticleReal f2; ///< cumulative distribution of second population + amrex::ParticleReal phi1; ///< potential generated by first population + amrex::ParticleReal phi2; ///< potential generated by second population + amrex::ParticleReal p1; ///< normalization constant of first population + amrex::ParticleReal p2; ///< normalization constant of second population + amrex::ParticleReal rmin; ///< minimum r value for tabulated cdf + amrex::ParticleReal rmax; ///< maximum r value for tabulated cdf + int nbins; ///< number of radial bins for tabulated cdf + amrex::ParticleReal cdf1[2001]; ///< tabulated cumulative distribution (first) + amrex::ParticleReal cdf2[2001]; ///< tabulated cumulative distribution (second) + amrex::ParticleReal Cintensity; ///< space charge intensity parameter + amrex::ParticleReal bg; ///< reference value of relativistic beta*gamma + amrex::ParticleReal k; ///< linear focusing strength (1/meters) + amrex::ParticleReal T1; ///< temperature k*T of the primary (core) population + amrex::ParticleReal T2; ///< temperature k*T of the secondary (halo) population + amrex::ParticleReal w; ///< weight of the secondary (halo) population + + Rprofile (amrex::ParticleReal kin, + amrex::ParticleReal T1in, + amrex::ParticleReal T2in, + amrex::ParticleReal p1in, + amrex::ParticleReal p2in, + amrex::ParticleReal win) + { + k = kin; + T1 = T1in; + T2 = T2in; + p1 = p1in; + p2 = p2in; + w = win; + } + }; + + /** Populate the radial CDF data. + * + * @param[in] bunch_charge the bunch charge in C + * @param[in] refpart the reference particle + * @param[inout] data the data to contain the radial CDF + */ + static void + generate_radial_dist (amrex::ParticleReal bunch_charge, RefPart const & refpart, Rprofile & data) + + { + using namespace amrex::literals; // for _rt and _prt + using namespace ablastr::constant::math; // for pi + + // Get relativistic beta*gamma + amrex::ParticleReal bg = refpart.beta_gamma(); + data.bg = bg; + + // Get reference particle rest energy (eV) and charge (q_e) + amrex::ParticleReal Erest = refpart.mass_MeV()*1.0e6; + amrex::ParticleReal q_e = refpart.charge_qe(); + + // Set space charge intensity + data.Cintensity = q_e*bunch_charge/(pow(bg,2)*Erest*ablastr::constant::SI::ep0); + + // Set minimum and maximum radius + amrex::ParticleReal r_scale = matched_scale_radius(data); + amrex::ParticleReal rmin = rin*r_scale; + amrex::ParticleReal rmax = rout*r_scale; + // amrex::PrintToFile("equilibrium_params.out") << r_scale << " " << data.Cintensity << "\n"; + + // Scale the parameters p1 and p2 + amrex::ParticleReal rt2pi = sqrt(2.0_prt*pi); + amrex::ParticleReal p_scale = pow(r_scale*rt2pi,-3); + data.p1 = data.p1*p_scale; + data.p2 = data.p2*p_scale; + + // store integration parameters + data.nbins = nsteps; + data.rmin = rmin; + data.rmax = rmax; + + // set initial conditions + data.f1 = 0.0_prt; + data.f2 = 0.0_prt; + data.phi1 = 0.0_prt; + data.phi2 = 0.0_prt; + data.cdf1[0] = 0.0_prt; + data.cdf2[0] = 0.0_prt; + + // integrate ODEs for the radial density + integrate(data,rmin,rmax,nsteps); + + // a search over normalization parameters p1, p2 can be added here + + } + + static amrex::ParticleReal + matched_scale_radius(Rprofile & data) + { + using namespace amrex::literals; // for _rt and _prt + using namespace ablastr::constant::math; // for pi + + amrex::ParticleReal k = data.k; + amrex::ParticleReal kT = (1.0_prt-data.w)*data.T1 + data.w*data.T2; + amrex::ParticleReal a = data.Cintensity/(4.0_prt*pi*5.0_prt*sqrt(5.0_prt)); + amrex::ParticleReal rscale = sqrt(kT + pow(a*k,2.0/3.0))/k; + + return rscale; + } + + static void + integrate ( + Rprofile & data, + amrex::ParticleReal const in, + amrex::ParticleReal const out, + int const steps + ) + { + using namespace amrex::literals; // for _rt and _prt + + // initialize numerical integration parameters + amrex::ParticleReal const tau = (out-in)/steps; + amrex::ParticleReal const half = tau/2.0_prt; + + // initialize the value of the independent variable + amrex::ParticleReal reval = in; + + // loop over integration steps + for (int j=0; j < steps; ++j) + { + // for a second-order Strang splitting + map1(half,data,reval); + map2(tau,data,reval); + map1(half,data,reval); + // store tabulated CDF + data.cdf1[j+1] = data.f1; + data.cdf2[j+1] = data.f2; + // write cumulative density to file for debugging + /* comment in for debugging + amrex::PrintToFile("density1.out") << reval << " " << data.f1 << "\n"; + amrex::PrintToFile("density2.out") << reval << " " << data.f2 << "\n"; + amrex::PrintToFile("phi1.out") << reval << " " << data.phi1 << "\n"; + amrex::PrintToFile("phi2.out") << reval << " " << data.phi2 << "\n"; + */ + } + + } + + static void + map1 ( + amrex::ParticleReal const tau, + Rprofile & data, + amrex::ParticleReal & reval + ) + { + using namespace amrex::literals; // for _rt and _prt + using namespace ablastr::constant::math; // for pi + + amrex::ParticleReal const f1 = data.f1; + amrex::ParticleReal const f2 = data.f2; + amrex::ParticleReal const phi1 = data.phi1; + amrex::ParticleReal const phi2 = data.phi2; + amrex::ParticleReal const r = reval; + + // Apply map to update phi1, phi2, and r: + reval = r + tau; + data.phi1 = phi1 + f1/(4.0_prt*pi*reval) - f1/(4.0_prt*pi*r); + data.phi2 = phi2 + f2/(4.0_prt*pi*reval) - f2/(4.0_prt*pi*r);; + data.f1 = f1; + data.f2 = f2; + + } + + static void + map2 ( + amrex::ParticleReal const tau, + Rprofile & data, + amrex::ParticleReal & reval + ) + { + using namespace amrex::literals; // for _rt and _prt + using namespace ablastr::constant::math; // for pi + + amrex::ParticleReal const f1 = data.f1; + amrex::ParticleReal const f2 = data.f2; + amrex::ParticleReal const phi1 = data.phi1; + amrex::ParticleReal const phi2 = data.phi2; + amrex::ParticleReal const r = reval; + amrex::ParticleReal const k = data.k; + amrex::ParticleReal const w = data.w; + amrex::ParticleReal const T1 = data.T1; + amrex::ParticleReal const T2 = data.T2; + + // Define space charge intensity parameters + amrex::ParticleReal const c1 = (1.0_prt-w)*data.Cintensity; + amrex::ParticleReal const c2 = w*data.Cintensity; + + // Define intermediate quantities + amrex::ParticleReal potential = 0.0_prt; + potential = pow(k*r,2.0)/2.0_prt + c1*phi1 + c2*phi2; + amrex::ParticleReal Pdensity1 = data.p1*exp(-potential/T1); + amrex::ParticleReal Pdensity2 = data.p2*exp(-potential/T2); + // amrex::ParticleReal Pdensity_tot = (1.0_prt-w)*Pdensity1 + w*Pdensity2; + // amrex::PrintToFile("Pdensity.out") << reval << " " << Pdensity_tot << "\n"; + + // Apply map to update f1 and f2: + data.phi1 = phi1; + data.phi2 = phi2; + data.f1 = f1 + tau*4.0_prt*pi*pow(r,2.0)*Pdensity1; + data.f2 = f2 + tau*4.0_prt*pi*pow(r,2.0)*Pdensity2; + reval = r; + } + }; + + struct Thermal + { + /** A stationary Thermal or Bi-thermal distribution + * + * Return sampling from a 6D bi-thermal distribution, a + * stationary solution of the Vlasov-Poisson system in + * an isotropic linear focusing channel. + * + * @param k linear focusing strength (1/meters) + * @param kT temperature k*T of the primary (core) population + * @param kT_halo temperature k*T of the secondary (halo) population + * @param normalize normalization constant of first population + * @param normalize_halo normalization constant of second population + * @param halo weight of the secondary (halo) population + */ + Thermal ( + amrex::ParticleReal k, + amrex::ParticleReal kT, + amrex::ParticleReal kT_halo, + amrex::ParticleReal normalize, + amrex::ParticleReal normalize_halo, + amrex::ParticleReal halo + ) + : m_k(k), + m_kT(kT), + m_T2(kT_halo), + m_normalize(normalize), + m_normalize_halo(normalize_halo), + m_halo(halo), + m_data(ThermalData::Rprofile(m_k, m_kT, m_T2, m_normalize, m_normalize_halo, m_halo)) + { + } + + /** Initialize the distribution. + * + * This in particular sets the m_data radial profile of the stationary beam + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize (amrex::ParticleReal bunch_charge, RefPart const & ref) + { + // Generate the struct "data" containing the radial CDF: + distribution::ThermalData::generate_radial_dist(bunch_charge, ref, m_data); + } + + /** Return 1 6D particle coordinate + * + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t + * @param px particle momentum in x + * @param py particle momentum in y + * @param pt particle momentum in t + * @param engine a random number engine (with associated state) + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + amrex::ParticleReal & AMREX_RESTRICT x, + amrex::ParticleReal & AMREX_RESTRICT y, + amrex::ParticleReal & AMREX_RESTRICT t, + amrex::ParticleReal & AMREX_RESTRICT px, + amrex::ParticleReal & AMREX_RESTRICT py, + amrex::ParticleReal & AMREX_RESTRICT pt, + amrex::RandomEngine const& engine) { + + using namespace amrex::literals; // for _rt and _prt + using namespace ablastr::constant::math; // for pi + + amrex::ParticleReal ln1,norm,u1,u2,uhalo; + amrex::ParticleReal g1,g2,g3,g4,g5,g6; + amrex::ParticleReal z,pz; + + // Use a Bernoulli random variable to select between core and halo: + // If u < w, the particle is in the halo population. + // If u >= w, the particle is in the core population. + uhalo = amrex::Random(engine); + + // Generate six standard normal random variables using Box-Muller: + u1 = amrex::Random(engine); + u2 = amrex::Random(engine); + ln1 = sqrt(-2_prt*log(u1)); + g1 = ln1*cos(2_prt*pi*u2); + g2 = ln1*sin(2_prt*pi*u2); + u1 = amrex::Random(engine); + u2 = amrex::Random(engine); + ln1 = sqrt(-2_prt*log(u1)); + g3 = ln1*cos(2_prt*pi*u2); + g4 = ln1*sin(2_prt*pi*u2); + u1 = amrex::Random(engine); + u2 = amrex::Random(engine); + ln1 = sqrt(-2_prt*log(u1)); + g5 = ln1*cos(2_prt*pi*u2); + g6 = ln1*sin(2_prt*pi*u2); + + // Scale the last three variables to produce the momenta: + amrex::ParticleReal kT = (uhalo > m_data.w) ? m_data.T1 : m_data.T2; // select core or halo value + px = sqrt(kT)*g4; + py = sqrt(kT)*g5; + pz = sqrt(kT)*g6; + + // Normalize the first three variables to produce uniform samples on the unit 3-sphere: + norm = sqrt(g1*g1+g2*g2+g3*g3); + g1 /= norm; + g2 /= norm; + g3 /= norm; + + // Collect the radial CDF returned by generate_radial_dist: + amrex::ParticleReal rmin = m_data.rmin; + amrex::ParticleReal rmax = m_data.rmax; + int const nbins = m_data.nbins; + constexpr int length = 2001; // Ideally, int const length = nbins + 1; + amrex::ParticleReal cdf[length]; + for (int n = 0; n < length; ++n) { + cdf[n] = (uhalo > m_data.w) ? m_data.cdf1[n] : m_data.cdf2[n]; //select core or halo CDF + } + for (int n = 0; n < length; ++n) { + cdf[n] = cdf[n]/cdf[nbins]; //rescale cdf to ensure the exact range [0,1] + } + + // Generate a radial coordinate from the CDF + amrex::ParticleReal u = amrex::Random(engine); + amrex::ParticleReal *ptr = amrex::lower_bound(cdf, cdf + nbins + 1, u); + int off = amrex::max(0, (int)(ptr - cdf - 1)); + amrex::ParticleReal tv = (u - cdf[off]) / (cdf[off + 1] - cdf[off]); + amrex::ParticleReal xv = (off + tv) / (float)(nbins); + amrex::ParticleReal r = rmin + (rmax - rmin) * xv; + + // Scale to produce samples with the correct radial density: + x = g1*r; + y = g2*r; + z = g3*r; + + // Transform logitudinal variables into the laboratory frame: + t = -z/(m_data.bg); + pt = -pz*(m_data.bg); + + (void) m_k; + (void) m_kT; + (void) m_T2; + (void) m_halo; + } + + private: + amrex::ParticleReal m_k; //! linear focusing strength (1/meters) + amrex::ParticleReal m_kT, m_T2; //! temperature of each particle population + amrex::ParticleReal m_normalize, m_normalize_halo; //! normalization constant of first/second population + amrex::ParticleReal m_halo; //! relative weight of halo population + ThermalData::Rprofile m_data; //! struct containing radial profile data + }; + +} // namespace distribution +} // namespace impactx + +#endif // IMPACTX_DISTRIBUTION_THERMAL diff --git a/src/particles/distribution/Triangle.H b/src/particles/distribution/Triangle.H index 0a761982a..45c3fbaaa 100644 --- a/src/particles/distribution/Triangle.H +++ b/src/particles/distribution/Triangle.H @@ -10,6 +10,8 @@ #ifndef IMPACTX_DISTRIBUTION_TRIANGLE #define IMPACTX_DISTRIBUTION_TRIANGLE +#include "particles/ReferenceParticle.H" + #include #include @@ -45,6 +47,17 @@ namespace impactx::distribution { } + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } + /** Return 1 6D particle coordinate * * @param x particle position in x diff --git a/src/particles/distribution/Waterbag.H b/src/particles/distribution/Waterbag.H index a9860e463..d41ed6eb2 100644 --- a/src/particles/distribution/Waterbag.H +++ b/src/particles/distribution/Waterbag.H @@ -10,6 +10,8 @@ #ifndef IMPACTX_DISTRIBUTION_WATERBAG #define IMPACTX_DISTRIBUTION_WATERBAG +#include "particles/ReferenceParticle.H" + #include #include @@ -47,6 +49,17 @@ namespace impactx::distribution { } + /** Initialize the distribution. + * + * Nothing to do here. + * + * @param bunch_charge charge of the beam in C + * @param ref the reference particle + */ + void initialize ([[maybe_unused]] amrex::ParticleReal bunch_charge, [[maybe_unused]] RefPart const & ref) + { + } + /** Return 1 6D particle coordinate * * @param x particle position in x diff --git a/src/particles/transformation/ToFixedT.H b/src/particles/transformation/ToFixedT.H index b6ce3fe5e..669a3264d 100644 --- a/src/particles/transformation/ToFixedT.H +++ b/src/particles/transformation/ToFixedT.H @@ -59,10 +59,13 @@ namespace impactx::transformation amrex::ParticleReal const y = p.pos(RealAoS::y); amrex::ParticleReal const t = p.pos(RealAoS::t); + // small tolerance to avoid NaN for pz<0: + constexpr amrex::ParticleReal tol = 1.0e-8_prt; + // compute value of reference pzd = beta*gamma amrex::ParticleReal const argd = -1.0_prt + pow(m_ptd, 2); - AMREX_ASSERT_WITH_MESSAGE(argd > 0.0_prt, "invalid pzd arg (<=0)"); - amrex::ParticleReal const pzdf = argd > 0.0_prt ? sqrt(argd) : 0.0_prt; + // AMREX_ASSERT_WITH_MESSAGE(argd > 0.0_prt, "invalid pzd arg (<=0)"); + amrex::ParticleReal const pzdf = argd > 0.0_prt ? sqrt(argd) : tol; // transform momenta to dynamic units (eg, so that momenta are // normalized by mc): @@ -72,8 +75,8 @@ namespace impactx::transformation // compute value of particle pz = beta*gamma amrex::ParticleReal const arg = -1.0_prt + pow(m_ptd+pt, 2) - pow(px, 2) - pow(py, 2); - AMREX_ASSERT_WITH_MESSAGE(arg > 0.0_prt, "invalid pz arg (<=0)"); - amrex::ParticleReal const pzf = arg > 0.0_prt ? sqrt(arg) : 0.0_prt; + // AMREX_ASSERT_WITH_MESSAGE(arg > 0.0_prt, "invalid pz arg (<=0)"); + amrex::ParticleReal const pzf = arg > 0.0_prt ? sqrt(arg) : tol; // transform position and momentum (from fixed s to fixed t) p.pos(RealAoS::x) = x + px*t/(m_ptd+pt); diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index 147dfa3cf..a5687d38a 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -92,6 +92,18 @@ void init_distribution(py::module& m) "A 6D Semi-Gaussian distribution (uniform in position, Gaussian in momentum)." ); + py::class_(md, "Thermal") + .def(py::init< + amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal, + amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal + >(), + py::arg("k"), py::arg("kT"), py::arg("kT_halo"), + py::arg("normalize"), py::arg("normalize_halo"), + py::arg("halo")=0.0, + "A stationary thermal or bithermal distribution\n\n" + "R. D. Ryne, J. Qiang, and A. Adelmann, in Proc. EPAC2004, pp. 1942-1944 (2004)" + ); + py::class_(md, "Triangle") .def(py::init< amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal,