Animations#

This is an example of making animations by updating the plot. First, create a simple line, a tracker and a particle set, and track it for a couple turns:

Hide code cell content 
%load_ext autoreload
%autoreload 2

import xtrack as xt
import xpart as xp
import xplt
import numpy as np
from matplotlib.animation import FuncAnimation
from IPython.display import display, HTML

xplt.apply_style()

Node = xt.Node

np.random.seed(43543557)
Hide code cell content 
## Generate a simple 6-fold symmetric FODO lattice

n = 6  # number of sections
elements = {
    "QF": xt.Multipole(length=0.3, knl=[0, +0.63]),
    "QD": xt.Multipole(length=0.3, knl=[0, -0.48]),
    "MU": xt.Multipole(length=0.5, knl=[np.pi / n], hxl=[np.pi / n]),
}
parts = {
    "a": [
        xt.Node(0.7, "QF"),
        xt.Node(1.4, "MU"),
        xt.Node(2.1, "QD"),
        xt.Node(2.8, "MU"),
    ],
    "b": [
        xt.Node(2.2, "MU"),
        xt.Node(2.9, "QD"),
        xt.Node(3.6, "MU"),
        xt.Node(4.3, "QF"),
    ],
}
nodes = [xt.Node(5.0 * i, "a" if i % 2 else "b", name=f"S{i+1}") for i in range(n)]

# sextupoles
for i in range(n):
    sx = xt.Multipole(length=0.2, knl=[0, 0, 0.5 * np.sin(2 * np.pi * (i / n))])
    nodes.append(xt.Node(0.2, sx, from_=f"S{i+1}", name=f"S{i+1}SX"))

line = xt.Line.from_sequence(
    nodes, length=5.0 * n, sequences=parts, elements=elements, auto_reorder=True
)
line.particle_ref = xp.Particles()
line.build_tracker();
Hide code cell content 
## Generate particles
nparticles = int(1e3)

# Transverse distribution (gaussian)
norm_emitt_x = 4e-6  # normalized 1-sigma emittance in m*rad (=beta*gamma*emitt_x)
norm_emitt_y = 1e-6  # normalized 1-sigma emittance in m*rad (=beta*gamma*emitt_y)
x, px = xp.generate_2D_gaussian(nparticles)
y, py = xp.generate_2D_gaussian(nparticles)

# Longitudinal distribution (coasting beam)
rel_momentum_spread = 1e-4  # relative momentum spread ( P/p0 - 1 )
zeta = line.get_length() * np.random.uniform(-0.5, 0.5, nparticles)
delta = rel_momentum_spread * xp.generate_2D_gaussian(nparticles)[0]

particles = line.build_particles(
    x_norm=x,
    px_norm=px,
    nemitt_x=norm_emitt_x,
    y_norm=y,
    py_norm=py,
    nemitt_y=norm_emitt_y,
    method="4d",  # for twiss (default is 6d, won't work without a cavity)
    zeta=zeta,
    delta=delta,
)
Hide code cell content 
## Generate tracking data
line.track(particles, num_turns=50, turn_by_turn_monitor=True)

Adjust some parameters to improve rendering performance of the animation:

# use a faster backend
xplt.mpl.rcParams["backend"] = "nbagg"
# increase buffer size
xplt.mpl.rcParams["animation.embed_limit"] = 100  # MB
# use a faster style
xplt.mpl.style.use("fast")

Phase space animation#

Use a FuncAnimation together with the update method to create an animation:

plot = xplt.PhaseSpacePlot(animated=True, std=True, mean=True)
particle_data = line.record_last_track


def animate(i):
    turn = particle_data.at_turn[0, i]
    plot.fig.suptitle(f"Turn {turn}")
    return plot.update(
        particle_data,
        mask=(slice(None), i),  # select all particles and a single turn
        autoscale=(i == 0),  # only on first frame, then freeze to avoid jitter
    )


anim = FuncAnimation(
    plot.fig, animate, frames=range(0, 50, 3), interval=100, blit=True
)

display(HTML(anim.to_jshtml()))
# anim.save('anim.gif', dpi=150)

# especially for larger animations, it is good practice to clean up:
plot.fig.clear()
# xplt.plt.close()
del plot
import gc

gc.collect();
On this page