add pythonization of transformation (#360)

* add pythonization of transformation

* add transformation.cpp to cmakelists.txt

* Cleanups

* add documentation on coordinate transformation

* add test of forward-inverse transformation

* clean markdown

* docs: fix section formatting

* code cleaning

* isort/black code cleaning

* fix isort/black incompatibility

* Apply suggestions from code review

Co-authored-by: Axel Huebl <[email protected]>

* code review suggestions

* more cleaning - remove amrex import

* Aesthetic Whitespaces Before Sections

* Move to Unit Test Location

* Removed unused variable

---------

Co-authored-by: Axel Huebl <[email protected]>

.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

.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/

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.

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.

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 `<http://www.physics.umd.edu/dsat/docs/MaryLieMan.pdf>`__
    :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`.

docs/source/usage/tests

@@ -0,0 +1 @@
+../../../tests

src/python/CMakeLists.txt

@@ -9,4 +9,5 @@ target_sources(pyImpactX
     ImpactX.cpp
     ImpactXParticleContainer.cpp
     ReferenceParticle.cpp
+    transformation.cpp
 )

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

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 <particles/ReferenceParticle.H>
+#include <particles/transformation/CoordinateTransformation.H>
+
+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_<transformation::Direction>(m, "TransformationDirection")
+        .value("to_fixed_s", transformation::Direction::to_fixed_s)
+        .value("to_fixed_t", transformation::Direction::to_fixed_t)
+        .export_values();
+}

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

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 <developers-testing>`.
+
+.. tab-set::
+
+   .. tab-item:: Python Script
+
+       .. literalinclude:: test_transformation.py
+          :language: python3
+          :caption: You can copy this file from ``tests/python/test_transformation.py``.