Generation of Fourier coefficients from on-axis data

This example illustrates the computation of Fourier coefficients that are used to represent on-axis field or gradient data for soft-edge elements.

Given data in the file onaxis_data.in, execution of the Python script python3 fcoef.py results in the following output:

fcoef.out: a file containing a list of cosine and sine Fourier coefficients that can be used in ImpactX to define a soft-edge element
onaxis_data.out: a file containing the reconstructed on-axis signal, together with its first and second derivatives

The signal is represented in the form:

\[g(z) = \frac{c_0}{2} + \sum_{j=1}^{n_{\max}} c_j \cos\!\left(\frac{2\pi j\,(z - z_{\mathrm{mid}})}{L}\right) + \sum_{j=1}^{n_{\max}} s_j \sin\!\left(\frac{2\pi j\,(z - z_{\mathrm{mid}})}{L}\right),\]

where \(z_{\mathrm{mid}} = (z_{\min} + z_{\max})/2\) is the longitudinal location of the midpoint, and \(L = z_{\max} - z_{\min}\) is the total length of the z-domain.

The benchmark test uses these coefficients to define and track through a soft-edge quadrupole element. Note that in practice, the ImpactX soft-edge elements also support specifying the on-axis z data as arguments directly as arguments, performing he Fourier coefficient calculation automatically.

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 as:

Python script: python3 run_quadrupole_fcoef.py or

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

Python: Script

Listing 64 You can copy this file from examples/fourier_coefficients/run_quadrupole_fcoef.py.

#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from fcoef import read_data, write_data

from impactx import ImpactX, distribution, elements, fourier_coefficients

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 2 GeV electron beam with an initial
# unnormalized rms emittance of 2 nm
kin_energy_MeV = 2.0e3  # reference energy
bunch_charge_C = 1.0e-9  # used with space charge
npart = 10000  # number of macro particles

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

#   particle bunch
distr = distribution.Waterbag(
    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 = 1  # number of slices per ds in the element

# read in the on-axis quadrupole gradient data
z, gradient_on_axis = read_data("onaxis_data.in")

# optional: compute and write coefficients to file (to visually compare)
ncoef = 25
cos_coeffs, sin_coeffs = fourier_coefficients(z, gradient_on_axis, ncoef)
write_data(cos_coeffs, sin_coeffs, z, "onaxis_data.out")

# lattice: construct SoftQuadrupole directly from on-axis field data
quad1 = elements.SoftQuadrupole(
    name="quad1",
    ds=0.2495,
    gscale=1.0,
    z=z,
    gradient_on_axis=gradient_on_axis,
    ncoef=ncoef,
    mapsteps=400,
    nslice=ns,
)

drift1 = elements.Drift(name="drift1", ds=0.25, nslice=ns)

# assign a fodo segment
sim.lattice.extend([monitor, drift1, quad1, drift1, monitor])

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

Analyze

We run the following script to analyze correctness:

Script analysis_quadrupole_fcoef.py

Listing 65 You can copy this file from examples/fourier_coefficients/analysis_quadrupole_fcoef.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 = 3.0 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
    [
        7.5451170454175073e-005,
        7.5441588239210947e-005,
        9.9775878164077539e-004,
        1.9959540393751392e-009,
        2.0175015289132990e-009,
        2.0013820193294972e-006,
    ],
    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 = 3.0 * num_particles**-0.5  # from random sampling of a smooth distribution
print(f"  rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
    [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
    [
        1.096989e-04,
        4.735185e-05,
        1.004119e-03,
        2.045074e-09,
        2.005126e-09,
        1.993852e-06,
    ],
    rtol=rtol,
    atol=atol,
)

Visualize

You can run the following script to visualize the reconstruction of the data on-axis:

Script plot_fcoef.py

Listing 66 You can copy this file from examples/fourier_coefficients/plot_fcoef.py.

#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#

import argparse

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MaxNLocator

# options to run this script
parser = argparse.ArgumentParser(
    description="Plot the reconstruction of on-axis data from Fourier coefficients."
)
parser.add_argument(
    "--save-png", action="store_true", help="non-interactive run: save to PNGs"
)
args = parser.parse_args()

data_in = np.loadtxt("onaxis_data.in")
z = data_in[:, 0]
g_onaxis = data_in[:, 1]
data_out = np.loadtxt("onaxis_data.out")
g_onaxis_reconstructed = data_out[:, 1]
gp_onaxis_reconstructed = data_out[:, 2]
gpp_onaxis_reconstructed = data_out[:, 3]
diff_onaxis = g_onaxis - g_onaxis_reconstructed

# print beam transverse size over steps
f = plt.figure(figsize=(9, 4.8))
ax1 = f.gca()
input_data = ax1.plot(z, g_onaxis, label=r"input data $g(z)$")
output_data = ax1.plot(z, g_onaxis_reconstructed, label=r"reconstructed data")
diff_data = ax1.plot(z, diff_onaxis, label=r"difference")

ax1.legend(
    handles=input_data + output_data + diff_data, loc="upper center", prop={"size": 11}
)
ax1.set_xlabel(r"$z$ [m]", fontsize=14)
ax1.set_ylabel(r"gradient", fontsize=14)
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.tight_layout()
if args.save_png:
    plt.savefig("fodo_sigma.png")
else:
    plt.show()

# print first derivative from reconstructed data
f = plt.figure(figsize=(9, 4.8))
ax1 = f.gca()
derivative1 = ax1.plot(z, gp_onaxis_reconstructed, label=r"$g'(z)$")

ax1.legend(handles=derivative1, loc="upper center", prop={"size": 11})
ax1.set_xlabel(r"$z$ [m]", fontsize=14)
ax1.set_ylabel(r"derivative", fontsize=14)
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.tight_layout()
if args.save_png:
    plt.savefig("fodo_polarization.png")
else:
    plt.show()

# print first and second derivatives from reconstructed data
f = plt.figure(figsize=(9, 4.8))
ax1 = f.gca()
derivative2 = ax1.plot(z, gpp_onaxis_reconstructed, label=r"$g''(z)$")

ax1.legend(handles=derivative2, loc="upper center", prop={"size": 11})
ax1.set_xlabel(r"$z$ [m]", fontsize=14)
ax1.set_ylabel(r"derivative", fontsize=14)
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.tight_layout()
if args.save_png:
    plt.savefig("fodo_polarization.png")
else:
    plt.show()