From 5638dfd4e4e26a2d3ce24ec73c411fc3d68e9319 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Mon, 13 Jan 2025 15:21:32 -0800 Subject: [PATCH] User-Defined Linear Map (#743) * [Draft] Transport and Covariance Matrices General structure. Co-authored-by: Chad Mitchell * Update InitElement.cpp Update initialization of LinearMap element. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update LinearMap.H Finish implementation of LinearMap element push. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update InitDistribution.cpp Initialize covariance matrix from distribution parameters. * Add Drift return linear map. * Use LinearMap struct in Drift.H. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Transport/Covariane: Use `SmallMatrix` * Add Python bindings for LinearMap element. * Update elements.cpp Remove unnecessary ,. * Update elements.cpp Correct naming in elements.cpp binding for LinearMap. * Update elements.cpp Correct argument declaration for Python bindings of LinearMap in elements.cpp. * Add LinearMap to 'Named' elements. * Update LinearMap.H Document "Name" parameter in LinearMap. * Modify LinearMap Python bindings. * Update src/initialization/InitDistribution.cpp * Update src/initialization/InitDistribution.cpp * Update InitDistribution.cpp Modify zero-initialization of CovarianceMatrix cv. * Leftover Improvements from first PR * Finalize Python * More Detailed Warning Include element name in output to user. * Start Examples * Start Documentation * Fix Bugs (incl. AMReX) Use parser for inputs. Needs upstream AMReX fix. * Docs Updates * `LinearMap`: `ds` bookkeeping And reference particle pushing (drifting), if thick. * Add thin linear map example. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Document example README. * Update run_map.py Add twiss * Update run_map.py Correct name of Rmat/R. * Update run_map.py Fix naming of R again. * Update input_map.in Fix number of periods. * Apply suggestions from code review Co-authored-by: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> * Add docs for matrix elements. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Doc Updates * Generalize CMake Args for inputs Taken over from generalized version we merged to WarpX. * Add symplectic warning in app element documentation. * Add symplectic warning in Python element documentation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove runtime warning. * Add support for nonzero ds to app input. * Update examples/CMakeLists.txt * FODO Analysis: Check s and gamma * __repr__: no select entries --------- Co-authored-by: Axel Huebl Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/source/usage/examples.rst | 3 +- docs/source/usage/parameters.rst | 19 ++ docs/source/usage/python.rst | 19 ++ docs/source/usage/workflows/add_element.rst | 31 ++- examples/CMakeLists.txt | 64 +++++- examples/fodo/analysis_fodo.py | 12 +- examples/fodo_userdef/README.rst | 85 ++++++++ examples/fodo_userdef/analysis_fodo.py | 1 + examples/fodo_userdef/input_fodo_userdef.in | 61 ++++++ examples/fodo_userdef/plot_fodo.py | 1 + examples/fodo_userdef/run_fodo_userdef.py | 89 +++++++++ examples/linear_map/README.rst | 58 ++++++ examples/linear_map/analysis_map.py | 154 +++++++++++++++ examples/linear_map/input_map.in | 58 ++++++ examples/linear_map/run_map.py | 104 ++++++++++ src/initialization/InitDistribution.cpp | 57 +++++- src/initialization/InitElement.cpp | 21 ++ src/particles/CovarianceMatrix.H | 24 +++ src/particles/PushAll.H | 4 + src/particles/elements/All.H | 6 +- src/particles/elements/Drift.H | 31 ++- src/particles/elements/LinearMap.H | 182 ++++++++++++++++++ .../elements/mixin/lineartransport.H | 52 +++++ src/python/elements.cpp | 58 ++++++ 24 files changed, 1173 insertions(+), 21 deletions(-) create mode 100644 examples/fodo_userdef/README.rst create mode 120000 examples/fodo_userdef/analysis_fodo.py create mode 100644 examples/fodo_userdef/input_fodo_userdef.in create mode 120000 examples/fodo_userdef/plot_fodo.py create mode 100755 examples/fodo_userdef/run_fodo_userdef.py create mode 100644 examples/linear_map/README.rst create mode 100755 examples/linear_map/analysis_map.py create mode 100644 examples/linear_map/input_map.in create mode 100755 examples/linear_map/run_map.py create mode 100644 src/particles/CovarianceMatrix.H create mode 100644 src/particles/elements/LinearMap.H create mode 100644 src/particles/elements/mixin/lineartransport.H diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 2602ac55c..8e911e9e6 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -31,10 +31,11 @@ Single Particle Dynamics examples/aperture/README.rst examples/iota_lens/README.rst examples/achromatic_spectrometer/README.rst + examples/fodo_userdef/README.rst examples/fodo_programmable/README.rst examples/dogleg/README.rst examples/coupled_optics/README.rst - + examples/linear_map/README.rst Collective Effects ------------------ diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index 9231df391..70d7437de 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -349,6 +349,25 @@ Lattice Elements * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane * ``.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``) + * ``linear_map`` for a custom, linear transport matrix. + + The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. + The transport matrix :math:`R` is defaulted to the identity matrix, so only matrix entries that differ from that need to be specified. + + The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m + and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. + + The internal tracking methods used by ImpactX are symplectic. However, if a user-defined linear map :math:`R` is provided, it is up to the user to ensure that the matrix :math:`R` is symplectic. Otherwise, this condition may be violated. + + This element requires these additional parameters: + + * ``.R(i,j)`` (``float``, ...) matrix entries + a 1-indexed, 6x6, linear transport map to multiply with the the phase space vector :math:`x,px,y,py,t,pt`. + * ``.ds`` (``float``, in meters) length associated with a user-defined linear element (defaults to 0) + * ``.dx`` (``float``, in meters) horizontal translation error + * ``.dy`` (``float``, in meters) vertical translation error + * ``.rotation`` (``float``, in degrees) rotation error in the transverse plane + * ``multipole`` for a thin multipole element. This requires these additional parameters: diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 16df118cc..889aca74d 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -611,6 +611,25 @@ This module provides elements for the accelerator lattice. :param unit: specification of units (``"dimensionless"`` in units of the magnetic rigidity of the reference particle or ``"T-m"``) :param name: an optional name for the element +.. py::class:: impactx.elements.LinearMap(R, dx=0, dy=0, rotation=0, name=None) + + A custom, linear transport matrix. + + The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. + + The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m + and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. + + The internal tracking methods used by ImpactX are symplectic. However, if a user-defined linear map :math:`R` is provided, it is + up to the user to ensure that the matrix :math:`R` is symplectic. Otherwise, this condition may be violated. + + :param R: a linear transport map to multiply with the the phase space vector :math:`(x,px,y,py,t,pt)`. + :param ds: length associated with a user-defined linear element (defaults to 0), in m + :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:class:: impactx.elements.Multipole(multipole, K_normal, K_skew, dx=0, dy=0, rotation=0, name=None) A general thin multipole element. diff --git a/docs/source/usage/workflows/add_element.rst b/docs/source/usage/workflows/add_element.rst index db80e21eb..fe8debaa3 100644 --- a/docs/source/usage/workflows/add_element.rst +++ b/docs/source/usage/workflows/add_element.rst @@ -10,6 +10,28 @@ The workflows described here apply both for thin kicks or thick elements. Thick elements can also use soft-edged fringe fields (see `existing soft-edged elements for implementation details `__). +.. _usage-workflows-add-element-linmap: + +Linear Map +---------- + +A custom linear element can be provided by specifying the 6x6 linear transport matrix :math:`R` as an input. +See the :ref:` example ` for Python and inputs file syntax to specify a custom linear element. + +The matrix elements :math:`R(i,j)` are indexed beginning with 1, so that :math:`i,j=1,2,3,4,5,6`. + +The matrix :math:`R` multiplies the phase space vector :math:`(x,px,y,py,t,pt)`, where coordinates :math:`(x,y,t)` have units of m +and momenta :math:`(px,py,pt)` are dimensionless. So, for example, :math:`R(1,1)` is dimensionless, and :math:`R(1,2)` has units of m. + + +.. note:: + + If a user-provided linear map is used, it is up to the user to ensure that the 6x6 transport matrix is symplectic. + If a more general form of user-defined transport is needed, the :ref:`Python Programmable Element ` and the :ref:`C++ Element ` provide a more general approach. + + +.. _usage-workflows-add-element-python: + Python Programmable Element --------------------------- @@ -30,14 +52,7 @@ Detailed examples that show usage of the programmable element are: Detailed particle computing interfaces are presented in the `pyAMReX examples `__. -Linear Map ----------- - -.. note:: - - We plan to add a simple, linear map element that can be configured in user input. - Follow `issue #538 `__ for progress. - +.. _usage-workflows-add-element-cxx: C++ Element ----------- diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 1783ad36d..ce1704b09 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -55,6 +55,30 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) # make a unique run directory file(MAKE_DIRECTORY ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name}) + # get input file/script and optional command-line arguments + separate_arguments(INPUTS_LIST UNIX_COMMAND "${input}") + list(GET INPUTS_LIST 0 INPUTS_FILE) + list(LENGTH INPUTS_LIST INPUTS_LIST_LENGTH) + if(INPUTS_LIST_LENGTH GREATER 1) + list(SUBLIST INPUTS_LIST 1 -1 INPUTS_ARGS) + list(JOIN INPUTS_ARGS " " INPUTS_ARGS) + else() + set(INPUTS_ARGS "") + endif() + cmake_path(SET INPUTS_FILE "${ImpactX_SOURCE_DIR}/${INPUTS_FILE}") + + # get analysis script and optional command-line arguments + separate_arguments(ANALYSIS_LIST UNIX_COMMAND "${analysis_script}") + list(GET ANALYSIS_LIST 0 ANALYSIS_FILE) + cmake_path(SET ANALYSIS_FILE "${CMAKE_CURRENT_SOURCE_DIR}/${ANALYSIS_FILE}") + list(LENGTH ANALYSIS_LIST ANALYSIS_LIST_LENGTH) + if(ANALYSIS_LIST_LENGTH GREATER 1) + list(SUBLIST ANALYSIS_LIST 1 -1 ANALYSIS_ARGS) + list(JOIN ANALYSIS_ARGS " " ANALYSIS_ARGS) + else() + set(ANALYSIS_ARGS "") + endif() + # test run set(THIS_WORKING_DIR ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name}) set(THIS_MPI_TEST_EXE) @@ -64,10 +88,10 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) set(THIS_Python_EXE) if(is_python) set(THIS_Python_EXE ${Python_EXECUTABLE}) - endif() - if(is_python) + # for argparse, do not pass command-line arguments as one quoted string + separate_arguments(INPUTS_ARGS UNIX_COMMAND "${INPUTS_ARGS}") add_test(NAME ${name}.run - COMMAND ${THIS_MPI_TEST_EXE} ${THIS_Python_EXE} ${ImpactX_SOURCE_DIR}/${input} + COMMAND ${THIS_MPI_TEST_EXE} ${THIS_Python_EXE} ${INPUTS_FILE} ${INPUTS_ARGS} WORKING_DIRECTORY ${THIS_WORKING_DIR} ) # TODO: @@ -77,12 +101,13 @@ function(add_impactx_test name input is_mpi analysis_script plot_script) else() add_test(NAME ${name}.run COMMAND - ${THIS_MPI_TEST_EXE} $ ${ImpactX_SOURCE_DIR}/${input} + ${THIS_MPI_TEST_EXE} $ ${INPUTS_FILE} amrex.abort_on_unused_inputs=1 amrex.throw_exception = 1 amrex.signal_handling = 0 impactx.always_warn_immediately=1 impactx.abort_on_warning_threshold=low + ${INPUTS_ARGS} WORKING_DIRECTORY ${THIS_WORKING_DIR} ) endif() @@ -210,6 +235,21 @@ add_impactx_test(FODO.MADX.py examples/fodo/plot_fodo.py ) +# Python: FODO Cell w/ custom linear element ################################# +# +add_impactx_test(FODO.userdef + examples/fodo_userdef/input_fodo_userdef.in + ON # ImpactX MPI-parallel + examples/fodo_userdef/analysis_fodo.py + examples/fodo_userdef/plot_fodo.py +) +add_impactx_test(FODO.userdef.py + examples/fodo_userdef/run_fodo_userdef.py + OFF # ImpactX MPI-parallel + examples/fodo_userdef/analysis_fodo.py + examples/fodo_userdef/plot_fodo.py +) + # Python: MPI-parallel FODO Cell ############################################## # add_impactx_test(FODO.py.MPI @@ -1056,3 +1096,19 @@ add_impactx_test(linac-segment.py examples/linac_segment/analysis_linac_segment.py OFF # no plot script yet ) + +# Iteration of a linear one-turn map ######################################### +# +# w/o space charge +add_impactx_test(linear-map + examples/linear_map/input_map.in + ON # ImpactX MPI-parallel + examples/linear_map/analysis_map.py + OFF # no plot script yet +) +add_impactx_test(linear-map.py + examples/linear_map/run_map.py + ON # ImpactX MPI-parallel + examples/linear_map/analysis_map.py + OFF # no plot script yet +) diff --git a/examples/fodo/analysis_fodo.py b/examples/fodo/analysis_fodo.py index f4e3e6b89..cb4b4dcdd 100755 --- a/examples/fodo/analysis_fodo.py +++ b/examples/fodo/analysis_fodo.py @@ -38,7 +38,8 @@ def get_moments(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() +beam_final = series.iterations[last_step].particles["beam"] +final = beam_final.to_df() # compare number of particles num_particles = 10000 @@ -74,9 +75,12 @@ def get_moments(beam): print("") print("Final Beam:") sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +s_ref = beam_final.get_attribute("s_ref") +gamma_ref = beam_final.get_attribute("gamma_ref") 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}" + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}\n" + f" s_ref={s_ref:e} gamma_ref={gamma_ref:e}" ) atol = 0.0 # ignored @@ -84,7 +88,7 @@ def get_moments(beam): print(f" rtol={rtol} (ignored: atol~={atol})") assert np.allclose( - [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t, s_ref, gamma_ref], [ 7.4790118496224206e-005, 7.5357525169680140e-005, @@ -92,6 +96,8 @@ def get_moments(beam): 1.9959539836392703e-009, 2.0175014668882125e-009, 2.0013820380883801e-006, + 3.000000, + 3.914902e003, ], rtol=rtol, atol=atol, diff --git a/examples/fodo_userdef/README.rst b/examples/fodo_userdef/README.rst new file mode 100644 index 000000000..657c94071 --- /dev/null +++ b/examples/fodo_userdef/README.rst @@ -0,0 +1,85 @@ +.. _examples-fodo-userdef: + +User-Defined Linear Element +=========================== + +This implements the same FODO cell as the :ref:`stable FODO cell example `. +However, in the example here we define *additional user-defined, custom linear elements* by providing a custom matrix. + +.. note:: + + Note that generally, if a user-provided linear map is used, the beam transport may not be symplectic. + For an even more general, user-defined element, see :ref:`the FODO Cell example that uses a Programmable Element `. + For more details, see :ref:`this section `. + +The matched Twiss parameters at entry are: + +* :math:`\beta_\mathrm{x} = 2.82161941` m +* :math:`\alpha_\mathrm{x} = -1.59050035` +* :math:`\beta_\mathrm{y} = 2.82161941` m +* :math:`\alpha_\mathrm{y} = 1.59050035` + +We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm. + +The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling. + +In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values. + + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_fodo_userdef.py`` or +* ImpactX **executable** using an input file: ``impactx input_fodo_userdef.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_fodo_userdef.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_userdef/run_fodo_userdef.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_fodo_userdef.in + :language: ini + :caption: You can copy this file from ``examples/fodo_userdef/input_fodo_userdef.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_fodo.py`` + + .. literalinclude:: analysis_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_userdef/analysis_fodo.py``. + + +Visualize +--------- + +You can run the following script to visualize the beam evolution over time: + +.. dropdown:: Script ``plot_fodo.py`` + + .. literalinclude:: plot_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo_userdef/plot_fodo.py``. + +.. figure:: https://user-images.githubusercontent.com/1353258/180287840-8561f6fd-278f-4856-abd8-04fbdb78c8ff.png + :alt: focusing, defocusing and preserved emittance in our FODO cell benchmark. + + FODO transversal beam width and emittance evolution + +.. figure:: https://user-images.githubusercontent.com/1353258/180287845-eb0210a7-2500-4aa9-844c-67fb094329d3.png + :alt: focusing, defocusing and phase space rotation in our FODO cell benchmark. + + FODO transversal beam width and phase space evolution diff --git a/examples/fodo_userdef/analysis_fodo.py b/examples/fodo_userdef/analysis_fodo.py new file mode 120000 index 000000000..dc8eb9737 --- /dev/null +++ b/examples/fodo_userdef/analysis_fodo.py @@ -0,0 +1 @@ +../fodo/analysis_fodo.py \ No newline at end of file diff --git a/examples/fodo_userdef/input_fodo_userdef.in b/examples/fodo_userdef/input_fodo_userdef.in new file mode 100644 index 000000000..420597477 --- /dev/null +++ b/examples/fodo_userdef/input_fodo_userdef.in @@ -0,0 +1,61 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 2.0e3 +beam.charge = 1.0e-9 +beam.particle = electron +beam.distribution = waterbag_from_twiss +beam.alphaX = -1.5905003499999992 +beam.alphaY = 1.5905003499999992 +beam.alphaT = 0.0 +beam.betaX = 2.8216194100262637 +beam.betaY = 2.8216194100262637 +beam.betaT = 0.5 +beam.emittX = 2e-09 +beam.emittY = 2e-09 +beam.emittT = 2e-06 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor drift1 monitor quad1 monitor drift2 monitor quad2 monitor drift1 monitor +lattice.nslice = 25 + +monitor.type = beam_monitor +monitor.backend = h5 + +drift1.type = linear_map +drift1.ds = 0.25 +drift1.R12 = 0.25 # ds +drift1.R34 = 0.25 # ds +drift1.R56 = 0.25 / 16.6464 # ds / (beta*gamma^2) + +quad1.type = quad +quad1.ds = 1.0 +quad1.k = 1.0 + +drift2.type = linear_map +drift2.ds = 0.5 +drift2.R12 = 0.5 # ds +drift2.R34 = 0.5 # ds +drift2.R56 = 0.5 / 16.6464 # ds / (beta*gamma^2) + +quad2.type = quad +quad2.ds = 1.0 +quad2.k = -1.0 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/fodo_userdef/plot_fodo.py b/examples/fodo_userdef/plot_fodo.py new file mode 120000 index 000000000..ce8494b77 --- /dev/null +++ b/examples/fodo_userdef/plot_fodo.py @@ -0,0 +1 @@ +../fodo/plot_fodo.py \ No newline at end of file diff --git a/examples/fodo_userdef/run_fodo_userdef.py b/examples/fodo_userdef/run_fodo_userdef.py new file mode 100755 index 000000000..dd15459ee --- /dev/null +++ b/examples/fodo_userdef/run_fodo_userdef.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2024 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Marco Garten +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, distribution, elements, twiss + +sim = ImpactX() + +# set numerical parameters and IO control +sim.particle_shape = 2 # B-spline order +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# load a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 2.0e3 # reference energy +bunch_charge_C = 1.0e-9 # used with space charge +npart = 10000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Waterbag( + **twiss( + beta_x=2.8216194100262637, + beta_y=2.8216194100262637, + beta_t=0.5, + emitt_x=2e-09, + emitt_y=2e-09, + emitt_t=2e-06, + alpha_x=-1.5905003499999992, + alpha_y=1.5905003499999992, + alpha_t=0.0, + ) +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# add a user-defined, linear element for the drifts +Iden = elements.LinearMap.Map6x6.identity() +R1, R2 = Iden, Iden + +ds1 = 0.25 +R1[1, 2] = ds1 +R1[3, 4] = ds1 +R1[5, 6] = ds1 / 16.6464 # ds / (beta*gamma^2) +drift1 = elements.LinearMap(name="drift1", R=R1, ds=ds1) + +ds2 = 0.5 +R2[1, 2] = ds2 +R2[3, 4] = ds2 +R2[5, 6] = ds2 / 16.6464 # ds / (beta*gamma^2) +drift2 = elements.LinearMap(name="drift2", R=R2, ds=ds2) + +# design the accelerator lattice) +ns = 25 # number of slices per ds in the element +fodo = [ + monitor, + drift1, + monitor, + elements.Quad(name="quad1", ds=1.0, k=1.0, nslice=ns), + monitor, + drift2, + monitor, + elements.Quad(name="quad2", ds=1.0, k=-1.0, nslice=ns), + monitor, + drift1, + monitor, +] +# assign a fodo segment +sim.lattice.extend(fodo) + +# run simulation +sim.evolve() + +# clean shutdown +sim.finalize() diff --git a/examples/linear_map/README.rst b/examples/linear_map/README.rst new file mode 100644 index 000000000..c6b72c746 --- /dev/null +++ b/examples/linear_map/README.rst @@ -0,0 +1,58 @@ +.. _examples-linear-map: + +Iteration of a User-Defined Linear Map +====================================== + +This example illustrates the application of a user-defined linear map via a matrix. + +Here, the linear map represents an abstract symplectic transformation of the beam in 6D phase space. +If desired, the user may interpret the matrix as the one-turn map of a storage ring or circular collider. + +The (fractional) tunes (Qx, Qy, Qt) of the map are given by (0.139, 0.219, 0.0250). +We use a 45.6 GeV electron beam that is invariant under the action of the linear map (matched). +The horizontal and vertical unnormalized emittances are 0.27 nm and 1.0 pm, respectively. + +These parameters are based on the `single-beam parameters of FCC-ee (Z-mode) `__. +(`backup `__). + +The second moments of the phase space variables should be unchanged under application of the map. + +In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values. + +In addition, the tunes associated with a single particle orbit are extracted, and must agree with the values given above. + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_map.py`` or +* ImpactX **executable** using an input file: ``impactx input_map.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_map.py + :language: python3 + :caption: You can copy this file from ``examples/linear_map/run_map.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_map.in + :language: ini + :caption: You can copy this file from ``examples/linear_map/input_map.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_map.py`` + + .. literalinclude:: analysis_map.py + :language: python3 + :caption: You can copy this file from ``examples/linear_map/analysis_map.py``. diff --git a/examples/linear_map/analysis_map.py b/examples/linear_map/analysis_map.py new file mode 100755 index 000000000..4de79069f --- /dev/null +++ b/examples/linear_map/analysis_map.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + +import numpy as np +import openpmd_api as io +from scipy.stats import moment + + +def get_moments(beam): + """Calculate standard deviations of beam position & momenta + and emittance values + + Returns + ------- + sigx, sigy, sigt, emittance_x, emittance_y, emittance_t + """ + sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev. + sigpx = moment(beam["momentum_x"], moment=2) ** 0.5 + sigy = moment(beam["position_y"], moment=2) ** 0.5 + sigpy = moment(beam["momentum_y"], moment=2) ** 0.5 + sigt = moment(beam["position_t"], moment=2) ** 0.5 + sigpt = moment(beam["momentum_t"], moment=2) ** 0.5 + + epstrms = beam.cov(ddof=0) + emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5 + emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5 + emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5 + + return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t) + + +# initial/final beam +series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) +last_step = list(series.iterations)[-1] +initial = series.iterations[1].particles["beam"].to_df() +final = series.iterations[last_step].particles["beam"].to_df() + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +print("Initial Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 6.363961030678928e-6, + 28.284271247461902e-9, + 0.0035, + 0.27e-9, + 1.0e-12, + 1.33e-6, + ], + rtol=rtol, + atol=atol, +) + +print("") +print("Final Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, emittance_x, emittance_y, emittance_t], + [ + 6.363961030678928e-6, + 28.284271247461902e-9, + 0.0035, + 0.27e-9, + 1.0e-12, + 1.33e-6, + ], + rtol=rtol, + atol=atol, +) + +# Specify time series for particle j +j = 5 +print(f"output for particle index = {j}") + +# Create array of TBT data values +x = [] +px = [] +y = [] +py = [] +t = [] +pt = [] +n = 0 +for k_i, i in series.iterations.items(): + beam = i.particles["beam"] + turn = beam.to_df() + x.append(turn["position_x"][j]) + px.append(turn["momentum_x"][j]) + y.append(turn["position_y"][j]) + py.append(turn["momentum_y"][j]) + t.append(turn["position_t"][j]) + pt.append(turn["momentum_t"][j]) + n = n + 1 + +# Output number of periods in data series +nturns = len(x) +print(f"number of periods = {nturns}") +print() + +# Approximate the tune and closed orbit using the 4-turn formula: + +# from x data only +argument = (x[0] - x[1] + x[2] - x[3]) / (2.0 * (x[1] - x[2])) +tunex = np.arccos(argument) / (2.0 * np.pi) +print(f"tune output from 4-turn formula, using x data = {tunex}") + +# from y data only +argument = (y[0] - y[1] + y[2] - y[3]) / (2.0 * (y[1] - y[2])) +tuney = np.arccos(argument) / (2.0 * np.pi) +print(f"tune output from 4-turn formula, using y data = {tuney}") + +# from t data only +argument = (t[0] - t[1] + t[2] - t[3]) / (2.0 * (t[1] - t[2])) +tunet = np.arccos(argument) / (2.0 * np.pi) +print(f"tune output from 4-turn formula, using t data = {tunet}") + +rtol = 1.0e-3 +print(f" rtol={rtol}") + +assert np.allclose( + [tunex, tuney, tunet], + [ + 0.139, + 0.219, + 0.0250, + ], + rtol=rtol, +) diff --git a/examples/linear_map/input_map.in b/examples/linear_map/input_map.in new file mode 100644 index 000000000..67fed5909 --- /dev/null +++ b/examples/linear_map/input_map.in @@ -0,0 +1,58 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 45.6e3 +beam.charge = 2.72370027e-8 #population 1.7e11 +beam.particle = electron +beam.distribution = waterbag_from_twiss +beam.alphaX = 0.0 +beam.alphaY = 0.0 +beam.alphaT = 0.0 +beam.betaX = 0.15 +beam.betaY = 0.8e-3 +beam.betaT = 9.210526315789473 +beam.emittX = 0.27e-09 +beam.emittY = 1e-12 +beam.emittT = 1.33e-6 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.periods = 5 +lattice.elements = monitor map1 + +monitor.type = beam_monitor +monitor.backend = h5 + +map1.type = linear_map +# horizontal plane +map1.R11 = 0.642252653176584 +map1.R12 = 0.114973951021402 +map1.R21 = -5.109953378728999 +map1.R22 = 0.642252653176584 +# vertical plane +map1.R33 = 0.193549468050860 +map1.R34 = 0.0007848724139547 +map1.R43 = -1226.363146804167548 +map1.R44 = 0.193549468050860 +# longitudinal plane +map1.R55 = 0.987688340595138 +map1.R56 = 1.440843756949495 +map1.R65 = -0.016984313347225 +map1.R66 = 0.987688340595138 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = false diff --git a/examples/linear_map/run_map.py b/examples/linear_map/run_map.py new file mode 100755 index 000000000..e65a77c2d --- /dev/null +++ b/examples/linear_map/run_map.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +# from elements import LinearTransport +import numpy as np + +from impactx import ImpactX, distribution, elements, twiss + +sim = ImpactX() + +# set numerical parameters and IO control +sim.particle_shape = 2 # B-spline order +sim.space_charge = False +# sim.diagnostics = False # benchmarking +sim.slice_step_diagnostics = True + +# domain decomposition & space charge mesh +sim.init_grids() + +# load a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 45.6e3 # reference energy +bunch_charge_C = 1.0e-9 # used with space charge +npart = 10000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV) + +# target beta functions (m) +beta_star_x = 0.15 +beta_star_y = 0.8e-3 +beta_star_t = 9.210526315789473 + +# particle bunch +distr = distribution.Waterbag( + **twiss( + beta_x=beta_star_x, + beta_y=beta_star_y, + beta_t=beta_star_t, + emitt_x=0.27e-09, + emitt_y=1.0e-12, + emitt_t=1.33e-06, + alpha_x=0.0, + alpha_y=0.0, + alpha_t=0.0, + ) +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# initialize the linear map +Iden = elements.LinearMap.Map6x6.identity() +Rmat = Iden + +# desired tunes +Qx = 0.139 +Qy = 0.219 +Qt = 0.0250 + +# desired phase advance +phi_x = 2.0 * np.pi * Qx +phi_y = 2.0 * np.pi * Qy +phi_t = 2.0 * np.pi * Qt + +# matrix elements for the horizontal plane +Rmat[1, 1] = np.cos(phi_x) +Rmat[1, 2] = beta_star_x * np.sin(phi_x) +Rmat[2, 1] = -np.sin(phi_x) / beta_star_x +Rmat[2, 2] = np.cos(phi_x) +# matrix elements for the vertical plane +Rmat[3, 3] = np.cos(phi_y) +Rmat[3, 4] = beta_star_y * np.sin(phi_y) +Rmat[4, 3] = -np.sin(phi_y) / beta_star_y +Rmat[4, 4] = np.cos(phi_y) +# matrix elements for the longitudinal plane +Rmat[5, 5] = np.cos(phi_t) +Rmat[5, 6] = beta_star_t * np.sin(phi_t) +Rmat[6, 5] = -np.sin(phi_t) / beta_star_t +Rmat[6, 6] = np.cos(phi_t) + +# design the accelerator lattice +map = [ + monitor, + elements.LinearMap(R=Rmat), +] + +sim.lattice.extend(map) + +# number of periods through the lattice +sim.periods = 4 + +# run simulation +sim.evolve() + +# clean shutdown +sim.finalize() diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp index 0b0efbed9..97806f752 100644 --- a/src/initialization/InitDistribution.cpp +++ b/src/initialization/InitDistribution.cpp @@ -10,6 +10,7 @@ #include "initialization/InitDistribution.H" #include "ImpactX.H" +#include "particles/CovarianceMatrix.H" #include "particles/ImpactXParticleContainer.H" #include "particles/distribution/All.H" @@ -32,6 +33,60 @@ namespace impactx { + /** Ignore the shape of a distribution and use the 2nd moments to create a covariance matrix + */ + CovarianceMatrix + create_covariance_matrix ( + distribution::KnownDistributions const & distr + ) + { + // zero out the 6x6 matrix + CovarianceMatrix cv{}; + + // initialize from 2nd order beam moments + std::visit([&](auto&& distribution) { + // quick hack + using Distribution = std::remove_cv_t< std::remove_reference_t< decltype(distribution)> >; + if constexpr (std::is_same::value || + std::is_same::value) + { + throw std::runtime_error("Empty and Thermal type cannot create Covariance matrices!"); + } else { + amrex::ParticleReal lambdaX = distribution.m_lambdaX; + amrex::ParticleReal lambdaY = distribution.m_lambdaY; + amrex::ParticleReal lambdaT = distribution.m_lambdaT; + amrex::ParticleReal lambdaPx = distribution.m_lambdaPx; + amrex::ParticleReal lambdaPy = distribution.m_lambdaPy; + amrex::ParticleReal lambdaPt = distribution.m_lambdaPt; + amrex::ParticleReal muxpx = distribution.m_muxpx; + amrex::ParticleReal muypy = distribution.m_muypy; + amrex::ParticleReal mutpt = distribution.m_mutpt; + + // use distribution inputs to populate a 6x6 covariance matrix + amrex::ParticleReal denom_x = 1.0 - muxpx*muxpx; + cv(1,1) = lambdaX*lambdaX / denom_x; + cv(1,2) = -lambdaX*lambdaPx*muxpx / denom_x; + cv(2,1) = cv(1,2); + cv(2,2) = lambdaPx*lambdaPx / denom_x; + + amrex::ParticleReal denom_y = 1.0 - muypy*muypy; + cv(3,3) = lambdaY*lambdaY / denom_y; + cv(3,4) = -lambdaY*lambdaPy*muypy / denom_y; + cv(4,3) = cv(3,4); + cv(4,4) = lambdaPy*lambdaPy / denom_y; + + amrex::ParticleReal denom_t = 1.0 - mutpt*mutpt; + cv(5,5) = lambdaT*lambdaT / denom_t; + cv(5,6) = -lambdaT*lambdaPt*mutpt / denom_t; + cv(6,5) = cv(5,6); + cv(6,6) = lambdaPt*lambdaPt / denom_t; + + } + }, distr); + + return cv; + } + void ImpactX::add_particles ( amrex::ParticleReal bunch_charge, @@ -95,7 +150,7 @@ namespace impactx amrex::ParticleReal * const AMREX_RESTRICT py_ptr = py.data(); amrex::ParticleReal * const AMREX_RESTRICT pt_ptr = pt.data(); - using Distribution = std::remove_reference_t< std::remove_cv_t >; + using Distribution = std::remove_reference_t< std::remove_cv_t >; // TODO: switch order ov remove_ ...? initialization::InitSingleParticleData const init_single_particle_data( distribution, x_ptr, y_ptr, t_ptr, px_ptr, py_ptr, pt_ptr); amrex::ParallelForRNG(npart_this_proc, init_single_particle_data); diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index d106a4b86..932ecdc3d 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -9,6 +9,9 @@ */ #include "ImpactX.H" #include "particles/elements/All.H" +#include "particles/elements/mixin/lineartransport.H" + +#include #include #include @@ -455,6 +458,24 @@ namespace detail read_element(sub_element_name, m_lattice, nslice_default, mapsteps_default); } } + } else if (element_type == "linear_map") + { + auto a = detail::query_alignment(pp_element); + + amrex::ParticleReal ds = 0.0; + pp_element.queryAdd("ds", ds); + + elements::LinearTransport::Map6x6 transport_map = elements::LinearTransport::Map6x6::Identity(); + + // safe to ParmParse inputs for reproducibility + for (int i=1; i<=6; ++i) { + for (int j=1; j<=6; ++j) { + std::string name = "R" + std::to_string(i) + std::to_string(j); + pp_element.queryAddWithParser(name.c_str(), transport_map(i, j)); + } + } + + m_lattice.emplace_back(LinearMap(transport_map, ds, a["dx"], a["dy"], a["rotation_degree"]) ); } else { amrex::Abort("Unknown type for lattice element " + element_name + ": " + element_type); } diff --git a/src/particles/CovarianceMatrix.H b/src/particles/CovarianceMatrix.H new file mode 100644 index 000000000..23af3e5aa --- /dev/null +++ b/src/particles/CovarianceMatrix.H @@ -0,0 +1,24 @@ +/* Copyright 2022-2024 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_COVARIANCE_MATRIX_H +#define IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H + +#include +#include + + +namespace impactx +{ + /** this is a 6x6 matrix */ + using CovarianceMatrix = amrex::SmallMatrix; + +} // namespace impactx::distribution + +#endif // IMPACTX_DISTRIBUTION_COVARIANCE_MATRIX_H diff --git a/src/particles/PushAll.H b/src/particles/PushAll.H index 3fec0f3bc..453f1949d 100644 --- a/src/particles/PushAll.H +++ b/src/particles/PushAll.H @@ -51,6 +51,10 @@ namespace impactx element(ref_part); } + // push covariance matrix + // TODO + // note: decide what to do for elements that have no covariance matrix + // loop over refinement levels int const nLevel = pc.finestLevel(); for (int lev = 0; lev <= nLevel; ++lev) diff --git a/src/particles/elements/All.H b/src/particles/elements/All.H index fbd0edc10..a8c0acd5d 100644 --- a/src/particles/elements/All.H +++ b/src/particles/elements/All.H @@ -20,21 +20,22 @@ #include "ConstF.H" #include "DipEdge.H" #include "Drift.H" +#include "Empty.H" #include "ExactDrift.H" #include "ExactSbend.H" #include "Kicker.H" +#include "LinearMap.H" #include "Marker.H" #include "Multipole.H" -#include "Empty.H" #include "NonlinearLens.H" #include "PlaneXYRot.H" #include "Programmable.H" +#include "PRot.H" #include "Quad.H" #include "RFCavity.H" #include "Sbend.H" #include "ShortRF.H" #include "Sol.H" -#include "PRot.H" #include "SoftSol.H" #include "SoftQuad.H" #include "TaperedPL.H" @@ -62,6 +63,7 @@ namespace impactx ExactDrift, ExactSbend, Kicker, + LinearMap, Marker, Multipole, NonlinearLens, diff --git a/src/particles/elements/Drift.H b/src/particles/elements/Drift.H index ef1e889af..cdc4b8e1f 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -16,9 +16,11 @@ #include "mixin/thick.H" #include "mixin/named.H" #include "mixin/nofinalize.H" +#include "mixin/lineartransport.H" #include #include +#include #include #include @@ -130,8 +132,8 @@ namespace impactx * @param[in,out] refpart reference particle */ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE - void operator() (RefPart & AMREX_RESTRICT refpart) const { - + void operator() (RefPart & AMREX_RESTRICT refpart) const + { using namespace amrex::literals; // for _rt and _prt // assign input reference particle values @@ -159,7 +161,32 @@ namespace impactx // advance integrated path length refpart.s = s + slice_ds; + } + + /** This function returns the linear transport map. + * + * @returns 6x6 transport matrix + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + elements::LinearTransport::Map6x6 + transport_map (RefPart & AMREX_RESTRICT refpart) const + { + using namespace amrex::literals; // for _rt and _prt + + // length of the current slice + amrex::ParticleReal const slice_ds = m_ds / nslice(); + + // access reference particle values to find beta*gamma^2 + amrex::ParticleReal const pt_ref = refpart.pt; + amrex::ParticleReal const betgam2 = std::pow(pt_ref, 2) - 1.0_prt; + + // assign linear map matrix elements + elements::LinearTransport::Map6x6 R = elements::LinearTransport::Map6x6::Identity(); + R(1,2) = slice_ds; + R(3,4) = slice_ds; + R(5,6) = slice_ds / betgam2; + return R; } }; diff --git a/src/particles/elements/LinearMap.H b/src/particles/elements/LinearMap.H new file mode 100644 index 000000000..472ae12be --- /dev/null +++ b/src/particles/elements/LinearMap.H @@ -0,0 +1,182 @@ +/* Copyright 2022-2024 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_ELEMENT_LINEAR_MAP_H +#define IMPACTX_ELEMENT_LINEAR_MAP_H + +#include "particles/ImpactXParticleContainer.H" +#include "mixin/alignment.H" +#include "mixin/beamoptic.H" +#include "mixin/lineartransport.H" +#include "mixin/named.H" +#include "mixin/nofinalize.H" + +#include +#include + +#include + + +namespace impactx +{ + struct LinearMap + : public elements::Named, + public elements::BeamOptic, + public elements::Alignment, + public elements::LinearTransport, + public elements::NoFinalize + { + static constexpr auto type = "LinearMap"; + using PType = ImpactXParticleContainer::ParticleType; + + /** A thin element that applies a user-provided linear transport map R + * to the 6-vector of phase space coordinates (x,px,y,py,t,pt). + * Thus x_final = R(1,1)*x + R(1,2)*px + R(1,3)*y + ..., + * px_final = R(2,1)*x + R(2,2)*px + R(2,3)*y + ..., etc. + * + * @param R user-provided transport map + * @param ds Segment length in m + * @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 + */ + LinearMap ( + LinearTransport::Map6x6 const & R, + amrex::ParticleReal ds = 0, + 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_transport_map = R; + m_ds = ds; + } + + /** Push all particles */ + using BeamOptic::operator(); + + /** This is a LinearMap functor, so that a variable of this type can be used like a + * LinearMap 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) + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() ( + amrex::ParticleReal & AMREX_RESTRICT x, + amrex::ParticleReal & AMREX_RESTRICT y, + amrex::ParticleReal & AMREX_RESTRICT t, + amrex::ParticleReal & AMREX_RESTRICT px, + amrex::ParticleReal & AMREX_RESTRICT py, + amrex::ParticleReal & AMREX_RESTRICT pt, + [[maybe_unused]] uint64_t & AMREX_RESTRICT idcpu, + [[maybe_unused]] RefPart const & refpart + ) const + { + using namespace amrex::literals; // for _rt and _prt + + // shift due to alignment errors of the element + shift_in(x, y, px, py); + + // input and output phase space vectors + amrex::SmallVector vectorin{ + x, px, y, py, t, pt + }; + + amrex::SmallVector const vectorout = m_transport_map * vectorin; + + // assign updated values + x = vectorout(1); + px = vectorout(2); + y = vectorout(3); + py = vectorout(4); + t = vectorout(5); + pt = vectorout(6); + + // undo shift due to alignment errors of the element + shift_out(x, y, px, py); + } + + /** This pushes the reference particle. + * + * @param[in,out] refpart reference particle + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + void operator() (RefPart & AMREX_RESTRICT refpart) const + { + if (m_ds > 0) // Drift + { + using namespace amrex::literals; // for _rt and _prt + + // assign input reference particle values + amrex::ParticleReal const x = refpart.x; + amrex::ParticleReal const px = refpart.px; + amrex::ParticleReal const y = refpart.y; + amrex::ParticleReal const py = refpart.py; + amrex::ParticleReal const z = refpart.z; + amrex::ParticleReal const pz = refpart.pz; + amrex::ParticleReal const t = refpart.t; + amrex::ParticleReal const pt = refpart.pt; + amrex::ParticleReal const s = refpart.s; + + // length of the current slice + amrex::ParticleReal const slice_ds = m_ds / nslice(); + + // assign intermediate parameter + amrex::ParticleReal const step = slice_ds /std::sqrt(std::pow(pt,2)-1.0_prt); + + // advance position and momentum (drift) + refpart.x = x + step*px; + refpart.y = y + step*py; + refpart.z = z + step*pz; + refpart.t = t - step*pt; + + // advance integrated path length + refpart.s = s + slice_ds; + } + // else nothing to do for a zero-length element + } + + /** Number of slices used for the application of space charge + * + * @return one, because we do not support slicing of this element + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + int nslice () const + { + return 1; + } + + /** Return the segment length + * + * @return by default zero, but users can set a corresponding ds for bookkeeping + */ + AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE + amrex::ParticleReal ds () const + { + return m_ds; + } + + LinearTransport::Map6x6 m_transport_map; // 6x6 transport map + amrex::ParticleReal m_ds; // finite ds allowed for bookkeeping, but we do not allow slicing + }; + +} // namespace impactx + +#endif // IMPACTX_ELEMENT_LINEAR_MAP_H diff --git a/src/particles/elements/mixin/lineartransport.H b/src/particles/elements/mixin/lineartransport.H new file mode 100644 index 000000000..9fce44ff6 --- /dev/null +++ b/src/particles/elements/mixin/lineartransport.H @@ -0,0 +1,52 @@ +/* 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: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#ifndef IMPACTX_ELEMENTS_MIXIN_LINEAR_TRANSPORT_H +#define IMPACTX_ELEMENTS_MIXIN_LINEAR_TRANSPORT_H + +#include "particles/ImpactXParticleContainer.H" + +#include + +#include +#include +#include +#include + + +namespace impactx::elements +{ + /** This is a helper class for lattice elements that can be expressed as linear transport maps. + */ + struct LinearTransport + { + /** ... + */ + LinearTransport ( + ) + { + } + + //LinearTransport () = default; + LinearTransport (LinearTransport const &) = default; + LinearTransport& operator= (LinearTransport const &) = default; + LinearTransport (LinearTransport&&) = default; + LinearTransport& operator= (LinearTransport&& rhs) = default; + + ~LinearTransport () = default; + + // 6x6 linear transport map + using Map6x6 = amrex::SmallMatrix; + // note: for most elements, R is returned by a member function. Some store it also internally as a member. + // Map6x6 m_transport_map; ///< linearized map + }; + +} // namespace impactx::elements + +#endif // IMPACTX_ELEMENTS_MIXIN_LINEAR_TRANSPORT_H diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 7d1488a55..ce712c370 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include @@ -215,6 +216,22 @@ void init_elements(py::module& m) ) ; + py::class_(mx, "LinearTransport") + .def(py::init<>(), + "Mixin class for linear transport approximation via matrices." + ) + // type of map + .def_property_readonly_static("Map6x6", + [](py::object /* lt */){ return py::type::of(); }, + "1-indexed, Fortran-ordered, 6x6 linear transport map type" + ) + // values of the map + //.def_property_readonly("R", + // [](elements::LinearTransport const & lt) { return lt.m_transport_map; }, + // "1-indexed, Fortran-ordered, 6x6 linear transport map values" + //) + ; + // diagnostics py::class_ py_BeamMonitor(me, "BeamMonitor"); @@ -1556,6 +1573,47 @@ void init_elements(py::module& m) ; register_beamoptics_push(py_TaperedPL); + py::class_ py_LinearMap(me, "LinearMap"); + py_LinearMap + .def("__repr__", + [](LinearMap const & linearmap) { + return element_name( + linearmap + ); + } + ) + .def(py::init< + elements::LinearTransport::Map6x6, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + amrex::ParticleReal, + std::optional + >(), + py::arg("R"), + py::arg("ds") = 0, + py::arg("dx") = 0, + py::arg("dy") = 0, + py::arg("rotation") = 0, + py::arg("name") = py::none(), + "(A user-provided linear map, represented as a 6x6 transport matrix.)" + ) + .def_property("R", + [](LinearMap & linearmap) { return linearmap.m_transport_map; }, + [](LinearMap & linearmap, elements::LinearTransport::Map6x6 R) { linearmap.m_transport_map = R; }, + "linear map as a 6x6 transport matrix" + ) + .def_property("ds", + [](LinearMap & linearmap) { return linearmap.m_ds; }, + [](LinearMap & linearmap, amrex::ParticleReal ds) { linearmap.m_ds = ds; }, + "segment length in m" + ) + .def_property_readonly("nslice", + [](LinearMap & linearmap) { return linearmap.nslice(); }, + "one, because we do not support slicing of this element" + ) + ; + register_beamoptics_push(py_LinearMap); // freestanding push function m.def("push", &Push,