Acceleration by RF Cavities

Beam accelerated through a sequence of 4 RF cavities (without space charge).

We use a 230 MeV electron beam with initial normalized rms emittance of 1 um.

The lattice and beam parameters are based on Example 2 of the IMPACT-Z examples folder:

https://github.com/impact-lbl/IMPACT-Z/tree/master/examples/Example2

The final target beam energy and beam moments are based on simulation in IMPACT-Z, without space charge.

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_rfcavity.py or

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

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

Python: Script

Listing 30 You can copy this file from examples/rfcavity/run_rfcavity.py.

#!/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 numpy as np

from impactx import ImpactX, distribution, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.space_charge = False
# sim.diagnostics = False  # benchmarking
sim.slice_step_diagnostics = False

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

# load a 230 MeV electron beam with an initial
# unnormalized rms emittance of 1 mm-mrad in all
# three phase planes
kin_energy_MeV = 230.0  # reference energy
bunch_charge_C = 1.0e-10  # used with space charge
npart = 10000  # number of macro particles (outside tests, use 1e5 or more)

#   reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

#   particle bunch
distr = distribution.Waterbag(
    lambdaX=0.352498964601e-3,
    lambdaY=0.207443478142e-3,
    lambdaT=0.70399950746e-4,
    lambdaPx=5.161852770e-6,
    lambdaPy=9.163582894e-6,
    lambdaPt=0.260528852031e-3,
    muxpx=0.5712386101751441,
    muypy=-0.514495755427526,
    mutpt=-5.05773e-10,
)
sim.add_particles(bunch_charge_C, distr, npart)

# design the accelerator lattice

# access RF cavity on-axis field data
data_in = np.loadtxt("onaxis_data.in")
z = data_in[:, 0]
ez_onaxis = data_in[:, 1]
ncoef = 25

#   Drift elements
dr1 = elements.Drift(name="dr1", ds=0.4, nslice=1)
dr2 = elements.Drift(name="dr2", ds=0.032997, nslice=1)
#   RF cavity element
rf = elements.RFCavity(
    name="rf",
    ds=1.31879807,
    escale=62.0,
    z=z,
    field_on_axis=ez_onaxis,
    ncoef=ncoef,
    freq=1.3e9,
    phase=85.5,
    mapsteps=100,
    nslice=4,
)

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

sim.lattice.extend(
    [
        monitor,
        dr1,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        monitor,
    ]
)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

Executable: Input File

Listing 31 You can copy this file from examples/rfcavity/input_rfcavity.in.

###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000  # outside tests, use 1e5 or more
beam.units = static
beam.kin_energy = 230
beam.charge = 1.0e-10
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 0.352498964601e-3
beam.lambdaY = 0.207443478142e-3
beam.lambdaT = 0.70399950746e-4
beam.lambdaPx = 5.161852770e-6
beam.lambdaPy = 9.163582894e-6
beam.lambdaPt = 0.260528852031e-3
beam.muxpx = 0.5712386101751441
beam.muypy = -0.514495755427526
beam.mutpt = -5.05773e-10


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor dr1 dr2 rf dr2 dr2 rf dr2 dr2 rf dr2 dr2 rf dr2 monitor

monitor.type = beam_monitor
monitor.backend = h5

dr1.type = drift
dr1.ds = 0.4
dr1.nslice = 1

dr2.type = drift
dr2.ds = 0.032997
dr1.nslice = 1

rf.type = rfcavity
rf.ds = 1.31879807
rf.escale = 62.0
rf.freq = 1.3e9
rf.phase = 85.5
rf.mapsteps = 100
rf.nslice = 4
rf.cos_coefficients =                    \
                0.1644024074311037       \
                -0.1324009958969339      \
                4.3443060026047219e-002  \
                8.5602654094946495e-002  \
                -0.2433578169042885      \
                0.5297150596779437       \
                0.7164884680963959       \
                -5.2579522442877296e-003 \
                -5.5025369142193678e-002 \
                4.6845673335028933e-002  \
                -2.3279346335638568e-002 \
                4.0800777539657775e-003  \
                4.1378326533752169e-003  \
                -2.5040533340490805e-003 \
                -4.0654981400000964e-003 \
                9.6630592067498289e-003  \
                -8.5275895985990214e-003 \
                -5.8078747006425020e-002 \
                -2.4044337836660403e-002 \
                1.0968240064697212e-002  \
                -3.4461179858301418e-003 \
                -8.1201564869443749e-004 \
                2.1438992904959380e-003  \
                -1.4997753525697276e-003 \
                1.8685171825676386e-004
rf.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.space_charge = false


###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false

Analyze

We run the following script to analyze correctness:

Script analysis_rfcavity.py

Listing 32 You can copy this file from examples/rfcavity/analysis_rfcavity.py.

#!/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 = 1.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],
    [
        4.29466150443e-4,
        2.41918588389e-4,
        7.0399951912e-5,
        2.21684103818e-9,
        2.21684103818e-9,
        1.83412186547e-8,
    ],
    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 = 1.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],
    [
        3.52596000000e-4,
        2.41775000000e-4,
        7.0417917357e-5,
        1.70893497973e-9,
        1.70893497973e-9,
        1.413901564889e-8,
    ],
    rtol=rtol,
    atol=atol,
)

Acceleration by RF Cavities (Reference Particle Tracking)

This test is identical to the test examples-rfcavity above, except that the code is run in reference orbit tracking mode.

It is used to validate the numerical integration of the reference energy gain in a chain of 4 RF cavities, and to demonstrate this tracking mode.

In this test, the initial and final reference values of \(s\) and \(\gamma\) must agree with nominal values.

Run

This example can be run either as:

• Python script: python3 run_rfcavity_ref_part.py or

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

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

Python: Script

Listing 33 You can copy this file from examples/rfcavity/run_rfcavity_ref_part.py.

#!/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 numpy as np

from impactx import ImpactX, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.space_charge = False
# sim.diagnostics = False  # benchmarking
sim.slice_step_diagnostics = True

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

# load a 230 MeV electron beam with an initial
# unnormalized rms emittance of 1 mm-mrad in all
# three phase planes
kin_energy_MeV = 230.0  # reference energy

#   reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

# design the accelerator lattice

# access RF cavity on-axis field data
data_in = np.loadtxt("onaxis_data.in")
z = data_in[:, 0]
ez_onaxis = data_in[:, 1]
ncoef = 25

#   Drift elements
dr1 = elements.Drift(name="dr1", ds=0.4, nslice=1)
dr2 = elements.Drift(name="dr2", ds=0.032997, nslice=1)
#   RF cavity element
rf = elements.RFCavity(
    name="rf",
    ds=1.31879807,
    escale=62.0,
    z=z,
    field_on_axis=ez_onaxis,
    ncoef=ncoef,
    freq=1.3e9,
    phase=85.5,
    mapsteps=100,
    nslice=4,
)

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

sim.lattice.extend(
    [
        monitor,
        dr1,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        monitor,
    ]
)

# run simulation
sim.track_reference(ref)

# clean shutdown
sim.finalize()

Executable: Input File

Listing 34 You can copy this file from examples/rfcavity/input_rfcavity_ref_part.in.

###############################################################################
# Particle Beam(s)
###############################################################################
beam.kin_energy = 230
beam.particle = electron


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor dr1 dr2 rf dr2 dr2 rf dr2 dr2 rf dr2 dr2 rf dr2 monitor

monitor.type = beam_monitor
monitor.backend = h5

dr1.type = drift
dr1.ds = 0.4
dr1.nslice = 1

dr2.type = drift
dr2.ds = 0.032997
dr1.nslice = 1

rf.type = rfcavity
rf.ds = 1.31879807
rf.escale = 62.0
rf.freq = 1.3e9
rf.phase = 85.5
rf.mapsteps = 100
rf.nslice = 4
rf.cos_coefficients =                    \
                0.1644024074311037       \
                -0.1324009958969339      \
                4.3443060026047219e-002  \
                8.5602654094946495e-002  \
                -0.2433578169042885      \
                0.5297150596779437       \
                0.7164884680963959       \
                -5.2579522442877296e-003 \
                -5.5025369142193678e-002 \
                4.6845673335028933e-002  \
                -2.3279346335638568e-002 \
                4.0800777539657775e-003  \
                4.1378326533752169e-003  \
                -2.5040533340490805e-003 \
                -4.0654981400000964e-003 \
                9.6630592067498289e-003  \
                -8.5275895985990214e-003 \
                -5.8078747006425020e-002 \
                -2.4044337836660403e-002 \
                1.0968240064697212e-002  \
                -3.4461179858301418e-003 \
                -8.1201564869443749e-004 \
                2.1438992904959380e-003  \
                -1.4997753525697276e-003 \
                1.8685171825676386e-004
rf.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.space_charge = false
algo.track = "reference_orbit"


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

Analyze

We run the following script to analyze correctness:

Script analysis_rfcavity_ref_part.py

Listing 35 You can copy this file from examples/rfcavity/analysis_rfcavity_ref_part.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 reference particle data
rbc = read_time_series("diags/ref_particle.*")

s = rbc["s"]
gamma = rbc["gamma"]

si = s.iloc[0]
gammai = gamma.iloc[0]

sf = s.iloc[-1]
gammaf = gamma.iloc[-1]

print("")
print("Initial Beam:")
print(f"  s_ref={si:e} gamma_ref={gammai:e}")

atol = 1.0e-4  # ignored
print(f"  atol={atol}")

assert np.allclose(
    [si, gammai],
    [
        0.000000,
        451.09877160930125,
    ],
    atol=atol,
)


print("")
print("Final Beam:")
print(f"  s_ref={sf:e} gamma_ref={gammaf:e}")

atol = 1.0e-4  # ignored
print(f"  atol={atol}")

assert np.allclose(
    [sf, gammaf],
    [
        5.9391682799999987,
        585.11989430128892,
    ],
    atol=atol,
)

Acceleration by RF Cavities (Using a Callback Hook)

This test is similar to the test above, except that it illustrates the use of the callback hook feature hook to provide an alternative method to set the RF cavity field amplitude(s) and phase(s).

The two functions get_synch_phase_Veff and get_phase_emax allow the user to convert between the pair of inputs (Veff, phase_synch) and (escale,phase). Here escale and phase are the documented ImpactX inputs in RFCavity, while Veff (\(V_{\rm eff}\)) and phase_synch (\(\phi_s\)) denote the effective voltage and synchronous phase of the cavity, defined here such that:

\[\Delta\gamma = V_{\rm eff}\cos(\phi_s).\]

The conversion is implemented under the approximation that the relativistic beta does not vary within the cavity. (E.g., this is an excellent approximation for large values of gamma.)

For the definition of effective voltage and synchronous phase, see:

• 20. Wangler, Principles of RF Linear Accelerators, John Wiley and Sons, Inc., New York, 1998, Sections 2.2 - 2.5

• 4. Uriot and N. Pichoff, TraceWin Manual, CEA/SACLAY - DRF/Irfu/DACM - Didiear URIOT, updated May 31, 2023, “Synchronous phase definitions”

The equivalent voltage here is multiplied by the particle charge \(q\) and divided by \(mc^2\), so this is a dimensionless quantity.

For each of the four RF cavities, the equivalent voltage and synchronous phase are computed. Those values are then converted back to the default phase and escale inputs, whose values are updated within the cavity.

When implemented correctly, the dynamics of the reference particle within these cavities should be identical to the dynamics seen in single reference particle tracking.

In this test, the initial and final reference values of \(s\) and \(\gamma\) must agree with nominal values.

Run

This example can be run as:

• Python script: python3 run_rfcavity_ref_part_hook.py

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

Python: Script

Listing 36 You can copy this file from examples/rfcavity/run_rfcavity_ref_part_hook.py.

#!/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 numpy as np
from scipy import constants as sconst

from impactx import ImpactX, elements, fourier_coefficients

sim = ImpactX()

# set numerical parameters and IO control
sim.space_charge = False
# sim.diagnostics = False  # benchmarking
sim.slice_step_diagnostics = True

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

# reference kinetic energy (initial)
kin_energy_MeV = 230.0  # reference energy

#   reference particle
ref = sim.beam.ref
ref.set_species("electron").set_kin_energy_MeV(kin_energy_MeV)

# design the accelerator lattice

# access RF cavity on-axis field data
data_in = np.loadtxt("onaxis_data.in")
z = data_in[:, 0]
ez_onaxis = data_in[:, 1]
ncoef = 25

cos_coeffs, sin_coeffs = fourier_coefficients(z, ez_onaxis, ncoef)


def get_effective_fourier_coeffs_AB(cos_coeffs, sin_coeffs, f, L, beta):
    w = 2.0 * np.pi * f
    k = w / sconst.c
    sinkL = np.sin(k * L / (2.0 * beta))

    j = np.arange(ncoef)
    signs = (-1.0) ** j
    denom = (k * L - 2 * j * np.pi * beta) * (k * L + 2 * j * np.pi * beta)
    Abasearr = signs / denom
    Bbasearr = signs * j / denom

    Acoeffs = -2.0 * Abasearr * k * L**2 * beta * cos_coeffs * sinkL
    Bcoeffs = 4.0 * Bbasearr * np.pi * L * beta**2 * sin_coeffs * sinkL
    Asum = np.sum(Acoeffs[1 : ncoef - 1]) + Acoeffs[0] / 2.0
    Bsum = np.sum(Bcoeffs[1 : ncoef - 1])

    A = Asum
    B = Bsum

    return A, B


def get_synch_phase_Veff(cos_coeffs, sin_coeffs, f, L, emax, beta, phase, t0):
    # predicted energy gain
    zmin = -L / (2.0 * beta)
    theta = phase * np.pi / 180.0
    w = 2.0 * np.pi * f
    k = w / sconst.c

    A, B = get_effective_fourier_coeffs_AB(cos_coeffs, sin_coeffs, f, L, beta)

    dpt = emax * (
        A * np.cos(k * (t0 - zmin / beta) + theta)
        + B * np.sin(k * (t0 - zmin / beta) + theta)
    )
    dgamma = -dpt

    Veff = emax * np.sqrt(A**2 + B**2)

    numerator = emax * (
        -B * np.cos(k * (t0 - zmin / beta) + theta)
        + A * np.sin(k * (t0 - zmin / beta) + theta)
    )

    synch_phase = np.arctan(numerator / dpt)

    return synch_phase, Veff, dgamma


def get_phase_emax(cos_coeffs, sin_coeffs, f, L, Veff, beta, synch_phase, t0):
    zmin = -L / (2.0 * beta)
    w = 2.0 * np.pi * f
    k = w / sconst.c

    A, B = get_effective_fourier_coeffs_AB(cos_coeffs, sin_coeffs, f, L, beta)

    term1 = -B - A * np.tan(synch_phase)
    term2 = -A + B * np.tan(synch_phase)
    phase = -k * (t0 - zmin / beta) + np.arctan2(term1, term2)
    phase_deg = np.degrees(np.mod(phase, 2 * np.pi))

    emax = Veff / np.sqrt(A**2 + B**2)

    return phase_deg, emax


#   Drift elements
dr1 = elements.Drift(name="dr1", ds=0.4, nslice=1)
dr2 = elements.Drift(name="dr2", ds=0.032997, nslice=1)

#   RF cavity element
rf = elements.RFCavity(
    name="rf",
    ds=1.31879807,
    escale=62.0,
    z=z,
    field_on_axis=ez_onaxis,
    ncoef=ncoef,
    freq=1.3e9,
    phase=85.5,
    mapsteps=100,
    nslice=4,
)


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

sim.lattice.extend(
    [
        monitor,
        dr1,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        monitor,
    ]
)


def hook_before_element(sim):
    element = sim.tracking_element
    if element.name == "rf":
        ref = sim.beam.ref
        beta = ref.beta
        f = element.freq
        L = element.ds
        emax = element.escale
        phase = element.phase
        t0 = ref.t
        print(
            f"  Beam at s={ref.s:.2f}m, t={ref.t:.2f}s, gamma at entry={ref.gamma:.2f}",
            flush=True,
        )
        # Extract the effective RF voltage and synchronous phase from the ImpactX RFCavity input parameters:
        synch_phase, Veff, dgamma = get_synch_phase_Veff(
            cos_coeffs, sin_coeffs, f, L, emax, beta, phase, t0
        )
        print(
            f"  RF cavity synchronous phase={synch_phase:.2f}, V effective={Veff:.2f}, predicted change in gamma={dgamma:.2f}",
            flush=True,
        )
        # From the effective RF voltage and synchronous phase, set the ImpactX RFCavity input parameters:
        updated_phase, updated_emax = get_phase_emax(
            cos_coeffs, sin_coeffs, f, L, Veff, beta, synch_phase, t0
        )
        element.phase = updated_phase
        element.escale = updated_emax
        print(
            f"  RF cavity updated (reset) values of phase={updated_phase:.2f} and emax={updated_emax:.2f}",
            flush=True,
        )


sim.hook["before_element"] = hook_before_element

# run simulation
sim.track_reference(ref)

# clean shutdown
sim.finalize()

Analyze

We run the following script to analyze correctness:

Script analysis_rfcavity_ref_part_hook.py

Listing 37 You can copy this file from examples/rfcavity/analysis_rfcavity_ref_part_hook.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 reference particle data
rbc = read_time_series("diags/ref_particle.*")

s = rbc["s"]
gamma = rbc["gamma"]

si = s.iloc[0]
gammai = gamma.iloc[0]

sf = s.iloc[-1]
gammaf = gamma.iloc[-1]

print("")
print("Initial Beam:")
print(f"  s_ref={si:e} gamma_ref={gammai:e}")

atol = 1.0e-4  # ignored
print(f"  atol={atol}")

assert np.allclose(
    [si, gammai],
    [
        0.000000,
        451.09877160930125,
    ],
    atol=atol,
)


print("")
print("Final Beam:")
print(f"  s_ref={sf:e} gamma_ref={gammaf:e}")

atol = 1.0e-4  # ignored
print(f"  atol={atol}")

assert np.allclose(
    [sf, gammaf],
    [
        5.9391682799999987,
        585.11989430128892,
    ],
    atol=atol,
)

Proton acceleration by RF Cavities (Using RF phase optimization)

This test is similar to the test above, except that the test is performed for protons (at the same energy). As a result, there is a large variation of relativistic beta (in the first RF cavity), and the analytical methods of the previous example cannot be used.

To treat this case, the callback hook feature hook is combined with an optimization loop (using scipy.minimize_scalar) to determine, for each RF cavity, the input phase setting that results in maximum energy gain. The phase is reset to this value (for each cavity).

The functions used for optimization of the RF phase are contained in the file phase_opt.py.

In this test, the initial and final reference values of \(s\) and \(\gamma\) must agree with nominal values.

Run

This example can be run as:

• Python script: python3 run_rfcavity_ref_part_opt.py

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

Python: Script

Listing 38 You can copy this file from examples/rfcavity/run_rfcavity_ref_part_opt.py.

#!/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 numpy as np
from phase_opt import optimize

from impactx import ImpactX, elements

sim = ImpactX()

# set numerical parameters and IO control
sim.space_charge = False
sim.slice_step_diagnostics = True

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

# reference kinetic energy (initial)
kin_energy_MeV = 230.0  # reference energy

#   reference particle
ref = sim.beam.ref
ref.set_species("proton").set_kin_energy_MeV(kin_energy_MeV)

# design the accelerator lattice

# access RF cavity on-axis field data
data_in = np.loadtxt("onaxis_data.in")
z = data_in[:, 0]
ez_onaxis = data_in[:, 1]
ncoef = 25

#   Drift elements
dr1 = elements.Drift(name="dr1", ds=0.4, nslice=1)
dr2 = elements.Drift(name="dr2", ds=0.032997, nslice=1)

#   RF cavity element
rf = elements.RFCavity(
    name="rf",
    ds=1.31879807,
    escale=20.0,
    z=z,
    field_on_axis=ez_onaxis,
    ncoef=ncoef,
    freq=1.3e9,
    phase=0.0,
    mapsteps=100,
    nslice=4,
)


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

sim.lattice.extend(
    [
        monitor,
        dr1,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        dr2,
        rf,
        dr2,
        monitor,
    ]
)


def hook_before_element(sim):
    element = sim.tracking_element
    if type(element) is elements.RFCavity:
        beam = sim.beam
        ref = beam.ref
        print(
            f"  Beam at s={ref.s:.2f}m, t={ref.t:.2f}s, gamma at entry={ref.gamma:.2f}",
            flush=True,
        )

        # Interpret input RF phase as measured relative to max accelerating phase:
        phase_shift = element.phase

        # Find RF phase that maximizes energy gain:
        phase_opt, e_gain = optimize(ref, element)

        # Reset input RF phase to appropriate value for tracking:
        element.phase = phase_opt + phase_shift
        print(
            f"  RF cavity updated (reset) values of phase={phase_opt:.2f}, energy gain (MeV) ={e_gain:.2f}",
            flush=True,
        )


sim.hook["before_element"] = hook_before_element

# run simulation
sim.track_reference(ref)

# clean shutdown
sim.finalize()

This script makes use of functions defined in the script phase_opt.py below, which can be re-used by the user:

Python: Script

Listing 39 You can copy this file from examples/rfcavity/phase_opt.py.

import numpy as np
from scipy.optimize import minimize_scalar

from impactx import push


def objective(parameter, ref, element):
    """
    A function that is evaluated by the optimizer.

    Parameters
    ----------
    parameter:
      rf cavity phase

    Returns
    -------
    Negative of the RF cavity energy gain in MeV (to minimize).
    """

    # adjust the RF cavity phase
    phase_opt = parameter
    element.phase = phase_opt

    # store the incoming energy and copy the reference particle
    KE_in = ref.kin_energy_MeV
    ref_copy = ref.copy()

    # push the copy of the reference particle
    push(ref_copy, element)
    KE_fin = ref_copy.kin_energy_MeV

    # evaluate the objective
    loss = KE_in - KE_fin

    if np.isnan(loss):
        loss = 1.0e99

    return loss


def optimize(ref, element):
    """
    Maximize the energy gain of a reference particle

    Using (ref) in an RFCavity (element): the optimization is performed by minimizing
    -(change in kinetic energy in MeV).

    Parameters
    ----------
    ref:
      the reference particle at RF entry
    element:
      the RF cavity element

    Returns
    -------
    The optimized phase and energy gain at that phase (phase_opt, e_gain).
    """

    # optimizer specific options
    options = {"maxiter": 2000, "disp": 1}

    # Call the optimizer
    res = minimize_scalar(
        objective,
        method="bounded",
        args=(ref, element),
        tol=1.0e-8,
        options=options,
        bounds=(-180, 180),
    )

    # Optimization result
    phase_opt = res.x
    e_gain = -1.0 * res.fun

    return phase_opt, e_gain

Analyze

We run the following script to analyze correctness:

Script analysis_rfcavity_ref_part_opt.py

Listing 40 You can copy this file from examples/rfcavity/analysis_rfcavity_ref_part_opt.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 reference particle data
rbc = read_time_series("diags/ref_particle.*")

s = rbc["s"]
gamma = rbc["gamma"]

si = s.iloc[0]
gammai = gamma.iloc[0]

sf = s.iloc[-1]
gammaf = gamma.iloc[-1]

print("")
print("Initial Beam:")
print(f"  s_ref={si:e} gamma_ref={gammai:e}")

atol = 1.0e-4  # ignored
print(f"  atol={atol}")

assert np.allclose(
    [si, gammai],
    [
        0.000000,
        1.2451314527015738,
    ],
    atol=atol,
)


print("")
print("Final Beam:")
print(f"  s_ref={sf:e} gamma_ref={gammaf:e}")

atol = 1.0e-4  # ignored
print(f"  atol={atol}")

assert np.allclose(
    [sf, gammaf],
    [
        5.9391682799999987,
        39.858859594214152,
    ],
    atol=atol,
)