Add MAD-X equivalent short RF model (#412)

* Examples for 3D space charge benchmarking - Modified the initial beam size in the IOTA lens benchmark example. - Added 2 benchmarks of 3D space charge for initial testing. - Add documentation for 2 benchmarks with space charge. - Add a benchmark example with space charge and periodic s-dependent focusing. - Added an s-dependent example using a Kurth beam without space charge. - Modified tolerance for IOTA lens benchmark example. Reduced tolerance to account for smaller initial beam size and improved preservation of invariants of motion. - Modified tolerances of space charge examples to allow CI tests to pass when space charge is not active. - Modified tolerance for space charge examples. These should fail unless space charge is turned on. * Update input_kurth_10nC.in Selected numerical values for amr.n_cell, lattice.nslice, and geometry.prob_relative. * Rename ShortRF and add MAD-X short RF model. * Delete input_kurth_10nC.in This file is not part of this PR. * Add reference particle push. * Add example. * Add documentation. * Compression README: Minor Formatting * Compression Analysis: Add Gamma Ref Template * Modified the example to include a small acceleration. * Update analysis script. * Update analysis_compression.py Relax tolerance. * Cleanup
ECP-WarpX · Aug 9, 2023 · d706a25 · d706a25
1 parent eae94da
commit d706a25
Show file tree

Hide file tree

Showing 15 changed files with 533 additions and 22 deletions.
diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst
@@ -26,6 +26,7 @@ This section allows you to **download input files** that correspond to different
    examples/positron_channel/README.rst
    examples/cyclotron/README.rst
    examples/cfbend/README.rst
+   examples/compression/README.rst
 
 
 Unit tests

diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst
@@ -409,7 +409,7 @@ Lattice Elements
 
             * ``<element_name>.nslice`` (``integer``) number of slices used for the application of space charge (default: ``1``)
 
-        * ``shortrf`` for a short RF (bunching) cavity element.
+        * ``buncher`` for a short RF cavity (linear) bunching element.
           This requires these additional parameters:
 
             * ``<element_name>.V`` (``float``, dimensionless) normalized voltage drop across the cavity
@@ -420,6 +420,23 @@ Lattice Elements
 
                 = 2*pi/(RF wavelength in m)
 
+        * ``shortrf`` for a short RF cavity element.
+          This requires these additional parameters:
+
+            * ``<element_name>.V`` (``float``, dimensionless) normalized voltage drop across the cavity
+
+                = (maximum energy gain in MeV) / (particle rest energy in MeV)
+
+            * ``<element_name>.freq`` (``float``, in Hz) the RF frequency
+
+            * ``<element_name>.phase`` (``float``, in degrees) the synchronous RF phase
+
+                phase = 0: maximum energy gain (on-crest)
+
+                phase = -90 deg:  zero energy gain for bunching
+
+                phase = 90 deg:  zero energy gain for debunching
+
         * ``uniform_acc_chromatic`` for a region of uniform acceleration, with chromatic effects included.
            The Hamiltonian is expanded through second order in the transverse variables (x,px,y,py), with the exact pt dependence retained.
            This requires these additional parameters:

diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -602,13 +602,21 @@ References:
    :param B: Magnetic field in Tesla; when B = 0 (default), the reference bending radius is defined by r0 = length / (angle in rad),   corresponding to a magnetic field of B = rigidity / r0; otherwise the reference bending radius is defined by r0 = rigidity / B.
    :param nslice: number of slices used for the application of space charge
 
-.. py:class:: impactx.elements.ShortRF(V, k)
+.. py:class:: impactx.elements.Buncher(V, k)
 
-   A short RF cavity element at zero crossing for bunching.
+   A short RF cavity element at zero crossing for bunching (MaryLie model).
 
    :param V: Normalized RF voltage drop V = Emax*L/(c*Brho)
    :param k: Wavenumber of RF in 1/m
 
+.. py:class:: impactx.elements.ShortRF(V, freq, phase)
+
+   A short RF cavity element (MAD-X model).
+
+   :param V: Normalized RF voltage V = maximum energy gain/(m*c^2)
+   :param freq: RF frequency in Hz
+   :param phase: RF synchronous phase in degrees (phase = 0 corresponds to maximum energy gain, phase = -90 corresponds go zero energy gain for bunching)
+
 .. py:class:: impactx.elements.ChrUniformAcc(ds, k, nslice=1)
 
    A region of constant Ez and Bz for uniform acceleration, with chromatic effects included.

diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -583,3 +583,21 @@ add_impactx_test(cfbend.py
     examples/cfbend/analysis_cfbend.py
     OFF  # no plot script yet
 )
+
+# Ballistic compression ########################################################
+#
+# w/o space charge
+add_impactx_test(compression
+    examples/compression/input_compression.in
+      ON  # ImpactX MPI-parallel
+      OFF  # ImpactX Python interface
+    examples/compression/analysis_compression.py
+    OFF  # no plot script yet
+)
+add_impactx_test(compression.py
+    examples/compression/run_compression.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/compression/analysis_compression.py
+    OFF  # no plot script yet
+)
diff --git a/examples/compression/README.rst b/examples/compression/README.rst
@@ -0,0 +1,46 @@
+.. _examples-compression:
+
+Ballistic Compression Using a Short RF Element
+==============================================
+
+A 20 MeV electron beam propagates through a short RF element near zero-crossing, inducing a head-tail energy correlation.
+This is followed by ballistic motion in a drift, which is used to compress the rms bunch length from 16 ps to 10 ps (compression of 5/3).
+
+The beam is not exactly on-crest (phase = -89.5 deg), so there is an energy gain of 4.5 MeV.
+
+The transverse emittance is sufficiently small that the horizontal and verticle beam size are essentially unchanged.  Due to RF curvature, there is some growth of the longitudinal emittance.
+
+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.
+
+
+Run
+---
+
+This example can be run as a Python script (``python3 run_compression.py``) or with an app with an input file (``impactx input_compression.in``).
+Each can also be prefixed with an `MPI executor <https://www.mpi-forum.org>`__, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+.. tab-set::
+
+   .. tab-item:: Python Script
+
+       .. literalinclude:: run_compression.py
+          :language: python3
+          :caption: You can copy this file from ``examples/compression/run_compression.py``.
+
+   .. tab-item:: App Input File
+
+       .. literalinclude:: input_compression.in
+          :language: ini
+          :caption: You can copy this file from ``examples/compression/input_compression.in``.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. dropdown:: Script ``analysis_compression.py``
+
+   .. literalinclude:: analysis_compression.py
+      :language: python3
+      :caption: You can copy this file from ``examples/compression/analysis_compression.py``.
diff --git a/examples/compression/analysis_compression.py b/examples/compression/analysis_compression.py
@@ -0,0 +1,118 @@
+#!/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)
+
+
+# openPMD data series at the beam monitors
+series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
+
+# first and last step
+final_step = list(series.iterations)[-1]
+first_it = series.iterations[1]
+final_it = series.iterations[final_step]
+
+# initial beam & reference particle gamma
+initial = first_it.particles["beam"].to_df()
+initial_gamma_ref = first_it.particles["beam"].get_attribute("gamma_ref")
+
+# final beam & reference particle gamma
+final = final_it.particles["beam"].to_df()
+final_gamma_ref = final_it.particles["beam"].get_attribute("gamma_ref")
+
+# 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}"
+)
+print(f"  gamma={initial_gamma_ref: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, initial_gamma_ref],
+    [
+        5.0e-04,
+        5.0e-04,
+        5.0e-03,
+        4.952764e-09,
+        5.028325e-09,
+        1.997821e-08,
+        40.1389432485322889,
+    ],
+    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}"
+)
+print(f"  gamma={final_gamma_ref: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, final_gamma_ref],
+    [
+        5.004995e-04,
+        5.005865e-04,
+        3.033949e-03,
+        4.067876e-09,
+        4.129937e-09,
+        6.432081e-05,
+        48.8654787469061860,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
diff --git a/examples/compression/input_compression.in b/examples/compression/input_compression.in
@@ -0,0 +1,42 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 10000
+beam.units = static
+beam.energy = 20.0
+beam.charge = 1.0e-9
+beam.particle = electron
+beam.distribution = waterbag
+beam.sigmaX = 0.5e-3
+beam.sigmaY = 0.5e-3
+beam.sigmaT = 5.0e-3
+beam.sigmaPx = 1.0e-5
+beam.sigmaPy = 1.0e-5
+beam.sigmaPt = 4.0e-6
+beam.muxpx = 0.0
+beam.muypy = 0.0
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor shortrf1 drift1 monitor
+
+monitor.type = beam_monitor
+monitor.backend = h5
+
+shortrf1.type = shortrf
+shortrf1.V = 1000.0
+shortrf1.freq = 1.3e9
+shortrf1.phase = -89.5
+
+drift1.type = drift
+drift1.ds = 1.7
+
+
+###############################################################################
+# Algorithms
+###############################################################################
+algo.particle_shape = 2
+algo.space_charge = false
diff --git a/examples/compression/run_compression.py b/examples/compression/run_compression.py
@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Marco Garten, Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import amrex.space3d as amr
+from impactx import ImpactX, 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 a 250 MeV proton beam with an initial
+# unnormalized rms emittance of 1 mm-mrad in all
+# three phase planes
+energy_MeV = 20.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(0.510998950).set_energy_MeV(energy_MeV)
+
+#   particle bunch
+distr = distribution.Waterbag(
+    sigmaX=0.5e-3,
+    sigmaY=0.5e-3,
+    sigmaT=5.0e-3,
+    sigmaPx=1.0e-5,
+    sigmaPy=1.0e-5,
+    sigmaPt=4.0e-6,
+)
+sim.add_particles(bunch_charge_C, distr, npart)
+
+# add beam diagnostics
+monitor = elements.BeamMonitor("monitor", backend="h5")
+
+# design the accelerator lattice
+sim.lattice.append(monitor)
+#   Short RF cavity element
+shortrf1 = elements.ShortRF(V=1000.0, freq=1.3e9, phase=-89.5)
+#   Drift element
+drift1 = elements.Drift(ds=1.7)
+
+sim.lattice.extend([shortrf1, drift1])
+
+sim.lattice.append(monitor)
+
+# run simulation
+sim.evolve()
+
+# clean shutdown
+del sim
+amr.finalize()
diff --git a/examples/fodo_rf/input_fodo_rf.in b/examples/fodo_rf/input_fodo_rf.in
@@ -34,7 +34,7 @@ quad1.k = 2.5
 drift1.type = drift
 drift1.ds = 1.0
 
-shortrf1.type = shortrf
+shortrf1.type = buncher
 shortrf1.V = 0.01
 shortrf1.k = 15.0
 

diff --git a/examples/fodo_rf/run_fodo_rf.py b/examples/fodo_rf/run_fodo_rf.py
@@ -53,7 +53,7 @@
 #   Drift element
 drift1 = elements.Drift(ds=1.0)
 #   Short RF cavity element
-shortrf1 = elements.ShortRF(V=0.01, k=15.0)
+shortrf1 = elements.Buncher(V=0.01, k=15.0)
 
 lattice_no_drifts = [quad1, shortrf1, quad2, shortrf1, quad1]
 #   set first lattice element

diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp
@@ -113,11 +113,18 @@ namespace detail
             pp_element.get("kt", kt);
             pp_element.queryAdd("nslice", nslice);
             m_lattice.emplace_back( ConstF(ds, kx, ky, kt, nslice) );
-        } else if (element_type == "shortrf") {
+        } else if (element_type == "buncher") {
             amrex::Real V, k;
             pp_element.get("V", V);
             pp_element.get("k", k);
-            m_lattice.emplace_back( ShortRF(V, k) );
+            m_lattice.emplace_back( Buncher(V, k) );
+        } else if (element_type == "shortrf") {
+            amrex::Real V, freq;
+            amrex::Real phase = -90.0;
+            pp_element.get("V", V);
+            pp_element.get("freq", freq);
+            pp_element.queryAdd("phase", phase);
+            m_lattice.emplace_back( ShortRF(V, freq, phase) );
         } else if (element_type == "multipole") {
             int m;
             amrex::Real k_normal, k_skew;