Add envelope tracking module

[austin-hoover]

This PR adds an orbit.envelope module for beam envelope tracking. The routines in this module are not space charge solvers; rather, they are methods to integrate ODEs describing the beam envelope dynamics for specified phase space distributions.

Background

We define an upright KV distribution function as

$$ f(J_x, J_y) \propto \delta \left( 1 - \frac{J_x}{\tilde{J}_x} - \frac{J_x}{\tilde{J}_y} \right), $$

where $J_x(x, x')$ and $J_y(y, y')$ are the Courant-Snyder (CS) invariants in a linear, uncoupled focusing lattice, $\tilde{J}_x$ and $\tilde{J}_y$ are maximum values of $J_x$ and $J_y$, and $\delta$ is the Dirac delta function. The distribution projects to a uniform density within an ellipse in the $x-y$ plane, with radii $c_x = 2 \sqrt{\langle{xx}\rangle}$ and $c_y = 2 \sqrt{\langle{yy}\rangle}$. The equations of motion for a particle in the beam are:

$$ \begin{align} x'' + k_x x &= \left[ \frac{2Q}{c_x (c_x + c_y)} \right] x, \\ y'' + k_y y &= \left[ \frac{2Q}{c_y (c_x + c_y)} \right] y, \end{align} $$

where $k_x$ is the linear external focusing strength in the $x$ plane, $k_y$ is the linear external focusing strength in the $y$ plane, and $Q$ is the beam perveance. The envelope radii evolve according to a set of coupled ODEs [Lund 2004]:

$$ \begin{align} c_x'' + k_x c_x &= -\frac{\varepsilon_x^2}{c_x^3} + \frac{2 Q}{c_x (c_x + c_y)}, \\ c_y'' + k_y c_y &= -\frac{\varepsilon_y^2}{c_y^3} + \frac{2 Q}{c_y (c_x + c_y)}, \end{align} $$

where $\varepsilon_x$ and $\varepsilon_y$ are the invariant (4-rms) emittances:

$$ \begin{align} \varepsilon_x &= 4 \sqrt{ \langle{xx}\rangle \langle{x'x'}\rangle - \langle{xx'}\rangle \langle{xx'}\rangle }, \\ \varepsilon_y &= 4 \sqrt{ \langle{yy}\rangle \langle{y'y'}\rangle - \langle{yy'}\rangle \langle{yy'}\rangle }. \end{align} $$

The Danilov distribution is a limiting case of the generalized KV distribution. We define the KV distribution as

$$ f(J_1, J_2) \propto \delta \left( 1 - \frac{J_1}{\tilde{J}_1} - \frac{J_2}{\tilde{J}_2} \right), $$

where $J_1(x, x', y, y')$ and $J_2(x, x', y, y')$ are generalized CS invariants in a linear focusing lattice. If we take the limit $\tilde{J}_2 \rightarrow 0$ (single mode), we find

$$ f(J_1, J_2) \propto \delta(J_2) \Theta(\tilde{J}_1 - J_1). $$

We define this limiting case as the Danilov distribution. The beam projects to a uniform density within a tilted ellipse. Equations for the beam envelope are found in [Danilov 2003] [Hoover 2021].

Implementation

The KV envelope equations can be integrated numerically. To work with PyORBIT, one needs to extract the linear focusing strength at each point in the lattice and then call an ODE solver; for example using scipy. We can instead use PyORBIT to perform the integration by tracking a Bunch with the first particle $x$ and $y$ coordinates representing the envelope radii $c_x$ and $c_y$. The nonlinear kicks are implemented by special envelope tracking nodes added to the lattice. The same can be done for the Danilov distribution. This method assumes the single-particle tracking through AccLattice object is linear; this is not true, even for drifts or quads, but is often approximately true, especially for particles with small energy offset and if fringe fields are ignored.

This approach also makes it easy to implement particle-core models, where we track test particles subject to the field of the beam envelope [Holmes 1998]. The forces acting on particles inside and outside the the beam are derived in [Parzen 2021]. We treat every additional particle in the beam (beyond the first one or two particles containing the envelope parameters) as a test particle.

Changes

1. Added orbit.core.envelope.KVEnvelopeTracker. Does the envelope tracking at the C++ level.
2. Added orbit.core.envelope.DanilovEnvelopeTracker. Does the envelope tracking at the C++ level.
3. Added orbit.envelope.KVEnvelope. Represents a KV distribution. Converts to and from Bunch objects (where the first particle represents the beam envelope parameters).
4. Added orbit.envelope.DanilovEnvelope. Represents a Danilov distribution. Converts to and from Bunch objects (where the first two particles represent the beam envelope parameters).
5. Added orbit.envelope.KVEnvelopeTracker. Adds envelope tracking nodes to the lattice and tracks KVEnvelope objects, optionally with test particles. Also has ability to find matched (periodic) envelope solution.
6. Added orbit.envelope.DanilovEnvelopeTracker. Adds envelope tracking nodes to the lattice and tracks DanilovEnvelope objects, optionally with test particles. Also has ability to find matched (periodic) envelope solution.
7. Added benchmarks and examples to /examples/Envelope directory.

Future extensions

1. Energy spread/dispersion not yet included [Holmes 1998].
2. Boundary conditions (beam pipe) not yet included.
3. Equations can be derived for generalized KV distribution, with any partition of $\tilde{J}_1$ and $\tilde{J}_2$ in arbitrary linear focusing lattices [Chung 2016].
4. A more general envelope solver could be implemented which tracks the covariance matrix $\Sigma$ as $M \Sigma M^T$, where $M$ is a transfer matrix (Envelope/covariance matrix tracking #51). The idea would be to have something like a CovMatrix object which is tracked, instead of using the technique above, where we track a Bunch.
5. There is an (unused?) module orbit.matching which uses the external ODE solver approach for KV distribution. We could integrate into orbit.envelope.

References

• [Lund 2004] https://doi.org/10.1103/PhysRevSTAB.7.024801
• [Danilov 2003] https://doi.org/10.1103/PhysRevSTAB.6.094202
• [Hoover 2021] https://doi.org/10.1103/PhysRevAccelBeams.24.044201
• [Holmes 1998] https://doi.org/10.1103/PhysRevSTAB.2.114202
• [Parzen 2001] https://arxiv.org/pdf/physics/0108040
• [Chung 2016] https://doi.org/10.1103/PhysRevLett.117.224801

[austin-hoover]

Examples

Benchmark: RMS evolution

https://github.com/austin-hoover/PyORBIT3/blob/envelope/examples/Envelope/bench_kv.py
https://github.com/austin-hoover/PyORBIT3/blob/envelope/examples/Envelope/bench_danilov.py

This example compares the envelope tracker to PIC (SpaceChargeCalc2p5D) predictions in a FODO lattice. In the following plots, black lines = envelope tracker and red dots = PIC.

KV:

Danilov:

Benchmark: particle-core evolution

https://github.com/austin-hoover/PyORBIT3/blob/envelope/examples/Envelope/bench_kv_particle.py
https://github.com/austin-hoover/PyORBIT3/blob/envelope/examples/Envelope/bench_danilov_particle.py

This example tracks test particles in the field of the beam. To do this in the PIC solver, we set the test-particle macrosizes to 0 so that they do not contribute to the electric field calculation. The following plots show the period-by-period $x-x'$ coordinates for a few different particles (each color is a different particle). Left = envelope, right = PIC. The grey ellipse shows the beam envelope.

KV:

Danilov:

Matching

https://github.com/austin-hoover/PyORBIT3/blob/envelope/examples/Envelope/match_kv.py
https://github.com/austin-hoover/PyORBIT3/blob/envelope/examples/Envelope/match_danilov.py

This example demonstrates matching of the beam envelope to a periodic lattice. The following plots show the beam envelope before/after matching.

KV:

Danilov:

Plotting

https://github.com/austin-hoover/PyORBIT3/blob/envelope/examples/Envelope/plot_kv.py
https://github.com/austin-hoover/PyORBIT3/blob/envelope/examples/Envelope/plot_danilov.py

We show the 2D pairwise projections of 4D phase space coordinates sampled from the distributions:

KV:

Danilov:

[austin-hoover]

Forgot to mention: the orbit.envelope module uses numpy, and for the matching it uses scipy, so they would need to be added to required packages.