ImpactX
ImpactX is an s-based beam dynamics code including space charge effects. This is the next generation of the IMPACT-Z code.
Please contact us with any questions on it or if you like to contribute to its development.
Contact us
If you are starting using ImpactX, or if you have a user question, please pop in our discussions page and get in touch with the community.
The ImpactX GitHub repo is the main communication platform. Have a look at the action icons on the top right of the web page: feel free to watch the repo if you want to receive updates, or to star the repo to support the project. For bug reports or to request new features, you can also open a new issue.
We also have a discussion page on which you can find already answered questions, add new questions, get help with installation procedures, discuss ideas or share comments.
Code of Conduct
Our Pledge
In the interest of fostering an open and welcoming environment, we as contributors and maintainers pledge to making participation in our project and our community a harassment-free experience for everyone, regardless of age, body size, disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, religion, or sexual identity and orientation.
Our Standards
Examples of behavior that contributes to creating a positive environment include:
Using welcoming and inclusive language
Being respectful of differing viewpoints and experiences
Gracefully accepting constructive criticism
Focusing on what is best for the community
Showing empathy towards other community members
Examples of unacceptable behavior by participants include:
The use of sexualized language or imagery and unwelcome sexual attention or advances
Trolling, insulting/derogatory comments, and personal or political attacks
Public or private harassment
Publishing others’ private information, such as a physical or electronic address, without explicit permission
Other conduct which could reasonably be considered inappropriate in a professional setting
Our Responsibilities
Project maintainers are responsible for clarifying the standards of acceptable behavior and are expected to take appropriate and fair corrective action in response to any instances of unacceptable behavior.
Project maintainers have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, or to ban temporarily or permanently any contributor for other behaviors that they deem inappropriate, threatening, offensive, or harmful.
Scope
This Code of Conduct applies both within project spaces and in public spaces when an individual is representing the project or its community. Examples of representing a project or community include using an official project e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event. Representation of a project may be further defined and clarified by project maintainers.
Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be reported by contacting the project team at warpx-coc@lbl.gov. All complaints will be reviewed and investigated and will result in a response that is deemed necessary and appropriate to the circumstances. The project team is obligated to maintain confidentiality with regard to the reporter of an incident. Further details of specific enforcement policies may be posted separately.
Project maintainers who do not follow or enforce the Code of Conduct in good faith may face temporary or permanent repercussions as determined by other members of the project’s leadership.
Attribution
This Code of Conduct is adapted from the Contributor Covenant, version 1.4, available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
For answers to common questions about this code of conduct, see https://www.contributor-covenant.org/faq
Acknowledge ImpactX
Please acknowledge the role that ImpactX played in your research.
In presentations
Note
TODO :-)
In publications
Please add the following sentence to your publications, it helps contributors keep in touch with the community and promote the project.
Plain text:
This research used the open-source particle-in-cell code ImpactX https://github.com/ECP-WarpX/impactx. We acknowledge all ImpactX contributors.
Latex:
\usepackage{hyperref}
This research used the open-source particle-in-cell code ImpactX \url{https://github.com/ECP-WarpX/impactx}.
We acknowledge all ImpactX contributors.
Main ImpactX Reference
If your project leads to a scientific publication, please consider citing the paper below.
Huebl A, Lehe R, Mitchell C E, Qiang J, Ryne R D, Sandberg R T, Vay JL. Next Generation Computational Tools for the Modeling and Design of Particle Accelerators at Exascale. 2022 North American Particle Accelerator Conference (NAPAC’22), TUYE2, pp. 302-306, 2022. arXiv:2208.02382, DOI:10.18429/JACoW-NAPAC2022-TUYE2
Further ImpactX References
Sandberg R T, Lehe R, Mitchell C E, Garten M, Myers A, Qiang J, Vay J-L and Huebl A. Synthesizing Particle-in-Cell Simulations Through Learning and GPU Computing for Hybrid Particle Accelerator Beamlines. Proc. of Platform for Advanced Scientific Computing (PASC’24), submitted, 2024. arXiv:2402.17248
Sandberg R T, Lehe R, Mitchell C E, Garten M, Qiang J, Vay J-L and Huebl A. Hybrid Beamline Element ML-Training for Surrogates in the ImpactX Beam-Dynamics Code. 14th International Particle Accelerator Conference (IPAC’23), WEPA101, 2023. DOI:10.18429/JACoW-IPAC2023-WEPA101
Huebl A, Lehe R, Zoni E, Shapoval O, Sandberg R T, Garten M, Formenti A, Jambunathan R, Kumar P, Gott K, Myers A, Zhang W, Almgren A, Mitchell C E, Qiang J, Sinn A, Diederichs S, Thevenet M, Grote D, Fedeli L, Clark T, Zaim N, Vincenti H, Vay JL. From Compact Plasma Particle Sources to Advanced Accelerators with Modeling at Exascale. Proceedings of the 20th Advanced Accelerator Concepts Workshop (AAC’22), in print, 2023. arXiv:2303.12873
Installation
Users
Our community is here to help. Please report installation problems in case you should get stuck.
Choose one of the installation methods below to get started:
HPC Systems
If want to use ImpactX on a specific high-performance computing (HPC) systems, jump directly to our HPC system-specific documentation.
Using the Conda Package
A package for ImpactX is available via the Conda package manager.
Tip
We recommend to configure your conda to use the faster libmamba
dependency solver.
conda update -y -n base conda
conda install -y -n base conda-libmamba-solver
conda config --set solver libmamba
We recommend to deactivate that conda self-activates its base
environment.
This avoids interference with the system and other package managers.
conda config --set auto_activate_base false
conda create -n impactx -c conda-forge impactx
conda activate impactx
Note
The impactx
conda package does not yet provide GPU support.
From Source with CMake
After installing the ImpactX dependencies, you can also install ImpactX from source with CMake:
# get the source code
git clone https://github.com/ECP-WarpX/impactx.git $HOME/src/impactx
cd $HOME/src/impactx
# configure
cmake -S . -B build
# optional: change configuration
ccmake build
# compile
# on Windows: --config Release
cmake --build build -j 4
# executables for ImpactX are now in build/bin/
We document the details in the developer installation.
Tips for macOS Users
Tip
Before getting started with package managers, please check what you manually installed in /usr/local
.
If you find entries in bin/
, lib/
et al. that look like you manually installed MPI, HDF5 or other software in the past, then remove those files first.
If you find software such as MPI in the same directories that are shown as symbolic links then it is likely you brew installed software before.
If you are trying annother package manager than brew
, run brew unlink … on such packages first to avoid software incompatibilities.
See also: A. Huebl, Working With Multiple Package Managers, Collegeville Workshop (CW20), 2020
Developers
CMake is our primary build system. If you are new to CMake, this short tutorial from the HEP Software foundation is the perfect place to get started. If you just want to use CMake to build the project, jump into sections 1. Introduction, 2. Building with CMake and 9. Finding Packages.
Dependencies
Before you start, you will need a copy of the ImpactX source code:
git clone https://github.com/ECP-WarpX/impactx.git $HOME/src/impactx
cd $HOME/src/impact
ImpactX depends on popular third party software.
On your development machine, follow the instructions here.
If you are on an HPC machine, follow the instructions here.
Dependencies
ImpactX depends on the following popular third party software. Please see installation instructions below.
a mature C++17 compiler, e.g., GCC 8.4+, Clang 7, NVCC 11.0, MSVC 19.15 or newer
AMReX: we automatically download and compile a copy
ABLASTR/WarpX: we automatically download and compile a copy
Optional dependencies include:
MPI 3.0+: for multi-node and/or multi-GPU execution
for on-node accelerated compute one of either:
OpenMP 3.1+: for threaded CPU execution or
CUDA Toolkit 11.0+ (11.3+ recommended): for Nvidia GPU support (see matching host-compilers) or
ROCm 5.2+ (5.5+ recommended): for AMD GPU support
openPMD-api 0.15.1+: we automatically download and compile a copy of openPMD-api for openPMD I/O support
CCache: to speed up rebuilds (For CUDA support, needs version 3.7.9+ and 4.2+ is recommended)
Ninja: for faster parallel compiles
-
see our
requirements.txt
file for compatible versions
If you are on a high-performance computing (HPC) system, then please see our separate HPC documentation.
For all other systems, we recommend to use a package dependency manager: Pick one of the installation methods below to install all dependencies for ImpactX development in a consistent manner.
Conda (Linux/macOS/Windows)
Conda/Mamba are cross-compatible, user-level package managers.
Tip
We recommend to configure your conda to use the faster libmamba
dependency solver.
conda update -y -n base conda
conda install -y -n base conda-libmamba-solver
conda config --set solver libmamba
We recommend to deactivate that conda self-activates its base
environment.
This avoids interference with the system and other package managers.
conda config --set auto_activate_base false
conda create -n impactx-cpu-mpich-dev -c conda-forge blaspp boost ccache cmake compilers git lapackpp "openpmd-api=*=mpi_mpich*" python numpy pandas scipy yt pkg-config matplotlib mamba ninja mpich pip virtualenv
conda activate impactx-cpu-mpich-dev
# compile ImpactX with -DImpactX_MPI=ON
# for pip, use: export IMPACTX_MPI=ON
conda create -n impactx-cpu-dev -c conda-forge blaspp boost ccache cmake compilers git lapackpp openpmd-api python numpy pandas scipy yt pkg-config matplotlib mamba ninja pip virtualenv
conda activate impactx-cpu-dev
# compile ImpactX with -DImpactX_MPI=OFF
# for pip, use: export IMPACTX_MPI=OFF
For OpenMP support, you will further need:
conda install -c conda-forge libgomp
conda install -c conda-forge llvm-openmp
Spack (Linux/macOS)
Spack is a user-level package manager. It is primarily written for Linux, with slightly less support for macOS, and future support for Windows.
First, download a WarpX Spack desktop development environment of your choice (which will also work for ImpactX). For most desktop developments, pick the OpenMP environment for CPUs unless you have a supported GPU.
Debian/Ubuntu Linux:
OpenMP:
system=ubuntu; compute=openmp
(CPUs)CUDA:
system=ubuntu; compute=cuda
(Nvidia GPUs)ROCm:
system=ubuntu; compute=rocm
(AMD GPUs)SYCL: todo (Intel GPUs)
macOS: first, prepare with
brew install gpg2; brew install gcc
OpenMP:
system=macos; compute=openmp
If you already installed Spack, we recommend to activate its binary caches for faster builds:
spack mirror add rolling https://binaries.spack.io/develop
spack buildcache keys --install --trust
Now install the WarpX/ImpactX dependencies in a new development environment:
# download environment file
curl -sLO https://raw.githubusercontent.com/ECP-WarpX/WarpX/development/Tools/machines/desktop/spack-${system}-${compute}.yaml
# create new development environment
spack env create impactx-${compute}-dev spack-${system}-${compute}.yaml
spack env activate impactx-${compute}-dev
# installation
spack install
python3 -m pip install jupyter matplotlib numpy openpmd-api openpmd-viewer pandas scipy virtualenv yt
In new terminal sessions, re-activate the environment with
spack env activate impactx-openmp-dev
again.
Replace openmp
with the equivalent you chose.
Compile ImpactX with -DImpactX_MPI=ON
.
For pip
, use export IMPACTX_MPI=ON
.
Brew (macOS/Linux)
Homebrew (Brew) is a user-level package manager primarily for Apple macOS, but also supports Linux.
brew update
brew tap openpmd/openpmd
brew install adios2 # for openPMD
brew install ccache
brew install cmake
brew install git
brew install hdf5-mpi # for openPMD
brew install libomp
brew unlink gcc
brew link --force libomp
brew install open-mpi
brew install openblas # for PSATD in RZ
brew install openpmd-api # for openPMD
If you also want to compile with PSATD in RZ, you need to manually install BLAS++ and LAPACK++:
sudo mkdir -p /usr/local/bin/
sudo curl -L -o /usr/local/bin/cmake-easyinstall https://raw.githubusercontent.com/ax3l/cmake-easyinstall/main/cmake-easyinstall
sudo chmod a+x /usr/local/bin/cmake-easyinstall
cmake-easyinstall --prefix=/usr/local git+https://github.com/icl-utk-edu/blaspp.git \
-Duse_openmp=OFF -Dbuild_tests=OFF -DCMAKE_VERBOSE_MAKEFILE=ON
cmake-easyinstall --prefix=/usr/local git+https://github.com/icl-utk-edu/lapackpp.git \
-Duse_cmake_find_lapack=ON -Dbuild_tests=OFF -DCMAKE_VERBOSE_MAKEFILE=ON
Compile ImpactX with -DImpactX_MPI=ON
.
For pip
, use export IMPACTX_MPI=ON
.
APT (Debian/Ubuntu Linux)
The Advanced Package Tool (APT) is a system-level package manager on Debian-based Linux distributions, including Ubuntu.
sudo apt update
sudo apt install build-essential ccache cmake g++ git libhdf5-openmpi-dev libopenmpi-dev pkg-config python3 python3-matplotlib python3-numpy python3-pandas python3-pip python3-scipy python3-venv
# optional:
# for CUDA, either install
# https://developer.nvidia.com/cuda-downloads (preferred)
# or, if your Debian/Ubuntu is new enough, use the packages
# sudo apt install nvidia-cuda-dev libcub-dev
# compile ImpactX with -DImpactX_MPI=ON
# for pip, use: export IMPACTX_MPI=ON
sudo apt update
sudo apt install build-essential ccache cmake g++ git libhdf5-dev pkg-config python3 python3-matplotlib python3-numpy python3-pandas python3-pip python3-scipy python3-venv
# optional:
# for CUDA, either install
# https://developer.nvidia.com/cuda-downloads (preferred)
# or, if your Debian/Ubuntu is new enough, use the packages
# sudo apt install nvidia-cuda-dev libcub-dev
# compile ImpactX with -DImpactX_MPI=OFF
# for pip, use: export IMPACTX_MPI=OFF
Note
Preparation: make sure you work with up-to-date Python tooling.
python3 -m pip install -U pip
python3 -m pip install -U build packaging setuptools wheel pytest
python3 -m pip install -U -r examples/requirements.txt
Compile
From the base of the ImpactX source directory, execute:
# find dependencies & configure
# see additional options below, e.g.
# -DCMAKE_INSTALL_PREFIX=$HOME/sw/impactX
cmake -S . -B build -DImpactX_PYTHON=ON
# compile, here we use four threads
cmake --build build -j 4
That’s all! ImpactX binaries are now in build/bin/
.
Most people execute these binaries directly or copy them out.
If you want to install the executables in a programmatic way, run this:
# for default install paths, you will need administrator rights, e.g. with sudo:
# this installs the application
cmake --build build --target install
# this installs the Python bindings via "python3 -m pip install ..."
cmake --build build --target pip_install -j 4
You can inspect and modify build options after running cmake -S . -B build
with either
ccmake build
or by adding arguments with -D<OPTION>=<VALUE>
to the first CMake call, e.g.:
cmake -S . -B build -DImpactX_PYTHON=ON -DImpactX_COMPUTE=CUDA
That’s it! You can now run a first example.
Developers could now change the ImpactX source code and then call the install lines again to refresh the installation.
Tip
If you do not develop with a user-level package manager, e.g., because you rely on a HPC system’s environment modules, then consider to set up a virtual environment via Python venv.
Otherwise, without a virtual environment, you likely need to add the CMake option -DPYINSTALLOPTIONS="--user"
.
Build Options
CMake Option |
Default & Values |
Description |
---|---|---|
|
ON/OFF |
Build tests |
|
RelWithDebInfo/Release/Debug |
Type of build, symbols & optimizations |
|
system-dependent path |
Install path prefix |
|
ON/OFF |
Print all compiler commands to the terminal during build |
|
ON/OFF |
Build the ImpactX executable application |
|
NOACC/OMP/CUDA/SYCL/HIP |
On-node, accelerated computing backend |
|
ON/OFF |
Compile ImpactX with interprocedural optimization (aka LTO) |
|
ON/OFF |
Multi-node support (message-passing) |
|
ON/OFF |
MPI thread-multiple support, i.e. for |
|
ON/OFF |
openPMD I/O (HDF5, ADIOS) |
|
SINGLE/DOUBLE |
Floating point precision (single/double) |
|
ON/OFF |
Python bindings |
|
(newest found) |
Path to Python executable |
ImpactX can be configured in further detail with options from AMReX, which are documented in the AMReX manual.
Developers might be interested in additional options that control dependencies of ImpactX. By default, the most important dependencies of ImpactX are automatically downloaded for convenience:
CMake Option |
Default & Values |
Description |
---|---|---|
|
ON/OFF |
Build shared libraries for dependencies |
|
ON/OFF |
Search and use CCache to speed up rebuilds. |
|
None |
Path to ABLASTR source directory (preferred if set) |
|
|
Repository URI to pull and build ABLASTR from |
|
we set and maintain a compatible commit |
Repository branch for |
|
ON/OFF |
Needs a pre-installed ABLASTR library if set to |
|
None |
Path to AMReX source directory (preferred if set) |
|
|
Repository URI to pull and build AMReX from |
|
we set and maintain a compatible commit |
Repository branch for |
|
ON/OFF |
Needs a pre-installed AMReX library if set to |
|
None |
Path to openPMD-api source directory (preferred if set) |
|
|
Repository URI to pull and build openPMD-api from |
|
we set and maintain a compatible commit |
Repository branch for |
|
ON/OFF |
Needs a pre-installed openPMD-api library if set to |
|
None |
Path to AMReX source directory (preferred if set) |
|
|
Repository URI to pull and build pyAMReX from |
|
we set and maintain a compatible commit |
Repository branch for |
|
ON/OFF |
Needs a pre-installed pyAMReX module if set to |
|
ON/OFF |
Build Python w/ interprocedural/link optimization (IPO/LTO) |
|
None |
Path to pybind11 source directory (preferred if set) |
|
|
Repository URI to pull and build pybind11 from |
|
we set and maintain a compatible commit |
Repository branch for |
|
ON/OFF |
Needs a pre-installed pybind11 library if set to |
For example, one can also build against a local AMReX copy.
Assuming AMReX’ source is located in $HOME/src/amrex
, add the cmake
argument -DImpactX_amrex_src=$HOME/src/amrex
.
Relative paths are also supported, e.g. -DImpactX_amrex_src=../amrex
.
Or build against an AMReX feature branch of a colleague.
Assuming your colleague pushed AMReX to https://github.com/WeiqunZhang/amrex/
in a branch new-feature
then pass to cmake
the arguments: -DImpactX_amrex_repo=https://github.com/WeiqunZhang/amrex.git -DImpactX_amrex_branch=new-feature
.
If you want to develop against local versions of ABLASTR (from WarpX) and AMReX at the same time, pass for instance -DImpactX_ablastr_src=$HOME/src/warpx -DImpactX_amrex_src=$HOME/src/amrex
.
You can speed up the install further if you pre-install these dependencies, e.g. with a package manager.
Set -DImpactX_<dependency-name>_internal=OFF
and add installation prefix of the dependency to the environment variable CMAKE_PREFIX_PATH.
Please see the introduction to CMake if this sounds new to you.
If you re-compile often, consider installing the Ninja build system.
Pass -G Ninja
to the CMake configuration call to speed up parallel compiles.
Configure Your Compiler
If you don’t want to use your default compiler, you can set the following environment variables. For example, using a Clang/LLVM:
export CC=$(which clang)
export CXX=$(which clang++)
If you also want to select a CUDA compiler:
export CUDACXX=$(which nvcc)
export CUDAHOSTCXX=$(which clang++)
Note
Please clean your build directory with rm -rf build/
after changing the compiler.
Now call cmake -S . -B build
(+ further options) again to re-initialize the build configuration.
Run
The ImpactX Python bindings, which provide the imports impactx
and amrex
(from pyAMReX), are automatically packaged and installed when calling the pip_install
CMake target.
An executable ImpactX application binary with the current compile-time options encoded in its file name will be created in build/bin/
.
Additionally, a symbolic link named impactx
can be found in that directory, which points to the last built ImpactX executable.
HPC
On selected high-performance computing (HPC) systems, ImpactX has documented or even pre-build installation routines. Follow the guide here instead of the generic installation routines for optimal stability and best performance.
impactx.profile
Use a impactx.profile
file to set up your software environment without colliding with other software.
Ideally, store that file directly in your $HOME/
and source it after connecting to the machine:
source $HOME/impactx.profile
We list example impactx.profile
files below, which can be used to set up ImpactX on various HPC systems.
HPC Systems
Perlmutter (NERSC)
The Perlmutter cluster is located at NERSC.
Introduction
If you are new to this system, please see the following resources:
Batch system: Slurm
-
$HOME
: per-user directory, use only for inputs, source and scripts; backed up (40GB)${CFS}/m3239/
: community file system for users in the projectm3239
(or equivalent); moderate performance (20TB default)$PSCRATCH
: per-user production directory; very fast for parallel jobs; purged every 8 weeks (20TB default)
Preparation
Use the following commands to download the ImpactX source code:
git clone https://github.com/ECP-WarpX/impactx.git $HOME/src/impactx
On Perlmutter, you can run either on GPU nodes with fast A100 GPUs (recommended) or CPU nodes.
We use system software modules, add environment hints and further dependencies via the file $HOME/perlmutter_gpu_impactx.profile
.
Create it now:
cp $HOME/src/impactx/docs/source/install/hpc/perlmutter-nersc/perlmutter_gpu_impactx.profile.example $HOME/perlmutter_gpu_impactx.profile
Edit the 2nd line of this script, which sets the export proj=""
variable.
Perlmutter GPU projects must end in ..._g
.
For example, if you are member of the project m3239
, then run nano $HOME/perlmutter_gpu_impactx.profile
and edit line 2 to read:
export proj="m3239_g"
Exit the nano
editor with Ctrl
+ O
(save) and then Ctrl
+ X
(exit).
Important
Now, and as the first step on future logins to Perlmutter, activate these environment settings:
source $HOME/perlmutter_gpu_impactx.profile
Finally, since Perlmutter does not yet provide software modules for some of our dependencies, install them once:
bash $HOME/src/impactx/docs/source/install/hpc/perlmutter-nersc/install_gpu_dependencies.sh
source ${CFS}/${proj%_g}/${USER}/sw/perlmutter/gpu/venvs/impactx/bin/activate
We use system software modules, add environment hints and further dependencies via the file $HOME/perlmutter_cpu_impactx.profile
.
Create it now:
cp $HOME/src/impactx/docs/source/install/hpc/perlmutter-nersc/perlmutter_cpu_impactx.profile.example $HOME/perlmutter_cpu_impactx.profile
Edit the 2nd line of this script, which sets the export proj=""
variable.
For example, if you are member of the project m3239
, then run nano $HOME/perlmutter_cpu_impactx.profile
and edit line 2 to read:
export proj="m3239"
Exit the nano
editor with Ctrl
+ O
(save) and then Ctrl
+ X
(exit).
Important
Now, and as the first step on future logins to Perlmutter, activate these environment settings:
source $HOME/perlmutter_cpu_impactx.profile
Finally, since Perlmutter does not yet provide software modules for some of our dependencies, install them once:
bash $HOME/src/impactx/docs/source/install/hpc/perlmutter-nersc/install_cpu_dependencies.sh
source ${CFS}/${proj}/${USER}/sw/perlmutter/cpu/venvs/impactx/bin/activate
Compilation
Use the following cmake commands to compile the application executable:
cd $HOME/src/impactx
rm -rf build_pm_gpu
cmake -S . -B build_pm_gpu -DImpactX_COMPUTE=CUDA
cmake --build build_pm_gpu -j 16
The ImpactX application executables are now in $HOME/src/impactx/build_pm_gpu/bin/
.
Additionally, the following commands will install ImpactX as a Python module:
cd $HOME/src/impactx
rm -rf build_pm_gpu_py
cmake -S . -B build_pm_gpu_py -DImpactX_COMPUTE=CUDA -DImpactX_APP=OFF -DImpactX_PYTHON=ON
cmake --build build_pm_gpu_py -j 16 --target pip_install
cd $HOME/src/impactx
rm -rf build_pm_cpu
cmake -S . -B build_pm_cpu -DImpactX_COMPUTE=OMP
cmake --build build_pm_cpu -j 16
The ImpactX application executables are now in $HOME/src/impactx/build_pm_cpu/bin/
.
Additionally, the following commands will install ImpactX as a Python module:
rm -rf build_pm_cpu_py
cmake -S . -B build_pm_cpu_py -DImpactX_COMPUTE=OMP -DImpactX_APP=OFF -DImpactX_PYTHON=ON
cmake --build build_pm_cpu_py -j 16 --target pip_install
Now, you can submit Perlmutter compute jobs for ImpactX Python (PICMI) scripts (example scripts).
Or, you can use the ImpactX executables to submit Perlmutter jobs (example inputs).
For executables, you can reference their location in your job script or copy them to a location in $PSCRATCH
.
Update ImpactX & Dependencies
If you already installed ImpactX in the past and want to update it, start by getting the latest source code:
cd $HOME/src/impactx
# read the output of this command - does it look ok?
git status
# get the latest ImpactX source code
git fetch
git pull
# read the output of these commands - do they look ok?
git status
git log # press q to exit
And, if needed,
update the perlmutter_gpu_impactx.profile or perlmutter_cpu_impactx files,
log out and into the system, activate the now updated environment profile as usual,
As a last step, clean the build directory rm -rf $HOME/src/impactx/build_pm_*
and rebuild ImpactX.
Running
The batch script below can be used to run a ImpactX simulation on multiple nodes (change -N
accordingly) on the supercomputer Perlmutter at NERSC.
This partition as up to 1536 nodes.
Replace descriptions between chevrons <>
by relevant values, for instance <input file>
could be plasma_mirror_inputs
.
Note that we run one MPI rank per GPU.
$HOME/src/impactx/docs/source/install/hpc/perlmutter-nersc/perlmutter_gpu.sbatch
.#!/bin/bash -l
# Copyright 2021-2023 Axel Huebl, Kevin Gott
#
# This file is part of ImpactX.
#
# License: BSD-3-Clause-LBNL
#SBATCH -t 00:10:00
#SBATCH -N 2
#SBATCH -J ImpactX
# note: <proj> must end on _g
#SBATCH -A <proj>
#SBATCH -q regular
# A100 40GB (most nodes)
#SBATCH -C gpu
# A100 80GB (256 nodes)
#S BATCH -C gpu&hbm80g
#SBATCH --exclusive
# ideally single:1, but NERSC cgroups issue
#SBATCH --gpu-bind=none
#SBATCH --ntasks-per-node=4
#SBATCH --gpus-per-node=4
#SBATCH -o ImpactX.o%j
#SBATCH -e ImpactX.e%j
# executable & inputs file or python interpreter & PICMI script here
EXE=./impactx
INPUTS=inputs
# pin to closest NIC to GPU
export MPICH_OFI_NIC_POLICY=GPU
# threads for OpenMP and threaded compressors per MPI rank
# note: 16 avoids hyperthreading (32 virtual cores, 16 physical)
export SRUN_CPUS_PER_TASK=16
# GPU-aware MPI optimizations
GPU_AWARE_MPI="amrex.use_gpu_aware_mpi=1"
# CUDA visible devices are ordered inverse to local task IDs
# Reference: nvidia-smi topo -m
srun --cpu-bind=cores bash -c "
export CUDA_VISIBLE_DEVICES=\$((3-SLURM_LOCALID));
${EXE} ${INPUTS} ${GPU_AWARE_MPI}" \
> output.txt
To run a simulation, copy the lines above to a file perlmutter_gpu.sbatch
and run
sbatch perlmutter_gpu.sbatch
to submit the job.
Perlmutter has 256 nodes that provide 80 GB HBM per A100 GPU.
In the A100 (40GB) batch script, replace -C gpu
with -C gpu&hbm80g
to use these large-memory GPUs.
The Perlmutter CPU partition as up to 3072 nodes, each with 2x AMD EPYC 7763 CPUs.
$HOME/src/impactx/docs/source/install/hpc/perlmutter-nersc/perlmutter_cpu.sbatch
.#!/bin/bash -l
# Copyright 2021-2023 ImpactX
#
# This file is part of ImpactX.
#
# Authors: Axel Huebl
# License: BSD-3-Clause-LBNL
#SBATCH -t 00:10:00
#SBATCH -N 2
#SBATCH -J ImpactX
#SBATCH -A <proj>
#SBATCH -q regular
#SBATCH -C cpu
#SBATCH --ntasks-per-node=16
#SBATCH --exclusive
#SBATCH -o ImpactX.o%j
#SBATCH -e ImpactX.e%j
# executable & inputs file or python interpreter & PICMI script here
EXE=./impactx
INPUTS=inputs_small
# each CPU node on Perlmutter (NERSC) has 64 hardware cores with
# 2x Hyperthreading/SMP
# https://en.wikichip.org/wiki/amd/epyc/7763
# https://www.amd.com/en/products/cpu/amd-epyc-7763
# Each CPU is made up of 8 chiplets, each sharing 32MB L3 cache.
# This will be our MPI rank assignment (2x8 is 16 ranks/node).
# threads for OpenMP and threaded compressors per MPI rank
export SRUN_CPUS_PER_TASK=16 # 8 cores per chiplet, 2x SMP
export OMP_PLACES=threads
export OMP_PROC_BIND=spread
srun --cpu-bind=cores \
${EXE} ${INPUTS} \
> output.txt
Post-Processing
For post-processing, most users use Python via NERSC’s Jupyter service (documentation).
As a one-time preparatory setup, log into Perlmutter via SSH and do not source the ImpactX profile script above. Create your own Conda environment and Jupyter kernel for post-processing:
module load python
conda config --set auto_activate_base false
# create conda environment
rm -rf $HOME/.conda/envs/impactx-pm-postproc
conda create --yes -n impactx-pm-postproc -c conda-forge mamba conda-libmamba-solver
conda activate impactx-pm-postproc
conda config --set solver libmamba
mamba install --yes -c conda-forge python ipykernel ipympl matplotlib numpy pandas yt openpmd-viewer openpmd-api h5py fast-histogram dask dask-jobqueue pyarrow
# create Jupyter kernel
rm -rf $HOME/.local/share/jupyter/kernels/impactx-pm-postproc/
python -m ipykernel install --user --name impactx-pm-postproc --display-name ImpactX-PM-PostProcessing
echo -e '#!/bin/bash\nmodule load python\nsource activate impactx-pm-postproc\nexec "$@"' > $HOME/.local/share/jupyter/kernels/impactx-pm-postproc/kernel-helper.sh
chmod a+rx $HOME/.local/share/jupyter/kernels/impactx-pm-postproc/kernel-helper.sh
KERNEL_STR=$(jq '.argv |= ["{resource_dir}/kernel-helper.sh"] + .' $HOME/.local/share/jupyter/kernels/impactx-pm-postproc/kernel.json | jq '.argv[1] = "python"')
echo ${KERNEL_STR} | jq > $HOME/.local/share/jupyter/kernels/impactx-pm-postproc/kernel.json
exit
When opening a Jupyter notebook on https://jupyter.nersc.gov, just select ImpactX-PM-PostProcessing
from the list of available kernels on the top right of the notebook.
Additional software can be installed later on, e.g., in a Jupyter cell using !mamba install -y -c conda-forge ...
.
Software that is not available via conda can be installed via !python -m pip install ...
.
Tip
Your HPC system is not in the list? Our instructions are nearly identical to installing WarpX, documented here.
Also, please do not hesitate to open an issue and together we can document how to run on your preferred system!
Usage
Run ImpactX
In order to run a new simulation:
create a new directory, where the simulation will be run
make sure the ImpactX executable is either copied into this directory or in your
PATH
environment variableadd an inputs file and on HPC systems a submission script to the directory
run
1. Run Directory
On Linux/macOS, this is as easy as this
mkdir -p <run_directory>
Where <run_directory>
by the actual path to the run directory.
2. Executable
If you installed ImpactX with a package manager, a impactx
-prefixed executable will be available as a regular system command to you.
Depending on the choosen build options, the name is suffixed with more details.
Try it like this:
impactx<TAB>
Hitting the <TAB>
key will suggest available ImpactX executables as found in your PATH
environment variable.
If you compiled the code yourself, the ImpactX executable is stored in the source folder under build/bin
.
We also create a symbolic link that is just called impactx
that points to the last executable you built, which can be copied, too.
Copy the executable to this directory:
cp build/bin/<impactx_executable> <run_directory>/
where <impactx_executable>
should be replaced by the actual name of the executable (see above) and <run_directory>
by the actual path to the run directory.
3. Inputs
Add an input file in the directory (see examples and parameters). This file contains the numerical and physical parameters that define the situation to be simulated.
On HPC systems, also copy and adjust a submission script that allocated computing nodes for you. Please reach out to us if you need help setting up a template that runs with ideal performance.
4. Run
Run the executable, e.g. with MPI:
cd <run_directory>
# run with an inputs file:
mpirun -np <n_ranks> ./impactx <input_file>
or
# run with a PICMI input script:
mpirun -np <n_ranks> python <python_script>
Here, <n_ranks>
is the number of MPI ranks used, and <input_file>
is the name of the input file (<python_script>
is the name of the PICMI script).
Note that the actual executable might have a longer name, depending on build options.
We used the copied executable in the current directory (./
); if you installed with a package manager, skip the ./
because ImpactX is in your PATH
.
On an HPC system, you would instead submit the job script at this point, e.g. sbatch <submission_script>
(SLURM on Perlmutter/NERSC) or bsub <submission_script>
(LSF on Summit/OLCF).
Tip
In the next sections, we will explain parameters of the <input_file>
.
You can overwrite all parameters inside this file also from the command line, e.g.:
mpirun -np 4 ./impactx <input_file> max_step=10 amr.n_cell=64 64 128
Examples
This section allows you to download input files that correspond to different physical situations or test different code features.
FODO Cell
Stable FODO cell with a zero-current phase advance of 67.8 degrees.
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 2.82161941\) m
\(\alpha_\mathrm{x} = -1.59050035\)
\(\beta_\mathrm{y} = 2.82161941\) m
\(\alpha_\mathrm{y} = 1.59050035\)
We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm.
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 \(\lambda_x\), \(\lambda_y\), \(\lambda_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.py
orImpactX executable using an input file:
impactx input_fodo.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/fodo/run_fodo.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 = False
# sim.diagnostics = False # benchmarking
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 = 2.0e3 # 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.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 = 25 # number of slices per ds in the element
fodo = [
monitor,
elements.Drift(ds=0.25, nslice=ns),
monitor,
elements.Quad(ds=1.0, k=1.0, nslice=ns),
monitor,
elements.Drift(ds=0.5, nslice=ns),
monitor,
elements.Quad(ds=1.0, k=-1.0, nslice=ns),
monitor,
elements.Drift(ds=0.25, nslice=ns),
monitor,
]
# assign a fodo segment
sim.lattice.extend(fodo)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/fodo/input_fodo.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 3.9984884770e-5
beam.lambdaY = 3.9984884770e-5
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.6623538760e-5
beam.lambdaPy = 2.6623538760e-5
beam.lambdaPt = 2.0e-3
beam.muxpx = -0.846574929020762
beam.muypy = 0.846574929020762
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 = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_fodo.py
examples/fodo/analysis_fodo.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 = 2.2 * 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 = 2.2 * 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.4790118496224206e-005,
7.5357525169680140e-005,
9.9775879288128088e-004,
1.9959539836392703e-009,
2.0175014668882125e-009,
2.0013820380883801e-006,
],
rtol=rtol,
atol=atol,
)
Visualize
You can run the following script to visualize the beam evolution over time:
Script plot_fodo.py
examples/fodo/plot_fodo.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import argparse
import glob
import re
import matplotlib.pyplot as plt
import openpmd_api as io
import pandas as pd
from matplotlib.ticker import MaxNLocator
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)
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')
# options to run this script
parser = argparse.ArgumentParser(description="Plot the FODO benchmark.")
parser.add_argument(
"--save-png", action="store_true", help="non-interactive run: save to PNGs"
)
args = parser.parse_args()
# 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()
ref_particle = read_time_series("diags/ref_particle.*")
# scaling to units
millimeter = 1.0e3 # m->mm
mrad = 1.0e3 # ImpactX uses "static units": momenta are normalized by the magnitude of the momentum of the reference particle p0: px/p0 (rad)
# mm_mrad = 1.e6
nm_rad = 1.0e9
# select a single particle by id
# particle_42 = beam[beam["id"] == 42]
# print(particle_42)
# steps & corresponding z
steps = list(series.iterations)
z = list(
map(lambda step: ref_particle[ref_particle["step"] == step].z.values[0], steps)
)
# print(f"z={z}")
# beam transversal size & emittance over steps
moments = list(
map(
lambda step: (
step,
get_moments(series.iterations[step].particles["beam"].to_df()),
),
steps,
)
)
# print(moments)
sigx = list(map(lambda step_val: step_val[1][0] * millimeter, moments))
sigy = list(map(lambda step_val: step_val[1][1] * millimeter, moments))
emittance_x = list(map(lambda step_val: step_val[1][3] * nm_rad, moments))
emittance_y = list(map(lambda step_val: step_val[1][4] * nm_rad, moments))
# print(sigx, sigy)
# print beam transversal size over steps
f = plt.figure(figsize=(9, 4.8))
ax1 = f.gca()
im_sigx = ax1.plot(z, sigx, label=r"$\lambda_x$")
im_sigy = ax1.plot(z, sigy, label=r"$\lambda_y$")
ax2 = ax1.twinx()
ax2.set_prop_cycle(None) # reset color cycle
im_emittance_x = ax2.plot(z, emittance_x, ":", label=r"$\epsilon_x$")
im_emittance_y = ax2.plot(z, emittance_y, ":", label=r"$\epsilon_y$")
ax1.legend(
handles=im_sigx + im_sigy + im_emittance_x + im_emittance_y, loc="lower center"
)
ax1.set_xlabel(r"$z$ [m]")
ax1.set_ylabel(r"$\lambda_{x,y}$ [mm]")
# ax2.set_ylabel(r"$\epsilon_{x,y}$ [mm-mrad]")
ax2.set_ylabel(r"$\epsilon_{x,y}$ [nm]")
ax2.set_ylim([1.5, 2.5])
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.tight_layout()
if args.save_png:
plt.savefig("fodo_lambda.png")
else:
plt.show()
# beam transversal scatter plot over steps
num_plots_per_row = len(steps)
fig, axs = plt.subplots(
3, num_plots_per_row, figsize=(9, 4.8), sharex="row", sharey="row"
)
ncol_ax = -1
for step in steps:
# plot initial distribution & at exit of each element
ncol_ax += 1
# x-y
ax = axs[(0, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_x.multiply(millimeter),
beam_at_step.position_y.multiply(millimeter),
s=0.01,
)
ax.set_title(f"$z={z[ncol_ax]}$ [m]")
ax.set_xlabel(r"$x$ [mm]")
# x-px
ax = axs[(1, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_x.multiply(millimeter),
beam_at_step.momentum_x.multiply(mrad),
s=0.01,
)
ax.set_xlabel(r"$x$ [mm]")
# y-py
ax = axs[(2, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_y.multiply(millimeter),
beam_at_step.momentum_y.multiply(mrad),
s=0.01,
)
ax.set_xlabel(r"$y$ [mm]")
axs[(0, 0)].set_ylabel(r"$y$ [mm]")
axs[(1, 0)].set_ylabel(r"$p_x$ [mrad]")
axs[(2, 0)].set_ylabel(r"$p_y$ [mrad]")
plt.tight_layout()
if args.save_png:
plt.savefig("fodo_scatter.png")
else:
plt.show()

FODO transversal beam width and emittance evolution

FODO transversal beam width and phase space evolution
Chicane
Berlin-Zeuthen magnetic bunch compression chicane: https://www.desy.de/csr/
All parameters can be found online. A 5 GeV electron bunch with normalized transverse rms emittance of 1 um undergoes longitudinal compression by a factor of 10 in a standard 4-bend chicane.
The emittances should be preserved, and the rms pulse length should decrease by the compression factor (10).
In this test, the initial and final values of \(\lambda_x\), \(\lambda_y\), \(\lambda_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_chicane.py
orImpactX executable using an input file:
impactx input_chicane.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/chicane/run_chicane.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 -*-
from impactx import ImpactX, distribution, elements
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 5 GeV electron beam with an initial
# normalized transverse rms emittance of 1 um
kin_energy_MeV = 5.0e3 # 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.Waterbag(
lambdaX=2.2951017632e-5,
lambdaY=1.3084093142e-5,
lambdaT=5.5555553e-8,
lambdaPx=1.598353425e-6,
lambdaPy=2.803697378e-6,
lambdaPt=2.000000000e-6,
muxpx=0.933345606203060,
muypy=0.933345606203060,
mutpt=0.999999961419755,
)
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
rc = 10.35 # bend radius (meters)
psi = 0.048345620280243 # pole face rotation angle (radians)
# Drift elements
dr1 = elements.Drift(ds=5.0058489435, nslice=ns)
dr2 = elements.Drift(ds=1.0, nslice=ns)
dr3 = elements.Drift(ds=2.0, nslice=ns)
# Bend elements
sbend1 = elements.Sbend(ds=0.50037, rc=-rc, nslice=ns)
sbend2 = elements.Sbend(ds=0.50037, rc=rc, nslice=ns)
# Dipole Edge Focusing elements
dipedge1 = elements.DipEdge(psi=-psi, rc=-rc, g=0.0, K2=0.0)
dipedge2 = elements.DipEdge(psi=psi, rc=rc, g=0.0, K2=0.0)
lattice_half = [sbend1, dipedge1, dr1, dipedge2, sbend2]
# assign a segment with the first half of the lattice
sim.lattice.append(monitor)
sim.lattice.extend(lattice_half)
sim.lattice.append(dr2)
lattice_half.reverse()
# extend the lattice by a reversed half
sim.lattice.extend(lattice_half)
sim.lattice.append(dr3)
sim.lattice.append(monitor)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/chicane/input_chicane.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 5.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 2.2951017632e-5
beam.lambdaY = 1.3084093142e-5
beam.lambdaT = 5.5555553e-8
beam.lambdaPx = 1.598353425e-6
beam.lambdaPy = 2.803697378e-6
beam.lambdaPt = 2.000000000e-6
beam.muxpx = 0.933345606203060
beam.muypy = 0.933345606203060
beam.mutpt = 0.999999961419755
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor sbend1 dipedge1 drift1 dipedge2 sbend2 drift2 \
sbend2 dipedge2 drift1 dipedge1 sbend1 drift3 monitor
lattice.nslice = 25
sbend1.type = sbend
sbend1.ds = 0.50037 # projected length 0.5 m, angle 2.77 deg
sbend1.rc = -10.35
drift1.type = drift
drift1.ds = 5.0058489435 # projected length 5 m
sbend2.type = sbend
sbend2.ds = 0.50037 # projected length 0.5 m, angle 2.77 deg
sbend2.rc = 10.35
drift2.type = drift
drift2.ds = 1.0
drift3.type = drift
drift3.ds = 2.0
dipedge1.type = dipedge # dipole edge focusing
dipedge1.psi = -0.048345620280243
dipedge1.rc = -10.35
dipedge1.g = 0.0
dipedge1.K2 = 0.0
dipedge2.type = dipedge
dipedge2.psi = 0.048345620280243
dipedge2.rc = 10.35
dipedge2.g = 0.0
dipedge2.K2 = 0.0
monitor.type = beam_monitor
monitor.backend = h5
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_chicane.py
examples/chicane/analysis_chicane.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 = 2.2 * 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],
[
6.4214719960819659e-005,
3.6603372435649773e-005,
1.9955175623579313e-004,
1.0198263116327677e-010,
1.0308359092878036e-010,
4.0035161705244885e-010,
],
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 = 2.2 * 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],
[
2.3928429374387210e-005,
8.4424535301423173e-005,
1.9976426324802290e-005,
1.0198281373761584e-010,
1.0308356090529235e-010,
4.0027996099961315e-010,
],
rtol=rtol,
atol=atol,
)
Visualize
You can run the following script to visualize the beam evolution over time:
Script plot_chicane.py
examples/chicane/plot_chicane.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import argparse
import glob
import re
import matplotlib.pyplot as plt
import openpmd_api as io
import pandas as pd
from matplotlib.ticker import MaxNLocator
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)
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')
# options to run this script
parser = argparse.ArgumentParser(description="Plot the chicane benchmark.")
parser.add_argument(
"--save-png", action="store_true", help="non-interactive run: save to PNGs"
)
args = parser.parse_args()
# 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()
ref_particle = read_time_series("diags/ref_particle.*")
# scaling to units
millimeter = 1.0e3 # m->mm
# for "t": the time coordinate is scaled by c, and therefore has units of length (m) by default, so we can label the axis ct (mm)
mrad = 1.0e3 # ImpactX uses "static units": momenta are normalized by the magnitude of the momentum of the reference particle p0: px/p0 (rad)
# mm_mrad = 1.e6
nm_rad = 1.0e9
# select a single particle by id
# particle_42 = beam[beam["id"] == 42]
# print(particle_42)
# steps & corresponding z
steps = list(series.iterations)
z = list(
map(lambda step: ref_particle[ref_particle["step"] == step].z.values[0], steps)
)
x = list(
map(lambda step: ref_particle[ref_particle["step"] == step].x.values[0], steps)
)
# print(f"z={z}")
# beam transversal size & emittance over steps
moments = list(
map(
lambda step: (
step,
get_moments(series.iterations[step].particles["beam"].to_df()),
),
steps,
)
)
# print(moments)
sigx = list(map(lambda step_val: step_val[1][0] * millimeter, moments))
sigt = list(map(lambda step_val: step_val[1][2] * millimeter, moments))
emittance_x = list(map(lambda step_val: step_val[1][3] * nm_rad, moments))
emittance_t = list(map(lambda step_val: step_val[1][5] * nm_rad, moments))
# print(sigx, sigt)
# print beam transversal size over steps
f, axs = plt.subplots(
2, 1, figsize=(9, 4.8), sharex=True, gridspec_kw={"height_ratios": [1, 2]}
)
ax0 = axs[0]
im_xz = ax0.plot(z, x, "--", lw=3, label=r"$x$")
ax0.legend(loc="upper right")
ax0.set_ylim([0, None])
ax0.set_ylabel(r"$x$ [m]")
ax1 = axs[1]
im_sigx = ax1.plot(z, sigx, label=r"$\lambda_x$")
im_sigt = ax1.plot(z, sigt, label=r"$\lambda_t$")
ax2 = ax1.twinx()
ax2.set_prop_cycle(None) # reset color cycle
im_emittance_x = ax2.plot(z, emittance_x, ":", label=r"$\epsilon_x$")
im_emittance_t = ax2.plot(z, emittance_t, ":", label=r"$\epsilon_t$")
ax1.legend(
handles=im_sigx + im_sigt + im_emittance_x + im_emittance_t, loc="upper right"
)
ax1.set_xlabel(r"$z$ [m]")
ax1.set_ylabel(r"$\lambda_{x,t}$ [mm]")
# ax2.set_ylabel(r"$\epsilon_{x,y}$ [mm-mrad]")
ax2.set_ylabel(r"$\epsilon_{x,t}$ [nm]")
ax1.set_ylim([0, None])
ax2.set_ylim([0, None])
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.tight_layout()
if args.save_png:
plt.savefig("chicane_lambda.png")
else:
plt.show()
# beam transversal scatter plot over steps
num_plots_per_row = len(steps)
fig, axs = plt.subplots(
4, num_plots_per_row, figsize=(9, 6.4), sharex="row", sharey="row"
)
ncol_ax = -1
for step in steps:
# plot initial distribution & at exit of each element
ncol_ax += 1
# t-pt
ax = axs[(0, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_t.multiply(millimeter),
beam_at_step.momentum_t.multiply(mrad),
s=0.01,
)
ax.set_xlabel(r"$ct$ [mm]")
z_unit = ""
if ncol_ax == num_plots_per_row - 1:
z_unit = " [m]"
ax.set_title(f"$z={z[ncol_ax]:.1f}${z_unit}")
# x-px
ax = axs[(1, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_x.multiply(millimeter),
beam_at_step.momentum_x.multiply(mrad),
s=0.01,
)
ax.set_xlabel(r"$x$ [mm]")
# t-x
ax = axs[(2, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_t.multiply(millimeter),
beam_at_step.position_x.multiply(millimeter),
s=0.01,
)
ax.set_xlabel(r"$ct$ [mm]")
# t-px
ax = axs[(3, ncol_ax)]
beam_at_step = series.iterations[step].particles["beam"].to_df()
ax.scatter(
beam_at_step.position_t.multiply(millimeter),
beam_at_step.momentum_x.multiply(mrad),
s=0.01,
)
ax.set_xlabel(r"$ct$ [mm]")
axs[(0, 0)].set_ylabel(r"$p_t$ [mrad]")
axs[(1, 0)].set_ylabel(r"$p_x$ [mrad]")
axs[(2, 0)].set_ylabel(r"$x$ [mm]")
axs[(3, 0)].set_ylabel(r"$p_x$ [mrad]")
plt.tight_layout()
if args.save_png:
plt.savefig("chicane_scatter.png")
else:
plt.show()

(top) Chicane floorplan. (bottom) Chicane beam width and emittance evolution.

Chicane beam width and emittance evolution
Constant Focusing Channel
Stationary beam in a constant focusing channel (without space charge).
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 1.0\) m
\(\alpha_\mathrm{x} = 0.0\)
\(\beta_\mathrm{y} = 1.0\) m
\(\alpha_\mathrm{y} = 0.0\)
We use a 2 GeV proton beam with initial unnormalized rms emittance of 1 um. The longitudinal beam parameters are chosen so that the bunch has radial symmetry when viewed in the beam rest frame.
The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. This fact is independent of the length of the channel. This is tested using the second moments of the distribution.
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_cfchannel.py
orImpactX executable using an input file:
impactx input_cfchannel.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/cfchannel/run_cfchannel.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 -*-
from impactx import ImpactX, distribution, elements
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 2 GeV proton beam with an initial
# normalized transverse rms emittance of 1 um
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.particle_container().ref_particle()
ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.0e-3,
lambdaY=1.0e-3,
lambdaT=3.369701494258956e-4,
lambdaPx=1.0e-3,
lambdaPy=1.0e-3,
lambdaPt=2.9676219145931020e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.extend(
[
monitor,
elements.ConstF(ds=2.0, kx=1.0, ky=1.0, kt=1.0),
monitor,
]
)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/cfchannel/input_cfchannel.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.0e-3
beam.lambdaY = 1.0e-3
beam.lambdaT = 3.369701494258956e-4
beam.lambdaPx = 1.0e-3
beam.lambdaPy = 1.0e-3
beam.lambdaPt = 2.9676219145931020e-3
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor constf1 monitor
monitor.type = beam_monitor
monitor.backend = h5
constf1.type = constf
constf1.ds = 2.0
constf1.kx = 1.0
constf1.ky = 1.0
constf1.kt = 1.0
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
Analyze
We run the following script to analyze correctness:
Script analysis_cfchannel.py
examples/cfchannel/analysis_cfchannel.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 # a big number
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],
[
1.0e-003,
1.0e-003,
3.369701494258956e-4,
1.0e-006,
1.0e-006,
1.0e-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 # a big number
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],
[
1.0e-003,
1.0e-003,
3.369701494258956e-4,
1.0e-006,
1.0e-006,
1.0e-006,
],
rtol=rtol,
atol=atol,
)
Constant Focusing Channel with Space Charge
RMS-matched beam in a constant focusing channel with space charge.
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 1.477305\) m
\(\alpha_\mathrm{x} = 0.0\)
\(\beta_\mathrm{y} = 1.477305\) m
\(\alpha_\mathrm{y} = 0.0\)
We use a 2 GeV proton beam with initial unnormalized rms emittance of 1 um. The longitudinal beam parameters are chosen so that the bunch has radial symmetry when viewed in the beam rest frame. The bunch charge is set to 10 nC, resulting in a transverse tune depression ratio of 0.67. The initial distribution used is a 6D waterbag.
The beam second moments should remain nearly unchanged, except for some small emittance growth due to nonlinear space charge. This is tested using the second moments of the distribution.
In this test, the initial and final values of \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), and :math:`
Run
This example can be run as a Python script (python3 run_cfchannel_10nC.py
) or as an app with an input file (impactx input_cfchannel_10nC.in
).
Each can also be prefixed with an MPI executor, such as mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/cfchannel/run_cfchannel_10nC.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 -*-
from impactx import ImpactX, distribution, elements
sim = ImpactX()
# set numerical parameters and IO control
sim.n_cell = [48, 48, 40] # [72, 72, 64] for increased precision
sim.particle_shape = 2 # B-spline order
sim.space_charge = True
sim.prob_relative = [3.0]
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = True
# domain decomposition & space charge mesh
sim.init_grids()
# load a 2 GeV proton beam with an initial
# normalized transverse rms emittance of 1 um
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 1.0e-8 # used with space charge
npart = 10000 # number of macro particles; use 1e5 for increased precision
# 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)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.2154443728379865788e-3,
lambdaY=1.2154443728379865788e-3,
lambdaT=4.0956844276541331005e-4,
lambdaPx=8.2274435782286157175e-4,
lambdaPy=8.2274435782286157175e-4,
lambdaPt=2.4415943602685364584e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
nslice = 50 # use 1e5 for increased precision
# design the accelerator lattice
sim.lattice.extend(
[
monitor,
elements.ConstF(ds=2.0, kx=1.0, ky=1.0, kt=1.0, nslice=nslice),
monitor,
]
)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/cfchannel/input_cfchannel_10nC.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
#beam.npart = 100000 # optional for increased precision
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-8
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.2154443728379865788e-3
beam.lambdaY = 1.2154443728379865788e-3
beam.lambdaT = 4.0956844276541331005e-4
beam.lambdaPx = 8.2274435782286157175e-4
beam.lambdaPy = 8.2274435782286157175e-4
beam.lambdaPt = 2.4415943602685364584e-3
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor constf1 monitor
lattice.nslice = 50
#lattice.nslice = 60 # optional for increased precision
monitor.type = beam_monitor
monitor.backend = h5
constf1.type = constf
constf1.ds = 2.0
constf1.kx = 1.0
constf1.ky = 1.0
constf1.kt = 1.0
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = true
amr.n_cell = 48 48 40
#amr.n_cell = 72 72 64 # optional for increased precision
geometry.prob_relative = 3.0
Analyze
We run the following script to analyze correctness:
Script analysis_cfchannel_10nC.py
examples/cfchannel/analysis_cfchannel_10nC.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],
[
1.2154443728379865788e-003,
1.2154443728379865788e-003,
4.0956844276541331005317e-004,
1.000000000e-006,
1.000000000e-006,
1.000000000e-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 = 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],
[
1.2154443728379865788e-003,
1.2154443728379865788e-003,
4.0956844276541331005317e-004,
1.000000000e-006,
1.000000000e-006,
1.000000000e-006,
],
rtol=rtol,
atol=atol,
)
Expanding Beam in Free Space
A coasting bunch expanding freely in free space under its own space charge.
We use a cold (zero emittance) 250 MeV electron bunch whose initial distribution is a uniformly-populated 3D ball of radius R0 = 1 mm when viewed in the bunch rest frame.
In the laboratory frame, the bunch is a uniformly-populated ellipsoid, which expands to twice its original size. This is tested using the second moments of the distribution.
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.
This test uses mesh-refinement to solve the space charge force. The coarse grid wraps the beam maximum extent by 300%, emulating “open boundary” conditions. The refined grid in level 1 spans 110% of the beam maximum extent. The grid spacing is adaptively adjusted as the beam evolves.
Run
This example can be run either as:
Python script:
python3 run_expanding.py
orImpactX executable using an input file:
impactx input_expanding.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/expanding/run_expanding.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.max_level = 1
sim.n_cell = [16, 16, 20]
sim.blocking_factor_x = [16]
sim.blocking_factor_y = [16]
sim.blocking_factor_z = [4]
sim.particle_shape = 2 # B-spline order
sim.space_charge = True
sim.dynamic_size = True
sim.prob_relative = [3.0, 1.1]
# beam diagnostics
# 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 = 250 # reference energy
bunch_charge_C = 1.0e-9 # used with space charge
npart = 10000 # number of macro particles (outside tests, use 1e5 or more)
# 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.Kurth6D(
lambdaX=4.472135955e-4,
lambdaY=4.472135955e-4,
lambdaT=9.12241869e-7,
lambdaPx=0.0,
lambdaPy=0.0,
lambdaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.extend([monitor, elements.Drift(ds=6.0, nslice=40), monitor])
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/expanding/input_expanding.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000 # outside tests, use 1e5 or more
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = kurth6d
beam.lambdaX = 4.472135955e-4
beam.lambdaY = 4.472135955e-4
beam.lambdaT = 9.12241869e-7
beam.lambdaPx = 0.0
beam.lambdaPy = 0.0
beam.lambdaPt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 monitor
lattice.nslice = 40
drift1.type = drift
drift1.ds = 6.0
monitor.type = beam_monitor
monitor.backend = h5
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = true
# Space charge solver with one MR level
amr.max_level = 1
amr.n_cell = 16 16 20
amr.blocking_factor_x = 16
amr.blocking_factor_y = 16
amr.blocking_factor_z = 4
geometry.prob_relative = 3.0 1.1
# Space charger solver without MR
#amr.max_level = 0
#amr.n_cell = 56 56 48
#geometry.prob_relative = 3.0
Analyze
We run the following script to analyze correctness:
Script analysis_expanding.py
examples/expanding/analysis_expanding.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.4721359550e-004,
4.4721359550e-004,
9.1224186858e-007,
0.0e-006,
0.0e-006,
0.0e-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 = 1.6 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[sigx, sigy, sigt],
[
8.9442719100e-004,
8.9442719100e-004,
1.8244837370e-006,
],
rtol=rtol,
atol=atol,
)
atol = 1.0e-8
rtol = 0.0 # ignored
assert np.allclose(
[emittance_x, emittance_y, emittance_t],
[
0.0,
0.0,
0.0,
],
rtol=rtol,
atol=atol,
)
Kurth Distribution in a Periodic Focusing Channel
Matched Kurth distribution in a periodic focusing channel (without space charge).
The distribution is radially symmetric in (x,y,t) space, and matched to a radially symmetric periodic linear focusing lattice with a phase advance of 121 degrees.
We use a 2 GeV proton beam with initial unnormalized rms emittance of 1 um in all three phase planes.
The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. This is tested using the second moments of the distribution.
In this test, the initial and final values of \(\lambda_x\), \(\lambda_y\), \(\lambda_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_kurth_periodic.py
orImpactX executable using an input file:
impactx input_kurth_periodic.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/kurth/run_kurth_periodic.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 -*-
from impactx import ImpactX, distribution, elements
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 2 GeV proton beam with an initial
# unnormalized rms emittance of 1 um in each
# coordinate plane
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 1.0e-8 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Kurth6D(
lambdaX=1.11e-3,
lambdaY=1.11e-3,
lambdaT=3.74036839224568e-4,
lambdaPx=9.00900900901e-4,
lambdaPy=9.00900900901e-4,
lambdaPt=2.6735334467940146e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
constf1 = elements.ConstF(ds=2.0, kx=0.7, ky=0.7, kt=0.7)
drift1 = elements.Drift(ds=1.0)
sim.lattice.extend([monitor, drift1, constf1, drift1, monitor])
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/kurth/input_kurth_periodic.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-8
beam.particle = proton
beam.distribution = kurth6d
beam.lambdaX = 1.11e-3
beam.lambdaY = 1.11e-3
beam.lambdaT = 3.74036839224568e-4
beam.lambdaPx = 9.00900900901e-4
beam.lambdaPy = 9.00900900901e-4
beam.lambdaPt = 2.6735334467940146e-3
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 constf1 drift1 monitor
monitor.type = beam_monitor
monitor.backend = h5
drift1.type = drift
drift1.ds = 1.0
constf1.type = constf
constf1.ds = 2.0
constf1.kx = 0.7
constf1.ky = 0.7
constf1.kt = 0.7
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
amr.n_cell = 40 40 32
Analyze
We run the following script to analyze correctness:
Script analysis_kurth_periodic.py
examples/kurth/analysis_kurth_periodic.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.7 * 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.11e-03,
1.11e-03,
3.74036839224568e-04,
1.000000000e-006,
1.000000000e-006,
1.000000000e-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 = 1.7 * 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.11e-03,
1.11e-03,
3.74036839224568e-04,
1.000000000e-006,
1.000000000e-006,
1.000000000e-006,
],
rtol=rtol,
atol=atol,
)
Kurth Distribution in a Periodic Focusing Channel with Space Charge
Matched Kurth distribution in a periodic focusing channel with space charge.
The distribution is radially symmetric in (x,y,t) space, and matched to a radially symmetric constant linear focusing.
We use a 2 GeV proton beam with initial unnormalized rms emittance of 1 um in all three phase planes. The bunch charge is set to 10 nC, depressing the phase advance from 121 degrees to 74 degrees.
The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. This is tested using the second moments of the distribution.
In this test, the initial and final values of \(\lambda_x\), \(\lambda_y\), \(\lambda_t\), \(\epsilon_x\), \(\epsilon_y\), and :math:`
Run
This example can be run as a Python script (python3 run_kurth_10nC_periodic.py
) or as an app with an input file (impactx input_kurth_10nC_periodic.in
).
Each can also be prefixed with an MPI executor, such as mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/kurth/run_kurth_10nC_periodic.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 -*-
from impactx import ImpactX, distribution, elements
sim = ImpactX()
# set numerical parameters and IO control
sim.n_cell = [48, 48, 40] # use [72, 72, 72] for increased precision
sim.particle_shape = 2 # B-spline order
sim.space_charge = True
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = True
# domain decomposition & space charge mesh
sim.init_grids()
# load a 2 GeV proton beam with an initial
# unnormalized rms emittance of 1 um in each
# coordinate plane
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 1.0e-8 # used with space charge
npart = 10000 # number of macro particles; use 1e5 for increased precision
# 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)
# particle bunch
distr = distribution.Kurth6D(
lambdaX=1.46e-3,
lambdaY=1.46e-3,
lambdaT=4.9197638312420749e-4,
lambdaPx=6.84931506849e-4,
lambdaPy=6.84931506849e-4,
lambdaPt=2.0326178944803812e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
nslice = 20 # use 30 for increased precision
constf1 = elements.ConstF(ds=2.0, kx=0.7, ky=0.7, kt=0.7, nslice=nslice)
drift1 = elements.Drift(ds=1.0, nslice=nslice)
sim.lattice.extend([monitor, drift1, constf1, drift1, monitor])
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/kurth/input_kurth_10nC_periodic.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
#beam.npart = 100000 #optional for increased precision
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-8
beam.particle = proton
beam.distribution = kurth6d
beam.lambdaX = 1.46e-3
beam.lambdaY = 1.46e-3
beam.lambdaT = 4.9197638312420749e-4
beam.lambdaPx = 6.84931506849e-4
beam.lambdaPy = 6.84931506849e-4
beam.lambdaPt = 2.0326178944803812e-3
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 constf1 drift1 monitor
lattice.nslice = 20
#lattice.nslice = 30 #optional for increased precision
monitor.type = beam_monitor
monitor.backend = h5
drift1.type = drift
drift1.ds = 1.0
constf1.type = constf
constf1.ds = 2.0
constf1.kx = 0.7
constf1.ky = 0.7
constf1.kt = 0.7
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = true
amr.n_cell = 48 48 40
#amr.n_cell = 72 72 72 # optional for increased precision
geometry.prob_relative = 3.0
Analyze
We run the following script to analyze correctness:
Script analysis_kurth_10nC_periodic.py
examples/kurth/analysis_kurth_10nC_periodic.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 = 2.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],
[
1.46e-3,
1.46e-3,
4.9197638312420749e-4,
1.000000000e-006,
1.000000000e-006,
1.000000000e-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 = 2.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],
[
1.46e-3,
1.46e-3,
4.9197638312420749e-4,
1.000000000e-006,
1.000000000e-006,
1.000000000e-006,
],
rtol=rtol,
atol=atol,
)
Quadrupole with Alignment Errors
A 2 GeV proton beam propagates through a single quadrupole with 3 mm horizontal misalignment and 30 degree rotation error.
The first and second moments of the particle distribution before and after the quadrupole should coincide with analytical predictions, to within the level expected due to noise due to statistical sampling.
In this test, the initial and final values of \(\mu_x\), \(\mu_y\), \(\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_alignment.py
orImpactX executable using an input file:
impactx input_alignment.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/alignment/run_alignment.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 = False
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = True
# domain decomposition & space charge mesh
sim.init_grids()
# load a 2 GeV proton beam
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 1.0e-9 # used with space charge
npart = 100000 # number of macro particles
# 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)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.16098260008648811e-3,
lambdaY=1.16098260008648811e-3,
lambdaT=1.0e-3,
lambdaPx=0.580491300043e-3,
lambdaPy=0.580491300043e-3,
lambdaPt=2.0e-3,
muxpx=0.0,
muypy=0.0,
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
lattice = [
monitor,
elements.Quad(ds=1.0, k=0.25, dx=0.003, dy=0.0, rotation=30.0, nslice=ns),
monitor,
]
sim.lattice.extend(lattice)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/alignment/input_alignment.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 100000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.16098260008648811e-3
beam.lambdaY = 1.16098260008648811e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 0.580491300043e-3
beam.lambdaPy = 0.580491300043e-3
beam.lambdaPt = 2.0e-3
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor quad_err monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
quad_err.type = quad
quad_err.ds = 1.0
quad_err.k = 0.25
quad_err.dx = 0.003
quad_err.rotation = 30.0
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_alignment.py
examples/alignment/analysis_alignment.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, tmean
def get_moments(beam):
"""Calculate standard deviations of beam position & momenta
and emittance values
Returns
-------
meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
"""
meanx = tmean(beam["position_x"])
meany = tmean(beam["position_y"])
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 (meanx, meany, 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 = 100000
assert num_particles == len(initial)
assert num_particles == len(final)
print("Initial Beam:")
meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(
initial
)
print(f" meanx={meanx:e} meany={meany:e}")
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 = 3.0 * num_particles ** (-0.5) * sigx
print(f" atol~={atol}")
assert np.allclose(
[meanx, meany],
[
0.0,
0.0,
],
atol=atol,
)
atol = 0.0 # ignored
rtol = 2.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.160982600086e-3,
1.160982600086e-3,
1.0e-3,
6.73940299e-7,
6.73940299e-7,
2.0e-6,
],
rtol=rtol,
atol=atol,
)
print("")
print("Final Beam:")
meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(
final
)
print(f" meanx={meanx:e} meany={meany:e}")
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 = 3.0 * num_particles ** (-0.5) * sigx
print(f" atol~={atol}")
assert np.allclose(
[meanx, meany],
[
1.79719761842e-4,
3.24815908981e-4,
],
atol=atol,
)
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.2372883901369e-3,
1.3772750218080e-3,
1.027364e-03,
7.39388142e-7,
7.39388142e-7,
2.0e-6,
],
rtol=rtol,
atol=atol,
)
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 \(\lambda_x\), \(\lambda_y\), \(\lambda_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
orImpactX 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.
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 -*-
from impactx import ImpactX, distribution, elements
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 = 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.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.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
# Drift elements
dr1 = elements.Drift(ds=0.4, nslice=1)
dr2 = elements.Drift(ds=0.032997, nslice=1)
# RF cavity element
rf = elements.RFCavity(
ds=1.31879807,
escale=62.0,
freq=1.3e9,
phase=85.5,
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,
],
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,
],
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.evolve()
# clean shutdown
sim.finalize()
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.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_rfcavity.py
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,
)
FODO Cell with RF
Stable FODO cell with short RF (buncher) cavities added for longitudinal focusing. The phase advance in all three phase planes is between 86-89 degrees.
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 9.80910407\) m
\(\alpha_\mathrm{x} = 0.0\)
\(\beta_\mathrm{y} = 1.31893788\) m
\(\alpha_\mathrm{y} = 0.0\)
\(\beta_\mathrm{t} = 4.6652668782\) m
\(\alpha_\mathrm{t} = 0.0\)
We use a 250 MeV proton beam with initial unnormalized rms emittance of 1 mm-mrad in all three phase planes.
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 \(\lambda_x\), \(\lambda_y\), \(\lambda_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_rf.py
orImpactX executable using an input file:
impactx input_fodo_rf.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/fodo_rf/run_fodo_rf.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 -*-
from impactx import ImpactX, distribution, elements
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 250 MeV proton beam with an initial
# unnormalized rms emittance of 1 mm-mrad in all
# three phase planes
kin_energy_MeV = 250.0 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=3.131948925200e-3,
lambdaY=1.148450209423e-3,
lambdaT=2.159922887089e-3,
lambdaPx=3.192900088357e-4,
lambdaPy=8.707386631090e-4,
lambdaPt=4.62979491526e-4,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.append(monitor)
# Quad elements
quad1 = elements.Quad(ds=0.15, k=2.5)
quad2 = elements.Quad(ds=0.3, k=-2.5)
# Drift element
drift1 = elements.Drift(ds=1.0)
# Short RF cavity element
shortrf1 = elements.Buncher(V=0.01, k=15.0)
lattice_no_drifts = [quad1, shortrf1, quad2, shortrf1, quad1]
# set first lattice element
sim.lattice.append(lattice_no_drifts[0])
# intersperse all remaining elements of the lattice with a drift element
for element in lattice_no_drifts[1:]:
sim.lattice.extend([drift1, element])
sim.lattice.append(monitor)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/fodo_rf/input_fodo_rf.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 3.131948925200e-3
beam.lambdaY = 1.148450209423e-3
beam.lambdaT = 2.159922887089e-3
beam.lambdaPx = 3.192900088357e-4
beam.lambdaPy = 8.707386631090e-4
beam.lambdaPt = 4.62979491526e-4
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor quad1 drift1 shortrf1 drift1 quad2 drift1 \
shortrf1 drift1 quad1 monitor
monitor.type = beam_monitor
monitor.backend = h5
quad1.type = quad
quad1.ds = 0.15
quad1.k = 2.5
drift1.type = drift
drift1.ds = 1.0
shortrf1.type = buncher
shortrf1.V = 0.01
shortrf1.k = 15.0
quad2.type = quad
quad2.ds = 0.3
quad2.k = -2.5
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
Analyze
We run the following script to analyze correctness:
Script analysis_fodo_rf.py
examples/fodo_rf/analysis_fodo_rf.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 = 2.2 * 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.145694e-03,
1.153344e-03,
2.155082e-03,
9.979770e-07,
1.008751e-06,
1.000691e-06,
],
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 = 2.2 * 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.112318e-03,
1.153322e-03,
2.166501e-03,
9.979770e-07,
1.008751e-06,
1.000691e-06,
],
rtol=rtol,
atol=atol,
)
FODO Cell, Chromatic
Stable FODO cell with a zero-current phase advance of 67.8 degrees, with chromatic focusing effects included.
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 2.82161941\) m
\(\alpha_\mathrm{x} = -1.59050035\)
\(\beta_\mathrm{y} = 2.82161941\) m
\(\alpha_\mathrm{y} = 1.59050035\)
We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm.
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_chr.py
orImpactX executable using an input file:
impactx input_fodo_chr.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/fodo/run_fodo_chr.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 = False
# sim.diagnostics = False # benchmarking
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 = 2.0e3 # 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.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 = 25 # number of slices per ds in the element
fodo = [
monitor,
elements.ChrDrift(ds=0.25, nslice=ns),
monitor,
elements.ChrQuad(ds=1.0, k=1.0, nslice=ns),
monitor,
elements.ChrDrift(ds=0.5, nslice=ns),
monitor,
elements.ChrQuad(ds=1.0, k=-1.0, nslice=ns),
monitor,
elements.ChrDrift(ds=0.25, nslice=ns),
monitor,
]
# assign a fodo segment
sim.lattice.extend(fodo)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/fodo/input_fodo_chr.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 3.9984884770e-5
beam.lambdaY = 3.9984884770e-5
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.6623538760e-5
beam.lambdaPy = 2.6623538760e-5
beam.lambdaPt = 2.0e-3
beam.muxpx = -0.846574929020762
beam.muypy = 0.846574929020762
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_chromatic
drift1.ds = 0.25
quad1.type = quad_chromatic
quad1.ds = 1.0
quad1.k = 1.0
drift2.type = drift_chromatic
drift2.ds = 0.5
quad2.type = quad_chromatic
quad2.ds = 1.0
quad2.k = -1.0
drift3.type = drift_chromatic
drift3.ds = 0.25
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_fodo_chr.py
examples/fodo/analysis_fodo_chr.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 = 2.2 * 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 = 2.2 * 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.4790118496224206e-005,
7.5357525169680140e-005,
9.9775879288128088e-004,
1.9959539836392703e-009,
2.0175014668882125e-009,
2.0013820380883801e-006,
],
rtol=rtol,
atol=atol,
)
Chain of thin multipoles
A series of thin multipoles (quad, sext, oct) with both normal and skew coefficients.
We use a 2 GeV electron beam.
The second moments of x, y, and t should be unchanged, but there is large emittance growth in the x and y phase planes.
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_multipole.py
orImpactX executable using an input file:
impactx input_multipole.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/multipole/run_multipole.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 -*-
from impactx import ImpactX, distribution, elements
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 2 GeV electron beam with an initial
# unnormalized rms emittance of nm
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 1.0e-9 # used without 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.Waterbag(
lambdaX=4.0e-3,
lambdaY=4.0e-3,
lambdaT=1.0e-3,
lambdaPx=3.0e-4,
lambdaPy=3.0e-4,
lambdaPt=2.0e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
multipole = [
monitor,
elements.Multipole(multipole=2, K_normal=3.0, K_skew=0.0),
elements.Multipole(multipole=3, K_normal=100.0, K_skew=-50.0),
elements.Multipole(multipole=4, K_normal=65.0, K_skew=6.0),
monitor,
]
# assign a fodo segment
sim.lattice.extend(multipole)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/multipole/input_multipole.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 4.0e-3
beam.lambdaY = 4.0e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 3.0e-4
beam.lambdaPy = 3.0e-4
beam.lambdaPt = 2.0e-3
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor thin_quadrupole thin_sextupole thin_octupole monitor
monitor.type = beam_monitor
monitor.backend = h5
thin_quadrupole.type = multipole
thin_quadrupole.multipole = 2 //Thin quadrupole
thin_quadrupole.k_normal = 3.0
thin_quadrupole.k_skew = 0.0
thin_sextupole.type = multipole
thin_sextupole.multipole = 3 //Thin sextupole
thin_sextupole.k_normal = 100.0
thin_sextupole.k_skew = -50.0
thin_octupole.type = multipole
thin_octupole.multipole = 4 //Thin octupole
thin_octupole.k_normal = 65.0
thin_octupole.k_skew = 6.0
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
Analyze
We run the following script to analyze correctness:
Script analysis_multipole.py
examples/multipole/analysis_multipole.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 = 2.2 * 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.017554e-03,
4.017044e-03,
9.977588e-04,
1.197572e-06,
1.210501e-06,
2.001382e-06,
],
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 = 2.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.017554e-03,
4.017044e-03,
9.977588e-04,
6.532644e-06,
6.630912e-06,
2.001382e-06,
],
rtol=rtol,
atol=atol,
)
A nonlinear focusing channel based on the IOTA nonlinear lens
A constant focusing channel with nonlinear focusing, using a string of thin IOTA nonlinear lens elements alternating with constant focusing elements.
We use a 2.5 MeV proton beam, corresponding to the nominal IOTA proton energy.
The two functions H (Hamiltonian) and I (the second invariant) should remain unchanged for all particles.
In this test, the initial and final values of \(\mu_H\), \(\sigma_H\), \(\mu_I\), \(\sigma_I\) must agree with nominal values.
Run
This example can be run either as:
Python script:
python3 run_iotalens.py
orImpactX executable using an input file:
impactx input_iotalens.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/iota_lens/run_iotalens.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 -*-
from impactx import ImpactX, distribution, elements
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()
# diagnostics: IOTA nonlinear lens invariants calculation
sim.set_diag_iota_invariants(alpha=0.0, beta=1.0, tn=0.4, cn=0.01)
# load a 2.5 MeV proton beam
kin_energy_MeV = 2.5 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=2.0e-3,
lambdaY=2.0e-3,
lambdaT=1.0e-3,
lambdaPx=3.0e-4,
lambdaPy=3.0e-4,
lambdaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# design the accelerator lattice
constEnd = elements.ConstF(ds=0.05, kx=1.0, ky=1.0, kt=1.0e-12)
nllens = elements.NonlinearLens(knll=4.0e-6, cnll=0.01)
const = elements.ConstF(ds=0.1, kx=1.0, ky=1.0, kt=1.0e-12)
num_lenses = 18
nllens_lattice = [constEnd] + [nllens, const] * (num_lenses - 1) + [nllens, constEnd]
# add elements to the lattice segment
sim.lattice.extend(nllens_lattice)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/iota_lens/input_iotalens.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.5
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 2.0e-3
beam.lambdaY = 2.0e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 3.0e-4
beam.lambdaPy = 3.0e-4
beam.lambdaPt = 0.0
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = const_end nllens foclens const_end
foclens.type = line
foclens.elements = const nllens
foclens.repeat = 17
nllens.type = nonlinear_lens
nllens.knll = 4.0e-6
nllens.cnll = 0.01
const_end.type = constf
const_end.ds = 0.05
const_end.kx = 1.0
const_end.ky = 1.0
const_end.kt = 1.0e-12
const.type = constf
const.ds = 0.1
const.kx = 1.0
const.ky = 1.0
const.kt = 1.0e-12
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.alpha = 0.0
diag.beta = 1.0
diag.tn = 0.4
diag.cn = 0.01
Analyze
We run the following script to analyze correctness:
Script analysis_iotalens.py
examples/iota_lens/analysis_iotalens.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import glob
import numpy as np
import pandas as pd
from scipy.stats import moment
def get_moments(beam):
"""Calculate mean and std dev of functions defining the IOTA invariants
Returns
-------
meanH, sigH, meanI, sigI
"""
meanH = np.mean(beam["H"])
sigH = moment(beam["H"], moment=2) ** 0.5
meanI = np.mean(beam["I"])
sigI = moment(beam["I"], moment=2) ** 0.5
return (meanH, sigH, meanI, sigI)
def read_all_files(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(
(
pd.read_csv(filename, delimiter=r"\s+")
for filename in glob.glob(file_pattern)
),
axis=0,
ignore_index=True,
).set_index("id")
# initial/final beam
initial = read_all_files("diags/nonlinear_lens_invariants_000000.*")
final = read_all_files("diags/nonlinear_lens_invariants_final.*")
# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)
print("Initial Beam:")
meanH, sigH, meanI, sigI = get_moments(initial)
print(f" meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
atol = 0.0 # a big number
rtol = 1.5 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[meanH, sigH, meanI, sigI],
[4.122650e-02, 4.235181e-02, 7.356057e-02, 8.793753e-02],
rtol=rtol,
atol=atol,
)
print("")
print("Final Beam:")
meanH, sigH, meanI, sigI = get_moments(final)
print(f" meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
atol = 0.0 # a big number
rtol = 1.5 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[meanH, sigH, meanI, sigI],
[4.122704e-02, 4.230576e-02, 7.348275e-02, 8.783157e-02],
rtol=rtol,
atol=atol,
)
# join tables on particle ID, so we can compare the same particle initial->final
beam_joined = final.join(initial, lsuffix="_final", rsuffix="_initial")
# add new columns: dH and dI
beam_joined["dH"] = (beam_joined["H_initial"] - beam_joined["H_final"]).abs()
beam_joined["dI"] = (beam_joined["I_initial"] - beam_joined["I_final"]).abs()
# print(beam_joined)
# particle-wise comparison of H & I initial to final
atol = 2.0e-3
rtol = 0.0 # large number
print()
print(f" atol={atol} (ignored: rtol~={rtol})")
print(f" dH_max={beam_joined['dH'].max()}")
assert np.allclose(beam_joined["dH"], 0.0, rtol=rtol, atol=atol)
atol = 3.0e-3
print(f" atol={atol} (ignored: rtol~={rtol})")
print(f" dI_max={beam_joined['dI'].max()}")
assert np.allclose(beam_joined["dI"], 0.0, rtol=rtol, atol=atol)
A nonlinear focusing channel based on the physical IOTA nonlinear magnet
A representation of the physical IOTA nonlinear magnetic element with realistic s-dependence, obtained using a sequence of nonlinear lenses and drifts equivalent to the use of a second-order symplectic integrator.
A thin linear lens is added at the exit of the nonlinear element, representing the ideal map associated with the remainder of the lattice.
We use a 2.5 MeV proton beam, corresponding to the nominal IOTA proton energy.
The two functions H (Hamiltonian) and I (the second invariant) are evaluated at the entrance to the nonlinear element, and then again after the thin lens (representing a single period). These values should be unchanged for all particles (to within acceptable tolerance).
In this test, the initial and final values of \(\mu_H\), \(\sigma_H\), \(\mu_I\), \(\sigma_I\) must agree with nominal values.
Run
This example can be run either as:
Python script:
python3 run_iotalens_sdep.py
orImpactX executable using an input file:
impactx input_iotalens_sdep.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/iota_lens/run_iotalens_sdep.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 math
from impactx import ImpactX, distribution, elements
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()
# diagnostics: IOTA nonlinear lens invariants calculation
sim.set_diag_iota_invariants(
alpha=1.376381920471173, beta=1.892632003628881, tn=0.4, cn=0.01
)
# load a 2.5 MeV proton beam
kin_energy_MeV = 2.5 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.397456296195e-003,
lambdaY=1.397456296195e-003,
lambdaT=1.0e-4,
lambdaPx=1.256184325020e-003,
lambdaPy=1.256184325020e-003,
lambdaPt=0.0,
muxpx=0.8090169943749474,
muypy=0.8090169943749474,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
sim.lattice.append(monitor)
# defining parameters of the nonlinear lens
lens_length = 1.8
num_lenses = 18
tune_advance = 0.3
c_parameter = 0.01
t_strength = 0.4
ds = lens_length / num_lenses
# drift elements
ds_half = ds / 2.0
dr = elements.Drift(ds=ds_half)
# define the nonlinear lens segments
for j in range(0, num_lenses):
s = -lens_length / 2.0 + ds_half + j * ds
beta_star = lens_length / 2.0 * 1.0 / math.tan(math.pi * tune_advance)
beta = beta_star * (
1.0 + (2.0 * s * math.tan(math.pi * tune_advance) / lens_length) ** 2
)
knll_s = t_strength * c_parameter**2 * ds / beta
cnll_s = c_parameter * math.sqrt(beta)
nllens = elements.NonlinearLens(knll=knll_s, cnll=cnll_s)
segments = [dr, nllens, dr]
sim.lattice.extend(segments)
# focusing lens
const = elements.ConstF(
ds=1.0e-8, kx=12060.113295833, ky=12060.113295833, kt=1.0e-12, nslice=1
)
sim.lattice.append(const)
sim.lattice.append(monitor)
# number of periods
sim.periods = 1
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/iota_lens/input_iotalens_sdep.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.5
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
#beam.lambdaX = 2.0e-3
#beam.lambdaY = 2.0e-3
#beam.lambdaT = 1.0e-3
#beam.lambdaPx = 3.0e-4
#beam.lambdaPy = 3.0e-4
#beam.lambdaPt = 0.0
#beam.muxpx = 0.0
#beam.muypy = 0.0
#beam.mutpt = 0.0
beam.lambdaX = 1.397456296195e-003
beam.lambdaY = 1.397456296195e-003
beam.lambdaT = 1.0e-4
beam.lambdaPx = 1.256184325020e-003
beam.lambdaPy = 1.256184325020e-003
beam.lambdaPt = 0.0
beam.muxpx = 0.8090169943749474
beam.muypy = 0.8090169943749474
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor lens const monitor
lens.type = line
lens.elements = = dr_end nll1 dr nll2 dr nll3 dr nll4 dr nll5 dr nll6 \
dr nll7 dr nll8 dr nll9 dr nll9 dr nll8 dr nll7 \
dr nll6 dr nll5 dr nll4 dr nll3 dr nll2 dr nll1 \
dr_end
# Nonlinear lens segments
nll1.type = nonlinear_lens
nll1.knll = 2.2742558121e-6
nll1.cnll = 0.013262040169952
nll2.type = nonlinear_lens
nll2.knll = 2.641786366e-6
nll2.cnll = 0.012304986694423
nll3.type = nonlinear_lens
nll3.knll = 3.076868353e-6
nll3.cnll = 0.011401855643727
nll4.type = nonlinear_lens
nll4.knll = 3.582606522e-6
nll4.cnll = 0.010566482535866
nll5.type = nonlinear_lens
nll5.knll = 4.151211157e-6
nll5.cnll = 0.009816181575902
nll6.type = nonlinear_lens
nll6.knll = 4.754946985e-6
nll6.cnll = 0.0091718544892154
nll7.type = nonlinear_lens
nll7.knll = 5.337102374e-6
nll7.cnll = 0.008657195579489
nll8.type = nonlinear_lens
nll8.knll = 5.811437818e-6
nll8.cnll = 0.008296371635942
nll9.type = nonlinear_lens
nll9.knll = 6.081693334e-6
nll9.cnll = 0.008109941789663
dr.type = drift
dr.ds = 0.1
dr_end.type = drift
dr_end.ds = 0.05
# Focusing of the external lattice
const.type = constf
const.ds = 1.0e-8
const.kx = 12060.113295833
const.ky = 12060.113295833
const.kt = 1.0e-12
const.nslice = 1
# Beam Monitor: Diagnostics
monitor.type = beam_monitor
monitor.backend = h5
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.alpha = 1.376381920471173
diag.beta = 1.892632003628881
diag.tn = 0.4
diag.cn = 0.01
Analyze
We run the following script to analyze correctness:
Script analysis_iotalens_sdep.py
examples/iota_lens/analysis_iotalens_sdep.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import glob
import numpy as np
import pandas as pd
from scipy.stats import moment
def get_moments(beam):
"""Calculate mean and std dev of functions defining the IOTA invariants
Returns
-------
meanH, sigH, meanI, sigI
"""
meanH = np.mean(beam["H"])
sigH = moment(beam["H"], moment=2) ** 0.5
meanI = np.mean(beam["I"])
sigI = moment(beam["I"], moment=2) ** 0.5
return (meanH, sigH, meanI, sigI)
def read_all_files(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(
(
pd.read_csv(filename, delimiter=r"\s+")
for filename in glob.glob(file_pattern)
),
axis=0,
ignore_index=True,
).set_index("id")
# initial/final beam
initial = read_all_files("diags/nonlinear_lens_invariants_000000.*")
final = read_all_files("diags/nonlinear_lens_invariants_final.*")
# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)
print("Initial Beam:")
meanH, sigH, meanI, sigI = get_moments(initial)
print(f" meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
atol = 0.0 # a big number
rtol = 3.0 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[meanH, sigH, meanI, sigI],
# [5.993291e-02, 3.426664e-02, 8.513875e-02, 7.022481e-02],
[6.016450e-02, 3.502064e-02, 8.560226e-02, 7.148169e-02],
rtol=rtol,
atol=atol,
)
print("")
print("Final Beam:")
meanH, sigH, meanI, sigI = get_moments(final)
print(f" meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
atol = 0.0 # a big number
rtol = 3.0 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[meanH, sigH, meanI, sigI],
# [5.993291e-02, 3.426664e-02, 8.513875e-02, 7.022481e-02],
[6.016450e-02, 3.502064e-02, 8.560226e-02, 7.148169e-02],
rtol=rtol,
atol=atol,
)
# join tables on particle ID, so we can compare the same particle initial->final
beam_joined = final.join(initial, lsuffix="_final", rsuffix="_initial")
# add new columns: dH and dI
beam_joined["dH"] = (beam_joined["H_initial"] - beam_joined["H_final"]).abs()
beam_joined["dI"] = (beam_joined["I_initial"] - beam_joined["I_final"]).abs()
# print(beam_joined)
# particle-wise comparison of H & I initial to final
Hrms = np.sqrt(sigH**2 + meanH**2)
Irms = np.sqrt(sigI**2 + meanI**2)
atol = 2.5e-3 * Hrms
rtol = 0.0 # large number
print()
print(f" atol={atol} (ignored: rtol~={rtol})")
print(f" dH_max={beam_joined['dH'].max()}")
assert np.allclose(beam_joined["dH"], 0.0, rtol=rtol, atol=atol)
atol = 3.5e-3 * Irms
rtol = 0.0
print()
print(f" atol={atol} (ignored: rtol~={rtol})")
print(f" dI_max={beam_joined['dI'].max()}")
assert np.allclose(beam_joined["dI"], 0.0, rtol=rtol, atol=atol)
The “bare” linear lattice of the Fermilab IOTA storage ring
The linear lattice of the IOTA storage ring, configured for operation with a 2.5 MeV proton beam.
The drift regions available for insertion of the special nonlinear magnetic element for integrable optics experiments are denoted dnll
.
The second moments of the particle distribution after a single turn should coincide with the initial section moments of the particle distribution, to within the level expected due to numerical particle noise. The example runs 5 turns.
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_iotalattice.py
orImpactX executable using an input file:
impactx input_iotalattice.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/iota_lattice/run_iotalattice.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Chad Mitchell, Axel Huebl
# 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 = False
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = True
# domain decomposition & space charge mesh
sim.init_grids()
# init particle beam
kin_energy_MeV = 2.5
bunch_charge_C = 1.0e-9 # used with space charge
npart = 10000
# 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)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.588960728035e-3,
lambdaY=2.496625268437e-3,
lambdaT=1.0e-3,
lambdaPx=2.8320397837724e-3,
lambdaPy=1.802433091137e-3,
lambdaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# init accelerator lattice
ns = 10 # number of slices per ds in the element
# Drift elements
dra1 = elements.Drift(ds=0.9125, nslice=ns)
dra2 = elements.Drift(ds=0.135, nslice=ns)
dra3 = elements.Drift(ds=0.725, nslice=ns)
dra4 = elements.Drift(ds=0.145, nslice=ns)
dra5 = elements.Drift(ds=0.3405, nslice=ns)
drb1 = elements.Drift(ds=0.3205, nslice=ns)
drb2 = elements.Drift(ds=0.14, nslice=ns)
drb3 = elements.Drift(ds=0.1525, nslice=ns)
drb4 = elements.Drift(ds=0.31437095, nslice=ns)
drc1 = elements.Drift(ds=0.42437095, nslice=ns)
drc2 = elements.Drift(ds=0.355, nslice=ns)
dnll = elements.Drift(ds=1.8, nslice=ns)
drd1 = elements.Drift(ds=0.62437095, nslice=ns)
drd2 = elements.Drift(ds=0.42, nslice=ns)
drd3 = elements.Drift(ds=1.625, nslice=ns)
drd4 = elements.Drift(ds=0.6305, nslice=ns)
dre1 = elements.Drift(ds=0.5305, nslice=ns)
dre2 = elements.Drift(ds=1.235, nslice=ns)
dre3 = elements.Drift(ds=0.8075, nslice=ns)
# Bend elements
rc30 = 0.822230996255981
sbend30 = elements.Sbend(ds=0.4305191429, rc=rc30)
edge30 = elements.DipEdge(psi=0.0, rc=rc30, g=0.058, K2=0.5)
rc60 = 0.772821121503940
sbend60 = elements.Sbend(ds=0.8092963858, rc=rc60)
edge60 = elements.DipEdge(psi=0.0, rc=rc60, g=0.058, K2=0.5)
# Quad elements
ds_quad = 0.21
qa1 = elements.Quad(ds=ds_quad, k=-8.78017699, nslice=ns)
qa2 = elements.Quad(ds=ds_quad, k=13.24451745, nslice=ns)
qa3 = elements.Quad(ds=ds_quad, k=-13.65151327, nslice=ns)
qa4 = elements.Quad(ds=ds_quad, k=19.75138652, nslice=ns)
qb1 = elements.Quad(ds=ds_quad, k=-10.84199727, nslice=ns)
qb2 = elements.Quad(ds=ds_quad, k=16.24844348, nslice=ns)
qb3 = elements.Quad(ds=ds_quad, k=-8.27411104, nslice=ns)
qb4 = elements.Quad(ds=ds_quad, k=-7.45719247, nslice=ns)
qb5 = elements.Quad(ds=ds_quad, k=14.03362243, nslice=ns)
qb6 = elements.Quad(ds=ds_quad, k=-12.23595641, nslice=ns)
qc1 = elements.Quad(ds=ds_quad, k=-13.18863768, nslice=ns)
qc2 = elements.Quad(ds=ds_quad, k=11.50601829, nslice=ns)
qc3 = elements.Quad(ds=ds_quad, k=-11.10445869, nslice=ns)
qd1 = elements.Quad(ds=ds_quad, k=-6.78179218, nslice=ns)
qd2 = elements.Quad(ds=ds_quad, k=5.19026998, nslice=ns)
qd3 = elements.Quad(ds=ds_quad, k=-5.8586173, nslice=ns)
qd4 = elements.Quad(ds=ds_quad, k=4.62460039, nslice=ns)
qe1 = elements.Quad(ds=ds_quad, k=-4.49607687, nslice=ns)
qe2 = elements.Quad(ds=ds_quad, k=6.66737146, nslice=ns)
qe3 = elements.Quad(ds=ds_quad, k=-6.69148177, nslice=ns)
# build lattice: first half, qe3, then mirror
# fmt: off
lattice_half = [
dra1, qa1, dra2, qa2, dra3, qa3, dra4, qa4, dra5,
edge30, sbend30, edge30, drb1, qb1, drb2, qb2, drb2, qb3,
drb3, dnll, drb3, qb4, drb2, qb5, drb2, qb6, drb4,
edge60, sbend60, edge60, drc1, qc1, drc2, qc2, drc2, qc3, drc1,
edge60, sbend60, edge60, drd1, qd1, drd2, qd2, drd3, qd3, drd2, qd4, drd4,
edge30, sbend30, edge30, dre1, qe1, dre2, qe2, dre3
]
# fmt:on
sim.lattice.append(monitor)
sim.lattice.extend(lattice_half)
sim.lattice.append(qe3)
lattice_half.reverse()
sim.lattice.extend(lattice_half)
sim.lattice.append(monitor)
# number of turns in the ring
sim.periods = 5
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/iota_lattice/input_iotalattice.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.5
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.588960728035e-3
beam.lambdaY = 2.496625268437e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.8320397837724e-3
beam.lambdaPy = 1.802433091137e-3
beam.lambdaPt = 0.0
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.periods = 5
lattice.elements = monitor first_half qe3 second_half monitor
# lines
first_half.type = line
first_half.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5 \
edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3 \
drb3 dnll drb3 qb4 drb2 qb5 drb2 qb6 drb4 \
edge60 sbend60 edge60 drc1 qc1 drc2 qc2 drc2 qc3 drc1 \
edge60 sbend60 edge60 drd1 qd1 drd2 qd2 drd3 qd3 drd2 qd4 drd4 \
edge30 sbend30 edge30 dre1 qe1 dre2 qe2 dre3
second_half.type = line
second_half.reverse = true
second_half.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5 \
edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3 \
drb3 dnll drb3 qb4 drb2 qb5 drb2 qb6 drb4 \
edge60 sbend60 edge60 drc1 qc1 drc2 qc2 drc2 qc3 drc1 \
edge60 sbend60 edge60 drd1 qd1 drd2 qd2 drd3 qd3 drd2 qd4 drd4 \
edge30 sbend30 edge30 dre1 qe1 dre2 qe2 dre3
# thick element splitting for space charge
lattice.nslice = 10
# Drift elements:
dra1.type = drift
dra1.ds = 0.9125
dra2.type = drift
dra2.ds = 0.135
dra3.type = drift
dra3.ds = 0.725
dra4.type = drift
dra4.ds = 0.145
dra5.type = drift
dra5.ds = 0.3405
drb1.type = drift
drb1.ds = 0.3205
drb2.type = drift
drb2.ds = 0.14
drb3.type = drift
drb3.ds = 0.1525
drb4.type = drift
drb4.ds = 0.31437095
drc1.type = drift
drc1.ds = 0.42437095
drc2.type = drift
drc2.ds = 0.355
dnll.type = drift
dnll.ds = 1.8
drd1.type = drift
drd1.ds = 0.62437095
drd2.type = drift
drd2.ds = 0.42
drd3.type = drift
drd3.ds = 1.625
drd4.type = drift
drd4.ds = 0.6305
dre1.type = drift
dre1.ds = 0.5305
dre2.type = drift
dre2.ds = 1.235
dre3.type = drift
dre3.ds = 0.8075
# Bend elements:
sbend30.type = sbend
sbend30.ds = 0.4305191429
sbend30.rc = 0.822230996255981
edge30.type = dipedge
edge30.psi = 0.0
edge30.rc = 0.822230996255981
edge30.g = 0.058
edge30.K2 = 0.5
sbend60.type = sbend
sbend60.ds = 0.8092963858
sbend60.rc = 0.772821121503940
edge60.type = dipedge
edge60.psi = 0.0
edge60.rc = 0.772821121503940
edge60.g = 0.058
edge60.K2 = 0.5
# Quad elements:
qa1.type = quad
qa1.ds = 0.21
qa1.k = -8.78017699
qa2.type = quad
qa2.ds = 0.21
qa2.k = 13.24451745
qa3.type = quad
qa3.ds = 0.21
qa3.k = -13.65151327
qa4.type = quad
qa4.ds = 0.21
qa4.k = 19.75138652
qb1.type = quad
qb1.ds = 0.21
qb1.k = -10.84199727
qb2.type = quad
qb2.ds = 0.21
qb2.k = 16.24844348
qb3.type = quad
qb3.ds = 0.21
qb3.k = -8.27411104
qb4.type = quad
qb4.ds = 0.21
qb4.k = -7.45719247
qb5.type = quad
qb5.ds = 0.21
qb5.k = 14.03362243
qb6.type = quad
qb6.ds = 0.21
qb6.k = -12.23595641
qc1.type = quad
qc1.ds = 0.21
qc1.k = -13.18863768
qc2.type = quad
qc2.ds = 0.21
qc2.k = 11.50601829
qc3.type = quad
qc3.ds = 0.21
qc3.k = -11.10445869
qd1.type = quad
qd1.ds = 0.21
qd1.k = -6.78179218
qd2.type = quad
qd2.ds = 0.21
qd2.k = 5.19026998
qd3.type = quad
qd3.ds = 0.21
qd3.k = -5.8586173
qd4.type = quad
qd4.ds = 0.21
qd4.k = 4.62460039
qe1.type = quad
qe1.ds = 0.21
qe1.k = -4.49607687
qe2.type = quad
qe2.ds = 0.21
qe2.k = 6.66737146
qe3.type = quad
qe3.ds = 0.21
qe3.k = -6.69148177
# Beam Monitor: Diagnostics
monitor.type = beam_monitor
monitor.backend = h5
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_iotalattice.py
examples/iota_lattice/analysis_iotalattice.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 # a big number
rtol = 1.9 * 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.595934e-03, 2.507263e-03, 9.977588e-04, 4.490896e-06, 4.539378e-06, 0.000000e00],
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 # a big number
rtol = 1.9 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[sigx, sigy, emittance_x, emittance_y, emittance_t],
[
1.579848e-03,
2.510900e-03,
4.490897e-06,
4.539378e-06,
0.0,
],
rtol=rtol,
atol=atol,
)
The full nonlinear lattice of the Fermilab IOTA storage ring
The full nonlinear lattice of the IOTA storage ring, configured for operation with a 2.5 MeV proton beam.
The special nonlinear magnetic element for integrable optics experiments is denoted nll
. To simplify analysis,
the lattice has been arranged so that nll appears as the first element in the sequence.
The two functions H (Hamiltonian) and I (the second invariant) are evaluated at the entrance to the nonlinear element. These values should be unchanged for all particles (to within acceptable tolerance), over the specified number of periods (default 5).
In this test, the initial and final values of \(\mu_H\), \(\sigma_H\), \(\mu_I\), \(\sigma_I\) must agree with nominal values.
Run
This example can be run either as:
Python script:
python3 run_iotalattice_sdep.py
orImpactX executable using an input file:
impactx input_iotalattice_sdep.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/iota_lattice/run_iotalattice_sdep.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Chad Mitchell, Axel Huebl
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
import math
from impactx import ImpactX, distribution, elements
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()
# diagnostics: IOTA nonlinear lens invariants calculation
sim.set_diag_iota_invariants(
alpha=1.376381920471173, beta=1.892632003628881, tn=0.4, cn=0.01
)
# init particle beam
energy_MeV = 2.5
bunch_charge_C = 1.0e-9 # used with space charge
npart = 10000
# reference particle
ref = sim.particle_container().ref_particle()
ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.865379469388e-003,
lambdaY=2.0192133150418e-003,
lambdaT=1.0e-3,
lambdaPx=1.402566720991e-003,
lambdaPy=9.57593913381e-004,
lambdaPt=0.0,
muxpx=0.482260919078473,
muypy=0.762127656873158,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# init accelerator lattice
ns = 10 # number of slices per ds in the element
# Drift elements
dra1 = elements.Drift(ds=0.9125, nslice=ns)
dra2 = elements.Drift(ds=0.135, nslice=ns)
dra3 = elements.Drift(ds=0.725, nslice=ns)
dra4 = elements.Drift(ds=0.145, nslice=ns)
dra5 = elements.Drift(ds=0.3405, nslice=ns)
drb1 = elements.Drift(ds=0.3205, nslice=ns)
drb2 = elements.Drift(ds=0.14, nslice=ns)
drb3 = elements.Drift(ds=0.1525, nslice=ns)
drb4 = elements.Drift(ds=0.31437095, nslice=ns)
drc1 = elements.Drift(ds=0.42437095, nslice=ns)
drc2 = elements.Drift(ds=0.355, nslice=ns)
dnll = elements.Drift(ds=1.8, nslice=ns)
drd1 = elements.Drift(ds=0.62437095, nslice=ns)
drd2 = elements.Drift(ds=0.42, nslice=ns)
drd3 = elements.Drift(ds=1.625, nslice=ns)
drd4 = elements.Drift(ds=0.6305, nslice=ns)
dre1 = elements.Drift(ds=0.5305, nslice=ns)
dre2 = elements.Drift(ds=1.235, nslice=ns)
dre3 = elements.Drift(ds=0.8075, nslice=ns)
# Bend elements
rc30 = 0.822230996255981
sbend30 = elements.Sbend(ds=0.4305191429, rc=rc30)
edge30 = elements.DipEdge(psi=0.0, rc=rc30, g=0.058, K2=0.5)
rc60 = 0.772821121503940
sbend60 = elements.Sbend(ds=0.8092963858, rc=rc60)
edge60 = elements.DipEdge(psi=0.0, rc=rc60, g=0.058, K2=0.5)
# Quad elements
ds_quad = 0.21
qa1 = elements.Quad(ds=ds_quad, k=-8.78017699, nslice=ns)
qa2 = elements.Quad(ds=ds_quad, k=13.24451745, nslice=ns)
qa3 = elements.Quad(ds=ds_quad, k=-13.65151327, nslice=ns)
qa4 = elements.Quad(ds=ds_quad, k=19.75138652, nslice=ns)
qb1 = elements.Quad(ds=ds_quad, k=-10.84199727, nslice=ns)
qb2 = elements.Quad(ds=ds_quad, k=16.24844348, nslice=ns)
qb3 = elements.Quad(ds=ds_quad, k=-8.27411104, nslice=ns)
qb4 = elements.Quad(ds=ds_quad, k=-7.45719247, nslice=ns)
qb5 = elements.Quad(ds=ds_quad, k=14.03362243, nslice=ns)
qb6 = elements.Quad(ds=ds_quad, k=-12.23595641, nslice=ns)
qc1 = elements.Quad(ds=ds_quad, k=-13.18863768, nslice=ns)
qc2 = elements.Quad(ds=ds_quad, k=11.50601829, nslice=ns)
qc3 = elements.Quad(ds=ds_quad, k=-11.10445869, nslice=ns)
qd1 = elements.Quad(ds=ds_quad, k=-6.78179218, nslice=ns)
qd2 = elements.Quad(ds=ds_quad, k=5.19026998, nslice=ns)
qd3 = elements.Quad(ds=ds_quad, k=-5.8586173, nslice=ns)
qd4 = elements.Quad(ds=ds_quad, k=4.62460039, nslice=ns)
qe1 = elements.Quad(ds=ds_quad, k=-4.49607687, nslice=ns)
qe2 = elements.Quad(ds=ds_quad, k=6.66737146, nslice=ns)
qe3 = elements.Quad(ds=ds_quad, k=-6.69148177, nslice=ns)
# Special (elliptic) nonlinear element:
# set basic parameters of the nonlinear element
lens_length = 1.8
num_lenses = 18
tune_advance = 0.3
c_parameter = 0.01
t_strength = 0.4
ds = lens_length / num_lenses
# build up the nonlinear lens in segments
ds_half = ds / 2.0
dr = elements.Drift(ds=ds_half, nslice=1)
nll = []
for j in range(0, num_lenses):
s = -lens_length / 2.0 + ds_half + j * ds
beta_star = lens_length / 2.0 * 1.0 / math.tan(math.pi * tune_advance)
beta = beta_star * (
1.0 + (2.0 * s * math.tan(math.pi * tune_advance) / lens_length) ** 2
)
knll_s = t_strength * c_parameter**2 * ds / beta
cnll_s = c_parameter * math.sqrt(beta)
nllens = elements.NonlinearLens(knll=knll_s, cnll=cnll_s)
segment = [dr, nllens, dr]
nll.extend(segment)
lattice_before_nll = [
dra1,
qa1,
dra2,
qa2,
dra3,
qa3,
dra4,
qa4,
dra5,
edge30,
sbend30,
edge30,
drb1,
qb1,
drb2,
qb2,
drb2,
qb3,
drb3,
]
lattice_after_nll = [
drb3,
qb4,
drb2,
qb5,
drb2,
qb6,
drb4,
edge60,
sbend60,
edge60,
drc1,
qc1,
drc2,
qc2,
drc2,
qc3,
drc1,
edge60,
sbend60,
edge60,
drd1,
qd1,
drd2,
qd2,
drd3,
qd3,
drd2,
qd4,
drd4,
edge30,
sbend30,
edge30,
dre1,
qe1,
dre2,
qe2,
dre3,
]
# build lattice: first half, qe3, then mirror
# modified to place nll as the first element
# fmt:on
sim.lattice.append(monitor)
sim.lattice.extend(nll)
sim.lattice.extend(lattice_after_nll)
sim.lattice.append(qe3)
lattice_after_nll.reverse()
sim.lattice.extend(lattice_after_nll)
sim.lattice.append(dnll)
lattice_before_nll.reverse()
sim.lattice.extend(lattice_before_nll)
lattice_before_nll.reverse()
sim.lattice.extend(lattice_before_nll)
sim.lattice.append(monitor)
# number of turns in the ring
sim.periods = 5
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/iota_lattice/input_iotalattice_sdep.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.5
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.865379469388e-003
beam.lambdaY = 2.0192133150418e-003
beam.lambdaT = 1.0e-3
beam.lambdaPx = 1.402566720991e-003
beam.lambdaPy = 9.57593913381e-004
beam.lambdaPt = 0.0
beam.muxpx = 0.482260919078473
beam.muypy = 0.762127656873158
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.periods = 5
lattice.elements = monitor nll after_nll qe3 after_nll_rev dnll before_nll_rev \
before_nll monitor
# lines
before_nll.type = line
before_nll.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5 \
edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3 \
drb3
after_nll.type = line
after_nll.elements = drb3 qb4 drb2 qb5 drb2 qb6 drb4 \
edge60 sbend60 edge60 drc1 qc1 drc2 qc2 drc2 qc3 drc1 \
edge60 sbend60 edge60 drd1 qd1 drd2 qd2 drd3 qd3 drd2 qd4 drd4 \
edge30 sbend30 edge30 dre1 qe1 dre2 qe2 dre3
before_nll_rev.type = line
before_nll_rev.reverse = true
before_nll_rev.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5 \
edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3 \
drb3
after_nll_rev.type = line
after_nll_rev.reverse = true
after_nll_rev.elements = drb3 qb4 drb2 qb5 drb2 qb6 drb4 \
edge60 sbend60 edge60 drc1 qc1 drc2 qc2 drc2 qc3 drc1 \
edge60 sbend60 edge60 drd1 qd1 drd2 qd2 drd3 qd3 drd2 qd4 drd4 \
edge30 sbend30 edge30 dre1 qe1 dre2 qe2 dre3
nll.type = line
nll.elements = = dr_end nll1 dr nll2 dr nll3 dr nll4 dr nll5 dr nll6 \
dr nll7 dr nll8 dr nll9 dr nll9 dr nll8 dr nll7 \
dr nll6 dr nll5 dr nll4 dr nll3 dr nll2 dr nll1 \
dr_end
# thick element splitting for space charge
lattice.nslice = 10
# Nonlinear lens segments
nll1.type = nonlinear_lens
nll1.knll = 2.2742558121e-6
nll1.cnll = 0.013262040169952
nll2.type = nonlinear_lens
nll2.knll = 2.641786366e-6
nll2.cnll = 0.012304986694423
nll3.type = nonlinear_lens
nll3.knll = 3.076868353e-6
nll3.cnll = 0.011401855643727
nll4.type = nonlinear_lens
nll4.knll = 3.582606522e-6
nll4.cnll = 0.010566482535866
nll5.type = nonlinear_lens
nll5.knll = 4.151211157e-6
nll5.cnll = 0.009816181575902
nll6.type = nonlinear_lens
nll6.knll = 4.754946985e-6
nll6.cnll = 0.0091718544892154
nll7.type = nonlinear_lens
nll7.knll = 5.337102374e-6
nll7.cnll = 0.008657195579489
nll8.type = nonlinear_lens
nll8.knll = 5.811437818e-6
nll8.cnll = 0.008296371635942
nll9.type = nonlinear_lens
nll9.knll = 6.081693334e-6
nll9.cnll = 0.008109941789663
# Drift elements:
dr.type = drift
dr.ds = 0.1
dr_end.type = drift
dr_end.ds = 0.05
dra1.type = drift
dra1.ds = 0.9125
dra2.type = drift
dra2.ds = 0.135
dra3.type = drift
dra3.ds = 0.725
dra4.type = drift
dra4.ds = 0.145
dra5.type = drift
dra5.ds = 0.3405
drb1.type = drift
drb1.ds = 0.3205
drb2.type = drift
drb2.ds = 0.14
drb3.type = drift
drb3.ds = 0.1525
drb4.type = drift
drb4.ds = 0.31437095
drc1.type = drift
drc1.ds = 0.42437095
drc2.type = drift
drc2.ds = 0.355
dnll.type = drift
dnll.ds = 1.8
drd1.type = drift
drd1.ds = 0.62437095
drd2.type = drift
drd2.ds = 0.42
drd3.type = drift
drd3.ds = 1.625
drd4.type = drift
drd4.ds = 0.6305
dre1.type = drift
dre1.ds = 0.5305
dre2.type = drift
dre2.ds = 1.235
dre3.type = drift
dre3.ds = 0.8075
# Bend elements:
sbend30.type = sbend
sbend30.ds = 0.4305191429
sbend30.rc = 0.822230996255981
edge30.type = dipedge
edge30.psi = 0.0
edge30.rc = 0.822230996255981
edge30.g = 0.058
edge30.K2 = 0.5
sbend60.type = sbend
sbend60.ds = 0.8092963858
sbend60.rc = 0.772821121503940
edge60.type = dipedge
edge60.psi = 0.0
edge60.rc = 0.772821121503940
edge60.g = 0.058
edge60.K2 = 0.5
# Quad elements:
qa1.type = quad
qa1.ds = 0.21
qa1.k = -8.78017699
qa2.type = quad
qa2.ds = 0.21
qa2.k = 13.24451745
qa3.type = quad
qa3.ds = 0.21
qa3.k = -13.65151327
qa4.type = quad
qa4.ds = 0.21
qa4.k = 19.75138652
qb1.type = quad
qb1.ds = 0.21
qb1.k = -10.84199727
qb2.type = quad
qb2.ds = 0.21
qb2.k = 16.24844348
qb3.type = quad
qb3.ds = 0.21
qb3.k = -8.27411104
qb4.type = quad
qb4.ds = 0.21
qb4.k = -7.45719247
qb5.type = quad
qb5.ds = 0.21
qb5.k = 14.03362243
qb6.type = quad
qb6.ds = 0.21
qb6.k = -12.23595641
qc1.type = quad
qc1.ds = 0.21
qc1.k = -13.18863768
qc2.type = quad
qc2.ds = 0.21
qc2.k = 11.50601829
qc3.type = quad
qc3.ds = 0.21
qc3.k = -11.10445869
qd1.type = quad
qd1.ds = 0.21
qd1.k = -6.78179218
qd2.type = quad
qd2.ds = 0.21
qd2.k = 5.19026998
qd3.type = quad
qd3.ds = 0.21
qd3.k = -5.8586173
qd4.type = quad
qd4.ds = 0.21
qd4.k = 4.62460039
qe1.type = quad
qe1.ds = 0.21
qe1.k = -4.49607687
qe2.type = quad
qe2.ds = 0.21
qe2.k = 6.66737146
qe3.type = quad
qe3.ds = 0.21
qe3.k = -6.69148177
# Beam Monitor: Diagnostics
monitor.type = beam_monitor
monitor.backend = h5
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
diag.alpha = 1.376381920471173
diag.beta = 1.892632003628881
diag.tn = 0.4
diag.cn = 0.01
Analyze
We run the following script to analyze correctness:
Script analysis_iotalattice_sdep.py
examples/iota_lattice/analysis_iotalattice_sdep.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import glob
import numpy as np
import pandas as pd
from scipy.stats import moment
def get_moments(beam):
"""Calculate mean and std dev of functions defining the IOTA invariants
Returns
-------
meanH, sigH, meanI, sigI
"""
meanH = np.mean(beam["H"])
sigH = moment(beam["H"], moment=2) ** 0.5
meanI = np.mean(beam["I"])
sigI = moment(beam["I"], moment=2) ** 0.5
return (meanH, sigH, meanI, sigI)
def read_all_files(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(
(
pd.read_csv(filename, delimiter=r"\s+")
for filename in glob.glob(file_pattern)
),
axis=0,
ignore_index=True,
).set_index("id")
# initial/final beam
initial = read_all_files("diags/nonlinear_lens_invariants_000000.*")
final = read_all_files("diags/nonlinear_lens_invariants_final.*")
# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)
print("Initial Beam:")
meanH, sigH, meanI, sigI = get_moments(initial)
print(f" meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
atol = 0.0 # a big number
rtol = 1.6 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[meanH, sigH, meanI, sigI],
[7.263202e-02, 4.454371e-02, 9.288060e-02, 8.211506e-02],
rtol=rtol,
atol=atol,
)
print("")
print("Final Beam:")
meanH, sigH, meanI, sigI = get_moments(final)
print(f" meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
atol = 0.0 # a big number
rtol = 1.6 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[meanH, sigH, meanI, sigI],
[7.263202e-02, 4.454371e-02, 9.288060e-02, 8.211506e-02],
rtol=rtol,
atol=atol,
)
# join tables on particle ID, so we can compare the same particle initial->final
beam_joined = final.join(initial, lsuffix="_final", rsuffix="_initial")
# add new columns: dH and dI
beam_joined["dH"] = (beam_joined["H_initial"] - beam_joined["H_final"]).abs()
beam_joined["dI"] = (beam_joined["I_initial"] - beam_joined["I_final"]).abs()
# print(beam_joined)
# particle-wise comparison of H & I initial to final
Hrms = np.sqrt(sigH**2 + meanH**2)
Irms = np.sqrt(sigI**2 + meanI**2)
atol = 4.2e-3 * Hrms
rtol = 0.0 # large number
print()
print(f" atol={atol} (ignored: rtol~={rtol})")
print(f" dH_max={beam_joined['dH'].max()}")
assert np.allclose(beam_joined["dH"], 0.0, rtol=rtol, atol=atol)
atol = 5.7e-3 * Irms
rtol = 0.0
print()
print(f" atol={atol} (ignored: rtol~={rtol})")
print(f" dI_max={beam_joined['dI'].max()}")
assert np.allclose(beam_joined["dI"], 0.0, rtol=rtol, atol=atol)
Solenoid channel
Proton beam undergoing 90 deg X-Y rotation in an ideal solenoid channel.
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 2.4321374875\) m
\(\alpha_\mathrm{x} = 0.0\)
\(\beta_\mathrm{y} = 2.4321374875\) m
\(\alpha_\mathrm{y} = 0.0\)
We use a 250 MeV proton beam with initial unnormalized rms emittance of 1 micron in the horizontal plane, and 2 micron in the vertical plane.
The solenoid magnetic field corresponds to B = 2 T.
The second moments of the particle distribution after the solenoid channel are rotated by 90 degrees: the final horizontal moments should coincide with the initial vertical moments, and vice-versa, 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_solenoid.py
orImpactX executable using an input file:
impactx input_solenoid.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/solenoid/run_solenoid.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 -*-
from impactx import ImpactX, distribution, elements
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 250 MeV proton beam with an initial
# horizontal rms emittance of 1 um and an
# initial vertical rms emittance of 2 um
kin_energy_MeV = 250.0 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.559531175539e-3,
lambdaY=2.205510139392e-3,
lambdaT=1.0e-3,
lambdaPx=6.41218345413e-4,
lambdaPy=9.06819680526e-4,
lambdaPt=1.0e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.extend(
[
monitor,
elements.Sol(ds=3.820395, ks=0.8223219329893234),
monitor,
]
)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/solenoid/input_solenoid.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.559531175539e-3
beam.lambdaY = 2.205510139392e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 6.41218345413e-4
beam.lambdaPy = 9.06819680526e-4
beam.lambdaPt = 1.0e-3
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor sol1 monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
sol1.type = solenoid
sol1.ds = 3.820395
sol1.ks = 0.8223219329893234
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_solenoid.py
examples/solenoid/analysis_solenoid.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.3 * 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.559531175539e-3,
2.205510139392e-3,
1.0e-3,
1.0e-6,
2.0e-6,
1.0e-6,
],
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.3 * 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],
[
2.205510139392e-3,
1.559531175539e-3,
6.404930308742e-3,
2.0e-6,
1.0e-6,
1.0e-6,
],
rtol=rtol,
atol=atol,
)
Drift using a Pole-Face Rotation
A drift that takes place in a rotated frame, using initial and final
applications of a pole rotation (prot
).
We use a 2 GeV electron beam.
The second moments of x, y, and t should be nearly unchanged.
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_rotation.py
orImpactX executable using an input file:
impactx input_rotation.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/rotation/run_rotation.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 -*-
from impactx import ImpactX, distribution, elements
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 2 GeV electron beam with an initial
# unnormalized rms emittance of nm
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 1.0e-9 # used without 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.Waterbag(
lambdaX=4.0e-3,
lambdaY=4.0e-3,
lambdaT=1.0e-3,
lambdaPx=3.0e-4,
lambdaPy=3.0e-4,
lambdaPt=2.0e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
rotated_drift = [
monitor,
elements.PRot(phi_in=0.0, phi_out=-5.0),
elements.Drift(ds=2.0, nslice=1),
elements.PRot(phi_in=-5.0, phi_out=0.0),
monitor,
]
# assign a lattice segment
sim.lattice.extend(rotated_drift)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/rotation/input_rotation.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 4.0e-3
beam.lambdaY = 4.0e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 3.0e-4
beam.lambdaPy = 3.0e-4
beam.lambdaPt = 2.0e-3
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor rotation1 drift1 rotation2 monitor
monitor.type = beam_monitor
monitor.backend = h5
drift1.type = drift
drift1.ds = 2.0
rotation1.type = prot
rotation1.phi_in = 0.0
rotation1.phi_out = -5.0
rotation2.type = prot
rotation2.phi_in = -5.0
rotation2.phi_out = 0.0
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
Analyze
We run the following script to analyze correctness:
Script analysis_rotation.py
examples/rotation/analysis_rotation.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.0e-03,
4.0e-03,
1.0e-03,
1.2e-06,
1.2e-06,
2.0e-06,
],
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],
[
4.061349e-03,
4.042332e-03,
1.002804e-03,
1.2e-06,
1.2e-06,
2.0e-06,
],
rtol=rtol,
atol=atol,
)
Soft-edge solenoid
Proton beam propagating through a 6 m region containing a soft-edge solenoid.
The solenoid model used is the default thin-shell model described in: P. Granum et al, “Efficient calculations of magnetic fields of solenoids for simulations,” NIMA 1034, 166706 (2022) DOI:10.1016/j.nima.2022.166706
The solenoid is a cylindrical current sheet with a length of 1 m and a radius of 0.1667 m, corresponding to an aspect ratio diameter/length = 1/3. The peak magnetic field on-axis is 3 T.
We use a 250 MeV proton beam with initial unnormalized rms emittance of 1 micron in the horizontal plane, and 2 micron in the vertical plane.
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_solenoid_softedge.py
orImpactX executable using an input file:
impactx input_solenoid_softedge.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/solenoid_softedge/run_solenoid_softedge.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Chad Mitchell, Axel Huebl
# 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 = False
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = False
# domain decomposition & space charge mesh
sim.init_grids()
# load a 250 MeV proton beam with an initial
# horizontal rms emittance of 1 um and an
# initial vertical rms emittance of 2 um
kin_energy_MeV = 250.0 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.559531175539e-3,
lambdaY=2.205510139392e-3,
lambdaT=1.0e-3,
lambdaPx=6.41218345413e-4,
lambdaPy=9.06819680526e-4,
lambdaPt=1.0e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# design the accelerator lattice
sol = elements.SoftSolenoid(
ds=6.0,
bscale=1.233482899483985,
cos_coefficients=[
0.350807812299706,
0.323554693720069,
0.260320578919415,
0.182848575294969,
0.106921016050403,
4.409581845710694e-002,
-9.416427163897508e-006,
-2.459452716865687e-002,
-3.272762575737291e-002,
-2.936414401076162e-002,
-1.995780078926890e-002,
-9.102893342953847e-003,
-2.456410658713271e-006,
5.788233017324325e-003,
8.040408292420691e-003,
7.480064552867431e-003,
5.230254569468851e-003,
2.447614547094685e-003,
-1.095525090532255e-006,
-1.614586867387170e-003,
-2.281365457438345e-003,
-2.148709081338292e-003,
-1.522541739363011e-003,
-7.185505862719508e-004,
-6.171194824600157e-007,
4.842109305036943e-004,
6.874508102002901e-004,
6.535550288205728e-004,
4.648795813759210e-004,
2.216564722797528e-004,
-4.100982995210341e-007,
-1.499332112463395e-004,
-2.151538438342482e-004,
-2.044590946652016e-004,
-1.468242784844341e-004,
],
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,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
],
mapsteps=200,
nslice=4,
)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
sim.lattice.extend(
[
monitor,
sol,
monitor,
]
)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/solenoid_softedge/input_solenoid_softedge.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.559531175539e-3
beam.lambdaY = 2.205510139392e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 6.41218345413e-4
beam.lambdaPy = 9.06819680526e-4
beam.lambdaPt = 1.0e-3
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor sol1 monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
sol1.type = solenoid_softedge
sol1.ds = 6.0
sol1.bscale = 1.233482899483985
sol1.mapsteps = 800
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_solenoid_softedge.py
examples/solenoid_softedge/analysis_solenoid_softedge.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Chad Mitchell, Axel Huebl
# 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 = 2.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.559531175539e-3,
2.205510139392e-3,
1.0e-3,
1.0e-6,
2.0e-6,
1.0e-6,
],
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 = 2.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],
[
2.425578906459e-3,
2.654015302646e-3,
9.985897906860e-3,
1.365357890e-6,
1.634641555e-6,
1.000000000e-6,
],
rtol=rtol,
atol=atol,
)
Soft-Edge Quadrupole
This is a modification of the test for a matched electron beam propagating through a stable FODO cell, in which the quadrupoles have been replaced with soft-edge quadrupole elements. The on-axis field profile in this example has been chosen to correspond to the hard-edge limit, so the two tests should coincide.
We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm.
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_quadrupole_softedge.py
orImpactX executable using an input file:
impactx input_quadrupole_softedge.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/quadrupole_softedge/run_quadrupole_softedge.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 = False
# sim.diagnostics = False # benchmarking
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 = 2.0e3 # 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.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
quad1 = elements.SoftQuadrupole(
ds=1.0,
gscale=1.0,
cos_coefficients=[2],
sin_coefficients=[0],
mapsteps=400,
nslice=ns,
)
quad2 = elements.SoftQuadrupole(
ds=1.0,
gscale=-1.0,
cos_coefficients=[2],
sin_coefficients=[0],
mapsteps=200,
nslice=ns,
)
drift1 = elements.Drift(ds=0.25, nslice=ns)
drift2 = elements.Drift(ds=0.5, nslice=ns)
# assign a fodo segment
sim.lattice.extend([monitor, drift1, quad1, drift2, quad2, drift1, monitor])
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/quadrupole_softedge/input_quadrupole_softedge.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 3.9984884770e-5
beam.lambdaY = 3.9984884770e-5
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.6623538760e-5
beam.lambdaPy = 2.6623538760e-5
beam.lambdaPt = 2.0e-3
beam.muxpx = -0.846574929020762
beam.muypy = 0.846574929020762
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 quad1 drift2 quad2 drift1 monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
drift1.type = drift
drift1.ds = 0.25
quad1.type = quadrupole_softedge
quad1.ds = 1.0
quad1.gscale = 1.0
quad1.cos_coefficients = 2.0
quad1.sin_coefficients = 0.0
quad1.mapsteps = 400
drift2.type = drift
drift2.ds = 0.5
quad2.type = quadrupole_softedge
quad2.ds = 1.0
quad2.gscale = -1.0
quad2.cos_coefficients = 2.0
quad2.sin_coefficients = 0.0
quad2.mapsteps = 400
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_quadrupole_softedge.py
examples/quadrupole_softedge/analysis_quadrupole_softedge.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 = 2.2 * 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 = 2.2 * 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.4790118496224206e-005,
7.5357525169680140e-005,
9.9775879288128088e-004,
1.9959539836392703e-009,
2.0175014668882125e-009,
2.0013820380883801e-006,
],
rtol=rtol,
atol=atol,
)
Positron Channel
Acceleration of a positron beam with large (10%) energy spread, from 10 GeV to 2.5 TeV, for parameters based on possible staging of a laser-wakefield accelerator.
The lattice consists of 250 periods, each consisting of a quadrupole triplet followed by 10 GeV energy gain in a uniform field.
We use a 190 pC positron beam with initial normalized rms emittance of 10 nm, rms beam size of 5 microns, and a triangular current pulse with end-to-end pulse length of 0.12 ps.
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_positron.py
orImpactX executable using an input file:
impactx input_positron.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/positron_channel/run_positron.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 = False
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = False
# domain decomposition & space charge mesh
sim.init_grids()
# load a 10 GeV positron beam with an initial
# normalized rms emittance of 10 nm
kin_energy_MeV = 10.0e3 # reference energy
bunch_charge_C = 190.0e-12 # 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.Triangle(
lambdaX=5.054566450e-6,
lambdaY=5.054566450e-6,
lambdaT=8.43732950e-7,
lambdaPx=1.01091329e-7,
lambdaPy=1.01091329e-7,
lambdaPt=1.0e-2,
muxpx=0.0,
muypy=0.0,
mutpt=0.995037190209989,
)
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
period = [
monitor,
elements.ChrQuad(ds=0.1, k=-6.674941, units=1, nslice=ns),
elements.ChrDrift(ds=0.3, nslice=ns),
elements.ChrQuad(ds=0.2, k=6.674941, units=1, nslice=ns),
elements.ChrDrift(ds=0.3, nslice=ns),
elements.ChrQuad(ds=0.1, k=-6.674941, units=1, nslice=ns),
elements.ChrDrift(ds=0.1, nslice=ns),
elements.ChrAcc(ds=1.8, ez=10871.950994502130424, bz=1.0e-12, nslice=ns),
elements.ChrDrift(ds=0.1, nslice=ns),
monitor,
]
sim.lattice.extend(period)
# number of periods through the lattice
sim.periods = 250
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/positron_channel/input_positron.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 10.0e3
beam.charge = 190.0e-12
beam.particle = positron
beam.distribution = triangle
beam.lambdaX = 5.054566450e-6
beam.lambdaY = 5.054566450e-6
beam.lambdaT = 8.43732950e-7
beam.lambdaPx = 1.01091329e-7
beam.lambdaPy = 1.01091329e-7
beam.lambdaPt = 1.0e-2
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.995037190209989
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.periods = 250
lattice.elements = monitor quad1 drift1 quad2 drift1 quad1 drift2 unifacc drift2 monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
quad1.type = quad_chromatic
quad1.ds = 0.1
quad1.k = -6.674941
quad1.units = 1
drift1.type = drift_chromatic
drift1.ds = 0.3
quad2.type = quad_chromatic
quad2.ds = 0.2
quad2.k = 6.674941
quad2.units = 1
drift2.type = drift_chromatic
drift2.ds = 0.1
unifacc.type = uniform_acc_chromatic
unifacc.ds = 1.8
unifacc.ez = 10871.950994502130424
unifacc.bz = 1.0e-12
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
#amr.n_cell = 56 56 48
#geometry.prob_relative = 1.0
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_positron.py
examples/positron_channel/analysis_positron.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
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 = 2.1 * 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],
[
5.05456645029333603e-006,
5.05456645029333603e-006,
8.47941120001532497e-006,
5.1097284000861952e-13,
5.1097284000861952e-13,
8.47941120001532497e-008,
],
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 = 2.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],
[
5.26e-006,
5.26e-006,
8.47941120001532497e-006,
2.0439953822082447e-15,
2.0439953822082447e-15,
3.3919371010764148e-10,
],
rtol=rtol,
atol=atol,
)
Cyclotron
This demonstrates a simple cyclotron as published by Ernest O. Lawrence and M. Stanley Livingston, The Production of High Speed Light Ions Without the Use of High Voltages, Phys. Rev. 40, 19 (1932). DOI: 10.1103/PhysRev.40.19
Run
This example can be run either as:
Python script:
python3 run_cyclotron.py
orImpactX executable using an input file:
impactx input_cyclotron.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/cyclotron/run_cyclotron.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 = False
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = True
# domain decomposition & space charge mesh
sim.init_grids()
# load initial beam
kin_energy_MeV = 4.0e-3 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.0e-3,
lambdaY=1.0e-3,
lambdaT=0.3,
lambdaPx=2.0e-4,
lambdaPy=2.0e-4,
lambdaPt=2.0e-5,
muxpx=-0.0,
muypy=0.0,
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
period = [
monitor,
elements.ChrAcc(ds=0.038, ez=1.12188308693e-4, bz=1.0e-14, nslice=ns),
elements.ExactSbend(ds=0.25, phi=180.0, B=1),
elements.ChrAcc(ds=0.038, ez=1.12188308693e-4, bz=1.0e-14, nslice=ns),
elements.ExactSbend(ds=0.25, phi=180.0, B=1),
monitor,
]
sim.lattice.extend(period)
# number of periods through the lattice
sim.periods = 150
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/cyclotron/input_cyclotron.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 4.0e-3
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.0e-3
beam.lambdaY = 1.0e-3
beam.lambdaT = 0.3
beam.lambdaPx = 2.0e-4
beam.lambdaPy = 2.0e-4
beam.lambdaPt = 2.0e-5
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor half half monitor
lattice.periods = 150
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
half.type = line
half.elements = gap bend
gap.type = uniform_acc_chromatic
gap.ds = 0.038
gap.ez = 1.12188308693e-4
gap.bz = 1.0e-14
bend.type = sbend_exact
bend.ds = 0.25
bend.phi = 180.0
bend.B = 1.0
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_cyclotron.py
examples/cyclotron/analysis_cyclotron.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
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.8 * 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.0e-3,
1.0e-3,
3.0e-1,
2.0e-7,
2.0e-7,
6.0e-6,
],
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 = 2.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.040898e-03,
2.221367e-03,
3.426366e-01,
1.151532e-08,
1.160797e-08,
3.440357e-07,
],
rtol=rtol,
atol=atol,
)
Visualize
Note
TODO :)
Combined Function Bend
A single combined function bending magnet (an ideal sector bend with an upright quadrupole field component added). The magnet parameters are based a single CSBEND element appearing in the ELEGANT input file for the ALS-U lattice.
The beam parameters are based on: C. Steier et al, “Status of the Conceptual Design of ALS-U”, IPAC2017, WEPAB104, DOI:10.18429/JACoW-IPAC2017-WEPAB104 (2017).
A 2 GeV electron bunch with normalized transverse rms emittance of 50 pm undergoes a 3.76 deg bend.
In this test, the initial and final values of \(\lambda_x\), \(\lambda_y\), \(\lambda_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_cfbend.py
orImpactX executable using an input file:
impactx input_cfbend.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/cfbend/run_cfbend.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Chad Mitchell, Axel Huebl
# 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 = False
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = True
# domain decomposition & space charge mesh
sim.init_grids()
# load a 5 GeV electron beam with an initial
# normalized transverse rms emittance of 1 um
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.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.Waterbag(
lambdaX=5.0e-6, # 5 um
lambdaY=8.0e-6, # 8 um
lambdaT=0.0599584916, # 200 ps
lambdaPx=2.5543422003e-9, # exn = 50 pm-rad
lambdaPy=1.5964638752e-9, # eyn = 50 pm-rad
lambdaPt=9.0e-4, # approximately dE/E
muxpx=0.0,
muypy=0.0,
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
bend = [
monitor,
elements.CFbend(ds=0.5, rc=7.613657587094493, k=-7.057403, nslice=ns),
monitor,
]
# assign a lattice segment
sim.lattice.extend(bend)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/cfbend/input_cfbend.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3 #2 GeV
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 5.0e-6 #5 um
beam.lambdaY = 8.0e-6 #8 um
beam.lambdaT = 0.0599584916 #200 ps
beam.lambdaPx = 2.5543422003e-9 #exn = 50 pm-rad
beam.lambdaPy = 1.5964638752e-9 #eyn = 50 pm-rad
beam.lambdaPt = 9.0e-4 #approximately dE/E
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor cfbend1 monitor
lattice.nslice = 25
cfbend1.type = cfbend
cfbend1.ds = 0.5 # projected length 0.5 m, angle 3.76 deg
cfbend1.rc = 7.613657587094493 # bending radius [m]
cfbend1.k = -7.057403 # (upright) quadrupole component [m^(-2)]
monitor.type = beam_monitor
monitor.backend = h5
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_cfbend.py
examples/cfbend/analysis_cfbend.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.8 * 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],
[
5.0e-6,
8.0e-6,
0.0599584916,
1.277171100130637e-14,
1.277171100130637e-14,
5.396264243e-005,
],
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.8 * 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.98301912761436476154e-005,
1.92110147814673522465e-006,
5.99584916026070271843e-002,
3.90166131470905952893e-010,
1.27717110013063679910e-014,
5.396264244141050920556e-005,
],
rtol=rtol,
atol=atol,
)
Ballistic Compression Using a Short RF Element
A 20 MeV electron beam propagates through a short RF element near zero-crossing, inducing a head-tail energy correlation. This is followed by ballistic motion in a drift, which is used to compress the rms bunch length from 16 ps to 10 ps (compression of 5/3).
The beam is not exactly on-crest (phase = -89.5 deg), so there is an energy gain of 4.5 MeV.
The transverse emittance is sufficiently small that the horizontal and verticle beam size are essentially unchanged. Due to RF curvature, there is some growth of the longitudinal emittance.
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_compression.py
orImpactX executable using an input file:
impactx input_compression.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/compression/run_compression.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 -*-
from impactx import ImpactX, distribution, elements
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 250 MeV proton beam with an initial
# unnormalized rms emittance of 1 mm-mrad in all
# three phase planes
kin_energy_MeV = 20.0 # 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.Waterbag(
lambdaX=0.5e-3,
lambdaY=0.5e-3,
lambdaT=5.0e-3,
lambdaPx=1.0e-5,
lambdaPy=1.0e-5,
lambdaPt=4.0e-6,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.append(monitor)
# Short RF cavity element
shortrf1 = elements.ShortRF(V=1000.0, freq=1.3e9, phase=-89.5)
# Drift element
drift1 = elements.Drift(ds=1.7)
sim.lattice.extend([shortrf1, drift1])
sim.lattice.append(monitor)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/compression/input_compression.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 20.0
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 0.5e-3
beam.lambdaY = 0.5e-3
beam.lambdaT = 5.0e-3
beam.lambdaPx = 1.0e-5
beam.lambdaPy = 1.0e-5
beam.lambdaPt = 4.0e-6
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor shortrf1 drift1 monitor
monitor.type = beam_monitor
monitor.backend = h5
shortrf1.type = shortrf
shortrf1.V = 1000.0
shortrf1.freq = 1.3e9
shortrf1.phase = -89.5
drift1.type = drift
drift1.ds = 1.7
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
Analyze
We run the following script to analyze correctness:
Script analysis_compression.py
examples/compression/analysis_compression.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)
# openPMD data series at the beam monitors
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
# first and last step
final_step = list(series.iterations)[-1]
first_it = series.iterations[1]
final_it = series.iterations[final_step]
# initial beam & reference particle gamma
initial = first_it.particles["beam"].to_df()
initial_gamma_ref = first_it.particles["beam"].get_attribute("gamma_ref")
# final beam & reference particle gamma
final = final_it.particles["beam"].to_df()
final_gamma_ref = final_it.particles["beam"].get_attribute("gamma_ref")
# 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}"
)
print(f" gamma={initial_gamma_ref:e}")
atol = 0.0 # ignored
rtol = 1.9 * 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, initial_gamma_ref],
[
5.0e-04,
5.0e-04,
5.0e-03,
4.952764e-09,
5.028325e-09,
1.997821e-08,
40.1389432485322889,
],
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}"
)
print(f" gamma={final_gamma_ref:e}")
atol = 0.0 # ignored
rtol = 1.9 * 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, final_gamma_ref],
[
5.004995e-04,
5.005865e-04,
3.033949e-03,
4.067876e-09,
4.129937e-09,
6.432081e-05,
48.8654787469061860,
],
rtol=rtol,
atol=atol,
)
Test of a Transverse Kicker
This test applies two transverse momentum kicks, first in the horizontal direction (2 mrad) and then in the vertical direction (3 mrad).
We use a 2 GeV electron beam.
The second beam moments should be unchanged, but the first beam moments corresponding to \(p_x\) and \(p_y\) should change according to the size of the kick.
In this test, the initial and final values of \(\lambda_x\), \(\lambda_y\), \(\lambda_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_kicker.py
orImpactX executable using an input file:
impactx input_kicker.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/kicker/run_kicker.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 -*-
from impactx import ImpactX, distribution, elements
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 2 GeV electron beam with an initial
# unnormalized rms emittance of nm
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 1.0e-9 # used without 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.Waterbag(
lambdaX=4.0e-3,
lambdaY=4.0e-3,
lambdaT=1.0e-3,
lambdaPx=3.0e-4,
lambdaPy=3.0e-4,
lambdaPt=2.0e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
kicklattice = [
monitor,
elements.Kicker(xkick=2.0e-3, ykick=0.0, units="dimensionless"),
elements.Kicker(xkick=0.0, ykick=3.0e-3, units="dimensionless"),
monitor,
]
# assign a lattice
sim.lattice.extend(kicklattice)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/kicker/input_kicker.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 4.0e-3
beam.lambdaY = 4.0e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 3.0e-4
beam.lambdaPy = 3.0e-4
beam.lambdaPt = 2.0e-3
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor hkick vkick monitor
monitor.type = beam_monitor
monitor.backend = h5
hkick.type = kicker
hkick.xkick = 2.0e-3 # 2 mrad horizontal kick
hkick.ykick = 0.0
vkick.type = kicker
vkick.xkick = 0.0
vkick.ykick = 3.0e-3 # 3 mrad vertical kick
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
Analyze
We run the following script to analyze correctness:
Script analysis_kicker.py
examples/kicker/analysis_kicker.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 describe, 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, meanpx, meanpy
"""
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
meanpx = describe(beam["momentum_x"]).mean
meanpy = describe(beam["momentum_y"]).mean
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, meanpx, meanpy)
# 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, meanpx, meanpy = get_moments(
initial
)
print(
f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e} meanpx={meanpx:e} meanpy={meanpy:e}"
)
print(
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)
atol = 0.0 # ignored
rtol = 5.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],
[
4.017554e-03,
4.017044e-03,
9.977588e-04,
1.197572e-06,
1.210501e-06,
2.001382e-06,
],
rtol=rtol,
atol=atol,
)
atol = rtol * emittance_x / sigx # relative to rms beam size
assert np.allclose(
[meanpx, meanpy],
[
0.0,
0.0,
],
atol=atol,
)
print("")
print("Final Beam:")
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, meanpx, meanpy = get_moments(
final
)
print(
f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e} meanpx={meanpx:e} meanpy={meanpy:e}"
)
print(
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)
atol = 0.0 # ignored
rtol = 5.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],
[
4.017554e-03,
4.017044e-03,
9.977588e-04,
1.197572e-06,
1.210501e-06,
2.001382e-06,
],
rtol=rtol,
atol=atol,
)
atol = rtol * emittance_x / sigx # relative to rms beam size
print(f" atol~={atol}")
assert np.allclose(
[meanpx, meanpy],
[
2.0e-3,
3.0e-3,
],
atol=atol,
)
Thin Dipole
This test involves tracking a 5 MeV proton beam through a 90 degree sector bend, using two different methods:
#. Using a sequence of drifts and thin kicks as described by: G. Ripken and F. Schmidt, “A Symplectic Six-Dimensional Thin-Lens Formalism for Tracking,” CERN/SL/95-12 (1995).
#. Using the exact nonlinear transfer map for a sector bend as described by: D. Bruhwiler et al, “Symplectic Propagation of the Map, Tangent Map and Tangent Map Derivative through Quadrupole and Combined-Function Dipole Magnets without Truncation,” THP41C, EPAC98, pp. 1171-1173 (1998).
To compare the two models, the beam is tracked using 1), followed by the inverse of 2). The beam moments and emittances should coincide with those of the initial distribution.
Run
This example can be run either as:
Python script:
python3 run_thin_dipole.py
orImpactX executable using an input file:
impactx input_thin_dipole.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/thin_dipole/run_thin_dipole.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 = False
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = True
# domain decomposition & space charge mesh
sim.init_grids()
# load initial beam
kin_energy_MeV = 5.0 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.0e-3,
lambdaY=1.0e-3,
lambdaT=0.3,
lambdaPx=2.0e-4,
lambdaPy=2.0e-4,
lambdaPt=2.0e-4,
muxpx=-0.0,
muypy=0.0,
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
segment = [
elements.Drift(ds=0.003926990816987, nslice=ns),
elements.ThinDipole(theta=0.45, rc=1.0),
elements.Drift(ds=0.003926990816987, nslice=ns),
]
bend = 200 * segment
inverse_bend = elements.ExactSbend(ds=-1.570796326794897, phi=-90.0)
sim.lattice.append(monitor)
sim.lattice.extend(bend)
sim.lattice.append(inverse_bend)
sim.lattice.append(monitor)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/thin_dipole/input_thin_dipole.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 5.0
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.0e-3
beam.lambdaY = 1.0e-3
beam.lambdaT = 0.3
beam.lambdaPx = 2.0e-4
beam.lambdaPy = 2.0e-4
beam.lambdaPt = 2.0e-4
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor bend inverse_bend monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
#90 degree sbend using drift-kick-drift:
bend.type = line
bend.elements = dr kick dr
bend.repeat = 200
dr.type = drift
dr.ds = 0.003926990816987
kick.type = thin_dipole
kick.theta = 0.45
kick.rc = 1.0
#inverse of a 90 degree sbend using exact nonlinear map:
inverse_bend.type = sbend_exact
inverse_bend.ds = -1.570796326794897
inverse_bend.phi = -90.0
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_thin_dipole.py
examples/thin_dipole/analysis_thin_dipole.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
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.8 * 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.0e-3,
1.0e-3,
3.0e-1,
2.0e-7,
2.0e-7,
6.0e-5,
],
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.8 * 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.0e-3,
1.0e-3,
3.0e-1,
2.0e-7,
2.0e-7,
6.0e-5,
],
rtol=rtol,
atol=atol,
)
Visualize
Note
TODO :)
Aperture Collimation
Proton beam undergoing collimation by a rectangular boundary aperture.
We use a 250 MeV proton beam with a horizontal rms beam size of 1.56 mm and a vertical rms beam size of 2.21 mm.
After a short drift of 0.123, the beam is scraped by a 1 mm x 1.5 mm rectangular aperture.
In this test, the initial values of \(\sigma_x\), \(\sigma_y\), \(\sigma_t\), \(\epsilon_x\), \(\epsilon_y\), and \(\epsilon_t\) must agree with nominal values. The test fails if:
any of the final coordinates for the valid (not lost) particles lie outside the aperture boundary or
any of the lost particles are inside the aperture boundary or
if the sum of lost and kept particles is not equal to the initial particles or
if the recorded position \(s\) for the lost particles does not coincide with the drift distance.
Run
This example can be run as a Python script (python3 run_aperture.py
) or with an app with an input file (impactx input_aperture.in
).
Each can also be prefixed with an MPI executor, such as mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/aperture/run_aperture.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
import amrex.space3d as amr
from impactx import ImpactX, distribution, elements
# work-around for https://github.com/ECP-WarpX/impactx/issues/499
pp_amrex = amr.ParmParse("amrex")
pp_amrex.add("the_arena_is_managed", 1)
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
sim.particle_lost_diagnostics_backend = "h5"
# domain decomposition & space charge mesh
sim.init_grids()
# load a 250 MeV proton beam with an initial
# horizontal rms emittance of 1 um and an
# initial vertical rms emittance of 2 um
kin_energy_MeV = 250.0 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.559531175539e-3,
lambdaY=2.205510139392e-3,
lambdaT=1.0e-3,
lambdaPx=6.41218345413e-4,
lambdaPy=9.06819680526e-4,
lambdaPt=1.0e-3,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.extend(
[
monitor,
elements.Drift(0.123),
elements.Aperture(xmax=1.0e-3, ymax=1.5e-3, shape="rectangular"),
monitor,
]
)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/aperture/input_aperture.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 1.559531175539e-3
beam.lambdaY = 2.205510139392e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 6.41218345413e-4
beam.lambdaPy = 9.06819680526e-4
beam.lambdaPt = 1.0e-3
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift collimator monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
drift.type = drift
drift.ds = 0.123
collimator.type = aperture
collimator.shape = rectangular
collimator.xmax = 1.0e-3
collimator.ymax = 1.5e-3
# work-around for https://github.com/ECP-WarpX/impactx/issues/499
amrex.the_arena_is_managed = 1
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
diag.backend = h5
Analyze
We run the following script to analyze correctness:
Script analysis_aperture.py
examples/aperture/analysis_aperture.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()
series_lost = io.Series("diags/openPMD/particles_lost.h5", io.Access.read_only)
particles_lost = series_lost.iterations[0].particles["beam"].to_df()
# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
# we lost particles in apertures
assert num_particles > len(final)
assert num_particles == len(particles_lost) + 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.8 * 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.559531175539e-3,
2.205510139392e-3,
1.0e-3,
1.0e-6,
2.0e-6,
1.0e-6,
],
rtol=rtol,
atol=atol,
)
# particle-wise comparison against the rectangular aperture boundary
xmax = 1.0e-3
ymax = 1.5e-3
# kept particles
dx = abs(final["position_x"]) - xmax
dy = abs(final["position_y"]) - ymax
print()
print(f" x_max={final['position_x'].max()}")
print(f" x_min={final['position_x'].min()}")
assert np.less_equal(dx.max(), 0.0)
print(f" y_max={final['position_y'].max()}")
print(f" y_min={final['position_y'].min()}")
assert np.less_equal(dy.max(), 0.0)
# lost particles
dx = abs(particles_lost["position_x"]) - xmax
dy = abs(particles_lost["position_y"]) - ymax
print()
print(f" x_max={particles_lost['position_x'].max()}")
print(f" x_min={particles_lost['position_x'].min()}")
assert np.greater_equal(dx.max(), 0.0)
print(f" y_max={particles_lost['position_y'].max()}")
print(f" y_min={particles_lost['position_y'].min()}")
assert np.greater_equal(dy.max(), 0.0)
# check that s is set correctly
lost_at_s = particles_lost["s_lost"]
drift_s = np.ones_like(lost_at_s) * 0.123
assert np.allclose(lost_at_s, drift_s)
Cold Beam in a FODO Channel with RF Cavities (and Space Charge)
This example is based on the subsection of the same name in: R. D. Ryne et al, “A Test Suite of Space-Charge Problems for Code Benchmarking”, in Proc. EPAC2004, Lucerne, Switzerland.
See additional documentation in: C. E. Mitchell et al, “ImpactX Modeling of Benchmark Tests for Space Charge Validation”, in Proc. HB2023, Geneva, Switzerland.
A cold (zero momentum spread), uniform density, 250 MeV, 143 pC proton bunch propagates in a FODO lattice with 700 MHz RF cavities added for longitudinal confinement. The on-axis profile of the RF electric field is given by:
The beam is matched to the 3D focusing, with space charge, using the rms envelope equations.
The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. This is tested using the second moments of the distribution.
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_rf_SC.py
orImpactX executable using an input file:
impactx input_fodo_rf_SC.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/epac2004_benchmarks/run_fodo_rf_SC.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 -*-
from impactx import ImpactX, distribution, elements
sim = ImpactX()
# set numerical parameters and IO control
sim.n_cell = [56, 56, 64]
sim.particle_shape = 2 # B-spline order
sim.space_charge = True
sim.dynamic_size = True
sim.prob_relative = [4.0]
# beam diagnostics
sim.slice_step_diagnostics = False
# domain decomposition & space charge mesh
sim.init_grids()
# beam parameters
kin_energy_MeV = 250.0 # reference energy
bunch_charge_C = 1.42857142857142865e-10 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Kurth6D(
lambdaX=9.84722273e-4,
lambdaY=6.96967278e-4,
lambdaT=4.486799242214e-03,
lambdaPx=0.0,
lambdaPy=0.0,
lambdaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.append(monitor)
# Quad elements
fquad = elements.Quad(ds=0.15, k=2.4669749766168163, nslice=6)
dquad = elements.Quad(ds=0.3, k=-2.4669749766168163, nslice=12)
# Drift element
dr = elements.Drift(ds=0.1, nslice=4)
# RF cavity elements
gapa1 = elements.RFCavity(
ds=1.0,
escale=0.042631556991578,
freq=7.0e8,
phase=45.0,
cos_coefficients=[
0.120864178375839,
-0.044057987631337,
-0.209107290958498,
-0.019831226655815,
0.290428111491964,
0.381974267375227,
0.276801212694382,
0.148265085353012,
0.068569351192205,
0.0290155855315696,
0.011281649986680,
0.004108501632832,
0.0014277644197320,
0.000474212125404,
0.000151675768439,
0.000047031436898,
0.000014154595193,
4.154741658e-6,
1.191423909e-6,
3.348293360e-7,
9.203061700e-8,
2.515007200e-8,
6.478108000e-9,
1.912531000e-9,
2.925600000e-10,
],
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,
],
mapsteps=100,
nslice=4,
)
gapb1 = elements.RFCavity(
ds=1.0,
escale=0.042631556991578,
freq=7.0e8,
phase=-1.0,
cos_coefficients=[
0.120864178375839,
-0.044057987631337,
-0.209107290958498,
-0.019831226655815,
0.290428111491964,
0.381974267375227,
0.276801212694382,
0.148265085353012,
0.068569351192205,
0.0290155855315696,
0.011281649986680,
0.004108501632832,
0.0014277644197320,
0.000474212125404,
0.000151675768439,
0.000047031436898,
0.000014154595193,
4.154741658e-6,
1.191423909e-6,
3.348293360e-7,
9.203061700e-8,
2.515007200e-8,
6.478108000e-9,
1.912531000e-9,
2.925600000e-10,
],
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,
],
mapsteps=100,
nslice=4,
)
lattice_no_drifts = [fquad, gapa1, dquad, gapb1, fquad]
# set first lattice element
sim.lattice.append(lattice_no_drifts[0])
# intersperse all remaining elements of the lattice with a drift element
for element in lattice_no_drifts[1:]:
sim.lattice.extend([dr, element])
sim.lattice.append(monitor)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/epac2004_benchmarks/input_fodo_rf_SC.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 250.0
beam.charge = 1.42857142857142865e-10
beam.particle = proton
beam.distribution = kurth6d
beam.lambdaX = 9.84722273e-4
beam.lambdaY = 6.96967278e-4
beam.lambdaT = 4.486799242214e-03
beam.lambdaPx = 0.0
beam.lambdaPy = 0.0
beam.lambdaPt = 0.0
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor fquad dr gapa1 dr dquad dr gapb1 dr fquad monitor
monitor.type = beam_monitor
monitor.backend = h5
dr.type = drift
dr.ds = 0.1
dr.nslice = 4
fquad.type = quad
fquad.ds = 0.15
fquad.k = 2.4669749766168163
fquad.nslice = 6
dquad.type = quad
dquad.ds = 0.3
dquad.k = -2.4669749766168163
dquad.nslice = 12
gapa1.type = rfcavity
gapa1.ds = 1.0
gapa1.escale = 0.042631556991578
gapa1.freq = 7.0e8
gapa1.phase = 45.0
gapa1.mapsteps = 100
gapa1.nslice = 10
gapa1.cos_coefficients = \
0.120864178375839 \
-0.044057987631337 \
-0.209107290958498 \
-0.019831226655815 \
0.290428111491964 \
0.381974267375227 \
0.276801212694382 \
0.148265085353012 \
0.068569351192205 \
0.0290155855315696 \
0.011281649986680 \
0.004108501632832 \
0.0014277644197320 \
0.000474212125404 \
0.000151675768439 \
0.000047031436898 \
0.000014154595193 \
4.154741658e-6 \
1.191423909e-6 \
3.348293360e-7 \
9.203061700e-8 \
2.515007200e-8 \
6.478108000e-9 \
1.912531000e-9 \
2.925600000e-10
gapa1.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
gapb1.type = rfcavity
gapb1.ds = 1.0
gapb1.escale = 0.042631556991578
gapb1.freq = 7.0e8
gapb1.phase = -1.0
gapb1.mapsteps = 100
gapb1.nslice = 10
gapb1.cos_coefficients = \
0.120864178375839 \
-0.044057987631337 \
-0.209107290958498 \
-0.019831226655815 \
0.290428111491964 \
0.381974267375227 \
0.276801212694382 \
0.148265085353012 \
0.068569351192205 \
0.0290155855315696 \
0.011281649986680 \
0.004108501632832 \
0.0014277644197320 \
0.000474212125404 \
0.000151675768439 \
0.000047031436898 \
0.000014154595193 \
4.154741658e-6 \
1.191423909e-6 \
3.348293360e-7 \
9.203061700e-8 \
2.515007200e-8 \
6.478108000e-9 \
1.912531000e-9 \
2.925600000e-10
gapb1.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.particle_shape = 2
algo.space_charge = true
amr.n_cell = 56 56 64
geometry.prob_relative = 4.0
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_fodo_rf_SC.py
examples/epac2004_benchmarks/analysis_fodo_rf_SC.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}")
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],
[9.84722273e-4, 6.96967278e-4, 4.486799242214e-03],
rtol=rtol,
atol=atol,
)
print(
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)
atol = 4.0e-8
print(f" atol={atol}")
assert np.allclose(
[emittance_x, emittance_y, emittance_t],
[0.0, 0.0, 0.0],
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}")
atol = 0.0 # ignored
rtol = 3.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],
[9.84722273e-4, 6.96967278e-4, 4.486799242214e-03],
rtol=rtol,
atol=atol,
)
print(
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)
atol = 4.0e-8
print(f" atol={atol}")
assert np.allclose(
[emittance_x, emittance_y, emittance_t],
[0.0, 0.0, 0.0],
atol=atol,
)
Thermal Beam in a Constant Focusing Channel (with Space Charge)
This example is based on the subsection of the same name in: R. D. Ryne et al, “A Test Suite of Space-Charge Problems for Code Benchmarking”, in Proc. EPAC2004, Lucerne, Switzerland.
See additional documentation in: C. E. Mitchell et al, “ImpactX Modeling of Benchmark Tests for Space Charge Validation”, in Proc. HB2023, Geneva, Switzerland.
This example illustrates a stationary solution of the Vlasov-Poisson equations with spherical symmetry (in the beam rest frame). The distribution represents a thermal equilibrium of the form:
where \(C\) and \(kT\) are constants, and \(H\) denotes the self-consistent Hamiltonian with space charge.
In this example, a 0.1 MeV, 143 pC proton bunch with \(kT=36\times 10^{-6}\) propagates in a constant focusing lattice with 3D isotropic focusing. (The isotropy is exact in the beam rest frame.)
The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. This is tested using the second moments of the distribution.
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_thermal.py
orImpactX executable using an input file:
impactx input_thermal.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/epac2004_benchmarks/run_thermal.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 -*-
from impactx import ImpactX, distribution, elements
sim = ImpactX()
# set numerical parameters and IO control
sim.n_cell = [56, 56, 64]
sim.particle_shape = 2 # B-spline order
sim.space_charge = True
sim.dynamic_size = True
sim.prob_relative = [4.0]
# beam diagnostics
sim.slice_step_diagnostics = False
# domain decomposition & space charge mesh
sim.init_grids()
# beam parameters
kin_energy_MeV = 0.1 # reference energy
bunch_charge_C = 1.4285714285714285714e-10 # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Thermal(
k=6.283185307179586,
kT=36.0e-6,
kT_halo=36.0e-6,
normalize=0.41604661,
normalize_halo=0.0,
w_halo=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.append(monitor)
constf = elements.Constf(
ds=10.0,
kx=6.283185307179586,
ky=6.283185307179586,
kt=6.283185307179586,
nslice=400,
)
# set first lattice element
sim.lattice.append(constf)
sim.lattice.append(monitor)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/epac2004_benchmarks/input_thermal.in
.###############################################################################
# Particle Beam(s)
###############################################################################
#beam.npart = 100000000 #full resolution
beam.npart = 10000
beam.units = static
beam.kin_energy = 0.1
beam.charge = 1.4285714285714285714e-10
beam.particle = proton
beam.distribution = thermal
beam.k = 6.283185307179586
beam.kT = 36.0e-6
beam.normalize = 0.41604661
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor constf1 monitor
monitor.type = beam_monitor
monitor.backend = h5
constf1.type = constf
constf1.ds = 10.0
constf1.kx = 6.283185307179586
constf1.ky = 6.283185307179586
constf1.kt = 6.283185307179586
constf1.nslice = 400 #full resolution
#constf1.nslice = 50
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = true
#amr.n_cell = 128 128 128 #full resolution
amr.n_cell = 64 64 64
geometry.prob_relative = 3.0
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_thermal.py
examples/epac2004_benchmarks/analysis_thermal.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.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],
[
2.569162e-03,
2.569557e-03,
1.757951e-01,
1.540773e-05,
1.541883e-05,
1.538814e-05,
],
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 = 6.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],
[
2.569162e-03,
2.569557e-03,
1.757951e-01,
1.540773e-05,
1.541883e-05,
1.538814e-05,
],
rtol=rtol,
atol=atol,
)
Bithermal Beam in a Constant Focusing Channel (with Space Charge)
This example is based on the subsection of the same name in: R. D. Ryne et al, “A Test Suite of Space-Charge Problems for Code Benchmarking”, in Proc. EPAC2004, Lucerne, Switzerland.
See additional documentation in: C. E. Mitchell et al, “ImpactX Modeling of Benchmark Tests for Space Charge Validation”, in Proc. HB2023, Geneva, Switzerland.
This example illustrates a stationary solution of the Vlasov-Poisson equations with spherical symmetry (in the beam rest frame). It provides a self-consistent model of a 3D bunch with a nontrivial core-halo distribution.
The distribution represents a bithermal stationary distribution of the form:
where \(c_j\), \(kT_j\) \((j=1,2)\) are constants, and \(H\) denotes the self-consistent Hamiltonian with space charge.
In this example, a 0.1 MeV, 143 pC proton bunch with \(kT_1=36\times 10^{-6}\) and \(kT_1=900\times 10^{-6}\) propagates in a constant focusing lattice with 3D isotropic focusing. (The isotropy is exact in the beam rest frame.) 5% of the total charge lies in the second (halo) population.
The particle distribution should remain unchanged, to within the level expected due to numerical particle noise. This is tested using the second moments of the distribution.
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_bithermal.py
orImpactX executable using an input file:
impactx input_bithermal.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/epac2004_benchmarks/run_bithermal.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 -*-
from impactx import ImpactX, distribution, elements
sim = ImpactX()
# set numerical parameters and IO control
# sim.n_cell = [128, 128, 128] # full resolution
sim.n_cell = [64, 64, 64]
sim.particle_shape = 2 # B-spline order
sim.space_charge = True
sim.dynamic_size = True
sim.prob_relative = [3.0]
# beam diagnostics
sim.slice_step_diagnostics = False
# domain decomposition & space charge mesh
sim.init_grids()
# beam parameters
kin_energy_MeV = 0.1 # reference energy
bunch_charge_C = 1.4285714285714285714e-10 # used with space charge
# npart = 100000000 # full resolution
npart = 10000 # number of macro particles
# 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)
# particle bunch
distr = distribution.Thermal(
k=6.283185307179586,
kT=36.0e-6,
kT_halo=900.0e-6,
normalize=0.54226,
normalize_halo=0.08195,
halo=0.05,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
# design the accelerator lattice
sim.lattice.append(monitor)
constf = elements.ConstF(
ds=10.0,
kx=6.283185307179586,
ky=6.283185307179586,
kt=6.283185307179586,
# nslice=400, # full resolution
nslice=50,
)
sim.lattice.append(constf)
sim.lattice.append(monitor)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/epac2004_benchmarks/input_bithermal.in
.###############################################################################
# Particle Beam(s)
###############################################################################
#beam.npart = 100000000 #full resolution
beam.npart = 10000
beam.units = static
beam.kin_energy = 0.1
beam.charge = 1.4285714285714285714e-10
beam.particle = proton
beam.distribution = thermal
beam.k = 6.283185307179586
beam.kT = 36.0e-6
beam.kT_halo = 900.0e-6
beam.halo = 0.05
beam.normalize = 0.54226
beam.normalize_halo = 0.08195
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor constf1 monitor
monitor.type = beam_monitor
monitor.backend = h5
constf1.type = constf
constf1.ds = 10.0
constf1.kx = 6.283185307179586
constf1.ky = 6.283185307179586
constf1.kt = 6.283185307179586
#constf1.nslice = 400 #full resolution
constf1.nslice = 50
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = true
#amr.n_cell = 128 128 128 #full resolution
amr.n_cell = 64 64 64
geometry.prob_relative = 3.0
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_bithermal.py
examples/epac2004_benchmarks/analysis_bithermal.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 = 4.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],
[
2.751162e-03,
2.751725e-03,
1.884003e-01,
2.449966e-05,
2.451077e-05,
2.444195e-05,
],
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 = 6000 * 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],
[
2.751162e-03,
2.751725e-03,
1.884003e-01,
2.449966e-05,
2.451077e-05,
2.444195e-05,
],
rtol=rtol,
atol=atol,
)
Visualize
You can run the following script to visualize the initial and final beam distribution:
Script plot_bithermal.py
examples/fodo/plot_bithermal.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import argparse
from math import pi
import matplotlib.pyplot as plt
import numpy as np
import openpmd_api as io
# options to run this script
parser = argparse.ArgumentParser(description="Plot the Bithermal benchmark.")
parser.add_argument(
"--save-png", action="store_true", help="non-interactive run: save to PNGs"
)
args = parser.parse_args()
# initial/final beam
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
last_step = list(series.iterations)[-1]
initial_beam = series.iterations[1].particles["beam"].to_df()
final_beam = series.iterations[last_step].particles["beam"].to_df()
# Constants
w1 = 0.95
w2 = 0.05
bg = 0.0146003
Min = 0.0
Max = 0.025
Np = 100000001
n = 300
# Function for radius calculation
def r(x, y, z):
return np.sqrt(x**2 + y**2 + z**2)
# Calculate radius and bin data
initial_radii = r(
bg * initial_beam["position_t"],
initial_beam["position_x"],
initial_beam["position_y"],
)
initial_hist, bin_edges = np.histogram(initial_radii, bins=n, range=(Min, Max))
initial_bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
final_radii = r(
bg * final_beam["position_t"], final_beam["position_x"], final_beam["position_y"]
)
final_hist, _ = np.histogram(final_radii, bins=n, range=(Min, Max))
# dr (m)
initial_r = initial_hist / (Np * (bin_edges[1] - bin_edges[0]))
final_r = final_hist / (Np * (bin_edges[1] - bin_edges[0]))
# Plotting
plt.figure(figsize=(10, 6))
plt.xscale("linear")
plt.yscale("log")
plt.xlim([Min, Max])
plt.ylim([0.1, 1.6e6])
plt.xlabel("r (m)", fontsize=30)
plt.xticks(fontsize=25)
plt.yticks(fontsize=25)
plt.grid(True)
# Plot the data
plt.plot(
initial_bin_centers,
initial_r / (4.0 * pi * (initial_bin_centers) ** 2),
label="Initial beam",
linewidth=2,
)
plt.plot(
initial_bin_centers,
final_r / (4.0 * pi * (initial_bin_centers) ** 2),
label="Final beam",
linewidth=2,
linestyle="dotted",
)
# Show plot
plt.legend(fontsize=20)
plt.tight_layout()
if args.save_png:
plt.savefig("bithermal.png")
else:
plt.show()

Initial and final beam distribution when running with full resolution (see inline comments in the input file/script). The bithermal distribution should stay static in this test.
15 Stage Laser-Plasma Accelerator Surrogate
This example models an electron beam accelerated through fifteen stages of laser-plasma accelerators with ideal plasma lenses providing the focusing between stages. For more details, see:
Sandberg R T, Lehe R, Mitchell C E, Garten M, Myers A, Qiang J, Vay J-L and Huebl A. Synthesizing Particle-in-Cell Simulations Through Learning and GPU Computing for Hybrid Particle Accelerator Beamlines. Proc. of Platform for Advanced Scientific Computing (PASC’24), submitted, 2024. arXiv:2402.17248
Sandberg R T, Lehe R, Mitchell C E, Garten M, Qiang J, Vay J-L and Huebl A. Hybrid Beamline Element ML-Training for Surrogates in the ImpactX Beam-Dynamics Code. 14th International Particle Accelerator Conference (IPAC’23), WEPA101, 2023. DOI:10.18429/JACoW-IPAC2023-WEPA101
A schematic with more information can be seen in the figure below:

Schematic of the 15 stages of laser-plasma accelerators.
The laser-plasma accelerator elements are modeled with neural networks as surrogates.
These networks are trained beforehand.
In this example, pre-trained neural networks are downloaded from a Zenodo archive and saved in the models
directory.
For more about how these neural network surrogate models were created,
see this description of a workflow for training neural networks from WarpX simulation data.
We use a 1 GeV electron beam with initial normalized rms emittance of 1 mm-mrad.
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 only be run with Python:
Python script:
python3 run_ml_surrogate_15_stage.py
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.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 argparse
import amrex.space3d as amr
try:
import cupy as cp
cupy_available = True
except ImportError:
cupy_available = False
import sys
import numpy as np
import scipy.optimize as opt
from impactx import (
Config,
CoordSystem,
ImpactX,
ImpactXParIter,
coordinate_transformation,
distribution,
elements,
)
from surrogate_model_definitions import surrogate_model
try:
import torch
except ImportError:
print("Warning: Cannot import PyTorch. Skipping test.")
sys.exit(0)
import zipfile
from urllib import request
parser = argparse.ArgumentParser()
parser.add_argument(
"--num_particles",
"-N",
type=int,
default=100000,
help="number of particles to use in beam",
)
parser.add_argument(
"--N_stages",
"-ns",
type=int,
default=15,
choices=range(1, 16),
help="number of LPA stages to simulate",
)
args = parser.parse_args()
if Config.have_gpu and cupy_available:
array = cp.array
stack = cp.stack
sqrt = cp.sqrt
device = torch.device("cuda")
if Config.gpu_backend == "SYCL":
print("Warning: SYCL GPU backend not yet implemented for Python")
else:
array = np.array
stack = np.stack
sqrt = np.sqrt
device = None
if device is not None:
print(f"device={device}")
else:
print("device set to default, cpu")
N_stage = args.N_stages
tune_by_x_or_y = "x"
npart = args.num_particles
ebeam_lpa_z0 = -107e-6
L_plasma = 0.28
L_transport = 0.03
L_stage_period = L_plasma + L_transport
drift_after_LPA = 43e-6
L_surrogate = abs(ebeam_lpa_z0) + L_plasma + drift_after_LPA
def download_and_unzip(url, data_dir):
request.urlretrieve(url, data_dir)
with zipfile.ZipFile(data_dir, "r") as zip_dataset:
zip_dataset.extractall()
data_url = "https://zenodo.org/records/10810754/files/models.zip?download=1"
download_and_unzip(data_url, "models.zip")
model_list = [
surrogate_model(f"models/beam_stage_{stage_i}_model.pt", device=device)
for stage_i in range(N_stage)
]
pp_amrex = amr.ParmParse("amrex")
pp_amrex.add("the_arena_init_size", 0)
pp_amrex.add("the_device_arena_init_size", 0)
sim = ImpactX()
# set numerical parameters and IO control
sim.particle_shape = 2 # B-spline order
sim.space_charge = False
sim.diagnostics = True # 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 1 nm
ref_u = 1957
energy_gamma = np.sqrt(1 + ref_u**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)
ref.z = ebeam_lpa_z0
pc = sim.particle_container()
distr = distribution.Gaussian(
lambdaX=0.75e-6,
lambdaY=0.75e-6,
lambdaT=0.1e-6,
lambdaPx=1.33 / energy_gamma,
lambdaPy=1.33 / energy_gamma,
lambdaPt=1e-8,
muxpx=0.0,
muypy=0.0,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
n_slice = 1
class LPASurrogateStage(elements.Programmable):
def __init__(self, stage_i, surrogate_model, surrogate_length, stage_start):
elements.Programmable.__init__(self)
self.stage_i = stage_i
self.surrogate_model = surrogate_model
self.surrogate_length = surrogate_length
self.stage_start = stage_start
self.push = self.surrogate_push
self.ds = surrogate_length
def surrogate_push(self, pc, step):
ref_part = pc.ref_particle()
ref_z_i = ref_part.z
ref_z_i_LPA = ref_z_i - self.stage_start
ref_z_f = ref_z_i + self.surrogate_length
ref_part_tensor = torch.tensor(
[
ref_part.x,
ref_part.y,
ref_z_i_LPA,
ref_part.px,
ref_part.py,
ref_part.pz,
],
device=device,
dtype=torch.float64,
)
ref_beta_gamma = torch.sqrt(torch.sum(ref_part_tensor[3:] ** 2))
ref_beta_gamma = ref_beta_gamma.to(device)
with torch.no_grad():
ref_part_model_final = self.surrogate_model(ref_part_tensor)
ref_uz_f = ref_part_model_final[5]
ref_beta_gamma_final = ref_uz_f
ref_part_final = torch.tensor(
[0, 0, ref_z_f, 0, 0, ref_uz_f], device=device, dtype=torch.float64
)
coordinate_transformation(pc, CoordSystem.t)
for lvl in range(pc.finest_level + 1):
for pti in ImpactXParIter(pc, level=lvl):
soa = pti.soa()
real_arrays = soa.get_real_data()
x = array(real_arrays[0], copy=False)
y = array(real_arrays[1], copy=False)
t = array(real_arrays[2], copy=False)
px = array(real_arrays[3], copy=False)
py = array(real_arrays[4], copy=False)
pt = array(real_arrays[5], copy=False)
data_arr = torch.tensor(
stack([x, y, t, px, py, pt], axis=1),
device=device,
dtype=torch.float64,
)
data_arr[:, 0] += ref_part.x
data_arr[:, 1] += ref_part.y
data_arr[:, 2] += ref_z_i_LPA
data_arr[:, 3:] *= ref_beta_gamma
data_arr[:, 3] += ref_part.px
data_arr[:, 4] += ref_part.py
data_arr[:, 5] += ref_part.pz
with torch.no_grad():
data_arr_post_model = self.surrogate_model(data_arr)
# z += stage start
data_arr_post_model[:, 2] += self.stage_start
# back to ref particle coordinates
for ii in range(3):
data_arr_post_model[:, ii] -= ref_part_final[ii]
data_arr_post_model[:, 3 + ii] -= ref_part_final[3 + ii]
data_arr_post_model[:, 3 + ii] /= ref_beta_gamma_final
x[:] = array(data_arr_post_model[:, 0])
y[:] = array(data_arr_post_model[:, 1])
t[:] = array(data_arr_post_model[:, 2])
px[:] = array(data_arr_post_model[:, 3])
py[:] = array(data_arr_post_model[:, 4])
pt[:] = array(data_arr_post_model[:, 5])
# TODO this part needs to be corrected for general geometry
# where the initial vector might not point in z
# and even if it does, bending elements may change the direction
ref_part.x = ref_part_final[0]
ref_part.y = ref_part_final[1]
ref_part.z = ref_part_final[2]
ref_gamma = torch.sqrt(1 + ref_beta_gamma_final**2)
ref_part.px = ref_part_final[3]
ref_part.py = ref_part_final[4]
ref_part.pz = ref_part_final[5]
ref_part.pt = -ref_gamma
ref_part.s += self.surrogate_length
ref_part.t += self.surrogate_length
coordinate_transformation(pc, CoordSystem.s)
## Done!
L_transport = 0.03
L_lens = 0.003
L_focal = 0.5 * L_transport
L_drift = 0.5 * (L_transport - L_lens)
K = np.sqrt(2.0 / L_focal / L_lens)
Kt = 1e-11 # number chosen arbitrarily since 0 isn't allowed
dL = 0
L_drift_minus_surrogate = L_drift
L_drift_1 = L_drift - drift_after_LPA - dL
L_drift_before_2nd_stage = abs(ebeam_lpa_z0)
L_drift_2 = L_drift - L_drift_before_2nd_stage + dL
def get_lattice_element_iter(sim, j):
assert (
0 <= j < len(sim.lattice)
), f"Argument j must be a nonnegative integer satisfying 0 <= j < {len(sim.lattice)}, not {j}"
i = 0
lat_it = sim.lattice.__iter__()
next(lat_it)
while i != j:
next(lat_it)
i += 1
return lat_it
def lens_eqn(k, lens_length, alpha, beta, gamma):
return np.tan(k * lens_length) + 2 * alpha / (k * beta - gamma / k)
k_list = []
class UpdateConstF(elements.Programmable):
def __init__(self, sim, stage_i, lattice_index, x_or_y):
elements.Programmable.__init__(self)
self.sim = sim
self.stage_i = stage_i
self.lattice_index = lattice_index
self.x_or_y = x_or_y
self.push = self.set_lens
def set_lens(self, step):
pc = self.sim.particle_container()
# get envelope parameters
rbc = pc.reduced_beam_characteristics()
alpha = rbc[f"alpha_{self.x_or_y}"]
beta = rbc[f"beta_{self.x_or_y}"]
gamma = (1 + alpha**2) / beta
# solve for k_new
sol = opt.root_scalar(
lens_eqn, bracket=[100, 300], args=(L_lens, alpha, beta, gamma)
)
k_new = sol.root
# set lens
self_it = get_lattice_element_iter(self.sim, self.lattice_index)
following_lens = next(self_it)
k_list.append(k_new)
following_lens.kx = k_new
following_lens.ky = k_new
lpa_stages = []
for i in range(N_stage):
lpa = LPASurrogateStage(i, model_list[i], L_surrogate, L_stage_period * i)
lpa.nslice = n_slice
lpa.ds = L_surrogate
lpa_stages.append(lpa)
monitor = elements.BeamMonitor("monitor")
for i in range(N_stage):
sim.lattice.extend(
[
monitor,
lpa_stages[i],
]
)
if i != N_stage - 1:
sim.lattice.extend(
[
monitor,
elements.Drift(ds=L_drift_1),
monitor,
UpdateConstF(
sim=sim, stage_i=i, lattice_index=5 + 9 * i, x_or_y=tune_by_x_or_y
),
elements.ConstF(ds=L_lens, kx=K, ky=K, kt=Kt),
monitor,
elements.Drift(ds=L_drift_2),
]
)
sim.lattice.extend([monitor])
sim.evolve()
sim.finalize()
del sim
This script requires some utility code for using the neural networks that is provided here:
Script surrogate_model_definitions.py
examples/pytorch_surrogate_model/surrogate_model_definitions.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Ryan Sandberg, Axel Huebl
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
from enum import Enum
import torch
from torch import nn
class Activation(Enum):
"""
Activation class provides an enumeration type for the supported activation layers
"""
ReLU = 1
Tanh = 2
PReLU = 3
Sigmoid = 4
def get_enum_type(type_to_test, EnumClass):
"""
Returns the enumeration type associated to type_to_test in EnumClass
Parameters
----------
type_to_test: EnumClass, int or str
object whose Enum class is to be obtained
EnumClass: Enum class
Enum class to test
"""
if isinstance(type_to_test, EnumClass): ## Useful ?
return type_to_test
if isinstance(type_to_test, int):
return EnumClass(type_to_test)
if isinstance(type_to_test, str):
return getattr(EnumClass, type_to_test)
class ConnectedNN(nn.Module):
"""
ConnectedNN is a class of fully connected neural networks
"""
def __init__(self, layers, device=None):
super().__init__()
self.stack = nn.Sequential(*layers)
if device is not None:
self.to(device)
def forward(self, x):
return self.stack(x)
class OneActNN(ConnectedNN):
"""
OneActNN is class of fully connected neural networks admitting only one activation function
"""
def __init__(self, n_in, n_out, n_hidden_nodes, n_hidden_layers, act, device=None):
self.n_in = n_in
self.n_out = n_out
self.n_hidden_layers = n_hidden_layers
self.n_hidden_nodes = n_hidden_nodes
self.act = act
layers = [nn.Linear(self.n_in, self.n_hidden_nodes)]
for ii in range(self.n_hidden_layers):
if self.act is Activation.ReLU:
layers += [nn.ReLU()]
if self.act is Activation.Tanh:
layers += [nn.Tanh()]
if self.act is Activation.PReLU:
layers += [nn.PReLU()]
if self.act is Activation.Sigmoid:
layers += [nn.Sigmoid()]
if ii < self.n_hidden_layers - 1:
layers += [nn.Linear(self.n_hidden_nodes, self.n_hidden_nodes)]
layers += [nn.Linear(self.n_hidden_nodes, self.n_out)]
super().__init__(layers, device)
class surrogate_model:
"""
Extend the functionality of the OneActNN class
This class is meant to act as a wrapper for the OneActNN class.
It provides a `__call__` function that normalizes input and returns dimensional output.
"""
def __init__(self, model_file, device=None):
self.device = device
if device is None:
model_dict = torch.load(model_file, map_location="cpu")
else:
model_dict = torch.load(model_file, map_location=device)
self.source_means = torch.tensor(
model_dict["source_means"], device=self.device, dtype=torch.float64
)
self.target_means = torch.tensor(
model_dict["target_means"], device=self.device, dtype=torch.float64
)
self.source_stds = torch.tensor(
model_dict["source_stds"], device=self.device, dtype=torch.float64
)
self.target_stds = torch.tensor(
model_dict["target_stds"], device=self.device, dtype=torch.float64
)
n_in = model_dict["model_state_dict"]["stack.0.weight"].shape[1]
final_layer_key = list(model_dict["model_state_dict"].keys())[-1]
n_out = model_dict["model_state_dict"][final_layer_key].shape[0]
n_hidden_nodes = model_dict["model_state_dict"]["stack.0.weight"].shape[0]
activation_type = model_dict["activation"]
activation = get_enum_type(activation_type, Activation)
if "n_hidden_layers" in model_dict.keys():
n_hidden_layers = model_dict["n_hidden_layers"]
else:
if activation is Activation.PReLU:
n_hidden_layers = int(
(len(model_dict["model_state_dict"].keys()) - 2) / 3
)
else:
n_hidden_layers = int(
len(model_dict["model_state_dict"].keys()) / 2 - 1
)
self.neural_network = OneActNN(
n_in=n_in,
n_out=n_out,
n_hidden_nodes=n_hidden_nodes,
n_hidden_layers=n_hidden_layers,
act=activation,
device=device,
)
self.neural_network.load_state_dict(model_dict["model_state_dict"])
self.neural_network.eval()
def __call__(self, data_arr):
data_arr -= self.source_means
data_arr /= self.source_stds
with torch.no_grad():
data_arr_post_model = self.neural_network(data_arr.float()).double()
data_arr_post_model *= self.target_stds
data_arr_post_model += self.target_means
return data_arr_post_model
Analyze
We run the following script to analyze correctness:
Script analyze_ml_surrogate_15_stage.py
examples/pytorch_surrogate_model/analyze_ml_surrogate_15_stage.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 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.bp", 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 = 100000
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 = 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],
[
7.494325e-07,
7.478916e-07,
9.976192e-08,
5.070297e-10,
5.080007e-10,
],
rtol=rtol,
atol=atol,
)
atol = 1.0e-6
print(f" atol~={atol}")
assert np.allclose([emittance_t], [0.0], 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 = 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.590999e-07,
1.634865e-07,
1.030930e-07,
5.031797e-12,
5.242205e-12,
2.049623e-11,
],
rtol=rtol,
atol=atol,
)
Visualize
You can run the following script to visualize the beam evolution over time:
Script visualize_ml_surrogate_15_stage.py
examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.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 argparse
import glob
import re
import numpy as np
import openpmd_api as io
import pandas as pd
from matplotlib import pyplot as plt
from scipy.constants import c, e, m_e
def read_all_files(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(
(
pd.read_csv(filename, delimiter=r"\s+")
for filename in glob.glob(file_pattern)
),
axis=0,
ignore_index=True,
).set_index("id")
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')
from enum import Enum
class TCoords(Enum):
REF = 1
GLOBAL = 2
def to_t(
ref_pz, ref_pt, data_arr_s, ref_z=None, coord_type=TCoords.REF
): # x, y, t, dpx, dpy, dpt):
"""Change to fixed t coordinates
Parameters
---
ref_pz: float, reference particle momentum in z
ref_pt: float, reference particle pt = -gamma
data_arr_s: Nx6 array-like structure containing fixed-s particle coordinates
ref_z: if transforming to global coordinates
coord_type: TCoords enum, (default is in ref coordinates) whether to get particle data relative to reference coordinate or in the global frame
"""
if type(data_arr_s) is pd.core.frame.DataFrame:
dx = data_arr_s["position_x"]
dy = data_arr_s["position_y"]
dt = data_arr_s["position_t"]
dpx = data_arr_s["momentum_x"]
dpy = data_arr_s["momentum_y"]
dpt = data_arr_s["momentum_t"]
elif type(data_arr_s) is np.ndarray:
assert (
data_arr_s.shape[1] == 6
), f"data_arr_s.shape={data_arr_s.shape} but data_arr_s must be an Nx6 array"
dx, dy, dt, dpx, dpy, dpt = data_arr_s.T
else:
raise Exception(
f"Incompatible input type {type(data_arr_s)} for data_arr_s, must be pandas DataFrame or Nx6 array-like object"
)
dx += ref_pz * dpx * dt / (ref_pt + ref_pz * dpt)
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
)
dt *= pz / (ref_pt + ref_pz * dpt)
if type(data_arr_s) is pd.core.frame.DataFrame:
data_arr_s["momentum_t"] = pz - ref_pz
dpt = data_arr_s["momentum_t"]
else:
dpt[:] = pz - ref_pz
if coord_type is TCoords.REF:
print("applying reference normalization")
dpt /= ref_pz
elif coord_type is TCoords.GLOBAL:
assert (
ref_z is not None
), "Reference particle z coordinate is required to transform to global coordinates"
print("target global coordinates")
dt += ref_z
dpx *= ref_pz
dpy *= ref_pz
dpt += ref_pz
# data_arr_t = np.column_stack([xt,yt,z,dpx,dpy,dpz])
return # modifies data_arr_s in place
def plot_beam_df(
beam_at_step,
axT,
unit=1e6,
unit_z=1e3,
unit_label="$\mu$m",
unit_z_label="mm",
alpha=1.0,
cmap=None,
color="k",
size=0.1,
t_offset=0.0,
label=None,
z_ticks=None,
):
ax = axT[0][0]
ax.scatter(
beam_at_step.position_x.multiply(unit),
beam_at_step.position_y.multiply(unit),
c=color,
s=size,
alpha=alpha,
cmap=cmap,
)
ax.set_xlabel(r"x (%s)" % unit_label)
ax.set_ylabel(r"y (%s)" % unit_label)
ax.axes.ticklabel_format(axis="both", style="sci", scilimits=(-2, 2))
###########
ax = axT[0][1]
ax.scatter(
beam_at_step.position_t.multiply(unit_z) - t_offset,
beam_at_step.position_x.multiply(unit),
c=color,
s=size,
alpha=alpha,
cmap=cmap,
)
ax.set_xlabel(r"%s" % unit_z_label)
ax.set_ylabel(r"x (%s)" % unit_label)
ax.axes.ticklabel_format(axis="y", style="sci", scilimits=(-2, 2))
if z_ticks is not None:
ax.set_xticks(z_ticks)
###########
ax = axT[0][2]
ax.scatter(
beam_at_step.position_t.multiply(unit_z) - t_offset,
beam_at_step.position_y.multiply(unit),
c=color,
s=size,
alpha=alpha,
cmap=cmap,
)
ax.set_xlabel(r"%s" % unit_z_label)
ax.set_ylabel(r"y (%s)" % unit_label)
ax.axes.ticklabel_format(axis="y", style="sci", scilimits=(-2, 2))
if z_ticks is not None:
ax.set_xticks(z_ticks)
############
##########
ax = axT[1][0]
ax.scatter(
beam_at_step.momentum_x,
beam_at_step.momentum_y,
c=color,
s=size,
alpha=alpha,
cmap=cmap,
)
ax.set_xlabel("px")
ax.set_ylabel("py")
ax.axes.ticklabel_format(axis="both", style="sci", scilimits=(-2, 2))
##########
ax = axT[1][1]
ax.scatter(
beam_at_step.momentum_t,
# beam_at_step.position_t.multiply(unit_z)-t_offset,
beam_at_step.momentum_x,
c=color,
s=size,
alpha=alpha,
cmap=cmap,
)
ax.set_xlabel("pt")
# ax.set_xlabel(r'%s'%unit_z_label)
ax.set_ylabel("px")
ax.axes.ticklabel_format(axis="both", style="sci", scilimits=(-2, 2))
##########
ax = axT[1][2]
ax.scatter(
beam_at_step.momentum_t,
# beam_at_step.position_t.multiply(unit_z)-t_offset,
beam_at_step.momentum_y,
c=color,
s=size,
alpha=alpha,
label=label,
cmap=cmap,
)
if label is not None:
ax.legend()
# ax.set_xlabel(r'%s'%unit_z_label)
ax.set_xlabel("pt")
ax.set_ylabel("py")
ax.axes.ticklabel_format(axis="both", style="sci", scilimits=(-2, 2))
############
############
##########
ax = axT[2][0]
ax.scatter(
beam_at_step.position_x.multiply(unit),
beam_at_step.momentum_x,
c=color,
s=size,
alpha=alpha,
cmap=cmap,
)
ax.set_xlabel(r"x (%s)" % unit_label)
ax.set_ylabel("px")
ax.axes.ticklabel_format(axis="both", style="sci", scilimits=(-2, 2))
############
ax = axT[2][1]
ax.scatter(
beam_at_step.position_y.multiply(unit),
beam_at_step.momentum_y,
c=color,
s=size,
alpha=alpha,
cmap=cmap,
)
ax.set_xlabel(r"y (%s)" % unit_label)
ax.set_ylabel("py")
ax.axes.ticklabel_format(axis="both", style="sci", scilimits=(-2, 2))
################
ax = axT[2][2]
ax.scatter(
beam_at_step.position_t.multiply(unit_z) - t_offset,
beam_at_step.momentum_t,
c=color,
s=size,
alpha=alpha,
cmap=cmap,
)
ax.set_xlabel(r"%s" % unit_z_label)
ax.set_ylabel("pt")
ax.axes.ticklabel_format(axis="y", style="sci", scilimits=(-2, 2))
if z_ticks is not None:
ax.set_xticks(z_ticks)
plt.tight_layout()
# done
# options to run this script
parser = argparse.ArgumentParser(description="Plot the ML surrogate benchmark.")
parser.add_argument(
"--save-png", action="store_true", help="non-interactive run: save to PNGs"
)
parser.add_argument(
"--num-stages", "-n", type=int, default=15, help="num stages to plot"
)
parser.add_argument(
"--stages_to_plot", "-s", type=int, help="num stages to plot", nargs="*"
)
args = parser.parse_args()
impactx_surrogate_reduced_diags = read_time_series(
"diags/reduced_beam_characteristics.*"
)
ref_gamma = np.sqrt(1 + impactx_surrogate_reduced_diags["ref_beta_gamma"] ** 2)
beam_gamma = (
ref_gamma
- impactx_surrogate_reduced_diags["pt_mean"]
* impactx_surrogate_reduced_diags["ref_beta_gamma"]
)
beam_u = np.sqrt(beam_gamma**2 - 1)
emit_x = impactx_surrogate_reduced_diags["emittance_x"]
emit_nx = emit_x * beam_u
emit_y = impactx_surrogate_reduced_diags["emittance_y"]
emit_ny = emit_y * beam_u
ix_slice = [0] + [2 + 9 * i for i in range(args.num_stages)]
############# plot moments ##############
fig, axT = plt.subplots(2, 2, figsize=(10, 8))
ymarker = "^"
######### emittance ##########
ax = axT[0][0]
scale = 1e6
ax.plot(
impactx_surrogate_reduced_diags["s"][ix_slice],
emit_nx[ix_slice] * scale,
"bo",
label="x",
)
ax.plot(
impactx_surrogate_reduced_diags["s"][ix_slice],
emit_ny[ix_slice] * scale,
"r",
marker=ymarker,
linestyle="None",
label="y",
)
ax.legend()
ax.set_xlabel("s (m)")
ax.set_ylabel(r"emittance (mm-mrad)")
######### energy ##########
ax = axT[0][1]
scale = m_e * c**2 / e * 1e-9
ax.plot(
impactx_surrogate_reduced_diags["s"][ix_slice],
beam_gamma[ix_slice] * scale,
"go",
)
ax.set_xlabel("s (m)")
ax.set_ylabel(r"mean energy (GeV)")
######### width ##########
ax = axT[1][0]
scale = 1e6
ax.plot(
impactx_surrogate_reduced_diags["s"][ix_slice],
impactx_surrogate_reduced_diags["sig_x"][ix_slice] * scale,
"bo",
label="x",
)
ax.plot(
impactx_surrogate_reduced_diags["s"][ix_slice],
impactx_surrogate_reduced_diags["sig_y"][ix_slice] * scale,
"r",
marker=ymarker,
linestyle="None",
label="y",
)
ax.legend()
ax.set_xlabel("s (m)")
ax.set_ylabel(r"beam width ($\mu$m)")
######### divergence ##########
ax = axT[1][1]
scale = 1e3
ax.semilogy(
impactx_surrogate_reduced_diags["s"][ix_slice],
impactx_surrogate_reduced_diags["sig_px"][ix_slice] * scale,
"bo",
label="x",
)
ax.semilogy(
impactx_surrogate_reduced_diags["s"][ix_slice],
impactx_surrogate_reduced_diags["sig_py"][ix_slice] * scale,
"r",
marker=ymarker,
linestyle="None",
label="y",
)
ax.legend()
ax.set_xlabel("s (m)")
ax.set_ylabel(r"divergence (mrad)")
plt.tight_layout()
if args.save_png:
plt.savefig("lpa_ml_surrogate_moments.png")
else:
plt.show()
######## plot phase spaces ###########
beam_impactx_surrogate_series = io.Series(
"diags/openPMD/monitor.bp", io.Access.read_only
)
impactx_surrogate_steps = list(beam_impactx_surrogate_series.iterations)
impactx_surrogate_ref_particle = read_time_series("diags/ref_particle.*")
millimeter = 1.0e3
micron = 1.0e6
N_stage = args.num_stages
impactx_stage_end_steps = [1] + [3 + 9 * i for i in range(N_stage)]
ise = impactx_stage_end_steps
# initial
step = 1
beam_at_step = beam_impactx_surrogate_series.iterations[step].particles["beam"].to_df()
ref_part_step = impactx_surrogate_ref_particle.loc[step]
ref_u = np.sqrt(ref_part_step["pt"] ** 2 - 1)
to_t(
ref_u,
ref_part_step["pt"],
beam_at_step,
ref_z=ref_part_step["z"],
coord_type=TCoords.GLOBAL,
)
t_offset = impactx_surrogate_ref_particle.loc[step, "t"] * micron
fig, axT = plt.subplots(3, 3, figsize=(10, 8))
fig.suptitle(f"initially, ct={impactx_surrogate_ref_particle.at[step,'t']:.2f} m")
plot_beam_df(
beam_at_step,
axT,
alpha=0.6,
color="red",
unit_z=1e6,
unit_z_label=r"$\xi$ ($\mu$m)",
t_offset=t_offset,
z_ticks=[-107.3, -106.6],
)
if args.save_png:
plt.savefig("initial_phase_spaces.png")
else:
plt.show()
####### final ###########
if args.stages_to_plot is not None:
for stage_i in args.stages_to_plot:
step = ise[stage_i]
beam_at_step = (
beam_impactx_surrogate_series.iterations[step].particles["beam"].to_df()
)
ref_part_step = impactx_surrogate_ref_particle.loc[step]
ref_u = np.sqrt(ref_part_step["pt"] ** 2 - 1)
to_t(
ref_u,
ref_part_step["pt"],
beam_at_step,
ref_z=ref_part_step["z"],
coord_type=TCoords.GLOBAL,
)
t_offset = impactx_surrogate_ref_particle.loc[step, "t"] * micron
fig, axT = plt.subplots(3, 3, figsize=(10, 8))
fig.suptitle(
f"stage {stage_i}, ct={impactx_surrogate_ref_particle.at[step,'t']:.2f} m"
)
plot_beam_df(
beam_at_step,
axT,
alpha=0.6,
color="red",
unit_z=1e6,
unit_z_label=r"$\xi$ ($\mu$m)",
t_offset=t_offset,
z_ticks=[-107.3, -106.6],
)
if args.save_png:
plt.savefig(f"stage_{stage_i-1}_phase_spaces.png")
else:
plt.show()

Evolution of electron beam moments through 15 stages of LPAs (via neural network surrogates).

Initial phase space projections going into 15 stage LPA (via neural network surrogates) simulation. Top row: spatial projections, middle row: momentum projections, bottom row: phase spaces.

Final phase space projections after 15 stage LPA (via neural network surrogates) simulation. Top row: spatial projections, middle row: momentum projections, bottom row: phase spaces.
Apochromatic Drift-Quadrupole Beamline
Electron beam matched to the 1st-order apochromatic drift-quadrupole beamline appearing in Fig. 4a of: C. A. Lindstrom and E. Adli, “Design of general apochromatic drift-quadrupole beam lines,” Phys. Rev. Accel. Beams 19, 071002 (2016).
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 0.325\) m
\(\alpha_\mathrm{x} = 0\)
\(\beta_\mathrm{y} = 0.325\) m
\(\alpha_\mathrm{y} = 0\)
We use a 100 GeV electron beam with an initially 6D Gaussian distribution of normalized rms emittance 1 micron and relative energy spread of 1%.
The second moments of the particle distribution after the focusing beamline should coincide with the second moments of the particle distribution before the beamline, to within the level expected due to noise due to statistical sampling. The emittance growth due to chromatic effects remain below 1%. In the absence of chromatic correction, the projected emittance growth is near 10%.
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_apochromatic.py
orImpactX executable using an input file:
impactx input_apochromatic.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/apochromatic/run_apochromatic.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 = False
# sim.diagnostics = False # benchmarking
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.0e3 # reference energy
bunch_charge_C = 1.0e-9 # used with space charge
npart = 100000 # 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=1.288697604e-6,
lambdaY=1.288697604e-6,
lambdaT=1.0e-6,
lambdaPx=3.965223396e-6,
lambdaPy=3.965223396e-6,
lambdaPt=0.01, # 1% energy spread
muxpx=0.0,
muypy=0.0,
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
# Drift elements
dr1 = elements.ChrDrift(ds=1.0, nslice=ns)
dr2 = elements.ChrDrift(ds=10.0, nslice=ns)
# Quad elements
q1 = elements.ChrQuad(ds=1.2258333333, k=0.5884, nslice=ns)
q2 = elements.ChrQuad(ds=1.5677083333, k=-0.7525, nslice=ns)
q3 = elements.ChrQuad(ds=1.205625, k=0.5787, nslice=ns)
q4 = elements.ChrQuad(ds=1.2502083333, k=-0.6001, nslice=ns)
q5 = elements.ChrQuad(ds=1.2502083333, k=0.6001, nslice=ns)
q6 = elements.ChrQuad(ds=1.205625, k=-0.5787, nslice=ns)
q7 = elements.ChrQuad(ds=1.5677083333, k=0.7525, nslice=ns)
q8 = elements.ChrQuad(ds=1.2258333333, k=-0.5884, nslice=ns)
lattice_line = [monitor, dr1, q1, q2, q3, dr2, q4, q5, dr2, q6, q7, q8, dr1, monitor]
# define the lattice
sim.lattice.extend(lattice_line)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/apochromatic/input_apochromatic.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 100000
beam.units = static
beam.kin_energy = 100.0e3 # 100 GeV nominal energy
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = gaussian
beam.lambdaX = 1.288697604e-6
beam.lambdaY = 1.288697604e-6
beam.lambdaT = 1.0e-6
beam.lambdaPx = 3.965223396e-6
beam.lambdaPy = 3.965223396e-6
beam.lambdaPt = 0.01 #1% energy spread
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor dr1 q1 q2 q3 dr2 q4 q5 dr2 q6 q7 q8 dr1 monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
dr1.type = drift_chromatic
dr1.ds = 1.0
dr2.type = drift_chromatic
dr2.ds = 10.0
q1.type = quad_chromatic
q1.ds = 1.2258333333
q1.k = 0.5884
q2.type = quad_chromatic
q2.ds = 1.5677083333
q2.k = -0.7525
q3.type = quad_chromatic
q3.ds = 1.205625
q3.k = 0.5787
q4.type = quad_chromatic
q4.ds = 1.2502083333
q4.k = -0.6001
q5.type = quad_chromatic
q5.ds = 1.2502083333
q5.k = 0.6001
q6.type = quad_chromatic
q6.ds = 1.205625
q6.k = -0.5787
q7.type = quad_chromatic
q7.ds = 1.5677083333
q7.k = 0.7525
q8.type = quad_chromatic
q8.ds = 1.2258333333
q8.k = -0.5884
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_apochromatic.py
examples/apochromatic/analysis_apochromatic.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["divergence_x"], moment=2) ** 0.5
sigy = moment(beam["position_y"], moment=2) ** 0.5
sigpy = moment(beam["divergence_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 = 100000
assert num_particles == len(initial)
assert num_particles == len(final)
scale = (
(1.0 - initial.momentum_t) ** 2
+ (initial.momentum_x) ** 2
+ (initial.momentum_y) ** 2
)
xp = initial.momentum_x / np.sqrt(scale)
initial["divergence_x"] = xp
yp = initial.momentum_y / np.sqrt(scale)
initial["divergence_y"] = yp
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],
[
1.288697604e-6,
1.288697604e-6,
1.0e-6,
5.10997388810014764e-12,
5.10997388810014764e-12,
1.0e-8,
],
rtol=rtol,
atol=atol,
)
scale = (
(1.0 - final.momentum_t) ** 2 + (final.momentum_x) ** 2 + (final.momentum_y) ** 2
)
xp = final.momentum_x / np.sqrt(scale)
final["divergence_x"] = xp
yp = final.momentum_y / np.sqrt(scale)
final["divergence_y"] = yp
print("")
print("Final Beam:")
sigx, sigy, sigt, emittance_xf, emittance_yf, emittance_tf = get_moments(final)
demittance_x = 100 * (emittance_xf - emittance_x) / emittance_x
demittance_y = 100 * (emittance_yf - emittance_y) / emittance_y
demittance_t = 100 * (emittance_tf - emittance_t) / emittance_t
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
print(
f" emittance change x (%)={demittance_x:e} emittance change y (%)={demittance_y:e} emittance change t (%)={demittance_t:e}"
)
atol = 0.0 # ignored
rtol = 26.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, demittance_x, demittance_y, emittance_t],
[
1.245e-6,
1.245e-6,
1.0e-6,
0.94,
0.94,
1.0e-8,
],
rtol=rtol,
atol=atol,
)
Apochromatic Drift-Plasma Lens Beamline
Electron beam matched to the 3rd-order apochromatic drift-plasma lens beamline appearing in Fig. 4b of: C. A. Lindstrom and E. Adli, “Design of general apochromatic drift-quadrupole beam lines,” Phys. Rev. Accel. Beams 19, 071002 (2016).
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 0.325\) m
\(\alpha_\mathrm{x} = 0\)
\(\beta_\mathrm{y} = 0.325\) m
\(\alpha_\mathrm{y} = 0\)
We use a 100 GeV electron beam with an initially 6D Gaussian distribution of normalized rms emittance 1 micron and relative energy spread of 1%.
The second moments of the particle distribution after the focusing beamline should coincide with the second moments of the particle distribution before the beamline, to within the level expected due to noise due to statistical sampling. The emittance growth due to chromatic effects remain below 1%. In the absence of chromatic correction, the projected emittance growth is near 10%.
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_apochromatic_pl.py
orImpactX executable using an input file:
impactx input_apochromatic_pl.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/apochromatic/run_apochromatic_pl.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 = False
# sim.diagnostics = False # benchmarking
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.0e3 # reference energy
bunch_charge_C = 1.0e-9 # used with space charge
npart = 100000 # 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=1.288697604e-6,
lambdaY=1.288697604e-6,
lambdaT=1.0e-6,
lambdaPx=3.965223396e-6,
lambdaPy=3.965223396e-6,
lambdaPt=0.01, # 1% energy spread
muxpx=0.0,
muypy=0.0,
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
# Drift elements
dr1 = elements.ChrDrift(ds=1.0, nslice=ns)
dr2 = elements.ChrDrift(ds=2.0, nslice=ns)
# Plasma lens elements
q1 = elements.ChrPlasmaLens(
ds=0.331817852986604588, k=996.147787384348956, units=1, nslice=ns
)
q2 = elements.ChrPlasmaLens(
ds=0.176038957633108457, k=528.485181135649785, units=1, nslice=ns
)
q3 = elements.ChrPlasmaLens(
ds=1.041842576046930486, k=3127.707468391874166, units=1, nslice=ns
)
q4 = elements.ChrPlasmaLens(
ds=0.334367090894399520, k=501.900417308233112, units=1, nslice=ns
)
q5 = elements.ChrPlasmaLens(
ds=1.041842576046930486, k=3127.707468391874166, units=1, nslice=ns
)
q6 = elements.ChrPlasmaLens(
ds=0.176038957633108457, k=528.485181135649785, units=1, nslice=ns
)
q7 = elements.ChrPlasmaLens(
ds=0.331817852986604588, k=996.147787384348956, units=1, nslice=ns
)
# q1 = elements.ChrPlasmaLens(ds=0.331817852986604588, k=2.98636067687944129, units=0, nslice=ns)
# q2 = elements.ChrPlasmaLens(ds=0.176038957633108457, k=1.584350618697976110, units=0, nslice=ns)
# q3 = elements.ChrPlasmaLens(ds=1.041842576046930486, k=9.37658318442237437, units=0, nslice=ns)
# q4 = elements.ChrPlasmaLens(ds=0.334367090894399520, k=1.50465190902479784, units=0, nslice=ns)
# q5 = elements.ChrPlasmaLens(ds=1.041842576046930486, k=9.37658318442237437, units=0, nslice=ns)
# q6 = elements.ChrPlasmaLens(ds=0.176038957633108457, k=1.584350618697976110, units=0, nslice=ns)
# q7 = elements.ChrPlasmaLens(ds=0.331817852986604588, k=2.98636067687944129, units=0, nslice=ns)
lattice_line = [
monitor,
dr1,
q1,
dr2,
q2,
dr2,
q3,
dr2,
q4,
dr2,
q5,
dr2,
q6,
dr2,
q7,
dr1,
monitor,
]
# define the lattice
sim.lattice.extend(lattice_line)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
examples/apochromatic/input_apochromatic_pl.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 100000
beam.units = static
beam.kin_energy = 100.0e3 # 100 GeV nominal energy
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = gaussian
beam.lambdaX = 1.288697604e-6
beam.lambdaY = 1.288697604e-6
beam.lambdaT = 1.0e-6
beam.lambdaPx = 3.965223396e-6
beam.lambdaPy = 3.965223396e-6
beam.lambdaPt = 0.01 #1% energy spread
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor dr1 q1 dr2 q2 dr2 q3 dr2 q4 dr2 q5 dr2 q6 dr2 q7 dr1 monitor
lattice.nslice = 1
monitor.type = beam_monitor
monitor.backend = h5
dr1.type = drift_chromatic
dr1.ds = 1.0
dr2.type = drift_chromatic
dr2.ds = 2.0
q1.type = plasma_lens_chromatic
q1.ds = 0.331817852986604588
q1.k = 996.147787384348956
q1.units = 1
#q1.k = 2.98636067687944129
q2.type = plasma_lens_chromatic
q2.ds = 0.176038957633108457
q2.k = 528.485181135649785
q2.units = 1
#q2.k = 1.584350618697976110
q3.type = plasma_lens_chromatic
q3.ds = 1.041842576046930486
q3.k = 3127.707468391874166
q3.units = 1
#q3.k = 9.37658318442237437
q4.type = plasma_lens_chromatic
q4.ds = 0.334367090894399520
q4.k = 501.900417308233112
q4.units = 1
#q4.k = 1.50465190902479784
q5.type = plasma_lens_chromatic
q5.ds = 1.041842576046930486
q5.k = 3127.707468391874166
q5.units = 1
#q5.k = 9.37658318442237437
q6.type = plasma_lens_chromatic
q6.ds = 0.176038957633108457
q6.k = 528.485181135649785
q6.units = 1
#q6.k = 1.584350618697976110
q7.type = plasma_lens_chromatic
q7.ds = 0.331817852986604588
q7.k = 996.147787384348956
q7.units = 1
#q7.k = 2.98636067687944129
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
Analyze
We run the following script to analyze correctness:
Script analysis_apochromatic_pl.py
examples/apochromatic/analysis_apochromatic_pl.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["divergence_x"], moment=2) ** 0.5
sigy = moment(beam["position_y"], moment=2) ** 0.5
sigpy = moment(beam["divergence_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 = 100000
assert num_particles == len(initial)
assert num_particles == len(final)
scale = (
(1.0 - initial.momentum_t) ** 2
+ (initial.momentum_x) ** 2
+ (initial.momentum_y) ** 2
)
xp = initial.momentum_x / np.sqrt(scale)
initial["divergence_x"] = xp
yp = initial.momentum_y / np.sqrt(scale)
initial["divergence_y"] = yp
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],
[
1.288697604e-6,
1.288697604e-6,
1.0e-6,
5.10997388810014764e-12,
5.10997388810014764e-12,
1.0e-8,
],
rtol=rtol,
atol=atol,
)
scale = (
(1.0 - final.momentum_t) ** 2 + (final.momentum_x) ** 2 + (final.momentum_y) ** 2
)
xp = final.momentum_x / np.sqrt(scale)
final["divergence_x"] = xp
yp = final.momentum_y / np.sqrt(scale)
final["divergence_y"] = yp
print("")
print("Final Beam:")
sigx, sigy, sigt, emittance_xf, emittance_yf, emittance_tf = get_moments(final)
demittance_x = 100 * abs(emittance_xf - emittance_x) / emittance_x
demittance_y = 100 * abs(emittance_yf - emittance_y) / emittance_y
demittance_t = 100 * abs(emittance_tf - emittance_t) / emittance_t
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
print(
f" emittance change x (%)={demittance_x:e} emittance change y (%)={demittance_y:e} emittance change t (%)={demittance_t:e}"
)
atol = 0.0 # ignored
rtol = 19.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],
[
1.283476e-06,
1.291507e-06,
1.0e-6,
],
rtol=rtol,
atol=atol,
)
atol = 2.0e-3 # from random sampling of a smooth distribution
print(f" atol={atol} %")
assert np.allclose(
[demittance_x, demittance_y, demittance_t],
[
0.0,
0.0,
0.0,
],
atol=atol,
)
Tune Calculation in a Periodic FODO Channel
This is identical to the FODO example, except that tracking for 100 periods is used to extract the horizontal tune.
Stable FODO cell with a zero-current phase advance of 67.8 degrees, corresponding to Qx = Qy = 0.1883.
The matched Twiss parameters at entry are:
\(\beta_\mathrm{x} = 2.82161941\) m
\(\alpha_\mathrm{x} = -1.59050035\)
\(\beta_\mathrm{y} = 2.82161941\) m
\(\alpha_\mathrm{y} = 1.59050035\)
We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm.
The horizontal tune of a single particle is obtained from period-by-period tracking data using several algorithms.
In this test, the computed horizontal tune must agree with the nominal value to within acceptable tolerance.
This example requires installation of PyNAFF.
Run
This example can be run either as:
Python script:
python3 run_fodo_tune.py
orImpactX executable using an input file:
impactx input_fodo_tune.in
For MPI-parallel runs, prefix these lines with mpiexec -n 4 ...
or srun -n 4 ...
, depending on the system.
examples/fodo_tune/run_fodo_tune.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
import amrex.space3d as amr
from impactx import ImpactX, distribution, elements
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 = 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.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.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
fodo = [
monitor,
elements.Drift(ds=0.25, nslice=ns),
elements.Quad(ds=1.0, k=1.0, nslice=ns),
elements.Drift(ds=0.5, nslice=ns),
elements.Quad(ds=1.0, k=-1.0, nslice=ns),
elements.Drift(ds=0.25, nslice=ns),
]
# assign a fodo segment
sim.lattice.extend(fodo)
# number of periods
sim.periods = 100
# run simulation
sim.evolve()
# clean shutdown
del sim
amr.finalize()
examples/fodo_tune/input_fodo_tune.in
.###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag
beam.lambdaX = 3.9984884770e-5
beam.lambdaY = 3.9984884770e-5
beam.lambdaT = 1.0e-3
beam.lambdaPx = 2.6623538760e-5
beam.lambdaPy = 2.6623538760e-5
beam.lambdaPt = 2.0e-3
beam.muxpx = -0.846574929020762
beam.muypy = 0.846574929020762
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.periods = 100
lattice.elements = monitor drift1 quad1 drift2 quad2 drift3
lattice.nslice = 1
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 = false
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Analyze
We run the following script to analyze correctness:
Script analysis_fodo_tune.py
examples/fodo_tune/analysis_fodo_tune.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
import PyNAFF as pnf
# Collect beam data series
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
# Specify time series for particle j
j = 5
print(f"output for particle index = {j}")
# Create array of TBT data values
x = []
px = []
n = 0
for k_i, i in series.iterations.items():
beam = i.particles["beam"]
turn = beam.to_df()
x.append(turn["position_x"][j])
px.append(turn["momentum_x"][j])
n = n + 1
# Output number of periods in data series
nturns = len(x)
print(f"number of periods = {nturns}")
print()
# Approximate the tune and closed orbit using the 4-turn formula:
# from x data only
argument = (x[0] - x[1] + x[2] - x[3]) / (2.0 * (x[1] - x[2]))
tune1 = np.arccos(argument) / (2.0 * np.pi)
print(f"tune output from 4-turn formula, using x data = {tune1}")
# from px data only
argument = (px[0] - px[1] + px[2] - px[3]) / (2.0 * (px[1] - px[2]))
tune2 = np.arccos(argument) / (2.0 * np.pi)
print(f"tune output from 4-turn formula, using px data = {tune2}")
# orbit offset
xco = (x[1] ** 2 - x[2] ** 2 + x[1] * x[3] - x[2] * x[0]) / (
3.0 * (x[1] - x[2]) + x[3] - x[0]
)
pxco = (px[1] ** 2 - px[2] ** 2 + px[1] * px[3] - px[2] * px[0]) / (
3.0 * (px[1] - px[2]) + px[3] - px[0]
)
print(f"orbit offset from 4-turn formula: (x[m], px[rad]) = ({xco},{pxco})")
# matched Twiss functions from the 4-turn formula:
C1 = px[0] * (x[1] - x[2]) + px[1] * (x[2] - x[0]) + px[2] * (x[0] - x[1])
C2 = -px[0] * (x[1] - x[2]) + px[1] * (x[2] + x[0]) + px[2] * (-x[0] - x[1])
C3 = px[0] * (-x[1] - x[2]) - px[1] * (x[2] - x[0]) + px[2] * (x[0] + x[1])
C4 = px[0] * (x[1] + x[2]) + px[1] * (-x[2] - x[0]) - px[2] * (x[0] - x[1])
C = -C1 * C2 * C3 * C4
alphax = (
px[0] ** 2 * (x[1] ** 2 - x[2] ** 2)
+ px[1] ** 2 * (x[2] ** 2 - x[0] ** 2)
+ px[2] ** 2 * (x[0] ** 2 - x[1] ** 2)
) / np.sqrt(C)
betax = (
2.0
* (
px[0] * x[0] * (x[2] ** 2 - x[1] ** 2)
+ px[1] * x[1] * (x[0] ** 2 - x[2] ** 2)
+ px[2] * x[2] * (x[1] ** 2 - x[0] ** 2)
)
/ np.sqrt(C)
)
alphax = np.sign(betax) * alphax # ensure selection of correct sign
betax = np.sign(betax) * betax # ensure selection of correct sign
print(f"Twiss from 4-turn formula: alphax, betax [m] = {alphax}, {betax}")
# Normalize TBT data using Twiss functions:
z = []
for n in range(0, nturns):
xn = x[n] / np.sqrt(betax)
pxn = px[n] * np.sqrt(betax) + x[n] * alphax / np.sqrt(betax)
z.append(complex(xn, -pxn))
print()
# Approximate the tune using average phase advance (APA) - 1 period
tune3 = np.angle(z[1] / z[0]) / (2.0 * np.pi)
print(f"tune from phase advance = {tune3}")
print()
# Approximate the tune by using NAFF on the entire data series:
output = pnf.naff(
x, turns=nturns, nterms=4, skipTurns=0, getFullSpectrum=True, window=1
)
print("tune output from NAFF, using x data:")
print(output[0, 1], output[1, 1])
output = pnf.naff(
px, turns=nturns, nterms=4, skipTurns=0, getFullSpectrum=True, window=1
)
print("tune output from NAFF, using px data:")
print(output[0, 1], output[1, 1])
output = pnf.naff(
z, turns=nturns, nterms=4, skipTurns=0, getFullSpectrum=True, window=1
)
tune4 = output[0, 1]
print(f"tune output from NAFF, using normalized x-ipx data: {tune4}")
rtol = 1.0e-3
print(f" rtol={rtol}")
assert np.allclose(
[alphax, betax, tune1, tune2, tune3, tune4],
[
-1.59050035,
2.82161941,
0.1883,
0.1883,
0.1883,
0.1883,
],
rtol=rtol,
)
Optimized Triplet
Optimization of focusing parameters for a quadrupole triplet. A 2 GeV electron beam is strongly focused from lattice initial parameters:
to final lattice parameters:
resulting in a reduction by a factor of 8.5 in the horizontal and vertical beam sizes.
Here, we start with a desired spatial layout of the triplet and find the quadrupole strengths through numerical optimization (by minimizing the L2 norm of alpha and beta) over multiple ImpactX simulations.
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 uses scipy.optimize (methods: Nelder-Mead or L-BFGS-B) to find the quadrupole strengths by minimizing the objective. Conventional optimization algorithms like this work best if there is only a global minima in the objective.
This example can be run as a Python script: python3 run_triplet.py
.
examples/optimize_triplet/run_triplet.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
import amrex.space3d as amr
import impactx
import numpy as np
from impactx import ImpactX, distribution, elements
from scipy.optimize import minimize
# Call MPI_Init and MPI_Finalize only once:
if impactx.Config.have_mpi:
from mpi4py import MPI # noqa
verbose = False
def build_lattice(parameters: tuple, write_particles: bool) -> list:
"""
Create the quadrupole triplet.
Parameters
----------
parameters: tuple
quadrupole strengths k of quad 1/3 and quad 2.
write_particles: bool
write the particles in a beam monitor at the beginning and
end of the simulation
Returns
-------
A lattice for ImpactX: a list of impactx.elements.
"""
q1_k, q2_k = parameters
ns = 10 # number of slices per ds in the element
# enforce a mirror symmetry of the triplet
line = [
elements.Drift(ds=2.7, nslice=ns),
elements.Quad(ds=0.1, k=q1_k, nslice=ns),
elements.Drift(ds=1.4, nslice=ns),
elements.Quad(ds=0.2, k=q2_k, nslice=ns),
elements.Drift(ds=1.4, nslice=ns),
elements.Quad(ds=0.1, k=q1_k, nslice=ns),
elements.Drift(ds=2.7, nslice=ns),
]
if write_particles:
monitor = elements.BeamMonitor("monitor", backend="h5")
line = [monitor] + line + [monitor]
return line
def run(parameters: tuple, write_particles=False, write_reduced=False) -> dict:
"""
Run an ImpactX simulation with a new set of lattice parameters.
Parameters
----------
parameters: tuple
quadrupole strengths k of quad 1/3 and quad 2.
write_particles: bool
write the particles in a beam monitor at the beginning and
end of the simulation
write_reduced: bool
write the reduced diagnositcs of ImpactX to a file.
Returns
-------
A dictionary with reduced diagnositcs of ImpactX, characterizing
the beam at the end of the simulation.
"""
pp_amrex = amr.ParmParse("amrex")
pp_amrex.add("verbose", 0)
sim = ImpactX()
if verbose is False:
sim.verbose = 0
# set numerical parameters and IO control
sim.particle_shape = 2 # B-spline order
sim.space_charge = False
sim.diagnostics = write_reduced
sim.slice_step_diagnostics = write_reduced
# domain decomposition & space charge mesh
sim.init_grids()
# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 5 nm
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 100.0e-12 # 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.Waterbag(
lambdaX=2.0e-4,
lambdaY=2.0e-4,
lambdaT=3.1622776602e-5,
lambdaPx=1.1180339887e-5,
lambdaPy=1.1180339887e-5,
lambdaPt=3.1622776602e-5,
muxpx=0.894427190999916,
muypy=-0.894427190999916,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# design the accelerator lattice
sim.lattice.extend(build_lattice(parameters, write_particles=write_particles))
# run simulation
sim.evolve()
# in situ calculate the reduced beam characteristics
beam = sim.particle_container()
rbc = beam.reduced_beam_characteristics()
# clean shutdown
sim.finalize()
return rbc
def objective(parameters: tuple) -> float:
"""
A function that is evaluated by the optimizer.
Parameters
----------
parameters: tuple
quadrupole strengths k of quad 1/3 and quad 2.
Returns
-------
The L2 norm of alpha and beta of the beam at the end of the
simulation.
"""
if verbose:
print(f"Run objective with parameters={parameters}...")
rbc = run(parameters, write_particles=False, write_reduced=False)
alpha_x, alpha_y, beta_x, beta_y = (
rbc["alpha_x"],
rbc["alpha_y"],
rbc["beta_x"],
rbc["beta_y"],
)
if verbose:
print(f"alpha_x={alpha_x}, alpha_y={alpha_y}, beta_x={beta_x}, beta_y={beta_y}")
alpha_beta_is = np.array([alpha_x, alpha_y, beta_x, beta_y])
beta_x_goal = 0.55
beta_y_goal = beta_x_goal
alpha_beta_goal = np.array([0, 0, beta_x_goal, beta_y_goal])
error = np.sum((alpha_beta_is - alpha_beta_goal) ** 2)
if np.isnan(error):
error = 1.0e99
return error
if __name__ == "__main__":
# Initial guess for the quadrople strengths
initial_quad_strengths = np.array([-3, 3])
# Bounds for values to test: (min, max)
positive = (0, None)
negative = (None, 0)
bounds = [negative, positive]
# optimizer specific values
# https://docs.scipy.org/doc/scipy/reference/optimize.minimize-neldermead.html
# https://docs.scipy.org/doc/scipy/reference/optimize.minimize-lbfgsb.html
options = {
"maxiter": 1000,
}
# Call the optimizer
res = minimize(
objective,
initial_quad_strengths,
method="Nelder-Mead",
# method="L-BFGS-B",
tol=1.0e-8,
options=options,
bounds=bounds,
)
# Print the optimization result
print("Optimal parameters for k:", res.x)
print("L2 norm of alpha & beta at the optimum:", res.fun)
# analytical result:
# k: -3.5, 2.75
# alpha & beta: 0, 0, 0.55, 0.55
# final run
rbc = run(res.x, write_particles=True, write_reduced=True)
alpha_x, alpha_y, beta_x, beta_y = (
rbc["alpha_x"],
rbc["alpha_y"],
rbc["beta_x"],
rbc["beta_y"],
)
print(f"alpha_x={alpha_x} alpha_y={alpha_y}\n beta_x={beta_x} beta_y={beta_y}")
This example uses Xopt (methods: Nelder-Mead or TuRBO) to find the quadrupole strengths by minimizing the objective.
Conventional optimization algorithms like Nelder-Mead
work best if there is only a global minima in the objective.
Machine-learning based, surrogate optimization works well for highly dimensional inputs and/or to find global minima in an objective that has potentially many local minima, where conventional optimizers can get stuck.
At the same time, the ML method Bayesian Optimization (BO) is prone to over-explore an objective (at the cost of finding a point closer to the global minima).
The variation Trust Region Bayesian Optimization (TuRBO)
was developed to narrow down further on found minima.
This example can be run as a Python script: python3 tests/python/test_xopt.py
.
tests/python/test_xopt.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 The ImpactX Community
#
# Authors: Axel Huebl
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
import importlib
import amrex.space3d as amr
import impactx
import numpy as np
import pytest
from impactx import ImpactX, distribution, elements
# configure the test
verbose = True
gen_name = "Nelder-Mead" # TuRBO or Nelder-Mead
max_steps = 60
def build_lattice(parameters: dict, write_particles: bool) -> list:
"""
Create the quadrupole triplet.
Parameters
----------
parameters: dict
quadrupole strengths k of quad 1/3 and quad 2.
write_particles: bool
write the particles in a beam monitor at the beginning and
end of the simulation
Returns
-------
A lattice for ImpactX: a list of impactx.elements.
"""
q1_k, q2_k = parameters["q1_k"], parameters["q2_k"]
ns = 10 # number of slices per ds in the element
# enforce a mirror symmetry of the triplet
line = [
elements.Drift(ds=2.7, nslice=ns),
elements.Quad(ds=0.1, k=q1_k, nslice=ns),
elements.Drift(ds=1.4, nslice=ns),
elements.Quad(ds=0.2, k=q2_k, nslice=ns),
elements.Drift(ds=1.4, nslice=ns),
elements.Quad(ds=0.1, k=q1_k, nslice=ns),
elements.Drift(ds=2.7, nslice=ns),
]
if write_particles:
monitor = elements.BeamMonitor("monitor", backend="h5")
line = [monitor] + line + [monitor]
return line
def run(parameters: dict, write_particles=False, write_reduced=False) -> dict:
"""
Run an ImpactX simulation with a new set of lattice parameters.
Parameters
----------
parameters: dict
quadrupole strengths k of quad 1/3 and quad 2.
write_particles: bool
write the particles in a beam monitor at the beginning and
end of the simulation
write_reduced: bool
write the reduced diagnositcs of ImpactX to a file.
Returns
-------
A dictionary with reduced diagnositcs of ImpactX, characterizing
the beam at the end of the simulation.
"""
pp_amrex = amr.ParmParse("amrex")
pp_amrex.add("verbose", 0)
sim = ImpactX()
sim.verbose = 0
# set numerical parameters and IO control
sim.particle_shape = 2 # B-spline order
sim.space_charge = False
sim.diagnostics = write_reduced
sim.slice_step_diagnostics = write_reduced
# domain decomposition & space charge mesh
sim.init_grids()
# load a 2 GeV electron beam with an initial
# unnormalized rms emittance of 5 nm
kin_energy_MeV = 2.0e3 # reference energy
bunch_charge_C = 100.0e-12 # 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.Waterbag(
lambdaX=2.0e-4,
lambdaY=2.0e-4,
lambdaT=3.1622776602e-5,
lambdaPx=1.1180339887e-5,
lambdaPy=1.1180339887e-5,
lambdaPt=3.1622776602e-5,
muxpx=0.894427190999916,
muypy=-0.894427190999916,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# design the accelerator lattice
sim.lattice.extend(build_lattice(parameters, write_particles=write_particles))
# run simulation
sim.evolve()
# in situ calculate the reduced beam characteristics
beam = sim.particle_container()
rbc = beam.reduced_beam_characteristics()
# clean shutdown
sim.finalize()
return rbc
def objective(parameters: dict) -> dict:
"""
A function that is evaluated by the optimizer.
Parameters
----------
parameters: dict
quadrupole strengths k of quad 1/3 and quad 2.
Returns
-------
The L2 norm of alpha and beta of the beam at the end of the
simulation.
"""
if verbose:
print(f"Run objective with parameters={parameters}...")
rbc = run(parameters, write_particles=False, write_reduced=False)
alpha_x, alpha_y, beta_x, beta_y = (
rbc["alpha_x"],
rbc["alpha_y"],
rbc["beta_x"],
rbc["beta_y"],
)
if verbose:
print(
f" -> alpha_x={alpha_x}, alpha_y={alpha_y}, beta_x={beta_x}, beta_y={beta_y}"
)
alpha_beta_is = np.array([alpha_x, alpha_y, beta_x, beta_y])
beta_x_goal = 0.55
beta_y_goal = beta_x_goal
alpha_beta_goal = np.array([0, 0, beta_x_goal, beta_y_goal])
error = np.sum((alpha_beta_is - alpha_beta_goal) ** 2)
# xopt will ignore NaN results, no cleaning needed
rbc["error"] = error
return rbc
@pytest.mark.skipif(
importlib.util.find_spec("xopt") is None, reason="xopt is not available"
)
def test_xopt():
from xopt import Xopt
from xopt.evaluator import Evaluator
from xopt.vocs import VOCS
if gen_name == "TuRBO":
from xopt.generators.bayesian import UpperConfidenceBoundGenerator as Generator
gen_args = {"turbo_controller": "optimize"}
elif gen_name == "Nelder-Mead":
from xopt.generators.scipy.neldermead import NelderMeadGenerator as Generator
gen_args = {}
else:
raise RuntimeError(f"Unconfigured generator named '{gen_name}'")
# Bounds for values to test: (min, max)
positive = [0, 10]
negative = [-10, 0]
# define variables and function objectives
vocs = VOCS(
variables={
"q1_k": negative,
"q2_k": positive,
},
objectives={"error": "MINIMIZE"},
)
# create Xopt evaluator, generator, and Xopt objects
evaluator = Evaluator(function=objective)
generator = Generator(vocs=vocs, **gen_args)
X = Xopt(evaluator=evaluator, generator=generator, vocs=vocs)
# Initial guess for the quadrople strengths
initial_quad_strengths = {
"q1_k": np.array([-3]),
"q2_k": np.array([3]),
}
if gen_name == "TuRBO":
# a few random guesses
# X.random_evaluate(3)
# a few somewhat educated guesses
X.evaluate_data(initial_quad_strengths)
elif gen_name == "Nelder-Mead":
# a few somewhat educated guesses
X.generator.initial_point = initial_quad_strengths
# run optimization for 60 steps (iterations)
for i in range(max_steps):
X.step()
# Print all trials
if verbose:
print(X.data)
# plot
# X.vocs.normalize_inputs(X.data).plot(*X.vocs.variable_names, kind="scatter")
# Select the best result
best_idx, best_error = X.vocs.select_best(X.data)
best_run = X.data.iloc[best_idx]
best_ks = best_run[["q1_k", "q2_k"]].to_dict(orient="index")[best_idx[0]]
# Print the optimization result
print("Optimal parameters for k:", best_ks)
print("L2 norm of alpha & beta at the optimum:", best_run["error"].values[0])
# analytical result:
# k: -3.5, 2.75
# alpha & beta: 0, 0, 0.55, 0.55
# final run w/ detailed I/O on
rbc = run(best_ks, write_particles=True, write_reduced=True)
alpha_x, alpha_y, beta_x, beta_y = (
rbc["alpha_x"],
rbc["alpha_y"],
rbc["beta_x"],
rbc["beta_y"],
)
print(f"alpha_x={alpha_x} alpha_y={alpha_y}\n beta_x={beta_x} beta_y={beta_y}")
if __name__ == "__main__":
# Call MPI_Init and MPI_Finalize only once:
if impactx.Config.have_mpi:
from mpi4py import MPI # noqa
test_xopt()
Analyze
We run the following script to analyze correctness:
Script analysis_triplet.py
examples/optimize_triplet/analysis_triplet.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
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 = 2.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],
[
4.47213595500e-004,
4.47213595500e-004,
3.16227766017e-005,
5.0e-009,
5.0e-009,
1.0e-009,
],
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 = 2.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],
[
5.2440442742e-005,
5.2440442742e-005,
3.16227766017e-005,
5.0e-009,
5.0e-009,
1.0e-009,
],
rtol=rtol,
atol=atol,
)
Visualize
You can run the following script to visualize the optimized beam evolution over time:
Script plot_triplet.py
examples/fodo/plot_triplet.py
.#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import argparse
import os
import matplotlib.pyplot as plt
import numpy as np
# options to run this script
parser = argparse.ArgumentParser(description="Plot the quadrupole triplet benchmark.")
parser.add_argument(
"--save-png", action="store_true", help="non-interactive run: save to PNGs"
)
args = parser.parse_args()
# read reduced diagnostics
rdc_name = "diags/reduced_beam_characteristics.0"
if os.path.exists(rdc_name):
data = np.loadtxt(rdc_name, skiprows=1)
else: # OpenMP
data = np.loadtxt(rdc_name + ".0", skiprows=1)
s = data[:, 1]
beta_x = data[:, 33]
beta_y = data[:, 34]
xMin = 0.0
xMax = 8.6
yMin = 0.0
yMax = 65.0
# Plotting
plt.figure(figsize=(10, 6))
plt.xscale("linear")
plt.yscale("linear")
plt.xlim([xMin, xMax])
# plt.ylim([yMin, yMax])
plt.xlabel("s (m)", fontsize=30)
plt.ylabel("CS Twiss beta (m)", fontsize=30)
plt.xticks(fontsize=25)
plt.yticks(fontsize=25)
plt.grid(True)
# Plot the data
plt.plot(s, beta_x, "b", label="Horizontal", linewidth=2, linestyle="solid")
plt.plot(s, beta_y, "r", label="Vertical", linewidth=2, linestyle="solid")
# Show plot
plt.legend(fontsize=20)
plt.tight_layout()
if args.save_png:
plt.savefig("triplet.png")
else:
plt.show()
CS Twiss beta
Unit tests
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.
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
For every change of the ImpactX code base, each of these examples and tests are continuously tested and benchmarked.
Parameters: Python
This documents on how to use ImpactX as a Python script (python3 run_script.py
).
General
- class impactx.ImpactX
This is the central simulation class.
- property particle_shape
Control the particle B-spline order.
The order of the shape factors (splines) for the macro-particles along all spatial directions: 1 for linear, 2 for quadratic, 3 for cubic. Low-order shape factors result in faster simulations, but may lead to more noisy results. High-order shape factors are computationally more expensive, but may increase the overall accuracy of the results. For production runs it is generally safer to use high-order shape factors, such as cubic order.
- Parameters
order (int) – B-spline order
1
,2
, or3
- property n_cell
The number of grid points along each direction on the coarsest level.
- property max_level
The maximum mesh-refinement level for the simulation.
- property finest_level
The currently finest level of mesh-refinement used. This is always less or equal to
max_level
.
- property domain
The physical extent of the full simulation domain, relative to the reference particle position, in meters. When set, turns
dynamic_size
toFalse
.Note: particles that move outside the simulation domain are removed.
- property prob_relative
This is a list with
amr.max_level
+ 1 entries.By default, we dynamically extract the minimum and maximum of the particle positions in the beam. The field mesh spans, per direction, multiple times the maximum physical extent of beam particles, as given by this factor. The beam minimum and maximum extent are symmetrically padded by the mesh. For instance,
1.2
means the mesh will span 10% above and 10% below the beam;1.0
means the beam is exactly covered with the mesh.Default:
[3.0 1.0 1.0 ...]
. When set, turnsdynamic_size
toTrue
.
- property dynamic_size
Use dynamic (
True
) resizing of the field mesh or static sizing (False
).
- property space_charge
Enable (
True
) or disable (False
) space charge calculations (default:False
).Whether to calculate space charge effects.
- property mlmg_relative_tolerance
Default:
1.e-7
The relative precision with which the electrostatic space-charge fields should be calculated. More specifically, the space-charge fields are computed with an iterative Multi-Level Multi-Grid (MLMG) solver. This solver can fail to reach the default precision within a reasonable time.
- property mlmg_absolute_tolerance
Default:
0
, which means: ignoredThe absolute tolerance with which the space-charge fields should be calculated in units of \(V/m^2\). More specifically, the acceptable residual with which the solution can be considered converged. In general this should be left as the default, but in cases where the simulation state changes very little between steps it can occur that the initial guess for the MLMG solver is so close to the converged value that it fails to improve that solution sufficiently to reach the
mlmg_relative_tolerance
value.
- property mlmg_max_iters
Default:
100
Maximum number of iterations used for MLMG solver for space-charge fields calculation. In case if MLMG converges but fails to reach the desired self_fields_required_precision, this parameter may be increased.
- property mlmg_verbosity
Default:
1
The verbosity used for MLMG solver for space-charge fields calculation. Currently MLMG solver looks for verbosity levels from 0-5. A higher number results in more verbose output.
- property diagnostics
Enable (
True
) or disable (False
) diagnostics generally (default:True
). Disabling this is mostly used for benchmarking.
- property slice_step_diagnostics
Enable (
True
) or disable (False
) diagnostics every slice step in elements (default:True
).By default, diagnostics is performed at the beginning and end of the simulation. Enabling this flag will write diagnostics every step and slice step.
- property diag_file_min_digits
The minimum number of digits (default:
6
) used for the step number appended to the diagnostic file names.
- set_diag_iota_invariants(alpha, beta, tn, cn)
Set the parameters of the IOTA nonlinear lens invariants diagnostics.
- Parameters
alpha (float) – Twiss alpha
beta (float) – Twiss beta (m)
tn (float) – dimensionless strength of the nonlinear insert
cn (float) – scale parameter of the nonlinear insert (m^[1/2])
- property particle_lost_diagnostics_backend
Diagnostics for particles lost in apertures. See the
BeamMonitor
element for backend values.
- init_grids()
Initialize AMReX blocks/grids for domain decomposition & space charge mesh.
This must come first, before particle beams and lattice elements are initialized.
- add_particles(charge_C, distr, npart)
Generate and add n particles to the particle container. Note: Set the reference particle properties (charge, mass, energy) first.
Will also resize the geometry based on the updated particle distribution’s extent and then redistribute particles in according AMReX grid boxes.
- Parameters
charge_C (float) – bunch charge (C)
distr – distribution function to draw from (object from
impactx.distribution
)npart (int) – number of particles to draw
- particle_container()
Access the beam particle container (
impactx.ParticleContainer
).
- property lattice
Access the elements in the accelerator lattice. See
impactx.elements
for lattice elements.
- property periods
The number of periods to repeat the lattice.
- property abort_on_warning_threshold
(optional) Set to “low”, “medium” or “high”. Cause the code to abort if a warning is raised that exceeds the warning threshold.
- property abort_on_unused_inputs
Set to
1
to cause the simulation to fail after its completion if there were unused parameters. (default:0
for false) It is mainly intended for continuous integration and automated testing to check that all tests and inputs are adapted to API changes.
- property always_warn_immediately
If set to
1
, ImpactX immediately prints every warning message as soon as it is generated. (default:0
for false) It is mainly intended for debug purposes, in case a simulation crashes before a global warning report can be printed.
- property verbose
Controls how much information is printed to the terminal, when running ImpactX.
0
for silent, higher is more verbose. Default is1
.
- evolve()
Run the main simulation loop for a number of steps.
- resize_mesh()
Resize the mesh
domain
based on thedynamic_size
and related parameters.
- class impactx.Config
Configuration information on ImpactX that were set at compile-time.
- property have_mpi
Indicates multi-process/multi-node support via the message-passing interface (MPI). Possible values:
True
/False
Note
Particle beam particles are not yet dynamically load balanced. Please see the progress in issue 198.
- property have_gpu
Indicates GPU support. Possible values:
True
/False
- property gpu_backend
Indicates the available GPU support. Possible values:
None
,"CUDA"
(for Nvidia GPUs),"HIP"
(for AMD GPUs) or"SYCL"
(for Intel GPUs).
- property have_omp
Indicates multi-threaded CPU support via OpenMP. Possible values:
True
/False`
Set the environment variable
OMP_NUM_THREADS
to control the number of threads.Warning
By default, OpenMP spawns as many threads as there are available virtual cores on a host. When MPI and OpenMP support are used at the same time, it can easily happen that one over-subscribes the available physical CPU cores. This will lead to a severe slow-down of the simulation.
By setting appropriate environment variables for OpenMP, ensure that the number of MPI processes (ranks) per node multiplied with the number of OpenMP threads is equal to the number of physical (or virtual) CPU cores. Please see our examples in the high-performance computing (HPC) on how to run efficiently in parallel environments such as supercomputers.
Particles
- class impactx.ParticleContainer
Beam Particles in ImpactX.
This class stores particles, distributed over MPI ranks.
- add_n_particles(x, y, t, px, py, pt, qm, bchchg)
Add new particles to the container for fixed s.
- Note: This can only be used after the initialization (grids) have
been created, meaning after the call to
ImpactX.init_grids()
has been made in the ImpactX class.
- Parameters
x – positions in x
y – positions in y
t – positions as time-of-flight in c*t
px – momentum in x
py – momentum in y
pt – momentum in t
qm – charge over mass in 1/eV
bchchg – total charge within a bunch in C
- ref_particle()
Access the reference particle (
impactx.RefPart
).- Returns
return a data reference to the reference particle
- Return type
- set_ref_particle(refpart)
Set reference particle attributes.
- Parameters
refpart (impactx.RefPart) – a reference particle to copy all attributes from
- reduced_beam_characteristics()
Compute reduced beam characteristics like the position and momentum moments of the particle distribution, as well as emittance and Twiss parameters.
- Returns
beam properties with string keywords
- Return type
dict
- min_and_max_positions()
Compute the min and max of the particle position in each dimension.
- Returns
x_min, y_min, z_min, x_max, y_max, z_max
- Return type
Tuple[float, float, float, float, float, float]
- mean_and_std_positions()
Compute the mean and std of the particle position in each dimension.
- Returns
x_mean, x_std, y_mean, y_std, z_mean, z_std
- Return type
Tuple[float, float, float, float, float, float]
- redistribute()
Redistribute particles in the current mesh in x, y, z.
- class impactx.RefPart
This struct stores the reference particle attributes stored in
impactx.ParticleContainer
.- property s
integrated orbit path length, in meters
- property x
horizontal position x, in meters
- property y
vertical position y, in meters
- property z
longitudinal position y, in meters
- property t
clock time * c in meters
- property px
momentum in x, normalized to mass*c, \(p_x = \gamma \beta_x\)
- property py
momentum in y, normalized to mass*c, \(p_x = \gamma \beta_x\)
- property pz
momentum in z, normalized to mass*c, \(p_x = \gamma \beta_x\)
- property pt
energy, normalized by rest energy, \(p_t = -\gamma\)
- property gamma
Read-only: Get reference particle relativistic gamma, \(\gamma = 1/\sqrt{1-\beta^2}\)
- property beta
Read-only: Get reference particle relativistic beta, \(\beta = v/c\)
- property beta_gamma
Read-only: Get reference particle \(\beta \cdot \gamma\)
- property qm_qeeV
Read-only: Get reference particle charge to mass ratio (elementary charge/eV)
- set_charge_qe(charge_qe)
Write-only: Set reference particle charge in (positive) elementary charges.
- set_mass_MeV(massE)
Write-only: Set reference particle rest mass (MeV/c^2).
- set_kin_energy_MeV(kin_energy_MeV)
Write-only: Set reference particle kinetic energy (MeV)
- load_file(madx_file)
Load reference particle information from a MAD-X file.
- Parameters
madx_file – file name to MAD-X file with a
BEAM
entry
Initial Beam Distributions
This module provides particle beam distributions that can be used to initialize particle beams in an impactx.ParticleContainer
.
- class impactx.distribution.Gaussian(lambdax, lambday, lambdat, lambdapx, lambdapy, lambdapt, muxpx=0.0, muypy=0.0, mutpt=0.0)
A 6D Gaussian distribution.
- Parameters
lambdax – phase space position axis intercept; for zero correlation, these are the related RMS sizes (in meters)
lambday – see lambdax
lambdat – see lambdax
lambdapx – phase space momentum axis intercept; for zero correlation, these are the related normalized RMS momenta (in radians)
lambdapy – see lambdapx
lambdapt – see lambdapx
muxpx – correlation length-momentum
muypy – see muxpx
mutpt – see muxpx
- class impactx.distribution.Kurth4D(lambdax, lambday, lambdat, lambdapx, lambdapy, lambdapt, muxpx=0.0, muypy=0.0, mutpt=0.0)
A 4D Kurth distribution transversely + a uniform distribution in t + a Gaussian distribution in pt.
- class impactx.distribution.Kurth6D(lambdax, lambday, lambdat, lambdapx, lambdapy, lambdapt, muxpx=0.0, muypy=0.0, mutpt=0.0)
A 6D Kurth distribution.
R. Kurth, Quarterly of Applied Mathematics vol. 32, pp. 325-329 (1978) C. Mitchell, K. Hwang and R. D. Ryne, IPAC2021, WEPAB248 (2021)
- class impactx.distribution.KVdist(lambdax, lambday, lambdat, lambdapx, lambdapy, lambdapt, muxpx=0.0, muypy=0.0, mutpt=0.0)
A K-V distribution transversely + a uniform distribution in t + a Gaussian distribution in pt.
- class impactx.distribution.None
This distribution sets all values to zero.
- class impactx.distribution.Semigaussian(lambdax, lambday, lambdat, lambdapx, lambdapy, lambdapt, muxpx=0.0, muypy=0.0, mutpt=0.0)
A 6D Semi-Gaussian distribution (uniform in position, Gaussian in momentum).
- class impactx.distribution.Triangle(lambdax, lambday, lambdat, lambdapx, lambdapy, lambdapt, muxpx=0.0, muypy=0.0, mutpt=0.0)
A triangle distribution for laser-plasma acceleration related applications.
A ramped, triangular current profile with a Gaussian energy spread (possibly correlated). The transverse distribution is a 4D waterbag.
- class impactx.distribution.Waterbag(lambdax, lambday, lambdat, lambdapx, lambdapy, lambdapt, muxpx=0.0, muypy=0.0, mutpt=0.0)
A 6D Waterbag distribution.
- class impactx.distribution.Thermal(k, kT, kT_halo, normalize, normalize_halo, halo)
A 6D stationary thermal or bithermal distribution.
Lattice Elements
This module provides elements for the accelerator lattice.
- class impactx.elements.KnownElementsList
An iterable,
list
-like type of elements.- clear()
Clear the list to become empty.
- extend(list)
Add a list of elements to the list.
- append(element)
Add a single element to the list.
- load_file(madx_file, nslice=1)
Load and append an accelerator lattice description from a MAD-X file.
- Parameters
madx_file – file name to MAD-X file with beamline elements
nslice – number of slices used for the application of space charge
- class impactx.elements.CFbend(ds, rc, k, dx=0, dy=0, rotation=0, nslice=1)
A combined function bending magnet. This is an ideal Sbend with a normal quadrupole field component.
- Parameters
ds – Segment length in m.
rc – Radius of curvature in m.
k – Quadrupole strength in m^(-2) (MADX convention) = (gradient in T/m) / (rigidity in T-m) k > 0 horizontal focusing k < 0 horizontal defocusing
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- class impactx.elements.ConstF(ds, kx, ky, kt, dx=0, dy=0, rotation=0, nslice=1)
A linear Constant Focusing element.
- Parameters
ds – Segment length in m.
kx – Focusing strength for x in 1/m.
ky – Focusing strength for y in 1/m.
kt – Focusing strength for t in 1/m.
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- property kx
focusing x strength in 1/m
- property ky
focusing y strength in 1/m
- property kt
focusing t strength in 1/m
- class impactx.elements.DipEdge(psi, rc, g, K2, dx=0, dy=0, rotation=0)
Edge focusing associated with bend entry or exit
This model assumes a first-order effect of nonzero gap. Here we use the linear fringe field map, given to first order in g/rc (gap / radius of curvature).
References:
Brown, SLAC Report No. 75 (1982).
Hwang and S. Y. Lee, PRAB 18, 122401 (2015).
- Parameters
psi – Pole face angle in rad
rc – Radius of curvature in m
g – Gap parameter in m
K2 – Fringe field integral (unitless)
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
- class impactx.elements.Drift(ds, dx=0, dy=0, rotation=0, nslice=1)
A drift.
- Parameters
ds – Segment length in m
nslice – number of slices used for the application of space charge
- class impactx.elements.ChrDrift(ds, dx=0, dy=0, rotation=0, nslice=1)
A drift with chromatic effects included. The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained.
- Parameters
ds – Segment length in m
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- class impactx.elements.ExactDrift(ds, dx=0, dy=0, rotation=0, nslice=1)
A drift using the exact nonlinear transfer map.
- Parameters
ds – Segment length in m
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- class impactx.elements.Kicker(xkick, ykick, units, dx=0, dy=0, rotation=0)
A thin transverse kicker.
- Parameters
xkick – horizontal kick strength (dimensionless OR T-m)
ykick – vertical kick strength (dimensionless OR T-m)
units – specification of units (
"dimensionless"
in units of the magnetic rigidity of the reference particle or"T-m"
)
- class impactx.elements.Multipole(multipole, K_normal, K_skew, dx=0, dy=0, rotation=0)
A general thin multipole element.
- Parameters
multipole – index m (m=1 dipole, m=2 quadrupole, m=3 sextupole etc.)
K_normal – Integrated normal multipole coefficient (1/meter^m)
K_skew – Integrated skew multipole coefficient (1/meter^m)
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
- class impactx.elements.NonlinearLens(knll, cnll, dx=0, dy=0, rotation=0)
Single short segment of the nonlinear magnetic insert element.
A thin lens associated with a single short segment of the nonlinear magnetic insert described by V. Danilov and S. Nagaitsev, PRSTAB 13, 084002 (2010), Sect. V.A. This element appears in MAD-X as type
NLLENS
.- Parameters
knll – integrated strength of the nonlinear lens (m)
cnll – distance of singularities from the origin (m)
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
- class impactx.elements.BeamMonitor(name, backend='default', encoding='g')
A beam monitor, writing all beam particles at fixed
s
to openPMD files.If the same element
name
is used multiple times, then an output series is created with multiple outputs.The I/O backend for openPMD data dumps.
bp
is the ADIOS2 I/O library,h5
is the HDF5 format, andjson
is a simple text format.json
only works with serial/single-rank jobs. By default, the first available backend in the order given above is taken.openPMD iteration encoding determines if multiple files are created for individual output steps or not. Variable based is an experimental feature with ADIOS2.
- Parameters
name – name of the series
backend – I/O backend, e.g.,
bp
,h5
,json
encoding – openPMD iteration encoding: (v)ariable based, (f)ile based, (g)roup based (default)
- class impactx.elements.Programmable
A programmable beam optics element.
This element can be programmed to receive callback hooks into Python functions.
- property beam_particles
This is a function hook for pushing all beam particles. This accepts a function or lambda with the following arguments:
- property ref_particle
This is a function hook for pushing the reference particle. This accepts a function or lambda with the following argument:
- class impactx.elements.Quad(ds, k, dx=0, dy=0, rotation=0, nslice=1)
A Quadrupole magnet.
- Parameters
ds – Segment length in m.
k – Quadrupole strength in m^(-2) (MADX convention) = (gradient in T/m) / (rigidity in T-m) k > 0 horizontal focusing k < 0 horizontal defocusing
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- class impactx.elements.ChrQuad(ds, k, units, dx=0, dy=0, rotation=0, nslice=1)
A Quadrupole magnet, with chromatic effects included. The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained.
- Parameters
ds – Segment length in m.
k –
- Quadrupole strength in m^(-2) (MADX convention, if units = 0)
= (gradient in T/m) / (rigidity in T-m)
- OR Quadrupole strength in T/m (MaryLie convention, if units = 1)
k > 0 horizontal focusing k < 0 horizontal defocusing
units – specification of units for quadrupole field strength
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- property k
quadrupole strength in 1/m^2 (or T/m)
- property units
unit specification for quad strength
- class impactx.elements.ChrPlasmaLens(ds, g, dx=0, dy=0, rotation=0, nslice=1)
An active cylindrically symmetric plasma lens, with chromatic effects included. The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained.
- Parameters
ds – Segment length in m.
k –
- focusing strength in m^(-2) (if units = 0)
= (azimuthal magnetic field gradient in T/m) / (rigidity in T-m)
OR azimuthal magnetic field gradient in T/m (if units = 1)
units – specification of units for plasma lens focusing strength
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- property k
plasma lens focusing strength in 1/m^2 (or T/m)
- property units
unit specification for plasma lens focusing strength
- class impactx.elements.ChrAcc(ds, ez, bz, dx=0, dy=0, rotation=0, nslice=1)
Acceleration in a uniform field Ez, with a uniform solenoidal field Bz.
The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained.
- Parameters
ds – Segment length in m
ez – electric field strength in m^(-1) = (charge * electric field Ez in V/m) / (m*c^2)
bz – magnetic field strength in m^(-1) = (charge * magnetic field Bz in T) / (m*c)
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- property ez
electric field strength in 1/m
- property bz
magnetic field strength in 1/m
- class impactx.elements.RFCavity(ds, escale, freq, phase, dx=0, dy=0, rotation=0, mapsteps=1, nslice=1)
A radiofrequency cavity.
- Parameters
ds – Segment length in m.
escale – scaling factor for on-axis RF electric field in 1/m = (peak on-axis electric field Ez in MV/m) / (particle rest energy in MeV)
freq – RF frequency in Hz
phase – RF driven phase in degrees
cos_coefficients – array of
float
cosine coefficients in Fourier expansion of on-axis electric field Ez (optional); default is a 9-cell TESLA superconducting cavity model from DOI:10.1103/PhysRevSTAB.3.092001cos_coefficients – array of
float
sine coefficients in Fourier expansion of on-axis electric field Ez (optional); default is a 9-cell TESLA superconducting cavity model from DOI:10.1103/PhysRevSTAB.3.092001dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
mapsteps – number of integration steps per slice used for map and reference particle push in applied fields
nslice – number of slices used for the application of space charge
- class impactx.elements.Sbend(ds, rc, dx=0, dy=0, rotation=0, nslice=1)
An ideal sector bend.
- Parameters
ds – Segment length in m.
rc – Radius of curvature in m.
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- class impactx.elements.ExactSbend(ds, phi, B, dx=0, dy=0, rotation=0, nslice=1)
An ideal sector bend using the exact nonlinear map. The model consists of a uniform bending field B_y with a hard edge. Pole faces are normal to the entry and exit velocity of the reference particle.
References:
Bruhwiler et al, in Proc. of EPAC 98, pp. 1171-1173 (1998).
Forest et al, Part. Accel. 45, pp. 65-94 (1994).
- Parameters
ds – Segment length in m.
phi – Bend angle in degrees.
B – Magnetic field in Tesla; when B = 0 (default), the reference bending radius is defined by r0 = length / (angle in rad), corresponding to a magnetic field of B = rigidity / r0; otherwise the reference bending radius is defined by r0 = rigidity / B.
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- class impactx.elements.Buncher(V, k, dx=0, dy=0, rotation=0)
A short RF cavity element at zero crossing for bunching (MaryLie model).
- Parameters
V – Normalized RF voltage drop V = Emax*L/(c*Brho)
k – Wavenumber of RF in 1/m
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
- class impactx.elements.ShortRF(V, freq, phase, dx=0, dy=0, rotation=0)
A short RF cavity element (MAD-X model).
- Parameters
V – Normalized RF voltage V = maximum energy gain/(m*c^2)
freq – RF frequency in Hz
phase – RF synchronous phase in degrees (phase = 0 corresponds to maximum energy gain, phase = -90 corresponds go zero energy gain for bunching)
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
- class impactx.elements.ChrUniformAcc(ds, k, dx=0, dy=0, rotation=0, nslice=1)
A region of constant Ez and Bz for uniform acceleration, with chromatic effects included. The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained.
- Parameters
ds – Segment length in m.
ez – Electric field strength in m^(-1) = (particle charge in C * field Ez in V/m) / (particle mass in kg * (speed of light in m/s)^2)
bz – Magnetic field strength in m^(-1) = (particle charge in C * field Bz in T) / (particle mass in kg * speed of light in m/s)
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- class impactx.elements.SoftSolenoid(ds, bscale, cos_coefficients, sin_coefficients, dx=0, dy=0, rotation=0, mapsteps=1, nslice=1)
A soft-edge solenoid.
- Parameters
ds – Segment length in m.
bscale – Scaling factor for on-axis magnetic field Bz in inverse meters
cos_coefficients – array of
float
cosine coefficients in Fourier expansion of on-axis magnetic field Bz (optional); default is a thin-shell model from DOI:10.1016/J.NIMA.2022.166706sin_coefficients – array of
float
sine coefficients in Fourier expansion of on-axis magnetic field Bz (optional); default is a thin-shell model from DOI:10.1016/J.NIMA.2022.166706dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
mapsteps – number of integration steps per slice used for map and reference particle push in applied fields
nslice – number of slices used for the application of space charge
- class impactx.elements.Sol(ds, ks, dx=0, dy=0, rotation=0, nslice=1)
An ideal hard-edge Solenoid magnet.
- Parameters
ds – Segment length in m.
ks – Solenoid strength in m^(-1) (MADX convention) in (magnetic field Bz in T) / (rigidity in T-m)
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
nslice – number of slices used for the application of space charge
- class impactx.elements.PRot(phi_in, phi_out)
Exact map for a pole-face rotation in the x-z plane.
- Parameters
phi_in – angle of the reference particle with respect to the longitudinal (z) axis in the original frame in degrees
phi_out – angle of the reference particle with respect to the longitudinal (z) axis in the rotated frame in degrees
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
- class impactx.elements.Aperture(xmax, ymax, shape='rectangular', dx=0, dy=0, rotation=0)
A thin collimator element, applying a transverse aperture boundary.
- Parameters
xmax – maximum allowed value of the horizontal coordinate (meter)
ymax – maximum allowed value of the vertical coordinate (meter)
shape – aperture boundary shape:
"rectangular"
(default) or"elliptical"
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
- property shape
aperture type (rectangular, elliptical)
- property xmax
maximum horizontal coordinate
- property ymax
maximum vertical coordinate
- class impactx.elements.SoftQuadrupole(ds, gscale, cos_coefficients, sin_coefficients, dx=0, dy=0, rotation=0, mapsteps=1, nslice=1)
A soft-edge quadrupole.
- Parameters
ds – Segment length in m.
gscale – Scaling factor for on-axis field gradient in inverse meters
cos_coefficients – array of
float
cosine coefficients in Fourier expansion of on-axis field gradient (optional); default is a tanh fringe field model based on http://www.physics.umd.edu/dsat/docs/MaryLieMan.pdfsin_coefficients – array of
float
sine coefficients in Fourier expansion of on-axis field gradient (optional); default is a tanh fringe field model based on http://www.physics.umd.edu/dsat/docs/MaryLieMan.pdfdx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
mapsteps – number of integration steps per slice used for map and reference particle push in applied fields
nslice – number of slices used for the application of space charge
- class impactx.elements.ThinDipole(theta, rc, dx=0, dy=0, rotation=0)
A general thin dipole element.
- Parameters
theta – Bend angle (degrees)
rc – Effective curvature radius (meters)
dx – horizontal translation error in m
dy – vertical translation error in m
rotation – rotation error in the transverse plane [degrees]
Reference:
Ripken and F. Schmidt, Thin-Lens Formalism for Tracking, CERN/SL/95-12 (AP), 1995.
Coordinate Transformation
- class impactx.TransformationDirection
Enumerated type indicating whether to transform to fixed \(s\) or fixed \(t\) coordinate system when applying
impactx.coordinate_transformation
.- Parameters
to_fixed_t –
to_fixed_s –
- impactx.coordinate_transformation(pc, direction)
Function to transform the coordinates of the particles in a particle container either to fixed \(t\) or to fixed \(s\).
- Parameters
pc –
impactx.particle_container
whose particle coordinates are to be transformed.direction – enumerated type
impactx.TransformationDirection
, indicates whether to transform to fixed \(s\) or fixed \(t\).
Parameters: Inputs File
This documents on how to use ImpactX with an inputs file (impactx input_file.in
).
Note
The AMReX parser (see Math parser and user-defined constants) is used for the right-hand-side of all input parameters that consist of one or more integers or floats, so expressions like <species_name>.density_max = "2.+1."
and/or using user-defined constants are accepted.
Overall simulation parameters
max_step
(integer
)The number of PIC cycles to perform.
stop_time
(float
; in seconds)The maximum physical time of the simulation. Can be provided instead of
max_step
. If bothmax_step
andstop_time
are provided, both criteria are used and the simulation stops when the first criterion is hit.
amrex.abort_on_out_of_gpu_memory
(0
or1
; default is1
for true)When running on GPUs, memory that does not fit on the device will be automatically swapped to host memory when this option is set to
0
. This will cause severe performance drops. Note that even with this set to1
ImpactX will not catch all out-of-memory events yet when operating close to maximum device memory. Please also see the documentation in AMReX.
amrex.the_arena_is_managed
(0
or1
; default is0
for false)When running on GPUs, device memory that is accessed from the host will automatically be transferred with managed memory. This is useful for convenience during development, but has sometimes severe performance and memory footprint implications if relied on (and sometimes vendor bugs). For all regular ImpactX operations, we therefore do explicit memory transfers without the need for managed memory and thus changed the AMReX default to false. Please also see the documentation in AMReX.
amrex.omp_threads
(system
,nosmt
or positive integer; default isnosmt
)An integer number can be set in lieu of the
OMP_NUM_THREADS
environment variable to control the number of OpenMP threads to use for theOMP
compute backend on CPUs. By default, we use thenosmt
option, which overwrites the OpenMP default of spawning one thread per logical CPU core, and instead only spawns a number of threads equal to the number of physical CPU cores on the machine. If set, the environment variableOMP_NUM_THREADS
takes precedence oversystem
andnosmt
, but not over integer numbers set in this option.
amrex.abort_on_unused_inputs
(0
or1
; default is0
for false)When set to
1
, this option causes the simulation to fail after its completion if there were unused parameters. It is mainly intended for continuous integration and automated testing to check that all tests and inputs are adapted to API changes.
impactx.always_warn_immediately
(0
or1
; default is0
for false)If set to
1
, ImpactX immediately prints every warning message as soon as it is generated. It is mainly intended for debug purposes, in case a simulation crashes before a global warning report can be printed.
impactx.abort_on_warning_threshold
(string:low
,medium
orhigh
) optionalOptional threshold to abort as soon as a warning is raised. If the threshold is set, warning messages with priority greater than or equal to the threshold trigger an immediate abort. It is mainly intended for debug purposes, and is best used with
impactx.always_warn_immediately=1
. For more information on the warning logger, see this section of the WarpX documentation.
impactx.verbose
(int:0
for silent, higher is more verbose; default is1
) optionalControls how much information is printed to the terminal, when running ImpactX.
Setting up the field mesh
ImpactX uses an AMReX grid of boxes to organize and parallelize the simulation domain. These boxes also contain a field mesh, if space charge calculations are enabled.
amr.n_cell
(3 integers) optional (default: 1 blocking_factor per MPI process)The number of grid points along each direction (on the coarsest level)
amr.max_level
(integer
, default:0
)When using mesh refinement, the number of refinement levels that will be used.
Use
0
in order to disable mesh refinement.
amr.ref_ratio
(integer
per refined level, default:2
)When using mesh refinement, this is the refinement ratio per level. With this option, all directions are fined by the same ratio.
amr.ref_ratio_vect
(3 integers for x,y,z per refined level)When using mesh refinement, this can be used to set the refinement ratio per direction and level, relative to the previous level.
Example: for three levels, a value of
2 2 4 8 8 16
refines the first level by 2-fold in x and y and 4-fold in z compared to the coarsest level (level 0/mother grid); compared to the first level, the second level is refined 8-fold in x and y and 16-fold in z.
Note
Field boundaries for space charge calculation are located at the outer ends of the field mesh. We currently assume Dirichlet boundary conditions with zero potential (a mirror charge). Thus, to emulate open boundaries, consider adding enough vacuum padding to the beam. This will be improved in future versions.
Note
Particles that move outside the simulation domain are removed.
geometry.dynamic_size
(boolean
) optional (default:true
for dynamic)Use dynamic (
true
) resizing of the field mesh, viageometry.prob_relative
, or static sizing (false
), viageometry.prob_lo
/geometry.prob_hi
.
geometry.prob_relative
(positivefloat
array withamr.max_level
entries, unitless) optional (default:3.0 1.0 1.0 ...
)By default, we dynamically extract the minimum and maximum of the particle positions in the beam. The field mesh spans, per direction, multiple times the maximum physical extent of beam particles, as given by this factor. The beam minimum and maximum extent are symmetrically padded by the mesh. For instance,
1.2
means the mesh will span 10% above and 10% below the beam;1.0
means the beam is exactly covered with the mesh.
geometry.prob_lo
andgeometry.prob_hi
(3 floats, in meters) optional (required ifgeometry.dynamic_size
isfalse
)The extent of the full simulation domain relative to the reference particle position. This can be used to explicitly size the simulation box and ignore
geometry.prob_relative
.This box is rectangular, and thus its extent is given here by the coordinates of the lower corner (
geometry.prob_lo
) and upper corner (geometry.prob_hi
). The first axis of the coordinates is x and the last is z.
Domain Boundary Conditions
Note
TODO :-)
Initial Beam Distributions
beam.npart
(integer
) number of weighted simulation particlesbeam.units
(string
) currently, onlystatic
is supported.beam.kin_energy
(float
, in MeV) beam kinetic energybeam.charge
(float
, in C) bunch chargebeam.particle
(string
) particle type: currently eitherelectron
,positron
orproton
beam.distribution
(string
)Indicates the initial distribution type. For additional information, consult the documentation on Beam Distribution Input. For all except the
thermal
distribution we allow input in two forms:Parameters that describe the phase space ellipse and position-momentum correlations:
beam.lambdaX
(float
, in meters) phase space ellipse intersection with Xbeam.lambdaY
(float
, in meters) phase space ellipse intersection with Ybeam.lambdaT
(float
, in meters) phase space ellipse intersection with T, normalized by multiplying with the speed of light cbeam.lambdaPx
(float
, in radians) phase space ellipse intersection with Pxbeam.lambdaPy
(float
, in radians) phase space ellipse intersection with Pybeam.lambdaPt
(float
, in radians) phase space ellipse intersection with Ptbeam.muxpx
(float
, dimensionless, default:0
) correlation X-Pxbeam.muypy
(float
, dimensionless, default:0
) correlation Y-Pybeam.mutpt
(float
, dimensionless, default:0
) correlation T-Pt
Courant-Snyder / Twiss parameters. To enable input via CS / Twiss parameters, add the suffix
_from_cs
orfrom_twiss
to the name of the distribution. Use the following parameters to characterize it:beam.alphaX
(float
, dimensionless, default:0
) CS / Twiss \(\alpha\) for Xbeam.alphaY
(float
, dimensionless, default:0
) CS / Twiss \(\alpha\) for Ybeam.alphaT
(float
, dimensionless, default:0
) CS / Twiss \(\alpha\) for Tbeam.betaX
(float
, in meters) CS / Twiss \(\beta\) for Xbeam.betaY
(float
, in meters) CS / Twiss \(\beta\) for Ybeam.betaT
(float
, in meters) CS / Twiss \(\beta\) for Tbeam.emittX
(float
, in meters times radian) geometric (unnormalized) emittance \(\epsilon\) in Xbeam.emittY
(float
, in meters times radian) geometric (unnormalized) emittance \(\epsilon\) in Ybeam.emittT
(float
, in meters times radian) geometric (unnormalized) emittance \(\epsilon\) in T
The following distributions are available:
waterbag
orwaterbag_from_cs
/waterbag_from_twiss
for initial Waterbag distribution.kurth6d
orkurth6d_from_cs
/kurth6d_from_twiss
for initial 6D Kurth distribution.gaussian
orgaussian_from_cs
/gaussian_from_twiss
for initial 6D Gaussian (normal) distribution.kvdist
orkvdist_from_cs
/kvdist_from_twiss
for initial K-V distribution in the transverse plane. The distribution is uniform in t and Gaussian in pt.kurth4d
orkurth4d_from_cs
/kurth4d_from_twiss
for initial 4D Kurth distribution in the transverse plane. The distribution is uniform in t and Gaussian in pt.semigaussian
orsemigaussian_from_cs
/semigaussian_from_twiss
for initial Semi-Gaussian distribution. The distribution is uniform within a cylinder in (x,y,z) and Gaussian in momenta (px,py,pt).triangle
ortriangle_from_cs
/triangle_from_twiss
a triangle distribution for laser-plasma acceleration related applications. A ramped, triangular current profile with a Gaussian energy spread (possibly correlated). The transverse distribution is a 4D waterbag.thermal
for a 6D stationary thermal or bithermal distribution. This distribution type is described, for example in: R. D. Ryne et al, “A Test Suite of Space-Charge Problems for Code Benchmarking”, in Proc. EPAC2004, Lucerne, Switzerland. C. E. Mitchell et al, “ImpactX Modeling of Benchmark Tests for Space Charge Validation”, in Proc. HB2023, Geneva, Switzerland. With additional parameters:beam.k
(float
, in inverse meters) external focusing strengthbeam.kT
(float
, dimensionless) temperature of core population= < p_x^2 > = < p_y^2 >, where all momenta are normalized by the reference momentum
beam.kT_halo
(float
, dimensionless, defaultkT
) temperature of halo populationbeam.normalize
(float
, dimensionless) normalizing constant for core populationbeam.normalize_halo
(float
, dimensionless) normalizing constant for halo populationbeam.halo
(float
, dimensionless) fraction of charge in halo
Lattice Elements
lattice.elements
(list of strings
) optional (default: no elements)A list of names (one name per lattice element), in the order that they appear in the lattice.
lattice.periods
(integer
) optional (default:1
)The number of periods to repeat the lattice.
lattice.reverse
(boolean
) optional (default:false
)Reverse the list of elements in the lattice. If
reverse
andperiods
both appear, thenreverse
is applied beforeperiods
.
lattice.nslice
(integer
) optional (default:1
)A positive integer specifying the number of slices used for the application of space charge in all elements; overwritten by element parameter “nslice”
<element_name>.type
(string
)Indicates the element type for this lattice element. This should be one of:
cfbend
for a combined function bending magnet. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.rc
(float
, in meters) the bend radius<element_name>.k
(float
, in inverse meters squared) the quadrupole strength= (magnetic field gradient in T/m) / (magnetic rigidity in T-m)
k > 0 horizontal focusing
k < 0 horizontal defocusing
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
drift
for a free drift. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
drift_chromatic
for a free drift, with chromatic effects included.The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained. This requires these additional parameters:
<element_name>.ds
(float
, in meters) the segment length<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
drift_exact
for a free drift, using the exact nonlinear map. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
quad
for a quadrupole. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.k
(float
, in inverse meters squared) the quadrupole strength= (magnetic field gradient in T/m) / (magnetic rigidity in T-m)
k > 0 horizontal focusing
k < 0 horizontal defocusing
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
quad_chromatic
for A Quadrupole magnet, with chromatic effects included.The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained. This requires these additional parameters:
<element_name>.ds
(float
, in meters) the segment length<element_name>.k
(float
, in inverse meters squared OR in T/m) the quadrupole strength= (magnetic field gradient in T/m) / (magnetic rigidity in T-m) - if units = 0
OR = magnetic field gradient in T/m - if units = 1
k > 0 horizontal focusing
k < 0 horizontal defocusing
<element_name>.units
(integer
) specification of units (default:0
)<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
quadrupole_softedge
for a soft-edge quadrupole. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.gscale
(float
, in inverse meters) Scaling factor for on-axis magnetic field gradient<element_name>.cos_coefficients
(array offloat
) cos coefficients in Fourier expansion of the on-axis field gradient (optional); default is a tanh fringe field model from MaryLie 3.0<element_name>.sin_coefficients
(array offloat
) sin coefficients in Fourier expansion of the on-axis field gradient (optional); default is a tanh fringe field model from MaryLie 3.0<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.mapsteps
(integer
) number of integration steps per slice used for map and reference particle push in applied fields(default:
1
)
<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
plasma_lens_chromatic
for an active cylindrically-symmetric plasma lens, with chromatic effects included.The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained. This requires these additional parameters:
<element_name>.ds
(float
, in meters) the segment length<element_name>.k
(float
, in inverse meters squared OR in T/m) the plasma lens focusing strength= (azimuthal magnetic field gradient in T/m) / (magnetic rigidity in T-m) - if units = 0
OR = azimuthal magnetic field gradient in T/m - if units = 1
<element_name>.units
(integer
) specification of units (default:0
)<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
sbend
for a bending magnet. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.rc
(float
, in meters) the bend radius<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
sbend_exact
for a bending magnet using the exact nonlinear map for the bend body. The map corresponds to the map described in:D. L. Bruhwiler et al, in Proc. of EPAC 98, pp. 1171-1173 (1998), E. Forest et al, Part. Accel. 45, pp. 65-94 (1994). The model consists of a uniform bending field B_y with a hard edge. Pole faces are normal to the entry and exit velocity of the reference particle. This requires these additional parameters:
<element_name>.ds
(float
, in meters) the segment length<element_name>.phi
(float
, in degrees) the bend angle<element_name>.B
(float
, in Tesla) the bend magnetic field; when B = 0 (default), the reference bending radius is defined by r0 = length / (angle in rad), corresponding to a magnetic field of B = rigidity / r0; otherwise the reference bending radius is defined by r0 = rigidity / B<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
solenoid
for an ideal hard-edge solenoid magnet. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.ks
(float
, in meters) Solenoid strength in m^(-1) (MADX convention)= (magnetic field Bz in T) / (rigidity in T-m)
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
solenoid_softedge
for a soft-edge solenoid. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.bscale
(float
, in inverse meters) Scaling factor for on-axis magnetic field Bz<element_name>.cos_coefficients
(array offloat
) cos coefficients in Fourier expansion of the on-axis magnetic field Bz (optional); default is a thin-shell model from DOI:10.1016/J.NIMA.2022.166706<element_name>.sin_coefficients
(array offloat
) sin coefficients in Fourier expansion of the on-axis magnetic field Bz (optional); default is a thin-shell model from DOI:10.1016/J.NIMA.2022.166706<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.mapsteps
(integer
) number of integration steps per slice used for map and reference particle push in applied fields (default:1
)<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
dipedge
for dipole edge focusing. This requires these additional parameters:<element_name>.psi
(float
, in radians) the pole face rotation angle<element_name>.rc
(float
, in meters) the bend radius<element_name>.g
(float
, in meters) the gap size<element_name>.K2
(float
, dimensionless) normalized field integral for fringe field<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane
constf
for a constant focusing element. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.kx
(float
, in 1/meters) the horizontal focusing strength<element_name>.ky
(float
, in 1/meters) the vertical focusing strength<element_name>.kt
(float
, in 1/meters) the longitudinal focusing strength<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
rfcavity
a radiofrequency cavity. This requires these additional parameters:<element_name>.ds
(float
, in meters) the segment length<element_name>.escale
(float
, in 1/m) scaling factor for on-axis RF electric field= (peak on-axis electric field Ez in MV/m) / (particle rest energy in MeV)
<element_name>.freq
(float
, in Hz) RF frequency<element_name>.phase
(float
, in degrees) RF driven phase<element_name>.cos_coefficients
(array offloat
) cosine coefficients in Fourier expansion of on-axis electric field Ez (optional); default is a 9-cell TESLA superconducting cavity model from DOI:10.1103/PhysRevSTAB.3.092001<element_name>.cos_coefficients
(array offloat
) sine coefficients in Fourier expansion of on-axis electric field Ez (optional); default is a 9-cell TESLA superconducting cavity model from DOI:10.1103/PhysRevSTAB.3.092001<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.mapsteps
(integer
) number of integration steps per slice used for map and reference particle push in applied fields (default:1
)<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
buncher
for a short RF cavity (linear) bunching element. This requires these additional parameters:<element_name>.V
(float
, dimensionless) normalized voltage drop across the cavity= (maximum voltage drop in Volts) / (speed of light in m/s * magnetic rigidity in T-m)
<element_name>.k
(float
, in 1/meters) the RF wavenumber= 2*pi/(RF wavelength in m)
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane
shortrf
for a short RF cavity element. This requires these additional parameters:<element_name>.V
(float
, dimensionless) normalized voltage drop across the cavity= (maximum energy gain in MeV) / (particle rest energy in MeV)
<element_name>.freq
(float
, in Hz) the RF frequency<element_name>.phase
(float
, in degrees) the synchronous RF phasephase = 0: maximum energy gain (on-crest)
phase = -90 deg: zero energy gain for bunching
phase = 90 deg: zero energy gain for debunching
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane
uniform_acc_chromatic
for a region of uniform acceleration, with chromatic effects included.The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained. This requires these additional parameters:
<element_name>.ds
(float
, in meters) the segment length<element_name>.ez
(float
, in inverse meters) the electric field strength= (particle charge in C * electric field Ez in V/m) / (particle mass in kg * (speed of light in m/s)^2)
<element_name>.bz
(float
, in inverse meters) the magnetic field strength= (particle charge in C * magnetic field Bz in T) / (particle mass in kg * speed of light in m/s)
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane<element_name>.nslice
(integer
) number of slices used for the application of space charge (default:1
)
multipole
for a thin multipole element. This requires these additional parameters:<element_name>.multipole
(integer
, dimensionless) order of multipole(m = 1) dipole, (m = 2) quadrupole, (m = 3) sextupole, etc.
<element_name>.k_normal
(float
, in 1/meters^m) integrated normal multipole coefficient (MAD-X convention)= 1/(magnetic rigidity in T-m) * (derivative of order m-1 of By with respect to x)
<element_name>.k_skew
(float
, in 1/meters^m) integrated skew multipole strength (MAD-X convention)<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane
nonlinear_lens
for a thin IOTA nonlinear lens element. This requires these additional parameters:<element_name>.knll
(float
, in meters) integrated strength of the lens segment (MAD-X convention)= dimensionless lens strength * c parameter**2 * length / Twiss beta
<element_name>.cnll
(float
, in meters) distance of the singularities from the origin (MAD-X convention)= c parameter * sqrt(Twiss beta)
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane
prot
for an exact pole-face rotation in the x-z plane. This requires these additional parameters:<element_name>.phi_in
(float
, in degrees) angle of the reference particle with respect to the longitudinal (z) axis in the original frame<element_name>.phi_out
(float
, in degrees) angle of the reference particle with respect to the longitudinal (z) axis in the rotated frame
kicker
for a thin transverse kicker. This requires these additional parameters:<element_name>.xkick
(float
, dimensionless OR in T-m) the horizontal kick strength<element_name>.ykick
(float
, dimensionless OR in T-m) the vertical kick strength<element_name>.units
(string
) specification of units:dimensionless
(default, in units of the magnetic rigidity of the reference particle) orT-m
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane
thin_dipole
for a thin dipole element. This requires these additional parameters:<element_name>.theta
(float
, in degrees) dipole bend angle<element_name>.rc
(float
, in meters) effective radius of curvature<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane
aperture
for a thin collimator element applying a transverse aperture boundary. This requires these additional parameters:<element_name>.xmax
(float
, in meters) maximum value of the horizontal coordinate<element_name>.ymax
(float
, in meters) maximum value of the vertical coordinate<element_name>.shape
(string
) shape of the aperture boundary:rectangular
(default) orelliptical
<element_name>.dx
(float
, in meters) horizontal translation error<element_name>.dy
(float
, in meters) vertical translation error<element_name>.rotation
(float
, in degrees) rotation error in the transverse plane
beam_monitor
a beam monitor, writing all beam particles at fixeds
to openPMD files. If the same element name is used multiple times, then an output series is created with multiple outputs.<element_name>.name
(string
, default value:<element_name>
)The output series name to use. By default, output is created under
diags/openPMD/<element_name>.<backend>
.<element_name>.backend
(string
, default value:default
)I/O backend for openPMD data dumps.
bp
is the ADIOS2 I/O library,h5
is the HDF5 format, andjson
is a simple text format.json
only works with serial/single-rank jobs. By default, the first available backend in the order given above is taken.<element_name>.encoding
(string
, default value:g
)openPMD iteration encoding: (v)ariable based, (f)ile based, (g)roup based (default) variable based is an experimental feature with ADIOS2.
line
a sub-lattice (line) of elements to append to the lattice.<element_name>.elements
(list of strings
) optional (default: no elements) A list of names (one name per lattice element), in the order that they appear in the lattice.<element_name>.reverse
(boolean
) optional (default:false
) Reverse the list of elements in the line before appending to the lattice.<element_name>.repeat
(integer
) optional (default:1
) Repeat the line multiple times before appending to the lattice. Note: Ifreverse
andrepeat
both appear, thenreverse
is applied beforerepeat
.
Distribution across MPI ranks and parallelization
amr.max_grid_size
(integer
) optional (default:128
)Maximum allowable size of each subdomain (expressed in number of grid points, in each direction). Each subdomain has its own ghost cells, and can be handled by a different MPI rank ; several OpenMP threads can work simultaneously on the same subdomain.
If
max_grid_size
is such that the total number of subdomains is larger that the number of MPI ranks used, than some MPI ranks will handle several subdomains, thereby providing additional flexibility for load balancing.When using mesh refinement, this number applies to the subdomains of the coarsest level, but also to any of the finer level.
Math parser and user-defined constants
ImpactX uses AMReX’s math parser that reads expressions in the input file. It can be used in all input parameters that consist of one or more integers or floats. Integer input expecting boolean, 0 or 1, are not parsed. Note that when multiple values are expected, the expressions are space delimited. For integer input values, the expressions are evaluated as real numbers and the final result rounded to the nearest integer. See this section of the AMReX documentation for a complete list of functions supported by the math parser.
ImpactX constants
ImpactX will provide a few pre-defined constants, that can be used for any parameter that consists of one or more floats.
Note
Develop, such as:
q_e |
elementary charge |
m_e |
electron mass |
m_p |
proton mass |
m_u |
unified atomic mass unit (Dalton) |
epsilon0 |
vacuum permittivity |
mu0 |
vacuum permeability |
clight |
speed of light |
pi |
math constant pi |
See in WarpX the file Source/Utils/WarpXConst.H
for the values.
User-defined constants
Users can define their own constants in the input file.
These constants can be used for any parameter that consists of one or more integers or floats.
User-defined constant names can contain only letters, numbers and the character _
.
The name of each constant has to begin with a letter. The following names are used
by ImpactX, and cannot be used as user-defined constants: x
, y
, z
, X
, Y
, t
.
The values of the constants can include the predefined ImpactX constants listed above as well as other user-defined constants.
For example:
my_constants.a0 = 3.0
my_constants.z_plateau = 150.e-6
my_constants.n0 = 1.e22
my_constants.wp = sqrt(n0*q_e**2/(epsilon0*m_e))
Coordinates
Besides, for profiles that depend on spatial coordinates (the plasma momentum distribution or the laser field, see below Particle initialization
and Laser initialization
), the parser will interpret some variables as spatial coordinates.
These are specified in the input parameter, i.e., density_function(x,y,z)
and field_function(X,Y,t)
.
The parser reads python-style expressions between double quotes, for instance
"a0*x**2 * (1-y*1.e2) * (x>0)"
is a valid expression where a0
is a
user-defined constant (see above) and x
and y
are spatial coordinates. The names are case sensitive. The factor
(x>0)
is 1
where x>0
and 0
where x<=0
. It allows the user to
define functions by intervals.
Alternatively the expression above can be written as if(x>0, a0*x**2 * (1-y*1.e2), 0)
.
Numerics and algorithms
algo.particle_shape
(integer
;1
,2
, or3
)The order of the shape factors (splines) for the macro-particles along all spatial directions: 1 for linear, 2 for quadratic, 3 for cubic. Low-order shape factors result in faster simulations, but may lead to more noisy results. High-order shape factors are computationally more expensive, but may increase the overall accuracy of the results. For production runs it is generally safer to use high-order shape factors, such as cubic order.
algo.space_charge
(boolean
, optional, default:false
)Whether to calculate space charge effects.
algo.mlmg_relative_tolerance
(float
, optional, default:1.e-7
)The relative precision with which the electrostatic space-charge fields should be calculated. More specifically, the space-charge fields are computed with an iterative Multi-Level Multi-Grid (MLMG) solver. This solver can fail to reach the default precision within a reasonable time.
algo.mlmg_absolute_tolerance
(float
, optional, default:0
, which means: ignored)The absolute tolerance with which the space-charge fields should be calculated in units of V/m^2. More specifically, the acceptable residual with which the solution can be considered converged. In general this should be left as the default, but in cases where the simulation state changes very little between steps it can occur that the initial guess for the MLMG solver is so close to the converged value that it fails to improve that solution sufficiently to reach the mlmg_relative_tolerance value.”
algo.mlmg_max_iters
(integer
, optional, default:100
)Maximum number of iterations used for MLMG solver for space-charge fields calculation. In case if MLMG converges but fails to reach the desired self_fields_required_precision, this parameter may be increased.
algo.mlmg_verbosity
(integer
, optional, default:1
)The verbosity used for MLMG solver for space-charge fields calculation. Currently MLMG solver looks for verbosity levels from 0-5. A higher number results in more verbose output.
Diagnostics and output
diag.enable
(boolean
, optional, default:true
) Enable or disable diagnostics generally. Disabling this is mostly used for benchmarking.This option is ignored for the openPMD output elements (remove them from the lattice to disable).
diag.slice_step_diagnostics
(boolean
, optional, default:false
) By default, diagnostics is performed at the beginning and end of the simulation. Enabling this flag will write diagnostics every step and slice stepdiag.file_min_digits
(integer
, optional, default:6
)The minimum number of digits used for the step number appended to the diagnostic file names.
diag.backend
(string
, default value:default
)Diagnostics for particles lost in apertures, stored as
diags/openPMD/particles_lost.*
at the end of the simulation. See thebeam_monitor
element for backend values.
Reduced Diagnostics
Reduced diagnostics allow the user to compute some reduced quantity (invariants of motion, particle temperature, max of a field, …) and write a small amount of data to text files. Reduced diagnostics are run in situ with the simulation.
Diagnostics related to integrable optics in the IOTA nonlinear magnetic insert element:
diag.alpha
(float
, unitless) Twiss alpha of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I. Horizontal and vertical values must be equal.diag.beta
(float
, meters) Twiss beta of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I. Horizontal and vertical values must be equal.diag.tn
(float
, unitless) dimensionless strength of the IOTA nonlinear magnetic insert element used for computing H and I.diag.cn
(float
, meters^(1/2)) scale factor of the IOTA nonlinear magnetic insert element used for computing H and I.
Checkpoints and restart
Note
ImpactX will support checkpoints/restart via AMReX.
The checkpoint capability can be turned with regular diagnostics: <diag_name>.format = checkpoint
.
amr.restart
(string)Name of the checkpoint file to restart from. Returns an error if the folder does not exist or if it is not properly formatted.
Intervals parser
Note
TODO :-)
ImpactX can parse time step interval expressions of the form start:stop:period
, e.g.
1:2:3, 4::, 5:6, :, ::10
.
A comma is used as a separator between groups of intervals, which we call slices.
The resulting time steps are the union set of all given slices.
White spaces are ignored.
A single slice can have 0, 1 or 2 colons :
, just as numpy slices, but with inclusive upper bound for stop
.
For 0 colon the given value is the period
For 1 colon the given string is of the type
start:stop
For 2 colons the given string is of the type
start:stop:period
Any value that is not given is set to default.
Default is 0
for the start, std::numeric_limits<int>::max()
for the stop and 1
for the
period.
For the 1 and 2 colon syntax, actually having values in the string is optional
(this means that ::5
, 100 ::10
and 100 :
are all valid syntaxes).
All values can be expressions that will be parsed in the same way as other integer input parameters.
Examples
something_intervals = 50
-> do something at timesteps 0, 50, 100, 150, etc. (equivalent tosomething_intervals = ::50
)something_intervals = 300:600:100
-> do something at timesteps 300, 400, 500 and 600.something_intervals = 300::50
-> do something at timesteps 300, 350, 400, 450, etc.something_intervals = 105:108,205:208
-> do something at timesteps 105, 106, 107, 108, 205, 206, 207 and 208. (equivalent tosomething_intervals = 105 : 108 : , 205 : 208 :
)something_intervals = :
orsomething_intervals = ::
-> do something at every timestep.something_intervals = 167:167,253:253,275:425:50
do something at timesteps 167, 253, 275, 325, 375 and 425.
This is essentially the python slicing syntax except that the stop is inclusive
(0:100
contains 100) and that no colon means that the given value is the period.
Note that if a given period is zero or negative, the corresponding slice is disregarded.
For example, something_intervals = -1
deactivates something
and
something_intervals = ::-1,100:1000:25
is equivalent to something_intervals = 100:1000:25
.
Workflows
This section collects typical user workflows and best practices for ImpactX.
Note
TODO: Add more workflows as in https://warpx.readthedocs.io/en/latest/usage/workflows.html
Data Analysis
Data Analysis
Beam Monitor
ImpactX provides a zero-sized beam monitor element that can be placed in lattices to output the particle beam at multiple positions in a lattice. Output is written in the standardized, open particle-mesh data schema (openPMD) and is compatible with many codes and data analysis frameworks.
For data analysis of openPMD data, see examples of many supported tools, Python libraries and frameworks. Exporting data to ASCII is possible, too.
See also WarpX’ documentation on openPMD.
Additional Beam Attributes
We add the following additional attributes on the openPMD beam
species at the monitor position.
Reference particle:
beta_ref
reference particle normalized velocity \(\beta = v/c\)gamma_ref
reference particle Lorentz factor \(\gamma = 1/\sqrt{1-\beta^2}\)s_ref
integrated orbit path length, in metersx_ref
horizontal position x, in metersy_ref
vertical position y, in metersz_ref
longitudinal position z, in meterst_ref
clock time * c in meterspx_ref
momentum in x, normalized to mass*c, \(p_x = \gamma \beta_x\)py_ref
momentum in y, normalized to mass*c, \(p_y = \gamma \beta_y\)pz_ref
momentum in z, normalized to mass*c, \(p_z = \gamma \beta_z\)pt_ref
energy, normalized by rest energy, \(p_t = -\gamma\)mass
reference rest mass, in kgcharge
reference charge, in C
Example to print the integrated orbit path length s
at each beam monitor position:
import openpmd_api as io
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
for k_i, i in series.iterations.items():
beam = i.particles["beam"]
s_ref = beam.get_attribute("s_ref")
print(f"step {k_i:>3}: s_ref={s_ref}")
Reduced Beam Characteristics
ImpactX calculates reduced beam characteristics like averaged positions, momenta, beam emittances and Courant-Snyder (Twiss) parameters during runtime.
These quantities are calculated before, after, and during each step of the simulation.
If diag.slice_step_diagnostics
is enabled, they will also be calculated during each slice of each beamline element.
The code writes out the values in an ASCII file prefixed reduced_beam_characteristics
containing the follow columns:
step
Iteration within the simulation
s
,ref_beta_gamma
Reference particle coordinate
s
(unit: meter) and relativistic momentum normalized by the particle mass and the speed of light (unit: dimensionless)
x_mean/min/max
,y_mean/min/max
,t_mean/min/max
Average / minimum / maximum beam particle position in the dimensions of
x
,y
(transverse coordinates, unit: meter), andt
(normalized time difference \(ct\), unit: meter)
sig_x
,sig_y
,sig_t
RMS of the average beam particle positions (unit: meter)
px_mean/min/max
,py_mean/min/max
,pt_mean/min/max
Average / minimum / maximum beam momenta normalized by reference particle momentum (unit: dimensionless, radians for transverse momenta)
sig_px
,sig_py
,sig_pt
RMS of the average beam momenta (energy difference for
pt
) (unit: dimensionless)
emittance_x
,emittance_y
,emittance_t
Normalized beam emittance (unit: meter)
alpha_x
,alpha_y
,alpha_t
Courant-Snyder (Twiss) alpha (unit: dimensionless)
beta_x
,beta_y
,beta_t
Courant-Snyder (Twiss) beta (unit: meter)
charge
Cumulated beam charge (unit: Coulomb)
Theory
Introduction
Concepts
Reference Trajectory
ImpactX is an s-based beam dynamics code with space charge. Particles are tracked using the longitudinal accelerator lattice position \(s\) as the independent dynamical variable. Particle phase space coordinates are specified relative to a nominal reference trajectory.
The reference trajectory plays an important role in ImpactX. In addition to specifying the nominal machine orbit, it specifies the relationship between the local beam coordinate system and the global lab coordinate system. A reference trajectory is favorable instead of, e.g., differences to the beam centroid, because:
the reference trajectory is the ideal single-particle orbit used as part of the optics design, so it is better that it can be computed independently of the beam distribution,
differences between the beam centroid and the reference particle are important when investigating the effects of misalignments and errors for beamline designs,
the fields are usually specified relative to the reference trajectory, and the dynamics becomes more nonlinear as one moves away from the reference trajectory, so knowing how far the beam particles are form the reference trajectory is important for accuracy,
we want to keep track of the reference trajectory in global coordinates, and that global information is not available in the particle data alone, so it would need to be specified in some other way.
The reference values of \(z=ct\) and \(s\) should be identical until reaching a bending element. (If the lattice contains no bending elements, then they should coincide.) Also, the reference value of \(ct\) coincides with the value of \(s\) in the ultrarelativistic limit. More generally, the derivative \(ds/d(ct) = \beta\), where the relativistic \(\beta = \sqrt{1-\frac{1}{p_t^2}}\).
Collective Effects
Collective effects from space charge of the particle beam are solved between steps in \(s\). One can set the number of slice steps through each lattice element for the application of the space charge push.
Note
Currently, space charge kicks are calculated in 3D. A 2D space-charge solver for purely transversal effects will be added in the future.
Coordinates and Units
Each particle in the beam is described at fixed \(s\) by a set of 6 canonical phase space variables (x [m], px, y [m], py, t [m], pt). Coordinates x and y denote the horizontal and vertical displacement from the reference particle, respectively, and describe motion in the plane transverse to the velocity of the reference particle. The longitudinal coordinate t denotes the difference between the arrival time of the particle and the arrival time of the reference particle, multiplied by the speed of light \(c\).
The momenta conjugate to x, y, and t are denoted px, py, and pt, respectively. These variables are normalized by the magnitude of the momentum of the reference particle, and are therefore dimensionless. In a region of zero vector potential, for example, \(p_x = \Delta(\beta_x\gamma)/(\beta_0\gamma_0)\), where \(\beta_0\) and \(\gamma_0\) denote the relativistic factors associated with the reference velocity. In a region of zero scalar potential, pt denotes the deviation from the reference energy normalized by the design momentum times the speed of light, so that \(p_t = \Delta(\gamma)/(\beta_0\gamma_0)\).
Unlike particles within the beam, the reference particle is described by a set of 8 phase space variables (x [m], px, y [m], py, z [m], pz, t [m], pt) that are specified in a global laboratory coordinate system (x,y,z). The momenta of the reference particle are normalized by \(mc\), so that \(p_x=\beta_x\gamma\), etc. A parameteric plot of the reference trajectory variables (x,z) allows the user to view the global geometry of the accelerator structure (footprint).
Beam Distribution Input
Phase space ellipse in the coordinate plane of position \(q\) (realized as \(x\), \(y\), and \(t\)) and associated conjugate momentum \(q\) (realized as \(p_x\), \(p_y\) and \(p_t\)).
Particle beam user input in ImpactX can be done in two ways.
The first option is to characterize the distribution via the intersections of the phase space ellipse with the coordinate axes and the correlation terms of the canonical coordinate pairs.
The units are \([\lambda_q] = \mathrm{m}\), \([\lambda_p] = \mathrm{rad}\), and \([\mu_{qp}] = 1\). To convert between normalized and unnormalized emittance, use the relation \(\epsilon_\mathrm{n} = (\beta\gamma)_\mathrm{ref} \cdot \epsilon\) which uses the momentum of the reference particle. Attention: Here, \((\beta\gamma)_\mathrm{ref}\) are the Lorentz variables for the reference particle momentum and not the Courant-Snyder parameters.
The second option is to specify the distribution via the Courant-Snyder / Twiss parameters \(\alpha\) and \(\beta\) along with the unnormalized (geometric, 1-RMS) emittance \(\epsilon\) for all the spatial coordinates. Recall the Courant-Snyder relation \(\gamma\beta - \alpha^2 = 1\) for conversion from \(\gamma\) values to our input conventions.
Distribution Sampling and the Covariance Matrix
In ImpactX, beam sampling is performed under the assumption that the initial beam distribution centroid (mean phase space vector) coincides with the phase space origin. The covariance matrix \(\Sigma\) is defined by \(\Sigma_{ij}=\langle{\zeta_i\zeta_j\rangle}\), where \(\zeta\) denotes the vector of phase space coordinates, and indices \(i,j\) specify the components of \(\zeta\).
Let \(P\) denote a phase space probability density with unit covariance matrix (i.e., equal to the identity matrix). To produce a phase space density with a target covariance matrix \(\Sigma\), we write \(\Sigma\) in terms of its (lower) Cholesky decomposition as:
where \(L\) is a lower triangular matrix.
Define a beam distribution function \(f\) by:
Then \(f\) has the desired covariance matrix \(\Sigma\). Samples from \(f\) are obtained by sampling from \(P\) and performing the linear transformation \(\zeta\mapsto L\zeta\).
Let \(P\) above denote a 2D probability distribution that is radially symmetric, in the sense that:
Here \(q\) denotes a position coordinate (e.g., \(x\), \(y\), or \(t\)) and \(p\) denotes the corresponding conjugate momentum.
Then the resulting distribution \(f\) has 2D elliptical symmetry, in the sense that:
The argument of \(G\) is a quadratic form in \((q,p)\), and it is convenient to express this quadratic form as:
Here \(\alpha\), \(\beta\), and \(\gamma\) denote the Courant-Snyder Twiss functions, and \(\epsilon\) denotes the rms (unnormalized) emittance.
The associated covariance matrix may be written explicitly in terms of the above parameters as:
Note: In the special case that \(\mu_{qp}=0\), we have \(\lambda_q=\sigma_q\) and \(\lambda_p=\sigma_p\), where \(\sigma_q=\langle{q^2\rangle}^{1/2}\) and \(\sigma_p=\langle{p^2\rangle}^{1/2}\).
Assumptions
This is an overview of physical assumptions implemented in the numerics of ImpactX.
Tracking and Lattice Optics
Tracking through lattice optics in ImpactX is performed by updating the canonical phase space variables (x,px,y,py,t,pt) using symplectic transport. The elements supported currently fall into one of the following categories:
zero-length (thin) elements, such as multipole kicks and coordinate transformations
ideal (thick) elements using a hard-edge fringe field approximation, such as drifts, quadrupoles, and dipoles
soft-edge elements described by \(s\)-dependent, user-provided field data, such as RF cavities
ML surrogate models using a trained neural network (not necessarily symplectic)
Transport may be performed using one of three possible levels of approximation to the underlying Hamiltonian:
linear transfer map (default): obtained by expanding the Hamiltonian through terms of degree 2 in the deviation of phase space variables from those of the reference particle
chromatic or paraxial approximation: obtained by expanding the Hamiltonian through terms of degree 2 in the transverse phase space variables, while retaining the nonlinear dependence on the energy variable pt
exact Hamiltonian: obtained using the exact nonlinear Hamiltonian
Space Charge (Poisson Solver)
velocity spread: when solving for space-charge effects, we assume that the relative spread of velocities of particles in the beam is negligible compared to the velocity of the reference particle, so that in the bunch frame (rest frame of the reference particle) particle velocities are nonrelativistic
electrostatic in the bunch frame: we assume there are no retardation effects and we solve the Poisson equation in the bunch frame
Development
Contribute to ImpactX
We welcome new contributors! Here is how to participate to the ImpactX development.
Git workflow
The ImpactX project uses git for version control. If you are new to git, you can follow one of these tutorials:
Configure your GitHub Account & Development Machine
First, let’s setup your Git environment and GitHub account.
Go to https://github.com/settings/profile and add your real name and affiliation
Go to https://github.com/settings/emails and add & verify the professional e-mails you want to be associated with.
Configure
git
on the machine you develop on to use the same spelling of your name and email:git config --global user.name "FIRSTNAME LASTNAME"
git config --global user.email EMAIL@EXAMPLE.com
Go to https://github.com/settings/keys and add the SSH public key of the machine you develop on. (Check out the GitHub guide to generating SSH keys or troubleshoot common SSH problems. )
Make your own fork
First, fork the ImpactX “mainline” repo on GitHub by pressing the Fork button on the top right of the page. A fork is a copy of ImpactX on GitHub, which is under your full control.
Then, we create local copies, for development:
# Clone the mainline ImpactX source code to your local computer.
# You cannot write to this repository, but you can read from it.
git clone git@github.com:ECP-WarpX/impactx.git
cd impactx
# rename what we just cloned: call it "mainline"
git remote rename origin mainline
# Add your own fork. You can get this address on your fork's Github page.
# Here is where you will publish new developments, so that they can be
# reviewed and integrated into "mainline" later on.
# "myGithubUsername" needs to be replaced with your user name on GitHub.
git remote add myGithubUsername git@github.com:myGithubUsername/impactx.git
Now you are free to play with your fork (for additional information, you can visit the Github fork help page).
Note
We only need to do the above steps for the first time.
Let’s Develop
You are all set! Now, the basic ImpactX development workflow is:
Implement your changes and push them on a new branch
branch_name
on your fork.Create a Pull Request from branch
branch_name
on your fork to branchdevelopment
on the main ImpactX repo.
Create a branch branch_name
(the branch name should reflect the piece of code you want to add, like fix-spectral-solver
) with
# start from an up-to-date development branch
git checkout development
git pull mainline development
# create a fresh branch
git checkout -b branch_name
and do the coding you want.
It is probably a good time to look at the AMReX documentation and at the Doxygen reference pages:
ImpactX Doxygen: https://impactx.readthedocs.io/en/latest/_static/doxyhtml
AMReX Doxygen: https://amrex-codes.github.io/amrex/doxygen
WarpX Doxygen: https://warpx.readthedocs.io/en/latest/_static/doxyhtml
Once you are done developing, add the files you created and/or modified to the git
staging area with
git add <file_I_created> <and_file_I_modified>
Build your changes
If you changed C++ files, then now is a good time to test those changes by compiling ImpactX locally. Follow the developer instructions in our manual to set up a local development environment, then compile and run ImpactX.
Commit & push your changes
Periodically commit your changes with
git commit
The commit message (between quotation marks) is super important in order to follow the developments during code-review and identify bugs. A typical format is:
This is a short, 40-character title
After a newline, you can write arbitray paragraphs. You
usually limit the lines to 70 characters, but if you don't, then
nothing bad will happen.
The most important part is really that you find a descriptive title
and add an empty newline after it.
For the moment, commits are on your local repo only. You can push them to your fork with
git push -u myGithubUsername branch_name
If you want to synchronize your branch with the development
branch (this is useful when the development
branch is being modified while you are working on branch_name
), you can use
git pull mainline development
and fix any conflict that may occur.
Submit a Pull Request
A Pull Request (PR) is the way to efficiently visualize the changes you made and to propose your new feature/improvement/fix to the ImpactX project.
Right after you push changes, a banner should appear on the Github page of your fork, with your branch_name
.
Click on the
compare & pull request
button to prepare your PR.It is time to communicate your changes: write a title and a description for your PR. People who review your PR are happy to know
what feature/fix you propose, and why
how you made it (added new/edited files, created a new class than inherits from…)
how you tested it and what was the output you got
and anything else relevant to your PR (attach images and scripts, link papers, etc.)
Press
Create pull request
. Now you can navigate through your PR, which highlights the changes you made.
Please DO NOT write large pull requests, as they are very difficult and time-consuming to review. As much as possible, split them into small, targeted PRs. For example, if find typos in the documentation open a pull request that only fixes typos. If you want to fix a bug, make a small pull request that only fixes a bug.
If you want to implement a feature and are not too sure how to split it, just open an issue about your plans and ping other ImpactX developers on it to chime in. Generally, write helper functionality first, test it and then write implementation code. Submit tests, documentation changes and implementation of a feature together for pull request review.
Even before your work is ready to merge, it can be convenient to create a PR (so you can use Github tools to visualize your changes).
In this case, please put the [WIP]
tag (for Work-In-Progress) at the beginning of the PR title.
You can also use the GitHub project tab in your fork to organize the work into separate tasks/PRs and share it with the ImpactX community to get feedback.
Include a test to your PR
A new feature is great, a working new feature is even better! Please test your code and add your test to the automated test suite. It’s the way to protect your work from adventurous developers.
Note
TOOD: Write a workflow how to add a test.
Include documentation about your PR
Now, let users know about your new feature by describing its usage in the ImpactX documentation.
Our documentation uses Sphinx, and it is located in docs/source/
.
Note
TODO: For instance, if you introduce a new runtime parameter in the input file, you can add it to Docs/source/running_cpp/parameters.rst.
If Sphinx is installed on your computer, you should be able to generate the html documentation with
make html
in docs/
. Then open docs/build/html/index.html
with your favorite web browser and look
for your changes.
Once your code is ready with documentation and automated test, congratulations!
You can create the PR (or remove the [WIP]
tag if you already created it).
Reviewers will interact with you if they have comments/questions.
Style and conventions
For indentation, ImpactX uses four spaces (no tabs)
Some text editors automatically modify the files you open. We recommend to turn on to remove trailing spaces and replace Tabs with 4 spaces.
The number of characters per line should be <100
Exception: in documentation files (
.rst
/.md
) use one sentence per line independent of its number of characters, which will allow easier edits.Space before and after assignment operator (
=
)To define a function , for e.g.,
myfunction()
use a space between the name of the function and the paranthesis -myfunction ()
. To call the function, the space is not required, i.e., just usemyfunction()
.The reason this is beneficial is that when we do a
git grep
to search formyfunction ()
, we can clearly see the locations wheremyfunction ()
is defined and wheremyfunction()
is called.Also, using
git grep "myfunction ()"
searches for files only in the git repo, which is more efficient compared to thegrep "myfunction ()"
command that searches through all the files in a directory, including plotfiles for example.It is recommended that style changes are not included in the PR where new code is added. This is to avoid any errors that may be introduced in a PR just to do style change.
ImpactX uses
CamelCase
convention for file names and class names, rather thansnake_case
.The names of all member variables should be prefixed with
m_
. This is particularly useful to avoid capturing member variables by value in a lambda function, which causes the whole object to be copied to GPU when running on a GPU-accelerated architecture. This convention should be used for all new piece of code, and it should be applied progressively to old code.#include
directives in C++ have a distinct order to avoid bugs, see the ImpactX repo structure for detailsFor all new code, we should avoid relying on
using namespace amrex;
and all amrex types should be prefixed withamrex::
. Inside limited scopes, AMReX type literals can be included withusing namespace amrex::literals;
. Ideally, old code should be modified accordingly.
Testing
Preparation
Prepare for running tests of ImpactX by building ImpactX from source.
In order to run our tests, you need to have a few Python packages installed:
python3 -m pip install -U pip
python3 -m pip install -U build packaging setuptools wheel pytest
python3 -m pip install -r examples/requirements.txt
Further Options
help:
ctest --test-dir build --help
list all tests:
ctest --test-dir build -N
only run tests that have “FODO” in their name:
ctest --test-dir build -R FODO
Documentation
Doxygen documentation
WarpX uses a Doxygen documentation. Whenever you create a new class, please document it where it is declared (typically in the header file):
/** A brief title
*
* few-line description explaining the purpose of my_class.
*
* If you are kind enough, also quickly explain how things in my_class work.
* (typically a few more lines)
*/
class my_class
{ ... }
Doxygen reads this docstring, so please be accurate with the syntax! See Doxygen manual for more information. Similarly, please document functions when you declare them (typically in a header file) like:
/** A brief title
*
* few-line description explaining the purpose of my_function.
*
* \param[in,out] my_int a pointer to an integer variable on which
* my_function will operate.
* \return what is the meaning and value range of the returned value
*/
int my_class::my_function(int* my_int);
An online version of this documentation is linked here.
Breathe documentation
Your Doxygen documentation is not only useful for people looking into the code, it is also part of the ImpactX online documentation based on Sphinx!
This is done using the Python module Breathe, that allows you to read Doxygen documentation dorectly in the source and include it in your Sphinx documentation, by calling Breathe functions.
For instance, the following line will get the Doxygen documentation for ImpactXParticleContainer
in src/particles/ImpactXParticleContainer.H
and include it to the html page generated by Sphinx:
-
class ImpactXParticleContainer : public amrex::ParticleContainer_impl<RealSoA::nattribs, IntSoA::nattribs>
Beam Particles in ImpactX
This class stores particles, distributed over MPI ranks.
Building the documentation
To build the documentation on your local computer, you will need to install Doxygen as well as the Python module breathe
.
First, change into docs/
and install the Python requirements:
cd docs/
pip install -U -r requirements.txt
You will also need Doxygen (macOS: brew install doxygen
; Ubuntu: sudo apt install doxygen
).
Then, to compile the documentation, use
make html
# This will first compile the Doxygen documentation (execute doxygen)
# and then build html pages from rst files using sphinx and breathe.
Open the created build/html/index.html
file with your favorite browser.
Rebuild and refresh as needed.
ImpactX Structure
Repo Organization
All the ImpactX source code is located in src/
.
All sub-directories have a pretty straightforward name.
Here is a visual representation of the repository structure.
Code organization
The main ImpactX class is ImpactX, implemented in src/ImpactX.cpp
.
Build System
ImpactX uses the CMake build system generator.
Each sub-folder contains a file CMakeLists.txt
with the names of the source files (.cpp
) that are added to the build.
Do not list header files (.H
) here.
C++ Includes
All ImpactX header files need to be specified relative to the src/
directory.
e.g.
#include "Utils/ImpactXConst.H"
files in the same directory as the including header-file can be included with
#include "FileName.H"
By default, in a MyName.cpp
source file we do not include headers already included in MyName.H
. Besides this exception, if a function or a class
is used in a source file, the header file containing its declaration must be included, unless the inclusion of a facade header is more appropriate. This is
sometimes the case for AMReX headers. For instance AMReX_GpuLaunch.H
is a façade header for AMReX_GpuLaunchFunctsC.H
and AMReX_GpuLaunchFunctsG.H
, which
contain respectively the CPU and the GPU implemetation of some methods, and which should not be included directly.
Whenever possible, forward declarations headers are included instead of the actual headers, in order to save compilation time (see dedicated section below). In ImpactX forward
declaration headers have the suffix *_fwd.H
, while in AMReX they have the suffix *Fwd.H
.
The include order (see PR #874 and PR #1947) and proper quotation marks are:
In a MyName.cpp
file:
#include "MyName.H"
(its header) then(further) ImpactX header files
#include "..."
thenImpactX forward declaration header files
#include "..._fwd.H"
AMReX header files
#include <...>
thenAMReX forward declaration header files
#include <...Fwd.H>
thenPICSAR header files
#include <...>
thenother third party includes
#include <...>
thenstandard library includes, e.g.
#include <vector>
In a MyName.H
file:
#include "MyName_fwd.H"
(the corresponding forward declaration header, if it exists) thenImpactX header files
#include "..."
thenImpactX forward declaration header files
#include "..._fwd.H"
AMReX header files
#include <...>
thenAMReX forward declaration header files
#include <...Fwd.H>
thenPICSAR header files
#include <...>
thenother third party includes
#include <...>
thenstandard library includes, e.g.
#include <vector>
Each of these groups of header files should ideally be sorted alphabetically, and a blank line should be placed between the groups.
For details why this is needed, please see PR #874, PR #1947, the LLVM guidelines, and include-what-you-use.
Forward Declaration Headers
Forward declarations can be used when a header file needs only to know that a given class exists, without any further detail (e.g., when only a pointer to an instance of
that class is used). Forward declaration headers are a convenient way to organize forward declarations. If a forward declaration is needed for a given class MyClass
, declared in MyClass.H
,
the forward declaration should appear in a header file named MyClass_fwd.H
, placed in the same folder containing MyClass.H
. As for regular header files, forward declaration headers must have
include guards. Below we provide a simple example:
MyClass_fwd.H
:
#ifndef MY_CLASS_FWD_H
#define MY_CLASS_FWD_H
class MyClass;
#endif //MY_CLASS_FWD_H
MyClass.H
:
#ifndef MY_CLASS_H
#define MY_CLASS_H
#include "MyClass_fwd.H"
#include "someHeader.H"
class MyClass{/* stuff */};
#endif //MY_CLASS_H
MyClass.cpp
:
#include "MyClass.H"
class MyClass{/* stuff */};
Usage: in SimpleUsage.H
#include "MyClass_fwd.H"
#include <memory>
/* stuff */
std::unique_ptr<MyClass> p_my_class;
/* stuff */
Implementation Details
Note
TODO :-)
C++ Objects & Functions
We generate the documentation of C++ objects and functions from our C++ source code by adding Doxygen strings.
Our Doxygen C++ API documentation in classic formatting is located here.
Python interface
Note
TODO :-)
Debugging the code
Sometimes, the code does not give you the result that you are expecting. This can be due to a variety of reasons, from misunderstandings or changes in the input parameters, system specific quirks, or bugs. You might also want to debug your code as you implement new features in ImpactX during development.
This section gives a step-by-step guidance on how to systematically check what might be going wrong.
Debugging Workflow
Try the following steps to debug a simulation:
Check the output text file, usually called
output.txt
: are there warnings or errors present?On an HPC system, look for the job output and error files, usually called
ImpactX.e...
andImpactX.o...
. Read long messages from the top and follow potential guidance.If your simulation already created output data files: Check if they look reasonable before the problem occurred; are the initial conditions of the simulation as you expected? Do you spot numerical artifacts or instabilities that could point to missing resolution or unexpected/incompatible numerical parameters?
Did the job output files indicate a crash? Check the
Backtrace.<mpirank>
files for the location of the code that triggered the crash. Backtraces are read from bottom (high-level) to top (most specific line that crashed).In case of a crash, Backtraces can be more detailed if you re-compile with debug flags: for example, try compiling with
-DCMAKE_BUILD_TYPE=RelWithDebInfo
(some slowdown) or even-DCMAKE_BUILD_TYPE=Debug
(this will make the simulation way slower) and rerun.If debug builds are too costly, try instead compiling with
-DAMReX_ASSERTIONS=ON
to activate more checks and rerun.If the problem looks like a memory violation, this could be from an invalid field or particle index access. Try compiling with
-DAMReX_BOUND_CHECK=ON
(this will make the simulation very slow), and rerun.If the problem looks like a random memory might be used, try initializing memory with signaling Not-a-Number (NaN) values through the runtime option
fab.init_snan = 1
. Further useful runtime options areamrex.fpe_trap_invalid
,amrex.fpe_trap_zero
andamrex.fpe_trap_overflow
(see details in the AMReX link below).On Nvidia GPUs, if you suspect the problem might be a race condition due to a missing host / device synchronization, set the environment variable
export CUDA_LAUNCH_BLOCKING=1
and rerun.Consider simplifying your input options and re-adding more options after having found a working baseline.
Fore more information, see also the AMReX Debugging Manual.
Last but not least: the community of ImpactX developers and users can help if you get stuck. Collect your above findings, describe where and what you are running and how you installed the code, describe the issue you are seeing with details and input files used and what you already tried. Can you reproduce the problem with a smaller setup (less parallelism and/or less resolution)? Report these details in a ImpactX GitHub issue.
Debuggers
See the AMReX debugger section on additional runtime parameters to
disable backtraces
rethrow exceptions
avoid AMReX-level signal handling
You will need to set those runtime options to work directly with debuggers.
Maintenance
Dependencies & Releases
Update ImpactX’ Core Dependencies
ImpactX has direct dependencies on AMReX and WarpX, which we periodically update.
The following scripts automate this workflow, in case one needs a newer commit of AMReX or WarpX between releases:
Note
./Tools/Release/updateAMReX.py
./Tools/Release/updateWarpX.py
Create a new ImpactX release
ImpactX has one release per month.
The version number is set at the beginning of the month and follows the format YY.MM
.
In order to create a GitHub release, you need to:
Create a new branch from
development
and update the version number in all source files. We usually wait for the AMReX release to be tagged first, then we also point to its tag.There is a script for updating core dependencies of ImpactX and the ImpactX version:
./Tools/Release/updateAMReX.py ./Tools/Release/updateWarpX.py ./Tools/Release/newVersion.shFor a ImpactX release, ideally a git tag of AMReX & WarpX shall be used instead of an unnamed commit.
Then open a PR, wait for tests to pass and then merge.
Local Commit (Optional): at the moment,
@ax3l
is managing releases and signs tags (naming:YY.MM
) locally with his GPG key before uploading them to GitHub.Publish: On the GitHub Release page, create a new release via
Draft a new release
. Either select the locally created tag or create one online (naming:YY.MM
) on the merged commit of the PR from step 1.In the release description, please specify the compatible versions of dependencies (see previous releases), and provide info on the content of the release. In order to get a list of PRs merged since last release, you may run
git log <last-release-tag>.. --format='- %s'Optional/future: create a
release-<version>
branch, write a changelog, and backport bug-fixes for a few days.
Epilogue
Glossary
In daily communication, we tend to abbreviate a lot of terms. It is important to us to make it easy to interact with the ImpactX community and thus, this list shall help to clarify often used terms.
Please see: https://warpx.readthedocs.io/en/latest/glossary.html
Funding and Acknowledgements
This work was supported by the Laboratory Directed Research and Development Program of Lawrence Berkeley National Laboratory under U.S. Department of Energy Contract No. DE-AC02-05CH11231.
ImpactX is supported by the CAMPA collaboration, a project of the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research and Office of High Energy Physics, Scientific Discovery through Advanced Computing (SciDAC) program.
We acknowledge all the contributors and users of the ImpactX community who participate to the code quality with valuable code improvement and important feedback.