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 \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), and \(\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.

Listing 160 You can copy this file from examples/fodo_space_charge/run_fodo_envelope_sc.py.
#!/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 = 0  # 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()
Listing 161 You can copy this file from examples/fodo_space_charge/input_fodo_envelope_sc.in.
###############################################################################
# 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.track = "envelope"
algo.space_charge = 2D

###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true

Analyze

We run the following script to analyze correctness:

Script analysis_fodo_envelope_sc.py
Listing 162 You can copy this file from examples/fodo_space_charge/analysis_fodo_envelope_sc.py.
#!/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"]
sigma_x = rbc["sigma_x"]
sigma_y = rbc["sigma_y"]
sigma_t = rbc["sigma_t"]
emittance_x = rbc["emittance_x"]
emittance_y = rbc["emittance_y"]
emittance_t = rbc["emittance_t"]

sigma_xi = sigma_x.iloc[0]
sigma_yi = sigma_y.iloc[0]
sigma_ti = sigma_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]
sigma_xf = sigma_x.iloc[length]
sigma_yf = sigma_y.iloc[length]
sigma_tf = sigma_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={sigma_xi:e} sigy={sigma_yi:e} sigt={sigma_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(
    [sigma_xi, sigma_yi, sigma_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={sigma_xf:e} sigy={sigma_yf:e} sigt={sigma_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(
    [sigma_xf, sigma_yf, sigma_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,
)

FODO Cell with 3D Gaussian Space Charge Using Particle Tracking

This example illustrates a 1 nC electron beam with a kinetic energy of 100 MeV in a FODO cell, with 3D Gaussian space charge included. The parameters are those described in:

The purpose of this example is to illustrate the use of particle tracking mode with 3D space charge from a Gaussian density distribution.

The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell with small noticeable growth.

In this test, the initial and final values of \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), and \(\epsilon_t\) must agree with nominal values.

Run

This example can be run either as:

  • Python script: python3 run_fodo_Gauss3D_sc.py or

  • ImpactX executable using an input file: impactx input_fodo_Gauss3D_sc.in

For MPI-parallel runs, prefix these lines with mpiexec -n 4 ... or srun -n 4 ..., depending on the system.

Listing 163 You can copy this file from examples/fodo_space_charge/run_fodo_Gauss3D_sc.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.particle_shape = 2  # B-spline order
sim.space_charge = "Gauss3D"
sim.slice_step_diagnostics = True

# domain decomposition & space charge mesh
sim.init_grids()

# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 2 nm
kin_energy_MeV = 100  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

#   reference particle
ref = sim.particle_container().ref_particle()
ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Gaussian(
    lambdaX=3.9984884770e-5,
    lambdaY=3.9984884770e-5,
    lambdaT=1.0e-3,
    lambdaPx=2.6623538760e-5,
    lambdaPy=2.6623538760e-5,
    lambdaPt=2.0e-3,
    muxpx=-0.846574929020762,
    muypy=0.846574929020762,
    mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")

# design the accelerator lattice)
ns = 25  # number of slices per ds in the element
fodo = [
    monitor,
    elements.ChrDrift(name="drift1", ds=0.25, nslice=ns),
    monitor,
    elements.ChrQuad(name="quad1", ds=1.0, k=1.0, nslice=ns),
    monitor,
    elements.ChrDrift(name="drift2", ds=0.5, nslice=ns),
    monitor,
    elements.ChrQuad(name="quad2", ds=1.0, k=-1.0, nslice=ns),
    monitor,
    elements.ChrDrift(name="drift3", ds=0.25, nslice=ns),
    monitor,
]
# assign a fodo segment
sim.lattice.extend(fodo)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()
Listing 164 You can copy this file from examples/fodo_space_charge/input_fodo_Gauss3D_sc.in.
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 1.0e2
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = gaussian
beam.lambdaX = 3.9984884770e-5
beam.lambdaY = beam.lambdaX
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.6623538760e-5
beam.lambdaPy = beam.lambdaPx
beam.lambdaPt = 2.0e-3
beam.muxpx = -0.846574929020762
beam.muypy = -beam.muxpx
beam.mutpt = 0.0


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift3 monitor
lattice.nslice = 25

monitor.type = beam_monitor
monitor.backend = h5

drift1.type = drift
drift1.ds = 0.25

quad1.type = quad
quad1.ds = 1.0
quad1.k = 1.0

drift2.type = drift
drift2.ds = 0.5

quad2.type = quad
quad2.ds = 1.0
quad2.k = -1.0

drift3.type = drift
drift3.ds = 0.25


###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = Gauss3D


###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true

Analyze

We run the following script to analyze correctness:

Script analysis_fodo_Gauss3D_sc.py
Listing 165 You can copy this file from examples/fodo_space_charge/analysis_fodo_Gauss3D_sc.py.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Ji Qiang
# License: BSD-3-Clause-LBNL
#

import numpy as np
import openpmd_api as io
from scipy.stats import moment


def get_moments(beam):
    """Calculate standard deviations of beam position & momenta
    and emittance values

    Returns
    -------
    sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
    """
    sigx = moment(beam["position_x"], moment=2) ** 0.5  # variance -> std dev.
    sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
    sigy = moment(beam["position_y"], moment=2) ** 0.5
    sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
    sigt = moment(beam["position_t"], moment=2) ** 0.5
    sigpt = moment(beam["momentum_t"], moment=2) ** 0.5

    epstrms = beam.cov(ddof=0)
    emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
    emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
    emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5

    return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)


# initial/final beam
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
last_step = list(series.iterations)[-1]
initial = series.iterations[1].particles["beam"].to_df()
beam_final = series.iterations[last_step].particles["beam"]
final = beam_final.to_df()

# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)

print("Initial Beam:")
sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti = get_moments(initial)
print(f"  sigx={sig_xi:e} sigy={sig_yi:e} sigt={sig_ti:e}")
print(
    f"  emittance_x={emittance_xi:e} emittance_y={emittance_yi:e} emittance_t={emittance_ti:e}"
)

atol = 0.0  # ignored
rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sig_xi, sig_yi, sig_ti, emittance_xi, emittance_yi, emittance_ti],
    [7.515765e-05, 7.511883e-05, 9.997395e-4, 2.001510e-09, 1.999755e-09, 1.999289e-06],
    rtol=rtol,
    atol=atol,
)


print("")
print("Final Beam:")
sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf = get_moments(initial)
print(f"  sigx={sig_xf:e} sigy={sig_yf:e} sigt={sig_tf:e}")
print(
    f"  emittance_x={emittance_xf:e} emittance_y={emittance_yf:e} emittance_t={emittance_tf:e}"
)

atol = 0.0  # ignored
rtol = 2.2 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sig_xf, sig_yf, sig_tf, emittance_xf, emittance_yf, emittance_tf],
    [
        7.51576586332169e-05,
        7.511883208451813e-05,
        0.0009997395499750136,
        2.0015106608723994e-09,
        1.999755254276969e-09,
        1.9992898444562777e-06,
    ],
    rtol=rtol,
    atol=atol,
)