Add a thin dipole element

docs/source/usage/examples.rst

@@ -28,6 +28,7 @@ This section allows you to **download input files** that correspond to different
    examples/cfbend/README.rst
    examples/compression/README.rst
    examples/kicker/README.rst
+   examples/thin_dipole/README.rst
 
 
 Unit tests

docs/source/usage/parameters.rst

@@ -509,6 +509,13 @@ Lattice Elements
 
             * ``<element_name>.units`` (``string``) specification of units: ``dimensionless`` (default, in units of the magnetic rigidity of the reference particle) or ``T-m``
 
+        * ``thin_dipole`` for a thin dipole element.
+          This requires these additional parameters:
+
+            * ``<element_name>.theta`` (``float``, in degrees) dipole bend angle
+
+            * ``<element_name>.rc`` (``float``, in meters) effective radius of curvature
+
         * ``beam_monitor`` a beam monitor, writing all beam particles at fixed ``s`` to openPMD files.
           If the same element name is used multiple times, then an output series is created with multiple outputs.
 

docs/source/usage/python.rst

@@ -725,6 +725,15 @@ References:
    :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
 
+.. py:class:: impactx.elements.ThinDipole(theta, rc)
+
+   A general thin dipole element.
+
+   :param theta: Bend angle (degrees)
+   :param rc: Effective curvature radius (meters)
+
+Reference:
+   * G. Ripken and F. Schmidt, Thin-Lens Formalism for Tracking, CERN/SL/95-12 (AP), 1995.
 
 Coordinate Transformation
 -------------------------

examples/CMakeLists.txt

@@ -639,3 +639,21 @@ add_impactx_test(hvkicker_madx.py
         examples/kicker/analysis_kicker.py
         OFF  # no plot script yet
 )
+
+# Thin dipole ########################################################
+#
+# w/o space charge
+add_impactx_test(thin_dipole
+    examples/thin_dipole/input_thin_dipole.in
+      ON  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/thin_dipole/analysis_thin_dipole.py
+    OFF  # no plot script yet
+)
+add_impactx_test(thin_dipole.py
+    examples/thin_dipole/run_thin_dipole.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/thin_dipole/analysis_thin_dipole.py
+    OFF  # no plot script yet
+)

examples/thin_dipole/README.rst

@@ -0,0 +1,61 @@
+.. _examples-thin-dipole:
+
+Thin Dipole
+===========
+
+This test involves tracking a 5 MeV proton beam through a 90 degree sector bend, using two different methods:
+
+#. Using a sequence of drifts and thin kicks as described by:
+G. Ripken and F. Schmidt, "A Symplectic Six-Dimensional Thin-Lens Formalism for Tracking," CERN/SL/95-12 (1995).
+
+#. Using the exact nonlinear transfer map for a sector bend as described by:
+D. Bruhwiler et al, "Symplectic Propagation of the Map, Tangent Map and Tangent Map Derivative through
+Quadrupole and Combined-Function Dipole Magnets without Truncation," THP41C, EPAC98, pp. 1171-1173 (1998).
+
+To compare the two models, the beam is tracked using 1), followed by the inverse of 2).
+The beam moments and emittances should coincide with those of the initial distribution.
+
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_thin_dipole.py`` or
+* ImpactX **executable** using an input file: ``impactx input_thin_dipole.in``
+
+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_thin_dipole.py
+          :language: python3
+          :caption: You can copy this file from ``examples/thin_dipole/run_thin_dipole.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_thin_dipole.in
+          :language: ini
+          :caption: You can copy this file from ``examples/thin_dipole/input_thin_dipole.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_thin_dipole.py``
+
+   .. literalinclude:: analysis_thin_dipole.py
+      :language: python3
+      :caption: You can copy this file from ``examples/thin_dipole/analysis_thin_dipole.py``.
+
+
+Visualize
+---------
+
+.. note::
+
+   TODO :)

examples/thin_dipole/analysis_thin_dipole.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 -*-
+
+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 = 1.8 * 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],
+    [
+        1.0e-3,
+        1.0e-3,
+        3.0e-1,
+        2.0e-7,
+        2.0e-7,
+        6.0e-5,
+    ],
+    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 = 1.8 * 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],
+    [
+        1.0e-3,
+        1.0e-3,
+        3.0e-1,
+        2.0e-7,
+        2.0e-7,
+        6.0e-5,
+    ],
+    rtol=rtol,
+    atol=atol,
+)

examples/thin_dipole/input_thin_dipole.in

@@ -0,0 +1,57 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.kin_energy = 5.0
+beam.charge = 1.0e-9
+beam.particle = proton
+beam.distribution = waterbag
+beam.sigmaX = 1.0e-3
+beam.sigmaY = 1.0e-3
+beam.sigmaT = 0.3
+beam.sigmaPx = 2.0e-4
+beam.sigmaPy = 2.0e-4
+beam.sigmaPt = 2.0e-4
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor bend inverse_bend monitor
+lattice.nslice = 1
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+#90 degree sbend using drift-kick-drift:
+bend.type = line
+bend.elements = dr kick dr
+bend.repeat = 200
+
+dr.type = drift
+dr.ds = 0.003926990816987
+
+kick.type = thin_dipole
+kick.theta = 0.45
+kick.rc = 1.0
+
+#inverse of a 90 degree sbend using exact nonlinear map:
+inverse_bend.type = sbend_exact
+inverse_bend.ds = -1.570796326794897
+inverse_bend.phi = -90.0
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
+
+
+###############################################################################
+# Diagnostics
+###############################################################################
+diag.slice_step_diagnostics = true

examples/thin_dipole/run_thin_dipole.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import amrex.space3d as amr
+from impactx import ImpactX, RefPart, distribution, elements
+
+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 initial beam
+kin_energy_MeV = 5.0  # 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(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
+
+#   particle bunch
+distr = distribution.Waterbag(
+    sigmaX=1.0e-3,
+    sigmaY=1.0e-3,
+    sigmaT=0.3,
+    sigmaPx=2.0e-4,
+    sigmaPy=2.0e-4,
+    sigmaPt=2.0e-4,
+    muxpx=-0.0,
+    muypy=0.0,
+    mutpt=0.0,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice)
+ns = 1  # number of slices per ds in the element
+segment = [
+    elements.Drift(ds=0.003926990816987, nslice=ns),
+    elements.ThinDipole(theta=0.45, rc=1.0),
+    elements.Drift(ds=0.003926990816987, nslice=ns),
+]
+bend = 200 * segment
+
+inverse_bend = elements.ExactSbend(ds=-1.570796326794897, phi=-90.0)
+
+sim.lattice.append(monitor)
+sim.lattice.extend(bend)
+sim.lattice.append(inverse_bend)
+sim.lattice.append(monitor)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amr.finalize()

src/initialization/InitElement.cpp

@@ -231,6 +231,11 @@ namespace detail
             pp_element.get("bz", bz);
             pp_element.queryAdd("nslice", nslice);
             m_lattice.emplace_back( ChrAcc(ds, ez, bz, nslice) );
+        } else if (element_type == "thin_dipole") {
+            amrex::Real theta, rc;
+            pp_element.get("theta", theta);
+            pp_element.get("rc", rc);
+            m_lattice.emplace_back( ThinDipole(theta, rc) );
         } else if (element_type == "kicker") {
             amrex::Real xkick, ykick;
             std::string units_str = "dimensionless";

src/particles/elements/All.H

@@ -33,6 +33,7 @@
 #include "PRot.H"
 #include "SoftSol.H"
 #include "SoftQuad.H"
+#include "ThinDipole.H"
 #include "diagnostics/openPMD.H"
 
 #include <variant>
@@ -64,7 +65,8 @@ namespace impactx
         ShortRF,
         SoftSolenoid,
         SoftQuadrupole,
-        Sol
+        Sol,
+        ThinDipole
     >;
 
 } // namespace impactx