Add nonlinear thin dipole fringe field map (#1158)

* Set up modified DipEdge inputs.
* Add intermediate variable definitions.
* Implement nonlinear map.
* Update default values for field integral inputs.
* Update documentation.
* Add entry/exit input flag.
* Fix enumeration type conversion.
* Add g=0 special case; make K2 optional.
* Remove debug print statement.
* Add scaling benchmark example.
* Fix c7 coefficient and invert 1 + delta.
* Fix one missing instance of 1 + delta.
* Update benchmark parameter values.
* Update tolerance.
* Update parameters.rst
Add some additional details to documentation.
* Update python.rst
Add additional documentation details.
* Vectorization Comment
Probably worth vectorizing now...

Author: cemitch99

docs/source/usage/examples.rst

@@ -42,7 +42,7 @@ Single Particle Dynamics
    examples/charge_sign/README.rst
    examples/symplectic_integration/README.rst
    examples/active_plasma_lens/README.rst
-
+   examples/edge_effects/README.rst
 
 Collective Effects
 ------------------

docs/source/usage/parameters.rst

@@ -322,12 +322,37 @@ This element requires these additional parameters:
 ``dipedge``
 ^^^^^^^^^^^
 
-``dipedge`` for dipole edge focusing. This requires these additional parameters:
+``dipedge`` for dipole edge focusing. The model here is based on:
+
+K. Hwang and S. Y. Lee, "Dipole fringe field map for compact synchrotrons," Phys. Rev. Accel. Beams 18, 122401 (2015)
+
+as represented in the explicit, symplectic form provided in:
+
+C. Mitchell and K. Hwang, "Explicit symplectic representations of nonlinear dipole fringe field maps," in Proc. NAPAC2025, TUP040, Sacramento, CA, 2025
+
+Here, ``g`` denotes the magnetic gap, which is a length scale that sets the rate of decay of the fringe field.  The values ``K0`` - ``K6`` denote dimensionless field integrals, describing the shape of the fringe field, as defined in eqs. (28-34) of the first reference above.  In particular, ``K2`` is the well-known fringe field parameter denoted ``FINT`` in MAD-X.  The default values of the field integrals ``K0`` - ``K6`` are those given in eq. (52), corresponding to a ``tanh`` (i.e. logistic) field profile.
+
+If ``model = "linear"``, then the linearized map is used.  This model is identical to:
+
+* K. L. Brown, SLAC Report No. 75 (1982)
+
+when expanded to first order in ``g/rc`` (gap / radius of curvature).
+
+This requires these additional parameters:
 
 * ``<element_name>.psi`` (``float``, in radians) the pole face rotation angle
 * ``<element_name>.rc`` (``float``, in meters) the bend radius
-* ``<element_name>.g`` (``float``, in meters) the gap size
-* ``<element_name>.K2`` (``float``, dimensionless) normalized field integral for fringe field
+* ``<element_name>.g`` (``float``, in meters) the full magnetic gap size
+* ``<element_name>.R`` (``float``, in meters) scale length for the field integrals (default: ``1 m``)
+* ``<element_name>.K0`` (``float``, dimensionless) normalized field integral for fringe field
+* ``<element_name>.K1`` (``float``, dimensionless) normalized field integral for fringe field
+* ``<element_name>.K2`` (``float``, dimensionless) normalized field integral for fringe field (FINT)
+* ``<element_name>.K3`` (``float``, dimensionless) normalized field integral for fringe field
+* ``<element_name>.K4`` (``float``, dimensionless) normalized field integral for fringe field
+* ``<element_name>.K5`` (``float``, dimensionless) normalized field integral for fringe field
+* ``<element_name>.K6`` (``float``, dimensionless) normalized field integral for fringe field
+* ``<element_name>.model`` (``string``) the fringe field model: ``linear`` (default) or ``nonlinear``
+* ``<element_name>.location`` (``string``) the fringe field edge location: ``entry`` (default) or ``exit``
 * ``<element_name>.dx`` (``float``, in meters) horizontal translation error
 * ``<element_name>.dy`` (``float``, in meters) vertical translation error
 * ``<element_name>.rotation`` (``float``, in degrees) rotation error in the transverse plane

docs/source/usage/python.rst

@@ -801,20 +801,40 @@ This module provides elements and methods for the accelerator lattice.
 
    Edge focusing associated with bend entry or exit
 
-   This model assumes a first-order effect of nonzero gap.
-   Here we use the linear fringe field map, given to first order in g/rc (gap / radius of curvature).
+   The model here is based on:
 
-   References:
-
-   * K. L. Brown, SLAC Report No. 75 (1982).
    * K. Hwang and S. Y. Lee, PRAB 18, 122401 (2015).
 
-   :param psi: Pole face angle in rad
-   :param rc: Radius of curvature in m
-   :param g: Gap parameter in m
-   :param K2: Fringe field integral (unitless)
-   :param dx: horizontal translation error in m
-   :param dy: vertical translation error in m
+   as represented in the explicit, symplectic form provided in:
+
+   * C. Mitchell and K. Hwang, in Proc. NAPAC2025, TUP040, Sacramento, CA (2025).
+
+   Here, ``g`` denotes the magnetic gap, which is a length scale that sets the rate of decay of the fringe field.  The values ``K0`` - ``K6`` denote
+   dimensionless field integrals, describing the shape of the fringe field, as defined in eqs. (28-34) of the first reference above.  In
+   particular, ``K2`` is the well-known fringe field parameter denoted ``FINT`` in MAD-X.  The default values of the field integrals ``K0`` - ``K6`` are
+   those given in eq. (52), corresponding to a ``tanh`` (i.e. logistic) field profile.
+
+   When ``model = "linear"``, the linearized map is used.  This model is identical to:
+
+   * K. L. Brown, SLAC Report No. 75 (1982)
+
+   when expanded to first order in ``g/rc`` (gap / radius of curvature).
+
+   :param psi: Pole face angle [radians]
+   :param rc: Radius of curvature [m]
+   :param g: Gap parameter [m]
+   :param R: Length scale used in fringe field integrals [m]
+   :param K0: Fringe field integral [unitless]
+   :param K1: Fringe field integral [unitless]
+   :param K2: Fringe field integral [unitless]
+   :param K3: Fringe field integral [unitless]
+   :param K4: Fringe field integral [unitless]
+   :param K5: Fringe field integral [unitless]
+   :param K6: Fringe field integral [unitless]
+   :param model: the fringe field model: ``linear`` (default) or ``nonlinear``
+   :param location: the fringe field edge location: ``entry`` (default) or ``exit``
+   :param dx: horizontal translation error [m]
+   :param dy: vertical translation error [m]
    :param rotation: rotation error in the transverse plane [degrees]
    :param name: an optional name for the element
 

examples/CMakeLists.txt

@@ -1698,3 +1698,14 @@ add_impactx_test(solenoid-softedge-solvable.py
     examples/solenoid_softedge/analysis_solenoid_softedge_solvable.py
     OFF  # no plot script yet
 )
+
+# Exactly-solvable (non-uniform) soft-edge solenoid ##############
+#
+file(COPY ${ImpactX_SOURCE_DIR}/examples/edge_effects/initial_coords.csv
+     DESTINATION ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/dipedge-nonlinear.py)
+add_impactx_test(dipedge-nonlinear.py
+    examples/edge_effects/run_dipedge.py
+    ON   # ImpactX MPI-parallel
+    examples/edge_effects/analysis_dipedge.py
+    OFF  # no plot script yet
+)

examples/edge_effects/README.rst

@@ -0,0 +1,45 @@
+.. _dipedge-nonlinear:
+
+Error scaling test for nonlinear dipole fringe field map
+========================================================
+
+This benchmark tests the use of the nonlinear ``DipEdge`` model for integrating through a dipole fringe field.
+
+Six distinct initial conditions are tested for a nominal proton beam with 800 MeV kinetic energy.  The values of the field integrals K0-K6 are set to the default
+values, as used in:
+
+K. Hwang and S. Y. Lee, "Dipole fringe field map for compact synchrotrons," Phys. Rev. Accel. Beams 18, 122401 (2015)
+
+The initial conditions are chosen with increasing distance from the origin in phase space.  The value of the Lie generator is a dynamical
+invariant of the ideal fringe field map.  In reality, there is an error in the final variables that scales with (x,px,y,py,t,pt,g) as degree 3.  As a result, the
+Lie generator is not an exact invariant of the numerically-computed fringe field map.
+
+In this test, the change in the Lie generator for each initial condition should coincide with its (small) nominal value.
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_dipedge.py``
+
+For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python: Script
+
+       .. literalinclude:: run_dipedge.py
+          :language: python3
+          :caption: You can copy this file from ``examples/edge_effects/run_dipedge.py``.
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_exact_quad.py``
+
+   .. literalinclude:: analysis_exact_quad.py
+      :language: python3
+      :caption: You can copy this file from ``examples/edge_effects/analysis_dipedge.py``.

examples/edge_effects/analysis_dipedge.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+
+import math
+
+import numpy as np
+import openpmd_api as io
+
+# initial/final beam
+series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
+last_step = list(series.iterations)[-1]
+initial = series.iterations[1].particles["beam"].to_df()
+beam_final = series.iterations[last_step].particles["beam"]
+final = beam_final.to_df()
+
+# Basic input parameters
+g = 1.0e-3
+phi = math.pi / 8.0
+rc = 10.0
+R = 1.0
+K0 = math.pi**2 / 6.0
+K3 = 1.0 / 6.0
+Kar = [K0, 0, 0, K3, 0, 0, 0]
+delta = 0.0
+
+# Derived quantities
+cs = math.cos(phi)
+sn = math.sin(phi)
+tn = sn / cs
+sc = 1.0 / cs
+
+# Lie generator coefficients
+c1 = g * Kar[1] / (rc * cs)
+c2 = sn * g**2 * Kar[0] / rc**2 * 1.0 / (2.0 * cs**3 * (1 + delta))
+c3 = g**2 / rc * Kar[0] / (cs**2 * (1 + delta))
+c4 = 1 / (1 + delta) * g / rc * Kar[1] * sn / cs**2
+c5 = sn / cs * 1.0 / (2 * rc)
+c6 = g * Kar[1] / rc * sn**2 / (4 * rc * (1 + delta) * cs**3)
+c7 = (
+    1
+    / (2 * cs**3 * (1 + delta))
+    * (g * Kar[1] / (2 * rc**2) + (1 + sn**2) * g / rc**2 * Kar[2])
+)
+c8 = 1 / 6 * tn**3 / (2 * rc**2 * (1 + delta))
+c9 = 1 / 2 * (tn * sc**2 / (2 * rc**2 * (1 + delta)))
+c10 = 1 / (2 * (1 + delta)) * tn**2 / rc
+c11 = 1 / (2 * rc * (1 + delta))
+c12 = 1 / 24 * (4 / cs - 8 / cs**3) * Kar[3] / (rc**2 * g * (1 + delta))
+c13 = sn**2 / (2 * cs**3) * g**2 / (rc * R) * Kar[4]
+c14 = 1 / 2 * sn / cs**3 * g / (rc * R) * Kar[5]
+c15 = Kar[6] / (rc * R) * 1 / cs**3
+
+xi = initial["position_x"]
+pxi = initial["momentum_x"]
+yi = initial["position_y"]
+pyi = initial["momentum_y"]
+ti = initial["position_t"]
+pti = initial["momentum_t"]
+
+Omega_initial = (
+    xi * c1
+    - xi * c2
+    + pxi * c3
+    + (xi * pxi - yi * pyi) * c4
+    + (xi**2 - yi**2) * c5
+    - xi**2 * c6
+    + yi**2 * c7
+    - xi**3 * c8
+    + xi * yi**2 * c9
+    + (xi**2 * pxi - yi**2 * pxi - 2 * xi * yi * pyi) * c10
+    - yi**2 * pxi * c11
+    + yi**4 * c12
+    + xi * c13
+    + (yi**2 - xi**2) * c14
+    + (xi * yi**2 / 2 - xi**3 / 6) * c15
+)
+
+xf = final["position_x"]
+pxf = final["momentum_x"]
+yf = final["position_y"]
+pyf = final["momentum_y"]
+tf = final["position_t"]
+ptf = final["momentum_t"]
+
+Omega_final = (
+    xf * c1
+    - xf * c2
+    + pxf * c3
+    + (xf * pxf - yf * pyf) * c4
+    + (xf**2 - yf**2) * c5
+    - xf**2 * c6
+    + yf**2 * c7
+    - xf**3 * c8
+    + xf * yf**2 * c9
+    + (xf**2 * pxf - yf**2 * pxf - 2 * xf * yf * pyf) * c10
+    - yf**2 * pxf * c11
+    + yf**4 * c12
+    + xf * c13
+    + (yf**2 - xf**2) * c14
+    + (xf * yf**2 / 2 - xf**3 / 6) * c15
+)
+
+Delta_Omega = (Omega_final - Omega_initial).abs()
+
+dx = (xf - xi).abs().max()
+dpx = (pxf - pxi).abs().max()
+dy = (yf - yi).abs().max()
+dpy = (pyf - pyi).abs().max()
+dt = (tf - ti).abs().max()
+dpt = (ptf - pti).abs().max()
+
+print("Change in the coordinates and momenta:")
+print("dx", dx)
+print("dpx", dpx)
+print("dy", dy)
+print("dpy", dpy)
+print("dt", dt)
+print("dpt", dpt)
+
+print("Change in the Lie generator, for each initial condition:")
+print(Delta_Omega)
+
+atol = 1.5e-10
+print(f"  atol={atol}")
+
+assert np.allclose(
+    [Delta_Omega.max()],
+    [
+        0.0,
+    ],
+    atol=atol,
+)

examples/edge_effects/initial_coords.csv

@@ -0,0 +1,8 @@
+ x px y py t pt
+0 1.0e-3 1.0e-3 1.0e-3 1.0e-3 1.0e-4 0
+1 2.0e-3 2.0e-3 2.0e-3 2.0e-3 2.0e-4 0
+2 4.0e-3 4.0e-3 4.0e-3 4.0e-3 4.0e-4 0
+3 8.0e-3 8.0e-3 8.0e-3 8.0e-3 8.0e-4 0
+4 1.6e-2 1.6e-2 1.6e-2 1.6e-2 1.6e-3 0
+5 3.2e-2 3.2e-2 3.2e-2 3.2e-2 3.2e-3 0
+6 6.4e-2 6.4e-2 6.4e-2 6.4e-2 6.4e-3 0