diff --git a/.github/workflows/dependencies/hip-openmpi.sh b/.github/workflows/dependencies/hip-openmpi.sh index eea95c482..a6e90e0d6 100755 --- a/.github/workflows/dependencies/hip-openmpi.sh +++ b/.github/workflows/dependencies/hip-openmpi.sh @@ -54,6 +54,7 @@ sudo apt-get install -y --no-install-recommends \ rocfft-dev \ rocprim-dev \ rocrand-dev \ + rocsparse-dev \ hiprand-dev # activate diff --git a/.github/workflows/dependencies/nvcc11-openmpi.sh b/.github/workflows/dependencies/nvcc11-openmpi.sh index 74a20c74d..5564b47ec 100755 --- a/.github/workflows/dependencies/nvcc11-openmpi.sh +++ b/.github/workflows/dependencies/nvcc11-openmpi.sh @@ -35,7 +35,8 @@ sudo apt-get install -y \ cuda-nvml-dev-11-3 \ cuda-nvtx-11-3 \ libcufft-dev-11-3 \ - libcurand-dev-11-3 + libcurand-dev-11-3 \ + libcusparse-dev-11-3 sudo ln -s cuda-11.3 /usr/local/cuda # cmake-easyinstall diff --git a/.github/workflows/windows.yml b/.github/workflows/windows.yml index 3e10db72c..66dc54d59 100644 --- a/.github/workflows/windows.yml +++ b/.github/workflows/windows.yml @@ -44,7 +44,8 @@ jobs: # https://github.com/actions/runner-images/issues/10004 CXXFLAGS: "/D_DISABLE_CONSTEXPR_MUTEX_CONSTRUCTOR" run: | - cmake -S fftw-3.3.10 -B build_fftw ` + cmake -S fftw-3.3.10 ` + -B build_fftw ` -DBUILD_SHARED_LIBS=OFF ` -DBUILD_TESTS=OFF ` -DCMAKE_BUILD_TYPE=RelWithDebInfo ` @@ -52,6 +53,16 @@ jobs: if(!$?) { Exit $LASTEXITCODE } cmake --build build_fftw --config RelWithDebInfo --target install --parallel 4 if(!$?) { Exit $LASTEXITCODE } + cmake -S fftw-3.3.10 ` + -B build_fftwf ` + -DBUILD_SHARED_LIBS=OFF ` + -DBUILD_TESTS=OFF ` + -DCMAKE_BUILD_TYPE=RelWithDebInfo ` + -DDISABLE_FORTRAN=ON ` + -DENABLE_FLOAT=ON + if(!$?) { Exit $LASTEXITCODE } + cmake --build build_fftwf --config RelWithDebInfo --target install --parallel 4 + if(!$?) { Exit $LASTEXITCODE } cmake -S hdf5-hdf5-1_12_2 -B build_hdf5 ` -DCMAKE_BUILD_TYPE=RelWithDebInfo ` @@ -83,6 +94,7 @@ jobs: CXXFLAGS: "/D_DISABLE_CONSTEXPR_MUTEX_CONSTRUCTOR" # FFTW3 install prefix path FFTW3_DIR: "C:/Program Files (x86)/fftw/" + FFTW3f_DIR: "C:/Program Files (x86)/fftw/" run: | $env:HDF5_DIR = "C:/Program Files/HDF_Group/HDF5/1.12.2/cmake/" cmake -S . -B build ` @@ -159,6 +171,17 @@ jobs: if errorlevel 1 exit 1 cmake --build build_fftw --config Release --target install --parallel 4 if errorlevel 1 exit 1 + cmake -S fftw-3.3.10 ^ + -B build_fftwf ^ + -G "Ninja" ^ + -DBUILD_TESTS=OFF ^ + -DBUILD_SHARED_LIBS=OFF ^ + -DCMAKE_BUILD_TYPE=Release ^ + -DDISABLE_FORTRAN=ON ^ + -DENABLE_FLOAT=ON + if errorlevel 1 exit 1 + cmake --build build_fftwf --config Release --target install --parallel 4 + if errorlevel 1 exit 1 cmake -S hdf5-hdf5-1_12_2 -B build_hdf5 ^ -G "Ninja" ^ @@ -191,6 +214,7 @@ jobs: env: # FFTW3 install prefix path FFTW3_DIR: "C:/Program Files (x86)/fftw/" + FFTW3f_DIR: "C:/Program Files (x86)/fftw/" run: | call "C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\vc\Auxiliary\build\vcvarsall.bat" x64 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a06b195cb..661ae7c66 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.8.6 + rev: v0.9.1 hooks: # Run the linter - id: ruff diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ac042f7b..308e7bc1c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,7 +1,7 @@ # Preamble #################################################################### # cmake_minimum_required(VERSION 3.24.0) -project(ImpactX VERSION 24.12) +project(ImpactX VERSION 25.01) include(${ImpactX_SOURCE_DIR}/cmake/ImpactXFunctions.cmake) @@ -265,7 +265,7 @@ if(ImpactX_PYTHON) impactx_enable_IPO(pyImpactX) else() # conditionally defined target in pybind11 - # https://github.com/pybind/pybind11/blob/v2.12.0/tools/pybind11Common.cmake#L397-L403 + # https://github.com/pybind/pybind11/blob/v2.13.0/tools/pybind11Common.cmake#L407-L413 target_link_libraries(pyImpactX PRIVATE pybind11::lto) endif() endif() diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index 588a74c38..5768d5ad1 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -58,6 +58,7 @@ macro(find_ablastr) set(CMAKE_POLICY_DEFAULT_CMP0077 NEW) set(ABLASTR_FFT ${ImpactX_FFT} CACHE BOOL "" FORCE) + set(AMReX_FFT ${ImpactX_FFT} CACHE BOOL "" FORCE) set(WarpX_APP OFF CACHE BOOL "" FORCE) set(WarpX_LIB OFF CACHE BOOL "" FORCE) @@ -127,7 +128,7 @@ macro(find_ablastr) set(COMPONENT_DIM 3D) set(COMPONENT_PRECISION ${ImpactX_PRECISION} P${ImpactX_PRECISION}) - find_package(ABLASTR 24.10 CONFIG REQUIRED COMPONENTS ${COMPONENT_DIM}) + find_package(ABLASTR 25.01 CONFIG REQUIRED COMPONENTS ${COMPONENT_DIM}) message(STATUS "ABLASTR: Found version '${ABLASTR_VERSION}'") endif() @@ -161,7 +162,7 @@ set(ImpactX_openpmd_src "" set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git" CACHE STRING "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)") -set(ImpactX_ablastr_branch "24.10" +set(ImpactX_ablastr_branch "25.01" CACHE STRING "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)") @@ -169,7 +170,7 @@ set(ImpactX_ablastr_branch "24.10" 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 "e64ffef57a7608d1d60f9abe738cc634e9c1272e" +set(ImpactX_amrex_branch "f6d3d70cadfc03036f5fdc4f89b7fcb225792716" 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 e9c27b6b4..7632c88b9 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 24.10 CONFIG REQUIRED) + find_package(pyAMReX 25.01 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 "8742a79c29b7db30c7287c8e33109c0d8be1277a" +set(ImpactX_pyamrex_branch "25.01" CACHE STRING "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)") diff --git a/cmake/dependencies/pybind11.cmake b/cmake/dependencies/pybind11.cmake index 3a6c01e99..a4462cfce 100644 --- a/cmake/dependencies/pybind11.cmake +++ b/cmake/dependencies/pybind11.cmake @@ -37,7 +37,7 @@ function(find_pybind11) mark_as_advanced(FETCHCONTENT_UPDATES_DISCONNECTED_FETCHEDpybind11) endif() else() - find_package(pybind11 2.12.0 CONFIG REQUIRED) + find_package(pybind11 2.13.0 CONFIG REQUIRED) message(STATUS "pybind11: Found version '${pybind11_VERSION}'") endif() endfunction() @@ -52,7 +52,7 @@ option(ImpactX_pybind11_internal "Download & build pybind11" ON) set(ImpactX_pybind11_repo "https://github.com/pybind/pybind11.git" CACHE STRING "Repository URI to pull and build pybind11 from if(ImpactX_pybind11_internal)") -set(ImpactX_pybind11_branch "v2.12.0" +set(ImpactX_pybind11_branch "v2.13.6" CACHE STRING "Repository branch for ImpactX_pybind11_repo if(ImpactX_pybind11_internal)") diff --git a/docs/source/conf.py b/docs/source/conf.py index 2101a45f3..61210e106 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -73,9 +73,9 @@ # built documents. # # The short X.Y version. -version = "24.12" +version = "25.01" # The full version, including alpha/beta/rc tags. -release = "24.12" +release = "25.01" # 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 7dd594a58..55313c347 100644 --- a/docs/source/install/dependencies.rst +++ b/docs/source/install/dependencies.rst @@ -28,7 +28,7 @@ Optional dependencies include: - see `optional I/O backends `__ - `CCache `__: to speed up rebuilds (For CUDA support, needs version 3.7.9+ and 4.2+ is recommended) - `Ninja `__: for faster parallel compiles -- `Python 3.8+ `__ +- `Python 3.9+ `__ - `mpi4py `__ - `numpy `__ 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 4d4c5be6b..fab2851d7 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -351,6 +351,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: @@ -815,7 +834,7 @@ Overall simulation parameters * ``amrex.the_arena_is_managed`` (``0`` or ``1``; default is ``0`` for false) When running on GPUs, device memory that is accessed from the host will automatically be transferred with managed memory. This is useful for convenience during development, but has sometimes severe performance and memory footprint implications if relied on (and sometimes vendor bugs). - For all regular ImpactX operations, we therefore do explicit memory transfers without the need for managed memory and thus changed the AMReX default to false. + For all regular ImpactX operations, we therefore do explicit memory transfers without the need for managed memory. `Please also see the documentation in AMReX `__. * ``amrex.omp_threads`` (``system``, ``nosmt`` or positive integer; default is ``nosmt``) diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 8bdcd1dbe..f782fe4a4 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -616,6 +616,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 7610f93cf..a7ae69480 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 @@ -1072,3 +1112,19 @@ add_impactx_test(aperture-thick.py examples/aperture/analysis_aperture_thick.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/aperture/input_aperture.in b/examples/aperture/input_aperture.in index cb28dee88..4a1c32a4d 100644 --- a/examples/aperture/input_aperture.in +++ b/examples/aperture/input_aperture.in @@ -35,9 +35,6 @@ collimator.shape = rectangular collimator.xmax = 1.0e-3 collimator.ymax = 1.5e-3 -# work-around for https://github.com/ECP-WarpX/impactx/issues/499 -amrex.the_arena_is_managed = 1 - ############################################################################### # Algorithms diff --git a/examples/aperture/run_absorber.py b/examples/aperture/run_absorber.py index 4c75592ae..e93c46026 100755 --- a/examples/aperture/run_absorber.py +++ b/examples/aperture/run_absorber.py @@ -6,13 +6,8 @@ # # -*- coding: utf-8 -*- -import amrex.space3d as amr from impactx import ImpactX, distribution, elements -# work-around for https://github.com/ECP-WarpX/impactx/issues/499 -pp_amrex = amr.ParmParse("amrex") -pp_amrex.add("the_arena_is_managed", 1) - sim = ImpactX() # set numerical parameters and IO control diff --git a/examples/aperture/run_aperture.py b/examples/aperture/run_aperture.py index b8b783edb..1d186db5e 100755 --- a/examples/aperture/run_aperture.py +++ b/examples/aperture/run_aperture.py @@ -6,13 +6,8 @@ # # -*- coding: utf-8 -*- -import amrex.space3d as amr from impactx import ImpactX, distribution, elements -# work-around for https://github.com/ECP-WarpX/impactx/issues/499 -pp_amrex = amr.ParmParse("amrex") -pp_amrex.add("the_arena_is_managed", 1) - sim = ImpactX() # set numerical parameters and IO control diff --git a/examples/aperture/run_aperture_periodic.py b/examples/aperture/run_aperture_periodic.py index 9523abcb5..c0568718d 100755 --- a/examples/aperture/run_aperture_periodic.py +++ b/examples/aperture/run_aperture_periodic.py @@ -6,13 +6,8 @@ # # -*- coding: utf-8 -*- -import amrex.space3d as amr from impactx import ImpactX, distribution, elements -# work-around for https://github.com/ECP-WarpX/impactx/issues/499 -pp_amrex = amr.ParmParse("amrex") -pp_amrex.add("the_arena_is_managed", 1) - sim = ImpactX() # set numerical parameters and IO control 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/initialize_from_array/transformation_utilities.py b/examples/initialize_from_array/transformation_utilities.py index adb33b2aa..8e92bc918 100644 --- a/examples/initialize_from_array/transformation_utilities.py +++ b/examples/initialize_from_array/transformation_utilities.py @@ -183,11 +183,11 @@ def __init__(self, x, y, z, px, py, pz, pt): def __repr__(self): mystr = "" for attr in self.attr_list: - mystr += f"self.{attr}={getattr(self,attr)}, " + mystr += f"self.{attr}={getattr(self, attr)}, " return mystr def __str__(self): mystr = "" for attr in self.attr_list: - mystr += f"self.{attr}={getattr(self,attr)}, " + mystr += f"self.{attr}={getattr(self, attr)}, " return mystr 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/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py index a49a024e3..3539993ae 100644 --- a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py +++ b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py @@ -276,9 +276,9 @@ def surrogate_push(self, pc, step, period): def get_lattice_element_iter(sim, j): - assert ( - 0 <= j < len(sim.lattice) - ), f"Argument j must be a nonnegative integer satisfying 0 <= j < {len(sim.lattice)}, not {j}" + assert 0 <= j < len(sim.lattice), ( + f"Argument j must be a nonnegative integer satisfying 0 <= j < {len(sim.lattice)}, not {j}" + ) i = 0 lat_it = sim.lattice.__iter__() next(lat_it) diff --git a/examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py b/examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py index a7ffd5097..6fada2371 100644 --- a/examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py +++ b/examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py @@ -88,9 +88,9 @@ def to_t( dpt = data_arr_s["momentum_t"] elif type(data_arr_s) is np.ndarray: - assert ( - data_arr_s.shape[1] == 6 - ), f"data_arr_s.shape={data_arr_s.shape} but data_arr_s must be an Nx6 array" + assert data_arr_s.shape[1] == 6, ( + f"data_arr_s.shape={data_arr_s.shape} but data_arr_s must be an Nx6 array" + ) dx, dy, dt, dpx, dpy, dpt = data_arr_s.T else: raise Exception( @@ -111,9 +111,9 @@ def to_t( print("applying reference normalization") dpt /= ref_pz elif coord_type is TCoords.GLOBAL: - assert ( - ref_z is not None - ), "Reference particle z coordinate is required to transform to global coordinates" + assert ref_z is not None, ( + "Reference particle z coordinate is required to transform to global coordinates" + ) print("target global coordinates") dt += ref_z dpx *= ref_pz @@ -422,7 +422,7 @@ def plot_beam_df( t_offset = impactx_surrogate_ref_particle.loc[step, "t"] * micron fig, axT = plt.subplots(3, 3, figsize=(10, 8)) -fig.suptitle(f"initially, ct={impactx_surrogate_ref_particle.at[step,'t']:.2f} m") +fig.suptitle(f"initially, ct={impactx_surrogate_ref_particle.at[step, 't']:.2f} m") plot_beam_df( beam_at_step, @@ -459,7 +459,7 @@ def plot_beam_df( t_offset = impactx_surrogate_ref_particle.loc[step, "t"] * micron fig, axT = plt.subplots(3, 3, figsize=(10, 8)) fig.suptitle( - f"stage {stage_i}, ct={impactx_surrogate_ref_particle.at[step,'t']:.2f} m" + f"stage {stage_i}, ct={impactx_surrogate_ref_particle.at[step, 't']:.2f} m" ) plot_beam_df( @@ -473,6 +473,6 @@ def plot_beam_df( z_ticks=[-107.3, -106.6], ) if args.save_png: - plt.savefig(f"stage_{stage_i-1}_phase_spaces.png") + plt.savefig(f"stage_{stage_i - 1}_phase_spaces.png") else: plt.show() diff --git a/setup.py b/setup.py index 2ddfa44d7..fa49892d6 100644 --- a/setup.py +++ b/setup.py @@ -230,7 +230,7 @@ def build_extension(self, ext): setup( name="impactx", # note PEP-440 syntax: x.y.zaN but x.y.z.devN - version="24.12", + version="25.01", packages=["impactx"], # Python sources: package_dir={"": "src/python"}, @@ -260,7 +260,7 @@ def build_extension(self, ext): ext_modules=cxx_modules, cmdclass=cmdclass, zip_safe=False, - python_requires=">=3.8", + python_requires=">=3.8", # left for CI, truly ">=3.9" tests_require=["numpy", "pandas", "pytest", "scipy"], install_requires=install_requires, # cmdclass={'test': PyTest}, @@ -278,12 +278,13 @@ def build_extension(self, ext): "Topic :: Software Development :: Libraries", "Programming Language :: C++", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", ( - "License :: OSI Approved :: " "BSD License" + "License :: OSI Approved :: BSD License" ), # TODO: use real SPDX: BSD-3-Clause-LBNL ], # new PEP 639 format 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 3913a6d58..5c59f3a33 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 @@ -479,6 +482,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 129f3412f..f08e10292 100644 --- a/src/particles/elements/Drift.H +++ b/src/particles/elements/Drift.H @@ -17,9 +17,11 @@ #include "mixin/thick.H" #include "mixin/named.H" #include "mixin/nofinalize.H" +#include "mixin/lineartransport.H" #include #include +#include #include #include @@ -140,8 +142,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 @@ -169,7 +171,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/particles/spacecharge/PoissonSolve.cpp b/src/particles/spacecharge/PoissonSolve.cpp index d9c9b3dd2..a414604b9 100644 --- a/src/particles/spacecharge/PoissonSolve.cpp +++ b/src/particles/spacecharge/PoissonSolve.cpp @@ -80,6 +80,7 @@ namespace impactx::spacecharge sorted_phi.emplace_back(&phi[lev]); } + const bool is_igf_2d = false; const bool do_single_precision_comms = false; const bool eb_enabled = false; ablastr::fields::computePhi( @@ -95,6 +96,7 @@ namespace impactx::spacecharge pc.GetParGDB()->boxArray(), ablastr::utils::enums::GridType::Collocated, is_solver_igf_on_lev0, + is_igf_2d, eb_enabled, do_single_precision_comms, rel_ref_ratio diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 46a80bb8f..56ec5bc43 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include @@ -215,7 +216,7 @@ void init_elements(py::module& m) ) ; - py::class_(mx, "Aperture") + py::class_(mx, "Aperture") .def(py::init<>(), "Mixin class for lattice elements with a transverse aperture." ) @@ -229,6 +230,21 @@ 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 @@ -1579,6 +1595,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, diff --git a/src/python/impactx/__init__.pyi b/src/python/impactx/__init__.pyi index 1e931b0d8..9d887e699 100644 --- a/src/python/impactx/__init__.pyi +++ b/src/python/impactx/__init__.pyi @@ -79,7 +79,7 @@ __author__: str = ( "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." ) __license__: str = "BSD-3-Clause-LBNL" -__version__: str = "24.12" +__version__: str = "25.01" s: impactx_pybind.CoordSystem # value = t: impactx_pybind.CoordSystem # value = cxx = impactx_pybind diff --git a/src/python/impactx/dashboard/Input/__init__.py b/src/python/impactx/dashboard/Input/__init__.py index e69de29bb..a0875b1bb 100644 --- a/src/python/impactx/dashboard/Input/__init__.py +++ b/src/python/impactx/dashboard/Input/__init__.py @@ -0,0 +1,14 @@ +from trame.widgets import vuetify as vuetify + +from ..trame_setup import setup_server +from .defaults import DashboardDefaults +from .generalFunctions import generalFunctions +from .trameFunctions import TrameFunctions + +__all__ = [ + "vuetify", + "DashboardDefaults", + "TrameFunctions", + "generalFunctions", + "setup_server", +] diff --git a/src/python/impactx/dashboard/Input/csrConfiguration/csrMain.py b/src/python/impactx/dashboard/Input/csrConfiguration/csrMain.py index 83b5e0710..924bfeddb 100644 --- a/src/python/impactx/dashboard/Input/csrConfiguration/csrMain.py +++ b/src/python/impactx/dashboard/Input/csrConfiguration/csrMain.py @@ -1,34 +1,7 @@ -from trame.widgets import vuetify - -from ...Input.trameFunctions import TrameFunctions -from ...trame_setup import setup_server -from ..generalFunctions import generalFunctions +from .. import TrameFunctions, setup_server, vuetify server, state, ctrl = setup_server() -# ----------------------------------------------------------------------------- -# Default State Variables -# ----------------------------------------------------------------------------- - -state.csr_bins = generalFunctions.get_default("csr_bins", "default_values") -state.csr_bins_error_message = "" - -# ----------------------------------------------------------------------------- -# -# ----------------------------------------------------------------------------- - - -@state.change("csr_bins") -def on_csr_bins_change(csr_bins, **kwargs): - error_message = generalFunctions.validate_against(csr_bins, "int", ["positive"]) - state.csr_bins_error_message = error_message - generalFunctions.update_simulation_validation_status() - - -# ----------------------------------------------------------------------------- -# UI -# ----------------------------------------------------------------------------- - class csrConfiguration: @staticmethod @@ -38,35 +11,16 @@ def card(): """ with vuetify.VCard(v_show="csr", style="width: 170px;"): - with vuetify.VCardTitle("CSR"): - vuetify.VSpacer() - TrameFunctions.create_refresh_button( - lambda: generalFunctions.reset_inputs("csr") - ) - vuetify.VIcon( - "mdi-information", - classes="ml-2", - click=lambda: generalFunctions.documentation("CSR"), - style="color: #00313C;", - ) - vuetify.VDivider() + TrameFunctions.input_section_header("CSR") with vuetify.VCardText(): with vuetify.VRow(classes="my-0"): with vuetify.VCol(classes="py-0"): - vuetify.VSelect( + TrameFunctions.select( label="Particle Shape", - v_model=("particle_shape",), - items=([1, 2, 3],), - dense=True, ) with vuetify.VRow(classes="my-0"): with vuetify.VCol(classes="py-0"): - vuetify.VTextField( + TrameFunctions.text_field( label="CSR Bins", - v_model=("csr_bins",), - error_messages=("csr_bins_error_message",), - type="number", - dense=True, - step=generalFunctions.get_default("csr_bins", "steps"), - __properties=["step"], + input=(ctrl.input_change, "['csr_bins']"), ) diff --git a/src/python/impactx/dashboard/Input/defaults.py b/src/python/impactx/dashboard/Input/defaults.py index 856833ac7..97a34706c 100644 --- a/src/python/impactx/dashboard/Input/defaults.py +++ b/src/python/impactx/dashboard/Input/defaults.py @@ -17,16 +17,17 @@ class DashboardDefaults: "mass_MeV": 0.51099895, "npart": 1000, "kin_energy": 2e3, + "kin_energy_MeV": 2e3, "kin_energy_unit": "MeV", "bunch_charge_C": 1e-9, } - DISTRIBUTION = { - "selected_distribution": "Waterbag", - "selected_distribution_type": "Twiss", + DISTRIBUTION_PARAMETERS = { + "distribution": "Waterbag", + "distribution_type": "Twiss", } - LATTICE = { + LATTICE_CONFIGURATION = { "selected_lattice_list": [], "selected_lattice": None, } @@ -70,13 +71,28 @@ class DashboardDefaults: DEFAULT_VALUES = { **SELECTION, **INPUT_PARAMETERS, - **DISTRIBUTION, - **LATTICE, + **DISTRIBUTION_PARAMETERS, + **LATTICE_CONFIGURATION, **SPACE_CHARGE, **CSR, **LISTS, } + TYPES = { + "npart": "int", + "kin_energy": "float", + "bunch_charge_C": "float", + "mass_MeV": "float", + "charge_qe": "int", + "csr_bins": "int", + } + + VALIDATION_CONDITION = { + "charge_qe": ["non_zero"], + "mass_MeV": ["positive"], + "csr_bins": ["positive"], + } + # If a parameter is not included in the dictionary, default step amount is 1. STEPS = { "mass_MeV": 0.1, diff --git a/src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py b/src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py index a37b8a32f..741d13e9c 100644 --- a/src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py +++ b/src/python/impactx/dashboard/Input/distributionParameters/distributionMain.py @@ -9,13 +9,16 @@ import inspect from distribution_input_helpers import twiss -from trame.widgets import vuetify from impactx import distribution -from ...Input.trameFunctions import TrameFunctions -from ...trame_setup import setup_server -from ..generalFunctions import generalFunctions +from .. import ( + DashboardDefaults, + TrameFunctions, + generalFunctions, + setup_server, + vuetify, +) from .distributionFunctions import DistributionFunctions server, state, ctrl = setup_server() @@ -24,39 +27,28 @@ # Helpful # ----------------------------------------------------------------------------- -DISTRIBUTIONS_MODULE_NAME = distribution - -state.listOfDistributions = generalFunctions.select_classes(DISTRIBUTIONS_MODULE_NAME) -state.listOfDistributionsAndParametersAndDefault = ( - generalFunctions.class_parameters_with_defaults(DISTRIBUTIONS_MODULE_NAME) +DISTRIBUTION_MODULE_NAME = distribution +DISTRIBUTION_LIST = generalFunctions.select_classes(DISTRIBUTION_MODULE_NAME) +DISTRIBUTION_PARAMETERS_AND_DEFAULTS = generalFunctions.class_parameters_with_defaults( + DISTRIBUTION_MODULE_NAME ) -# ----------------------------------------------------------------------------- -# Defaults -# ----------------------------------------------------------------------------- - -state.selected_distribution = generalFunctions.get_default( - "selected_distribution", "default_values" -) -state.selected_distribution_type = generalFunctions.get_default( - "selected_distribution_type", "default_values" -) state.selected_distribution_parameters = [] -state.distributionTypeDisabled = False +state.distribution_type_disable = False # ----------------------------------------------------------------------------- # Main Functions # ----------------------------------------------------------------------------- -def populate_distribution_parameters(selected_distribution): +def populate_distribution_parameters(): """ Populates distribution parameters based on the selected distribution. :param selected_distribution (str): The name of the selected distribution whose parameters need to be populated. """ - if state.selected_distribution_type == "Twiss": + if state.distribution_type == "Twiss": sig = inspect.signature(twiss) state.selected_distribution_parameters = [ { @@ -77,10 +69,8 @@ def populate_distribution_parameters(selected_distribution): ] else: # when type == 'Quadratic Form' - selected_distribution_parameters = ( - state.listOfDistributionsAndParametersAndDefault.get( - selected_distribution, [] - ) + selected_distribution_parameters = DISTRIBUTION_PARAMETERS_AND_DEFAULTS.get( + state.distribution, [] ) state.selected_distribution_parameters = [ @@ -103,26 +93,6 @@ def populate_distribution_parameters(selected_distribution): return state.selected_distribution_parameters -def update_distribution_parameters( - parameterName, parameterValue, parameterErrorMessage -): - """ - Updates the value of a distribution parameter and its error message. - - :param parameterName (str): The name of the parameter to update. - :param parameterValue: The new value for the parameter. - :param parameterErrorMessage: The error message related to the parameter's value. - """ - - for param in state.selected_distribution_parameters: - if param["parameter_name"] == parameterName: - param["parameter_default_value"] = parameterValue - param["parameter_error_message"] = parameterErrorMessage - - generalFunctions.update_simulation_validation_status() - state.dirty("selected_distribution_parameters") - - # ----------------------------------------------------------------------------- # Write to file functions # ----------------------------------------------------------------------------- @@ -134,10 +104,10 @@ def distribution_parameters(): initialized with the appropriate parameters provided by the user. """ - distribution_name = state.selected_distribution + distribution_name = state.distribution parameters = DistributionFunctions.convert_distribution_parameters_to_valid_type() - if state.selected_distribution_type == "Twiss": + if state.distribution_type == "Twiss": twiss_params = twiss(**parameters) distr = getattr(distribution, distribution_name)(**twiss_params) else: @@ -151,20 +121,25 @@ def distribution_parameters(): # ----------------------------------------------------------------------------- -@state.change("selected_distribution") -def on_distribution_name_change(selected_distribution, **kwargs): - if selected_distribution == "Thermal": - state.selected_distribution_type = "Quadratic Form" - state.distributionTypeDisabled = True - state.dirty("selected_distribution_type") +@state.change("distribution") +def on_distribution_name_change(distribution, **kwargs): + if distribution == "Thermal" or distribution == "Empty": + state.distribution_type = "" + state.distribution_type_disable = True + state.dirty("distribution_type") else: - state.distributionTypeDisabled = False - populate_distribution_parameters(selected_distribution) + type_list_default = DashboardDefaults.LISTS["distribution_type_list"] + type_default = DashboardDefaults.DISTRIBUTION_PARAMETERS["distribution_type"] + + if state.distribution_type not in type_list_default: + state.distribution_type = type_default + state.distribution_type_disable = False -@state.change("selected_distribution_type") + +@state.change("distribution_type") def on_distribution_type_change(**kwargs): - populate_distribution_parameters(state.selected_distribution) + populate_distribution_parameters() @ctrl.add("updateDistributionParameters") @@ -172,7 +147,13 @@ def on_distribution_parameter_change(parameter_name, parameter_value, parameter_ parameter_value, input_type = generalFunctions.determine_input_type(parameter_value) error_message = generalFunctions.validate_against(parameter_value, parameter_type) - update_distribution_parameters(parameter_name, parameter_value, error_message) + for param in state.selected_distribution_parameters: + if param["parameter_name"] == parameter_name: + param["parameter_default_value"] = parameter_value + param["parameter_error_message"] = error_message + + generalFunctions.update_simulation_validation_status() + state.dirty("selected_distribution_parameters") # ----------------------------------------------------------------------------- @@ -192,47 +173,29 @@ def card(): """ with vuetify.VCard(style="width: 340px; height: 300px"): - with vuetify.VCardTitle("Distribution Parameters"): - vuetify.VSpacer() - TrameFunctions.create_refresh_button( - lambda: generalFunctions.reset_inputs("distribution") - ) - vuetify.VIcon( - "mdi-information", - style="color: #00313C;", - click=lambda: generalFunctions.documentation("BeamDistributions"), - ) - vuetify.VDivider() + TrameFunctions.input_section_header("Distribution Parameters") with vuetify.VCardText(): with vuetify.VRow(): with vuetify.VCol(cols=6): - vuetify.VCombobox( + TrameFunctions.select( label="Select Distribution", - v_model=("selected_distribution",), - items=("listOfDistributions",), - dense=True, + v_model_name="distribution", + items=(DISTRIBUTION_LIST,), ) with vuetify.VCol(cols=6): - vuetify.VSelect( - v_model=("selected_distribution_type",), + TrameFunctions.select( label="Type", - items=( - generalFunctions.get_default( - "distribution_type_list", "default_values" - ), - ), - dense=True, - disabled=("distributionTypeDisabled",), + v_model_name="distribution_type", + disabled=("distribution_type_disable",), ) with vuetify.VRow(classes="my-2"): for i in range(3): with vuetify.VCol(cols=4, classes="py-0"): with vuetify.VRow( - v_for="(parameter, index) in selected_distribution_parameters" + v_for="(parameter, index) in selected_distribution_parameters", + v_if=f"index % 3 == {i}", ): - with vuetify.VCol( - v_if=f"index % 3 == {i}", classes="py-1" - ): + with vuetify.VCol(classes="py-1"): vuetify.VTextField( label=("parameter.parameter_name",), v_model=("parameter.parameter_default_value",), diff --git a/src/python/impactx/dashboard/Input/generalFunctions.py b/src/python/impactx/dashboard/Input/generalFunctions.py index 49280e8ec..125b360c8 100644 --- a/src/python/impactx/dashboard/Input/generalFunctions.py +++ b/src/python/impactx/dashboard/Input/generalFunctions.py @@ -30,11 +30,11 @@ def documentation(section_name): :param section_name (str): The name of the documentation section to open. """ url_dict = { - "LatticeElements": "https://impactx.readthedocs.io/en/latest/usage/python.html#lattice-elements", - "BeamDistributions": "https://impactx.readthedocs.io/en/latest/usage/python.html#initial-beam-distributions", - "pythonParameters": "https://impactx.readthedocs.io/en/latest/usage/python.html#general", - "space_charge_documentation": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#space-charge", - "CSR": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#coherent-synchrotron-radiation-csr", + "input_parameters": "https://impactx.readthedocs.io/en/latest/usage/python.html#general", + "lattice_configuration": "https://impactx.readthedocs.io/en/latest/usage/python.html#lattice-elements", + "distribution_parameters": "https://impactx.readthedocs.io/en/latest/usage/python.html#initial-beam-distributions", + "space_charge": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#space-charge", + "csr": "https://impactx.readthedocs.io/en/latest/usage/parameters.html#coherent-synchrotron-radiation-csr", } url = url_dict.get(section_name) @@ -154,16 +154,18 @@ def update_simulation_validation_status(): ) # Check for errors in input card - if state.npart_validation: - error_details.append(f"Number of Particles: {state.npart_validation}") - if state.kin_energy_validation: - error_details.append(f"Kinetic Energy: {state.kin_energy_validation}") - if state.bunch_charge_C_validation: - error_details.append(f"Bunch Charge: {state.bunch_charge_C_validation}") - if state.charge_qe_validation: - error_details.append(f"Ref. Particle Charge: {state.charge_qe_validation}") - if state.mass_MeV_validation: - error_details.append(f"Ref. Particle Mass: {state.mass_MeV}") + if state.npart_error_message: + error_details.append(f"Number of Particles: {state.npart_error_message}") + if state.kin_energy_error_message: + error_details.append(f"Kinetic Energy: {state.kin_energy_error_message}") + if state.bunch_charge_C_error_message: + error_details.append(f"Bunch Charge: {state.bunch_charge_C_error_message}") + if state.charge_qe_error_message: + error_details.append( + f"Ref. Particle Charge: {state.charge_qe_error_message}" + ) + if state.mass_MeV_error_message: + error_details.append(f"Ref. Particle Mass: {state.mass_MeV_error_message}") if state.selected_lattice_list == []: error_details.append("LatticeListIsEmpty") @@ -343,9 +345,9 @@ def reset_inputs(input_section): if input_section.upper() in possible_section_names: state.update(getattr(DashboardDefaults, input_section.upper())) - if input_section == "distribution": + if input_section == "distribution_parameters": state.dirty("selected_distribution_type") - elif input_section == "lattice": + elif input_section == "lattice_configuration": state.selected_lattice_list = [] elif input_section == "space_charge": state.dirty("max_level") diff --git a/src/python/impactx/dashboard/Input/inputParameters/inputMain.py b/src/python/impactx/dashboard/Input/inputParameters/inputMain.py index 7441ae284..cfe5eec01 100644 --- a/src/python/impactx/dashboard/Input/inputParameters/inputMain.py +++ b/src/python/impactx/dashboard/Input/inputParameters/inputMain.py @@ -6,11 +6,13 @@ License: BSD-3-Clause-LBNL """ -from trame.widgets import vuetify - -from ...Input.trameFunctions import TrameFunctions -from ...trame_setup import setup_server -from ..generalFunctions import generalFunctions +from .. import ( + DashboardDefaults, + TrameFunctions, + generalFunctions, + setup_server, + vuetify, +) from .inputFunctions import InputFunctions server, state, ctrl = setup_server() @@ -20,10 +22,13 @@ # ----------------------------------------------------------------------------- -@ctrl.add("on_input_change") -def validate_and_convert_to_correct_type( - value, desired_type, state_name, validation_name, conditions=None -): +@ctrl.add("input_change") +def validate_and_convert_to_correct_type(state_name): + value = getattr(state, state_name) + desired_type = DashboardDefaults.TYPES[state_name] + validation_name = f"{state_name}_error_message" + conditions = DashboardDefaults.VALIDATION_CONDITION.get(state_name, None) + validation_result = generalFunctions.validate_against( value, desired_type, conditions ) @@ -65,48 +70,13 @@ class InputParameters: User-Input section for beam properties. """ - def __init__(self): - state.particle_shape = generalFunctions.get_default( - "particle_shape", "default_values" - ) - state.npart = generalFunctions.get_default("npart", "default_values") - state.kin_energy = generalFunctions.get_default("kin_energy", "default_values") - state.kin_energy_MeV = state.kin_energy - state.bunch_charge_C = generalFunctions.get_default( - "bunch_charge_C", "default_values" - ) - state.kin_energy_unit = generalFunctions.get_default( - "kin_energy_unit", "default_values" - ) - state.old_kin_energy_unit = generalFunctions.get_default( - "kin_energy_unit", "default_values" - ) - state.charge_qe = generalFunctions.get_default("charge_qe", "default_values") - state.mass_MeV = generalFunctions.get_default("mass_MeV", "default_values") - - state.npart_validation = [] - state.kin_energy_validation = [] - state.bunch_charge_C_validation = [] - state.mass_MeV_validation = [] - state.charge_qe_validation = [] - def card(self): """ Creates UI content for beam properties. """ with vuetify.VCard(style="width: 340px; height: 350px"): - with vuetify.VCardTitle("Input Parameters"): - vuetify.VSpacer() - TrameFunctions.create_refresh_button( - lambda: generalFunctions.reset_inputs("input_parameters") - ) - vuetify.VIcon( - "mdi-information", - style="color: #00313C;", - click=lambda: generalFunctions.documentation("pythonParameters"), - ) - vuetify.VDivider() + TrameFunctions.input_section_header("Input Parameters") with vuetify.VCardText(): with vuetify.VRow(classes="py-2"): with vuetify.VCol(cols=6, classes="py-0"): @@ -123,95 +93,42 @@ def card(self): ) with vuetify.VRow(classes="my-2"): with vuetify.VCol(cols=6, classes="py-0"): - vuetify.VTextField( + TrameFunctions.text_field( label="Ref. Particle Charge", - v_model=("charge_qe",), - suffix=generalFunctions.get_default("charge_qe", "units"), - type="number", - step=generalFunctions.get_default("charge_qe", "steps"), - __properties=["step"], - dense=True, - error_messages=("charge_qe_validation",), - change=( - ctrl.on_input_change, - "[$event, 'int','charge_qe','charge_qe_validation', ['non_zero']]", - ), + v_model_name="charge_qe", + input=(ctrl.input_change, "['charge_qe']"), ) with vuetify.VCol(cols=6, classes="py-0"): - vuetify.VTextField( + TrameFunctions.text_field( label="Ref. Particle Mass", - v_model=("mass_MeV",), - suffix=generalFunctions.get_default("mass_MeV", "units"), - type="number", - step=generalFunctions.get_default("mass_MeV", "steps"), - __properties=["step"], - dense=True, - error_messages=("mass_MeV_validation",), - change=( - ctrl.on_input_change, - "[$event, 'float','mass_MeV','mass_MeV_validation', ['positive']]", - ), + v_model_name="mass_MeV", + input=(ctrl.input_change, "['mass_MeV']"), ) with vuetify.VRow(classes="my-0"): with vuetify.VCol(cols=12, classes="py-0"): - vuetify.VTextField( - v_model=("npart",), + TrameFunctions.text_field( label="Number of Particles", - error_messages=("npart_validation",), - change=( - ctrl.on_input_change, - "[$event, 'int','npart','npart_validation']", - ), - type="number", - step=generalFunctions.get_default("npart", "steps"), - __properties=["step"], - dense=True, + v_model_name="npart", + input=(ctrl.input_change, "['npart']"), ) with vuetify.VRow(classes="my-2"): with vuetify.VCol(cols=8, classes="py-0"): - vuetify.VTextField( - v_model=("kin_energy",), + TrameFunctions.text_field( label="Kinetic Energy", - error_messages=("kin_energy_validation",), - change=( - ctrl.on_input_change, - "[$event, 'float','kin_energy','kin_energy_validation']", - ), - type="number", - step=generalFunctions.get_default("kin_energy", "steps"), - __properties=["step"], - dense=True, + v_model_name="kin_energy", + input=(ctrl.input_change, "['kin_energy']"), classes="mr-2", ) with vuetify.VCol(cols=4, classes="py-0"): - vuetify.VSelect( - v_model=("kin_energy_unit",), + TrameFunctions.select( label="Unit", - items=( - generalFunctions.get_default( - "kin_energy_unit_list", "default_values" - ), - ), - change=(ctrl.kin_energy_unit_change, "[$event]"), - dense=True, + v_model_name="kin_energy_unit", + input=(ctrl.kin_energy_unit_change, "[$event]"), ) with vuetify.VRow(classes="my-2"): with vuetify.VCol(cols=12, classes="py-0"): - vuetify.VTextField( + TrameFunctions.text_field( label="Bunch Charge", - v_model=("bunch_charge_C",), - error_messages=("bunch_charge_C_validation",), - change=( - ctrl.on_input_change, - "[$event, 'float','bunch_charge_C','bunch_charge_C_validation']", - ), - type="number", - step=generalFunctions.get_default( - "bunch_charge_C", "steps" - ), - __properties=["step"], - dense=True, - suffix=generalFunctions.get_default( - "bunch_charge_C", "units" - ), + v_model_name="bunch_charge_C", + input=(ctrl.input_change, "['bunch_charge_C']"), ) diff --git a/src/python/impactx/dashboard/Input/latticeConfiguration/latticeMain.py b/src/python/impactx/dashboard/Input/latticeConfiguration/latticeMain.py index 9b5b80a8f..8a25f3f23 100644 --- a/src/python/impactx/dashboard/Input/latticeConfiguration/latticeMain.py +++ b/src/python/impactx/dashboard/Input/latticeConfiguration/latticeMain.py @@ -6,13 +6,9 @@ License: BSD-3-Clause-LBNL """ -from trame.widgets import vuetify - from impactx import elements -from ...Input.trameFunctions import TrameFunctions -from ...trame_setup import setup_server -from ..generalFunctions import generalFunctions +from .. import TrameFunctions, generalFunctions, setup_server, vuetify server, state, ctrl = setup_server() @@ -33,11 +29,8 @@ # Default # ----------------------------------------------------------------------------- -state.selected_lattice = generalFunctions.get_default( - "selected_lattice", "default_values" -) state.selected_lattice_list = [] -state.nsliceDefaultValue = generalFunctions.get_default("n_slice", "default_values") +state.nslice = "" # ----------------------------------------------------------------------------- # Main Functions @@ -220,7 +213,7 @@ def update_default_value(parameter_name, new_value): # ----------------------------------------------------------------------------- -# ContentSetup +# UI # ----------------------------------------------------------------------------- @@ -235,25 +228,16 @@ def card(): Creates UI content for lattice configuration. """ - with vuetify.VDialog(v_model=("showDialog", False), width="1200px"): - LatticeConfiguration.dialog_lattice_elementList() + with vuetify.VDialog(v_model=("showDialog", False)): + LatticeConfiguration.dialog_configuration_list() - with vuetify.VDialog(v_model=("showDialog_settings", False), width="500px"): - LatticeConfiguration.dialog_lattice_settings() + with vuetify.VDialog( + v_model=("lattice_configuration_dialog_settings", False), width="500px" + ): + LatticeConfiguration.dialog_settings() with vuetify.VCard(style="width: 696px;"): - with vuetify.VCardTitle("Lattice Configuration"): - vuetify.VSpacer() - TrameFunctions.create_refresh_button( - lambda: generalFunctions.reset_inputs("lattice") - ) - vuetify.VIcon( - "mdi-information", - classes="ml-2", - click=lambda: generalFunctions.documentation("LatticeElements"), - style="color: #00313C;", - ) - vuetify.VDivider() + TrameFunctions.input_section_header("Lattice Configuration") with vuetify.VCardText(): with vuetify.VRow(align="center", no_gutters=True): with vuetify.VCol(cols=10): @@ -276,7 +260,7 @@ def card(): with vuetify.VCol(cols="auto"): vuetify.VIcon( "mdi-cog", - click="showDialog_settings = true", + click="lattice_configuration_dialog_settings = true", ) with vuetify.VRow(): with vuetify.VCol(): @@ -293,144 +277,98 @@ def card(): click="showDialog = true", ) vuetify.VDivider() - with vuetify.VContainer(fluid=True): - with vuetify.VRow( - v_for="(latticeElement, index) in selected_lattice_list", - align="center", - no_gutters=True, - style="min-width: 1500px;", - ): - with vuetify.VCol(cols="auto", classes="pa-2"): - vuetify.VIcon( - "mdi-menu-up", - click=( - ctrl.move_latticeElementIndex_up, - "[index]", - ), - ) - vuetify.VIcon( - "mdi-menu-down", - click=( - ctrl.move_latticeElementIndex_down, - "[index]", - ), - ) - vuetify.VIcon( - "mdi-delete", - click=( - ctrl.deleteLatticeElement, - "[index]", - ), - ) - vuetify.VChip( - v_text=("latticeElement.name",), - dense=True, - classes="mr-2", - style="justify-content: center", - ) - with vuetify.VCol( - v_for="(parameter, parameterIndex) in latticeElement.parameters", - cols="auto", - classes="pa-2", - ): - vuetify.VTextField( - label=("parameter.parameter_name",), - v_model=( - "parameter.parameter_default_value", - ), - change=( - ctrl.updateLatticeElementParameters, - "[index, parameter.parameter_name, $event, parameter.parameter_type]", - ), - error_messages=( - "parameter.parameter_error_message", - ), - dense=True, - style="width: 100px;", - ) + LatticeConfiguration.configuration_list() + + # ----------------------------------------------------------------------------- + # Dialogs + # ----------------------------------------------------------------------------- @staticmethod - def dialog_lattice_elementList(): + def dialog_configuration_list(): """ - Displays the content shown on the dialog - box for lattice configuration. + Displays the configuration for lattice elements + shown in a dialog box. """ with vuetify.VCard(): - with vuetify.VCardTitle("Elements", classes="text-subtitle-2 pa-3"): + with vuetify.VCardTitle("Elements", classes="headline d-flex align-center"): vuetify.VSpacer() + with vuetify.VBtn(icon=True, click="showDialog = false"): + vuetify.VIcon("mdi-close") vuetify.VDivider() - with vuetify.VContainer(fluid=True): - with vuetify.VRow( - v_for="(latticeElement, index) in selected_lattice_list", - align="center", - no_gutters=True, - style="min-width: 1500px;", - ): - with vuetify.VCol(cols="auto", classes="pa-2"): - vuetify.VIcon( - "mdi-delete", - click=(ctrl.deleteLatticeElement, "[index]"), - ) - vuetify.VChip( - v_text=("latticeElement.name",), - dense=True, - classes="mr-2", - style="justify-content: center", - ) - with vuetify.VCol( - v_for="(parameter, parameterIndex) in latticeElement.parameters", - cols="auto", - classes="pa-2", - ): - vuetify.VTextField( - label=("parameter.parameter_name",), - v_model=("parameter.parameter_default_value",), - change=( - ctrl.updateLatticeElementParameters, - "[index, parameter.parameter_name, $event, parameter.parameter_type]", - ), - error_messages=("parameter.parameter_error_message",), - dense=True, - style="width: 100px;", - ) + LatticeConfiguration.configuration_list() @staticmethod - def dialog_lattice_settings(): + def dialog_settings(): """ - Creates UI content for lattice configuration - settings. + Provides controls for lattice element configuration, + allowing dashboard users to define parameter defaults. """ + dialog_name = "lattice_configuration_dialog_tab_settings" + + TrameFunctions.create_dialog_tabs(dialog_name, 1, ["Defaults"]) + with vuetify.VTabsItems(v_model=(dialog_name, 0)): + with vuetify.VTabItem(): + with vuetify.VCardText(): + with vuetify.VRow(): + with vuetify.VCol(cols=3): + TrameFunctions.text_field( + label="nslice", + v_model_name="nslice", + change=( + ctrl.nsliceDefaultChange, + "['nslice', $event]", + ), + ) + + # ----------------------------------------------------------------------------- + # lattice_configuration_lsit + # ----------------------------------------------------------------------------- - with vuetify.VCard(): - with vuetify.VTabs(v_model=("tab", "Settings")): - vuetify.VTab("Settings") - # vuetify.VTab("Variable Referencing") - vuetify.VDivider() - with vuetify.VTabsItems(v_model="tab"): - with vuetify.VTabItem(): - with vuetify.VContainer(fluid=True): - with vuetify.VRow(no_gutters=True, align="center"): - with vuetify.VCol(no_gutters=True, cols="auto"): - vuetify.VListItem( - "nslice", classes="ma-0 pl-0 font-weight-bold" - ) - with vuetify.VCol(no_gutters=True): - vuetify.VTextField( - v_model=("nsliceDefaultValue",), - change=( - ctrl.nsliceDefaultChange, - "['nslice', $event]", - ), - type="number", - step=generalFunctions.get_default( - "nslice", "steps" - ), - __properties=["step"], - placeholder="Value", - dense=True, - outlined=True, - hide_details=True, - style="max-width: 75px", - classes="ma-0 pa-0", - ) + @staticmethod + def configuration_list(): + """ + Displays the configuration for lattice elements. + """ + with vuetify.VContainer(fluid=True): + with vuetify.VRow( + v_for="(latticeElement, index) in selected_lattice_list", + align="center", + no_gutters=True, + style="min-width: 1500px;", + ): + with vuetify.VCol(cols="auto", classes="pa-2"): + vuetify.VIcon( + "mdi-menu-up", + click=(ctrl.move_latticeElementIndex_up, "[index]"), + ) + vuetify.VIcon( + "mdi-menu-down", + click=(ctrl.move_latticeElementIndex_down, "[index]"), + ) + vuetify.VIcon( + "mdi-delete", + click=(ctrl.deleteLatticeElement, "[index]"), + ) + vuetify.VChip( + v_text=("latticeElement.name",), + dense=True, + classes="mr-2", + style="justify-content: center", + ) + with vuetify.VCol( + v_for="(parameter, parameterIndex) in latticeElement.parameters", + cols="auto", + classes="pa-2", + ): + vuetify.VTextField( + label=("parameter.parameter_name",), + v_model=("parameter.parameter_default_value",), + change=( + ctrl.updateLatticeElementParameters, + "[index, parameter.parameter_name, $event, parameter.parameter_type]", + ), + error_messages=("parameter.parameter_error_message",), + dense=True, + style="width: 100px;", + ) diff --git a/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeFunctions.py b/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeFunctions.py index 514978b85..ddcf36f57 100644 --- a/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeFunctions.py +++ b/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeFunctions.py @@ -26,7 +26,7 @@ def validate_prob_relative_fields(index, prob_relative_value): if index == 0: if poisson_solver == "multigrid": - if prob_relative_value < 3: + if prob_relative_value <= 3: error_message = "Must be greater than 3." elif poisson_solver == "fft": if prob_relative_value <= 1: @@ -58,10 +58,10 @@ def validate_n_cell_and_blocking_factor(direction): blocking_factor_value, "int", ["non_zero", "positive"] ) - setattr(state, f"error_message_n_cell_{direction}", "; ".join(n_cell_errors)) + setattr(state, f"n_cell_{direction}_error_message", "; ".join(n_cell_errors)) setattr( state, - f"error_message_blocking_factor_{direction}", + f"blocking_factor_{direction}_error_message", "; ".join(blocking_factor_errors), ) @@ -71,6 +71,6 @@ def validate_n_cell_and_blocking_factor(direction): if n_cell_value % blocking_factor_value != 0: setattr( state, - f"error_message_n_cell_{direction}", + f"n_cell_{direction}_error_message", "Must be a multiple of blocking factor.", ) diff --git a/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeMain.py b/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeMain.py index c9a920b49..5d1dc914b 100644 --- a/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeMain.py +++ b/src/python/impactx/dashboard/Input/space_charge_configuration/spaceChargeMain.py @@ -1,8 +1,4 @@ -from trame.widgets import vuetify - -from ...Input.trameFunctions import TrameFunctions -from ...trame_setup import setup_server -from ..generalFunctions import generalFunctions +from .. import TrameFunctions, generalFunctions, setup_server, vuetify from .spaceChargeFunctions import SpaceChargeFunctions server, state, ctrl = setup_server() @@ -11,67 +7,42 @@ # Default # ----------------------------------------------------------------------------- -state.dynamic_size = generalFunctions.get_default("dynamic_size", "default_values") -state.max_level = generalFunctions.get_default("max_level", "default_values") -state.particle_shape = generalFunctions.get_default("particle_shape", "default_values") -state.poisson_solver = generalFunctions.get_default("poisson_solver", "default_values") - state.prob_relative = [] state.prob_relative_fields = [] - state.n_cell = [] -state.n_cell_x = generalFunctions.get_default("n_cell_x", "default_values") -state.n_cell_y = generalFunctions.get_default("n_cell_y", "default_values") -state.n_cell_z = generalFunctions.get_default("n_cell_z", "default_values") - -state.blocking_factor_x = generalFunctions.get_default( - "blocking_factor_x", "default_values" -) -state.blocking_factor_y = generalFunctions.get_default( - "blocking_factor_y", "default_values" -) -state.blocking_factor_z = generalFunctions.get_default( - "blocking_factor_z", "default_values" -) - -state.mlmg_relative_tolerance = generalFunctions.get_default( - "mlmg_relative_tolerance", "default_values" -) -state.mlmg_absolute_tolerance = generalFunctions.get_default( - "mlmg_absolute_tolerance", "default_values" -) -state.mlmg_max_iters = generalFunctions.get_default("mlmg_max_iters", "default_values") -state.mlmg_verbosity = generalFunctions.get_default("mlmg_verbosity", "default_values") - -state.error_message_mlmg_relative_tolerance = "" -state.error_message_mlmg_absolute_tolerance = "" -state.error_message_mlmg_max_iters = "" -state.error_message_mlmg_verbosity = "" + # ----------------------------------------------------------------------------- # Helper functions # ----------------------------------------------------------------------------- -def populate_prob_relative_fields(max_level): - num_prob_relative_fields = int(max_level) + 1 +def populate_prob_relative_fields(): + tot_num_prob_relative_fields = int(state.max_level) + 1 fft_first_field_value = generalFunctions.get_default( "prob_relative_first_value_fft", "default_values" ) multigrid_first_field_value = generalFunctions.get_default( "prob_relative_first_value_multigrid", "default_values" ) - + first_field_value = 0 if state.poisson_solver == "fft": - state.prob_relative = [fft_first_field_value] + [0.0] * ( - num_prob_relative_fields - 1 - ) + first_field_value = fft_first_field_value elif state.poisson_solver == "multigrid": - state.prob_relative = [multigrid_first_field_value] + [0.0] * ( - num_prob_relative_fields - 1 - ) + first_field_value = multigrid_first_field_value + + if state.prob_relative: + size = len(state.prob_relative) + num_of_extra_fields = tot_num_prob_relative_fields - size + + if size < tot_num_prob_relative_fields: + state.prob_relative.extend([0.0] * (num_of_extra_fields)) + elif size > tot_num_prob_relative_fields: + state.prob_relative = state.prob_relative[:tot_num_prob_relative_fields] else: - state.prob_relative = [0.0] * num_prob_relative_fields + state.prob_relative = [first_field_value] + [0.0] * ( + tot_num_prob_relative_fields - 1 + ) state.prob_relative_fields = [ { @@ -81,7 +52,7 @@ def populate_prob_relative_fields(max_level): ), "step": generalFunctions.get_default("prob_relative", "steps"), } - for i in range(num_prob_relative_fields) + for i in range(tot_num_prob_relative_fields) ] @@ -93,9 +64,9 @@ def update_blocking_factor_and_n_cell(category, kwargs): direction = state_name.split("_")[-1] SpaceChargeFunctions.validate_n_cell_and_blocking_factor(direction) - n_cell_error = getattr(state, f"error_message_n_cell_{direction}") + n_cell_error = getattr(state, f"n_cell_{direction}_error_message") blocking_factor_error = getattr( - state, f"error_message_blocking_factor_{direction}" + state, f"blocking_factor_{direction}_error_message" ) if not n_cell_error: @@ -119,7 +90,7 @@ def update_blocking_factor_and_n_cell(category, kwargs): # ----------------------------------------------------------------------------- @state.change("poisson_solver") def on_poisson_solver_change(poisson_solver, **kwargs): - populate_prob_relative_fields(state.max_level) + populate_prob_relative_fields() state.dirty("prob_relative_fields") generalFunctions.update_simulation_validation_status() @@ -132,7 +103,7 @@ def on_space_charge_change(space_charge, **kwargs): @state.change("max_level") def on_max_level_change(max_level, **kwargs): - populate_prob_relative_fields(max_level) + populate_prob_relative_fields() generalFunctions.update_simulation_validation_status() @@ -185,6 +156,14 @@ def on_update_prob_relative_call(index, value): # ----------------------------------------------------------------------------- +def multigrid_settings(): + vuetify.VIcon( + "mdi-cog", + v_if="poisson_solver == 'multigrid'", + click="space_charge_dialog_settings = true", + ) + + class SpaceChargeConfiguration: @staticmethod def card(): @@ -192,66 +171,29 @@ def card(): Creates UI content for space charge configuration """ - with vuetify.VDialog(v_model=("showSpaceChargeDialog", False), width="500px"): - SpaceChargeConfiguration.dialog_space_charge_settings() + with vuetify.VDialog( + v_model=("space_charge_dialog_settings", False), width="500px" + ): + SpaceChargeConfiguration.dialog_settings() with vuetify.VCard(v_show="space_charge", style="width: 340px;"): - with vuetify.VCardTitle("Space Charge"): - vuetify.VSpacer() - vuetify.VIcon( - "mdi-cog", - classes="ml-2", - v_if="poisson_solver == 'multigrid'", - click="showSpaceChargeDialog = true", - style="cursor: pointer;", - ) - TrameFunctions.create_refresh_button( - lambda: generalFunctions.reset_inputs("space_charge") - ) - vuetify.VIcon( - "mdi-information", - classes="ml-2", - click=lambda: generalFunctions.documentation( - "space_charge_documentation" - ), - style="color: #00313C;", - ) - vuetify.VDivider() + TrameFunctions.input_section_header( + "Space Charge", additional_components=multigrid_settings + ) with vuetify.VCardText(): with vuetify.VRow(classes="my-0"): with vuetify.VCol(cols=5, classes="py-0"): - vuetify.VSelect( + TrameFunctions.select( label="Poisson Solver", - v_model=("poisson_solver",), - items=( - generalFunctions.get_default( - "poisson_solver_list", "default_values" - ), - ), - dense=True, hide_details=True, ) with vuetify.VCol(cols=4, classes="py-0"): - vuetify.VSelect( + TrameFunctions.select( label="Particle Shape", - v_model=("particle_shape",), - items=( - generalFunctions.get_default( - "particle_shape_list", "default_values" - ), - ), - dense=True, ) with vuetify.VCol(cols=3, classes="py-0"): - vuetify.VSelect( + TrameFunctions.select( label="Max Level", - v_model=("max_level",), - items=( - generalFunctions.get_default( - "max_level_list", "default_values" - ), - ), - dense=True, ) with vuetify.VCol(classes="pa-0"): vuetify.VListItemSubtitle( @@ -261,14 +203,10 @@ def card(): with vuetify.VRow(classes="my-0"): for direction in ["x", "y", "z"]: with vuetify.VCol(cols=4, classes="py-0"): - vuetify.VTextField( - placeholder=direction, - v_model=(f"n_cell_{direction}",), - error_messages=(f"error_message_n_cell_{direction}",), - type="number", - step=generalFunctions.get_default("n_cell", "steps"), - __properties=["step"], - dense=True, + TrameFunctions.text_field( + label="", + v_model_name=f"n_cell_{direction}", + prefix=f"{direction}:", style="margin-top: -5px", ) with vuetify.VCol(classes="pa-0"): @@ -279,18 +217,10 @@ def card(): with vuetify.VRow(classes="my-0"): for direction in ["x", "y", "z"]: with vuetify.VCol(cols=4, classes="py-0"): - vuetify.VTextField( - placeholder=direction, - v_model=(f"blocking_factor_{direction}",), - error_messages=( - f"error_message_blocking_factor_{direction}", - ), - type="number", - step=generalFunctions.get_default( - "blocking_factor", "steps" - ), - __properties=["step"], - dense=True, + TrameFunctions.text_field( + label="", + prefix=f"{direction}:", + v_model_name=f"blocking_factor_{direction}", style="margin-top: -5px", ) with vuetify.VCol(classes="pa-0"): @@ -316,78 +246,33 @@ def card(): ) @staticmethod - def dialog_space_charge_settings(): + def dialog_settings(): """ Creates UI content for space charge configuration settings. """ - with vuetify.VCard(): - with vuetify.VTabs( - v_model=("space_charge_tab", "Advanced Multigrid Settings") - ): - vuetify.VTab("Settings") - vuetify.VDivider() - with vuetify.VTabsItems(v_model="space_charge_tab"): - with vuetify.VTabItem(): - with vuetify.VContainer(fluid=True): - with vuetify.VRow( - classes="my-2", v_if="poisson_solver == 'multigrid'" - ): - with vuetify.VCol(cols=6, classes="py-0"): - vuetify.VTextField( - label="MLMG Relative Tolerance", - v_model=("mlmg_relative_tolerance",), - error_messages=( - "error_message_mlmg_relative_tolerance", - ), - type="number", - step=generalFunctions.get_default( - "mlmg_relative_tolerance", "steps" - ), - __properties=["step"], - dense=True, - ) - with vuetify.VCol(cols=6, classes="py-0"): - vuetify.VTextField( - label="MLMG Absolute Tolerance", - v_model=("mlmg_absolute_tolerance",), - error_messages=( - "error_message_mlmg_absolute_tolerance", - ), - suffix=generalFunctions.get_default( - "mlmg_absolute_tolerance", "units" - ), - type="number", - step=generalFunctions.get_default( - "mlmg_absolute_tolerance", "steps" - ), - __properties=["step"], - dense=True, - ) - with vuetify.VRow( - classes="my-0", v_if="poisson_solver == 'multigrid'" - ): - with vuetify.VCol(cols=6, classes="py-0"): - vuetify.VTextField( - label="MLMG Max Iterations", - v_model=("mlmg_max_iters",), - error_messages=("error_message_mlmg_max_iters",), - type="number", - step=generalFunctions.get_default( - "mlmg_max_iters", "steps" - ), - __properties=["step"], - dense=True, - ) - with vuetify.VCol(cols=6, classes="py-0"): - vuetify.VTextField( - label="MLMG Verbosity", - v_model=("mlmg_verbosity",), - error_messages=("error_message_mlmg_verbosity",), - type="number", - step=generalFunctions.get_default( - "mlmg_verbosity", "steps" - ), - __properties=["step"], - dense=True, - ) + dialog_name = "space_charge_dialog_tab_settings" + TrameFunctions.create_dialog_tabs( + dialog_name, 1, ["Advanced Multigrid Settings"] + ) + with vuetify.VTabsItems(v_model=("dialog_name", 0)): + with vuetify.VTabItem(): + with vuetify.VCardText(): + with vuetify.VRow(): + with vuetify.VCol(): + TrameFunctions.text_field( + label="MLMG Relative Tolerance", + ) + with vuetify.VCol(): + TrameFunctions.text_field( + label="MLMG Absolute Tolerance", + ) + with vuetify.VRow(): + with vuetify.VCol(): + TrameFunctions.text_field( + label="MLMG Max Iters", + ) + with vuetify.VCol(): + TrameFunctions.text_field( + label="MLMG Verbosity", + ) diff --git a/src/python/impactx/dashboard/Input/trameFunctions.py b/src/python/impactx/dashboard/Input/trameFunctions.py index ef9c58246..d540fa7f3 100644 --- a/src/python/impactx/dashboard/Input/trameFunctions.py +++ b/src/python/impactx/dashboard/Input/trameFunctions.py @@ -8,6 +8,7 @@ from trame.widgets import vuetify +from ..Input.generalFunctions import generalFunctions from ..trame_setup import setup_server server, state, ctrl = setup_server() @@ -43,13 +44,79 @@ def create_route(route_title, mdi_icon): vuetify.VListItemTitle(route_title) @staticmethod - def create_refresh_button(reset_function_name): + def select(label, v_model_name=None, items=None, **kwargs): + # in place for now as some variables are not in same format + if v_model_name is None: + v_model_name = label.lower().replace(" ", "_") + + if items is None: + items = ( + generalFunctions.get_default(f"{v_model_name}_list", "default_values"), + ) + + vuetify.VSelect( + label=label, + v_model=(f"{v_model_name}",), + items=items, + dense=True, + **kwargs, + ) + + @staticmethod + def text_field(label, v_model_name=None, **kwargs): + if v_model_name is None: + v_model_name = label.lower().replace(" ", "_") + + vuetify.VTextField( + label=label, + v_model=(f"{v_model_name}",), + error_messages=(f"{v_model_name}_error_message", []), + type="number", + step=generalFunctions.get_default(f"{v_model_name}", "steps"), + suffix=generalFunctions.get_default(f"{v_model_name}", "units"), + __properties=["step"], + dense=True, + **kwargs, + ) + + @staticmethod + def documentation_icon(section_name, **kwargs): + vuetify.VIcon( + "mdi-information", + style="color: #00313C;", + click=lambda: generalFunctions.documentation(section_name), + ) + + @staticmethod + def refresh_icon(section_name): """ Creates a standardized refresh button. :param reset_function: The reset function to call when clicked. """ - return vuetify.VIcon( + vuetify.VIcon( "mdi-refresh", style="color: #00313C;", - click=reset_function_name, + click=lambda: generalFunctions.reset_inputs(section_name), ) + + @staticmethod + def input_section_header(section_name, additional_components=None): + documentation_name = section_name.lower().replace(" ", "_") + with vuetify.VCardTitle(section_name): + vuetify.VSpacer() + if additional_components: + additional_components() + TrameFunctions.refresh_icon(documentation_name) + TrameFunctions.documentation_icon(documentation_name) + vuetify.VDivider() + + @staticmethod + def create_dialog_tabs(name: str, num_tabs: int, tab_names: list[str]): + if len(tab_names) != num_tabs: + raise ValueError("Number of tab names must match number of tabs_names") + + with vuetify.VCard(): + with vuetify.VTabs(v_model=(f"{name}", 0)): + for tab_name in tab_names: + vuetify.VTab(tab_name) + vuetify.VDivider() diff --git a/src/python/impactx/dashboard/Toolbar/exportTemplate.py b/src/python/impactx/dashboard/Toolbar/exportTemplate.py index c01b53e82..1fbf042e0 100644 --- a/src/python/impactx/dashboard/Toolbar/exportTemplate.py +++ b/src/python/impactx/dashboard/Toolbar/exportTemplate.py @@ -14,15 +14,15 @@ def build_distribution_list(): Generates an instance of distribution inputs as a string for exporting purposes. """ - distribution_name = state.selected_distribution + distribution_name = state.distribution parameters = DistributionFunctions.convert_distribution_parameters_to_valid_type() - indentation = " " * (8 if state.selected_distribution_type == "Twiss" else 4) + indentation = " " * (8 if state.distribution_type == "Twiss" else 4) distribution_parameters = ",\n".join( f"{indentation}{key}={value}" for key, value in parameters.items() ) - if state.selected_distribution_type == "Twiss": + if state.distribution_type == "Twiss": return ( f"distr = distribution.{distribution_name}(\n" f" **twiss(\n" @@ -32,9 +32,7 @@ def build_distribution_list(): ) else: return ( - f"distr = distribution.{distribution_name}(\n" - f"{distribution_parameters},\n" - f")" + f"distr = distribution.{distribution_name}(\n{distribution_parameters},\n)" ) @@ -45,7 +43,7 @@ def build_lattice_list(): """ lattice_elements = ",\n ".join( - f'elements.{element["name"]}(' + f"elements.{element['name']}(" + ", ".join( f"{key}={value}" for key, value in parameter_input_checker_for_lattice(element).items() diff --git a/src/python/impactx/dashboard/Toolbar/toolbarMain.py b/src/python/impactx/dashboard/Toolbar/toolbarMain.py index 810a3ac62..c3288510e 100644 --- a/src/python/impactx/dashboard/Toolbar/toolbarMain.py +++ b/src/python/impactx/dashboard/Toolbar/toolbarMain.py @@ -96,36 +96,22 @@ def dashboard_info(): class Toolbars: """ - Builds section toolbars for various pages. + Builds toolbar for dashboard. """ @staticmethod - def input_toolbar(): - """ - Builds toolbar for the 'Input' page. - """ - - (ToolbarElements.dashboard_info(),) - vuetify.VSpacer() - ToolbarElements.reset_inputs_button() - ToolbarElements.export_input_data() - - @staticmethod - def run_toolbar(): - """ - Builds toolbar for the 'Run' page. - """ - - (ToolbarElements.dashboard_info(),) - (vuetify.VSpacer(),) - (ToolbarElements.run_simulation_button(),) - - @staticmethod - def analyze_toolbar(): - """ - Builds toolbar for the 'Analyze' page. - """ - - (ToolbarElements.dashboard_info(),) - vuetify.VSpacer() - ToolbarElements.plot_options() + def dashboard_toolbar(toolbar_name: str) -> None: + toolbar_name = toolbar_name.lower() + if toolbar_name == "input": + (ToolbarElements.dashboard_info(),) + vuetify.VSpacer() + ToolbarElements.reset_inputs_button() + ToolbarElements.export_input_data() + elif toolbar_name == "run": + (ToolbarElements.dashboard_info(),) + (vuetify.VSpacer(),) + (ToolbarElements.run_simulation_button(),) + elif toolbar_name == "analyze": + (ToolbarElements.dashboard_info(),) + vuetify.VSpacer() + ToolbarElements.plot_options() diff --git a/src/python/impactx/dashboard/__main__.py b/src/python/impactx/dashboard/__main__.py index 4f15fa468..dd973d6fd 100644 --- a/src/python/impactx/dashboard/__main__.py +++ b/src/python/impactx/dashboard/__main__.py @@ -72,11 +72,11 @@ def application(): layout.title.hide() with layout.toolbar: with vuetify.Template(v_if="$route.path == '/Analyze'"): - Toolbars.analyze_toolbar() + Toolbars.dashboard_toolbar("analyze") with vuetify.Template(v_if="$route.path == '/Input'"): - Toolbars.input_toolbar() + Toolbars.dashboard_toolbar("input") with vuetify.Template(v_if="$route.path == '/Run'"): - Toolbars.run_toolbar() + Toolbars.dashboard_toolbar("run") with layout.drawer as drawer: drawer.width = 200 diff --git a/src/python/impactx/dashboard/start.py b/src/python/impactx/dashboard/start.py index 54740cd9f..2d541a53b 100644 --- a/src/python/impactx/dashboard/start.py +++ b/src/python/impactx/dashboard/start.py @@ -6,6 +6,7 @@ License: BSD-3-Clause-LBNL """ +from .Input.defaults import DashboardDefaults from .trame_setup import setup_server server, state, ctrl = setup_server() @@ -16,9 +17,20 @@ # ----------------------------------------------------------------------------- +def initialize_states(): + """ + Initializes all dashboard state values upon call. + + The issue as of now is it initialize all at once instead of by section. + """ + for name, value in DashboardDefaults.DEFAULT_VALUES.items(): + setattr(state, name, value) + + def main(): """ Launches Trame application server """ + initialize_states() server.start() return 0 diff --git a/src/python/impactx/impactx_pybind/__init__.pyi b/src/python/impactx/impactx_pybind/__init__.pyi index 8a5e7b4bd..1367006a2 100644 --- a/src/python/impactx/impactx_pybind/__init__.pyi +++ b/src/python/impactx/impactx_pybind/__init__.pyi @@ -46,6 +46,8 @@ class Config: have_gpu: typing.ClassVar[bool] = False have_mpi: typing.ClassVar[bool] = True have_omp: typing.ClassVar[bool] = True + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... class CoordSystem: """ @@ -61,6 +63,8 @@ class CoordSystem: ] # value = {'s': , 't': } s: typing.ClassVar[CoordSystem] # value = t: typing.ClassVar[CoordSystem] # value = + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __eq__(self, other: typing.Any) -> bool: ... def __getstate__(self) -> int: ... def __hash__(self) -> int: ... @@ -77,6 +81,8 @@ class CoordSystem: def value(self) -> int: ... class ImpactX: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def DistributionMap( self, lev: int ) -> amrex.space3d.amrex_3d_pybind.DistributionMapping: ... @@ -371,6 +377,8 @@ class ImpactX: class ImpactXParConstIter( amrex.space3d.amrex_3d_pybind.ParConstIter_pureSoA_8_0_default ): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... @typing.overload def __init__( self, @@ -399,6 +407,8 @@ class ImpactXParConstIter( """ class ImpactXParIter(amrex.space3d.amrex_3d_pybind.ParIter_pureSoA_8_0_default): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... @typing.overload def __init__( self, @@ -431,6 +441,8 @@ class ImpactXParticleContainer( ): const_iterator = ImpactXParConstIter iterator = ImpactXParIter + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def add_n_particles( self, x: amrex.space3d.amrex_3d_pybind.PODVector_real_std, @@ -525,6 +537,8 @@ class ImpactXParticleContainer( """ class RefPart: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... @staticmethod def load_file(ref: RefPart, madx_file): """ @@ -720,6 +734,6 @@ __author__: str = ( "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." ) __license__: str = "BSD-3-Clause-LBNL" -__version__: str = "24.12" +__version__: str = "25.01" s: CoordSystem # value = t: CoordSystem # value = diff --git a/src/python/impactx/impactx_pybind/distribution.pyi b/src/python/impactx/impactx_pybind/distribution.pyi index 105a1c40a..2b284687b 100644 --- a/src/python/impactx/impactx_pybind/distribution.pyi +++ b/src/python/impactx/impactx_pybind/distribution.pyi @@ -17,12 +17,16 @@ __all__ = [ ] class Empty: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__(self) -> None: """ Sets all values to zero. """ class Gaussian: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, lambdaX: float, @@ -40,6 +44,8 @@ class Gaussian: """ class KVdist: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, lambdaX: float, @@ -58,6 +64,8 @@ class KVdist: """ class Kurth4D: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, lambdaX: float, @@ -76,6 +84,8 @@ class Kurth4D: """ class Kurth6D: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, lambdaX: float, @@ -96,6 +106,8 @@ class Kurth6D: """ class Semigaussian: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, lambdaX: float, @@ -113,6 +125,8 @@ class Semigaussian: """ class Thermal: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, k: float, @@ -129,6 +143,8 @@ class Thermal: """ class Triangle: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, lambdaX: float, @@ -149,6 +165,8 @@ class Triangle: """ class Waterbag: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, lambdaX: float, diff --git a/src/python/impactx/impactx_pybind/elements/__init__.pyi b/src/python/impactx/impactx_pybind/elements/__init__.pyi index 0229e0c5b..950e24875 100644 --- a/src/python/impactx/impactx_pybind/elements/__init__.pyi +++ b/src/python/impactx/impactx_pybind/elements/__init__.pyi @@ -46,6 +46,8 @@ __all__ = [ ] class Aperture(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, xmax: float, @@ -116,6 +118,8 @@ class Aperture(mixin.Named, mixin.Thin, mixin.Alignment): def ymax(self, arg1: float) -> None: ... class BeamMonitor(mixin.Thin): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, name: str, @@ -179,6 +183,8 @@ class BeamMonitor(mixin.Thin): def tn(self, arg1: float) -> None: ... class Buncher(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, V: float, @@ -217,6 +223,8 @@ class Buncher(mixin.Named, mixin.Thin, mixin.Alignment): def k(self, arg1: float) -> None: ... class CFbend(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -257,6 +265,8 @@ class CFbend(mixin.Named, mixin.Thick, mixin.Alignment): def rc(self, arg1: float) -> None: ... class ChrAcc(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -297,6 +307,8 @@ class ChrAcc(mixin.Named, mixin.Thick, mixin.Alignment): def ez(self, arg1: float) -> None: ... class ChrDrift(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -321,6 +333,8 @@ class ChrDrift(mixin.Named, mixin.Thick, mixin.Alignment): """ class ChrPlasmaLens(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -361,6 +375,8 @@ class ChrPlasmaLens(mixin.Named, mixin.Thick, mixin.Alignment): def unit(self, arg1: int) -> None: ... class ChrQuad(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -401,6 +417,8 @@ class ChrQuad(mixin.Named, mixin.Thick, mixin.Alignment): def unit(self, arg1: int) -> None: ... class ConstF(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -449,6 +467,8 @@ class ConstF(mixin.Named, mixin.Thick, mixin.Alignment): def ky(self, arg1: float) -> None: ... class DipEdge(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, psi: float, @@ -503,6 +523,8 @@ class DipEdge(mixin.Named, mixin.Thin, mixin.Alignment): def rc(self, arg1: float) -> None: ... class Drift(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -527,6 +549,8 @@ class Drift(mixin.Named, mixin.Thick, mixin.Alignment): """ class Empty(mixin.Thin): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__(self) -> None: """ This element does nothing. @@ -543,6 +567,8 @@ class Empty(mixin.Thin): """ class ExactDrift(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -567,6 +593,8 @@ class ExactDrift(mixin.Named, mixin.Thick, mixin.Alignment): """ class ExactSbend(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -607,6 +635,8 @@ class ExactSbend(mixin.Named, mixin.Thick, mixin.Alignment): def phi(self, arg1: float) -> None: ... class Kicker(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, xkick: float, @@ -646,6 +676,8 @@ class Kicker(mixin.Named, mixin.Thin, mixin.Alignment): def ykick(self, arg1: float) -> None: ... class KnownElementsList: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... @typing.overload def __init__(self) -> None: ... @typing.overload @@ -779,6 +811,8 @@ class KnownElementsList: """ class Marker(mixin.Named, mixin.Thin): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__(self, arg0: str) -> None: """ This named element does nothing. @@ -795,6 +829,8 @@ class Marker(mixin.Named, mixin.Thin): """ class Multipole(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, multipole: int, @@ -841,6 +877,8 @@ class Multipole(mixin.Named, mixin.Thin, mixin.Alignment): def multipole(self, arg1: float) -> None: ... class NonlinearLens(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, knll: float, @@ -879,6 +917,8 @@ class NonlinearLens(mixin.Named, mixin.Thin, mixin.Alignment): def knll(self, arg1: float) -> None: ... class PRot(mixin.Named, mixin.Thin): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__(self, phi_in: float, phi_out: float, name: str | None = None) -> None: """ An exact pole-face rotation in the x-z plane. Both angles are in degrees. @@ -909,6 +949,8 @@ class PRot(mixin.Named, mixin.Thin): def phi_out(self, arg1: float) -> None: ... class PlaneXYRot(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, angle: float, @@ -941,6 +983,8 @@ class PlaneXYRot(mixin.Named, mixin.Thin, mixin.Alignment): class Programmable(mixin.Named): ds: float nslice: int + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float = 0.0, nslice: int = 1, name: str | None = None ) -> None: @@ -999,6 +1043,8 @@ class Programmable(mixin.Named): def threadsafe(self, arg1: bool) -> None: ... class Quad(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -1031,6 +1077,8 @@ class Quad(mixin.Named, mixin.Thick, mixin.Alignment): def k(self, arg1: float) -> None: ... class RFCavity(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -1089,6 +1137,8 @@ class RFCavity(mixin.Named, mixin.Thick, mixin.Alignment): def phase(self, arg1: float) -> None: ... class Sbend(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -1121,6 +1171,8 @@ class Sbend(mixin.Named, mixin.Thick, mixin.Alignment): def rc(self, arg1: float) -> None: ... class ShortRF(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, V: float, @@ -1167,6 +1219,8 @@ class ShortRF(mixin.Named, mixin.Thin, mixin.Alignment): def phase(self, arg1: float) -> None: ... class SoftQuadrupole(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -1209,6 +1263,8 @@ class SoftQuadrupole(mixin.Named, mixin.Thick, mixin.Alignment): def mapsteps(self, arg1: int) -> None: ... class SoftSolenoid(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -1259,6 +1315,8 @@ class SoftSolenoid(mixin.Named, mixin.Thick, mixin.Alignment): def unit(self, arg1: float) -> None: ... class Sol(mixin.Named, mixin.Thick, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, ds: float, @@ -1291,6 +1349,8 @@ class Sol(mixin.Named, mixin.Thick, mixin.Alignment): def ks(self, arg1: float) -> None: ... class TaperedPL(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, k: float, @@ -1343,6 +1403,8 @@ class TaperedPL(mixin.Named, mixin.Thin, mixin.Alignment): def unit(self, arg1: int) -> None: ... class ThinDipole(mixin.Named, mixin.Thin, mixin.Alignment): + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__( self, theta: float, diff --git a/src/python/impactx/impactx_pybind/elements/mixin.pyi b/src/python/impactx/impactx_pybind/elements/mixin.pyi index d6822303b..ff84c2db0 100644 --- a/src/python/impactx/impactx_pybind/elements/mixin.pyi +++ b/src/python/impactx/impactx_pybind/elements/mixin.pyi @@ -7,6 +7,8 @@ from __future__ import annotations __all__ = ["Alignment", "Named", "Thick", "Thin"] class Alignment: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__(self) -> None: """ Mixin class for lattice elements with horizontal/vertical alignment errors. @@ -34,6 +36,8 @@ class Alignment: def rotation(self, arg1: float) -> None: ... class Named: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... @property def has_name(self) -> bool: ... @property @@ -45,6 +49,8 @@ class Named: def name(self, arg1: str) -> None: ... class Thick: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__(self, ds: float, nslice: float = 1) -> None: """ Mixin class for lattice elements with finite length. @@ -65,6 +71,8 @@ class Thick: def nslice(self, arg1: int) -> None: ... class Thin: + @staticmethod + def _pybind11_conduit_v1_(*args, **kwargs): ... def __init__(self) -> None: """ Mixin class for lattice elements with zero length. diff --git a/tests/python/conftest.py b/tests/python/conftest.py index 884bb1677..7b9234f68 100644 --- a/tests/python/conftest.py +++ b/tests/python/conftest.py @@ -32,8 +32,6 @@ def amrex_init(tmpdir): "amrex.signal_handling=0", # abort GPU runs if out-of-memory instead of swapping to host RAM "amrex.abort_on_out_of_gpu_memory=1", - # do not rely on implicit host-device memory transfers - "amrex.the_arena_is_managed=0", ] ) yield