diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index f038a2bf8..71e685f19 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -52,6 +52,7 @@ Space Charge examples/cfchannel/README.rst examples/kurth/README.rst examples/epac2004_benchmarks/README.rst + examples/fodo_space_charge/README.rst Coherent Synchrotron Radiation (CSR) """""""""""""""""""""""""""""""""""" diff --git a/docs/source/usage/how_to_run.rst b/docs/source/usage/how_to_run.rst index e5ab89196..5cda59106 100644 --- a/docs/source/usage/how_to_run.rst +++ b/docs/source/usage/how_to_run.rst @@ -13,13 +13,13 @@ ImpactX can be run using any of three distinct tracking modes. ImpactX's most p Additionally, ImpactX provides two simplified tracking modes to aid scientists through every step, from beamline inception to operation: tracking of the beam envelope (6x6 covariance matrix) through linearized transport maps, or only tracking of the reference particle orbit. -================== =============== =============== ================== +================== =============== =============== =================== Mode Use Case Generality Collective Effects -================== =============== =============== ================== -Particle Tracking Full Dynamics Most general Supported -Envelope Tracking Rapid Scans Linearized `Soon `__ +================== =============== =============== =================== +Particle Tracking Full Dynamics Most general Supported (3D only) +Envelope Tracking Rapid Scans Linearized Supported (2D only) Reference Tracking Early Design Reference orbit No -================== =============== =============== ================== +================== =============== =============== =================== .. _usage_run-user-interface: diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index f2399550b..3243bd52c 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -42,6 +42,9 @@ Initial Beam Distributions * ``beam.charge`` (``float``, in C) bunch charge +* ``beam.current`` (``float``, in A) + beam current, used only if ``algo.space_charge = "2D"`` + * ``beam.particle`` (``string``) particle type: currently either ``electron``, ``positron`` or ``proton`` @@ -714,12 +717,24 @@ Space Charge Space charge kicks are applied in between slices of thick :ref:`lattice elements `. See there ``nslice`` option on lattice elements for slicing. -* ``algo.space_charge`` (``boolean``, optional, default: ``false``) +* ``algo.space_charge`` (``string``, optional) + + The physical model of space charge used. + + ImpactX uses an AMReX grid of boxes to organize and parallelize space charge simulation domain. + These boxes also contain a field mesh, if space charge calculations are enabled. + + Options: + + * ``"false"`` (default): space charge effects are not calculated. + + * ``"2D"``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched. + + Currently, this model is supported only in envelope mode (when ``algo.track = "envelope"``). - Whether to calculate space charge effects. + * ``"3D"``: Space charge forces are computed in three dimensions, assuming the beam is bunched. -ImpactX uses an AMReX grid of boxes to organize and parallelize space charge simulation domain. -These boxes also contain a field mesh, if space charge calculations are enabled. + Currently, this model is supported only in particle mode (when ``algo.track = "particles"``). * ``amr.n_cell`` (3 integers) optional (default: 1 `blocking_factor `__ per MPI process) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 5f3932038..57477e48c 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -62,9 +62,19 @@ Collective Effects & Overall Simulation Parameters .. py:property:: space_charge - Enable (``True``) or disable (``False``) space charge calculations (default: ``False``). + The physical model of space charge used. + + Options: + + * ``False`` (default): space charge effects are not calculated. + + * ``"2D"``: Space charge forces are computed in the plane ``(x,y)`` transverse to the reference particle velocity, assuming the beam is long and unbunched. + + Currently, this model is supported only in envelope mode (when ``algo.track = "envelope"``). + + * ``"3D"``: Space charge forces are computed in three dimensions, assuming the beam is bunched. - Whether to calculate space charge effects. + Currently, this model is supported only in particle mode (when ``algo.track = "particles"``). .. py:property:: poisson_solver @@ -180,13 +190,14 @@ Collective Effects & Overall Simulation Parameters :param distr: distribution function to draw from (object from :py:mod:`impactx.distribution`) :param int npart: number of particles to draw - .. py:method:: init_envelope(ref, distr) + .. py:method:: init_envelope(ref, distr, intensity=None) Envelope tracking mode: Create a 6x6 covariance matrix from a distribution and then initialize the the simulation for envelope tracking relative to a reference particle. :param ref: the reference particle (object from :py:class:`impactx.RefPart`) :param distr: distribution function (object from :py:mod:`impactx.distribution`) + :param float intensity: the beam intensity, given as bunch charge (C) for 3D or beam current (A) for 2D space charge .. py:method:: particle_container() diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 370f3e6ce..0c711c446 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1269,3 +1269,19 @@ add_impactx_test(examples-scraping.py examples/scraping_beam/analysis_scraping.py OFF # no plot script yet ) + +# FODO cell with 2D space charge using envelope tracking ###################### +# +# with space charge +add_impactx_test(FODO.envelope.sc + examples/fodo_space_charge/input_fodo_envelope_sc.in + ON # ImpactX MPI-parallel + examples/fodo_space_charge/analysis_fodo_envelope_sc.py + OFF # no plot script yet +) +add_impactx_test(FODO.envelope.sc.py + examples/fodo_space_charge/run_fodo_envelope_sc.py + OFF # ImpactX MPI-parallel + examples/fodo_space_charge/analysis_fodo_envelope_sc.py + OFF # no plot script yet +) diff --git a/examples/cfchannel/input_cfchannel_10nC_fft.in b/examples/cfchannel/input_cfchannel_10nC_fft.in index 5b93893b9..1fedfc541 100644 --- a/examples/cfchannel/input_cfchannel_10nC_fft.in +++ b/examples/cfchannel/input_cfchannel_10nC_fft.in @@ -37,7 +37,7 @@ constf1.kt = 1.0 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D algo.poisson_solver = "fft" amr.n_cell = 48 48 40 diff --git a/examples/cfchannel/input_cfchannel_10nC_mlmg.in b/examples/cfchannel/input_cfchannel_10nC_mlmg.in index 36d39811c..a009a33d8 100644 --- a/examples/cfchannel/input_cfchannel_10nC_mlmg.in +++ b/examples/cfchannel/input_cfchannel_10nC_mlmg.in @@ -37,7 +37,7 @@ constf1.kt = 1.0 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D amr.n_cell = 48 48 40 #amr.n_cell = 72 72 64 # optional for increased precision diff --git a/examples/cfchannel/run_cfchannel_10nC_fft.py b/examples/cfchannel/run_cfchannel_10nC_fft.py index 6349438b4..19f352d8f 100755 --- a/examples/cfchannel/run_cfchannel_10nC_fft.py +++ b/examples/cfchannel/run_cfchannel_10nC_fft.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [48, 48, 40] # [72, 72, 64] for increased precision sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.poisson_solver = "fft" sim.prob_relative = [1.1] # sim.diagnostics = False # benchmarking diff --git a/examples/cfchannel/run_cfchannel_10nC_mlmg.py b/examples/cfchannel/run_cfchannel_10nC_mlmg.py index 45f057edc..dc2506f65 100755 --- a/examples/cfchannel/run_cfchannel_10nC_mlmg.py +++ b/examples/cfchannel/run_cfchannel_10nC_mlmg.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [48, 48, 40] # [72, 72, 64] for increased precision sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.prob_relative = [3.0] # sim.diagnostics = False # benchmarking sim.slice_step_diagnostics = True diff --git a/examples/epac2004_benchmarks/input_bithermal.in b/examples/epac2004_benchmarks/input_bithermal.in index 928b80467..2345b1415 100644 --- a/examples/epac2004_benchmarks/input_bithermal.in +++ b/examples/epac2004_benchmarks/input_bithermal.in @@ -36,7 +36,7 @@ constf1.nslice = 50 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D #amr.n_cell = 128 128 128 #full resolution amr.n_cell = 64 64 64 diff --git a/examples/epac2004_benchmarks/input_fodo_rf_SC.in b/examples/epac2004_benchmarks/input_fodo_rf_SC.in index 6e6cf533f..67229db67 100644 --- a/examples/epac2004_benchmarks/input_fodo_rf_SC.in +++ b/examples/epac2004_benchmarks/input_fodo_rf_SC.in @@ -117,7 +117,7 @@ gapb1.sin_coefficients = 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 +algo.space_charge = 3D amr.n_cell = 56 56 64 geometry.prob_relative = 4.0 diff --git a/examples/epac2004_benchmarks/input_thermal.in b/examples/epac2004_benchmarks/input_thermal.in index d6ba68b6d..10bc41518 100644 --- a/examples/epac2004_benchmarks/input_thermal.in +++ b/examples/epac2004_benchmarks/input_thermal.in @@ -33,7 +33,7 @@ constf1.nslice = 400 #full resolution # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D #amr.n_cell = 128 128 128 #full resolution amr.n_cell = 64 64 64 diff --git a/examples/epac2004_benchmarks/run_bithermal.py b/examples/epac2004_benchmarks/run_bithermal.py index 008a3a7bd..aba07705e 100755 --- a/examples/epac2004_benchmarks/run_bithermal.py +++ b/examples/epac2004_benchmarks/run_bithermal.py @@ -14,7 +14,7 @@ # 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.space_charge = "3D" sim.dynamic_size = True sim.prob_relative = [3.0] diff --git a/examples/epac2004_benchmarks/run_fodo_rf_SC.py b/examples/epac2004_benchmarks/run_fodo_rf_SC.py index 07c25eedb..8a6f2a84d 100755 --- a/examples/epac2004_benchmarks/run_fodo_rf_SC.py +++ b/examples/epac2004_benchmarks/run_fodo_rf_SC.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [56, 56, 64] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.dynamic_size = True sim.prob_relative = [4.0] diff --git a/examples/epac2004_benchmarks/run_thermal.py b/examples/epac2004_benchmarks/run_thermal.py index c20dc1b8d..f8fe36503 100755 --- a/examples/epac2004_benchmarks/run_thermal.py +++ b/examples/epac2004_benchmarks/run_thermal.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [56, 56, 64] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.dynamic_size = True sim.prob_relative = [4.0] diff --git a/examples/expanding_beam/input_expanding_fft.in b/examples/expanding_beam/input_expanding_fft.in index 9ba39c481..45676b5de 100644 --- a/examples/expanding_beam/input_expanding_fft.in +++ b/examples/expanding_beam/input_expanding_fft.in @@ -32,7 +32,7 @@ monitor.backend = h5 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D algo.poisson_solver = "fft" # Space charge solver with one MR level diff --git a/examples/expanding_beam/input_expanding_mlmg.in b/examples/expanding_beam/input_expanding_mlmg.in index a31c6ed97..cb2693e2b 100644 --- a/examples/expanding_beam/input_expanding_mlmg.in +++ b/examples/expanding_beam/input_expanding_mlmg.in @@ -32,7 +32,7 @@ monitor.backend = h5 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D # Space charge solver with one MR level amr.max_level = 1 diff --git a/examples/expanding_beam/run_expanding_fft.py b/examples/expanding_beam/run_expanding_fft.py index 92fa1269c..e34ef060f 100755 --- a/examples/expanding_beam/run_expanding_fft.py +++ b/examples/expanding_beam/run_expanding_fft.py @@ -18,7 +18,7 @@ sim.blocking_factor_z = [4] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.poisson_solver = "fft" sim.dynamic_size = True sim.prob_relative = [1.2, 1.1] diff --git a/examples/expanding_beam/run_expanding_mlmg.py b/examples/expanding_beam/run_expanding_mlmg.py index 6485d95eb..a53246e99 100755 --- a/examples/expanding_beam/run_expanding_mlmg.py +++ b/examples/expanding_beam/run_expanding_mlmg.py @@ -18,7 +18,7 @@ sim.blocking_factor_z = [4] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.dynamic_size = True sim.prob_relative = [3.0, 1.1] diff --git a/examples/fodo_space_charge/README.rst b/examples/fodo_space_charge/README.rst new file mode 100644 index 000000000..0486f4c79 --- /dev/null +++ b/examples/fodo_space_charge/README.rst @@ -0,0 +1,53 @@ +.. _examples-fodo-envelope-sc: + +FODO Cell with 2D Space Charge using Envelope Tracking +====================================================== + +This example illustrates a 0.5 A proton beam with a kinetic energy of 6.7 MeV in a FODO cell, +with 2D space charge included. The parameters are those described in: + +R.D. Ryne et al, `"A Test Suite of Space-Charge Problems for Code Benchmarking," `__ +in Proc. EPAC 2004, Lucerne, Switzerland: KV Beam in a FODO Channel + +The purpose of this example is to illustrate the use of envelope tracking mode with 2D space charge. + +The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling. + +In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_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_envelope_sc.py`` or +* ImpactX **executable** using an input file: ``impactx input_fodo_envelope_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_envelope_sc.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_space_charge/run_fodo_envelope_sc.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_fodo_envelope_sc.in + :language: ini + :caption: You can copy this file from ``examples/fodo_space_charge/input_fodo_envelope_sc.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_fodo_envelope_sc.py`` + + .. literalinclude:: analysis_fodo_envelope_sc.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_space_charge/analysis_fodo_envelope_sc.py``. diff --git a/examples/fodo_space_charge/analysis_fodo_envelope_sc.py b/examples/fodo_space_charge/analysis_fodo_envelope_sc.py new file mode 100755 index 000000000..da2a8fba8 --- /dev/null +++ b/examples/fodo_space_charge/analysis_fodo_envelope_sc.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + +import glob +import re + +import numpy as np +import pandas as pd + + +def read_file(file_pattern): + for filename in glob.glob(file_pattern): + df = pd.read_csv(filename, delimiter=r"\s+") + if "step" not in df.columns: + step = int(re.findall(r"[0-9]+", filename)[0]) + df["step"] = step + yield df + + +def read_time_series(file_pattern): + """Read in all CSV files from each MPI rank (and potentially OpenMP + thread). Concatenate into one Pandas dataframe. + + Returns + ------- + pandas.DataFrame + """ + return pd.concat( + read_file(file_pattern), + axis=0, + ignore_index=True, + ) # .set_index('id') + + +# read reduced diagnostics +rbc = read_time_series("diags/reduced_beam_characteristics.*") + +s = rbc["s"] +sig_x = rbc["sig_x"] +sig_y = rbc["sig_y"] +sig_t = rbc["sig_t"] +emittance_x = rbc["emittance_x"] +emittance_y = rbc["emittance_y"] +emittance_t = rbc["emittance_t"] + +sig_xi = sig_x.iloc[0] +sig_yi = sig_y.iloc[0] +sig_ti = sig_t.iloc[0] +emittance_xi = emittance_x.iloc[0] +emittance_yi = emittance_y.iloc[0] +emittance_ti = emittance_t.iloc[0] + +length = len(s) - 1 + +sf = s.iloc[length] +sig_xf = sig_x.iloc[length] +sig_yf = sig_y.iloc[length] +sig_tf = sig_t.iloc[length] +emittance_xf = emittance_x.iloc[length] +emittance_yf = emittance_y.iloc[length] +emittance_tf = emittance_t.iloc[length] + + +print("Initial Beam:") +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 = 1.0e-3 # 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], + [ + 8.590000e-04, + 8.590000e-04, + 7.071068e-07, + 1.000000e-06, + 1.000000e-06, + 1.000000e-12, + ], + rtol=rtol, + atol=atol, +) + + +print("") +print("Final Beam:") +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 = 1.0e-3 # 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], + [ + 8.590000e-04, + 8.590000e-04, + 4.140854e-05, + 1.000000e-06, + 1.000000e-06, + 1.000000e-12, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/fodo_space_charge/input_fodo_envelope_sc.in b/examples/fodo_space_charge/input_fodo_envelope_sc.in new file mode 100644 index 000000000..0ee52f44f --- /dev/null +++ b/examples/fodo_space_charge/input_fodo_envelope_sc.in @@ -0,0 +1,51 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.kin_energy = 6.7 +beam.current = 0.5 +beam.particle = proton +beam.distribution = kvdist_from_twiss +beam.alphaX = 2.4685083 +beam.alphaY = -beam.alphaX +beam.alphaT = 0.0 +beam.betaX = 0.737881 +beam.betaY = 0.737881 +beam.betaT = 0.5 +beam.emittX = 1.0e-6 +beam.emittY = beam.emittX +beam.emittT = 1.0e-12 + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor drift1 quad1 drift2 quad2 drift1 monitor +lattice.nslice = 50 + +monitor.type = beam_monitor +monitor.backend = h5 + +drift1.type = drift +drift1.ds = 7.44e-2 + +quad1.type = quad +quad1.ds = 6.10e-2 +quad1.k = -103.12574100336 + +drift2.type = drift +drift2.ds = 14.88e-2 + +quad2.type = quad +quad2.ds = 6.10e-2 +quad2.k = -quad1.k + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.track = "envelope" +algo.space_charge = 2D + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/fodo_space_charge/run_fodo_envelope_sc.py b/examples/fodo_space_charge/run_fodo_envelope_sc.py new file mode 100755 index 000000000..929a24e43 --- /dev/null +++ b/examples/fodo_space_charge/run_fodo_envelope_sc.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, distribution, elements, twiss + +sim = ImpactX() + +# set numerical parameters and IO control +sim.particle_shape = 2 # B-spline order +sim.space_charge = "2D" +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# model a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 6.7 # reference energy + +# 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) + +# beam current in A +beam_current_A = 0.5 + +# particle bunch +distr = distribution.KVdist( + **twiss( + beta_x=0.737881, + beta_y=0.737881, + beta_t=0.5, + emitt_x=1.0e-6, + emitt_y=1.0e-6, + emitt_t=1.0e-12, + alpha_x=2.4685083, + alpha_y=-2.4685083, + alpha_t=0.0, + ) +) + +sim.init_envelope(ref, distr, beam_current_A) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice) +ns = 50 # number of slices per ds in the element +fodo = [ + monitor, + elements.Drift(name="drift1", ds=7.44e-2, nslice=ns), + elements.Quad(name="quad1", ds=6.10e-2, k=-103.12574100336, nslice=ns), + elements.Drift(name="drift2", ds=14.88e-2, nslice=ns), + elements.Quad(name="quad2", ds=6.10e-2, k=103.12574100336, nslice=ns), + elements.Drift(name="drift3", ds=7.44e-2, nslice=ns), + monitor, +] +# assign a fodo segment +sim.lattice.extend(fodo) + +# run simulation +sim.track_envelope() + +# clean shutdown +sim.finalize() diff --git a/examples/kurth/input_kurth_10nC_periodic.in b/examples/kurth/input_kurth_10nC_periodic.in index 20360cea6..dcac80da4 100644 --- a/examples/kurth/input_kurth_10nC_periodic.in +++ b/examples/kurth/input_kurth_10nC_periodic.in @@ -40,7 +40,7 @@ constf1.kt = 0.7 # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D amr.n_cell = 48 48 40 #amr.n_cell = 72 72 72 # optional for increased precision diff --git a/examples/kurth/run_kurth_10nC_periodic.py b/examples/kurth/run_kurth_10nC_periodic.py index 8c94035e3..6c2a92dc5 100755 --- a/examples/kurth/run_kurth_10nC_periodic.py +++ b/examples/kurth/run_kurth_10nC_periodic.py @@ -13,7 +13,7 @@ # set numerical parameters and IO control sim.n_cell = [48, 48, 40] # use [72, 72, 72] for increased precision sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" # sim.diagnostics = False # benchmarking sim.slice_step_diagnostics = True diff --git a/examples/linac_segment/input_linac_segment.in b/examples/linac_segment/input_linac_segment.in index 3a858ce53..217e48d3a 100644 --- a/examples/linac_segment/input_linac_segment.in +++ b/examples/linac_segment/input_linac_segment.in @@ -359,7 +359,7 @@ RF1b.sin_coefficients = \ # Algorithms ############################################################################### algo.particle_shape = 2 -algo.space_charge = true +algo.space_charge = 3D # Space charge using IGF solver algo.poisson_solver = "fft" diff --git a/examples/linac_segment/run_linac_segment.py b/examples/linac_segment/run_linac_segment.py index 392280056..ab5d76675 100644 --- a/examples/linac_segment/run_linac_segment.py +++ b/examples/linac_segment/run_linac_segment.py @@ -17,7 +17,7 @@ # sim.n_cell = [64, 64, 64] #use this for high-resolution runs sim.n_cell = [32, 32, 32] sim.particle_shape = 2 # B-spline order -sim.space_charge = True +sim.space_charge = "3D" sim.poisson_solver = "fft" sim.dynamic_size = True sim.prob_relative = [1.1, 1.1] diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 3bccb50fe..b20bbe268 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -12,6 +12,7 @@ add_subdirectory(diagnostics) add_subdirectory(elements) add_subdirectory(initialization) add_subdirectory(particles) +add_subdirectory(envelope) add_subdirectory(tracking) diff --git a/src/elements/Aperture.H b/src/elements/Aperture.H index a305a0eee..a94bfdfc1 100644 --- a/src/elements/Aperture.H +++ b/src/elements/Aperture.H @@ -12,7 +12,6 @@ #include "particles/ImpactXParticleContainer.H" -#include "particles/CovarianceMatrix.H" #include "mixin/alignment.H" #include "mixin/beamoptic.H" #include "mixin/thin.H" diff --git a/src/envelope/CMakeLists.txt b/src/envelope/CMakeLists.txt new file mode 100644 index 000000000..cdffe61e8 --- /dev/null +++ b/src/envelope/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(spacecharge) diff --git a/src/envelope/spacecharge/CMakeLists.txt b/src/envelope/spacecharge/CMakeLists.txt new file mode 100644 index 000000000..dc9cf8000 --- /dev/null +++ b/src/envelope/spacecharge/CMakeLists.txt @@ -0,0 +1,4 @@ +target_sources(lib + PRIVATE + EnvelopeSpaceChargePush.cpp +) diff --git a/src/envelope/spacecharge/EnvelopeSpaceChargePush.H b/src/envelope/spacecharge/EnvelopeSpaceChargePush.H new file mode 100644 index 000000000..34ba1883b --- /dev/null +++ b/src/envelope/spacecharge/EnvelopeSpaceChargePush.H @@ -0,0 +1,42 @@ +/* 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: Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ENVELOPESPACECHARGEPUSH_H +#define IMPACTX_ENVELOPESPACECHARGEPUSH_H + +#include "particles/ImpactXParticleContainer.H" +#include "particles/CovarianceMatrix.H" + +#include // for AMREX_RESTRICT +#include + + +namespace impactx::envelope::spacecharge +{ + /** This function pushes the 6x6 beam covariance matrix for a slice + * of length ds, using the linear space charge fields in an rms + * equivalent 2D ellipse, as determined from the beam covariance matrix. + * Note: This is a reduced model of 2D space charge. + * + * @param[in] refpart reference particle + * @param[inout] cm covariance matrix + * @param[in] current beam current [A] + * @param[in] ds step size [m] + */ + void + space_charge2D_push ( + RefPart const & AMREX_RESTRICT refpart, + Map6x6 & AMREX_RESTRICT cm, + amrex::ParticleReal current, + amrex::ParticleReal ds + ); + +} // namespace impactx + +#endif // IMPACTX_ENVELOPESPACECHARGEPUSH_H diff --git a/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp new file mode 100644 index 000000000..4796ade10 --- /dev/null +++ b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp @@ -0,0 +1,71 @@ +/* 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: Marco Garten, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "EnvelopeSpaceChargePush.H" + +#include + +#include // for Real +#include +#include + +#include + + +namespace impactx::envelope::spacecharge +{ + void + space_charge2D_push ( + RefPart const & AMREX_RESTRICT refpart, + Map6x6 & AMREX_RESTRICT cm, + amrex::ParticleReal current, + amrex::ParticleReal ds + ) + { + using namespace amrex::literals; + + // skip calculations for trivial case + if (current == 0_prt) { return; } + + // initialize the linear transport map + Map6x6 R = Map6x6::Identity(); + + // physical constants and reference quantities + using ablastr::constant::SI::c; + using ablastr::constant::SI::ep0; + using ablastr::constant::math::pi; + amrex::ParticleReal const mass = refpart.mass; + amrex::ParticleReal const charge = refpart.charge; + amrex::ParticleReal const pt_ref = refpart.pt; + amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1_prt; + amrex::ParticleReal const betgam = std::sqrt(betgam2); + amrex::ParticleReal const betgam3 = std::pow(betgam,3); + + // evaluate the beam space charge perveance from current + amrex::ParticleReal const IA = 4_prt * pi * ep0 * mass * std::pow(c,3) / charge; + amrex::ParticleReal const Kpv = std::abs(current / IA) * 2_prt / betgam3; + + // evaluate the linear transfer map + amrex::ParticleReal const sigma2 = cm(1,1) * cm(3,3) - cm(1,3) * cm(1,3); + amrex::ParticleReal const sigma = std::sqrt(sigma2); + amrex::ParticleReal const D = (sigma + cm(1,1)) * (sigma + cm(3,3)) - cm(1,3) * cm(1,3); + amrex::ParticleReal const coeff = Kpv * ds / (2_prt * D); + + R(2,1) = coeff * (sigma + cm(3,3)); + R(2,3) = coeff * (-cm(1,3)); + R(4,1) = R(2,3); + R(4,3) = coeff * (sigma + cm(1,1)); + + // update the beam covariance matrix + cm = R * cm * R.transpose(); + + } + + +} // namespace impactx::spacecharge diff --git a/src/initialization/Algorithms.H b/src/initialization/Algorithms.H new file mode 100644 index 000000000..6a28f0903 --- /dev/null +++ b/src/initialization/Algorithms.H @@ -0,0 +1,40 @@ +/* Copyright 2022-2025 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Axel Huebl, Chad Mitchell + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ALGORITHMS_H +#define IMPACTX_ALGORITHMS_H + +#include + +#include + + +namespace impactx +{ + /** Space Charge: Implemented algorithms */ + AMREX_ENUM(SpaceChargeAlgo, + False, /**< Disabled: no space charge calculation */ + True_3D, /**< 3D beam distribution */ + True_2D /**< averaged 2D transverse beam distribution with a current along s */ + ); + + /** Return the currently active space charge algorithm */ + SpaceChargeAlgo + get_space_charge_algo (); + + /** A user-friendly string + * + * Returns a string that is similar to the spelling in user input options. + */ + std::string + to_string (SpaceChargeAlgo sca); + +} // namespace impactx + +#endif // IMPACTX_ALGORITHMS_H diff --git a/src/initialization/Algorithms.cpp b/src/initialization/Algorithms.cpp new file mode 100644 index 000000000..007f5bc40 --- /dev/null +++ b/src/initialization/Algorithms.cpp @@ -0,0 +1,91 @@ +/* Copyright 2022-2025 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Axel Huebl, Chad Mitchell + * License: BSD-3-Clause-LBNL + */ +#include "initialization/Algorithms.H" + +#include + +#include + +#include // for std::transform +#include +#include + + +namespace impactx +{ + SpaceChargeAlgo + get_space_charge_algo () + { + amrex::ParmParse const pp_algo("algo"); + std::string space_charge; + bool has_space_charge = pp_algo.query("space_charge", space_charge); + + if (!has_space_charge) { return SpaceChargeAlgo::False; } + + // TODO: at some point, remove backwards compatibility to pre 25.03 + std::string space_charge_lower = space_charge; + std::transform(space_charge.begin(), space_charge.end(), space_charge_lower.begin(), + [](unsigned char c){ return std::tolower(c); }); + if (space_charge == "1" || space_charge_lower == "true" || space_charge_lower == "on") + { + ablastr::warn_manager::WMRecordWarning( + "algo.space_charge", + "The option algo.space_charge = true is deprecated and will be removed in a future version of ImpactX. " + "Please use algo.space_charge = 3D instead.", + ablastr::warn_manager::WarnPriority::high + ); + return SpaceChargeAlgo::True_3D; + } + if (space_charge == "0") + { + ablastr::warn_manager::WMRecordWarning( + "algo.space_charge", + "The option algo.space_charge = 0 is deprecated and will be removed in a future version of ImpactX. " + "Please use algo.space_charge = false instead.", + ablastr::warn_manager::WarnPriority::high + ); + return SpaceChargeAlgo::False; + } + + if (space_charge_lower == "false" || space_charge_lower == "off") + { + return SpaceChargeAlgo::False; + } + else if (space_charge == "3D") + { + return SpaceChargeAlgo::True_3D; + } + else if (space_charge == "2D") + { + return SpaceChargeAlgo::True_2D; + } + else + { + throw std::runtime_error("algo.space_charge = " + space_charge + " is not a valid option"); + } + } + + std::string + to_string (SpaceChargeAlgo sca) + { + std::string const str = amrex::getEnumNameString(sca); + + // strip True_ prefix, which we only added because var names / enum member names + // cannot start with a number + std::string const true_prefix = "True_"; + if (str.find(true_prefix) == 0) + { + return str.substr(true_prefix.length()); + } + + return str; + } + +} // namespace impactx diff --git a/src/initialization/AmrCoreData.H b/src/initialization/AmrCoreData.H index b4d4a312f..1a457a6c4 100644 --- a/src/initialization/AmrCoreData.H +++ b/src/initialization/AmrCoreData.H @@ -89,7 +89,7 @@ namespace impactx::initialization std::optional m_ref; /** The 6x6 covariance matrix for the beam envelope, relative to the reference particle */ - std::optional m_cm; + std::optional m_env; } track_envelope; diff --git a/src/initialization/CMakeLists.txt b/src/initialization/CMakeLists.txt index b0a828b52..745af67b3 100644 --- a/src/initialization/CMakeLists.txt +++ b/src/initialization/CMakeLists.txt @@ -1,5 +1,6 @@ target_sources(lib PRIVATE + Algorithms.cpp AmrCoreData.cpp InitAMReX.cpp InitAmrCore.cpp diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H index 90e9bb108..f1f615b42 100644 --- a/src/initialization/InitDistribution.H +++ b/src/initialization/InitDistribution.H @@ -17,6 +17,7 @@ #include // for AMREX_RESTRICT #include +#include // for std::optional #include // for std::move @@ -39,10 +40,14 @@ namespace impactx::initialization read_distribution (amrex::ParmParse const & pp_dist); /** Ignore the shape of a distribution and use the 2nd moments to create a covariance matrix + * + * @param distr the distribution to draw from + * @param intensity the beam charge (in C, 3D space charge) or current (in A, 2D space charge) */ - CovarianceMatrix - create_covariance_matrix ( - distribution::KnownDistributions const & distr + Envelope + create_envelope ( + distribution::KnownDistributions const & distr, + std::optional intensity = std::nullopt ); /** Initialize a single particle's data using the given distribution diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 63d34a2cb..9142f0e68 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -10,6 +10,7 @@ #include "initialization/InitDistribution.H" #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "particles/CovarianceMatrix.H" #include "particles/ImpactXParticleContainer.H" #include "particles/distribution/All.H" @@ -185,9 +186,10 @@ namespace impactx return dist; } - CovarianceMatrix - initialization::create_covariance_matrix ( - distribution::KnownDistributions const & distr + Envelope + initialization::create_envelope ( + distribution::KnownDistributions const & distr, + std::optional intensity ) { // zero out the 6x6 matrix @@ -233,7 +235,11 @@ namespace impactx } }, distr); - return cv; + Envelope env; + if (intensity) { env.set_beam_intensity(intensity.value()); } + env.set_covariance_matrix(cv); + + return env; } void @@ -313,13 +319,11 @@ namespace impactx ref.qm_ratio_SI(), bunch_charge * rel_part_this_proc); - bool space_charge = false; - amrex::ParmParse pp_algo("algo"); - pp_algo.queryAdd("space_charge", space_charge); + auto space_charge = get_space_charge_algo(); // For pure tracking simulations, we keep the particles split equally // on all MPI ranks, and ignore spatial "RealBox" extents of grids. - if (space_charge) { + if (space_charge != SpaceChargeAlgo::False) { // Resize the mesh to fit the spatial extent of the beam and then // redistribute particles, so they reside on the MPI rank that is // responsible for the respective spatial particle position. @@ -473,7 +477,24 @@ namespace impactx { amr_data->track_envelope.m_ref = initialization::read_reference_particle(pp_dist); auto dist = initialization::read_distribution(pp_dist); - amr_data->track_envelope.m_cm = impactx::initialization::create_covariance_matrix(dist); + + + amrex::ParticleReal intensity = 0.0; // bunch charge (C) for 3D model, beam current (A) for 2D model + + auto space_charge = get_space_charge_algo(); + if (space_charge == SpaceChargeAlgo::True_3D) + { + //pp_dist.get("charge", intensity); + //amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist, intensity); + throw std::runtime_error("3D space charge model not yet implemented in envelope mode."); + } else if (space_charge == SpaceChargeAlgo::True_2D) + { + pp_dist.get("current", intensity); + amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist, intensity); + } else + { + amr_data->track_envelope.m_env = impactx::initialization::create_envelope(dist); + } } else if (track == "reference_orbit") { diff --git a/src/initialization/InitMeshRefinement.H b/src/initialization/InitMeshRefinement.H index a8b1be634..3b680cc12 100644 --- a/src/initialization/InitMeshRefinement.H +++ b/src/initialization/InitMeshRefinement.H @@ -9,6 +9,8 @@ */ #pragma once +#include "initialization/Algorithms.H" + #include #include @@ -32,8 +34,7 @@ namespace impactx::initialization amrex::ParmParse pp_amr("amr"); amrex::ParmParse pp_geometry("geometry"); - bool space_charge = false; - pp_algo.queryAdd("space_charge", space_charge); + auto space_charge = get_space_charge_algo(); std::string poisson_solver = "multigrid"; pp_algo.queryAdd("poisson_solver", poisson_solver); @@ -41,7 +42,7 @@ namespace impactx::initialization int max_level = 0; pp_amr.queryWithParser("max_level", max_level); - if (max_level > 1 && !space_charge) + if (max_level > 1 && space_charge != SpaceChargeAlgo::False) throw std::runtime_error( "Mesh-refinement (amr.max_level>=0) is only supported with " "space charge modeling (algo.space_charge=1)."); @@ -51,7 +52,7 @@ namespace impactx::initialization prob_relative[0] = 3.0; // top/bottom pad the beam on the lowest level by default by its width pp_geometry.queryarr("prob_relative", prob_relative); - if (prob_relative[0] < 3.0 && space_charge && poisson_solver == "multigrid") + if (prob_relative[0] < 3.0 && space_charge != SpaceChargeAlgo::False && poisson_solver == "multigrid") ablastr::warn_manager::WMRecordWarning( "ImpactX::read_mr_prob_relative", "Dynamic resizing of the mesh uses a geometry.prob_relative " diff --git a/src/initialization/InitMeshRefinement.cpp b/src/initialization/InitMeshRefinement.cpp index 7fd142710..cfddce123 100644 --- a/src/initialization/InitMeshRefinement.cpp +++ b/src/initialization/InitMeshRefinement.cpp @@ -8,6 +8,7 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "initialization/InitAmrCore.H" #include "particles/ImpactXParticleContainer.H" #include "particles/distribution/Waterbag.H" @@ -82,10 +83,8 @@ namespace detail BL_PROFILE("ImpactX::ResizeMesh"); { - amrex::ParmParse pp_algo("algo"); - bool space_charge = false; - pp_algo.query("space_charge", space_charge); - if (!space_charge) + auto space_charge = get_space_charge_algo(); + if (space_charge == SpaceChargeAlgo::False) ablastr::warn_manager::WMRecordWarning( "ImpactX::ResizeMesh", "This is a simulation without space charge. " diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H index b55ed66df..a67a489ad 100644 --- a/src/particles/CovarianceMatrix.H +++ b/src/particles/CovarianceMatrix.H @@ -19,9 +19,63 @@ namespace impactx /** this is a 6x6 matrix */ using Map6x6 = amrex::SmallMatrix; - /** the covariance matrix is 6x6 */ + /** this is the 6x6 covariance matrix */ using CovarianceMatrix = Map6x6; + + /** This struct stores the beam envelope attributes, including the 6x6 beam + * covariance matrix. Used during envelope tracking mode. + */ + struct Envelope + { + CovarianceMatrix m_env; ///< the 6x6 beam covariance matrix + amrex::ParticleReal m_beam_intensity = 0.0; ///< optional: charge in A (for 3D space charge) or current in A (for 2D space charge) + + /** Set envelope beam charge/current for 3D/2D space charge + * + * @param intensity beam charge (C) in 3D or beam current (A) in 2D + */ + Envelope & + set_beam_intensity (amrex::ParticleReal const intensity) + { + m_beam_intensity = intensity; + + return *this; + } + + /** Get envelope beam charge/current for 3D/2D space charge + * + * @returns 3D: beam charge in C; 2D: beam current in A + */ + amrex::ParticleReal + beam_intensity () const + { + return m_beam_intensity; + } + + /** Set 6x6 covariance matrix for envelope tracking + * + * @param covariance_matrix beam 6x6 covariance matrix + */ + Envelope & + set_covariance_matrix (CovarianceMatrix const & covariance_matrix) + { + m_env = covariance_matrix; + + return *this; + } + + /** Get the 6x6 covariance matrix for envelope tracking + */ + CovarianceMatrix + covariance_matrix () const + { + return m_env; + } + + }; + + } // namespace impactx::distribution #endif // IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H diff --git a/src/particles/spacecharge/HandleSpacecharge.cpp b/src/particles/spacecharge/HandleSpacecharge.cpp index 4cc13a73a..9612b236f 100644 --- a/src/particles/spacecharge/HandleSpacecharge.cpp +++ b/src/particles/spacecharge/HandleSpacecharge.cpp @@ -9,6 +9,7 @@ */ #include "HandleSpacecharge.H" +#include "initialization/Algorithms.H" #include "initialization/AmrCoreData.H" #include "particles/ImpactXParticleContainer.H" #include "particles/spacecharge/ForceFromSelfFields.H" @@ -32,12 +33,10 @@ namespace impactx::particles::spacecharge { BL_PROFILE("impactx::particles::wakefields::HandleSpacecharge") - amrex::ParmParse const pp_algo("algo"); - bool space_charge = false; - pp_algo.query("space_charge", space_charge); + auto space_charge = get_space_charge_algo(); // turn off if disabled by user - if (!space_charge) { return; } + if (space_charge == SpaceChargeAlgo::False) { return; } // turn off if less than 2 particles if (amr_data->track_particles.m_particle_container->TotalNumberOfParticles(true, false) < 2) { return; } diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 01e39021f..8ba1fd0fe 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -16,7 +16,10 @@ #if defined(AMREX_DEBUG) || defined(DEBUG) # include #endif +#include #include +#include +#include namespace py = pybind11; @@ -227,14 +230,32 @@ void init_ImpactX (py::module& m) "Enable or disable eigenemittance diagnostic calculations (default: disabled)." ) .def_property("space_charge", - [](ImpactX & /* ix */) { - return detail::get_or_throw("algo", "space_charge"); - }, - [](ImpactX & /* ix */, bool const enable) { - amrex::ParmParse pp_algo("algo"); - pp_algo.add("space_charge", enable); - }, - "Enable or disable space charge calculations (default: enabled)." + [](ImpactX & /* ix */) -> std::string { + return detail::get_or_throw("algo", "space_charge"); + }, + [](ImpactX & /* ix */, std::variant space_charge_v) { + if (std::holds_alternative(space_charge_v)) { + amrex::ParmParse pp_algo("algo"); + if (std::get(space_charge_v)) { + // TODO: boolean True is deprecated since 25.03, remove some time after + py::print("sim.space_charge = True is deprecated, please use space_charge = \"3D\""); + pp_algo.add("space_charge", std::string("3D")); + } else { + // map boolean False to "false" / off + pp_algo.add("space_charge", std::string("false")); + } + } + else + { + std::string const space_charge = std::get(space_charge_v); + if (space_charge != "false" && space_charge != "off" && space_charge != "2D" && space_charge != "3D") { + throw std::runtime_error("Space charge model must be 2D or 3D but is: " + space_charge); + } + amrex::ParmParse pp_algo("algo"); + pp_algo.add("space_charge", space_charge); + } + }, + "The model to be used when calculating space charge effects. Either off, 2D, or 3D." ) .def_property("poisson_solver", [](ImpactX & /* ix */) { @@ -438,11 +459,11 @@ void init_ImpactX (py::module& m) .def("init_beam_distribution_from_inputs", &ImpactX::initBeamDistributionFromInputs) .def("init_lattice_elements_from_inputs", &ImpactX::initLatticeElementsFromInputs) .def("init_envelope", - [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr) { + [](ImpactX & ix, RefPart ref, distribution::KnownDistributions distr, std::optional intensity) { ix.amr_data->track_envelope.m_ref = ref; - ix.amr_data->track_envelope.m_cm = initialization::create_covariance_matrix(distr); + ix.amr_data->track_envelope.m_env = initialization::create_envelope(distr, intensity); }, - py::arg("ref"), py::arg("distr"), + py::arg("ref"), py::arg("distr"), py::arg("intensity") = py::none(), "Envelope tracking mode:" "Create a 6x6 covariance matrix from a distribution and then initialize " "the the simulation for envelope tracking relative to a reference particle." diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index 2659a5f76..80d754b2c 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -131,5 +131,12 @@ void init_distribution(py::module& m) "A 6D Waterbag distribution" ); - m.def("create_covariance_matrix", &initialization::create_covariance_matrix); + py::class_(m, "Envelope") + .def(py::init<>()) + .def(py::init()) + .def_property("envelope", &Envelope::covariance_matrix, &Envelope::set_covariance_matrix) + .def_property("beam_intensity", &Envelope::beam_intensity, &Envelope::set_beam_intensity) + ; + + m.def("create_envelope", &initialization::create_envelope); } diff --git a/src/tracking/envelope.cpp b/src/tracking/envelope.cpp index 603c011e3..5bdb66a54 100644 --- a/src/tracking/envelope.cpp +++ b/src/tracking/envelope.cpp @@ -8,18 +8,24 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "initialization/InitAmrCore.H" #include "particles/ImpactXParticleContainer.H" #include "particles/Push.H" +#include "envelope/spacecharge/EnvelopeSpaceChargePush.H" #include "diagnostics/DiagnosticOutput.H" +#include + #include #include #include #include #include +#include #include +#include namespace impactx @@ -29,6 +35,8 @@ namespace impactx { BL_PROFILE("ImpactX::track_envelope"); + using namespace amrex::literals; + // verbosity amrex::ParmParse pp_impactx("impactx"); int verbose = 1; @@ -46,12 +54,14 @@ namespace impactx { throw std::runtime_error("track_envelope: Reference particle not set."); } - if (!amr_data->track_envelope.m_cm.has_value()) + if (!amr_data->track_envelope.m_env.has_value()) { throw std::runtime_error("track_envelope: Envelope (covariance matrix) not set."); } auto & ref = amr_data->track_envelope.m_ref.value(); - auto & cm = amr_data->track_envelope.m_cm.value(); + auto & env = amr_data->track_envelope.m_env.value(); + auto & cm = env.m_env; + auto & current = env.m_beam_intensity; // output of init state amrex::ParmParse pp_diag("diag"); @@ -75,21 +85,31 @@ namespace impactx } amrex::ParmParse const pp_algo("algo"); - bool space_charge = false; - pp_algo.query("space_charge", space_charge); - AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!space_charge, "Space charge not yet implemented for envelope tracking."); + auto space_charge = get_space_charge_algo(); if (verbose > 0) { - amrex::Print() << " Space Charge effects: " << space_charge << "\n"; + amrex::Print() << " Space Charge effects: " << to_string(space_charge) << "\n"; + } + if (space_charge == SpaceChargeAlgo::True_3D) + { + throw std::runtime_error("3D space charge effects are not yet implemented for envelope tracking."); + } + if (space_charge != SpaceChargeAlgo::False && current == 0_prt) { + ablastr::warn_manager::WMRecordWarning( + "algo.space_charge", + "Space charge calculations are enabled but zero beam current was provided. " + "Skipping space charge calculations.", + ablastr::warn_manager::WarnPriority::high + ); } bool csr = false; pp_algo.query("csr", csr); - AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!csr, "CSR not yet implemented for envelope tracking."); if (verbose > 0) { amrex::Print() << " CSR effects: " << csr << "\n"; } + AMREX_ALWAYS_ASSERT_WITH_MESSAGE(!csr, "CSR effects are not yet implemented for envelope tracking."); // periods through the lattice int num_periods = 1; @@ -123,6 +143,12 @@ namespace impactx << " slice_step=" << slice_step << "\n"; } + if (space_charge != SpaceChargeAlgo::False) + { + // push Covariance Matrix in 2D space charge fields + envelope::spacecharge::space_charge2D_push(ref, cm, current, slice_ds); + } + std::visit([&ref, &cm](auto&& element) { // push reference particle in global coordinates @@ -131,7 +157,7 @@ namespace impactx element(ref); } - // push Covariance Matrix + // push Covariance Matrix in external fields element(cm, ref); }, element_variant); diff --git a/src/tracking/particles.cpp b/src/tracking/particles.cpp index 655944caf..8247cd0ac 100644 --- a/src/tracking/particles.cpp +++ b/src/tracking/particles.cpp @@ -8,6 +8,7 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "initialization/InitAmrCore.H" #include "particles/CollectLost.H" #include "particles/ImpactXParticleContainer.H" @@ -68,13 +69,16 @@ namespace impactx } - amrex::ParmParse const pp_algo("algo"); - bool space_charge = false; - pp_algo.query("space_charge", space_charge); + auto space_charge = get_space_charge_algo(); if (verbose > 0) { - amrex::Print() << " Space Charge effects: " << space_charge << "\n"; + amrex::Print() << " Space Charge effects: " << to_string(space_charge) << "\n"; + } + if (space_charge == SpaceChargeAlgo::True_2D) + { + throw std::runtime_error("2D space charge effects are not yet implemented for particle tracking."); } + amrex::ParmParse const pp_algo("algo"); bool csr = false; pp_algo.query("csr", csr); if (verbose > 0) { diff --git a/src/tracking/reference.cpp b/src/tracking/reference.cpp index 372a12917..5e42651c1 100644 --- a/src/tracking/reference.cpp +++ b/src/tracking/reference.cpp @@ -8,6 +8,7 @@ * License: BSD-3-Clause-LBNL */ #include "ImpactX.H" +#include "initialization/Algorithms.H" #include "initialization/InitAmrCore.H" #include "particles/ImpactXParticleContainer.H" #include "particles/Push.H" @@ -20,6 +21,7 @@ #include #include +#include namespace impactx @@ -60,6 +62,20 @@ namespace impactx } + auto space_charge = get_space_charge_algo(); + if (space_charge != SpaceChargeAlgo::False) + { + throw std::runtime_error("Space charge effects cannot be modeled for single particle tracking."); + } + + amrex::ParmParse const pp_algo("algo"); + bool csr = false; + pp_algo.query("csr", csr); + if (!csr) + { + throw std::runtime_error("CSR effects cannot be modeled for single particle tracking."); + } + // periods through the lattice int num_periods = 1; amrex::ParmParse("lattice").queryAddWithParser("periods", num_periods);