Particle tracking

Particle tracking can be performed by referring to the various methods of the lattice elements or similarly the lattice itself. For example linear optics tracking can be done via the linear instance method:

from importlib import resources
from dipas.build import from_file
from dipas.elements import Kicker
import dipas.test.sequences
import torch
torch.manual_seed(1)

with resources.path(dipas.test.sequences, 'cryring.seq') as path:
    lattice = from_file(path)

lattice = lattice.makethin({Kicker: 3})  # Need to make thin for tracking.

particles = 0.001 * torch.randn(6, 1000)  # 1000 particles
print(particles[[0, 2], :].std(dim=1))

tracked = lattice.linear(particles)
print(tracked[[0, 2], :].std(dim=1))

tensor([0.0010, 0.0010])
tensor([0.0083, 0.0014])

This tracks one turn through the lattice. By default no aperture checks are performed. We can enable aperture checks by setting the parameter aperture=True:

particles = 0.01 * torch.randn(6, 1000)
tracked = lattice.linear(particles, aperture=True)
print(tracked.shape)

torch.Size([6, 43])

So we lost most of the particles in this case. To get an idea of where they were lost, we can instruct the tracking method to record the loss:

tracked, loss = lattice.linear(particles, aperture=True, recloss=True)
print(tracked.shape)

torch.Size([6, 43])

Setting recloss=True records the loss values at each element and adds them as a separate return value in form of a dict, mapping element labels to loss values. The loss values themselves are determined by the particular aperture type (see elements.aperture_types). The loss value is computed for each particle arriving at the entrance of an element. If the loss value is greater than zero the particle lost, otherwise it is tracked further. Let’s see the loss values for the first ten elements:

for label, loss_val in list(loss.items())[:10]:
    print(f'{label}: {len(loss_val)}')

p_0: 1000
drift_0: 1000
p_lp2end: 998
drift_1: 998
yr01lb3: 998
drift_2: 998
yr01lb4: 997
drift_3: 994
yr01df3: 989
drift_4: 989

That means all 1,000 particles arrived at the entrance of element p_0 (which is a marker) and thus also arrives at element drift_0. Note that even though drift_0 is a k1 = 0 quadrupole, serving as an aperture-checked drift in MADX, the tracking here performs aperture checks also for Drift spaces. Since at the next marker only 998 particles arrive, this means we lost two particles at the previous element. We can confirm that by checking the loss values greater than zero:

l_drift_0 = loss['drift_0']
print(l_drift_0[l_drift_0 > 0])

tensor([0.0012, 0.0029])

Instead of returning a loss history we can also ask for an accumulated version of the loss value. This will sum the loss values which are greater than zero at every element:

tracked, loss = lattice.linear(particles, aperture=True, recloss='sum')
print(tracked.shape)
print(loss)

torch.Size([6, 43])
tensor(151.6977)

This is helpful for particle loss optimization because if our lattice contained optimization parameters, we could inject the corresponding gradients via loss.backward().

We can also use more fine-grained control over the loss history by specifying one or more multi-element selectors that will be matched against elements (these multi-element selectors are str, re.Pattern or lattice element types).

from dipas.elements import SBend

tracked, loss = lattice.linear(particles, aperture=True, recloss=SBend)
for k, v in loss.items():
    print(k, len(v))

yr01mh 961
yr02mh 351
yr03mh 107
yr04mh 91
yr05mh 89
yr06mh 86
yr07mh 69
yr08mh 66
yr09mh 63
yr10mh 60
yr11mh 60
yr12mh 60

Again the lengths of the loss values indicate how many particles arrived at a particular element. Using a wildcard expression we can record the loss at all the quadrupoles for example:

tracked, loss = lattice.linear(particles, aperture=True, recloss='yr*qs*')
print(len(loss))
print(set(type(lattice[label]) for label in loss))

18
{<class 'dipas.elements.Quadrupole'>}

The same options are available for observing particle coordinates at specific elements. For that purpose we can use the observe parameter. We can provide similar values as for recloss (except for "sum" which doesn’t make sense here):

tracked, locations = lattice.linear(particles, aperture=True, observe='yr*qs*')
print(len(locations))
print(set(type(lattice[label]) for label in locations))

18
{<class 'dipas.elements.Quadrupole'>}

By inspecting the shape of the corresponding position we can see how many particles were successfully tracked through an element, i.e. made it to the element’s exit. This number is the number of particles that arrived at an element (the len(loss_value)) minus the number of particles that were lost at the element (len(loss_value[loss_value > 0])). The loss is computed at the entrance of an element and the coordinates are recorded at the exit of elements:

tracked, locations, loss = lattice.linear(particles, aperture=True, observe='yr*qs*', recloss='yr*qs*')
print(loss['yr02qs1'].shape[-1])
print(len(loss['yr02qs1'][loss['yr02qs1'] > 0]))
print(locations['yr02qs1'].shape[-1])
print(loss['yr02qs1'].shape[-1] - len(loss['yr02qs1'][loss['yr02qs1'] > 0]) == locations['yr02qs1'].shape[-1])

724
81
643
True

Irrespective of the tracking method used (e.g. linear in the above examples), drift spaces will always be tracked through by using the exact solutions to the equations of motion (referred to by the exact tracking method). If this behavior is undesired and drift spaces should use the specified tracking method instead of exact this can be done by specifying the parameter exact_drift=False.

print(lattice.linear(particles, exact_drift=False).std(dim=1))
print(lattice.linear(particles, exact_drift=True).std(dim=1))

tensor([0.0863, 0.0151, 0.0141, 0.0089, 3.2844, 0.0103])
tensor([0.0861, 0.0151, 0.0141, 0.0089, 3.2826, 0.0103])