diff --git a/.editorconfig b/.editorconfig index cf750da3a..7e2d70bc2 100644 --- a/.editorconfig +++ b/.editorconfig @@ -35,3 +35,6 @@ indent_size = unset # isort config force_sort_within_sections = true known_first_party = amrex,impactx,picmistandard,pywarpx,warpx +# same as the "black" multi-line import style +multi_line_output = 3 +include_trailing_comma = True diff --git a/.gitignore b/.gitignore index 7e8820fd6..d45913ed6 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,7 @@ dist/ docs/amrex-doxygen-web.tag.xml docs/warpx-doxygen-web.tag.xml docs/impactx-doxygen-web.tag.xml +docs/openpmd-api-doxygen-web.tag.xml docs/doxyhtml/ docs/doxyxml/ docs/source/_static/ diff --git a/docs/source/conf.py b/docs/source/conf.py index 606130271..77611018d 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -81,7 +81,7 @@ # # This is also used if you do content translation via gettext catalogs. # Usually you set "language" from the command line for these cases. -language = None +language = "en" # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 8d52cf987..fb4860dd3 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -3,7 +3,7 @@ Examples ======== -This section allows you to **download input files** that correspond to different physical situations. +This section allows you to **download input files** that correspond to different physical situations or test different code features. .. toctree:: :maxdepth: 1 @@ -25,4 +25,13 @@ This section allows you to **download input files** that correspond to different examples/quadrupole_softedge/README.rst examples/positron_channel/README.rst -For every change of the ImpactX code base, each of these examples are continuously tested and benchmarked. + +Unit tests +---------- + +.. toctree:: + :maxdepth: 1 + + tests/python/transformation.rst + +For every change of the ImpactX code base, each of these examples and tests are continuously tested and benchmarked. diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 1ca25267b..3a55b95b2 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -637,3 +637,22 @@ This module provides elements for the accelerator lattice. (optional); default is a tanh fringe field model based on ``__ :param mapsteps: number of integration steps per slice used for map and reference particle push in applied fields :param nslice: number of slices used for the application of space charge + + +Coordinate Transformation +------------------------- + +.. py:class:: impactx.TransformationDirection + + Enumerated type indicating whether to transform to fixed :math:`s` or fixed :math:`t` coordinate system when applying ``impactx.coordinate_transformation``. + + :param to_fixed_t: + :param to_fixed_s: + +Function +.. py:method:: impactx.coordinate_transformation(pc, direction) + + Function to transform the coordinates of the particles in a particle container either to fixed :math:`t` or to fixed :math:`s`. + + :param pc: ``impactx.particle_container`` whose particle coordinates are to be transformed. + :param direction: enumerated type ``impactx.TransformationDirection``, indicates whether to transform to fixed :math:`s` or fixed :math:`t`. diff --git a/docs/source/usage/tests b/docs/source/usage/tests new file mode 120000 index 000000000..3cafd37fe --- /dev/null +++ b/docs/source/usage/tests @@ -0,0 +1 @@ +../../../tests \ No newline at end of file diff --git a/src/python/CMakeLists.txt b/src/python/CMakeLists.txt index 55ee28f66..d658332a8 100644 --- a/src/python/CMakeLists.txt +++ b/src/python/CMakeLists.txt @@ -9,4 +9,5 @@ target_sources(pyImpactX ImpactX.cpp ImpactXParticleContainer.cpp ReferenceParticle.cpp + transformation.cpp ) diff --git a/src/python/pyImpactX.cpp b/src/python/pyImpactX.cpp index f07bee0d7..3846c917d 100644 --- a/src/python/pyImpactX.cpp +++ b/src/python/pyImpactX.cpp @@ -21,6 +21,7 @@ void init_elements(py::module&); void init_ImpactX(py::module&); void init_impactxparticlecontainer(py::module&); void init_refparticle(py::module&); +void init_transformation(py::module&); PYBIND11_MODULE(impactx_pybind, m) { // make sure AMReX types are known @@ -43,6 +44,7 @@ PYBIND11_MODULE(impactx_pybind, m) { init_elements(m); init_refparticle(m); init_impactxparticlecontainer(m); + init_transformation(m); init_ImpactX(m); // API runtime version diff --git a/src/python/transformation.cpp b/src/python/transformation.cpp new file mode 100644 index 000000000..1eaef9cb5 --- /dev/null +++ b/src/python/transformation.cpp @@ -0,0 +1,25 @@ +/* Copyright 2021-2023 The ImpactX Community + * + * Authors: Ryan Sandberg, Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include "pyImpactX.H" + +#include +#include + +namespace py = pybind11; +using namespace impactx; + +void init_transformation(py::module& m) +{ + m.def("coordinate_transformation", + &transformation::CoordinateTransformation, + "Transform coordinates from fixed s to fixed to or vice versa." + ); + + py::enum_(m, "TransformationDirection") + .value("to_fixed_s", transformation::Direction::to_fixed_s) + .value("to_fixed_t", transformation::Direction::to_fixed_t) + .export_values(); +} diff --git a/tests/python/test_transformation.py b/tests/python/test_transformation.py new file mode 100644 index 000000000..8f0e0cb10 --- /dev/null +++ b/tests/python/test_transformation.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Ryan Sandberg, Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import numpy as np + +from impactx import ( + Config, + ImpactX, + ImpactXParIter, + RefPart, + TransformationDirection, + coordinate_transformation, + distribution, + elements, +) + + +def test_transformation(): + """ + This test ensures s->t and t->s transformations + do round-trip. + """ + sim = ImpactX() + + # set numerical parameters and IO control + sim.particle_shape = 2 # B-spline order + sim.space_charge = False + # sim.diagnostics = False # benchmarking + sim.slice_step_diagnostics = True + + # domain decomposition & space charge mesh + sim.init_grids() + + # load a 1 GeV electron beam with an initial + # unnormalized rms emittance of 2 nm + energy_MeV = 1e3 # reference energy + energy_gamma = energy_MeV / 0.510998950 + bunch_charge_C = 1.0e-9 # used with space charge + npart = 10000 # number of macro particles + + # reference particle + pc = sim.particle_container() + ref = pc.ref_particle() + ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_energy_MeV(energy_MeV) + + # particle bunch + distr = distribution.Gaussian( + sigmaX=3e-6, + sigmaY=3e-6, + sigmaT=1e-2, + sigmaPx=1.33 / energy_gamma, + sigmaPy=1.33 / energy_gamma, + sigmaPt=100 / energy_gamma, + muxpx=-0.5, + muypy=0.4, + mutpt=0.8, + ) + sim.add_particles(bunch_charge_C, distr, npart) + + rbc_s0 = pc.reduced_beam_characteristics() + coordinate_transformation(pc, TransformationDirection.to_fixed_t) + rbc_t = pc.reduced_beam_characteristics() + coordinate_transformation(pc, TransformationDirection.to_fixed_s) + rbc_s = pc.reduced_beam_characteristics() + + # clean shutdown + del sim + + # assert that forward-inverse transformation of the beam leaves beam unchanged + atol = 1e-14 + rtol = 1e-10 + for key, val in rbc_s0.items(): + if not np.isclose(val, rbc_s[key], rtol=rtol, atol=atol): + print(f"initial[{key}]={val}, final[{key}]={rbc_s[key]} not equal") + assert np.isclose(val, rbc_s[key], rtol=rtol, atol=atol) + # assert that the t-based beam is different, at least in the following keys: + large_st_diff_keys = [ + "beta_x", + "beta_y", + "emittance_y", + "emittance_x", + "sig_y", + "sig_x", + "t_mean", + ] + for key in large_st_diff_keys: + rel_error = (rbc_s0[key] - rbc_t[key]) / rbc_s0[key] + assert abs(rel_error) > 1 diff --git a/tests/python/transformation.rst b/tests/python/transformation.rst new file mode 100644 index 000000000..23515eef1 --- /dev/null +++ b/tests/python/transformation.rst @@ -0,0 +1,28 @@ +.. _tests-transformation: + +Transformation +============== + +Test the t/s transformations on an electron beam. + +We use a long electron beam, :math:`L_z=1` cm, with significant correlations in :math:`x-px`, :math:`y-py`, and :math:`t-pt`. +The beam has average energy 1 GeV. + +This tests that the t/s transforms are inverses of each other +Specifically, in this test the :math:`t`- and :math:`s`-coordinates of the beam must differ substantially +and the forward-inverse transformed coordinates must agree with the initial coordinates. +That is, we require that :math:`to_fixed_s` ( :math:`to_fixed_t` (initial beam)) = initial beam. + + +Run +--- + +This file is run from :ref:`pytest `. + +.. tab-set:: + + .. tab-item:: Python Script + + .. literalinclude:: test_transformation.py + :language: python3 + :caption: You can copy this file from ``tests/python/test_transformation.py``.