diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index c7e89a338..282b8e588 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -28,7 +28,7 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v5 + uses: actions/checkout@v6 - uses: actions/setup-python@v6 name: Install Python diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 84792b4e0..a209014ab 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -30,14 +30,14 @@ jobs: docker-images: true swap-storage: false - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - name: install dependencies run: | .github/workflows/dependencies/nvcc11-openmpi.sh - name: CCache Cache - uses: actions/cache@v4 + uses: actions/cache@v5 # - once stored under a key, they become immutable (even if local cache path content changes) # - for a refresh the key has to change, e.g., hash of a tracked file in the key with: @@ -64,4 +64,4 @@ jobs: -DImpactX_PRECISION=SINGLE \ -DAMReX_CUDA_ERROR_CROSS_EXECUTION_SPACE_CALL=ON \ -DAMReX_CUDA_ERROR_CAPTURE_THIS=ON - cmake --build build -j 4 + cmake --build build -j 3 diff --git a/.github/workflows/hip.yml b/.github/workflows/hip.yml index 73403ef24..64a2ff02a 100644 --- a/.github/workflows/hip.yml +++ b/.github/workflows/hip.yml @@ -15,14 +15,14 @@ jobs: CMAKE_GENERATOR: Ninja if: github.event.pull_request.draft == false steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - name: install dependencies shell: bash run: .github/workflows/dependencies/hip-openmpi.sh 6.3.2 - name: CCache Cache - uses: actions/cache@v4 + uses: actions/cache@v5 # - once stored under a key, they become immutable (even if local cache path content changes) # - for a refresh the key has to change, e.g., hash of a tracked file in the key with: diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml index 8f69f0f86..b134b14bd 100644 --- a/.github/workflows/macos.yml +++ b/.github/workflows/macos.yml @@ -14,7 +14,7 @@ jobs: env: HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK: TRUE steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - uses: actions/setup-python@v6 name: Install Python with: @@ -46,7 +46,7 @@ jobs: python3 -m pip install --upgrade pipx python3 -m pipx install openPMD-validator - name: CCache Cache - uses: actions/cache@v4 + uses: actions/cache@v5 # - once stored under a key, they become immutable (even if local cache path content changes) # - for a refresh the key has to change, e.g., hash of a tracked file in the key with: @@ -81,7 +81,7 @@ jobs: - name: Upload dashboard screenshots if: always() - uses: actions/upload-artifact@v5 + uses: actions/upload-artifact@v6 with: name: dashboard-screenshots-macos path: | diff --git a/.github/workflows/source.yml b/.github/workflows/source.yml index a8a0d425e..faea1554c 100644 --- a/.github/workflows/source.yml +++ b/.github/workflows/source.yml @@ -19,7 +19,7 @@ jobs: runs-on: ubuntu-24.04 steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - name: Non-ASCII characters run: .github/workflows/source/hasNonASCII diff --git a/.github/workflows/stubs.yml b/.github/workflows/stubs.yml index dc13f4e63..520fca1a8 100644 --- a/.github/workflows/stubs.yml +++ b/.github/workflows/stubs.yml @@ -18,10 +18,10 @@ jobs: OMP_NUM_THREAD: 4 steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 if: github.event_name != 'push' || github.repository != 'BLAST-ImpactX/impactx' || github.ref != 'refs/heads/development' - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 if: github.event_name == 'push' && github.repository == 'BLAST-ImpactX/impactx' && github.ref == 'refs/heads/development' with: token: ${{ secrets.IMPACTX_PUSH_TOKEN }} @@ -41,7 +41,7 @@ jobs: python3 -m pip install -U pybind11-stubgen pre-commit - name: Set Up Cache - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: ~/.cache/ccache key: ccache-${{ github.workflow }}-${{ github.job }}-git-${{ github.sha }} diff --git a/.github/workflows/tooling.yml b/.github/workflows/tooling.yml index 24aa85f2e..edbc88061 100644 --- a/.github/workflows/tooling.yml +++ b/.github/workflows/tooling.yml @@ -12,14 +12,14 @@ jobs: runs-on: ubuntu-22.04 if: github.event.pull_request.draft == false steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - name: install dependencies run: | .github/workflows/dependencies/clang-san-openmpi.sh - name: CCache Cache - uses: actions/cache@v4 + uses: actions/cache@v5 # - once stored under a key, they become immutable (even if local cache path content changes) # - for a refresh the key has to change, e.g., hash of a tracked file in the key with: diff --git a/.github/workflows/ubuntu.yml b/.github/workflows/ubuntu.yml index 129199cd7..d7489bfdd 100644 --- a/.github/workflows/ubuntu.yml +++ b/.github/workflows/ubuntu.yml @@ -27,14 +27,14 @@ jobs: docker-images: true swap-storage: false - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - name: install dependencies run: | .github/workflows/dependencies/gcc-openmpi.sh - name: CCache Cache - uses: actions/cache@v4 + uses: actions/cache@v5 # - once stored under a key, they become immutable (even if local cache path content changes) # - for a refresh the key has to change, e.g., hash of a tracked file in the key with: @@ -71,7 +71,7 @@ jobs: - name: Upload dashboard screenshots if: always() - uses: actions/upload-artifact@v5 + uses: actions/upload-artifact@v6 with: name: dashboard-screenshots-ubuntu path: | @@ -121,14 +121,14 @@ jobs: docker-images: true swap-storage: false - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - name: install dependencies run: | .github/workflows/dependencies/gcc.sh - name: CCache Cache - uses: actions/cache@v4 + uses: actions/cache@v5 # - once stored under a key, they become immutable (even if local cache path content changes) # - for a refresh the key has to change, e.g., hash of a tracked file in the key with: diff --git a/.github/workflows/windows.yml b/.github/workflows/windows.yml index 6a298ea71..dc2328830 100644 --- a/.github/workflows/windows.yml +++ b/.github/workflows/windows.yml @@ -12,13 +12,13 @@ jobs: runs-on: windows-latest if: github.event.pull_request.draft == false steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - uses: actions/setup-python@v6 name: Install Python with: python-version: '3.12' - name: CCache Cache - uses: actions/cache@v4 + uses: actions/cache@v5 # - once stored under a key, they become immutable (even if local cache path content changes) # - for a refresh the key has to change, e.g., hash of a tracked file in the key with: @@ -130,13 +130,13 @@ jobs: runs-on: windows-2025 if: github.event.pull_request.draft == false steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - uses: actions/setup-python@v6 name: Install Python with: python-version: '3.11' - name: CCache Cache - uses: actions/cache@v4 + uses: actions/cache@v5 # - once stored under a key, they become immutable (even if local cache path content changes) # - for a refresh the key has to change, e.g., hash of a tracked file in the key with: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 108128261..045794549 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -66,7 +66,7 @@ repos: # Python: Ruff linter & formatter # https://docs.astral.sh/ruff/ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.4 + rev: v0.14.9 hooks: # Run the linter - id: ruff-check diff --git a/CMakeLists.txt b/CMakeLists.txt index ba6d6b41c..3f19ca963 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,7 +1,7 @@ # Preamble #################################################################### # cmake_minimum_required(VERSION 3.24.0) -project(ImpactX VERSION 25.11) +project(ImpactX VERSION 25.12) include(${ImpactX_SOURCE_DIR}/cmake/ImpactXFunctions.cmake) diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index 4a65f26f4..1ae36f8fd 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -139,7 +139,7 @@ macro(find_ablastr) set(COMPONENT_DIM 3D) set(COMPONENT_PRECISION ${ImpactX_PRECISION} P${ImpactX_PRECISION}) - find_package(ABLASTR 25.11 CONFIG REQUIRED COMPONENTS ${COMPONENT_DIM}) + find_package(ABLASTR 25.12 CONFIG REQUIRED COMPONENTS ${COMPONENT_DIM}) message(STATUS "ABLASTR: Found version '${ABLASTR_VERSION}'") endif() @@ -178,7 +178,7 @@ set(ImpactX_openpmd_src "" set(ImpactX_ablastr_repo "https://github.com/BLAST-WarpX/warpx.git" CACHE STRING "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)") -set(ImpactX_ablastr_branch "25.11" +set(ImpactX_ablastr_branch "25.12" CACHE STRING "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)") @@ -186,7 +186,7 @@ set(ImpactX_ablastr_branch "25.11" set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git" CACHE STRING "Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)") -set(ImpactX_amrex_branch "" +set(ImpactX_amrex_branch "4dad1664d7467ae00c133c1425c1a300003fa885" CACHE STRING "Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)") diff --git a/cmake/dependencies/pyAMReX.cmake b/cmake/dependencies/pyAMReX.cmake index 0e8f5e170..cd4b3a099 100644 --- a/cmake/dependencies/pyAMReX.cmake +++ b/cmake/dependencies/pyAMReX.cmake @@ -59,7 +59,7 @@ function(find_pyamrex) endif() elseif(NOT ImpactX_pyamrex_internal) # TODO: MPI control - find_package(pyAMReX 25.11 CONFIG REQUIRED) + find_package(pyAMReX 25.12 CONFIG REQUIRED) message(STATUS "pyAMReX: Found version '${pyAMReX_VERSION}'") endif() endfunction() @@ -74,7 +74,7 @@ option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON) set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git" CACHE STRING "Repository URI to pull and build pyamrex from if(ImpactX_pyamrex_internal)") -set(ImpactX_pyamrex_branch "25.11" +set(ImpactX_pyamrex_branch "74e213e3186371d5a7412af3d2a5d6040fab51ce" CACHE STRING "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)") diff --git a/docs/README.md b/docs/README.md index 6594a790f..510ab6e28 100644 --- a/docs/README.md +++ b/docs/README.md @@ -6,10 +6,13 @@ This explains how to generate the documentation for ImpactX, and contribute to i ### Installing the requirements -Install the Python requirements for compiling the documentation: -``` -pip install sphinx sphinx_rtd_theme -``` +`cd` into the ImpactX `docs/` directory. +Install the Python requirements for compiling the documentation with **one** of the following package managers: + +- pip: `pip install -U -r requirements.txt` +- uv: `uv pip install -U -r requirements.txt` +- conda: `conda create -n readthedocs -f conda.yml && conda activate readthedocs` +- spack: `spack env activate -d . && spack install` ### Compiling the documentation diff --git a/docs/source/conf.py b/docs/source/conf.py index 673402a69..f31e04bdb 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -96,9 +96,9 @@ def download_with_headers(url, filename): # built documents. # # The short X.Y version. -version = "25.11" +version = "25.12" # The full version, including alpha/beta/rc tags. -release = "25.11" +release = "25.12" # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. diff --git a/docs/source/install/dependencies.rst b/docs/source/install/dependencies.rst index 46bb1a2ba..0c28973ab 100644 --- a/docs/source/install/dependencies.rst +++ b/docs/source/install/dependencies.rst @@ -32,11 +32,12 @@ Optional dependencies include: - `Ninja `__: for faster parallel compiles - `Python 3.9+ `__ + - `matplotlib 3.3+ `__ - `mpi4py `__ - `numpy `__ + - `openPMD-api `__ - `pals-schema `__ - `quantiphy `__ - - `openPMD-api `__ - see our ``requirements.txt`` file for compatible versions - web browser/Jupyter Dashboard: `trame `__ diff --git a/docs/source/install/hpc/perlmutter-nersc/perlmutter_gpu_impactx.profile.example b/docs/source/install/hpc/perlmutter-nersc/perlmutter_gpu_impactx.profile.example index b64f68587..56b813842 100644 --- a/docs/source/install/hpc/perlmutter-nersc/perlmutter_gpu_impactx.profile.example +++ b/docs/source/install/hpc/perlmutter-nersc/perlmutter_gpu_impactx.profile.example @@ -37,10 +37,10 @@ fi # an alias to request an interactive batch node for one hour # for parallel execution, start on the batch node: srun -alias getNode="salloc -N 1 --ntasks-per-node=4 -t 1:00:00 -q interactive -C gpu --gpu-bind=single:1 -c 32 -G 4 -A $proj" +alias getNode="salloc -N 1 --ntasks-per-node=4 -t 1:00:00 -q interactive -C gpu --gpu-bind=none -c 32 -G 4 -A $proj" # an alias to run a command on a batch node for up to 30min # usage: runNode -alias runNode="srun -N 1 --ntasks-per-node=4 -t 0:30:00 -q interactive -C gpu --gpu-bind=single:1 -c 32 -G 4 -A $proj" +alias runNode="srun -N 1 --ntasks-per-node=4 -t 0:30:00 -q interactive -C gpu --gpu-bind=none -c 32 -G 4 -A $proj" # necessary to use CUDA-Aware MPI and run a job export CRAY_ACCEL_TARGET=nvidia80 diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 64d7529be..4d4bc3609 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -30,6 +30,7 @@ Single Particle Dynamics examples/kicker/README.rst examples/thin_dipole/README.rst examples/aperture/README.rst + examples/polygon_aperture/README.rst examples/iota_lens/README.rst examples/achromatic_spectrometer/README.rst examples/fodo_userdef/README.rst diff --git a/docs/source/usage/howto/add_element.rst b/docs/source/usage/howto/add_element.rst index f32b274a5..2f594d939 100644 --- a/docs/source/usage/howto/add_element.rst +++ b/docs/source/usage/howto/add_element.rst @@ -98,7 +98,7 @@ As a last step, we expose our C++ beamline elements to Python in `src/python/ele :language: cpp :dedent: 4 :start-at: py::class_ py_Drift(me, "Drift"); - :end-at: register_envelope_push(py_Drift); + :end-at: register_push(py_Drift); Pull requests that added a new element and can be taken as examples are: diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 37be058b6..14def9ae2 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -138,6 +138,16 @@ Initial Beam Distributions * ``beam.normalize_halo`` (``float``, dimensionless) normalizing constant for halo population * ``beam.halo`` (``float``, dimensionless) fraction of charge in halo +Initial Spin Distributions +-------------------------- + + The specification of an initial particle spin distribution is optional, and is required only if spin tracking is used. + The default distribution type is the von Mises-Fisher distribution, uniquely determined by the input polarization vector. + The polarization vector provided by the user must lie within the unit ball. + +* ``beam.polarization_x`` (``float``, dimensionless) mean value of the spin vector x-component +* ``beam.polarization_y`` (``float``, dimensionless) mean value of the spin vector y-component +* ``beam.polarization_z`` (``float``, dimensionless) mean value of the spin vector z-component .. _running-cpp-parameters-lattice: @@ -338,19 +348,21 @@ If ``model = "linear"``, then the linearized map is used. This model is identic when expanded to first order in ``g/rc`` (gap / radius of curvature). +By comparison, note that the MAD-X DIPEDGE element uses as input the half-gap ``HGAP = g/2``, and sets the default value ``FINT = 0`` (while the corresponding default value of ``K2`` is set to 1). + This requires these additional parameters: * ``.psi`` (``float``, in radians) the pole face rotation angle * ``.rc`` (``float``, in meters) the bend radius * ``.g`` (``float``, in meters) the full magnetic gap size * ``.R`` (``float``, in meters) scale length for the field integrals (default: ``1 m``) -* ``.K0`` (``float``, dimensionless) normalized field integral for fringe field -* ``.K1`` (``float``, dimensionless) normalized field integral for fringe field -* ``.K2`` (``float``, dimensionless) normalized field integral for fringe field (FINT) -* ``.K3`` (``float``, dimensionless) normalized field integral for fringe field -* ``.K4`` (``float``, dimensionless) normalized field integral for fringe field -* ``.K5`` (``float``, dimensionless) normalized field integral for fringe field -* ``.K6`` (``float``, dimensionless) normalized field integral for fringe field +* ``.K0`` (``float``, dimensionless) normalized field integral for fringe field (default: ``pi/6``) +* ``.K1`` (``float``, dimensionless) normalized field integral for fringe field (default: ``0``) +* ``.K2`` (``float``, dimensionless) normalized field integral for fringe field (FINT, default: ``1``) +* ``.K3`` (``float``, dimensionless) normalized field integral for fringe field (default: ``1/6``) +* ``.K4`` (``float``, dimensionless) normalized field integral for fringe field (default: ``0``) +* ``.K5`` (``float``, dimensionless) normalized field integral for fringe field (default: ``0``) +* ``.K6`` (``float``, dimensionless) normalized field integral for fringe field (default: ``0``) * ``.model`` (``string``) the fringe field model: ``linear`` (default) or ``nonlinear`` * ``.location`` (``string``) the fringe field edge location: ``entry`` (default) or ``exit`` * ``.dx`` (``float``, in meters) horizontal translation error @@ -552,6 +564,25 @@ This requires these additional parameters: * ``.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``) +``polygon_aperture`` +^^^^^^^^^^^^^^^^^^^^ + +``polygon_aperture`` for a thin collimator element applying a transverse polygon aperture boundary defined by :math:`(x,y)` coordinates +and optional radius below which all particles are transmitted. The vertices must define a closed curve and be ordered in the counter-clockwise direction. +The first and last vertices must be identical. These parameters define the element: + +* ``.vertices_x`` (``float``, in meters) array of horizontal locations of aperture vertices +* ``.vertices_y`` (``float``, in meters) array of vertical locations of aperture vertices +* ``.min_radius2`` (``float``, in meters-squared) optional minimum radius-squared of a circle fully inscribed within the polygon. Particles with + radius-squared less than this value are transmitted by the aperture and the polygon calculation is skipped. (default ``0``) +* ``.repeat_x`` (``float``, in meters) horizontal period for repeated aperture masking (inactive by default) +* ``.repeat_y`` (``float``, in meters) vertical period for repeated aperture masking (inactive by default) +* ``.shift_odd_x`` (``bool``) for hexagonal/triangular mask patterns: horizontal shift of every 2nd (odd) vertical period by repeat_x / 2. Use alignment offsets dx,dy to move whole mask as needed. +* ``.action`` (``string``) action of the aperture domain: ``transmit`` (default) or ``absorb`` +* ``.dx`` (``float``, in meters) horizontal translation error +* ``.dy`` (``float``, in meters) vertical translation error +* ``.rotation`` (``float``, in degrees) rotation error in the transverse plane + ``prot`` ^^^^^^^^ @@ -810,6 +841,8 @@ Currently, this only supports openPMD files from our ``beam_monitor``. Distribution type of particles in the source. currently, only ``"openPMD"`` is supported * ``.openpmd_path`` (``string``) path to the openPMD series +* ``.active_once`` (``boolean``, default: ``true``) + Inject particles only for the first lattice period. ``tapered_pl`` @@ -903,6 +936,7 @@ See there ``nslice`` option on lattice elements for slicing. When running in envelope mode (when ``algo.track = "envelope"``), this model currently assumes that `` = = = 0``. * ``"Gauss3D"``: Calculate 3D space charge forces as if the beam was a Gaussian distribution. + * ``"Gauss2p5D"``: Calculate 2.5D space charge forces as if the beam was a transverse Gaussian distribution. These models are supported only in particle tracking mode (when ``algo.track = "particles"``). diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 7a741c977..36ac6bb63 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -233,7 +233,7 @@ Collective Effects & Overall Simulation Parameters This must come first, before particle beams and lattice elements are initialized. - .. py:method:: add_particles(charge_C, distr, npart) + .. py:method:: add_particles(charge_C, distr, npart, spinv=None) Particle tracking mode: Generate and add n particles to the particle container. Note: Set the reference particle properties (charge, mass, energy) first. @@ -243,6 +243,7 @@ Collective Effects & Overall Simulation Parameters :param float charge_C: bunch charge (C) :param distr: distribution function to draw from (object from :py:mod:`impactx.distribution`) :param int npart: number of particles to draw + :param SpinvMF spinv: optional spin distribution .. py:method:: init_envelope(ref, distr, intensity=None) @@ -584,8 +585,8 @@ Particles :param madx_file: file name to MAD-X file with a ``BEAM`` entry -Initial Beam Distributions --------------------------- +Initial Beam Phase Space Distributions +-------------------------------------- This module provides particle beam distributions that can be used to initialize particle beams in an :py:class:`impactx.ParticleContainer`. @@ -665,6 +666,25 @@ For the input from Twiss parameters in Python, please use the helper function `` A 6D stationary thermal or bithermal distribution. +Initial Beam Spin Distribution +------------------------------ + +.. py:class:: impactx.distribution.SpinvMF(mux, muy, muz) + + A von Mises-Fisher (vMF) distribution on the unit 2-sphere. + + This is used for initializing particle spin. There is a natural bijective correspondence between vMF distributions and mean (polarization) vectors. + + The algorithm used here is a simplification of the algorithm described in: + C. Pinzon and K. Jung, "Fast Python sampler of the von Mises Fisher distribution", in the special case of the 2-sphere. Additional references used include: + + - K. V. Mardia and P. E. Jupp, Directional Statistics, Wiley, 1999; + - S. Kang and H-S. Oh, "Novel sampling method for the von Mises-Fisher distribution", Stat. and Comput. 34, 106 (2024), `DOI:10.1007/s11222-024-10419-3 `__ + + :param mux: x component of the unit vector specifying the mean direction + :param muy: y component of the unit vector specifying the mean direction + :param muz: z component of the unit vector specifying the mean direction + Lattice Elements ---------------- @@ -815,7 +835,7 @@ This module provides elements and methods for the accelerator lattice. focusing t strength in 1/m -.. py:class:: impactx.elements.DipEdge(psi, rc, g, K2, dx=0, dy=0, rotation=0, name=None) +.. py:class:: impactx.elements.DipEdge(psi, rc, g, R=1, K0=pi/6, K1=0, K2=1, K3=1/6, K4=0, K5=0, K6=0, model="linear", location="entry", dx=0, dy=0, rotation=0, name=None) Edge focusing associated with bend entry or exit @@ -838,6 +858,9 @@ This module provides elements and methods for the accelerator lattice. when expanded to first order in ``g/rc`` (gap / radius of curvature). + By comparison, note that the MAD-X DIPEDGE element uses as input the half-gap ``HGAP = g/2``, and sets the default value ``FINT = 0`` (while + the corresponding default value of ``K2`` is set to 1). + :param psi: Pole face angle [radians] :param rc: Radius of curvature [m] :param g: Gap parameter [m] @@ -1099,6 +1122,7 @@ This module provides elements and methods for the accelerator lattice. :param distribution: Distribution type of particles in the source. currently, only ``"openPMD"`` is supported :param openpmd_path: path to the openPMD series + :param active_once: Inject particles only for the first lattice period. Default: ``True`` :param name: an optional name for the element .. py:class:: impactx.elements.Programmable(ds=0.0, nslice=1, name=None) @@ -1468,6 +1492,30 @@ This module provides elements and methods for the accelerator lattice. maximum vertical coordinate +.. py:class:: impactx.elements.PolygonAperture(vertices_x, vertices_y, min_radius2=0.0, repeat_x, repeat_y, shift_odd_x, action="transmit", dx=0, dy=0, rotation=0, name=None) + + This element defines a thin collimator element applying a transverse polygon aperture boundary defined by :math:`(x,y)` coordinates + and optional radius below which all particles are transmitted. The vertices must define a closed curve and be ordered in the counter-clockwise direction. + The first and last vertices must be identical. + + :param vertices_x: sequence of aperture boundary :math:`x` coordinates in m + :param vertices_y: sequence of aperture boundary :math:`y` coordinates in m + :param min_radius2: radius-squared of a circle fully inscribed by the polygon aperture (default 0) (meters-squared) + :param repeat_x: horizontal period for repeated aperture masking (inactive by default) (meter) + :param repeat_y: vertical period for repeated aperture masking (inactive by default) (meter) + :param shift_odd_x: for hexagonal/triangular mask patterns: horizontal shift of every 2nd (odd) vertical period by repeat_x / 2. Use alignment offsets dx,dy to move whole mask as needed. + :param action: aperture domain action: ``"transmit"`` (default) or ``"absorb"`` + :param dx: horizontal translation error in m + :param dy: vertical translation error in m + :param rotation: rotation error in the transverse plane [degrees] + :param name: an optional name for the element + + .. py:property:: min_radius2 + + radius-squared of a fully inscribed circle. Particles with radius-squared less than this value are transmitted by the aperture and the polygon calculation is skipped. + + aperture type (transmit, absorb) + .. py:class:: impactx.elements.SoftQuadrupole(ds, gscale, cos_coefficients, sin_coefficients, dx=0, dy=0, rotation=0, aperture_x=0, aperture_y=0, mapsteps=1, nslice=1, name=None) A soft-edge quadrupole. diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index f1cd0f7e8..388104738 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1057,6 +1057,44 @@ add_impactx_test(absorber.py OFF # no plot script yet ) +# Polygon Aperture ####################################################################### +# +add_impactx_test(polygon_aperture + examples/polygon_aperture/input_polygon_aperture.in + ON # ImpactX MPI-parallel + examples/polygon_aperture/analysis_polygon_aperture.py + examples/polygon_aperture/plot_polygon_aperture.py +) +add_impactx_test(polygon_aperture.py + examples/polygon_aperture/run_polygon_aperture.py + ON # ImpactX MPI-parallel + examples/polygon_aperture/analysis_polygon_aperture.py + examples/polygon_aperture/plot_polygon_aperture.py +) +add_impactx_test(polygon_aperture_absorb.py + examples/polygon_aperture/run_polygon_aperture_absorb.py + ON # ImpactX MPI-parallel + OFF + OFF +) +add_impactx_test(polygon_aperture_absorb_offset.py + examples/polygon_aperture/run_polygon_aperture_absorb_offset.py + ON # ImpactX MPI-parallel + OFF + OFF +) +add_impactx_test(polygon_aperture_absorb_rotate.py + examples/polygon_aperture/run_polygon_aperture_absorb_rotate.py + ON # ImpactX MPI-parallel + OFF + OFF +) +add_impactx_test(polygon_aperture_absorb_rotate_offset.py + examples/polygon_aperture/run_polygon_aperture_absorb_rotate_offset.py + ON # ImpactX MPI-parallel + OFF + OFF +) # Apochromat drift-quad example ########################################################## # @@ -1659,50 +1697,71 @@ add_impactx_test(fodo-pals.py # Active Plasma Lens tests #################################### # -# ChrPlasmaLens, k = 0 (acting like a drift) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_zero.py) - -add_impactx_test(APL_ChrPlasmaLens_zero.py - examples/active_plasma_lens/run_APL_ChrPlasmaLens_zero.py +# ChrPlasmaLens/tracking, k = 0 (acting like a drift) +add_impactx_test(APL_ChrPlasmaLens_tracking_zero.py + examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_zero.py OFF # No MPI needed - examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_zero.py - examples/active_plasma_lens/plot_APL_ChrPlasmaLens_zero.py + examples/active_plasma_lens/analysis_APL_tracking_zero.py + examples/active_plasma_lens/plot_APL_zero.py ) - -# ChrPlasmaLens, mg = -1000T/m (focusing) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_focusing.py) - -add_impactx_test(APL_ChrPlasmaLens_focusing.py - examples/active_plasma_lens/run_APL_ChrPlasmaLens_focusing.py +# ConstF/tracking, k = 0 (acting like a drift) +add_impactx_test(APL_ConstF_tracking_zero.py + examples/active_plasma_lens/run_APL_ConstF_tracking_zero.py + OFF # No MPI needed + examples/active_plasma_lens/analysis_APL_tracking_zero.py + examples/active_plasma_lens/plot_APL_zero.py +) +# ConstF/envelope, k = 0 (acting like a drift) +add_impactx_test(APL_ConstF_envelope_zero.py + examples/active_plasma_lens/run_APL_ConstF_envelope_zero.py OFF # No MPI needed - examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_focusing.py - examples/active_plasma_lens/plot_APL_ChrPlasmaLens_focusing.py + examples/active_plasma_lens/analysis_APL_envelope_zero.py + examples/active_plasma_lens/plot_APL_zero.py ) -# ChrPlasmaLens, mg = 1000T/m (defocusing) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/run_APL.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py) -file(COPY ${ImpactX_SOURCE_DIR}/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py - DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/APL_ChrPlasmaLens_defocusing.py) +# ChrPlasmaLens/tracking, mg = -1000T/m (focusing) +add_impactx_test(APL_ChrPlasmaLens_tracking_focusing.py + examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_focusing.py + OFF # No MPI needed + examples/active_plasma_lens/analysis_APL_tracking_focusing.py + examples/active_plasma_lens/plot_APL_focusing.py +) +# ConstF/tracking, mg = -1000T/m (focusing) +add_impactx_test(APL_ConstF_tracking_focusing.py + examples/active_plasma_lens/run_APL_ConstF_tracking_focusing.py + OFF # No MPI needed + examples/active_plasma_lens/analysis_APL_tracking_focusing.py + examples/active_plasma_lens/plot_APL_focusing.py +) +# ConstF/envelope, mg = -1000T/m (focusing) +add_impactx_test(APL_ConstF_envelope_focusing.py + examples/active_plasma_lens/run_APL_ConstF_envelope_focusing.py + OFF # No MPI needed + examples/active_plasma_lens/analysis_APL_envelope_focusing.py + examples/active_plasma_lens/plot_APL_focusing.py +) -add_impactx_test(APL_ChrPlasmaLens_defocusing.py - examples/active_plasma_lens/run_APL_ChrPlasmaLens_defocusing.py +# ChrPlasmaLens/tracking, mg = 1000T/m (defocusing) +add_impactx_test(APL_ChrPlasmaLens_tracking_defocusing.py + examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_defocusing.py + OFF # No MPI needed + examples/active_plasma_lens/analysis_APL_tracking_defocusing.py + examples/active_plasma_lens/plot_APL_defocusing.py +) +# ConstF/tracking, mg = 1000T/m (defocusing) +add_impactx_test(APL_ConstF_tracking_defocusing.py + examples/active_plasma_lens/run_APL_ConstF_tracking_defocusing.py + OFF # No MPI needed + examples/active_plasma_lens/analysis_APL_tracking_defocusing.py + examples/active_plasma_lens/plot_APL_defocusing.py +) +# ConstF/envelope, mg = 1000T/m (defocusing) +add_impactx_test(APL_ConstF_envelope_defocusing.py + examples/active_plasma_lens/run_APL_ConstF_envelope_defocusing.py OFF # No MPI needed - examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_defocusing.py - examples/active_plasma_lens/plot_APL_ChrPlasmaLens_defocusing.py - ) + examples/active_plasma_lens/analysis_APL_envelope_defocusing.py + examples/active_plasma_lens/plot_APL_defocusing.py +) # Exactly-solvable (non-uniform) soft-edge solenoid ############## # @@ -1756,3 +1815,18 @@ add_impactx_test(dipedge-nonlinear.py examples/edge_effects/analysis_dipedge.py OFF # no plot script yet ) + +# Spin Sampling from a vMF Distribution ###################### +# without space charge +add_impactx_test(distgen-spinvmf + examples/distgen/input_kurth4d_spin.in + ON # ImpactX MPI-parallel + examples/distgen/analysis_kurth4d_spin.py + OFF # no plot script yet +) +add_impactx_test(distgen-spinvmf.py + examples/distgen/run_kurth4d_spin.py + ON # ImpactX MPI-parallel + examples/distgen/analysis_kurth4d_spin.py + OFF # no plot script yet +) diff --git a/examples/active_plasma_lens/APL_analytical_alpha_REF.png b/examples/active_plasma_lens/APL_analytical_alpha_REF.png new file mode 100644 index 000000000..ab165afa1 Binary files /dev/null and b/examples/active_plasma_lens/APL_analytical_alpha_REF.png differ diff --git a/examples/active_plasma_lens/APL_analytical_sqrtBeta_REF.png b/examples/active_plasma_lens/APL_analytical_sqrtBeta_REF.png new file mode 100644 index 000000000..5066e3e12 Binary files /dev/null and b/examples/active_plasma_lens/APL_analytical_sqrtBeta_REF.png differ diff --git a/examples/active_plasma_lens/APL_defocusing-sigma_REF.png b/examples/active_plasma_lens/APL_defocusing-sigma_REF.png new file mode 100644 index 000000000..0e389b5cd Binary files /dev/null and b/examples/active_plasma_lens/APL_defocusing-sigma_REF.png differ diff --git a/examples/active_plasma_lens/APL_focusing-sigma_REF.png b/examples/active_plasma_lens/APL_focusing-sigma_REF.png new file mode 100644 index 000000000..1866a0703 Binary files /dev/null and b/examples/active_plasma_lens/APL_focusing-sigma_REF.png differ diff --git a/examples/active_plasma_lens/APL_zero-sigma_REF.png b/examples/active_plasma_lens/APL_zero-sigma_REF.png new file mode 100644 index 000000000..274edb45e Binary files /dev/null and b/examples/active_plasma_lens/APL_zero-sigma_REF.png differ diff --git a/examples/active_plasma_lens/README.rst b/examples/active_plasma_lens/README.rst index 025d0f2de..605b5b8b9 100644 --- a/examples/active_plasma_lens/README.rst +++ b/examples/active_plasma_lens/README.rst @@ -6,40 +6,241 @@ Active Plasma Lens These examples demonstrate the effect of an Active Plasma Lens (APL) on the beam. The lattice contains this element and nothing else. The length of the element is 20 mm, and it can be run in no-field, focusing, and defocusing mode. +It is implemented with two different elements, ```ChrPlasmaLens``` and ```ConstF```; +for ```ConstF``` we test both particle tracking and envelope tracking. -We use a 200 MeV electron beam with an initial normalized rms emittance of 10 um. -The beam is set to have :math:`\alpha = 0` in the middle of the lens in the case of no field. -The beam size in the middle of the lens is set to 10 µm for the no-field examples (in order to have a strongly parabolic :math:`\beta`-function within the lens), and 100 µm for the focusing and defocusing examples. -A :math:`\sigma_{pt} = 10^{-3}` is also assumed. -Before the simulation, this beam is back-propagated to the lens entry assuming zero field. +We use a 200 MeV electron beam with an initial normalized rms emittance of 10 µm. +The initial Twiss parameters for the simulations are set by first assuming the beam +to have :math:`\alpha = 0` in the middle of the lens, and then the we backwards-propagate +this analytically to the start of the lens, under the assumption of no field. +The beam is then forward-propagated using ImpactX and the different cases for +element strength and simulation method. + +The beam size in the middle of the lens is set to 10 µm for the no-field examples +in order to have a strongly parabolic :math:`\beta`-function within the lens, +and 100 µm for the focusing and defocusing examples. +The :math:`\beta` and :math:`\gamma`-function in the middle of the lens is calculated +from this beam size and the assumed emittance. +A :math:`\sigma_{pt} = 10^{-3}` is also assumed for the tracking examples. Run --- -This example can be run as -* ``python3 run_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking) -* ``python3 run_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking) -* ``python3 run_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking) +This example can be run as Python scripts, with the names +indicating the element used (``ChrPlasmaLens`` or ``ConstF``), +the field in the lens (``zero``, ``focusing``, or ``defocusing``), +and simulation type (``tracking`` or ``envelope``). + +These both run the simulation, and produce analytical reference parameters +which are used for comparison by the analysis scripts. + +* ``python3 run_APL_ChrPlasmaLens_tracking_zero.py`` +* ``python3 run_APL_ChrPlasmaLens_tracking_focusing.py`` +* ``python3 run_APL_ChrPlasmaLens_tracking_defocusing.py`` +* ``python3 run_APL_ConstF_tracking_zero.py`` +* ``python3 run_APL_ConstF_tracking_focusing.py`` +* ``python3 run_APL_ConstF_tracking_defocusing.py`` +* ``python3 run_APL_ConstF_envelope_zero.py`` +* ``python3 run_APL_ConstF_envelope_focusing.py`` +* ``python3 run_APL_ConstF_envelope_defocusing.py`` + +These all use the library ``run_APL.py`` internally to setup and and run the simulations. + +.. dropdown:: ``run_APL.py`` + + .. literalinclude:: run_APL.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL.py``. + +The scripts used to start the simulations: + +.. dropdown:: ``run_APL_ChrPlasmaLens_tracking_zero.py`` + + .. literalinclude:: run_APL_ChrPlasmaLens_tracking_zero.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_zero.py``. + +.. dropdown:: ``run_APL_ChrPlasmaLens_tracking_focusing.py`` + + .. literalinclude:: run_APL_ChrPlasmaLens_tracking_focusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_focusing.py``. + +.. dropdown:: ``run_APL_ChrPlasmaLens_tracking_defocusing.py`` + + .. literalinclude:: run_APL_ChrPlasmaLens_tracking_defocusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_defocusing.py``. + +.. dropdown:: ``run_APL_ConstF_tracking_zero.py`` + + .. literalinclude:: run_APL_ConstF_tracking_zero.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ConstF_tracking_zero.py``. -These all use the library ``run_APL.py`` internally to create the simulations. +.. dropdown:: ``run_APL_ConstF_tracking_focusing.py`` + + .. literalinclude:: run_APL_ConstF_tracking_focusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ConstF_tracking_focusing.py``. + +.. dropdown:: ``run_APL_ConstF_tracking_defocusing.py`` + + .. literalinclude:: run_APL_ConstF_tracking_defocusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ConstF_tracking_defocusing.py``. + +.. dropdown:: ``run_APL_ConstF_envelope_zero.py`` + + .. literalinclude:: run_APL_ConstF_envelope_zero.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ConstF_envelope_zero.py``. + +.. dropdown:: ``run_APL_ConstF_envelope_focusing.py`` + + .. literalinclude:: run_APL_ConstF_envelope_focusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ConstF_envelope_focusing.py``. + +.. dropdown:: ``run_APL_ConstF_envelope_defocusing.py`` + + .. literalinclude:: run_APL_ConstF_envelope_defocusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/run_APL_ConstF_envelope_defocusing.py``. Analyze ------- -We run the following scripts to analyze correctness of the output: -* ``python3 analysis_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking) -* ``python3 analysis_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking) -* ``python3 analysis_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking) +The following scripts can be used to analyze correctness of the output, +by comparing it to a reference output that is produced and outputed to +the standard output (terminal) from the run scripts. + +The output should be the same accross elements (``ConstF`` or ``ChrPlasmaLens``), +but depend on the field in the lens (``zero``, ``focusing``, or ``defocusing``), +and simulation type (``tracking`` or ``envelope``). +The analysis scripts are therefore the same for both element types. + +All analysis scripts look at the output from most recent simulation run in +the current working directory, i.e. the ``diags`` folder. + +* ``python3 analysis_APL_tracking_zero.py`` +* ``python3 analysis_APL_tracking_focusing.py`` +* ``python3 analysis_APL_tracking_defocusing.py`` +* ``python3 analysis_APL_envelope_zero.py`` +* ``python3 analysis_APL_envelope_focusing.py`` +* ``python3 analysis_APL_envelope_defocusing.py`` + +These all use the library ``analysis_APL.py`` internally to load data etc: + +.. dropdown:: ``analysis_APL.py`` + + .. literalinclude:: analysis_APL.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/analysis_APL.py``. + +The analysis scripts are: + +.. dropdown:: ``analysis_APL_tracking_zero.py`` + + .. literalinclude:: analysis_APL_tracking_zero.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/analysis_APL_tracking_zero.py``. + +.. dropdown:: ``analysis_APL_tracking_focusing.py`` + + .. literalinclude:: analysis_APL_tracking_focusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/analysis_APL_tracking_focusing.py``. + +.. dropdown:: ``analysis_APL_tracking_defocusing.py`` + + .. literalinclude:: analysis_APL_tracking_defocusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/analysis_APL_tracking_defocusing.py``. + +.. dropdown:: ``analysis_APL_envelope_zero.py`` + + .. literalinclude:: analysis_APL_envelope_zero.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/analysis_APL_envelope_zero.py``. + +.. dropdown:: ``analysis_APL_envelope_focusing.py`` + + .. literalinclude:: analysis_APL_envelope_focusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/analysis_APL_envelope_focusing.py``. + +.. dropdown:: ``analysis_APL_envelope_defocusing.py`` + + .. literalinclude:: analysis_APL_envelope_defocusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/analysis_APL_envelope_defocusing.py``. -These all use the library ``analysis_APL_ChrPlasmaLens.py`` internally. Visualize --------- -You can run the following scripts to visualize the beam evolution over time (e.g. :math:`s`): -* ``python3 s_APL_ChrPlasmaLens_zero.py`` (no field, ``ChrPlasmaLens``, tracking)plot -* ``python3 plot_APL_ChrPlasmaLens_focusing.py`` (focusing field, ``ChrPlasmaLens``, tracking) -* ``python3 plot_APL_ChrPlasmaLens_defocusing.py`` (defocusing field, ``ChrPlasmaLens``, tracking) +You can run the following scripts to visualize the beam evolution over time (e.g. :math:`s`), +and compare to analytical expectation. +Here, For this, the output format is identical accross the element- and simulation-types, +only depending on the selected lens field (``zero``, ``focusing``, or ``defocusing``). + +* ``python3 plot_APL_zero.py`` +* ``python3 plot_APL_focusing.py`` +* ``python3 plot_APL_defocusing.py`` + +These all use the library ``plot_APL.py`` internally, +as well as ``run_APL.py`` and ``analysis_APL.py`` which are described above. + +.. dropdown:: ``plot_APL.py`` + + .. literalinclude:: plot_APL.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/plot_APL.py``. + +The plotting scripts are given below, together with their graphical output. +For the plots, output from the ``ChrPlasmaLens_tracking`` simulations were used, +which shows some error in the envelope due to statistical fluctuations in the initial particle distribution. + +.. dropdown:: ``plot_APL_zero.py`` and output figures: + + .. literalinclude:: plot_APL_zero.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/plot_APL_zero.py``. + + .. figure:: APL_zero-sigma_REF.png + :alt: The calculated :math:`\sigma` values throughout the lens without field, for both planes and a separate analytical estimate. + +.. dropdown:: ``plot_APL_focusing.py`` and output figures: + + .. literalinclude:: plot_APL_focusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/plot_APL_focusing.py``. + + .. figure:: APL_focusing-sigma_REF.png + :alt: The calculated :math:`\sigma` values throughout the lens with focusing field, for both planes and a separate analytical estimate. + +.. dropdown:: ``plot_APL_defocusing.py`` and output figures: + + .. literalinclude:: plot_APL_defocusing.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/plot_APL_defocusing.py``. + + .. figure:: APL_defocusing-sigma_REF.png + :alt: The calculated :math:`\sigma` values throughout the lens with defocusing field, for both planes and a separate analytical estimate. + +Additionally, it is also possible to run ``python3 plot_APL_analytical.py``, +which plots the expected Twiss :math:`\alpha` and :math:`\beta` functions at the end of the lens +as a function of the lens gradient. This uses the stand-alone Twiss propagation function +``analytic_final_estimate()`` from ``run_APL.py``. + +.. dropdown:: ``plot_APL_analytical.py`` and output figures: + + .. literalinclude:: plot_APL_analytical.py + :language: python3 + :caption: You can copy this file from ``examples/active_plasma_lens/plot_APL_analytical.py``. -These all use the library ``plot_APL_ChrPlasmaLens.py`` internally. + .. figure:: APL_analytical_sqrtBeta_REF.png + :alt: The analytically computed :math:`\sqrt{\beta}` value at the end of the lens (proportional to beam size), as a function of gradient. -Additionally, it is also possible to run ``python3 plot_APL_ChrPlasmaLens_analytical.py``, which plots the expected Twiss :math:`\alpha` and :math:`\beta` functions at the end of the lens as a function of the lens gradient. This uses the stand-alone Twiss propagation function ``analytic_final_estimate()`` from ``run_APL.py``. + .. figure:: APL_analytical_alpha_REF.png + :alt: The analytically computed :math:`\alpha` value at the end of the lens, as a function of gradient. diff --git a/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py b/examples/active_plasma_lens/analysis_APL.py similarity index 70% rename from examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py rename to examples/active_plasma_lens/analysis_APL.py index 5b98d33c1..7470bf64e 100644 --- a/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens.py +++ b/examples/active_plasma_lens/analysis_APL.py @@ -6,6 +6,8 @@ # # -*- coding: utf-8 -*- +import os + import openpmd_api as io from scipy.stats import moment @@ -55,10 +57,44 @@ def get_twiss(beam): def get_beams(): "Load the initial and final beam from last simulation" - series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) + + getFname = "diags/openPMD/monitor.h5" + print(f"** get_beams(): Loading {os.path.abspath(getFname)}") + + series = io.Series(getFname, io.Access.read_only) last_step = list(series.iterations)[-1] initial = series.iterations[1].particles["beam"].to_df() beam_final = series.iterations[last_step].particles["beam"] final = beam_final.to_df() return (initial, beam_final, final) + + +# Load data from envelope simulation +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 + """ + + import glob + import re + + import pandas as pd + + def read_file(file_pattern): + for filename in glob.glob(file_pattern): + df = pd.read_csv(filename, delimiter=r"\s+") + if "step" not in df.columns: + step = int(re.findall(r"[0-9]+", filename)[0]) + df["step"] = step + yield df + + return pd.concat( + read_file(file_pattern), + axis=0, + ignore_index=True, + ) # .set_index('id') diff --git a/examples/active_plasma_lens/analysis_APL_envelope_defocusing.py b/examples/active_plasma_lens/analysis_APL_envelope_defocusing.py new file mode 100755 index 000000000..d8f0e07ee --- /dev/null +++ b/examples/active_plasma_lens/analysis_APL_envelope_defocusing.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import numpy as np +from analysis_APL import read_time_series + +rbc = read_time_series("diags/reduced_beam_characteristics.*") + +print("Initial Beam:") + +sigx = rbc["sigma_x"].iloc[0] +sigy = rbc["sigma_y"].iloc[0] +sigt = rbc["sigma_t"].iloc[0] +emittance_x = rbc["emittance_x"].iloc[0] +emittance_y = rbc["emittance_y"].iloc[0] +emittance_t = rbc["emittance_t"].iloc[0] + +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}" +) + +betax = rbc["beta_x"].iloc[0] +betay = rbc["beta_y"].iloc[0] +alphax = rbc["alpha_x"].iloc[0] +alphay = rbc["alpha_y"].iloc[0] + +print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") + +atol = 0.0 # ignored +rtol = 1e-5 +print(f" rtol={rtol} (ignored: atol~={atol})") + +# Compare initial beam to analytical values +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 0.00010003246877770656, + 0.00010003246877770656, + 0.001, + 2.548491664266332e-08, + 2.548491664266332e-08, + 1e-06, + ], + rtol=rtol, + atol=atol, +) + +print("") +print("Final Beam:") + +sigx = rbc["sigma_x"].iloc[-1] +sigy = rbc["sigma_y"].iloc[-1] +sigt = rbc["sigma_t"].iloc[-1] +emittance_x = rbc["emittance_x"].iloc[-1] +emittance_y = rbc["emittance_y"].iloc[-1] +emittance_t = rbc["emittance_t"].iloc[-1] + +s_ref = rbc["s"].iloc[-1] + +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}\n" + f" s_ref={s_ref:e}" +) + +betax = rbc["beta_x"].iloc[-1] +betay = rbc["beta_y"].iloc[-1] +alphax = rbc["alpha_x"].iloc[-1] +alphay = rbc["alpha_y"].iloc[-1] +print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") + +atol = 0.0 # ignored +rtol = 1e-5 +print(f" rtol={rtol} (ignored: atol~={atol})") + +# Compare final beam to analytical values +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref], + [ + 0.0001314429025974998, + 0.0001314429025974998, + 0.001, + 2.548491664266332e-08, + 2.548491664266332e-08, + 1e-06, + 20e-3, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/active_plasma_lens/analysis_APL_envelope_focusing.py b/examples/active_plasma_lens/analysis_APL_envelope_focusing.py new file mode 100755 index 000000000..e22a3f93f --- /dev/null +++ b/examples/active_plasma_lens/analysis_APL_envelope_focusing.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import numpy as np +from analysis_APL import read_time_series + +rbc = read_time_series("diags/reduced_beam_characteristics.*") + +print("Initial Beam:") + +sigx = rbc["sigma_x"].iloc[0] +sigy = rbc["sigma_y"].iloc[0] +sigt = rbc["sigma_t"].iloc[0] +emittance_x = rbc["emittance_x"].iloc[0] +emittance_y = rbc["emittance_y"].iloc[0] +emittance_t = rbc["emittance_t"].iloc[0] + +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}" +) + +betax = rbc["beta_x"].iloc[0] +betay = rbc["beta_y"].iloc[0] +alphax = rbc["alpha_x"].iloc[0] +alphay = rbc["alpha_y"].iloc[0] + +print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") + +atol = 0.0 # ignored +rtol = 1e-5 +print(f" rtol={rtol} (ignored: atol~={atol})") + +# Compare initial beam to analytical values +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 0.00010003246877770656, + 0.00010003246877770656, + 0.001, + 2.548491664266332e-08, + 2.548491664266332e-08, + 1e-06, + ], + rtol=rtol, + atol=atol, +) + +print("") +print("Final Beam:") + +sigx = rbc["sigma_x"].iloc[-1] +sigy = rbc["sigma_y"].iloc[-1] +sigt = rbc["sigma_t"].iloc[-1] +emittance_x = rbc["emittance_x"].iloc[-1] +emittance_y = rbc["emittance_y"].iloc[-1] +emittance_t = rbc["emittance_t"].iloc[-1] + +s_ref = rbc["s"].iloc[-1] + +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}\n" + f" s_ref={s_ref:e}" +) + +betax = rbc["beta_x"].iloc[-1] +betay = rbc["beta_y"].iloc[-1] +alphax = rbc["alpha_x"].iloc[-1] +alphay = rbc["alpha_y"].iloc[-1] +print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") + +atol = 0.0 # ignored +rtol = 1e-5 +print(f" rtol={rtol} (ignored: atol~={atol})") + +# Compare final beam to analytical values +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref], + [ + 7.161196476484095e-05, + 7.161196476484095e-05, + 0.001, + 2.548491664266332e-08, + 2.548491664266332e-08, + 1e-06, + 20e-3, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/active_plasma_lens/analysis_APL_envelope_zero.py b/examples/active_plasma_lens/analysis_APL_envelope_zero.py new file mode 100755 index 000000000..3fd3095a7 --- /dev/null +++ b/examples/active_plasma_lens/analysis_APL_envelope_zero.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import numpy as np +from analysis_APL import read_time_series + +rbc = read_time_series("diags/reduced_beam_characteristics.*") + +print("Initial Beam:") + +sigx = rbc["sigma_x"].iloc[0] +sigy = rbc["sigma_y"].iloc[0] +sigt = rbc["sigma_t"].iloc[0] +emittance_x = rbc["emittance_x"].iloc[0] +emittance_y = rbc["emittance_y"].iloc[0] +emittance_t = rbc["emittance_t"].iloc[0] + +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}" +) + +betax = rbc["beta_x"].iloc[0] +betay = rbc["beta_y"].iloc[0] +alphax = rbc["alpha_x"].iloc[0] +alphay = rbc["alpha_y"].iloc[0] + +print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") + +atol = 0.0 # ignored +rtol = 1e-5 +print(f" rtol={rtol} (ignored: atol~={atol})") + +# Compare initial beam to analytical values +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 2.737665020201518e-05, + 2.737665020201518e-05, + 0.001, + 2.548491664266332e-08, + 2.548491664266332e-08, + 1e-06, + ], + rtol=rtol, + atol=atol, +) + +print("") +print("Final Beam:") + +sigx = rbc["sigma_x"].iloc[-1] +sigy = rbc["sigma_y"].iloc[-1] +sigt = rbc["sigma_t"].iloc[-1] +emittance_x = rbc["emittance_x"].iloc[-1] +emittance_y = rbc["emittance_y"].iloc[-1] +emittance_t = rbc["emittance_t"].iloc[-1] + +s_ref = rbc["s"].iloc[-1] + +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}\n" + f" s_ref={s_ref:e}" +) + +betax = rbc["beta_x"].iloc[-1] +betay = rbc["beta_y"].iloc[-1] +alphax = rbc["alpha_x"].iloc[-1] +alphay = rbc["alpha_y"].iloc[-1] +print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") + +atol = 0.0 # ignored +rtol = 1e-5 +print(f" rtol={rtol} (ignored: atol~={atol})") + +# Compare final beam to analytical values +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref], + [ + 2.737665020201518e-05, + 2.737665020201518e-05, + 0.001, + 2.548491664266332e-08, + 2.548491664266332e-08, + 1e-06, + 20e-3, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_defocusing.py b/examples/active_plasma_lens/analysis_APL_tracking_defocusing.py similarity index 83% rename from examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_defocusing.py rename to examples/active_plasma_lens/analysis_APL_tracking_defocusing.py index 4184b1c9a..2b802bfb5 100755 --- a/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_defocusing.py +++ b/examples/active_plasma_lens/analysis_APL_tracking_defocusing.py @@ -7,13 +7,13 @@ # -*- coding: utf-8 -*- import numpy as np -from analysis_APL_ChrPlasmaLens import get_beams, get_moments, get_twiss +from analysis_APL import get_beams, get_moments, get_twiss # initial/final beam (initial, beam_final, final) = get_beams() # compare number of particles -num_particles = 10000 +num_particles = 100000 assert num_particles == len(initial) assert num_particles == len(final) @@ -28,7 +28,7 @@ print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") atol = 0.0 # ignored -rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +rtol = 2.5 * num_particles**-0.5 # from random sampling of a smooth distribution print(f" rtol={rtol} (ignored: atol~={atol})") # Compare to analytical values @@ -62,10 +62,7 @@ print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") atol = 0.0 # ignored -# rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution -rtol = ( - 2.9 * num_particles**-0.5 -) # from random sampling of a smooth distribution -- tolerance increased here +rtol = 2.5 * num_particles**-0.5 # from random sampling of a smooth distribution print(f" rtol={rtol} (ignored: atol~={atol})") # Compare to analytical values @@ -75,8 +72,8 @@ 0.0001314429025974998, 0.0001314429025974998, 0.001, - 2.514662e-08, - 2.514662e-08, + 2.548491664266332e-08, + 2.548491664266332e-08, 1e-06, 20e-3, 3.923902e02, diff --git a/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_focusing.py b/examples/active_plasma_lens/analysis_APL_tracking_focusing.py similarity index 90% rename from examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_focusing.py rename to examples/active_plasma_lens/analysis_APL_tracking_focusing.py index 48550a0a7..e926d4f99 100755 --- a/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_focusing.py +++ b/examples/active_plasma_lens/analysis_APL_tracking_focusing.py @@ -7,13 +7,13 @@ # -*- coding: utf-8 -*- import numpy as np -from analysis_APL_ChrPlasmaLens import get_beams, get_moments, get_twiss +from analysis_APL import get_beams, get_moments, get_twiss # initial/final beam (initial, beam_final, final) = get_beams() # compare number of particles -num_particles = 10000 +num_particles = 100000 assert num_particles == len(initial) assert num_particles == len(final) @@ -28,7 +28,7 @@ print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") atol = 0.0 # ignored -rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +rtol = 2.5 * num_particles**-0.5 # from random sampling of a smooth distribution print(f" rtol={rtol} (ignored: atol~={atol})") # Compare to analytical values @@ -62,7 +62,7 @@ print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") atol = 0.0 # ignored -rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +rtol = 2.5 * num_particles**-0.5 # from random sampling of a smooth distribution print(f" rtol={rtol} (ignored: atol~={atol})") # Compare to analytical values diff --git a/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_zero.py b/examples/active_plasma_lens/analysis_APL_tracking_zero.py similarity index 87% rename from examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_zero.py rename to examples/active_plasma_lens/analysis_APL_tracking_zero.py index 645310989..61eabb2fe 100755 --- a/examples/active_plasma_lens/analysis_APL_ChrPlasmaLens_zero.py +++ b/examples/active_plasma_lens/analysis_APL_tracking_zero.py @@ -7,13 +7,13 @@ # -*- coding: utf-8 -*- import numpy as np -from analysis_APL_ChrPlasmaLens import get_beams, get_moments, get_twiss +from analysis_APL import get_beams, get_moments, get_twiss # initial/final beam (initial, beam_final, final) = get_beams() # compare number of particles -num_particles = 10000 +num_particles = 100000 assert num_particles == len(initial) assert num_particles == len(final) @@ -28,10 +28,10 @@ print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") atol = 0.0 # ignored -rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +rtol = 2.5 * num_particles**-0.5 # from random sampling of a smooth distribution print(f" rtol={rtol} (ignored: atol~={atol})") -# Compare to analytical values +# Compare initial beam to analytical values assert np.allclose( [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], [ @@ -46,7 +46,6 @@ atol=atol, ) - print("") print("Final Beam:") sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) @@ -62,10 +61,10 @@ print(f" betax={betax}[m],betay={betay}[m],alphax={alphax},alphay={alphay}") atol = 0.0 # ignored -rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +rtol = 2.5 * num_particles**-0.5 # from random sampling of a smooth distribution print(f" rtol={rtol} (ignored: atol~={atol})") -# Compare to analytical values +# Compare final beam to analytical values assert np.allclose( [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref, gamma_ref], [ diff --git a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py b/examples/active_plasma_lens/plot_APL.py similarity index 70% rename from examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py rename to examples/active_plasma_lens/plot_APL.py index 7ad3e5cd5..967a415f8 100644 --- a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens.py +++ b/examples/active_plasma_lens/plot_APL.py @@ -6,38 +6,10 @@ # # -*- coding: utf-8 -*- -import glob -import re import matplotlib.pyplot as plt -import pandas as pd from matplotlib.ticker import MaxNLocator - -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') - - # 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) diff --git a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens_analytical.py b/examples/active_plasma_lens/plot_APL_analytical.py similarity index 70% rename from examples/active_plasma_lens/plot_APL_ChrPlasmaLens_analytical.py rename to examples/active_plasma_lens/plot_APL_analytical.py index f07869a70..77ee976bc 100755 --- a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens_analytical.py +++ b/examples/active_plasma_lens/plot_APL_analytical.py @@ -9,10 +9,21 @@ ## Plots analytical Twiss parameters # as a function of APL gradient [T/m] +import argparse + import matplotlib.pyplot as plt import numpy as np from run_APL import analytic_final_estimate +# options to run this script, this one is used by the CTest harness +parser = argparse.ArgumentParser( + description="Plot the ChrPlasmaLens_focusing and ConstF_tracking_focusing benchmarks." +) +parser.add_argument( + "--save-png", action="store_true", help="non-interactive run: save to PNGs" +) +args = parser.parse_args() + rigidity_Tm = -0.6688305274603505 # [T*m], 200MeV electrons APL_length = 20e-3 # [m] @@ -37,14 +48,9 @@ plt.figure(1) plt.plot(g_def, beta) -plt.xlabel("g_APL [T/m]") -plt.ylabel("sqrt(beta)") plt.figure(2) plt.plot(g_def, alpha) -plt.xlabel("g_APL [T/m]") -plt.ylabel("alpha") - # positive g (defocusing, rigidity is negative) beta = [] @@ -60,13 +66,18 @@ plt.figure(1) plt.plot(g_def, beta) -plt.xlabel("g_APL [T/m]") -plt.ylabel("sqrt(beta)") +plt.xlabel(r"$g_{APL}$ [T/m]") +plt.ylabel(r"$\sqrt{\beta}$") plt.grid() +if args.save_png: + plt.savefig("APL_analytical_sqrtBeta.png") + plt.figure(2) plt.plot(g_def, alpha) -plt.xlabel("g_APL [T/m]") -plt.ylabel("alpha") +plt.xlabel("$g_{APL}$ [T/m]") +plt.ylabel(r"$\alpha$") plt.grid() - -plt.show() +if args.save_png: + plt.savefig("APL_analytical_alpha.png") +else: + plt.show() diff --git a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens_defocusing.py b/examples/active_plasma_lens/plot_APL_defocusing.py similarity index 56% rename from examples/active_plasma_lens/plot_APL_ChrPlasmaLens_defocusing.py rename to examples/active_plasma_lens/plot_APL_defocusing.py index 606ec861d..4142d62b1 100755 --- a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens_defocusing.py +++ b/examples/active_plasma_lens/plot_APL_defocusing.py @@ -8,11 +8,13 @@ import argparse -from plot_APL_ChrPlasmaLens import millimeter, plot_sigmas, plt, read_time_series +from analysis_APL import read_time_series +from plot_APL import millimeter, plot_sigmas, plt +from run_APL import analytic_sigma_function # options to run this script, this one is used by the CTest harness parser = argparse.ArgumentParser( - description="Plot the ChrPlasmaLens_focusing benchmark." + description="Plot the ChrPlasmaLens_focusing and ConstF_tracking_focusing benchmarks." ) parser.add_argument( "--save-png", action="store_true", help="non-interactive run: save to PNGs" @@ -27,11 +29,20 @@ # Plot beam transverse sizes plot_sigmas(rbc) +# Analytical estimates # Start/end plt.axhline(0.0001314429025974998 * millimeter, ls="--", color="k") plt.axhline(100e-6 * millimeter, ls="--", color="k") -plt.axvline(10e-3, ls="--", color="k") +# plt.axvline(10e-3, ls="--", color="k") +# As function of s +(s, sigma) = analytic_sigma_function(+1000, 100e-6) +plt.plot(s, sigma * millimeter, ls="--", color="green", label="Analytical") + +plt.legend(loc="center left") +plt.title(r"Defocusing e$^-$, 200 MeV, $g$ = 1000 [T/m]") +plt.tight_layout() + if args.save_png: - plt.savefig("APL_ChrPlasmaLens_defocusing-sigma.png") + plt.savefig("APL_defocusing-sigma.png") else: plt.show() diff --git a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens_focusing.py b/examples/active_plasma_lens/plot_APL_focusing.py similarity index 56% rename from examples/active_plasma_lens/plot_APL_ChrPlasmaLens_focusing.py rename to examples/active_plasma_lens/plot_APL_focusing.py index 10529d54f..970ffae36 100755 --- a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens_focusing.py +++ b/examples/active_plasma_lens/plot_APL_focusing.py @@ -8,11 +8,13 @@ import argparse -from plot_APL_ChrPlasmaLens import millimeter, plot_sigmas, plt, read_time_series +from analysis_APL import read_time_series +from plot_APL import millimeter, plot_sigmas, plt +from run_APL import analytic_sigma_function # options to run this script, this one is used by the CTest harness parser = argparse.ArgumentParser( - description="Plot the ChrPlasmaLens_focusing benchmark." + description="Plot the ChrPlasmaLens_focusing and ConstF_tracking_focusing benchmarks." ) parser.add_argument( "--save-png", action="store_true", help="non-interactive run: save to PNGs" @@ -27,11 +29,20 @@ # Plot beam transverse sizes plot_sigmas(rbc) +# Analytical estimates # Start/end plt.axhline(7.161196476484095e-05 * millimeter, ls="--", color="k") plt.axhline(100e-6 * millimeter, ls="--", color="k") -plt.axvline(10e-3, ls="--", color="k") +# plt.axvline(10e-3, ls="--", color="k") +# As function of s +(s, sigma) = analytic_sigma_function(-1000, 100e-6) +plt.plot(s, sigma * millimeter, ls="--", color="green", label="Analytical") + +plt.legend(loc="center left") +plt.title(r"Focusing e$^-$, 200 MeV, $g$ = -1000 [T/m]") +plt.tight_layout() + if args.save_png: - plt.savefig("APL_ChrPlasmaLens_focusing-sigma.png") + plt.savefig("APL_focusing-sigma.png") else: plt.show() diff --git a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens_zero.py b/examples/active_plasma_lens/plot_APL_zero.py similarity index 57% rename from examples/active_plasma_lens/plot_APL_ChrPlasmaLens_zero.py rename to examples/active_plasma_lens/plot_APL_zero.py index d6c9e16b4..ebcb07b15 100755 --- a/examples/active_plasma_lens/plot_APL_ChrPlasmaLens_zero.py +++ b/examples/active_plasma_lens/plot_APL_zero.py @@ -8,29 +8,40 @@ import argparse -from plot_APL_ChrPlasmaLens import millimeter, plot_sigmas, plt, read_time_series +from analysis_APL import read_time_series +from plot_APL import millimeter, plot_sigmas, plt +from run_APL import analytic_sigma_function # options to run this script, this one is used by the CTest harness -parser = argparse.ArgumentParser(description="Plot the ChrPlasmaLens_zero benchmark.") +parser = argparse.ArgumentParser( + description="Plot the ChrPlasmaLens_zero and ConstF_tracking_zero benchmarks." +) parser.add_argument( "--save-png", action="store_true", help="non-interactive run: save to PNGs" ) args = parser.parse_args() -# import matplotlib.pyplot as plt - # read reduced diagnostics rbc = read_time_series("diags/reduced_beam_characteristics.*") # Plot beam transverse sizes plot_sigmas(rbc) +# Analytical estimates # Start/end plt.axhline(2.737665020201518e-05 * millimeter, ls="--", color="k") # mid plt.axhline(10e-6 * millimeter, ls="--", color="k") plt.axvline(10e-3, ls="--", color="k") +# As function of s +(s, sigma) = analytic_sigma_function(0.0, 10e-6) +plt.plot(s, sigma * millimeter, ls="--", color="green", label="Analytical") + +plt.legend(loc="center") +plt.title(r"No-field e$^-$, 200 MeV, $g$ = 0 [T/m]") +plt.tight_layout() + if args.save_png: - plt.savefig("APL_ChrPlasmaLens_zero-sigma.png") + plt.savefig("APL_zero-sigma.png") else: plt.show() diff --git a/examples/active_plasma_lens/run_APL.py b/examples/active_plasma_lens/run_APL.py index a7f76f2a3..82fb9c5c8 100644 --- a/examples/active_plasma_lens/run_APL.py +++ b/examples/active_plasma_lens/run_APL.py @@ -9,7 +9,11 @@ import numpy as np -from impactx import ImpactX, distribution, elements +# Physical reference parameters +APL_length = 20e-3 # [m] +kin_energy_MeV = 200 # [MeV] reference energy +mass_MeV = 0.510998950 # [MeV] +bunch_charge_C = 1.0e-9 # used with space charge def run_APL_tracking( @@ -17,9 +21,18 @@ def run_APL_tracking( ): """ Run a plasma lens tracking simulation with the given APL gradient APL_g [T/m], sigma_pt [-], and sigma_mid [m]. - Can use lensType='ChrPlasmaLens' | 'ConstK' | 'ChrDrift' (expect APL_g = 0) | 'ChrQuad' (only horizontal plane valid) + Can use lensType='ChrPlasmaLens' | 'ConstF' | 'ChrDrift' (expects APL_g = 0) | 'ChrQuad' (only horizontal plane valid) """ + from impactx import ImpactX, distribution, elements + + print("", flush=True) + print( + f"*** run_APL_tracking({APL_g}, {sigpt_0}, {sigma_mid}, '{lensType}') :", + flush=True, + ) + print("", flush=True) + sim = ImpactX() # set numerical parameters and IO control @@ -31,65 +44,133 @@ def run_APL_tracking( sim.init_grids() ## Physics parameters for test (APL_g from input arguments) - APL_length = 20e-3 # [m] # Load a 200 MeV electron beam with alpha=0 (x and y) # in the center of the APL - kin_energy_MeV = 200 # reference energy - bunch_charge_C = 1.0e-9 # 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(kin_energy_MeV) + ref.set_charge_qe(-1.0).set_mass_MeV(mass_MeV).set_kin_energy_MeV(kin_energy_MeV) - # Midpoint parameters - alpha_mid = 0.0 - # sigma_mid = 10e-6 # [m] - emitn = 10e-6 # [m] - emitg = emitn / ref.beta_gamma - beta_mid = sigma_mid**2 / emitg - gamma_mid = 1 / beta_mid # [1/m] - print( - f"sigma_mid = {sigma_mid} [m], beta_mid = {beta_mid} [m], gamma_mid = {gamma_mid} [m], alpha_mid = {alpha_mid}" - ) - print( - f"emitn = {emitn} [m], emitg = {emitg} [m], ref.beta_gamma = {ref.beta_gamma}, ref.rigidity_Tm = {ref.rigidity_Tm} [T*m]" + (emitg, beta_0, gamma_0, mu_0, alpha_0) = do_analytic_backprop( + sigma_mid, APL_g, ref.beta_gamma, ref.rigidity_Tm ) - print() - # Back-propagate 1/2 lens length as in vacuum, - # from symmetry point in the middle of the lens to the start of the lens - assert alpha_mid == 0.0 - beta_0 = beta_mid + (APL_length / 2) ** 2 / beta_mid - alpha_0 = +APL_length / 2 / beta_mid - gamma_0 = gamma_mid - sigma_0 = math.sqrt(emitg * beta_0) - sigmap_0 = math.sqrt(emitg * gamma_0) - mu_0 = alpha_0 / math.sqrt(beta_0 * gamma_0) + # Longitudinal parameters (sigpt_0 [-] from input arguments) + sigt_0 = 1e-3 # [m] + emit_t = math.sqrt(sigt_0**2 * sigpt_0**2 - 0**2) print( - f"sigma_0 = {sigma_0} [m], beta_0 = {beta_0} [m], alpha_0 = {alpha_0}, sigmap_0 = {sigmap_0}" + f"sigt_0 = {sigt_0} [m], sigpt_0 = {sigpt_0} [-], emit_t = {emit_t}", flush=True ) - print() + print(flush=True) - # Forward-propagate that through the focusing/defocusing lens - # from the beginning, ignoring energy spread - # [Doesn't give same result as ChrPlasmaLens and ChrQuad for some reason, even when sigpt_0 = 0] - (beta_end, alpha_end, gamma_end) = analytic_final_estimate( - APL_g, ref.rigidity_Tm, APL_length, beta_0, alpha_0 + # particle bunch + distr = distribution.Gaussian( + lambdaX=math.sqrt(emitg / gamma_0), + lambdaY=math.sqrt(emitg / gamma_0), + lambdaT=sigt_0, # OK for mutpt=0 + lambdaPx=math.sqrt(emitg / beta_0), + lambdaPy=math.sqrt(emitg / beta_0), + lambdaPt=sigpt_0, # OK for mutpt=0 + muxpx=mu_0, + muypy=mu_0, + mutpt=0.0, ) + npart = 100000 # number of macro particles + sim.add_particles(bunch_charge_C, distr, npart) + + # create the accelerator lattice + + # Plasma lens parameters for ConstF + APL_k = APL_g / ref.rigidity_Tm + APL_k_sqrt = np.sign(APL_k) * np.sqrt(np.abs(APL_k)) + print(f"APL_g = {APL_g} [T/m], APL_k = {APL_k} [1/m^2]", flush=True) + print(flush=True) + + ns = 40 # number of slices per ds in the element + monitor = elements.BeamMonitor("monitor", backend="h5") + APL = None + if lensType == "ChrPlasmaLens": + APL = elements.ChrPlasmaLens( + name="APL", ds=APL_length, k=APL_g, unit=1, nslice=ns + ) + + elif lensType == "ConstF": + APL = elements.ConstF( + name="APL", ds=APL_length, kx=APL_k_sqrt, ky=APL_k_sqrt, kt=0.0, nslice=ns + ) + + elif lensType == "ChrDrift": + # For comparison with k=0 + assert float(APL_g) == 0.0 + APL = elements.ChrDrift(name="APL", ds=APL_length, nslice=ns) + + elif lensType == "ChrQuad": + # For comparison with k != 0, single plane + APL = elements.ChrQuad(name="APL", ds=APL_length, k=APL_g, unit=1, nslice=ns) + else: + raise ValueError(f"Unknown lensType {lensType}") + + lattice = [ + monitor, + APL, + monitor, + ] + + # run simulation + sim.lattice.extend(lattice) + sim.track_particles() + + # clean shutdown + sim.finalize() + + +def run_APL_envelope( + APL_g: float, sigpt_0: float, sigma_mid: float, lensType: str = "ChrPlasmaLens" +): + """ + Run a plasma lens envelope simulation with the given APL gradient APL_g [T/m], sigma_pt [-], and sigma_mid [m]. + Can specify lensType='ChrPlasmaLens' | 'ConstF' | 'ChrDrift' (expects APL_g = 0) | 'ChrQuad' (only horizontal plane valid) + Note that for lensType='ChrPlasmaLens' | 'ChrDrift' | 'ChrQuad', the envelope simulation is not yet implemented. + """ + from impactx import ImpactX, distribution, elements + + print("", flush=True) print( - f"beta_end = {beta_end} [m], alpha_end = {alpha_end} [-], gamma_end = {gamma_end} [1/m]" + f"*** run_APL_envelope({APL_g}, {sigpt_0}, {sigma_mid}, '{lensType}') :", + flush=True, + ) + print("", flush=True) + + sim = ImpactX() + + # set numerical parameters and IO control + sim.space_charge = False + # sim.diagnostics = False # benchmarking + sim.slice_step_diagnostics = True + + # domain decomposition & space charge mesh + sim.init_grids() + + ## Physics parameters for test (APL_g from input arguments) + + # Load a 200 MeV electron beam with alpha=0 (x and y) + # in the center of the APL + # reference particle + ref = sim.particle_container().ref_particle() + ref.set_charge_qe(-1.0).set_mass_MeV(mass_MeV).set_kin_energy_MeV(kin_energy_MeV) + + (emitg, beta_0, gamma_0, mu_0, alpha_0) = do_analytic_backprop( + sigma_mid, APL_g, ref.beta_gamma, ref.rigidity_Tm ) - sigma_end = np.sqrt(emitg * beta_end) - sigmap_end = np.sqrt(emitg * gamma_end) - print(f"sigma_end = {sigma_end} [m], sigmap_end = {sigmap_end} [-]") - # print() # Longitudinal parameters (sigpt_0 [-] from input arguments) sigt_0 = 1e-3 # [m] emit_t = math.sqrt(sigt_0**2 * sigpt_0**2 - 0**2) - print(f"sigt_0 = {sigt_0} [m], sigpt_0 = {sigpt_0} [-], emit_t = {emit_t}") - print() + print( + f"sigt_0 = {sigt_0} [m], sigpt_0 = {sigpt_0} [-], emit_t = {emit_t}", flush=True + ) + print(flush=True) # particle bunch distr = distribution.Gaussian( @@ -103,15 +184,16 @@ def run_APL_tracking( muypy=mu_0, mutpt=0.0, ) - npart = 10000 # number of macro particles - sim.add_particles(bunch_charge_C, distr, npart) + # npart = 100000 # number of macro particles + # sim.add_particles(bunch_charge_C, distr, npart) # create the accelerator lattice # Plasma lens parameters for ConstF APL_k = APL_g / ref.rigidity_Tm APL_k_sqrt = np.sign(APL_k) * np.sqrt(np.abs(APL_k)) - print(f"APL_g = {APL_g} [T/m], APL_k = {APL_k} [1/m^2]") + print(f"APL_g = {APL_g} [T/m], APL_k = {APL_k} [1/m^2]", flush=True) + print(flush=True) ns = 40 # number of slices per ds in the element monitor = elements.BeamMonitor("monitor", backend="h5") @@ -121,8 +203,8 @@ def run_APL_tracking( name="APL", ds=APL_length, k=APL_g, unit=1, nslice=ns ) - elif lensType == "ConstK": - APL = elements.ConstK( + elif lensType == "ConstF": + APL = elements.ConstF( name="APL", ds=APL_length, kx=APL_k_sqrt, ky=APL_k_sqrt, kt=0.0, nslice=ns ) @@ -142,21 +224,109 @@ def run_APL_tracking( APL, monitor, ] - # assign a fodo segment - sim.lattice.extend(lattice) # run simulation - sim.track_particles() + sim.lattice.extend(lattice) + sim.init_envelope(ref, distr) + sim.track_envelope() # clean shutdown sim.finalize() -def analytic_final_estimate(APL_g, rigidity_Tm, APL_length, beta_0, alpha_0): +def get_beta_gamma(Ek: float, m0: float): + "Get beta*gamma given Ek and m_0c^2 in the same units (e.g. MeV)" + E0 = Ek + m0 + gamma_rel = E0 / m0 + beta_rel = np.sqrt( + (1.0 - 1.0 / np.sqrt(gamma_rel)) * (1.0 + 1.0 / np.sqrt(gamma_rel)) + ) + return beta_rel * gamma_rel + + +def get_rigidity(P0: float): + "Convert momentum P0 [eV/c] to magnetic rigidity [T*m] for electron" + return -P0 / 299792458 + + +def do_analytic_backprop(sigma_mid: float, APL_g: float, beta_gamma: float, rigidity): + """ + Make analytical back-propagation from lens midpoint to entry, + taking the beam sigma at the midpoint [m] and APL gradient [T/m], + along with the assumed beta_gamma and rigitdity [T*m]. + + Also print analytical estimates of initial, midpoint, and final Twiss parameters; + these can be used to compare simulation parameters with for tests. + + Returns: + (emitg, beta_0, gamma_0, mu_0, alpha_0), + which defines the beam parameters at the lens entry. + """ + + # Midpoint parameters + alpha_mid = 0.0 + # sigma_mid = 10e-6 # [m] + emitn = 10e-6 # [m] + emitg = emitn / beta_gamma # [m] + beta_mid = sigma_mid**2 / emitg # [m] + gamma_mid = 1 / beta_mid # [1/m] + print( + f"sigma_mid = {sigma_mid} [m], beta_mid = {beta_mid} [m], gamma_mid = {gamma_mid} [m], alpha_mid = {alpha_mid}", + flush=True, + ) + print( + f"emitn = {emitn} [m], emitg = {emitg} [m], beta_gamma = {beta_gamma}, rigidity = {rigidity} [T*m]", + flush=True, + ) + print(flush=True) + + # Back-propagate 1/2 lens length as in vacuum, + # from symmetry point in the middle of the lens to the start of the lens + assert alpha_mid == 0.0 + beta_0 = beta_mid + (APL_length / 2) ** 2 / beta_mid + alpha_0 = +APL_length / 2 / beta_mid + gamma_0 = gamma_mid + sigma_0 = math.sqrt(emitg * beta_0) + sigmap_0 = math.sqrt(emitg * gamma_0) + mu_0 = alpha_0 / math.sqrt(beta_0 * gamma_0) + print( + f"sigma_0 = {sigma_0} [m], beta_0 = {beta_0} [m], alpha_0 = {alpha_0}, sigmap_0 = {sigmap_0}", + flush=True, + ) + print(flush=True) + + # Forward-propagate that through the focusing/defocusing lens + # from the beginning, ignoring energy spread + (beta_end, alpha_end, gamma_end) = analytic_final_estimate( + APL_g, rigidity, APL_length, beta_0, alpha_0 + ) + + print( + f"beta_end = {beta_end} [m], alpha_end = {alpha_end} [-], gamma_end = {gamma_end} [1/m]", + flush=True, + ) + sigma_end = np.sqrt(emitg * beta_end) + sigmap_end = np.sqrt(emitg * gamma_end) + print(f"sigma_end = {sigma_end} [m], sigmap_end = {sigmap_end} [-]", flush=True) + print(flush=True) + + return (emitg, beta_0, gamma_0, mu_0, alpha_0) + + +def analytic_final_estimate( + APL_g: float, + rigidity: float, + APL_length: float, + beta_0: float, + alpha_0: float, + printk=True, +): "Analytical estimates of the beam Twiss parameters after the Plasma Lens" - k = APL_g / rigidity_Tm + k = APL_g / rigidity + if printk: + print(f"k = {k} [1/m^2]", flush=True) + print("", flush=True) - print(f"k = {k} [1/m^2]") if k > 0: M = np.asarray( [ @@ -188,10 +358,41 @@ def analytic_final_estimate(APL_g, rigidity_Tm, APL_length, beta_0, alpha_0): # Do the Twiss propagation B0 = np.asarray([[beta_0, -alpha_0], [-alpha_0, (1 + alpha_0**2) / beta_0]]) B = M @ B0 @ M.T - # print(B) + # print(B, flush=True) beta_end = B[0, 0] alpha_end = -B[0, 1] gamma_end = B[1, 1] return (beta_end, alpha_end, gamma_end) + + +def analytic_sigma_function(APL_g: float, sigma_mid: float): + """ + Do a plasma lens analytical envelope calculation with + the given APL gradient APL_g [T/m] and sigma_mid [m]. + """ + + E_MeV = kin_energy_MeV + mass_MeV + P0_MeV_c = np.sqrt(E_MeV**2 - mass_MeV**2) + + beta_gamma = get_beta_gamma(E_MeV, mass_MeV) + rigidity = get_rigidity(P0_MeV_c * 1e6) # [T*m] + + (emitg, beta_0, gamma_0, mu_0, alpha_0) = do_analytic_backprop( + sigma_mid, APL_g, beta_gamma, rigidity + ) + + s = np.linspace(0, APL_length) + sigma = np.empty_like(s) + + for i, s_ in enumerate(s): + (beta_i, alpha_i, gamma_i) = analytic_final_estimate( + APL_g, rigidity, s_, beta_0, alpha_0, printk=False + ) + sigma[i] = np.sqrt(beta_i * emitg) + + print(s) + print(sigma) + + return (s, sigma) diff --git a/examples/active_plasma_lens/run_APL_ChrPlasmaLens_defocusing.py b/examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_defocusing.py similarity index 84% rename from examples/active_plasma_lens/run_APL_ChrPlasmaLens_defocusing.py rename to examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_defocusing.py index 6664ba895..4cbd9ff5f 100755 --- a/examples/active_plasma_lens/run_APL_ChrPlasmaLens_defocusing.py +++ b/examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_defocusing.py @@ -8,6 +8,6 @@ from run_APL import run_APL_tracking -# Run the ChrPlasmaLens APL test in defocusing mode +# Run the ChrPlasmaLens/tracking APL test in defocusing mode # (rigiditiy is negative. Gradient given in [T/m]) run_APL_tracking(+1000, 1.0e-3, 100e-6, lensType="ChrPlasmaLens") diff --git a/examples/active_plasma_lens/run_APL_ChrPlasmaLens_focusing.py b/examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_focusing.py similarity index 85% rename from examples/active_plasma_lens/run_APL_ChrPlasmaLens_focusing.py rename to examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_focusing.py index 62aa04323..68ed9b5cd 100755 --- a/examples/active_plasma_lens/run_APL_ChrPlasmaLens_focusing.py +++ b/examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_focusing.py @@ -8,7 +8,7 @@ from run_APL import run_APL_tracking -# Run the ChrPlasmaLens APL test in focusing mode +# Run the ChrPlasmaLens/tracking APL test in focusing mode # (rigiditiy is also negative. Gradient given in [T/m]) run_APL_tracking(-1000, 1.0e-3, 100e-6, lensType="ChrPlasmaLens") diff --git a/examples/active_plasma_lens/run_APL_ChrPlasmaLens_zero.py b/examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_zero.py similarity index 75% rename from examples/active_plasma_lens/run_APL_ChrPlasmaLens_zero.py rename to examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_zero.py index 2a2c510a3..6c947cd7d 100755 --- a/examples/active_plasma_lens/run_APL_ChrPlasmaLens_zero.py +++ b/examples/active_plasma_lens/run_APL_ChrPlasmaLens_tracking_zero.py @@ -8,8 +8,8 @@ from run_APL import run_APL_tracking -# Run the ChrPlasmaLens APL test in no-field mode +# Run the ChrPlasmaLens/tracking APL test in no-field mode run_APL_tracking(0.0, 1e-3, 10e-6, lensType="ChrPlasmaLens") -# Gives the same output -- for creating the analysis file +# Gives the same output # run_APL_tracking(0.0, 1e-3, 10e-6, lensType='ChrDrift') diff --git a/examples/active_plasma_lens/run_APL_ConstF_envelope_defocusing.py b/examples/active_plasma_lens/run_APL_ConstF_envelope_defocusing.py new file mode 100755 index 000000000..af2049a1c --- /dev/null +++ b/examples/active_plasma_lens/run_APL_ConstF_envelope_defocusing.py @@ -0,0 +1,13 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from run_APL import run_APL_envelope + +# Run the ConstF/envelope APL test in defocusing mode +# (rigiditiy is negative. Gradient given in [T/m]) +run_APL_envelope(+1000, 1.0e-3, 100e-6, lensType="ConstF") diff --git a/examples/active_plasma_lens/run_APL_ConstF_envelope_focusing.py b/examples/active_plasma_lens/run_APL_ConstF_envelope_focusing.py new file mode 100755 index 000000000..9481fa69f --- /dev/null +++ b/examples/active_plasma_lens/run_APL_ConstF_envelope_focusing.py @@ -0,0 +1,13 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from run_APL import run_APL_envelope + +# Run the ConstF/envelope APL test in focusing mode +# (rigiditiy is negative. Gradient given in [T/m]) +run_APL_envelope(-1000, 1.0e-3, 100e-6, lensType="ConstF") diff --git a/examples/active_plasma_lens/run_APL_ConstF_envelope_zero.py b/examples/active_plasma_lens/run_APL_ConstF_envelope_zero.py new file mode 100755 index 000000000..f7ead30ec --- /dev/null +++ b/examples/active_plasma_lens/run_APL_ConstF_envelope_zero.py @@ -0,0 +1,12 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from run_APL import run_APL_envelope + +# Run the ConstF/envelope APL test in no-field mode +run_APL_envelope(0.0, 1e-3, 10e-6, lensType="ConstF") diff --git a/examples/active_plasma_lens/run_APL_ConstF_tracking_defocusing.py b/examples/active_plasma_lens/run_APL_ConstF_tracking_defocusing.py new file mode 100755 index 000000000..4123b32ed --- /dev/null +++ b/examples/active_plasma_lens/run_APL_ConstF_tracking_defocusing.py @@ -0,0 +1,13 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from run_APL import run_APL_tracking + +# Run the ConstF/tracking APL test in defocusing mode +# (rigiditiy is negative. Gradient given in [T/m]) +run_APL_tracking(+1000, 1.0e-3, 100e-6, lensType="ConstF") diff --git a/examples/active_plasma_lens/run_APL_ConstF_tracking_focusing.py b/examples/active_plasma_lens/run_APL_ConstF_tracking_focusing.py new file mode 100755 index 000000000..18b451299 --- /dev/null +++ b/examples/active_plasma_lens/run_APL_ConstF_tracking_focusing.py @@ -0,0 +1,15 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from run_APL import run_APL_tracking + +# Run the ConstF/tracking APL test in focusing mode +# (rigiditiy is also negative. Gradient given in [T/m]) +run_APL_tracking(-1000, 1.0e-3, 100e-6, lensType="ConstF") + +# about -4118 T/m is fun diff --git a/examples/active_plasma_lens/run_APL_ConstF_tracking_zero.py b/examples/active_plasma_lens/run_APL_ConstF_tracking_zero.py new file mode 100755 index 000000000..d1f6d91d6 --- /dev/null +++ b/examples/active_plasma_lens/run_APL_ConstF_tracking_zero.py @@ -0,0 +1,15 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Kyrre Sjobak +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from run_APL import run_APL_tracking + +# Run the ConstF/tracking APL test in no-field mode +run_APL_tracking(0.0, 1e-3, 10e-6, lensType="ConstF") + +# Gives the same output -- for creating the analysis file +# run_APL_tracking(0.0, 1e-3, 10e-6, lensType='ChrDrift') diff --git a/examples/distgen/README.rst b/examples/distgen/README.rst index c81081456..7bf4e659b 100644 --- a/examples/distgen/README.rst +++ b/examples/distgen/README.rst @@ -219,3 +219,50 @@ We run the following script to analyze correctness: .. literalinclude:: analysis_semigaussian.py :language: python3 :caption: You can copy this file from ``examples/distgen/analysis_semigaussian.py``. + + +.. _examples-distgen-spinvmf: + +Spin Sampling from a vMF Distribution +===================================== + +This tests the sampling of initial particle spin from a von Mises-Fisher distribution, given an initial input polarization. +The phase space distribution coincides with the 4D Kurth distribution used in examples-distgen-kurth4d. + +In this test, the initial and final values of of the mean spin 3-vector (i.e., the polarization vector) must agree with nominal values. + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_kurth4d_spin.py`` or +* ImpactX **executable** using an input file: ``impactx input_kurth4d_spin.in`` + +For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python: Script + + .. literalinclude:: run_kurth4d_spin.py + :language: python3 + :caption: You can copy this file from ``examples/distgen/run_kurth4d_spin.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_kurth4d_spin.in + :language: ini + :caption: You can copy this file from ``examples/distgen/input_kurth4d_spin.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_kurth4d_spin.py`` + + .. literalinclude:: analysis_kurth4d_spin.py + :language: python3 + :caption: You can copy this file from ``examples/distgen/analysis_kurth4d_spin.py``. diff --git a/examples/distgen/analysis_kurth4d_spin.py b/examples/distgen/analysis_kurth4d_spin.py new file mode 100755 index 000000000..87f49e8e8 --- /dev/null +++ b/examples/distgen/analysis_kurth4d_spin.py @@ -0,0 +1,74 @@ +#!/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 + + +def get_polarization(beam): + """Calculate polarization vector, given by the mean values of spin components. + + Returns + ------- + polarization_x, polarization_y, polarization_z + """ + polarization_x = np.mean(beam["spin_x"]) + polarization_y = np.mean(beam["spin_y"]) + polarization_z = np.mean(beam["spin_z"]) + + return (polarization_x, polarization_y, polarization_z) + + +# 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:") +polarization_x, polarization_y, polarization_z = get_polarization(initial) +print( + f" polarization_x={polarization_x:e} polarization_y={polarization_y:e} polarization_z={polarization_z:e}" +) + +atol = 1.3 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" atol={atol}") + +assert np.allclose( + [polarization_x, polarization_y, polarization_z], + [ + 0.7, + 0.0, + 0.0, + ], + atol=atol, +) + +print("") +print("Final Beam:") +polarization_x, polarization_y, polarization_z = get_polarization(final) +print( + f" polarization_x={polarization_x:e} polarization_y={polarization_y:e} polarization_z={polarization_z:e}" +) + +atol = 1.3 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" atol={atol}") + +assert np.allclose( + [polarization_x, polarization_y, polarization_z], + [ + 0.7, + 0.0, + 0.0, + ], + atol=atol, +) diff --git a/examples/distgen/input_kurth4d_spin.in b/examples/distgen/input_kurth4d_spin.in new file mode 100644 index 000000000..f2794e960 --- /dev/null +++ b/examples/distgen/input_kurth4d_spin.in @@ -0,0 +1,41 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 2.0e3 +beam.charge = 1.0e-9 +beam.particle = proton +beam.distribution = kurth4d +beam.lambdaX = 1.0e-3 +beam.lambdaY = beam.lambdaX +beam.lambdaT = 1.0e-3 +beam.lambdaPx = 1.0e-3 +beam.lambdaPy = beam.lambdaPx +beam.lambdaPt = 2.0e-3 +beam.muxpx = 0.0 +beam.muypy = 0.0 +beam.mutpt = 0.0 +beam.polarization_x = 0.7 +beam.polarization_y = 0.0 +beam.polarization_z = 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.0e-4 + + +############################################################################### +# Algorithms +############################################################################### +algo.space_charge = false diff --git a/examples/distgen/run_kurth4d_spin.py b/examples/distgen/run_kurth4d_spin.py new file mode 100755 index 000000000..093d80cd9 --- /dev/null +++ b/examples/distgen/run_kurth4d_spin.py @@ -0,0 +1,68 @@ +#!/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.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.Kurth4D( + lambdaX=1.0e-3, + lambdaY=1.0e-3, + lambdaT=1.0e-3, + lambdaPx=1.0e-3, + lambdaPy=1.0e-3, + lambdaPt=2.0e-3, + muxpx=-0.0, + muypy=0.0, + mutpt=0.0, +) +spinv = distribution.SpinvMF( + 0.7, + 0.0, + 0.0, +) + +sim.add_particles(bunch_charge_C, distr, npart, spinv) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice) +constf = [ + monitor, + elements.ConstF(name="constf1", ds=2.0, kx=1.0, ky=1.0, kt=1.0e-4), + monitor, +] + +# assign a constf segment +sim.lattice.extend(constf) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/examples/edge_effects/README.rst b/examples/edge_effects/README.rst index cf92bca7c..6bb4c2f45 100644 --- a/examples/edge_effects/README.rst +++ b/examples/edge_effects/README.rst @@ -38,8 +38,8 @@ Analyze We run the following script to analyze correctness: -.. dropdown:: Script ``analysis_exact_quad.py`` +.. dropdown:: Script ``analysis_dipedge.py`` - .. literalinclude:: analysis_exact_quad.py + .. literalinclude:: analysis_dipedge.py :language: python3 :caption: You can copy this file from ``examples/edge_effects/analysis_dipedge.py``. diff --git a/examples/initialize_from_array/README.rst b/examples/initialize_from_array/README.rst index 359692693..a67c3c4be 100644 --- a/examples/initialize_from_array/README.rst +++ b/examples/initialize_from_array/README.rst @@ -48,6 +48,15 @@ This example can **only** be run with **Python**: For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. +.. attention:: + + In MPI-parallel simulations, ``pc.add_n_particles(...)`` is local to the MPI rank, spatial locality does not matter. + Thus, you can add particles at any MPI rank, e.g., equally chuncked up for perfect load balancing. + + You do NOT want to add the same unique particle at multiple MPI ranks. + + When ImpactX needs to sort particles spatially, it will redistribute them over MPI ranks automatically during tracking. + .. literalinclude:: run_from_array.py :language: python3 :caption: You can copy this file from ``examples/initialize_from_array/run_from_array.py``. diff --git a/examples/initialize_from_array/run_from_array.py b/examples/initialize_from_array/run_from_array.py index 831345260..7f64505cd 100644 --- a/examples/initialize_from_array/run_from_array.py +++ b/examples/initialize_from_array/run_from_array.py @@ -117,6 +117,17 @@ dx_podv, dy_podv, dt_podv, dpx_podv, dpy_podv, dpt_podv, qm_eev, w=w_podv ) +# Note for MPI-parallel simulations: +# `pc.add_n_particles(...)` is local to the MPI rank, spatial +# locality does not matter. Thus, you can add particles at any +# MPI rank, e.g., equally chuncked up for perfect load balancing. +# +# You do NOT want to add the same unique particle at multiple +# MPI ranks. +# +# When ImpactX needs to sort particles spatially, it will +# redistribute them over MPI ranks automatically during tracking. + # build the accelerator lattice monitor = elements.BeamMonitor("monitor", backend="h5") sim.lattice.extend( diff --git a/examples/polygon_aperture/README.rst b/examples/polygon_aperture/README.rst new file mode 100644 index 000000000..f60652b76 --- /dev/null +++ b/examples/polygon_aperture/README.rst @@ -0,0 +1,96 @@ +.. _examples-polygon-aperture: + +Polygon Aperture +================ + +A 2D transverse aperture of a closed polygon defined by the :math:`x` and +:math:`y` coordinates of the vertices. +The vertices must be ordered in the counter-clockwise direction and must close, +i.e. the +first and last coordinates must be the same. + +This example takes a 800 MeV proton beam generated as a waterbag +distribution with :math:`\sigma_x`, :math:`\sigma_y` both equal to +2 mm impinging directly the mask. + +Several variations are given, with the mask either transmitting or +blocking the particles, also with option rotation and transverse offset. + +Run +--- + +This example of a transmitting mask can be run **either** as: + +* **Python** script: ``python3 run_polygon_aperture.py`` or +* ImpactX **executable** using an input file: ``impactx input_polygon_aperture.in`` + +For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python: Script + + .. literalinclude:: run_polygon_aperture.py + :language: python3 + :caption: You can copy this file from ``examples/polygon_aperture/run_polygon_aperture.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_polygon_aperture.in + :language: ini + :caption: You can copy this file from ``examples/polygon_aperture/input_polygon_aperture.in``. + +Other examples are + +.. tab-set:: + + .. tab-item:: Aperture that absorbs + + .. literalinclude:: run_polygon_aperture_absorb.py + :language: python3 + :caption: You can copy this file from ``examples/polygon_aperture/run_polygon_aperture_absorb.py``. + + .. tab-item:: Aperture with rotation that absorbs + + .. literalinclude:: run_polygon_aperture_absorb_rotate.py + :language: python3 + :caption: You can copy this file from ``examples/polygon_aperture/run_polygon_aperture_absorb_rotate.py``. + + .. tab-item:: Aperture with offset that absorbs + + .. literalinclude:: run_polygon_aperture_absorb_offset.py + :language: python3 + :caption: You can copy this file from run_polygon_aperture_absorb_offset.py + + .. tab-item:: Aperture with offset and rotation that absorbs + + .. literalinclude:: run_polygon_aperture_absorb_rotate_offset.py + :language: python3 + :caption: You can copy this file from ``examples/polygon_aperture/run_polygon_aperture_absorb_rotate_offset.py``. + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_polygon_aperture.py`` + + .. literalinclude:: analysis_polygon_aperture.py + :language: python3 + :caption: You can copy this file from ``examples/polygon_aperture/analysis_polygon_aperture.py``. + +The number of surviving particles is printed and checked. + +Visualize +--------- + +You can run the following script to visualize aperture effect: + +.. dropdown:: Script ``plot_polygon_aperture.py`` + + .. literalinclude:: plot_polygon_aperture.py + :language: python3 + :caption: You can copy this file from ``examples/polygon_aperture/plot_polygon_aperture.py``. + +.. figure:: polygon_aperture.png + :alt: Initial and transmitted particles through the example polygon aperture. diff --git a/examples/polygon_aperture/analysis_polygon_aperture.py b/examples/polygon_aperture/analysis_polygon_aperture.py new file mode 100755 index 000000000..5a7e751f5 --- /dev/null +++ b/examples/polygon_aperture/analysis_polygon_aperture.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Eric G. Stern, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + + +import numpy as np +import openpmd_api as io + +# 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() +beam_final = series.iterations[last_step].particles["beam"] +final = beam_final.to_df() + +# compare number of particles +num_particles = len(initial.momentum_x) +assert num_particles == len(initial) +assert num_particles != len(final) + +print("Initial Beam: ", len(initial), " particles") +print("Final Beam: ", len(final), " particles") + +# Make sure no particles are outside of the aperture in the final particle set +abs_x_final = abs(final["position_x"]).to_numpy() +abs_y_final = abs(final["position_y"]).to_numpy() + +N = abs_x_final.shape[0] +insides = np.zeros(N, dtype=np.bool) +for i in range(N): + insides[i] = ((abs_y_final[i] < 0.5e-3) and (abs_x_final[i] < 1.5e-3)) or ( + (abs_x_final[i] < 0.5e-3) and (abs_y_final[i] < 1.5e-3) + ) + + +outsides = ~insides +ninside = insides.sum() +noutside = outsides.sum() + +assert ninside == len(final) +assert noutside == 0 diff --git a/examples/polygon_aperture/input_polygon_aperture.in b/examples/polygon_aperture/input_polygon_aperture.in new file mode 100644 index 000000000..81c210bb9 --- /dev/null +++ b/examples/polygon_aperture/input_polygon_aperture.in @@ -0,0 +1,46 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 50000 +beam.units = static +beam.kin_energy = 800.0 +beam.charge = 1.0e-9 +beam.particle = proton +beam.distribution = waterbag +beam.lambdaX = 2.0e-3 +beam.lambdaY = beam.lambdaX +beam.lambdaT = 0.4 +beam.lambdaPx = 4.0e-4 +beam.lambdaPy = beam.lambdaPx +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 pa monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.backend = h5 + +pa.type = polygon_aperture +pa.vertices_x = 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3 +pa.vertices_y = 0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 +pa.min_radius2 = 5.0e-7 +pa.action = "transmit" # this is the default transmit particles within the polygon +#pa.action = "absorb" # alternatively absorb particles within the polygon + +############################################################################### +# Algorithms +############################################################################### +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/polygon_aperture/input_polygon_aperture_absorb.in b/examples/polygon_aperture/input_polygon_aperture_absorb.in new file mode 100644 index 000000000..76f8442cd --- /dev/null +++ b/examples/polygon_aperture/input_polygon_aperture_absorb.in @@ -0,0 +1,45 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 50000 +beam.units = static +beam.kin_energy = 800.0 +beam.charge = 1.0e-9 +beam.particle = proton +beam.distribution = waterbag +beam.lambdaX = 2.0e-3 +beam.lambdaY = beam.lambdaX +beam.lambdaT = 0.4 +beam.lambdaPx = 4.0e-4 +beam.lambdaPy = beam.lambdaPx +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 pa monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.backend = h5 + +pa.type = polygon_aperture +pa.vertices_x = 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3 +pa.vertices_y = 0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 +pa.min_radius2 = 5.0e-7 +pa.action = "absorb" # absorb particles within the polygon + +############################################################################### +# Algorithms +############################################################################### +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/polygon_aperture/input_polygon_aperture_absorb_offset.in b/examples/polygon_aperture/input_polygon_aperture_absorb_offset.in new file mode 100644 index 000000000..513876042 --- /dev/null +++ b/examples/polygon_aperture/input_polygon_aperture_absorb_offset.in @@ -0,0 +1,47 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 50000 +beam.units = static +beam.kin_energy = 800.0 +beam.charge = 1.0e-9 +beam.particle = proton +beam.distribution = waterbag +beam.lambdaX = 2.0e-3 +beam.lambdaY = beam.lambdaX +beam.lambdaT = 0.4 +beam.lambdaPx = 4.0e-4 +beam.lambdaPy = beam.lambdaPx +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 pa monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.backend = h5 + +pa.type = polygon_aperture +pa.vertices_x = 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3 +pa.vertices_y = 0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 +pa.min_radius2 = 5.0e-7 +pa.action = "absorb" # absorb particles within the polygon +pa.dx = 0.0006 +pa.dy = -0.0012 + +############################################################################### +# Algorithms +############################################################################### +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/polygon_aperture/input_polygon_aperture_absorb_rotate.in b/examples/polygon_aperture/input_polygon_aperture_absorb_rotate.in new file mode 100644 index 000000000..8beda9935 --- /dev/null +++ b/examples/polygon_aperture/input_polygon_aperture_absorb_rotate.in @@ -0,0 +1,46 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 50000 +beam.units = static +beam.kin_energy = 800.0 +beam.charge = 1.0e-9 +beam.particle = proton +beam.distribution = waterbag +beam.lambdaX = 2.0e-3 +beam.lambdaY = beam.lambdaX +beam.lambdaT = 0.4 +beam.lambdaPx = 4.0e-4 +beam.lambdaPy = beam.lambdaPx +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 pa monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.backend = h5 + +pa.type = polygon_aperture +pa.vertices_x = 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3 +pa.vertices_y = 0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 +pa.min_radius2 = 5.0e-7 +pa.action = "absorb" # absorb particles within the polygon +pa.rotation = 30.0 # degrees? + +############################################################################### +# Algorithms +############################################################################### +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/polygon_aperture/input_polygon_aperture_absorb_rotate_offset.in b/examples/polygon_aperture/input_polygon_aperture_absorb_rotate_offset.in new file mode 100644 index 000000000..6a41a1500 --- /dev/null +++ b/examples/polygon_aperture/input_polygon_aperture_absorb_rotate_offset.in @@ -0,0 +1,47 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 50000 +beam.units = static +beam.kin_energy = 800.0 +beam.charge = 1.0e-9 +beam.particle = proton +beam.distribution = waterbag +beam.lambdaX = 2.0e-3 +beam.lambdaY = beam.lambdaX +beam.lambdaT = 0.4 +beam.lambdaPx = 4.0e-4 +beam.lambdaPy = beam.lambdaPx +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 pa monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.backend = h5 + +pa.type = polygon_aperture +pa.vertices_x = 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3 +pa.vertices_y = 0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 +pa.min_radius2 = 5.0e-7 +pa.action = "absorb" # absorb particles within the polygon +pa.rotation = 30.0 # degrees? +pa.dx = 0.0006 +pa.dy = -0.0012 +############################################################################### +# Algorithms +############################################################################### +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/polygon_aperture/plot_polygon_aperture.py b/examples/polygon_aperture/plot_polygon_aperture.py new file mode 100755 index 000000000..93f83ff06 --- /dev/null +++ b/examples/polygon_aperture/plot_polygon_aperture.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2025 ImpactX contributors +# Authors: Eric G. Stern, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + +import argparse + +import matplotlib.pyplot as plt +import openpmd_api as io + +# options to run this script +parser = argparse.ArgumentParser(description="Plot action of the polygon aperture.") +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() + + +f, axs = plt.subplots(1, 2, figsize=(8, 4), constrained_layout=True) + +axs[0].scatter(initial["position_x"] * 1.0e3, initial["position_y"] * 1.0e3) +axs[0].set_title("initial") +axs[0].set_xlabel(r"$x$ [mm]") +axs[0].set_ylabel(r"$y$ [mm]") +axs[0].set_xlim([-5.5, 5.5]) +axs[0].set_ylim([-5.5, 5.5]) + +axs[1].scatter(final["position_x"] * 1.0e3, final["position_y"] * 1.0e3) +axs[1].set_title("final") +axs[1].set_xlabel(r"$x$ [mm]") +axs[1].set_ylabel(r"$y$ [mm]") +axs[1].set_xlim([-5.5, 5.5]) +axs[1].set_ylim([-5.3, 5.3]) + + +plt.tight_layout() +if args.save_png: + plt.savefig("polygon_aperture.png") +else: + plt.show() diff --git a/examples/polygon_aperture/polygon_aperture.png b/examples/polygon_aperture/polygon_aperture.png new file mode 100644 index 000000000..e0ff6e324 Binary files /dev/null and b/examples/polygon_aperture/polygon_aperture.png differ diff --git a/examples/polygon_aperture/run_polygon_aperture.py b/examples/polygon_aperture/run_polygon_aperture.py new file mode 100755 index 000000000..057a1aa96 --- /dev/null +++ b/examples/polygon_aperture/run_polygon_aperture.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Eric G. Stern, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from scipy.constants import c, eV, m_p + +from impactx import ImpactX, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = 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 = 800.0 # reference energy 800 MeV proton +bunch_charge_C = 1.0e-9 # used with space charge +npart = 50000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(1.0e-6 * m_p * c**2 / eV).set_kin_energy_MeV( + kin_energy_MeV +) + +# particle bunch +distr = distribution.Waterbag( + lambdaX=2.0e-3, + lambdaY=2.0e-3, + lambdaT=0.4, + lambdaPx=4.0e-4, + lambdaPy=4.0e-4, + 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") + +vertices_x = [ + float(u) + for u in "0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3".split() +] +vertices_y = [ + float(u) + for u in "0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3".split() +] +mr2 = 2 * 0.5e-3**2 + +aperture = elements.PolygonAperture(vertices_x, vertices_y, action="transmit") + +print(aperture.to_dict()) +assert aperture.min_radius2 == 0.0 +aperture.min_radius2 = mr2 +assert abs(aperture.min_radius2 / mr2 - 1.0) < 1.0e-15 + +# design the accelerator lattice) +ns = 1 # number of slices per ds in the element +channel = [ + monitor, + aperture, + monitor, +] + +sim.lattice.extend(channel) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/examples/polygon_aperture/run_polygon_aperture_absorb.py b/examples/polygon_aperture/run_polygon_aperture_absorb.py new file mode 100644 index 000000000..85c13faed --- /dev/null +++ b/examples/polygon_aperture/run_polygon_aperture_absorb.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Eric G. Stern, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from scipy.constants import c, eV, m_p + +from impactx import ImpactX, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = 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 = 800.0 # reference energy 800 MeV proton +bunch_charge_C = 1.0e-9 # used with space charge +npart = 50000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(1.0e-6 * m_p * c**2 / eV).set_kin_energy_MeV( + kin_energy_MeV +) + +# particle bunch +distr = distribution.Waterbag( + lambdaX=2.0e-3, + lambdaY=2.0e-3, + lambdaT=0.4, + lambdaPx=4.0e-4, + lambdaPy=4.0e-4, + 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") + +vertices_x = [ + float(u) + for u in "0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3".split() +] +vertices_y = [ + float(u) + for u in "0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3".split() +] +mr2 = 2 * 0.5e-3**2 + +aperture = elements.PolygonAperture(vertices_x, vertices_y, action="absorb") + +print(aperture.to_dict()) +assert aperture.min_radius2 == 0.0 +aperture.min_radius2 = mr2 +assert abs(aperture.min_radius2 / mr2 - 1.0) < 1.0e-15 + +# design the accelerator lattice) +ns = 1 # number of slices per ds in the element +channel = [ + monitor, + aperture, + monitor, +] + +sim.lattice.extend(channel) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/examples/polygon_aperture/run_polygon_aperture_absorb_offset.py b/examples/polygon_aperture/run_polygon_aperture_absorb_offset.py new file mode 100644 index 000000000..b081e11d4 --- /dev/null +++ b/examples/polygon_aperture/run_polygon_aperture_absorb_offset.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Eric G. Stern, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from scipy.constants import c, eV, m_p + +from impactx import ImpactX, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = 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 = 800.0 # reference energy 800 MeV proton +bunch_charge_C = 1.0e-9 # used with space charge +npart = 50000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(1.0e-6 * m_p * c**2 / eV).set_kin_energy_MeV( + kin_energy_MeV +) + +# particle bunch +distr = distribution.Waterbag( + lambdaX=2.0e-3, + lambdaY=2.0e-3, + lambdaT=0.4, + lambdaPx=4.0e-4, + lambdaPy=4.0e-4, + 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") + +vertices_x = [ + float(u) + for u in "0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3".split() +] +vertices_y = [ + float(u) + for u in "0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3".split() +] +mr2 = 2 * 0.5e-3**2 + +aperture = elements.PolygonAperture( + vertices_x, vertices_y, action="absorb", dx=0.0006, dy=-0.0012 +) + +print(aperture.to_dict()) +assert aperture.min_radius2 == 0.0 +aperture.min_radius2 = mr2 +assert abs(aperture.min_radius2 / mr2 - 1.0) < 1.0e-15 + +# design the accelerator lattice) +ns = 1 # number of slices per ds in the element +channel = [ + monitor, + aperture, + monitor, +] + +sim.lattice.extend(channel) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/examples/polygon_aperture/run_polygon_aperture_absorb_rotate.py b/examples/polygon_aperture/run_polygon_aperture_absorb_rotate.py new file mode 100644 index 000000000..f999e527c --- /dev/null +++ b/examples/polygon_aperture/run_polygon_aperture_absorb_rotate.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Eric G. Stern, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from scipy.constants import c, eV, m_p + +from impactx import ImpactX, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = 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 = 800.0 # reference energy 800 MeV proton +bunch_charge_C = 1.0e-9 # used with space charge +npart = 50000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(1.0e-6 * m_p * c**2 / eV).set_kin_energy_MeV( + kin_energy_MeV +) + +# particle bunch +distr = distribution.Waterbag( + lambdaX=2.0e-3, + lambdaY=2.0e-3, + lambdaT=0.4, + lambdaPx=4.0e-4, + lambdaPy=4.0e-4, + 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") + +vertices_x = [ + float(u) + for u in "0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3".split() +] +vertices_y = [ + float(u) + for u in "0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3".split() +] +mr2 = 2 * 0.5e-3**2 + +aperture = elements.PolygonAperture( + vertices_x, vertices_y, action="absorb", rotation=30.0 +) + +print(aperture.to_dict()) +assert aperture.min_radius2 == 0.0 +aperture.min_radius2 = mr2 +assert abs(aperture.min_radius2 / mr2 - 1.0) < 1.0e-15 + +# design the accelerator lattice) +ns = 1 # number of slices per ds in the element +channel = [ + monitor, + aperture, + monitor, +] + +sim.lattice.extend(channel) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/examples/polygon_aperture/run_polygon_aperture_absorb_rotate_offset.py b/examples/polygon_aperture/run_polygon_aperture_absorb_rotate_offset.py new file mode 100644 index 000000000..00ba3369c --- /dev/null +++ b/examples/polygon_aperture/run_polygon_aperture_absorb_rotate_offset.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Eric G. Stern, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from scipy.constants import c, eV, m_p + +from impactx import ImpactX, distribution, elements + +sim = ImpactX() + +# set numerical parameters and IO control +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = 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 = 800.0 # reference energy 800 MeV proton +bunch_charge_C = 1.0e-9 # used with space charge +npart = 50000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(1.0e-6 * m_p * c**2 / eV).set_kin_energy_MeV( + kin_energy_MeV +) + +# particle bunch +distr = distribution.Waterbag( + lambdaX=2.0e-3, + lambdaY=2.0e-3, + lambdaT=0.4, + lambdaPx=4.0e-4, + lambdaPy=4.0e-4, + 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") + +vertices_x = [ + float(u) + for u in "0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3 1.5e-3 1.5e-3 0.5e-3".split() +] +vertices_y = [ + float(u) + for u in "0.5e-3 1.5e-3 1.5e-3 0.5e-3 0.5e-3 -0.5e-3 -0.5e-3 -1.5e-3 -1.5e-3 -0.5e-3 -0.5e-3 0.5e-3 0.5e-3".split() +] +mr2 = 2 * 0.5e-3**2 + +aperture = elements.PolygonAperture( + vertices_x, vertices_y, action="absorb", dx=0.0006, dy=-0.0012, rotation=30.0 +) + +print(aperture.to_dict()) +assert aperture.min_radius2 == 0.0 +aperture.min_radius2 = mr2 +assert abs(aperture.min_radius2 / mr2 - 1.0) < 1.0e-15 + +# design the accelerator lattice) +ns = 1 # number of slices per ds in the element +channel = [ + monitor, + aperture, + monitor, +] + +sim.lattice.extend(channel) + +# run simulation +sim.track_particles() + +# clean shutdown +sim.finalize() diff --git a/examples/requirements.txt b/examples/requirements.txt index 006ad04b6..71cab47a5 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -2,7 +2,7 @@ # Work-Around for https://github.com/nkarast/PyNAFF/pull/10 # pynaff git+https://github.com/ax3l/PyNAFF.git@fix-complex-warning -matplotlib +matplotlib>=3.3 numpy openpmd-api pandas diff --git a/examples/solenoid_restart/analysis_solenoid.py b/examples/solenoid_restart/analysis_solenoid.py index 7d1d10f82..eb9f5fc93 100755 --- a/examples/solenoid_restart/analysis_solenoid.py +++ b/examples/solenoid_restart/analysis_solenoid.py @@ -79,7 +79,7 @@ def get_moments(beam): ) atol = 0.0 # ignored -rtol = 2.3 * num_particles**-0.5 # from random sampling of a smooth distribution +rtol = 1.4 * num_particles**-0.5 # from random sampling of a smooth distribution print(f" rtol={rtol} (ignored: atol~={atol})") assert np.allclose( diff --git a/setup.py b/setup.py index da7941762..74f707905 100644 --- a/setup.py +++ b/setup.py @@ -237,7 +237,7 @@ def build_extension(self, ext): setup( name="impactx", # note PEP-440 syntax: x.y.zaN but x.y.z.devN - version="25.11", + version="25.12", packages=["impactx"], # Python sources: package_dir={"": "src/python"}, diff --git a/src/ImpactX.H b/src/ImpactX.H index 531dfc95f..90c257ccc 100644 --- a/src/ImpactX.H +++ b/src/ImpactX.H @@ -85,12 +85,14 @@ namespace impactx * @param bunch_charge bunch charge (C) * @param distr distribution function to draw from (object) * @param npart number of particles to draw + * @param spin_distr optional spin distribution */ void add_particles ( amrex::ParticleReal bunch_charge, distribution::KnownDistributions distr, - amrex::Long npart + amrex::Long npart, + std::optional spin_distr = std::nullopt ); /** Validate the simulation is ready to run the particle tracking loop via @see track_particles diff --git a/src/elements/All.H b/src/elements/All.H index 4005e848b..3205c1bf8 100644 --- a/src/elements/All.H +++ b/src/elements/All.H @@ -32,6 +32,7 @@ #include "Multipole.H" #include "NonlinearLens.H" #include "PlaneXYRot.H" +#include "PolygonAperture.H" #include "Programmable.H" #include "PRot.H" #include "Quad.H" @@ -76,6 +77,7 @@ namespace impactx::elements Multipole, NonlinearLens, PlaneXYRot, + PolygonAperture, Programmable, PRot, Quad, diff --git a/src/elements/CMakeLists.txt b/src/elements/CMakeLists.txt index f4c7dad33..dc935d00b 100644 --- a/src/elements/CMakeLists.txt +++ b/src/elements/CMakeLists.txt @@ -1,6 +1,7 @@ target_sources(lib PRIVATE Aperture.cpp + PolygonAperture.cpp Programmable.cpp Source.cpp ) diff --git a/src/elements/ChrPlasmaLens.H b/src/elements/ChrPlasmaLens.H index b50723051..54b483282 100644 --- a/src/elements/ChrPlasmaLens.H +++ b/src/elements/ChrPlasmaLens.H @@ -4,7 +4,7 @@ * * This file is part of ImpactX. * - * Authors: Chad Mitchell, Axel Huebl + * Authors: Chad Mitchell, Axel Huebl, Kyrre Sjobak * License: BSD-3-Clause-LBNL */ #ifndef IMPACTX_CHRPLASMALENS_H diff --git a/src/elements/ConstF.H b/src/elements/ConstF.H index 3b33a806d..3eb845927 100644 --- a/src/elements/ConstF.H +++ b/src/elements/ConstF.H @@ -4,7 +4,7 @@ * * This file is part of ImpactX. * - * Authors: Chad Mitchell, Axel Huebl + * Authors: Chad Mitchell, Axel Huebl, Kyrre Sjobak * License: BSD-3-Clause-LBNL */ #ifndef IMPACTX_CONSTF_H @@ -102,19 +102,60 @@ namespace impactx::elements amrex::ParticleReal const betgam2 = powi<2>(refpart.pt) - 1.0_prt; // trigo - auto const [sin_kxds, cos_kxds] = amrex::Math::sincos(m_kx * slice_ds); - m_cos_kxds = cos_kxds; - m_const_x = -m_kx * sin_kxds; - auto const [sin_kyds, cos_kyds] = amrex::Math::sincos(m_ky * slice_ds); - m_cos_kyds = cos_kyds; - m_const_y = -m_ky * sin_kyds; + // Note: The convention here is m_kx[1/m] = sign(k)*sqrt(abs(k)) + // while k[1/m^2] is in the usual quadrupole-like convention + if (m_kx > 0_prt) + { + auto const [sin_kxds, cos_kxds] = amrex::Math::sincos(m_kx * slice_ds); + m_cos_kxds = cos_kxds; + m_const_x = -m_kx * sin_kxds; + m_sincx = sin_kxds / m_kx; + } + else if (m_kx < 0_prt) + { + auto const sinh_kxds = std::sinh(-m_kx*slice_ds); + auto const cosh_kxds = std::cosh(-m_kx*slice_ds); + m_cos_kxds = cosh_kxds; + m_const_x = -m_kx * sinh_kxds; + m_sincx = -sinh_kxds / m_kx; + } + else //m_kx == 0 + { + m_cos_kxds = 1.0_prt; + m_const_x = 0.0_prt; + m_sincx = slice_ds; + } + + if (m_ky > 0_prt) + { + auto const [sin_kyds, cos_kyds] = amrex::Math::sincos(m_ky * slice_ds); + m_cos_kyds = cos_kyds; + m_const_y = -m_ky * sin_kyds; + m_sincy = sin_kyds / m_ky; + } + else if (m_ky < 0_prt) + { + auto const sinh_kyds = std::sinh(-m_ky*slice_ds); + auto const cosh_kyds = std::cosh(-m_ky*slice_ds); + m_cos_kyds = cosh_kyds; + m_const_y = -m_ky * sinh_kyds; + m_sincy = -sinh_kyds / m_ky; + } + else //m_ky == 0 + { + m_cos_kyds = 1.0_prt; + m_const_y = 0.0_prt; + m_sincy = slice_ds; + } + + if (m_kt < 0_prt) + { + throw std::runtime_error(std::string(type) + ": must have kt >= 0!"); + } auto const [sin_ktds, cos_ktds] = amrex::Math::sincos(m_kt * slice_ds); m_cos_ktds = cos_ktds; m_const_t = -m_kt * betgam2 * sin_ktds; - - // intermediate quantities - to avoid division by zero - m_sincx = m_kx > 0 ? std::sin(m_kx * slice_ds) / m_kx : slice_ds; - m_sincy = m_ky > 0 ? std::sin(m_ky * slice_ds) / m_ky : slice_ds; + // intermediate quantity - to avoid division by zero amrex::ParticleReal const sinct = m_kt > 0 ? std::sin(m_kt * slice_ds) / m_kt : slice_ds; m_const_pt = sinct / betgam2; } @@ -242,29 +283,66 @@ namespace impactx::elements amrex::ParticleReal const betgam2 = powi<2>(refpart.pt) - 1.0_prt; // trigo - auto const [sin_kxds, cos_kxds] = amrex::Math::sincos(m_kx * slice_ds); - amrex::ParticleReal const_x = -m_kx * sin_kxds; - auto const [sin_kyds, cos_kyds] = amrex::Math::sincos(m_ky * slice_ds); - amrex::ParticleReal const_y = -m_ky * sin_kyds; + // Note: The convention here is m_kx[1/m] = sign(k)*sqrt(abs(k)) + // while k[1/m^2] is in the usual quadrupole-like convention + amrex::ParticleReal sin_kxds, cos_kxds, const_x, sincx; + if (m_kx > 0_prt) { + std::tie(sin_kxds, cos_kxds) = amrex::Math::sincos(m_kx * slice_ds); + const_x = -m_kx * sin_kxds; + sincx = sin_kxds / m_kx; + } + else if (m_kx < 0_prt) { + sin_kxds = std::sinh(-m_kx*slice_ds); + cos_kxds = std::cosh(-m_kx*slice_ds); + const_x = -m_kx * sin_kxds; + sincx = -sin_kxds / m_kx; + } + else { //m_kx == 0 + cos_kxds = 1.0_prt; + const_x = 0.0_prt; + sincx = slice_ds; + } + + amrex::ParticleReal sin_kyds, cos_kyds, const_y, sincy; + if (m_ky > 0_prt) { + std::tie(sin_kyds, cos_kyds) = amrex::Math::sincos(m_ky * slice_ds); + const_y = -m_ky * sin_kyds; + sincy = sin_kyds / m_ky; + } + else if (m_ky < 0_prt) { + sin_kyds = std::sinh(-m_ky*slice_ds); + cos_kyds = std::cosh(-m_ky*slice_ds); + const_y = -m_ky * sin_kyds; + sincy = -sin_kyds / m_ky; + } + else { //m_ky == 0 + cos_kyds = 1.0_prt; + const_y = 0.0_prt; + sincy = slice_ds; + } + + if (m_kt < 0_prt) { + throw std::runtime_error(std::string(type) + ": must have kt >= 0!"); + } auto const [sin_ktds, cos_ktds] = amrex::Math::sincos(m_kt * slice_ds); amrex::ParticleReal const_t = -m_kt * betgam2 * sin_ktds; - // intermediate quantities - to avoid division by zero - amrex::ParticleReal sincx = m_kx > 0 ? std::sin(m_kx * slice_ds) / m_kx : slice_ds; - amrex::ParticleReal sincy = m_ky > 0 ? std::sin(m_ky * slice_ds) / m_ky : slice_ds; amrex::ParticleReal const sinct = m_kt > 0 ? std::sin(m_kt * slice_ds) / m_kt : slice_ds; amrex::ParticleReal const_pt = sinct / betgam2; // assign linear map matrix elements Map6x6 R = Map6x6::Identity(); + R(1,1) = cos_kxds; R(1,2) = sincx; R(2,1) = const_x; R(2,2) = cos_kxds; + R(3,3) = cos_kyds; R(3,4) = sincy; R(4,3) = const_y; R(4,4) = cos_kyds; + R(5,5) = cos_ktds; R(5,6) = const_pt; R(6,5) = const_t; diff --git a/src/elements/PolygonAperture.H b/src/elements/PolygonAperture.H new file mode 100644 index 000000000..06ed1f9dc --- /dev/null +++ b/src/elements/PolygonAperture.H @@ -0,0 +1,331 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Eric G. Stern, Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_POLYGONAPERTURE_H +#define IMPACTX_POLYGONAPERTURE_H + + +#include "particles/ImpactXParticleContainer.H" +#include "mixin/alignment.H" +#include "mixin/beamoptic.H" +#include "mixin/thin.H" +#include "mixin/lineartransport.H" +#include "mixin/named.H" +#include "mixin/nofinalize.H" + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include + + +namespace impactx::elements +{ + + /** Dynamic data for the PolygonAperture elements + * + * Since we copy the element to the device, we cannot store this data on the element itself. + * But we can store pointers to this data with the element and keep a lookup table here, + * which we clean up in the end. + */ + namespace PolygonApertureData + { + //! last used id for a created ExactMultipole + inline int next_id = 0; + + //! host: normal multipole coefficients of the magnetic field + inline std::map> h_vertices_x = {}; + //! host: skew multipole coefficients of the magnetic field + inline std::map> h_vertices_y = {}; + + //! device: normal multipole coefficients of the magnetic field + inline std::map> d_vertices_x = {}; + //! device: skew multipole coefficients of the magnetic field + inline std::map> d_vertices_y = {}; + + } // namespace PolygonApertureData + + struct PolygonAperture + : public mixin::Named, + public mixin::BeamOptic, + public mixin::LinearTransport, + public mixin::Thin, + public mixin::Alignment, + public mixin::NoFinalize + // public amrex::simd::Vectorized + { + static constexpr auto type = "PolygonAperture"; + using PType = ImpactXParticleContainer::ParticleType; + + enum Action + { + transmit, + absorb + }; + + static std::string + action_name (Action const & action); + + /** A thin collimator element that applies a transverse aperture boundary defined + * by points of a polygon vertices. + * Particles outside the boundary are considered lost. + * + * @param action specify action of domain (transmit/absorb) + * @param vertices_x array of coordinates of the horizontal vertex positions (m) + * @param vertices_y array of coordinates of the vertical vertex positions (m) + * @param min_radius2 a minimum radius**2 below which all particles are transmitted + * @param repeat_x horizontal period for repeated masking, optional (m) + * @param repeat_y vertical period for repeated masking, optional (m) + * @param shift_odd_x for hexagonal/triangular mask patterns: horizontal shift of every 2nd (odd) vertical period by repeat_x / 2. Use alignment offsets dx,dy to move whole mask as needed. + * @param dx horizontal translation error in m + * @param dy vertical translation error in m + * @param rotation_degree rotation error in the transverse plane [degrees] + * @param name a user defined and not necessarily unique name of the element + */ + PolygonAperture ( + std::vector< amrex::ParticleReal> vertices_x, + std::vector< amrex::ParticleReal> vertices_y, + amrex::ParticleReal min_radius2, + amrex::ParticleReal repeat_x, + amrex::ParticleReal repeat_y, + bool shift_odd_x, + Action action, + amrex::ParticleReal dx = 0, + amrex::ParticleReal dy = 0, + amrex::ParticleReal rotation_degree = 0, + std::optional name = std::nullopt + ) + // + : Named(std::move(name)), + Alignment(dx, dy, rotation_degree), + m_action(action), m_repeat_x(repeat_x), m_repeat_y(repeat_y), m_shift_odd_x(shift_odd_x), + m_id(PolygonApertureData::next_id), + m_min_radius2(min_radius2) { + + using namespace amrex::literals; // for _rt and _prt + + if (m_repeat_y == 0_prt && m_shift_odd_x) { + throw std::runtime_error("PolygonAperture: shift_odd_x can only be used if repeat_y is set, too!"); + } + + // next created MulPolygonAperture has another id for its data + PolygonApertureData::next_id++; + + m_nvert = vertices_x.size(); + // there must be the same number of horizontal and vertical vertices + if (m_nvert != vertices_y.size()) { + throw std::runtime_error("PolygonAperture: horizontal and vertical vertices arrays must have same length!"); + } + + // verify that first and last vertices are the same + if ((vertices_x[0] != vertices_x[m_nvert-1]) || + (vertices_y[0] != vertices_y[m_nvert-1])) { + throw std::runtime_error("PolygonAperture: first and last vertex must be exactly equal!"); + } + + // host data + PolygonApertureData::h_vertices_x[m_id] = vertices_x; + PolygonApertureData::h_vertices_y[m_id] = vertices_y; + m_vertices_x_h_data = PolygonApertureData::h_vertices_x[m_id].data(); + m_vertices_y_h_data = PolygonApertureData::h_vertices_y[m_id].data(); + + // device data + PolygonApertureData::d_vertices_x.emplace(m_id, amrex::Gpu::DeviceVector(m_nvert)); + PolygonApertureData::d_vertices_y.emplace(m_id, amrex::Gpu::DeviceVector(m_nvert)); + amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, + vertices_x.begin(), vertices_x.end(), + PolygonApertureData::d_vertices_x[m_id].begin()); + amrex::Gpu::copyAsync(amrex::Gpu::hostToDevice, + vertices_y.begin(), vertices_y.end(), + PolygonApertureData::d_vertices_y[m_id].begin()); + amrex::Gpu::streamSynchronize(); + + // low-level objects we can use on device + m_vertices_x_d_data = PolygonApertureData::d_vertices_x[m_id].data(); + m_vertices_y_d_data = PolygonApertureData::d_vertices_y[m_id].data(); + + } + + /** Push all particles */ + using BeamOptic::operator(); + + /** Compute and cache the constants for the push. + * + * In particular, used to pre-compute and cache variables that are + * independent of the individually tracked particle. + * + * @param refpart reference particle (unused) + */ + void compute_constants ([[maybe_unused]] RefPart const & refpart) + { + Alignment::compute_constants(refpart); + } + + /** This is an aperture functor, so that a variable of this type can be used like an + * aperture function. + * + * @param x particle position in x + * @param y particle position in y + * @param t particle position in t (unused) + * @param px particle momentum in x + * @param py particle momentum in y + * @param pt particle momentum in t (unused) + * @param idcpu particle global index + * @param refpart reference particle (unused) + */ + template + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + T_Real & AMREX_RESTRICT x, + T_Real & AMREX_RESTRICT y, + [[maybe_unused]] T_Real & AMREX_RESTRICT t, + T_Real & AMREX_RESTRICT px, + T_Real & AMREX_RESTRICT py, + [[maybe_unused]] T_Real & AMREX_RESTRICT pt, + T_IdCpu & AMREX_RESTRICT idcpu, + [[maybe_unused]] RefPart const & AMREX_RESTRICT refpart + ) const + { + // a complex type with two T_Real + using Complex = amrex::GpuComplex; + + using namespace amrex::literals; // for _rt and _prt + namespace stdx = amrex::simd::stdx; + using amrex::Math::powi; + using std::fmod; + using std::abs; + using amrex::arg; + using VBool = decltype(x>0_prt); + + // shift due to alignment errors of the element + shift_in(x, y, px, py); + + // define intermediate variables + amrex::ParticleReal const dx = m_repeat_x * 0.5_prt; + amrex::ParticleReal const dy = m_repeat_y * 0.5_prt; + + // shift every 2nd row by half of m_repeat_x + T_Real sx = x; + T_Real const sy = y; + VBool const shift_x = m_shift_odd_x == false ? VBool{false} : fmod(floor((y+dy)/m_repeat_y), 2) != T_Real{0}; +#ifdef AMREX_USE_SIMD + amrex::simd::stdx::where(shift_x, sx) += dx; +#else + sx += shift_x ? dx : 0_prt; +#endif + + // if the aperture is periodic, shift sx,sy coordinates to the fundamental domain + T_Real u = (m_repeat_x == 0_prt) ? sx : (fmod(abs(sx)+dx, m_repeat_x)-dx); + T_Real v = (m_repeat_y == 0_prt) ? sy : (fmod(abs(sy)+dy, m_repeat_y)-dy); + + // compare against the aperture boundary + VBool inside_aperture; + + // calculate whether point is inside polygon or not + T_Real r2 = powi<2>(u) + powi<2>(v); + // are we inside the minimum radius? + if (r2 < m_min_radius2) { + inside_aperture = VBool(true); // yes! + } else{ + +#if AMREX_DEVICE_COMPILE + amrex::ParticleReal* vertices_x = m_vertices_x_d_data; + amrex::ParticleReal* vertices_y = m_vertices_y_d_data; +#else + amrex::ParticleReal* vertices_x = m_vertices_x_h_data; + amrex::ParticleReal* vertices_y = m_vertices_y_h_data; +#endif + + // no, we have to do the polygon calculation + + /* + * The algorithms works as follows: + * + * You have a polygon (in 2D) described by points V_1, ..., V_n arranged + * counter-clockwise. You have a particle with coordinates w = (x,y). + * Get the difference vector from the particle to each of + * the vertices in turn expressed as a complex number z_k, k=1,n. Now the angle + * between the particle and vertex i and vertex j is + * + * theta_{ij} = arg(\bar{z_i} z_j) + * This is easily demonstrated by expressing the complex numbers in exponential + * polar notation and also the angle is positive in the counter-clockwise direction. + * + * Take the sum of all the angles. The angle sum for particle inside the polygon will + * be 2 pi. The sum of angles of particle outside of the polygon will be 0. + */ + T_Real thetasum = 0.0; + Complex v1conj, v2; + + for (size_t i = 0; i < m_nvert-1; ++i) { + v1conj = Complex(u - vertices_x[i], -(v-vertices_y[i])); + v2 = Complex(u - vertices_x[i+1], v-vertices_y[i+1]); + thetasum += arg(v2 * v1conj); + } + inside_aperture = abs(thetasum/(2*ablastr::constant::math::pi)) > 1.0e-12_prt; + } + + // For transmit (default): invalidate if outside aperture + // For absorb: invalidate if inside aperture + auto const invalid_mask = m_action == Action::transmit ? !inside_aperture : inside_aperture; + amrex::ParticleIDWrapper{idcpu}.make_invalid(invalid_mask); + + // undo shift due to alignment errors of the element + shift_out(x, y, px, py); + } + + /** This pushes the reference particle. */ + using Thin::operator(); + + /** This pushes the covariance matrix. */ + using LinearTransport::operator(); + + /** This function returns the linear transport map. + * + * @param[in] refpart reference particle + * @returns 6x6 transport matrix + */ + AMREX_GPU_HOST AMREX_FORCE_INLINE + Map6x6 + transport_map ([[maybe_unused]] RefPart const & AMREX_RESTRICT refpart) const + { + throw std::runtime_error(std::string(type) + ": Envelope tracking is not yet implemented!"); + return Map6x6::Identity(); + } + + /////////////////////////////////////////////////////////////// + // class data + + Action m_action; //! action type (transmit, absorb) + amrex::ParticleReal m_repeat_x; //! horizontal period for repeated masking + amrex::ParticleReal m_repeat_y; //! vertical period for repeated masking + bool m_shift_odd_x; //! for hexagonal/triangular mask patterns: horizontal shift of every 2nd (odd) vertical period by m_repeat_x / 2. Use Alignment offsets to move whole mask as needed. + + int m_id; //! unique PolygonAperture id used for data lookup map + + amrex::ParticleReal m_min_radius2; // minimum radius in which all pass squared + + std::vector::size_type m_nvert = 0; //! number of vertices + amrex::ParticleReal* m_vertices_x_h_data = nullptr; //! non-owning pointer to host x vertices + amrex::ParticleReal* m_vertices_y_h_data = nullptr; //! non-owning pointer to host vertices + amrex::ParticleReal* m_vertices_x_d_data = nullptr; //! non-owning pointer to device x vertices + amrex::ParticleReal* m_vertices_y_d_data = nullptr; //! non-owning pointer to device y vertices + + }; + +} // namespace impactx + +#endif // IMPACTX_POLYGONAPERTURE_H diff --git a/src/elements/PolygonAperture.cpp b/src/elements/PolygonAperture.cpp new file mode 100644 index 000000000..68a851140 --- /dev/null +++ b/src/elements/PolygonAperture.cpp @@ -0,0 +1,27 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Eric G. Stern, Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "PolygonAperture.H" + +#include +#include + +std::string +impactx::elements::PolygonAperture::action_name (Action const & action) +{ + switch (action) + { + case PolygonAperture::Action::transmit : // default + return "transmit"; + case PolygonAperture::Action::absorb : + return "absorb"; + default: + throw std::runtime_error("Unknown action"); + } +} diff --git a/src/elements/Source.H b/src/elements/Source.H index 817b333c1..051550af4 100644 --- a/src/elements/Source.H +++ b/src/elements/Source.H @@ -35,16 +35,19 @@ namespace impactx::elements * * @param distribution Must read "openPMD" for this constructor * @param openpmd_path path to openPMD series as accepted by openPMD::Series + * @param active_once inject particles only for the first lattice period * @param name a user defined and not necessarily unique name of the element */ Source ( std::string distribution, std::string openpmd_path, + bool active_once = true, std::optional name = std::nullopt ) : Named(std::move(name)), m_distribution(distribution), - m_series_name(std::move(openpmd_path)) + m_series_name(std::move(openpmd_path)), + m_active_once(active_once) { if (distribution != "openPMD") { throw std::runtime_error("Only 'openPMD' distribution is supported if openpmd_path is provided!"); @@ -62,7 +65,7 @@ namespace impactx::elements void operator() ( ImpactXParticleContainer & pc, [[maybe_unused]] int step, - [[maybe_unused]] int period + int period ); /** This function returns the linear transport map. @@ -86,6 +89,7 @@ namespace impactx::elements std::string m_distribution; //! Distribution type of particles in the source std::string m_series_name; //! openPMD filename + bool m_active_once = true; //! Inject particles only for the first lattice period }; } // namespace impactx diff --git a/src/elements/Source.cpp b/src/elements/Source.cpp index 2984906ef..ca4a548fb 100644 --- a/src/elements/Source.cpp +++ b/src/elements/Source.cpp @@ -26,9 +26,14 @@ namespace impactx::elements Source::operator() ( ImpactXParticleContainer & pc, int, - int + int period ) { + // Check if only active for lattice period zero (0), e.g., in rings + if (m_active_once && period > 0) { + return; + } + #ifdef ImpactX_USE_OPENPMD auto series = io::Series(m_series_name, io::Access::READ_ONLY # if openPMD_HAVE_MPI==1 @@ -45,7 +50,7 @@ namespace impactx::elements std::string const species_name = "beam"; io::ParticleSpecies beam = iteration.particles[species_name]; // TODO: later we can make the bunch charge an option (i.e., allow rescaling a distribution) - amrex::ParticleReal bunch_charge = beam.getAttribute("charge_C").get(); + // amrex::ParticleReal bunch_charge = beam.getAttribute("charge_C").get(); auto const scalar = openPMD::RecordComponent::SCALAR; auto const getComponentRecord = [&beam](std::string comp_name) { @@ -68,8 +73,6 @@ namespace impactx::elements int const nleft = npart - navg * nprocs; std::uint64_t const npart_this_proc = (myproc < nleft) ? navg+1 : navg; // add 1 to each proc until distributed std::uint64_t npart_before_this_proc = (myproc < nleft) ? (navg+1) * myproc : navg * myproc + nleft; - auto const rel_part_this_proc = - amrex::ParticleReal(npart_this_proc) / amrex::ParticleReal(npart); // alloc data for particle attributes std::map> pinned_SoA; @@ -113,11 +116,14 @@ namespace impactx::elements // finalize distributions and deallocate temporary device global memory amrex::Gpu::streamSynchronize(); - // TODO: at this point, we ignore the "id", "qm" and "weighting" in the file. Could be improved + // TODO: at this point, we ignore the "id" and "qm" in the file. Could be improved + pc.AddNParticles(d_SoA["position_x"], d_SoA["position_y"], d_SoA["position_t"], d_SoA["momentum_x"], d_SoA["momentum_y"], d_SoA["momentum_t"], pc.GetRefParticle().qm_ratio_SI(), - bunch_charge * rel_part_this_proc); + std::nullopt, + d_SoA["weighting"], + d_SoA["spin_x"], d_SoA["spin_y"], d_SoA["spin_z"]); #else // ImpactX_USE_OPENPMD amrex::ignore_unused(pc); diff --git a/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp index 3c2779399..ddacaf9f0 100644 --- a/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp +++ b/src/envelope/spacecharge/EnvelopeSpaceChargePush.cpp @@ -40,8 +40,8 @@ namespace impactx::envelope::spacecharge Map6x6 R = Map6x6::Identity(); // physical constants and reference quantities - using ablastr::constant::SI::c; - using ablastr::constant::SI::ep0; + constexpr amrex::ParticleReal c = ablastr::constant::SI::c_v; + constexpr amrex::ParticleReal ep0 = ablastr::constant::SI::epsilon_0_v; using ablastr::constant::math::pi; amrex::ParticleReal const mass = refpart.mass; amrex::ParticleReal const charge = refpart.charge; @@ -88,8 +88,8 @@ namespace impactx::envelope::spacecharge Map6x6 R = Map6x6::Identity(); // physical constants and reference quantities - using ablastr::constant::SI::c; - using ablastr::constant::SI::ep0; + constexpr amrex::ParticleReal c = ablastr::constant::SI::c_v; + constexpr amrex::ParticleReal ep0 = ablastr::constant::SI::epsilon_0_v; using ablastr::constant::math::pi; amrex::ParticleReal const mass = refpart.mass; amrex::ParticleReal const charge = refpart.charge; diff --git a/src/initialization/InitDistribution.H b/src/initialization/InitDistribution.H index 4f1aa5874..63425ea6a 100644 --- a/src/initialization/InitDistribution.H +++ b/src/initialization/InitDistribution.H @@ -129,6 +129,69 @@ namespace impactx::initialization amrex::ParticleReal* const AMREX_RESTRICT m_part_pt; }; + /** Initialize a single particle's spin using the SpinvMF distribution + * + * Note: we usually would just write a C++ lambda below in ParallelForRNG. But, due to restrictions + * in NVCC as of 11.5, we cannot write a lambda in a lambda as we also std::visit the element + * types of our lattice elements list. + * error #3206-D: An extended __device__ lambda cannot be defined inside a generic lambda expression("operator()"). + * Thus, we fall back to writing a C++ functor here, instead of nesting two lambdas. + * + * Nvidia bug report: 3458976 + * Minimal demonstrator: https://cuda.godbolt.org/z/39e4q53Ye + */ + struct InitSingleParticleSpin + { + /** Constructor taking in pointers to particle data + * + * @param spinv spin distribution function to call + * @param part_sx the array to the particle spin in x (sx) + * @param part_sy the array to the particle spin in y (sy) + * @param part_sz the array to the particle spin in z (sz) + */ + InitSingleParticleSpin ( + distribution::SpinvMF spinv, + amrex::ParticleReal* AMREX_RESTRICT part_sx, + amrex::ParticleReal* AMREX_RESTRICT part_sy, + amrex::ParticleReal* AMREX_RESTRICT part_sz + ) + : m_spinv(std::move(spinv)), + m_part_sx(part_sx), m_part_sy(part_sy), m_part_sz(part_sz) + { + } + + InitSingleParticleSpin () = delete; + InitSingleParticleSpin (InitSingleParticleSpin const &) = default; + InitSingleParticleSpin (InitSingleParticleSpin &&) = default; + ~InitSingleParticleSpin () = default; + + /** Initialize the spin for a single particle + * + * @param i particle index in the current box + * @param engine a random number engine (with associated state) + */ + AMREX_GPU_DEVICE AMREX_FORCE_INLINE + void + operator() ( + amrex::Long i, + amrex::RandomEngine const & engine + ) const + { + m_spinv( + m_part_sx[i], + m_part_sy[i], + m_part_sz[i], + engine + ); + } + + private: + distribution::SpinvMF const m_spinv; + amrex::ParticleReal* const AMREX_RESTRICT m_part_sx; + amrex::ParticleReal* const AMREX_RESTRICT m_part_sy; + amrex::ParticleReal* const AMREX_RESTRICT m_part_sz; + }; + /** Initialize the input parameters for all distributions that read phase space ellipse parameters from the input * diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 511faf7b6..ada50354f 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -14,6 +14,7 @@ #include "particles/CovarianceMatrix.H" #include "particles/ImpactXParticleContainer.H" #include "particles/distribution/All.H" +#include "particles/distribution/SpinvMF.H" #include "particles/SplitEqually.H" #include @@ -46,18 +47,21 @@ namespace impactx amrex::ParticleReal qe; // charge (elementary charge) amrex::ParticleReal massE; // MeV/c^2 + constexpr amrex::ParticleReal m_e = ablastr::constant::SI::m_e_v; + constexpr amrex::ParticleReal m_p = ablastr::constant::SI::m_p_v; + constexpr amrex::ParticleReal MeV_inv_c2 = ablastr::constant::SI::MeV_inv_c2_v; amrex::ParticleReal gyromagnetic_anomaly; // anomalous magnetic dipole moment if (particle_type == "electron") { qe = -1.0; - massE = ablastr::constant::SI::m_e / ablastr::constant::SI::MeV_invc2; + massE = m_e / MeV_inv_c2; gyromagnetic_anomaly = 0.00115965218062; } else if (particle_type == "positron") { qe = 1.0; - massE = ablastr::constant::SI::m_e / ablastr::constant::SI::MeV_invc2; + massE = m_e / MeV_inv_c2; gyromagnetic_anomaly = 0.00115965218062; } else if (particle_type == "proton") { qe = 1.0; - massE = ablastr::constant::SI::m_p / ablastr::constant::SI::MeV_invc2; + massE = m_p / MeV_inv_c2; gyromagnetic_anomaly = 1.7928473446; } else if (particle_type == "Hminus") { qe = -1.0; @@ -71,7 +75,7 @@ namespace impactx ablastr::warn_manager::WarnPriority::low ); qe = -1.0; - massE = ablastr::constant::SI::m_e / ablastr::constant::SI::MeV_invc2; + massE = m_e / MeV_inv_c2; gyromagnetic_anomaly = 0.00115965218062; } @@ -301,7 +305,8 @@ namespace impactx ImpactX::add_particles ( amrex::ParticleReal bunch_charge, distribution::KnownDistributions distr, - amrex::Long npart + amrex::Long npart, + std::optional spin_distr ) { BL_PROFILE("ImpactX::add_particles"); @@ -343,6 +348,7 @@ namespace impactx // alloc data for particle attributes amrex::Gpu::DeviceVector x, y, t; amrex::Gpu::DeviceVector px, py, pt; + std::optional> sx, sy, sz; x.resize(npart_this_proc); y.resize(npart_this_proc); t.resize(npart_this_proc); @@ -350,6 +356,13 @@ namespace impactx py.resize(npart_this_proc); pt.resize(npart_this_proc); + bool const has_spin = spin_distr.has_value(); + if (has_spin) { + sx = amrex::Gpu::DeviceVector(npart_this_proc); + sy = amrex::Gpu::DeviceVector(npart_this_proc); + sz = amrex::Gpu::DeviceVector(npart_this_proc); + } + std::visit([&](auto&& distribution){ // initialize distributions distribution.initialize(bunch_charge, ref); @@ -360,6 +373,9 @@ namespace impactx amrex::ParticleReal * const AMREX_RESTRICT px_ptr = px.data(); amrex::ParticleReal * const AMREX_RESTRICT py_ptr = py.data(); amrex::ParticleReal * const AMREX_RESTRICT pt_ptr = pt.data(); + amrex::ParticleReal * const AMREX_RESTRICT sx_ptr = has_spin ? sx->data() : nullptr; + amrex::ParticleReal * const AMREX_RESTRICT sy_ptr = has_spin ? sy->data() : nullptr; + amrex::ParticleReal * const AMREX_RESTRICT sz_ptr = has_spin ? sz->data() : nullptr; using Distribution = std::decay_t; @@ -379,6 +395,7 @@ namespace impactx npart_this_thread = thread_chunk.size; #endif + // phase space init initialization::InitSingleParticleData const init_single_particle_data( distribution, x_ptr + my_offset, @@ -388,8 +405,18 @@ namespace impactx py_ptr + my_offset, pt_ptr + my_offset ); - amrex::ParallelForRNG(npart_this_thread, init_single_particle_data); + + // spin init + if (has_spin) { + initialization::InitSingleParticleSpin const init_single_particle_spin( + spin_distr.value(), + sx_ptr + my_offset, + sy_ptr + my_offset, + sz_ptr + my_offset + ); + amrex::ParallelForRNG(npart_this_thread, init_single_particle_spin); + } } // finalize distributions and deallocate temporary device global memory @@ -397,9 +424,14 @@ namespace impactx distribution.finalize(); }, distr); - amr_data->track_particles.m_particle_container->AddNParticles(x, y, t, px, py, pt, - ref.qm_ratio_SI(), - bunch_charge * rel_part_this_proc); + amr_data->track_particles.m_particle_container->AddNParticles( + x, y, t, + px, py, pt, + ref.qm_ratio_SI(), + bunch_charge * rel_part_this_proc, + std::nullopt, + sx, sy, sz + ); auto space_charge = get_space_charge_algo(); @@ -554,15 +586,29 @@ namespace impactx std::string unit_type; // System of units pp_dist.get("units", unit_type); + // phase space distribution distribution::KnownDistributions dist = initialization::read_distribution(pp_dist); std::string distribution; pp_dist.get("distribution", distribution); + // spin distribution + amrex::ParticleReal polarization_x = 0.0_prt, + polarization_y = 0.0_prt, + polarization_z = 0.0_prt; + bool const has_sx = pp_dist.queryWithParser("polarization_x", polarization_x); + bool const has_sy = pp_dist.queryWithParser("polarization_y", polarization_y); + bool const has_sz = pp_dist.queryWithParser("polarization_z", polarization_z); + + std::optional spin_dist; + if (has_sx || has_sy || has_sz) { + spin_dist = distribution::SpinvMF(polarization_x, polarization_y, polarization_z); + } + amrex::Long npart = 0; // Number of simulation particles if (distribution != "empty") { pp_dist.getWithParser("npart", npart); - add_particles(bunch_charge, dist, npart); + add_particles(bunch_charge, dist, npart, spin_dist); } // print information on the initialized beam diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index ef106caab..36b53274d 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -611,7 +611,10 @@ element_name) ); pp_element.get("openpmd_path", openpmd_path); } - m_lattice.emplace_back( Source(distribution, openpmd_path, element_name) ); + bool active_once = true; + pp_element.queryAdd("active_once", active_once); + + m_lattice.emplace_back( Source(distribution, openpmd_path, active_once, element_name) ); } else if (element_type == "line") { // Parse the lattice elements for the sub-lattice in the line @@ -651,6 +654,38 @@ element_name) ); } m_lattice.emplace_back(LinearMap(transport_map, ds, a["dx"], a["dy"], a["rotation_degree"]) ); + } else if (element_type == "polygon_aperture") + { + auto a = detail::query_alignment(pp_element); + + amrex::ParticleReal repeat_x = 0.0; + amrex::ParticleReal repeat_y = 0.0; + bool shift_odd_x = false; + std::string action_str = "transmit"; + + std::vector vertices_x = {0.0}; + std::vector vertices_y = {0.0}; + amrex::ParticleReal min_radius2 = 0.0; + + detail::queryAddResize(pp_element, "vertices_x", vertices_x); + detail::queryAddResize(pp_element, "vertices_y", vertices_y); + pp_element.queryAddWithParser("min_radius2", min_radius2); + + pp_element.queryAddWithParser("repeat_x", repeat_x); + pp_element.queryAddWithParser("repeat_y", repeat_y); + pp_element.queryAdd("shift_odd_x", shift_odd_x); // https://github.com/AMReX-Codes/amrex/issues/4535 + pp_element.queryAdd("action", action_str); + + //AMREX_ALWAYS_ASSERT_WITH_MESSAGE(action_str == "transmit" || action_str == "absorb", + // element_name + ".action must be \"transmit\" or \"absorb\""); + + PolygonAperture::Action action = action_str == "transmit" ? + PolygonAperture::Action::transmit : + PolygonAperture::Action::absorb; + + m_lattice.emplace_back(PolygonAperture(vertices_x, vertices_y, min_radius2, repeat_x, repeat_y, + shift_odd_x, action, a["dx"], a["dy"], a["rotation_degree"], element_name) ); + } else { amrex::Abort("Unknown type for lattice element " + element_name + ": " + element_type); } diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H index d518b5bd6..950e931a7 100644 --- a/src/particles/ImpactXParticleContainer.H +++ b/src/particles/ImpactXParticleContainer.H @@ -51,6 +51,9 @@ namespace impactx px, ///< momentum in x, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t) py, ///< momentum in y, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t) pt, ///< energy deviation, scaled by speed of light * the magnitude of the reference momentum [unitless] (at fixed s) + sx, ///< spin vector x-component [unitless] (at fixed s or t) + sy, ///< spin vector y-component [unitless] (at fixed s or t) + sz, ///< spin vector z-component [unitless] (at fixed s or t) qm, ///< charge to mass ratio, in q_e/m_e [q_e/eV] w, ///< particle weight, number of real particles represented by this macroparticle [unitless] nattribs ///< the number of attributes above (always last) @@ -63,9 +66,9 @@ namespace impactx }; //! named labels for fixed s - static constexpr auto names_s = { "position_x", "position_y", "position_t", "momentum_x", "momentum_y", "momentum_t", "qm", "weighting" }; + static constexpr auto names_s = { "position_x", "position_y", "position_t", "momentum_x", "momentum_y", "momentum_t", "spin_x", "spin_y", "spin_z", "qm", "weighting" }; //! named labels for fixed t - static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "qm", "weighting" }; + static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "spin_x", "spin_y", "spin_z", "qm", "weighting" }; static_assert(names_s.size() == nattribs); static_assert(names_t.size() == nattribs); }; @@ -179,6 +182,9 @@ namespace impactx * @param qm charge over mass in 1/eV * @param bunch_charge total charge within a bunch in C * @param w weight of each particle: how many real particles to represent + * @param sx spin component in x + * @param sy spin component in y + * @param sz spin component in z */ void AddNParticles ( @@ -190,7 +196,10 @@ namespace impactx amrex::Gpu::DeviceVector const & pt, amrex::ParticleReal qm, std::optional bunch_charge = std::nullopt, - std::optional> w = std::nullopt + std::optional> w = std::nullopt, + std::optional> sx = std::nullopt, + std::optional> sy = std::nullopt, + std::optional> sz = std::nullopt ); /** Register storage for lost particles diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp index bb6ce3e4a..9a6cbe955 100644 --- a/src/particles/ImpactXParticleContainer.cpp +++ b/src/particles/ImpactXParticleContainer.cpp @@ -242,28 +242,40 @@ namespace impactx amrex::Gpu::DeviceVector const & pt, amrex::ParticleReal qm, std::optional bunch_charge, - std::optional> w + std::optional> w, + std::optional> sx, + std::optional> sy, + std::optional> sz ) { BL_PROFILE("ImpactX::AddNParticles"); using namespace amrex::literals; // for _rt and _prt + // number of particles to add + std::size_t const np_s = x.size(); + + // input validation bool const has_w = w.has_value(); if (!(bunch_charge.has_value() ^ has_w)) { throw std::runtime_error("AddNParticles: Exactly one of bunch_charge or w must be provided!"); } - AMREX_ALWAYS_ASSERT(x.size() == y.size()); - AMREX_ALWAYS_ASSERT(x.size() == t.size()); - AMREX_ALWAYS_ASSERT(x.size() == px.size()); - AMREX_ALWAYS_ASSERT(x.size() == py.size()); - AMREX_ALWAYS_ASSERT(x.size() == pt.size()); - if (has_w) { AMREX_ALWAYS_ASSERT(x.size() == w->size()); } - - // number of particles to add - amrex::Long const np = x.size(); + bool const has_spin = sx.has_value(); + + AMREX_ALWAYS_ASSERT(np_s == y.size()); + AMREX_ALWAYS_ASSERT(np_s == t.size()); + AMREX_ALWAYS_ASSERT(np_s == px.size()); + AMREX_ALWAYS_ASSERT(np_s == py.size()); + AMREX_ALWAYS_ASSERT(np_s == pt.size()); + if (has_spin) { + AMREX_ALWAYS_ASSERT(sy.has_value()); + AMREX_ALWAYS_ASSERT(sz.has_value()); + AMREX_ALWAYS_ASSERT(np_s == sy.value().size()); + AMREX_ALWAYS_ASSERT(np_s == sz.value().size()); + } + if (has_w) { AMREX_ALWAYS_ASSERT(np_s == w->size()); } // we add particles to lev 0, grid 0 int lid = 0, gid = 0; @@ -288,7 +300,8 @@ namespace impactx DefineAndReturnParticleTile(lid, gid, ithr); } - amrex::Long pid = ParticleType::NextID(); + amrex::Long const pid = ParticleType::NextID(); + amrex::Long const np = np_s; ParticleType::NextID(pid + np); AMREX_ALWAYS_ASSERT_WITH_MESSAGE( pid + np < amrex::LongParticleIds::LastParticleID, @@ -331,6 +344,9 @@ namespace impactx amrex::ParticleReal * const AMREX_RESTRICT px_arr = soa[RealSoA::px].dataPtr(); amrex::ParticleReal * const AMREX_RESTRICT py_arr = soa[RealSoA::py].dataPtr(); amrex::ParticleReal * const AMREX_RESTRICT pt_arr = soa[RealSoA::pt].dataPtr(); + amrex::ParticleReal * const AMREX_RESTRICT sx_arr = soa[RealSoA::sx].dataPtr(); + amrex::ParticleReal * const AMREX_RESTRICT sy_arr = soa[RealSoA::sy].dataPtr(); + amrex::ParticleReal * const AMREX_RESTRICT sz_arr = soa[RealSoA::sz].dataPtr(); amrex::ParticleReal * const AMREX_RESTRICT qm_arr = soa[RealSoA::qm].dataPtr(); amrex::ParticleReal * const AMREX_RESTRICT w_arr = soa[RealSoA::w ].dataPtr(); @@ -342,6 +358,9 @@ namespace impactx amrex::ParticleReal const * const AMREX_RESTRICT px_ptr = px.data(); amrex::ParticleReal const * const AMREX_RESTRICT py_ptr = py.data(); amrex::ParticleReal const * const AMREX_RESTRICT pt_ptr = pt.data(); + amrex::ParticleReal const * const AMREX_RESTRICT sx_ptr = has_spin ? sx.value().data() : nullptr; + amrex::ParticleReal const * const AMREX_RESTRICT sy_ptr = has_spin ? sy.value().data() : nullptr; + amrex::ParticleReal const * const AMREX_RESTRICT sz_ptr = has_spin ? sz.value().data() : nullptr; amrex::ParticleReal const * const AMREX_RESTRICT w_ptr = has_w ? w->data() : nullptr; amrex::ParticleReal const bunch_charge_value = has_w ? 0_prt : bunch_charge.value(); @@ -358,6 +377,16 @@ namespace impactx py_arr[old_np+i] = py_ptr[my_offset+i]; pt_arr[old_np+i] = pt_ptr[my_offset+i]; + if (has_spin) { + sx_arr[old_np+i] = sx_ptr[my_offset+i]; + sy_arr[old_np+i] = sy_ptr[my_offset+i]; + sz_arr[old_np+i] = sz_ptr[my_offset+i]; + } else { + sx_arr[old_np+i] = 0_prt; + sy_arr[old_np+i] = 0_prt; + sz_arr[old_np+i] = 0_prt; + } + qm_arr[old_np+i] = qm; w_arr[old_np+i] = has_w ? w_ptr[my_offset+i] : bunch_charge_value / ablastr::constant::SI::q_e/np; }); diff --git a/src/particles/ReferenceParticle.H b/src/particles/ReferenceParticle.H index 3d0c7bd46..f4e0536ed 100644 --- a/src/particles/ReferenceParticle.H +++ b/src/particles/ReferenceParticle.H @@ -128,7 +128,7 @@ namespace impactx { using namespace amrex::literals; - constexpr amrex::ParticleReal inv_MeV_invc2 = 1.0_prt / ablastr::constant::SI::MeV_invc2; + constexpr amrex::ParticleReal inv_MeV_invc2 = 1.0_prt / ablastr::constant::SI::MeV_inv_c2_v; return amrex::ParticleReal(mass * inv_MeV_invc2); } @@ -146,7 +146,7 @@ namespace impactx AMREX_ASSERT_WITH_MESSAGE(massE != 0.0_prt, "set_mass_MeV: Mass cannot be zero!"); - mass = massE * ablastr::constant::SI::MeV_invc2; + mass = massE * ablastr::constant::SI::MeV_inv_c2_v; // re-scale pt and pz if (pt != 0.0_prt) diff --git a/src/particles/distribution/All.H b/src/particles/distribution/All.H index 09606cec6..5e676fef6 100644 --- a/src/particles/distribution/All.H +++ b/src/particles/distribution/All.H @@ -19,12 +19,14 @@ #include "Thermal.H" #include "Triangle.H" #include "Waterbag.H" +#include "SpinvMF.H" #include namespace impactx::distribution { + // Phase space distributions using KnownDistributions = std::variant< Empty, /* must be first, so KnownDistributions creates a default constructor */ Gaussian, diff --git a/src/particles/distribution/SpinvMF.H b/src/particles/distribution/SpinvMF.H new file mode 100644 index 000000000..da074941d --- /dev/null +++ b/src/particles/distribution/SpinvMF.H @@ -0,0 +1,157 @@ +/* Copyright 2022-2023 The Regents of the University of California, through Lawrence + * Berkeley National Laboratory (subject to receipt of any required + * approvals from the U.S. Dept. of Energy). All rights reserved. + * + * This file is part of ImpactX. + * + * Authors: Chad Mitchell, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_DISTRIBUTION_SPINVMF +#define IMPACTX_DISTRIBUTION_SPINVMF + +#include + +#include +#include +#include + +#include + + +namespace impactx::distribution +{ + struct SpinvMF + { + /** A von Mises-Fisher (vMF) distribution on the unit 2-sphere. + * + * This is used for initializing particle spin. There is a natural + * bijective correspondence between vMF distributions and mean + * (polarization) vectors. + * + * The algorithm used here is a simplification of the algorithm described in: + * C. Pinzon and K. Jung, "Fast Python sampler of the von Mises Fisher distribution", in the special + * case of the 2-sphere. Additional references used include: + * + * K. V. Mardia and P. E. Jupp, Directional Statistics, Wiley, 1999; + * S. Kang and H-S. Oh, "Novel sampling method for the von Mises-Fisher distribution", Stat. + * and Comput. 34, 106 (2024), https://doi.org/10.1007/s11222-024-10419-3. + * + * Return sampling from a vMF distribution. + * + * @param mux,muy,muz components of a unit vector specifying the mean direction + */ + SpinvMF ( + amrex::ParticleReal mux, + amrex::ParticleReal muy, + amrex::ParticleReal muz + ) + : m_muX(mux), m_muY(muy), m_muZ(muz) + { + // magnitude of the polarization vector (= polarization magnitude) + amrex::ParticleReal const pmag = std::sqrt( + mux * mux + + muy * muy + + muz * muz + ); + + m_kappa = inverse_Langevin(pmag); + } + + /** Return 1 3D spin (unit) vector + * + * @param sx particle spin component in x + * @param sy particle spin component in y + * @param sz particle spin component in z + * @param engine a random number engine (with associated state) + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + amrex::ParticleReal & AMREX_RESTRICT sx, + amrex::ParticleReal & AMREX_RESTRICT sy, + amrex::ParticleReal & AMREX_RESTRICT sz, + amrex::RandomEngine const & engine + ) const + { + using namespace amrex::literals; + using amrex::Math::powi; + using ablastr::constant::math::pi; + + // Normalize the mean direction vector just in case: + amrex::ParticleReal mu_norm = std::sqrt(powi<2>(m_muX)+powi<2>(m_muY)+powi<2>(m_muZ)); + amrex::ParticleReal mux = m_muX/mu_norm; + amrex::ParticleReal muy = m_muY/mu_norm; + amrex::ParticleReal muz = m_muZ/mu_norm; + + // Evaluate constants independent of particle: + amrex::ParticleReal const muperp = std::sqrt(powi<2>(mux)+powi<2>(muy)); + amrex::ParticleReal const a = (muperp==0) ? 0_prt : mux/muperp; + amrex::ParticleReal const b = (muperp==0) ? 1_prt : muy/muperp; + + // Sample two iid random variables: + amrex::ParticleReal x1 = amrex::Random(engine); + amrex::ParticleReal x2 = amrex::Random(engine); + + // Basic transformation of the random variables: + amrex::ParticleReal c = std::cos(2_prt*pi*x1); + amrex::ParticleReal s = std::sin(2_prt*pi*x1); + amrex::ParticleReal t; + // The form of the function below is obtained from eq (30) of: + // D. Frisch and U. D. Hanebeck, "Deterministic Von Mises-Fisher Sampling on the Sphere Using + // Fibonacci Lattices", IEEE, 2023, DOI: 10.1109/SDF-MFI59545.2023.10361396 + if(m_kappa > 0) { + t = 1_prt + std::log1p(x2 * std::expm1(-2_prt * m_kappa)) / m_kappa; + } else { + t = 1_prt - 2_prt * x2; + } + + // Evaluation of the spin components + amrex::ParticleReal u = std::sqrt(1_prt-powi<2>(t)); + sx = u * (b*c + a*muz*s) + t*mux; + sy = u * (-a*c + b*muz*s) + t*muy; + sz = u * (-muperp*s) + t*muz; + + } + + + /** This function evaluates the inverse Langevin function, in order to return + * the value of concentration (kappa) required to produce a given polarization magnitude. + * It is used to prepare for spin sampling using the vMF distribution. + * + * Here, we use an approximate expression due to R. Petrosyan. See: + * R. Petrosyan, Rheologica Acta 56, 21-26, 2016; doi:10.1007/s00397-016-0977-9 + * R. Jedynak, Mathematics and Mechanics of Solids 24, 1-25, 2018; doi:10.1177/1081286518811395 + * + * @param pmag polarization magnitude + * @returns required value of concentration kappa + */ + AMREX_GPU_HOST AMREX_FORCE_INLINE + static amrex::ParticleReal + inverse_Langevin (amrex::ParticleReal pmag) + { + using namespace amrex::literals; // for _rt and _prt + using amrex::Math::powi; + + AMREX_ASSERT_WITH_MESSAGE(pmag >= 0_prt, "Polarization magnitude must be greater or equal 0."); + AMREX_ASSERT_WITH_MESSAGE(pmag < 1_prt, "Polarization magnitude must be less than 1."); + + // Evaluation of the inverse Langevin function. Here, we use the approximate expression + // due to R. Petrosyan. See: + // R. Petrosyan, Rheologica Acta 56, 21-26, 2016; doi:10.1007/s00397-016-0977-9 + // R. Jedynak, Mathematics and Mechanics of Solids 24, 1-25, 2018; doi:10.1177/1081286518811395 + + amrex::ParticleReal concentration = + 3_prt * pmag + + powi<2>(pmag) / 5_prt * std::sin(7_prt * pmag * 0.5_prt) + + powi<3>(pmag) / (1_prt - pmag); + + return concentration; + } + + amrex::ParticleReal m_muX, m_muY, m_muZ; //! components of a unit vector specifying the mean direction + amrex::ParticleReal m_kappa; //! concentration parameter + }; + +} // namespace impactx::distribution + +#endif // IMPACTX_DISTRIBUTION_SPINVMF diff --git a/src/particles/distribution/Thermal.H b/src/particles/distribution/Thermal.H index 17d655c10..8d2a2936c 100644 --- a/src/particles/distribution/Thermal.H +++ b/src/particles/distribution/Thermal.H @@ -89,6 +89,7 @@ namespace distribution using namespace amrex::literals; // for _rt and _prt using namespace ablastr::constant::math; // for pi using amrex::Math::powi; + constexpr amrex::ParticleReal ep0 = ablastr::constant::SI::epsilon_0_v; // Get relativistic beta*gamma amrex::ParticleReal bg = refpart.beta_gamma(); @@ -99,7 +100,7 @@ namespace distribution amrex::ParticleReal q_e = refpart.charge_qe(); // Set space charge intensity - m_Cintensity = q_e*bunch_charge/(powi<2>(bg)*Erest*ablastr::constant::SI::ep0); + m_Cintensity = q_e*bunch_charge/(powi<2>(bg)*Erest*ep0); // Set minimum and maximum radius amrex::ParticleReal r_scale = matched_scale_radius(); diff --git a/src/particles/spacecharge/PoissonSolve.cpp b/src/particles/spacecharge/PoissonSolve.cpp index 4d367e3d6..367162bf9 100644 --- a/src/particles/spacecharge/PoissonSolve.cpp +++ b/src/particles/spacecharge/PoissonSolve.cpp @@ -150,7 +150,7 @@ namespace impactx::particles::spacecharge // fix side effect on rho from previous call for (int lev=0; lev<=finest_level; lev++) { using namespace ablastr::constant::SI; - rho[lev].mult(-1._rt * ep0); + rho[lev].mult(-1._rt * epsilon_0); } // We may need to copy phi from phi_2d diff --git a/src/particles/transformation/ToFixedT.H b/src/particles/transformation/ToFixedT.H index 8963a1235..ca649b761 100644 --- a/src/particles/transformation/ToFixedT.H +++ b/src/particles/transformation/ToFixedT.H @@ -60,7 +60,12 @@ namespace impactx::transformation using amrex::Math::powi; // small tolerance to avoid NaN for pz<0: - constexpr amrex::ParticleReal tol = 1.0e-8_prt; + constexpr amrex::ParticleReal tol = +#ifdef AMREX_SINGLE_PRECISION_PARTICLES + 1.0e-6_prt; +#else + 1.0e-12_prt; +#endif // compute value of reference pzd = beta*gamma amrex::ParticleReal const argd = -1.0_prt + powi<2>(m_ptd); diff --git a/src/particles/wakefields/WakeConvolution.H b/src/particles/wakefields/WakeConvolution.H index 20ecae3bb..b70d91d14 100644 --- a/src/particles/wakefields/WakeConvolution.H +++ b/src/particles/wakefields/WakeConvolution.H @@ -112,7 +112,7 @@ namespace impactx::particles::wakefields using namespace amrex::literals; using amrex::Math::powi; - amrex::Real const rc = powi<2>(ablastr::constant::SI::q_e) / (4_rt * amrex::Real(M_PI) * ablastr::constant::SI::ep0 * ablastr::constant::SI::m_e * powi<2>(ablastr::constant::SI::c)); + amrex::Real const rc = powi<2>(ablastr::constant::SI::q_e) / (4_rt * ablastr::constant::math::pi * ablastr::constant::SI::epsilon_0 * ablastr::constant::SI::m_e * powi<2>(ablastr::constant::SI::c)); amrex::Real const kappa = (2_rt * rc * ablastr::constant::SI::m_e * powi<2>(ablastr::constant::SI::c)) / std::pow(3_rt, 1_rt/3_rt) / std::pow(R, 2_rt/3_rt); return -3_rt/2_rt * kappa / bin_size * (unit_step(s + bin_size / 2_rt) * std::pow(std::abs(s + bin_size / 2_rt), 2_rt/3_rt) - unit_step(s - bin_size / 2_rt) * std::pow(std::abs(s - bin_size / 2_rt), 2_rt/3_rt)); diff --git a/src/python/ImpactX.cpp b/src/python/ImpactX.cpp index 339bcd89a..7add1f6da 100644 --- a/src/python/ImpactX.cpp +++ b/src/python/ImpactX.cpp @@ -285,7 +285,11 @@ void init_ImpactX (py::module& m) amrex::ParmParse pp_algo("algo"); if (std::get(space_charge_v)) { // TODO: boolean True is deprecated since 25.03, remove some time after - py::print("sim.space_charge = True is deprecated, please use space_charge = \"3D\""); + py::warnings::warn( + "sim.space_charge = True is deprecated, please use space_charge = \"3D\"", + PyExc_DeprecationWarning, + 2 + ); pp_algo.add("space_charge", std::string("3D")); } else { // map boolean False to "false" / off @@ -617,6 +621,7 @@ void init_ImpactX (py::module& m) .def("add_particles", &ImpactX::add_particles, py::arg("bunch_charge"), py::arg("distr"), py::arg("npart"), + py::arg("spin_distr") = py::none(), "Particle tracking mode:" "Generate and add n particles to the particle container.\n\n" "Will also resize the geometry based on the updated particle\n" @@ -626,7 +631,12 @@ void init_ImpactX (py::module& m) .def("evolve", /** TODO: deprecated API. Only for internal use. Remove after a few releases. */ [](ImpactX & ix) { - py::print("Warning: evolve() is deprecated and will soon be removed. Use track_particles() instead."); + py::warnings::warn( + "Warning: evolve() is deprecated and will soon be removed. " + "Use track_particles() instead.", + PyExc_DeprecationWarning, + 2 + ); ix.evolve(); }, "Run the main simulation loop." diff --git a/src/python/ImpactXParticleContainer.cpp b/src/python/ImpactXParticleContainer.cpp index 93f097886..8ff439627 100644 --- a/src/python/ImpactXParticleContainer.cpp +++ b/src/python/ImpactXParticleContainer.cpp @@ -65,12 +65,12 @@ void init_impactxparticlecontainer(py::module& m) // simpler particle iterator loops: return types of this particle box // note: overwritten to return ImpactX instead of (py)AMReX iterators .def_property_readonly_static( - "iterator", + "Iterator", [](py::object /* pc */){ return py::type::of(); }, "ImpactX iterator for particle boxes" ) .def_property_readonly_static( - "const_iterator", + "ConstIterator", [](py::object /* pc */){ return py::type::of(); }, "ImpactX constant iterator for particle boxes (read-only)" ) @@ -84,6 +84,7 @@ void init_impactxparticlecontainer(py::module& m) py::arg("x"), py::arg("y"), py::arg("t"), py::arg("px"), py::arg("py"), py::arg("pt"), py::arg("qm"), py::arg("bunch_charge")=py::none(), py::arg("w")=py::none(), + py::arg("sx")=py::none(), py::arg("sy")=py::none(), py::arg("sz")=py::none(), "Add new particles to the container for fixed s.\n\n" "Either the total charge (bunch_charge) or the weight of each\n" "particle (w) must be provided.\n\n" @@ -99,6 +100,9 @@ void init_impactxparticlecontainer(py::module& m) ":param qm: charge over mass in 1/eV\n" ":param bunch_charge: total charge within a bunch in C" ":param w: weight of each particle: how many real particles to represent" + ":param sx: spin component in x\n" + ":param sy: spin component in y\n" + ":param sz: spin component in z\n" ) .def("ref_particle", py::overload_cast<>(&ImpactXParticleContainer::GetRefParticle), @@ -122,11 +126,11 @@ void init_impactxparticlecontainer(py::module& m) ) .def("reduced_beam_characteristics", [](ImpactXParticleContainer & pc) { - ablastr::warn_manager::WMRecordWarning( - "reduced_beam_characteristics", + py::warnings::warn( "WARNING: reduced_beam_characteristics() is deprecated. " "Use beam_moments() instead.", - ablastr::warn_manager::WarnPriority::medium + PyExc_DeprecationWarning, + 2 ); return diagnostics::reduced_beam_characteristics(pc); }, diff --git a/src/python/distribution.cpp b/src/python/distribution.cpp index a074f6805..7f3ac79f1 100644 --- a/src/python/distribution.cpp +++ b/src/python/distribution.cpp @@ -187,6 +187,21 @@ void init_distribution(py::module& m) "A 6D Waterbag distribution" ); + py::class_(md, "SpinvMF") + .def(py::init< + amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal + >(), + py::arg("mux"), py::arg("muy"), py::arg("muz"), + "A von Mises-Fisher (vMF) distribution on the unit 2-sphere, for particle spin." + ) + .def_static("inverse_Langevin", + &distribution::SpinvMF::inverse_Langevin, + py::arg("pmag"), + "This function evaluates the inverse Langevin function, in order to return " + "the value of concentration (kappa) required to produce a given polarization magnitude." + ) + ; + py::class_(m, "Envelope") .def(py::init<>()) .def(py::init()) diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 7e0dff7e4..5967ac3dc 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -1,6 +1,6 @@ /* Copyright 2021-2023 The ImpactX Community * - * Authors: Axel Huebl + * Authors: Axel Huebl, Eric G. Stern * License: BSD-3-Clause-LBNL */ #include "pyImpactX.H" @@ -1611,6 +1611,108 @@ void init_elements(py::module& m) ; register_push(py_PlaneXYRot); + py::class_ py_PolygonAperture(me, "PolygonAperture"); + py_PolygonAperture + .def("__repr__", + [](PolygonAperture const & polygon_aperture) { + return element_name( + polygon_aperture, + std::make_pair("action", polygon_aperture.action_name(polygon_aperture.m_action)) + ); + } + ) + .def("to_dict", + [](PolygonAperture const & polygon_aperture) { + using namespace amrex::literals; + return element_dict( + polygon_aperture, + std::make_pair("vertices_x", PolygonApertureData::h_vertices_x[polygon_aperture.m_id]), + std::make_pair("vertices_y", PolygonApertureData::h_vertices_y[polygon_aperture.m_id]), + std::make_pair("min_radius2", polygon_aperture.m_min_radius2), + std::make_pair("action", polygon_aperture.m_action), + std::make_pair("repeat_x", polygon_aperture.m_repeat_x), + std::make_pair("repeat_y", polygon_aperture.m_repeat_y), + std::make_pair("shift_odd_x", polygon_aperture.m_shift_odd_x) + ); + } + ) + .def(py::init([]( + std::vector vertices_x, + std::vector vertices_y, + amrex::ParticleReal min_radius2, + amrex::ParticleReal repeat_x, + amrex::ParticleReal repeat_y, + bool shift_odd_x, + std::string const & action, + amrex::ParticleReal dx, + amrex::ParticleReal dy, + amrex::ParticleReal rotation_degree, + std::optional name + ) + { + PolygonAperture::Action pa_action; + if (action == "transmit") { + pa_action = PolygonAperture::Action::transmit; + } else if (action == "absorb") { + pa_action = PolygonAperture::Action::absorb; + } else + throw std::runtime_error(R"(action must be "transmit" or "absorb")"); + + return new PolygonAperture(vertices_x, vertices_y, min_radius2, repeat_x, repeat_y, shift_odd_x, pa_action, dx, dy, rotation_degree, name); + }), + py::arg("vertices_x"), + py::arg("vertices_y"), + py::arg("min_radius2")=0.0, + py::arg("repeat_x") = 0, + py::arg("repeat_y") = 0, + py::arg("shift_odd_x") = false, + py::arg("action") = "transmit", + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("name") = py::none(), + "A short collimator element described by a polygon with vertices given by their x and y coordinates." + ) + .def_property("action", + [](PolygonAperture & ap) + { + return ap.action_name(ap.m_action); + }, + [](PolygonAperture & ap, std::string const & action) + { + if (action != "transmit" && action != "absorb") + throw std::runtime_error(R"(action must be "transmit" or "absorb")"); + + ap.m_action = action == "transmit" ? + PolygonAperture::Action::transmit : + PolygonAperture::Action::absorb; + }, + "action type (transmit, absorb)" + ) + .def_property("min_radius2", + [](PolygonAperture & pa) { return pa.m_min_radius2; }, + [](PolygonAperture & pa, amrex::ParticleReal mr2) { pa.m_min_radius2 = mr2; }, + "All particles with radius squared smaller than min_radius2 pass the aperture" + ) + .def_property("repeat_x", + [](PolygonAperture & ap) { return ap.m_repeat_x; }, + [](PolygonAperture & ap, amrex::ParticleReal repeat_x) { ap.m_repeat_x = repeat_x; }, + "horizontal period for repeated aperture masking" + ) + .def_property("repeat_y", + [](PolygonAperture & ap) { return ap.m_repeat_y; }, + [](PolygonAperture & ap, amrex::ParticleReal repeat_y) { ap.m_repeat_y = repeat_y; }, + "vertical period for repeated aperture masking" + ) + .def_property("shift_odd_x", + [](PolygonAperture & ap) { return ap.m_shift_odd_x; }, + [](PolygonAperture & ap, bool shift_odd_x) { ap.m_shift_odd_x = shift_odd_x; }, + "for hexagonal/triangular mask patterns: horizontal shift of every 2nd (odd) vertical period by repeat_x / 2. " + "Use alignment offsets dx,dy to move whole mask as needed." + ) + ; + register_push(py_PolygonAperture); + py::class_(me, "Programmable", py::dynamic_attr()) .def("__repr__", [](Programmable const & prg) { @@ -2088,7 +2190,8 @@ void init_elements(py::module& m) return element_name( src, std::make_pair("distribution", src.m_distribution), - std::make_pair("path", src.m_series_name) + std::make_pair("path", src.m_series_name), + std::make_pair("active_once", src.m_active_once) ); } ) @@ -2097,17 +2200,20 @@ void init_elements(py::module& m) return element_dict( src, std::make_pair("distribution", src.m_distribution), - std::make_pair("path", src.m_series_name) + std::make_pair("path", src.m_series_name), + std::make_pair("active_once", src.m_active_once) ); } ) .def(py::init< std::string, std::string, + bool, std::optional >(), py::arg("distribution"), py::arg("openpmd_path"), + py::arg("active_once") = true, py::arg("name") = py::none(), "A particle source." ) @@ -2121,6 +2227,11 @@ void init_elements(py::module& m) [](Source & src, std::string series_name) { src.m_series_name = series_name; }, "Path to openPMD series as accepted by openPMD_api.Series" ) + .def_property("active_once", + [](Source & src) { return src.m_active_once; }, + [](Source & src, bool actice_once) { src.m_active_once = actice_once; }, + "Inject particles only for the first lattice period." + ) ; register_push(py_Source); diff --git a/src/python/impactx/__init__.pyi b/src/python/impactx/__init__.pyi index 6e454177f..5efa6912e 100644 --- a/src/python/impactx/__init__.pyi +++ b/src/python/impactx/__init__.pyi @@ -95,6 +95,6 @@ __author__: str = ( "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." ) __license__: str = "BSD-3-Clause-LBNL" -__version__: str = "25.11" +__version__: str = "25.12" s: impactx_pybind.CoordSystem # value = t: impactx_pybind.CoordSystem # value = diff --git a/src/python/impactx/dashboard/Input/visualization/lattice/statistics.py b/src/python/impactx/dashboard/Input/visualization/lattice/statistics.py index 7735b3566..ca20af311 100644 --- a/src/python/impactx/dashboard/Input/visualization/lattice/statistics.py +++ b/src/python/impactx/dashboard/Input/visualization/lattice/statistics.py @@ -69,9 +69,10 @@ def update_length_statistics() -> None: @staticmethod def update_element_counts() -> dict[str, int]: """ - Computes the element counts in the lattice list. + Computes the element counts in the lattice list + and stores them in descending order by count. - :return: Dictionary of element names and their counts, sorted by count descending. + :return: Dictionary of element counts indexed by element name. """ counts = {} for element in state.selected_lattice_list: diff --git a/src/python/impactx/dashboard/Toolbar/examples_loader/loader.py b/src/python/impactx/dashboard/Toolbar/examples_loader/loader.py index 8344fd69c..969933cbb 100644 --- a/src/python/impactx/dashboard/Toolbar/examples_loader/loader.py +++ b/src/python/impactx/dashboard/Toolbar/examples_loader/loader.py @@ -16,15 +16,15 @@ DASHBOARD_EXAMPLES = { "fodo/run_fodo.py", - # "chicane/run_chicane_csr.py", - # "fodo_space_charge/run_fodo_envelope_sc.py", - # "apochromatic/run_apochromatic.py", - # "kurth/run_kurth_10nC_periodic.py", - # "expanding_beam/run_expanding_fft.py", - # "expanding_beam/run_expanding_envelope.py", - # "iota_lattice/run_iotalattice.py", - # "cyclotron/run_cyclotron.py", - # "dogleg/run_dogleg.py", + "chicane/run_chicane_csr.py", + "fodo_space_charge/run_fodo_envelope_sc.py", + "apochromatic/run_apochromatic.py", + # "kurth/run_kurth_10nC_periodic.py", - running into recursion issues + "expanding_beam/run_expanding_fft.py", + "expanding_beam/run_expanding_envelope.py", + "iota_lattice/run_iotalattice.py", + "cyclotron/run_cyclotron.py", + "dogleg/run_dogleg.py", } @@ -64,7 +64,7 @@ def _get_example_content(file_name: str) -> dict: Retrieve the selected ImpactX example file and populate the UI with its values. """ - impactx_directory = DashboardExamplesLoader.get_impactx_path() + impactx_directory = GeneralFunctions.get_impactx_root_dir() impactx_example_file_path = impactx_directory / "examples" / file_name file_content_as_str = impactx_example_file_path.read_text() diff --git a/src/python/impactx/dashboard/Toolbar/file_imports/python/__init__.py b/src/python/impactx/dashboard/Toolbar/file_imports/python/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/python/impactx/dashboard/Toolbar/file_imports/python/helper.py b/src/python/impactx/dashboard/Toolbar/file_imports/python/helper.py index 24acc39d4..3e61f7959 100644 --- a/src/python/impactx/dashboard/Toolbar/file_imports/python/helper.py +++ b/src/python/impactx/dashboard/Toolbar/file_imports/python/helper.py @@ -53,6 +53,10 @@ def parse_single_inputs(content: str) -> dict: pattern_match = re.search(pattern.format(parameter_name), content) if pattern_match: value = ast.literal_eval(pattern_match.group(1)) + + if parameter_name == "space_charge" and isinstance(value, bool): + value = str(value).lower() + reference_dictionary[parameter_name] = value break @@ -61,8 +65,14 @@ def parse_single_inputs(content: str) -> dict: r"\bkin_energy_MeV\s*=\s*([^#\n]+)", content ) if kin_energy_pattern_match: - kin_energy_value = kin_energy_pattern_match.group(1) - reference_dictionary["kin_energy_on_ui"] = kin_energy_value + kin_energy_value = kin_energy_pattern_match.group(1).strip() + + try: + parsed_value = ast.literal_eval(kin_energy_value) + except (ValueError, SyntaxError): + parsed_value = reference_dictionary.get("kin_energy_MeV") + + reference_dictionary["kin_energy_on_ui"] = parsed_value return reference_dictionary @@ -71,11 +81,24 @@ def parse_list_inputs(content: str) -> dict: """ Parses list-based simulation parameters from the simulation file content. + Some of the dashboard inputs are lists, such as n_cell = [16, 16, 16]. + We have to look inside of the brackets to retrieve the individual values. + :param content: The content of the ImpactX simulation file. """ dictionary = {} - list_inputs = ["n_cell", "prob_relative"] - list_parsing = "{} = (\\[.*?\\])" + list_inputs = [ + "n_cell", + "prob_relative", + "blocking_factor_x", + "blocking_factor_y", + "blocking_factor_z", + ] + + # The below regex will capture everything inside of the brackets + # it assumes that the list will be placed in a single line (not multiple). + # Example: n_cell = [16, 16, 16] --> [16, 16, 16] + list_parsing = r"{}\s*=\s*(\[[^\]]*\])" for input_name in list_inputs: match = re.search(list_parsing.format(input_name), content) @@ -86,9 +109,13 @@ def parse_list_inputs(content: str) -> dict: for i, dim in enumerate(["x", "y", "z"]): dictionary[f"n_cell_{dim}"] = values[i] - if input_name == "prob_relative": + elif input_name == "prob_relative": dictionary["prob_relative"] = values + elif input_name.startswith("blocking_factor_"): + # Convert [16] -> 16 + dictionary[input_name] = values[0] + return dictionary @staticmethod @@ -134,33 +161,6 @@ def extract_parameters(distribution_type, parsing_pattern): return dictionary - @staticmethod - def parse_lattice_elements(content: str) -> dict: - """ - Parses lattice elements from the simulation file content. - - :param content: The content of the ImpactX simulation file. - """ - - dictionary = {"lattice_elements": []} - used_variables = set() - - lattice_elements = re.findall(r"elements\.(\w+)\((.*?)\)", content) - - for element_name, element_parameter in lattice_elements: - element = {"element": element_name, "parameters": {}} - - parameter_pairs = re.findall(r"(\w+)=([^,\)]+)", element_parameter) - for parameter_name, parameter_value in parameter_pairs: - parameter_value_cleaned = parameter_value.strip("'\"") - element["parameters"][parameter_name] = parameter_value_cleaned - used_variables.add(parameter_value_cleaned) - - dictionary["lattice_elements"].append(element) - - dictionary["used_lattice_variables"] = used_variables - return dictionary - @staticmethod def parse_variables(content: str, used_vars: set) -> dict: """ diff --git a/src/python/impactx/dashboard/Toolbar/file_imports/python/lattice_helper.py b/src/python/impactx/dashboard/Toolbar/file_imports/python/lattice_helper.py new file mode 100644 index 000000000..27d341270 --- /dev/null +++ b/src/python/impactx/dashboard/Toolbar/file_imports/python/lattice_helper.py @@ -0,0 +1,323 @@ +import re +from typing import Any, Dict, List + + +class DashboardLatticeParser: + """ + Helper class to parse the lattice configuration from Impactx .py-compatible simulation files. + """ + + def __init__(self, content: str): + self._flatten_cache: Dict[str, List[str]] = {} + self._variable_assignments_cache: Dict[str, Dict[str, str]] = {} + self._content = content + self._content_hash: str = str(hash(content)) + + def parse(self) -> Dict[str, Any]: + """ + Extracts the lattice configuration from provided ImpactX simulation file. + + Example return: + { + "lattice_elements": [ + {"name": "Drift", "parameters": {"ds": "0.5"}}, + {"name": "Quad", "parameters": {"k": "quad_strength", "ds": "0.3"}} + ] + } + + :return: Parsed dictionary containing 'lattice_elements'. + """ + lattice_order = self.collect_lattice_operations(debug=False) + + expanded_elements = [] + for operation in lattice_order: + operation_type = operation["type"] + operation_arg = operation["argument"] + + match operation_type: + case "extend": + expanded_elements.extend(self._flatten(operation_arg)) + case "append": + expanded_elements.append(operation_arg) + case "reverse": + self._apply_reverse(operation_arg) + case _: + print(f"Warning: Unsupported operation type: {operation_type}") + + clean_lattice_list = self.replace_variables(expanded_elements) + clean_lattice_list_str = "\n".join(clean_lattice_list) + result = self.parse_cleaned_lattice(clean_lattice_list_str) + + return result + + def parse_cleaned_lattice(self, content: str) -> Dict[str, List[Dict[str, Any]]]: + """ + Parses the lattice elements from the ImpactX simulation file content. + + Extracts element names and their parameters from constructor calls in the format: + elements.ElementName(param1=value1, param2=value2, ...) + + EX: + elements.Drift(ds=1.0) + + Results in: + { + "lattice_elements": [ + { + "name": "Drift", + "parameters": {"ds": "1.0"} + } + ] + } + + :param content: The content of the ImpactX simulation file. + :return: A dictionary containing the parsed lattice elements. + """ + + dictionary = {"lattice_elements": []} + + element_pattern = r"elements\.(\w+)\((.*?)\)" # EX: elements.Drift(...) + lattice_elements = re.findall(element_pattern, content, re.DOTALL) + + for element_name, parameter in lattice_elements: + element = {"name": element_name, "parameters": {}} + + # CHANGE: Updated parameter pattern to handle multiline and whitespace around = + # OLD: r"(\w+)=([^,\)]+)" + # NEW: r"(\w+)\s*=\s*([^,\)\n]+)" with re.MULTILINE flag + parameter_pattern = r"(\w+)\s*=\s*([^,\)\n]+)" # EX: ds=1.0, k1=0.5 + all_parameters = re.findall(parameter_pattern, parameter) + + for parameter_name, value in all_parameters: + element["parameters"][parameter_name] = value.strip("'\"") + + dictionary["lattice_elements"].append(element) + + return dictionary + + def collect_lattice_operations(self, debug: bool = False) -> List[Dict[str, str]]: + """ + Collects lattice operations (sim.lattice.append(), sim.lattice.extend(), and variable.reverse() calls) + in the order they appear in the content. + + EX: + sim.lattice.append(monitor) ; sim.lattice.extend([drift1, quad1]) ; lattice_half.reverse() + + Results the following (in order): + [ + {"type": "append", "argument": "monitor"}, + {"type": "extend", "argument": "[drift1, quad1]"} + {"type": "reverse", "argument": "lattice_half"} + ] + + :param debug: Whether to print the collected operations. + :return: List of dictionaries with 'type' and 'argument' keys. + """ + operations = [] + + # Captures the operation type (append or extend) and its argument + # (everything inside the parentheses, up to the last closing parenthesis on the line). + lattice_call_pattern = r"sim\.lattice\.(append|extend)\((.*)\)" + + # Store sim.lattice.append and sim.lattice.extend calls + for match in re.finditer(lattice_call_pattern, self._content): + operation, arg = match.groups() + operations.append( + (match.start(), {"type": operation, "argument": arg.strip()}) + ) + + # Store .reverse() calls + reverse_pattern = r"(\w+)\.reverse\(\)" + for match in re.finditer(reverse_pattern, self._content): + operations.append( + (match.start(), {"type": "reverse", "argument": match.group(1)}) + ) + + # important: sort operations by their position in the content + # since the for loops can be executed in any order + operations = [type for _, type in sorted(operations, key=lambda x: x[0])] + + if debug: + self.print_lattice_operations(operations) + + return operations + + def _get_variable_assignments(self) -> Dict[str, str]: + """ + Helper function to extract all variable list assignments from content. + Caches the result to avoid re-parsing the same content. + + EX: + content = ''' + drift1 = elements.Drift(ds=1.0) + cell = [drift1, quad1] + line = [cell, monitor] + ''' + _get_variable_assignments() + + Results in: + { + "cell": "drift1, quad1", + "line": "cell, monitor" + } + + """ + if self._content_hash in self._variable_assignments_cache: + return self._variable_assignments_cache[self._content_hash] + + variable_assignments = {} + var_assignment_pattern = r"(\w+)\s*=\s*\[(.*?)\]" + + for match in re.finditer(var_assignment_pattern, self._content, re.DOTALL): + var_name = match.group(1) + list_content = match.group(2) + variable_assignments[var_name] = list_content + + self._variable_assignments_cache[self._content_hash] = variable_assignments + return variable_assignments + + def _flatten(self, variable_name: str, debug: bool = False) -> List[str]: + """ + Recursively expands a varible name to replace it with its set of elements. + Utilizes caching to avoid redundant parsing. + + EX: + content = ''' + cell = [drift1, quad1] + line = [cell, cell] + sim.lattice.extend(line) + ''' + _flatten(content, "line") + + Results in: + ["drift1", "quad1", "drift1", "quad1"] + + :param variable_name: Name of the specific variable to expand (e.g. "line") + :param debug: Whether to print the expanded list. + :return: List of individual element names with all nesting resolved. + """ + # check cache first + cache_key = f"{self._content_hash}:{variable_name}" + if cache_key in self._flatten_cache: + return self._flatten_cache[cache_key] + + # Get variable assignments (cached) + variable_assignments = self._get_variable_assignments() + + # Check if the input is an inline list like "[monitor, elements.Drift(...)]" + if variable_name.startswith("[") and variable_name.endswith("]"): + list_contents = variable_name[1:-1] + # split on commas that are NOT inside parentheses + list_to_flatten = [ + element.strip() + for element in re.split(r",\s*(?![^()]*\))", list_contents) + if element.strip() + ] + else: + if variable_name not in variable_assignments: + self._flatten_cache[cache_key] = [variable_name] + return [variable_name] # It's not a list, it's a single element + + list_content = variable_assignments[variable_name] + list_to_flatten = [ + item.strip() for item in list_content.split(",") if item.strip() + ] + + expanded = [] + for item in list_to_flatten: + # recursively expand each item + sub_items = self._flatten(item, debug) + expanded.extend(sub_items) + + # Cache the result + self._flatten_cache[cache_key] = expanded + + return expanded + + def replace_variables(self, raw_lattice: List[str]) -> List[str]: + """ + This function is called to simplify the lattice list by replacing variable names with their corresponding constructor calls. + + EX: + (input) + drift1 = elements.Drift(ds=1.0) + quad1 = elements.Quad(k=0.5) + raw_lattice = ["drift1", "quad1"] + (output) + raw_lattice = ["elements.Drift(ds=1.0)", "elements.Quad(k=0.5)"] + + :param raw_lattice: List of lattice element variable names or constructor calls, e.g. ["drift1", "quad1"]. + :return: List with variable names replaced by their corresponding constructor calls. + """ + element_mapping = {} + + ellement_assignment_pattern = r"^\s*(\w+)\s*=\s*(elements\.\w+\(.*?\))" + all_element_assignments = re.findall( + ellement_assignment_pattern, self._content, re.MULTILINE | re.DOTALL + ) + + for var_name, element in all_element_assignments: + element_mapping[var_name] = element + + if not element_mapping: + return raw_lattice + + # Replace each item in the list + # later can be optimized by not iterating over the whole raw_lattice list + return [element_mapping.get(item, item) for item in raw_lattice] + + def extract_lattice_inputs(self, parsed_lattice: Dict[str, Any]) -> List[str]: + """ + Extracts all parameter values from parsed lattice data. + """ + used_variables = set() + + for element in parsed_lattice.get("lattice_elements", []): + for value in element["parameters"].values(): + used_variables.add(str(value).strip()) + + return list(used_variables) + + def _apply_reverse(self, var_name: str) -> None: + """ + Reverse the order of elements inside a list variable (in place). + + Example: + 1.) line = [a, b, c] + 2.) line.reverse() + 3.) result is -> [c, b, a] + + :param var_name: Name of the variable to reverse. + """ + + variable_assignments = self._get_variable_assignments() + if var_name not in variable_assignments: + return + + # Take the string inside the brackets, e.g. "a, b, c". + list_content = variable_assignments[var_name] + + # Split that string into individual items. + # We split on commas, but the regex makes sure we do not split commas + # that are inside parentheses (e.g., elements.Quad(k=0.5, ds=1.0)). + raw_items = re.split(r",\s*(?![^()]*\))", list_content) + + items = [item.strip() for item in raw_items if item.strip()] + items.reverse() + variable_assignments[var_name] = ", ".join(items) + + # Clear the flatten cache so future lookups rebuild the list using the new reversed order. + self._flatten_cache.clear() + + # ----------------------------------------------------------------------------- + # Debug methods + # ----------------------------------------------------------------------------- + + def print_lattice_operations(self, operations: List[Dict[str, str]]) -> None: + """ + Prints all lattice operations in the order they appear. + """ + + print("Full lattice operation sequence (in order):") + for operation in operations: + print(f" {operation['type']}({operation['argument']})") diff --git a/src/python/impactx/dashboard/Toolbar/file_imports/python/parser.py b/src/python/impactx/dashboard/Toolbar/file_imports/python/parser.py index 745435c70..7ab13fb0d 100644 --- a/src/python/impactx/dashboard/Toolbar/file_imports/python/parser.py +++ b/src/python/impactx/dashboard/Toolbar/file_imports/python/parser.py @@ -8,6 +8,7 @@ from .... import state from .helper import DashboardParserHelper +from .lattice_helper import DashboardLatticeParser state.import_file = False state.import_file_details = None @@ -63,22 +64,22 @@ def parse_impactx_simulation_file(file) -> None: file_content = DashboardParserHelper.import_file_content(file, state) + lattice_parser = DashboardLatticeParser(file_content) single_input_contents = DashboardParserHelper.parse_single_inputs(file_content) list_input_contents = DashboardParserHelper.parse_list_inputs(file_content) distribution_contents = DashboardParserHelper.parse_distribution(file_content) - lattice_element_contents = DashboardParserHelper.parse_lattice_elements( - file_content - ) - used_vars = lattice_element_contents.pop("used_lattice_variables", set()) + lattice_contents = lattice_parser.parse() + used_inputs = lattice_parser.extract_lattice_inputs(lattice_contents) variable_contents = DashboardParserHelper.parse_variables( - file_content, used_vars + file_content, used_inputs ) + parsed_values_dictionary = { **single_input_contents, **list_input_contents, **distribution_contents, - **lattice_element_contents, + **lattice_contents, "variables": variable_contents, } diff --git a/src/python/impactx/dashboard/Toolbar/file_imports/ui_populator.py b/src/python/impactx/dashboard/Toolbar/file_imports/ui_populator.py index 4440e79c8..36cc6846d 100644 --- a/src/python/impactx/dashboard/Toolbar/file_imports/ui_populator.py +++ b/src/python/impactx/dashboard/Toolbar/file_imports/ui_populator.py @@ -19,9 +19,11 @@ def on_import_file_change(import_file, **kwargs): state.importing_file = True DashboardParser.file_details(import_file) populate_impactx_simulation_file_to_ui(import_file) - except Exception: + except Exception as error: state.import_file_error = True - state.import_file_error_message = "Unable to parse" + state.import_file_error_message = ( + f"Unable to parse because of the following error: {error}" + ) finally: state.importing_file = False @@ -80,7 +82,7 @@ def _populate_lattice_config_to_ui(parsed_data): state.selected_lattice_list = [] for lattice_element_index, element in enumerate(parsed_data): - parsed_element = element["element"] + parsed_element = element["name"] parsed_parameters = element["parameters"] state.selected_lattice = parsed_element diff --git a/src/python/impactx/extensions/KnownElementsList.py b/src/python/impactx/extensions/KnownElementsList.py index abaae6b47..69ab252b0 100644 --- a/src/python/impactx/extensions/KnownElementsList.py +++ b/src/python/impactx/extensions/KnownElementsList.py @@ -125,7 +125,7 @@ def select( kind=None, name=None, ): - """Apply filtering to this filtered list. + r"""Apply filtering to this filtered list. This method applies additional filtering to an already filtered list, maintaining references to the original elements and enabling chaining. @@ -364,7 +364,7 @@ def select( kind=None, name=None, ) -> FilteredElementsList: - """Filter elements by type and name with OR-based logic. + r"""Filter elements by type and name with OR-based logic. This method supports filtering elements by their type and/or name using keyword arguments. Returns references to original elements, allowing modification and chaining. diff --git a/src/python/impactx/impactx_pybind/__init__.pyi b/src/python/impactx/impactx_pybind/__init__.pyi index b83e9f743..ab84e45ac 100644 --- a/src/python/impactx/impactx_pybind/__init__.pyi +++ b/src/python/impactx/impactx_pybind/__init__.pyi @@ -527,6 +527,7 @@ class ImpactX: | elements.Multipole | elements.NonlinearLens | elements.PlaneXYRot + | elements.PolygonAperture | elements.Programmable | elements.PRot | elements.Quad @@ -608,8 +609,8 @@ class ImpactXParIter(amrex.space3d.amrex_3d_pybind.ParIter_pureSoA_8_0_default): class ImpactXParticleContainer( amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default ): - const_iterator = ImpactXParConstIter - iterator = ImpactXParIter + ConstIterator = ImpactXParConstIter + Iterator = ImpactXParIter def add_n_particles( self, x: amrex.space3d.amrex_3d_pybind.PODVector_real_std, @@ -965,6 +966,7 @@ def push( | elements.Multipole | elements.NonlinearLens | elements.PlaneXYRot + | elements.PolygonAperture | elements.Programmable | elements.PRot | elements.Quad @@ -1011,6 +1013,7 @@ def push( | elements.Multipole | elements.NonlinearLens | elements.PlaneXYRot + | elements.PolygonAperture | elements.Programmable | elements.PRot | elements.Quad @@ -1033,6 +1036,6 @@ __author__: str = ( "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." ) __license__: str = "BSD-3-Clause-LBNL" -__version__: str = "25.11" +__version__: str = "25.12" s: CoordSystem # value = t: CoordSystem # value = diff --git a/src/python/impactx/impactx_pybind/elements/__init__.pyi b/src/python/impactx/impactx_pybind/elements/__init__.pyi index 2a6e4030a..f7ecbc295 100644 --- a/src/python/impactx/impactx_pybind/elements/__init__.pyi +++ b/src/python/impactx/impactx_pybind/elements/__init__.pyi @@ -39,6 +39,7 @@ __all__: list[str] = [ "NonlinearLens", "PRot", "PlaneXYRot", + "PolygonAperture", "Programmable", "Quad", "QuadEdge", @@ -1382,6 +1383,7 @@ class KnownElementsList: | impactx.impactx_pybind.elements.Multipole | impactx.impactx_pybind.elements.NonlinearLens | impactx.impactx_pybind.elements.PlaneXYRot + | impactx.impactx_pybind.elements.PolygonAperture | impactx.impactx_pybind.elements.Programmable | impactx.impactx_pybind.elements.PRot | impactx.impactx_pybind.elements.Quad @@ -1424,6 +1426,7 @@ class KnownElementsList: | impactx.impactx_pybind.elements.Multipole | impactx.impactx_pybind.elements.NonlinearLens | impactx.impactx_pybind.elements.PlaneXYRot + | impactx.impactx_pybind.elements.PolygonAperture | impactx.impactx_pybind.elements.Programmable | impactx.impactx_pybind.elements.PRot | impactx.impactx_pybind.elements.Quad @@ -1466,6 +1469,7 @@ class KnownElementsList: | impactx.impactx_pybind.elements.Multipole | impactx.impactx_pybind.elements.NonlinearLens | impactx.impactx_pybind.elements.PlaneXYRot + | impactx.impactx_pybind.elements.PolygonAperture | impactx.impactx_pybind.elements.Programmable | impactx.impactx_pybind.elements.PRot | impactx.impactx_pybind.elements.Quad @@ -1509,6 +1513,7 @@ class KnownElementsList: | impactx.impactx_pybind.elements.Multipole | impactx.impactx_pybind.elements.NonlinearLens | impactx.impactx_pybind.elements.PlaneXYRot + | impactx.impactx_pybind.elements.PolygonAperture | impactx.impactx_pybind.elements.Programmable | impactx.impactx_pybind.elements.PRot | impactx.impactx_pybind.elements.Quad @@ -2064,6 +2069,94 @@ class PlaneXYRot(mixin.Named, mixin.Thin, mixin.Alignment): @angle.setter def angle(self, arg1: typing.SupportsFloat) -> None: ... +class PolygonAperture(mixin.Named, mixin.Thin, mixin.Alignment): + def __init__( + self, + vertices_x: collections.abc.Sequence[typing.SupportsFloat], + vertices_y: collections.abc.Sequence[typing.SupportsFloat], + min_radius2: typing.SupportsFloat = 0.0, + repeat_x: typing.SupportsFloat = 0, + repeat_y: typing.SupportsFloat = 0, + shift_odd_x: bool = False, + action: str = "transmit", + dx: typing.SupportsFloat = 0, + dy: typing.SupportsFloat = 0, + rotation: typing.SupportsFloat = 0, + name: str | None = None, + ) -> None: + """ + A short collimator element described by a polygon with vertices given by their x and y coordinates. + """ + def __repr__(self) -> str: ... + @typing.overload + def push( + self, + pc: impactx.impactx_pybind.ImpactXParticleContainer, + step: typing.SupportsInt = 0, + period: typing.SupportsInt = 0, + ) -> None: + """ + Push first the reference particle, then all other particles. + """ + @typing.overload + def push( + self, + cm: amrex.space3d.amrex_3d_pybind.SmallMatrix_6x6_F_SI1_double, + ref: impactx.impactx_pybind.RefPart, + ) -> None: + """ + Linear push of the covariance matrix through an element. Expects that the reference particle was advanced first. + """ + def to_dict( + self, + ) -> dict[ + str, + float + | int + | int + | str + | list[float] + | list[int] + | list[int] + | amrex.space3d.amrex_3d_pybind.SmallMatrix_6x6_F_SI1_double + | None, + ]: ... + @property + def action(self) -> str: + """ + action type (transmit, absorb) + """ + @action.setter + def action(self, arg1: str) -> None: ... + @property + def min_radius2(self) -> float: + """ + All particles with radius squared smaller than min_radius2 pass the aperture + """ + @min_radius2.setter + def min_radius2(self, arg1: typing.SupportsFloat) -> None: ... + @property + def repeat_x(self) -> float: + """ + horizontal period for repeated aperture masking + """ + @repeat_x.setter + def repeat_x(self, arg1: typing.SupportsFloat) -> None: ... + @property + def repeat_y(self) -> float: + """ + vertical period for repeated aperture masking + """ + @repeat_y.setter + def repeat_y(self, arg1: typing.SupportsFloat) -> None: ... + @property + def shift_odd_x(self) -> bool: + """ + for hexagonal/triangular mask patterns: horizontal shift of every 2nd (odd) vertical period by repeat_x / 2. Use alignment offsets dx,dy to move whole mask as needed. + """ + @shift_odd_x.setter + def shift_odd_x(self, arg1: bool) -> None: ... + class Programmable(mixin.Named): def __init__( self, @@ -2693,7 +2786,11 @@ class Sol(mixin.Named, mixin.Thick, mixin.Alignment, mixin.PipeAperture): class Source(mixin.Named, mixin.Thin): def __init__( - self, distribution: str, openpmd_path: str, name: str | None = None + self, + distribution: str, + openpmd_path: str, + active_once: bool = True, + name: str | None = None, ) -> None: """ A particle source. @@ -2733,6 +2830,13 @@ class Source(mixin.Named, mixin.Thin): | None, ]: ... @property + def active_once(self) -> bool: + """ + Inject particles only for the first lattice period. + """ + @active_once.setter + def active_once(self, arg1: bool) -> None: ... + @property def distribution(self) -> str: """ Distribution type of particles in the source diff --git a/src/python/impactx/impactx_pybind/elements/transformation.pyi b/src/python/impactx/impactx_pybind/elements/transformation.pyi index 68905f1d1..b9aabb6f6 100644 --- a/src/python/impactx/impactx_pybind/elements/transformation.pyi +++ b/src/python/impactx/impactx_pybind/elements/transformation.pyi @@ -36,6 +36,7 @@ def insert_element_every_ds( | impactx.impactx_pybind.elements.Multipole | impactx.impactx_pybind.elements.NonlinearLens | impactx.impactx_pybind.elements.PlaneXYRot + | impactx.impactx_pybind.elements.PolygonAperture | impactx.impactx_pybind.elements.Programmable | impactx.impactx_pybind.elements.PRot | impactx.impactx_pybind.elements.Quad diff --git a/src/python/impactx/plot/Survey.py b/src/python/impactx/plot/Survey.py index 234a4cf67..33e48db53 100644 --- a/src/python/impactx/plot/Survey.py +++ b/src/python/impactx/plot/Survey.py @@ -84,9 +84,13 @@ def plot_survey( if "Quad" in el_type: height = copysign(0.8, el_dict["k"] * charge_qe) if "Sbend" in el_type: - if ref is None: + if ref is not None: height = copysign(0.8, element.rc(ref)) else: # guess + if el_type == "Sbend": + el_dict["phi"] = ( + el_dict["ds"] / (2 * np.pi * el_dict["rc"]) * 360 + ) # calculate bending angle (in degrees) and add to dict height = copysign(0.8, el_dict["phi"]) # TODO: sign dependent, read m_p_scale # if el_type == "Kicker": @@ -127,6 +131,6 @@ def plot_survey( ax.set_ylim(-1, 1) ax.set_yticks([]) - ax.set_aspect(1 / 1.618) # golden ratio + ax.set_box_aspect(1 / 6) # some nice aspect ratio return ax diff --git a/src/python/pyImpactX.H b/src/python/pyImpactX.H index 80d5632e6..e2017da29 100644 --- a/src/python/pyImpactX.H +++ b/src/python/pyImpactX.H @@ -14,6 +14,7 @@ // include // not yet used #include #include +#include #include #include diff --git a/tests/python/dashboard/conftest.py b/tests/python/dashboard/conftest.py index bc3a40737..3b2c52a0d 100644 --- a/tests/python/dashboard/conftest.py +++ b/tests/python/dashboard/conftest.py @@ -1,6 +1,11 @@ +import pathlib +import sys + import pytest from seleniumbase import SB +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[3] / "src" / "python")) + from .utils import ( DashboardTester, save_failure_screenshot, diff --git a/tests/python/dashboard/test_impactx_examples_lattice.py b/tests/python/dashboard/test_impactx_examples_lattice.py new file mode 100644 index 000000000..abd03c7e7 --- /dev/null +++ b/tests/python/dashboard/test_impactx_examples_lattice.py @@ -0,0 +1,114 @@ +""" +This file is part of ImpactX + +Copyright 2025 ImpactX contributors +Authors: Parthib Roy +License: BSD-3-Clause-LBNL +""" + +from .utils import DashboardTester + + +class DashboardExamples: + def __init__(self, dashboard: DashboardTester): + self.dashboard = dashboard + self.sb = dashboard.sb + + def _load_example(self, example_path): + self.dashboard.load_example(example_path) + + def _clear(self): + self.sb.click("#reset_all_inputs_button") + + def _begin(self, example_name: str): + print(f"Starting {example_name} test.") + + def _end(self, example_name: str): + print(f"Successfully ending {example_name} test.") + + def _test_lattice(self, inputs: list): + for element_id, expected_value in inputs: + actual_value = self.sb.get_value(element_id) + + # Try numeric comparison if expected_value is a number + if isinstance(expected_value, (int, float)): + actual_value = float(actual_value) + assert actual_value == expected_value, ( + f"{element_id}: expected {expected_value}, got {actual_value}" + ) + else: + assert actual_value == expected_value, ( + f"{element_id}: expected '{expected_value}', got '{actual_value}'" + ) + + def _run_example(self, name: str, path: str, test_data: list): + self._begin(name) + self._clear() + self._load_example(path) + self._test_lattice(test_data) + self._end(name) + + def dogleg_lattice(self): + LATTICE_CONFIGURATION = [ + # sbend1 + ("#name2", "sbend1"), + ("#ds2", "lb"), + ("#rc2", "-rc"), + ("#nslice2", "ns"), + # dipedge1 + ("#name3", "dipedge1"), + ("#psi3", "-psi"), + ("#rc3", "-rc"), + # dr1 + ("#name4", "dr1"), + ("#ds4", 5.0058489435), + ("#nslice4", "ns"), + # dipedge2 + ("#name5", "dipedge2"), + ("#psi5", "psi"), + ("#rc5", "rc"), + # sbend2 + ("#name6", "sbend2"), + ("#ds6", "lb"), + ("#rc6", "rc"), + ("#nslice6", "ns"), + # dr2 + ("#name7", "dr2"), + ("#ds7", 0.5), + ("#nslice7", "ns"), + ] + self._run_example("dogleg", "dogleg/run_dogleg.py", LATTICE_CONFIGURATION) + + def chicane_lattice(self): + """ + Utilizes lattice_half.reverse() + """ + LATTICE_CONFIGURATION = [ + # FIRST LATTICE_HALF + ("#name2", "sbend1"), + ("#name3", "dipedge1"), + ("#name5", "dr1"), + ("#name7", "dipedge2"), + ("#name8", "sbend2"), + ("#name10", "dr2"), + ("#name12", "sbend2"), + ("#name13", "dipedge2"), + ("#name15", "dr1"), + ("#name17", "dipedge1"), + ("#name18", "sbend1"), + ("#name20", "dr3"), + ] + + self._run_example("chicane", "chicane/run_chicane.py", LATTICE_CONFIGURATION) + + +def test_examples(dashboard) -> None: + """ + Exercise a subset of example imports against the running dashboard. + """ + + test_examples = DashboardExamples(dashboard) + + # test lattice builds w/example importing + test_examples.dogleg_lattice() # uses lattice.append() and lattice.extend() + test_examples.chicane_lattice() # uses .reverse() diff --git a/tests/python/dashboard/test_lattice_stats.py b/tests/python/dashboard/test_lattice_stats.py new file mode 100644 index 000000000..91ef23022 --- /dev/null +++ b/tests/python/dashboard/test_lattice_stats.py @@ -0,0 +1,79 @@ +""" +This file is part of ImpactX + +Copyright 2025 ImpactX contributors +Authors: Parthib Roy +License: BSD-3-Clause-LBNL +""" + + +def test_lattice_stats(dashboard): + """ + Tests the lattice statistics to ensure accuracy and indirectly + tests the dashboard's python parser to correctly import all inputs + from a loaded Python file. + """ + + FODO_EXPECTED_RESULTS = { + "total_elements": 11, + "total_length": "3.0m", + "max_length": "1.0m", + "avg_length": "0.6m", + "min_length": "0.25m", + "total_steps": 125, + "periods": 1, + "element_counts": {"beammonitor": 6, "drift": 3, "quad": 2}, + } + + APOCHROMATIC_EXPECTED_RESULTS = { + "total_elements": 14, + "total_length": "32.5m", + "total_steps": 300, + "periods": 1, + "element_counts": {"chrquad": 8, "chrdrift": 4, "beammonitor": 2}, + } + + CHICANE_CSR_EXPECTED_RESULTS = { + "total_elements": 14, + "total_length": "15.01m", + "total_steps": 200, + "periods": 1, + "element_counts": {"drift": 4, "dipedge": 4, "sbend": 4, "beammonitor": 2}, + } + + DOGLEG_REVERSE_EXPECTED_RESULTS = { + "total_elements": 8, + "total_length": "6.51m", + "total_steps": 100, + "periods": 1, + "element_counts": {"beammonitor": 2, "drift": 2, "sbend": 2, "dipedge": 2}, + } + + IOLATTICE_EXPECTED_RESULTS = { + "total_elements": 117, + "total_length": "39.97m", + # locally gives 919, pytest gives 2283 when all 5 examples are tested, passes if ran alone + # "total_steps": 918, + "periods": 5, + "element_counts": { + "drift": 52, + "quad": 39, + "dipedge": 16, + "sbend": 8, + "beammonitor": 2, + }, + } + + # Begin Testing + EXAMPLES = [ + ("fodo/run_fodo.py", FODO_EXPECTED_RESULTS), + ("apochromatic/run_apochromatic.py", APOCHROMATIC_EXPECTED_RESULTS), + ("chicane/run_chicane_csr.py", CHICANE_CSR_EXPECTED_RESULTS), + ("dogleg/run_dogleg_reverse.py", DOGLEG_REVERSE_EXPECTED_RESULTS), + ("iota_lattice/run_iotalattice.py", IOLATTICE_EXPECTED_RESULTS), + ] + + for example, expected in EXAMPLES: + dashboard.load_example(example) + for state_name, expected_value in expected.items(): + dashboard.assert_state(state_name, expected_value) diff --git a/tests/python/dashboard/test_python_import.py b/tests/python/dashboard/test_python_import.py index ff143b20e..e9de74b88 100644 --- a/tests/python/dashboard/test_python_import.py +++ b/tests/python/dashboard/test_python_import.py @@ -55,13 +55,15 @@ def test_python_import(dashboard): ] LATTICE_CONFIGURATION = [ - ("#ds1", 0.25), - ("#ds2", 1), - ("#k2", 1.0), - ("#ds3", 0.5), - ("#ds4", 1.0), - ("#k4", -1.0), - ("#ds5", 0.25), + ("#name1", "drift1"), + ("#name2", "quad1"), + ("#name3", "drift2"), + ("#name4", "quad2"), + ("#name5", "drift3"), + ("#name6", "drift3"), + ("#name7", "quad2"), + ("#name8", "drift2"), + ("#ds8", 0.5), ] # Check state parameters @@ -72,7 +74,16 @@ def test_python_import(dashboard): # Check input values for element_id, expected_value in DISTRIBUTION_VALUES + LATTICE_CONFIGURATION: - actual_value = float(dashboard.sb.get_value(element_id)) - assert actual_value == pytest.approx(expected_value, **APPROX_TOL), ( - f"{element_id}: expected {expected_value}, got {actual_value}" - ) + actual_value = dashboard.sb.get_value(element_id) + + if isinstance(expected_value, str): + # Compare strings directly + assert actual_value == expected_value, ( + f"{element_id}: expected '{expected_value}', got '{actual_value}'" + ) + else: + # Compare numbers with tolerance + actual_value = float(actual_value) + assert actual_value == pytest.approx(expected_value, **APPROX_TOL), ( + f"{element_id}: expected {expected_value}, got {actual_value}" + ) diff --git a/tests/python/dashboard/testdata/example.py b/tests/python/dashboard/testdata/example.py index 2ed31d1b3..26eb93414 100644 --- a/tests/python/dashboard/testdata/example.py +++ b/tests/python/dashboard/testdata/example.py @@ -32,15 +32,20 @@ ) ns = 25 + +quad1 = elements.Quad(ds=1.0, k=1.0, nslice=ns, name="quad1") lattice_configuration = [ - elements.Drift(ds=0.25, nslice=ns, name="drift1"), - elements.Quad(ds=1.0, k=1.0, nslice=ns, name="quad1"), elements.Drift(ds=0.5, nslice=ns, name="drift2"), elements.Quad(ds=1.0, k=-1.0, nslice=ns, name="quad2"), elements.Drift(ds=0.25, nslice=ns, name="drift3"), ] +sim.lattice.append(elements.Drift(ds=0.25, nslice=ns, name="drift1")) +sim.lattice.append(quad1) +sim.lattice.extend(lattice_configuration) +lattice_configuration.reverse() sim.lattice.extend(lattice_configuration) + sim.periods = 2 # Simulate diff --git a/tests/python/dashboard/utils.py b/tests/python/dashboard/utils.py index e811b2635..8bbb34ec0 100644 --- a/tests/python/dashboard/utils.py +++ b/tests/python/dashboard/utils.py @@ -15,6 +15,7 @@ import pytest from selenium.common.exceptions import TimeoutException +REPO_ROOT = Path(__file__).resolve().parents[3] TIMEOUT = 120 APPROX_TOL = {"rel": 1e-12, "abs": 1e-12} @@ -23,7 +24,7 @@ def start_dashboard() -> subprocess.Popen[str]: """ Starts the impactx-dashboard in a subprocess. """ - repo_root = get_impactx_root_dir() + repo_root = REPO_ROOT working_directory = os.path.normpath( os.path.join(repo_root, "src", "python", "impactx") ) @@ -38,21 +39,6 @@ def start_dashboard() -> subprocess.Popen[str]: ) -def get_impactx_root_dir(): - """ - Locates the ImpactX source directory. - - Looks for the first parent directory named 'impactx' that contains a '.git' folder. - """ - - current_directory = Path(__file__).resolve() - - for parent_dir in current_directory.parents: - if parent_dir.name == "impactx" and (parent_dir / ".git").is_dir(): - return parent_dir - return None - - def wait_for_interaction_ready(sb, timeout=TIMEOUT): """ Waits until the dashboard has fully loaded and is ready for interaction. @@ -97,7 +83,7 @@ def __init__(self, sb): self.sb = sb # Set up the examples directory path once - impactx_directory = Path(get_impactx_root_dir()) + impactx_directory = REPO_ROOT self.examples_directory = impactx_directory / "examples" self.testing_directory = impactx_directory / "tests" / "python" / "dashboard" diff --git a/tests/python/test_dataframe.py b/tests/python/test_dataframe.py index aa45271a4..f776c4a25 100644 --- a/tests/python/test_dataframe.py +++ b/tests/python/test_dataframe.py @@ -107,6 +107,9 @@ def test_df_pandas(save_png=True): "momentum_x", "momentum_y", "momentum_t", + "spin_x", + "spin_y", + "spin_z", "qm", "weighting", ] diff --git a/tests/python/test_impactx.py b/tests/python/test_impactx.py index c1dce0f46..644cd3043 100755 --- a/tests/python/test_impactx.py +++ b/tests/python/test_impactx.py @@ -11,7 +11,7 @@ import pytest from conftest import basepath -from impactx import ImpactX, distribution, elements +from impactx import Config, ImpactX, distribution, elements # FIXME in AMReX via https://github.com/AMReX-Codes/amrex/pull/3727 # def test_impactx_module(): @@ -26,8 +26,16 @@ def validate_fodo(beam): """see examples/fodo/analysis_fodo.py""" num_particles = beam.total_number_of_particles() assert num_particles == 10000 - atol = 0.0 # ignored - rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution + if Config.precision == "SINGLE": + atol = 0.0 # ignored + rtol = ( + 2.5 * num_particles**-0.5 + ) # from random sampling of a smooth distribution + else: + atol = 0.0 # ignored + rtol = ( + 2.2 * num_particles**-0.5 + ) # from random sampling of a smooth distribution # in situ calculate the reduced beam characteristics rbc = beam.beam_moments() @@ -160,7 +168,7 @@ def test_impactx_nofile(): # simulate full lattice but keep beam global position sim.track_particles() - assert np.allclose([ref.s], [7.0]) + assert ref.s == pytest.approx(7.0) # finalize simulation sim.finalize() diff --git a/tests/python/test_lattice_select.py b/tests/python/test_lattice_select.py index 709457bf0..4a5bb872b 100644 --- a/tests/python/test_lattice_select.py +++ b/tests/python/test_lattice_select.py @@ -521,9 +521,9 @@ def test_complex_scenarios(): quad.k *= 1.1 # Increase strength by 10% # Verify modifications affected original elements - assert lattice[1].k == 1.1 # qf1 - assert lattice[5].k == -1.1 # qd1 - assert lattice[7].k == 1.1 # qf2 + assert lattice[1].k == pytest.approx(1.1) # qf1 + assert lattice[5].k == pytest.approx(-1.1) # qd1 + assert lattice[7].k == pytest.approx(1.1) # qf2 # Scenario 3b: Find all quadrupoles and modify their strength (type-based) all_quads_type = lattice.select(kind=elements.Quad) @@ -546,8 +546,8 @@ def test_complex_scenarios(): bend.ds *= 0.9 # Decrease length by 10% # Verify modifications - assert lattice[3].ds == 0.9 # bend1 - assert lattice[9].ds == 0.9 # bend2 + assert lattice[3].ds == pytest.approx(0.9) # bend1 + assert lattice[9].ds == pytest.approx(0.9) # bend2 # Scenario 4b: Find all bends and modify their length (type-based) all_bends_type = lattice.select(kind=elements.Sbend) diff --git a/tests/python/test_transformation.py b/tests/python/test_transformation.py index cab4bdcf3..ea4933798 100644 --- a/tests/python/test_transformation.py +++ b/tests/python/test_transformation.py @@ -9,7 +9,13 @@ import numpy as np import pytest -from impactx import CoordSystem, ImpactX, coordinate_transformation, distribution +from impactx import ( + Config, + CoordSystem, + ImpactX, + coordinate_transformation, + distribution, +) def test_transformation(): @@ -75,8 +81,12 @@ def test_transformation(): sim.finalize() # assert that forward-inverse transformation of the beam leaves beam unchanged - atol = 1e-14 - rtol = 1e-10 + if Config.precision == "SINGLE": + atol = 1e-6 + rtol = 1e-6 + else: + 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")