Transformation

Test the t/s transformations on an electron beam.

We use a long electron beam, \(L_z=1\) cm, with significant correlations in \(x-px\), \(y-py\), and \(t-pt\). The beam has average energy 1 GeV.

This tests that the t/s transforms are inverses of each other Specifically, in this test the \(t\)- and \(s\)-coordinates of the beam must differ substantially and the forward-inverse transformed coordinates must agree with the initial coordinates. That is, we require that direction=CoordSystem.s (direction=CoordSystem.t (initial beam)) = initial beam.

Run

This file is run from pytest.

Python Script

Listing 119 You can copy this file from tests/python/test_transformation.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 numpy as np
import pytest
from impactx import CoordSystem, ImpactX, coordinate_transformation, distribution


def test_transformation():
    """
    This test ensures s->t and t->s transformations
    do round-trip.
    """
    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()

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

    #   reference particle
    pc = sim.particle_container()
    ref = pc.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=3e-6,
        lambdaY=3e-6,
        lambdaT=1e-2,
        lambdaPx=1.33 / energy_gamma,
        lambdaPy=1.33 / energy_gamma,
        lambdaPt=100 / energy_gamma,
        muxpx=-0.5,
        muypy=0.4,
        mutpt=0.8,
    )
    sim.add_particles(bunch_charge_C, distr, npart)
    rbc_s0 = pc.reduced_beam_characteristics()

    # this must fail: we cannot transform from s to s
    with pytest.raises(Exception):
        coordinate_transformation(pc, direction=CoordSystem.s)

    # transform to t
    coordinate_transformation(pc, direction=CoordSystem.t)
    rbc_t = pc.reduced_beam_characteristics()

    # this must fail: we cannot transform from t to t
    with pytest.raises(Exception):
        coordinate_transformation(pc, direction=CoordSystem.t)

    # transform back to s
    coordinate_transformation(pc, direction=CoordSystem.s)
    rbc_s = pc.reduced_beam_characteristics()

    # finalize simulation
    sim.finalize()
    del sim

    # assert that forward-inverse transformation of the beam leaves beam unchanged
    atol = 1e-14
    rtol = 1e-10
    for key, val in rbc_s0.items():
        if not np.isclose(val, rbc_s[key], rtol=rtol, atol=atol):
            print(f"initial[{key}]={val}, final[{key}]={rbc_s[key]} not equal")
        assert np.isclose(val, rbc_s[key], rtol=rtol, atol=atol)
    # assert that the t-based beam is different, at least in the following keys:
    large_st_diff_keys = [
        "beta_x",
        "beta_y",
        "emittance_y",
        "emittance_x",
        "sig_y",
        "sig_x",
        "t_mean",
    ]
    for key in large_st_diff_keys:
        rel_error = (rbc_s0[key] - rbc_t[key]) / rbc_s0[key]
        assert abs(rel_error) > 1