Initialize a Beam from Arrays

Note

This example is an early draft of this workflow and will be simplified in future versions of ImpactX, adding direct support for NumPy and CuPy arrays.

This example demonstrates how a beam can be initalized in ImpactX from array-like structures. This allows various applications of interest, such as using a beam from a different simulation, initializing a beam from file, or creating a custom distribution. This example includes a set of utilities for transforming the beam to the fixed-s coordinates of ImpactX.

In this example, a custom beam is specified at fixed t, transformed to fixed s, and then loaded in ImpactX. The custom beam is a ring in x-y, with radius r=2 mm, radial width \(\sigma_r = 5\ \mu\) and \(p_y\) with momentum width \(\sigma_p=10\); and chirped in z-pz with bunch length \(\sigma_z=1\) mm, mean energy about 10 GeV, 1% uncorrelated energy spread, and z-pz covariance of -0.18.

In specifying the beam at fixed t and transforming to fixed s, it is assumed that the local and global coordinate frames align. That is, the beam transverse directions x and y are the global x and y directions and the beam z/t direction is the global z direction. The transformation utility function reproduces the t-to-s and s-to-t transformations done internally in ImpactX given this assumption that the beam and global coordinate systems align. These utility functions are provided in the following script:

Script transformation_utilities.py

Listing 85 You can copy this file from examples/initialize_from_array/transformation_utilities.py.

import numpy as np


def to_ref_part_t_from_global_t(ref_part, x, y, z, px, py, pz):
    """
    Transform from global coordinates to reference particle coordinates

    This function takes particle bunch data and returns the bunch phase space positions relative to a reference particle.
    It is assumed that the local and global coordinate frames align.
    That is, the beam transverse directions x and y are the global x and y directions
    and the beam z direction is the global z direction.

    Parameters
    ----------
    ref_part: reference particle object, either from ImpactX or from the class defined in this file, MyRefPart
    x: array-like, global beam particle x-positions
    y: array-like, global beam particle y-positions
    z: array-like, global beam particle z-positions
    px: array-like, global beam particle x-momenta
    py: array-like, global beam particle y-momenta
    pz: array-like, global beam particle z-momenta

    Returns
    -------
    dx: array-like, beam particle x-positions relative to reference particle x value, ``ref_part.x``
    dy: array-like, beam particle y-positions relative to reference particle y value, ``ref_part.y``
    dz: array-like, beam particle z-positions relative to reference particle z value, ``ref_part.z``
    dpx: array-like, beam particle x-momenta relative to reference particle x value, ``ref_part.px``
    dpy: array-like, beam particle y-momenta relative to reference particle y value, ``ref_part.py``
    dpz: array-like, beam particle z-momenta relative to reference particle z value, ``ref_part.pz``
    """
    dx = x - ref_part.x
    dy = y - ref_part.y
    dz = z - ref_part.z
    dpx = px - ref_part.px
    dpy = py - ref_part.py
    dpz = pz - ref_part.pz

    dpx /= ref_part.pz
    dpy /= ref_part.pz
    dpz /= ref_part.pz

    return dx, dy, dz, dpx, dpy, dpz


def to_global_t_from_ref_part_t(ref_part, dx, dy, dz, dpx, dpy, dpz):
    """
    Transform from reference particle to global coordinates.

    This function takes particle bunch data relative to a reference particle
    and returns all particle data in the global coordinate frame.
    It is assumed that the local and global coordinate frames align.
    That is, the beam transverse directions x and y are the global x and y directions
    and the beam z direction is the global z direction.

    Parameters
    ----------
    ref_part: reference particle object, either from ImpactX or from the class defined in this file, MyRefPart
    dx: array-like, beam particle x-positions relative to reference particle x value, ``ref_part.x``
    dy: array-like, beam particle y-positions relative to reference particle y value, ``ref_part.y``
    dz: array-like, beam particle z-positions relative to reference particle z value, ``ref_part.z``
    dpx: array-like, beam particle x-momenta relative to reference particle x value, ``ref_part.px``
    dpy: array-like, beam particle y-momenta relative to reference particle y value, ``ref_part.py``
    dpz: array-like, beam particle z-momenta relative to reference particle z value, ``ref_part.pz``

    Returns
    -------
    x: array-like, global beam particle x-positions
    y: array-like, global beam particle y-positions
    z: array-like, global beam particle z-positions
    px: array-like, global beam particle x-momenta
    py: array-like, global beam particle y-momenta
    pz: array-like, global beam particle z-momenta
    """
    x = dx + ref_part.x
    y = dy + ref_part.y
    z = dz + ref_part.z

    px = ref_part.px + ref_part.pz * dpx
    py = ref_part.py + ref_part.pz * dpy
    pz = ref_part.pz + ref_part.pz * dpz

    return x, y, z, px, py, pz


def to_s_from_t(ref_part, dx, dy, dz, dpx, dpy, dpz):  # data_arr_t):
    """
    Transform from fixed-s to fixed-t coordinates

    This function takes particle bunch data relative to a reference particle
    at a fixed time t and returns data at a fixed longitudinal position s.
    That is, spatial distance is transformed to time delay from the reference particle.

    Parameters
    ----------
    ref_part: reference particle object, either from ImpactX or from the class defined in this file, MyRefPart
    dx: array-like, beam particle x-positions relative to reference particle at fixed t
    dy: array-like, beam particle y-positions relative to reference particle at fixed t
    dz: array-like, beam particle z-positions relative to reference particle at fixed t
    dpx: array-like, beam particle x-momenta relative to reference particle at fixed t
    dpy: array-like, beam particle y-momenta relative to reference particle at fixed t
    dpz: array-like, beam particle z-momenta relative to reference particle at fixed t

    Returns
    -------
    dxs: array-like, beam particle x-positions relative to reference particle at fixed s
    dys: array-like, beam particle y-positions relative to reference particle at fixed s
    dt: array-like, beam particle time delay relative to reference particle at fixed s
    dpx: array-like, beam particle x-momenta relative to reference particle at fixed s
    dpy: array-like, beam particle y-momenta relative to reference particle at fixed s
    dpz: array-like, beam particle t-momenta (-gamma) relative to reference particle at fixed s
    """
    ref_pz = ref_part.pz
    ref_pt = ref_part.pt
    dxs = dx - ref_pz * dpx * dz / (ref_pz + ref_pz * dpz)
    dys = dy - ref_pz * dpy * dz / (ref_pz + ref_pz * dpz)
    pt = -np.sqrt(
        1 + (ref_pz + ref_pz * dpz) ** 2 + (ref_pz * dpx) ** 2 + (ref_pz * dpy) ** 2
    )
    dt = pt * dz / (ref_pz + ref_pz * dpz)
    dpt = (pt - ref_pt) / ref_pz
    return dxs, dys, dt, dpx, dpy, dpt


def to_t_from_s(ref_part, dx, dy, dt, dpx, dpy, dpt):  # data_arr_t):
    """
    Transform from fixed-t to fixed-s coordinates

    This function takes particle bunch data relative to a reference particle
    at a fixed longitudinal position s and returns data at a fixed time t.
    That is, time delay is transformed to spatial distance from the reference particle.

    Parameters
    ----------
    ref_part: reference particle object, either from ImpactX or from the class defined in this file, MyRefPart
    dx: array-like, beam particle x-positions relative to reference particle at fixed s
    dy: array-like, beam particle y-positions relative to reference particle at fixed s
    dz: array-like, beam particle time delay relative to reference particle at fixed s
    dpx: array-like, beam particle x-momenta relative to reference particle at fixed s
    dpy: array-like, beam particle y-momenta relative to reference particle at fixed s
    dpt: array-like, beam particle t-momenta (-gamma) relative to reference particle at fixed s

    Returns
    -------
    dxt: array-like, beam particle x-positions relative to reference particle at fixed t
    dyt: array-like, beam particle y-positions relative to reference particle at fixed t
    dt: array-like, beam particle z-positions relative to reference particle at fixed t
    dpx: array-like, beam particle x-momenta relative to reference particle at fixed t
    dpy: array-like, beam particle y-momenta relative to reference particle at fixed t
    dpz: array-like, beam particle z-momenta relative to reference particle at fixed t
    """
    ref_pz = ref_part.pz
    ref_pt = ref_part.pt
    dxt = dx + ref_pz * dpx * dt / (ref_pt + ref_pz * dpt)
    dyt = dy + ref_pz * dpy * dt / (ref_pt + ref_pz * dpt)

    pz = np.sqrt(
        -1 + (ref_pt + ref_pz * dpt) ** 2 - (ref_pz * dpx) ** 2 - (ref_pz * dpy) ** 2
    )
    dz = dt * pz / (ref_pt + ref_pz * dpt)
    dpz = (pz - ref_pz) / ref_pz
    return dxt, dyt, dz, dpx, dpy, dpz


class MyRefPart:
    """Struct containing reference particle data

    This class replicates the data storage of the ImpactX reference particle.
    It is used in coordinate transformations when an ImpactX reference particle isn't available,
    so the transformation syntax works in the context of pyImpactX
    """

    def __init__(self, x, y, z, px, py, pz, pt):
        self.attr_list = ["x", "y", "z", "px", "py", "pz", "pt"]
        self.x = x
        self.y = y
        self.z = z
        self.px = px
        self.py = py
        self.pz = pz
        self.pt = pt

    def __repr__(self):
        mystr = ""
        for attr in self.attr_list:
            mystr += f"self.{attr}={getattr(self,attr)}, "
        return mystr

    def __str__(self):
        mystr = ""
        for attr in self.attr_list:
            mystr += f"self.{attr}={getattr(self,attr)}, "
        return mystr

Run

This example can only be run with Python:

• Python script: python3 run_from_array.py

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

Listing 86 You can copy this file from examples/initialize_from_array/run_from_array.py.

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

import amrex.space3d as amr
import numpy as np
import transformation_utilities as pycoord
from impactx import Config, ImpactX, elements

################

N_part = int(1e5)
beam_radius = 2e-3
sigr = 500e-6
sigpx = 10
sigpy = 10
px = np.random.normal(0, sigpx, N_part)
py = np.random.normal(0, sigpy, N_part)
theta = 2 * np.pi * np.random.rand(N_part)
r = abs(np.random.normal(beam_radius, sigr, N_part))
x = r * np.cos(theta)
y = r * np.sin(theta)
z_mean = 0
pz_mean = 2e4
z_std = 1e-3
pz_std = 2e2
zpz_std = -0.18
zpz_cov_list = [[z_std**2, zpz_std], [zpz_std, pz_std**2]]
z, pz = np.random.default_rng().multivariate_normal([0, 0], zpz_cov_list, N_part).T
pz += pz_mean

sim = ImpactX()

# set numerical parameters and IO control
sim.particle_shape = 2  # B-spline order
sim.space_charge = False
# sim.diagnostics = False  # benchmarking
sim.slice_step_diagnostics = True

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

energy_gamma = np.sqrt(1 + pz_mean**2)
energy_MeV = 0.510998950 * energy_gamma  # reference energy
bunch_charge_C = 10.0e-15  # used with space charge

#   reference particle
ref = sim.particle_container().ref_particle()
ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(energy_MeV)
qm_eev = -1.0 / 0.510998950 / 1e6  # electron charge/mass in e / eV
ref.z = 0

pc = sim.particle_container()

dx, dy, dz, dpx, dpy, dpz = pycoord.to_ref_part_t_from_global_t(
    ref, x, y, z, px, py, pz
)
dx, dy, dt, dpx, dpy, dpt = pycoord.to_s_from_t(ref, dx, dy, dz, dpx, dpy, dpz)

if not Config.have_gpu:  # initialize using cpu-based PODVectors
    dx_podv = amr.PODVector_real_std()
    dy_podv = amr.PODVector_real_std()
    dt_podv = amr.PODVector_real_std()
    dpx_podv = amr.PODVector_real_std()
    dpy_podv = amr.PODVector_real_std()
    dpt_podv = amr.PODVector_real_std()
else:  # initialize on device using arena/gpu-based PODVectors
    dx_podv = amr.PODVector_real_arena()
    dy_podv = amr.PODVector_real_arena()
    dt_podv = amr.PODVector_real_arena()
    dpx_podv = amr.PODVector_real_arena()
    dpy_podv = amr.PODVector_real_arena()
    dpt_podv = amr.PODVector_real_arena()

for p_dx in dx:
    dx_podv.push_back(p_dx)
for p_dy in dy:
    dy_podv.push_back(p_dy)
for p_dt in dt:
    dt_podv.push_back(p_dt)
for p_dpx in dpx:
    dpx_podv.push_back(p_dpx)
for p_dpy in dpy:
    dpy_podv.push_back(p_dpy)
for p_dpt in dpt:
    dpt_podv.push_back(p_dpt)

pc.add_n_particles(
    dx_podv, dy_podv, dt_podv, dpx_podv, dpy_podv, dpt_podv, qm_eev, bunch_charge_C
)

monitor = elements.BeamMonitor("monitor", backend="h5")
sim.lattice.extend(
    [
        monitor,
        elements.Drift(ds=0.01),
        monitor,
    ]
)

sim.evolve()

# clean shutdown
sim.finalize()

Analyze

We run the following script to analyze correctness:

Script analyze_from_array.py

Listing 87 You can copy this file from examples/initialize_from_array/analyze_from_array.py.

#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Ryan Sandberg
# License: BSD-3-Clause-LBNL
#

import impactx_read_utilities as ix_read
import numpy as np
import openpmd_api as io
import transformation_utilities as coord

print("Initial Beam:")

beam_series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
ref_series = ix_read.read_time_series("diags/ref_particle.*")

num_particles = int(1e5)
millimeter = 1.0e3
micron = 1.0e6
step = 1
beam_at_step = beam_series.iterations[step].particles["beam"].to_df()
ref_part_step = ref_series.loc[step]
ref_part = myref = coord.MyRefPart(
    ref_part_step["x"],
    ref_part_step["y"],
    ref_part_step["z"],
    ref_part_step["px"],
    ref_part_step["py"],
    ref_part_step["pz"],
    ref_part_step["pt"],
)
dx_s = beam_at_step["position_x"]
dy_s = beam_at_step["position_y"]
dt = beam_at_step["position_t"]
dpx_s = beam_at_step["momentum_x"]
dpy_s = beam_at_step["momentum_y"]
dpt = beam_at_step["momentum_t"]

dx_t, dy_t, dz, dpx_t, dpy_t, dpz = coord.to_t_from_s(
    ref_part, dx_s, dy_s, dt, dpx_s, dpy_s, dpt
)
x1, y1, z1, px1, py1, pz1 = coord.to_global_t_from_ref_part_t(
    ref_part, dx_t, dy_t, dz, dpx_t, dpy_t, dpz
)

sigx = x1.std()
sigy = y1.std()
sigz = z1.std()
sigpx = px1.std()
sigpy = py1.std()
sigpz = pz1.std()
print(f"sig x={sigx:.2e}, sig y={sigy:.2e}, sig z={sigz:.2e}")
print(f"sig px={sigpx:.2e}, sig py={sigpy:.2e}, sig pz={sigpz:.2e}")

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

assert np.allclose(
    [sigx, sigy, sigz, sigpx, sigpy, sigpz],
    [1.46e-3, 1.46e-3, 1.0e-3, 10, 10, 200],
    rtol=rtol,
    atol=atol,
)

This uses the following utilities to read ImpactX output:

Script impactx_read_utilities.py

Listing 88 You can copy this file from examples/initialize_from_array/impactx_read_utilities.py.

import glob
import re

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')

Visualize

We run the following script to visualize the ImpactX output and confirm the beam is properly initialized:

Script visualize_from_array.py

Listing 89 You can copy this file from examples/initialize_from_array/visualize_from_array.py.

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

import impactx_read_utilities as ix_plt
import openpmd_api as io
import transformation_utilities as coord
from matplotlib import pyplot as plt

######## plot phase spaces ###########
beam_series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
beam_steps = list(beam_series.iterations)
ref_series = ix_plt.read_time_series("diags/ref_particle.*")

millimeter = 1.0e3
micron = 1.0e6
step = 1
beam_at_step = beam_series.iterations[step].particles["beam"].to_df()
ref_part_step = ref_series.loc[step]
ref_part = coord.MyRefPart(
    ref_part_step["x"],
    ref_part_step["y"],
    ref_part_step["z"],
    ref_part_step["px"],
    ref_part_step["py"],
    ref_part_step["pz"],
    ref_part_step["pt"],
)

dx_s = beam_at_step["position_x"]
dy_s = beam_at_step["position_y"]
dt = beam_at_step["position_t"]
dpx_s = beam_at_step["momentum_x"]
dpy_s = beam_at_step["momentum_y"]
dpt = beam_at_step["momentum_t"]

dx_t, dy_t, dz, dpx_t, dpy_t, dpz = coord.to_t_from_s(
    ref_part, dx_s, dy_s, dt, dpx_s, dpy_s, dpt
)
x, y, z, px, py, pz = coord.to_global_t_from_ref_part_t(
    ref_part, dx_t, dy_t, dz, dpx_t, dpy_t, dpz
)

fig, axT = plt.subplots(1, 3, figsize=(8, 4))

ax = axT[0]
ax.hist2d(x * millimeter, y * millimeter, bins=200)
ax.set_xlabel("x (mm)")
ax.set_ylabel("y (mm)")

ax = axT[1]
ax.hist2d(px, py, bins=200)
ax.set_xlabel("px")
ax.set_ylabel("py")
ax = axT[2]
ax.hist2d(z * millimeter, pz, bins=200)
ax.set_xlabel("z (mm)")
ax.set_ylabel("pz")
ax.ticklabel_format(axis="y", style="sci", scilimits=(-1, 1))

plt.tight_layout()
plt.savefig("phase_space.png")

The resulting phase space snapshots are shown in the following figure:

Fig. 6 [fig:custom_beam] Phase space snapshots of custom beam from arrays, showing the ring in x-y, Gaussian in px-py, and linear correlation in z-pz.