Skip to content

Implement envelope tracker#112

Draft
austin-hoover wants to merge 136 commits into
PyORBIT-Collaboration:mainfrom
austin-hoover:env-tracker
Draft

Implement envelope tracker#112
austin-hoover wants to merge 136 commits into
PyORBIT-Collaboration:mainfrom
austin-hoover:env-tracker

Conversation

@austin-hoover

@austin-hoover austin-hoover commented Apr 24, 2026

Copy link
Copy Markdown
Contributor

This PR implements an envelope/centroid tracker (Issue #51). We track the covariance matrix $\mathbf{\Sigma} = \langle \mathbf{x} \mathbf{x}^T \rangle$ and centroid $\mathbf{\mu} = \langle \mathbf{x} \rangle$, where $\mathbf{x} = [x, x', y, y', z, \Delta E ]^T$ is the 6D phase space vector.

Evolution equations:

$$ \begin{align} \mathbf{x} &\rightarrow \mathbf{M} \mathbf{x} + \mathbf{u} , \\ \mathbf{\mu} &\rightarrow \mathbf{M} \mathbf{\mu} + \mathbf{u} , \\ \mathbf{\Sigma} &\rightarrow \mathbf{M} \mathbf{\Sigma} \mathbf{M}^T . \end{align} $$

We track the $7 \times 7$ matrix $\mathbf{S} = \langle \mathbf{y} \mathbf{y}^T \rangle$, where $\mathbf{y} = [x, x', y, y', z, dE, 1]^T$:

$$ \mathbf{S} = \langle \mathbf{y} \mathbf{y}^T \rangle = \begin{bmatrix} \langle \mathbf{x} \mathbf{x}^T \rangle & \langle \mathbf{x} \rangle \\ \langle \mathbf{x} \rangle^T & 1 \end{bmatrix} = \begin{bmatrix} \mathbf{\Sigma} + \mathbf{\mu}\mathbf{\mu}^T & \mathbf{\mu} \\ \mathbf{\mu}^T & 1 \end{bmatrix} . $$

Evolution equations for $\mathbf{y}$:

$$ \begin{align} \mathbf{y} &\rightarrow \mathbf{N} \mathbf{y} , \\ \mathbf{S} &\rightarrow \mathbf{N} \mathbf{S} \mathbf{N}^T , \end{align} $$

where $\mathbf{N}$ is defined as

$$ \mathbf{N} = \begin{bmatrix} \mathbf{M} & \mathbf{u} \\ \mathbf{0} & 1 \end{bmatrix} . $$

Accelerator nodes and child nodes are mapped to transfer matrices at the Python level. An error is raised if the node is not recognized. Space charge is handled by assuming a uniform charge density within an ellipsoid in the $x$-$y$-$z$ plane (for 3D space charge) or within an ellipse in the $x$-$y$ plane (for 2D space charge).

@austin-hoover

Copy link
Copy Markdown
Contributor Author

Thanks! that's very helpful. If matrix multiplication is dominating then there's not all that much to do. I added some checks to skip the matrix multiplication tilt angle is zero in TiltTEAPOT. This speeds up because all teapot nodes have TiltTEAPOT child nodes at entrance and exit. Building the matrices also takes time, but I don't see a way around this to handle time-dependent nodes in rings.

@austin-hoover

Copy link
Copy Markdown
Contributor Author

I've added a function to generate the matrix for RF gap nodes, copied from orbit.core.linac.MatrixRFGap. This function updates the energy of the synchronous particle.

Each envelope tracking step needs to track the synchronous particle. For most elements this will just be increasing the time by sync_part.time(length / (sync_part.beta() * speed_of_light). RF gaps will update the synchonous particle energy. Not sure of the best way to handle this.

@woodtp

woodtp commented Jun 24, 2026

Copy link
Copy Markdown
Contributor

One other thing you could try that may or may not be worth the effort is using Numba to do some JIT compilation where possible. You'll eat an initial cost to do the compilation the first time you run the tracker, but the byte code can be cached so that subsequent runs are faster (or have the potential to be). You would probably need to factor out types made accessible from bindings, e.g., the synchronous particle instance, and instead directly pass in the raw values you need to the kernels you want to be JIT compilable.

@austin-hoover

Copy link
Copy Markdown
Contributor Author

WIth these small modifications:

Envelope:

Time per turn: 0.011063957214355468
         675520 function calls (675419 primitive calls) in 0.277 seconds

   Ordered by: internal time
   List reduced from 439 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       25    0.054    0.002    0.254    0.010 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:250(track)
    92025    0.036    0.000    0.036    0.000 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:124(apply_transfer_matrix)
    16000    0.018    0.000    0.024    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/lib/_twodim_base_impl.py:176(eye)
    72600    0.015    0.000    0.015    0.000 /Users/46h/repo/PyORBIT3/py/orbit/lattice/AccNode.py:200(getChildNodes)
    11600    0.012    0.000    0.054    0.000 /Users/46h/repo/PyORBIT3/py/orbit/matrix_lattice/analytic.py:51(drift_matrix)
     2600    0.012    0.000    0.022    0.000 /Users/46h/repo/PyORBIT3/py/orbit/matrix_lattice/analytic.py:60(quad_matrix)
     4000    0.011    0.000    0.042    0.000 /Users/46h/repo/PyORBIT3/py/orbit/teapot/teapot.py:924(matrix)
        5    0.010    0.002    0.010    0.002 {built-in method _imp.create_dynamic}
    30400    0.010    0.000    0.014    0.000 /Users/46h/repo/PyORBIT3/py/orbit/teapot/teapot.py:1548(matrix)
    16000    0.010    0.000    0.042    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/_core/numeric.py:2231(identity)

Bunch (10,000 particles):

Time per turn: 0.04294008255004883
         1197659 function calls (1086109 primitive calls) in 1.073 seconds

   Ordered by: internal time
   List reduced from 147 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    15000    0.227    0.000    0.227    0.000 {built-in method orbit.core.teapot_base.drift}
126775/15225    0.202    0.000    1.066    0.000 /Users/46h/repo/PyORBIT3/py/orbit/lattice/AccNode.py:307(trackActions)
    10150    0.099    0.000    0.099    0.000 {built-in method orbit.core.teapot_base.wrapbunch}
     1125    0.080    0.000    0.080    0.000 {method 'analyzeBunch' of 'BunchTwissAnalysis' objects}
     3450    0.053    0.000    0.053    0.000 {built-in method orbit.core.teapot_base.multp}
   386175    0.052    0.000    0.849    0.000 /Users/46h/repo/PyORBIT3/py/orbit/lattice/AccActionsContainer.py:49(performActions)
     2600    0.041    0.000    0.041    0.000 {built-in method orbit.core.teapot_base.quad2}
     2600    0.037    0.000    0.037    0.000 {built-in method orbit.core.teapot_base.quad1}
     1800    0.033    0.000    0.033    0.000 {built-in method orbit.core.teapot_base.bend4}
   132625    0.032    0.000    0.797    0.000 /Users/46h/repo/PyORBIT3/py/orbit/teapot/teapot.py:160(track)

@austin-hoover

Copy link
Copy Markdown
Contributor Author

With 2D space charge:

Envelope:

Time per turn: 0.06159739494323731
         3340984 function calls (3340883 primitive calls) in 1.540 seconds

   Ordered by: internal time
   List reduced from 482 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    21075    0.133    0.000    1.154    0.000 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:137(sc_transfer_matrix_2d)
    21075    0.102    0.000    0.211    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/linalg/_linalg.py:1536(eigh)
    79225    0.090    0.000    0.121    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/lib/_twodim_base_impl.py:176(eye)
    21075    0.088    0.000    0.174    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/linalg/_linalg.py:557(inv)
       25    0.086    0.003    1.517    0.061 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:250(track)
    42150    0.083    0.000    0.114    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/_core/numeric.py:910(outer)
   113100    0.078    0.000    0.078    0.000 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:124(apply_transfer_matrix)
    42150    0.070    0.000    0.188    0.000 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:115(cov)
    21075    0.039    0.000    0.043    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/linalg/_linalg.py:3009(_multi_dot_three)
    21075    0.038    0.000    0.067    0.000 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:26(build_diag_matrix_from_xyz_eig)

Bunch (100,000 particles, 64 x 64 x 1 grid):

Time per turn: 2.021827611923218
         1325513 function calls (1200163 primitive calls) in 50.548 seconds

   Ordered by: internal time
   List reduced from 167 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    13800   42.838    0.003   42.838    0.003 {method 'trackBunch' of 'SpaceChargeCalc2p5D' objects}
    15000    2.407    0.000    2.407    0.000 {built-in method orbit.core.teapot_base.drift}
    10150    1.007    0.000    1.007    0.000 {built-in method orbit.core.teapot_base.wrapbunch}
     1125    0.861    0.001    0.861    0.001 {method 'analyzeBunch' of 'BunchTwissAnalysis' objects}
     3450    0.550    0.000    0.550    0.000 {built-in method orbit.core.teapot_base.multp}
     2600    0.433    0.000    0.433    0.000 {built-in method orbit.core.teapot_base.quad2}
     2600    0.387    0.000    0.387    0.000 {built-in method orbit.core.teapot_base.quad1}
     1800    0.340    0.000    0.340    0.000 {built-in method orbit.core.teapot_base.bend4}
     1800    0.320    0.000    0.320    0.000 {built-in method orbit.core.teapot_base.bend3}
     1800    0.298    0.000    0.298    0.000 {built-in method orbit.core.teapot_base.bend2}

@austin-hoover

austin-hoover commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

Some improvement with 2D space charge:

Time per turn: 0.033572874069213866
         1343797 function calls (1343696 primitive calls) in 0.839 seconds

   Ordered by: internal time
   List reduced from 443 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    22650    0.255    0.000    0.416    0.000 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:122(sc_transfer_matrix_2d)
   116250    0.131    0.000    0.131    0.000 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:108(apply_transfer_matrix)
    85525    0.086    0.000    0.117    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/lib/_twodim_base_impl.py:176(eye)
       25    0.079    0.003    0.813    0.033 /Users/46h/repo/PyORBIT3/py/orbit/envelope/envelope.py:236(track)
    62875    0.037    0.000    0.156    0.000 /Users/46h/miniforge3/envs/pyorbit/lib/python3.13/site-packages/numpy/_core/numeric.py:2231(identity)
62890/62883    0.021    0.000    0.044    0.000 <frozen importlib._bootstrap>:1390(_handle_fromlist)
    85525    0.020    0.000    0.020    0.000 {built-in method numpy.zeros}
    75750    0.018    0.000    0.018    0.000 /Users/46h/repo/PyORBIT3/py/orbit/lattice/AccNode.py:200(getChildNodes)
    13175    0.015    0.000    0.062    0.000 /Users/46h/repo/PyORBIT3/py/orbit/matrix_lattice/analytic.py:51(drift_matrix)
     2600    0.012    0.000    0.022    0.000 /Users/46h/repo/PyORBIT3/py/orbit/matrix_lattice/analytic.py:60(quad_matrix)

No time saved because to get cov you need to call outer(mu, mu) and
subtract from moment matrix. This is more straightforward
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants