Particles

Introduction

In Xsuite, collections of particles for tracking simulations are generated using the xpart module. Such collections are stored as instances of the xpart.Particles class. All quantities stored by the Particles objects are described in the Particles class documentation.

The following sections illustrate:

  • How to create Particles objects on CPU or GPU, providing the coordinates in the form of arrays or using the xpart generators to obtain specific distributions (e.g. Gaussian, halo, pencil);

  • How to copy Particles objects (optionally across contexts, e.g GPU to CPU);

  • How to transform Particle objects into dictionaries or pandas dataframes and back;

  • How to merge Particles objects;

  • How to filter Particles objects to select a subset of particles satisfying a logical condition defined by the user.

Building particles with the Particles class

If all the coordinates of the particles are known, a Particles object can be created directly with the xpart.Particles class. For example:

import xpart as xp
import xobjects as xo

ctx = xo.ContextCpu()

particles = xp.Particles(_context=ctx,
        mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, # 7 TeV
        x=[1e-3, 0], px=[1e-6, -1e-6], y=[0, 1e-3], py=[2e-6, 0],
        zeta=[1e-2, 2e-2], delta=[0, 1e-4])

# Complete source: xpart/examples/particles_generation/000_basics.py

The build_particles function

It is often convenient to generate new Particles objects starting from a given reference particle, which defines the particle type (charge and mass) and the reference energy and momentum. This can be accomplished using the xpart.build_particles() function or its alias Line.build_particles, which feature three different modes illustrated in the following.

The set mode

By default, or if mode="set" is passed to the function, only reference quantities including mass0, q0, p0c, gamma0, etc. are taken from the provided reference particle. Particles coordinates, instead, are set according to the provided input x, px, y, py, zeta, delta (with zero assumed as default). For example:

import xpart as xp
import xobjects as xo

# Build a reference particle
p0 = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3)

# Choose a context
ctx = xo.ContextCpu()

# Built a set of three particles with different y coordinates
particles = xp.build_particles(_context=ctx, particle_ref=p0, y=[1,2,3])

# Inspect
print(particles.p0c[1]) # gives 7e12
print(particles.x[1]) # gives 0.0
print(particles.y[1]) # gives 2.0

# Complete source: xpart/examples/particles_generation/001a_build_particles_set.py

Equivalently one can use the Line.build_particles function (automatically infers context and reference particle from the line):

import json

import xpart as xp
import xobjects as xo
import xtrack as xt

ctx = xo.ContextCpu() # choose a context

# Load machine model and built a line
filename = ('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
with open(filename, 'r') as fid:
    line = xt.Line.from_dict(json.load(fid)['line'])
line.build_tracker(_context=ctx)

# Attach a reference particle to the line
line.particle_ref = xp.Particles(p0c=7e12, mass0=xp.PROTON_MASS_EV, q0=1, x =1 , y=3)

# Built a set of three particles with different y coordinates
# (context and particle_ref are taken from the line)
particles = line.build_particles(y=[1,2,3])

# Complete source: xpart/examples/particles_generation/001at_build_particles_set_with_tracker.py

The shift mode

If mode="shift" is passed to the function, reference quantities including quantities including mass0, q0, p0c, gamma0, etc. are taken from the provided reference particle, and the other coordinates are set from the reference particle and shifted according to the provided input x, px, y, py, zeta, delta (with zero assumed as default). For example:

import xpart as xp
import xobjects as xo

# Build a reference particle
p0 = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3)

# Choose a context
ctx = xo.ContextCpu()

# Built a set of three particles with different y coordinates
particles = xp.build_particles(mode='shift', particle_ref=p0, y=[1,2,3],
                               _context=ctx)

# Inspect
print(particles.p0c[1]) # gives 7e12
print(particles.x[1]) # gives 1.0
print(particles.y[1]) # gives 5.0

# Complete source: xpart/examples/particles_generation/001b_build_particles_shift.py

Equivalently one can use the line.build_particles` function (automatically infers context and reference particle from the line):

import json

import xpart as xp
import xobjects as xo
import xtrack as xt

ctx = xo.ContextCpu() # choose a context

# Load machine model and built a line
filename = ('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
with open(filename, 'r') as fid:
    line = xt.Line.from_dict(json.load(fid)['line'])
line.build_tracker(_context=ctx)

# Attach a reference particle to the line
line.particle_ref = xp.Particles(p0c=7e12, mass0=xp.PROTON_MASS_EV, q0=1, x=1 , y=3)

# Built a set of three particles with different y coordinates
# (context and particle_ref are taken from the line)
particles = line.build_particles(mode='shift', y=[1,2,3])

# Inspect
print(particles.p0c[1]) # gives 7e12
print(particles.x[1]) # gives 1.0
print(particles.y[1]) # gives 5.0

# Complete source: xpart/examples/particles_generation/001bt_build_particles_shift_with_tracker.py

The normalized_transverse mode

If mode="normalized_transverse" is passed to the function or if any of the input x_norm, px_norm, y_norm, py_norm is provided, the transverse coordinates are computed from normalized values x_norm, px_norm, y_norm, py_norm (with zero assumed as default) using the closed-orbit information and the linear transfer map obtained from the line argument or provided by the user. Reference quantities including mass0, q0, p0c, gamma0, etc. are taken from the provided reference particle. The longitudinal coordinates are set according to the provided input zeta, delta (zero is assumed as default). For example:

import json

import xobjects as xo
import xpart as xp
import xtrack as xt

# Choose a context
ctx = xo.ContextCpu()

# Load machine model (from pymask)
filename = ('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
with open(filename, 'r') as fid:
    dct = json.load(fid)
line = xt.Line.from_dict(dct['line'])
line.build_tracker(_context=ctx)

# Attach a reference particle to the line
line.particle_ref = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3)

# Built a set of three particles with different x coordinates
particles = line.build_particles(
                               zeta=0, delta=1e-3,
                               x_norm=[1,0,-1], # in sigmas
                               px_norm=[0,1,0], # in sigmas
                               nemitt_x=3e-6, nemitt_y=3e-6)

# Complete source: xpart/examples/particles_generation/001c_build_particles_normalized.py

Generating particles distributions

For several applications it is convenient to generate the transverse coordinates in the normalized phase space and then transform them to physical coordinates. Xpart provides functions to generate independently particles distributions in the three dimensions, which are then combined using the xpart.build_particles() function. This is illustrated by the following examples.

Example: Pencil beam

The following example shows how to generate a distribution often used for collimation studies, which combines:

  • A Gaussian distribution in (x, px);

  • A pencil distribution in (y, py);

  • A Gaussian distribution matched to the non-linear bucket in (zeta, delta).

import json
import numpy as np

import xpart as xp
import xtrack as xt

num_particles = 10000
nemitt_x = 2.5e-6
nemitt_y = 3e-6

# Load machine model (from pymask)
filename = ('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
with open(filename, 'r') as fid:
    input_data = json.load(fid)
line=xt.Line.from_dict(input_data['line'])
line.particle_ref = xp.Particles.from_dict(input_data['particle'])
line.build_tracker()

# Horizontal plane: generate gaussian distribution in normalized coordinates
x_in_sigmas, px_in_sigmas = xp.generate_2D_gaussian(num_particles)

# Vertical plane: generate pencil distribution in normalized coordinates
pencil_cut_sigmas = 6.
pencil_dr_sigmas = 0.7
y_in_sigmas, py_in_sigmas, r_points, theta_points = xp.generate_2D_pencil(
                             num_particles=num_particles,
                             pos_cut_sigmas=pencil_cut_sigmas,
                             dr_sigmas=pencil_dr_sigmas,
                             side='+-')

# Longitudinal plane: generate gaussian distribution matched to bucket 
zeta, delta = xp.generate_longitudinal_coordinates(
        num_particles=num_particles, distribution='gaussian',
        sigma_z=10e-2, line=line)

# Build particles:
#    - scale with given emittances
#    - transform to physical coordinates (using 1-turn matrix)
#    - handle dispersion
#    - center around the closed orbit
particles = line.build_particles(
            zeta=zeta, delta=delta,
            x_norm=x_in_sigmas, px_norm=px_in_sigmas,
            y_norm=y_in_sigmas, py_norm=py_in_sigmas,
            nemitt_x=nemitt_x, nemitt_y=nemitt_y)

# Absolute coordinates can be inspected in particle.x, particles.px, etc.

# Tracking can be done with:
# tracker.track(particles, num_turns=10)

# Complete source: xpart/examples/particles_generation/003_pencil.py
_images/pencil.png

Particle distribution in normalized coordinates (left) and physical coordinates (right). See the full code generating the image.

Example: Halo beam

The following example shows how to generate a distribution, which combines:

  • A halo distribution with an azimuthal cut in (x, px);

  • All particles on the closed orbit in (y, py);

  • All particles in the same point in (zeta, delta);

import json
import numpy as np

import xpart as xp
import xtrack as xt

num_particles = 10000
nemitt_x = 2.5e-6
nemitt_y = 3e-6

# Load machine model (from pymask)
filename = ('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
with open(filename, 'r') as fid:
    input_data = json.load(fid)
line = xt.Line.from_dict(input_data['line'])
line.particle_ref = xp.Particles.from_dict(input_data['particle'])
line.build_tracker()

# Horizontal plane: generate cut halo distribution
(x_in_sigmas, px_in_sigmas, r_points, theta_points
    )= xp.generate_2D_uniform_circular_sector(
                                          num_particles=num_particles,
                                          r_range=(0.6, 0.9), # sigmas
                                          theta_range=(0.25*np.pi, 1.75*np.pi))

# Vertical plane: all particles on the closed orbit
y_in_sigmas = 0.
py_in_sigmas = 0.

# Longitudinal plane: all particles off momentum by 1e-3
zeta = 0.
delta = 1e-3

# Build particles:
#    - scale with given emittances
#    - transform to physical coordinates (using 1-turn matrix)
#    - handle dispersion
#    - center around the closed orbit
particles = line.build_particles(
            zeta=zeta, delta=delta,
            x_norm=x_in_sigmas, px_norm=px_in_sigmas,
            y_norm=y_in_sigmas, py_norm=py_in_sigmas,
            nemitt_x=nemitt_x, nemitt_y=nemitt_y)

# Absolute coordinates can be inspected in particle.x, particles.px, etc.

# Tracking can be done with:
# line.track(particles, num_turns=10)

# Complete source: xpart/examples/particles_generation/002_halo.py
_images/halo.png

Example: Gaussian bunch

The function xpart.generate_matched_gaussian_bunch() can be used to generate a bunch having Gaussian distribution in all coordinates and matched to the non-linead RF bucket, as illustrated by the following example:

import json

import numpy as np
from scipy.constants import e as qe
from scipy.constants import m_p

import xpart as xp
import xtrack as xt

bunch_intensity = 1e11
sigma_z = 22.5e-2
n_part = int(5e5)
nemitt_x = 2e-6
nemitt_y = 2.5e-6

filename = ('../../../xtrack/test_data/sps_w_spacecharge'
            '/line_no_spacecharge_and_particle.json')
with open(filename, 'r') as fid:
    ddd = json.load(fid)
line = xt.Line.from_dict(ddd['line'])
line.particle_ref = xp.Particles.from_dict(ddd['particle'])

line.build_tracker()

particles = xp.generate_matched_gaussian_bunch(
         num_particles=n_part, total_intensity_particles=bunch_intensity,
         nemitt_x=nemitt_x, nemitt_y=nemitt_y, sigma_z=sigma_z,
         line=line)

# Complete source: xpart/examples/particles_generation/004_generate_gaussian.py
_images/gaussian.png

Matching distribution at custom location in the ring

The functions xtrack.Line.generate_matched_gaussian_bunch() can be used to match a particle distribution at a custom location in the ring, as illustrated by the following example:

import json

import xpart as xp
import xtrack as xt

# Load machine model and build tracker
filename = ('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
with open(filename, 'r') as fid:
    input_data = json.load(fid)
line = xt.Line.from_dict(input_data['line'])
line.particle_ref = xp.Particles.from_dict(input_data['particle'])
line.build_tracker()

# Match distribution at a given element
particles = tracker.build_particles(x_norm=[0,1,2], px_norm=[0,0,0], # in sigmas
                   nemitt_x=2.5e-6, nemitt_y=2.5e-6,
                   at_element='ip2')

# Match distribution at a given s position (100m downstream of ip6)
particles = tracker.build_particles(x_norm=[0,1,2], px_norm=[0,0,0], # in sigmas
                   nemitt_x=2.5e-6, nemitt_y=2.5e-6,
                   at_element='ip6',
                   match_at_s=tracker.line.get_s_position('ip6') + 100
                   )

# Complete source: xpart/examples/particles_generation/006_match_at_element.py

Copying a Particles object (optionally across contexts)

The copy method allows making copies of a Particles object within the same context or in another context. It can be used for example to transfer Particles objects to/from GPU, as shown by the following example:

import xpart as xp
import xobjects as xo

p1 = xp.Particles(x=[1,2,3])

# Make a copy of p1 in the same context
p2 = p1.copy()

# Alter p1
p1.x += 10

# Inspect
print(p1.x) # gives [11. 12. 13.]
print(p2.x) # gives [1. 2. 3.]

# Copy across contexts
ctxgpu = xo.ContextCupy()
p3 = p1.copy(_context=ctxgpu)

# Inspect
print(p3.x[2]) # gives 13

# Complete source: xpart/examples/merge_copy_filter/001_copy.py

Saving and loading Particles objects to/from dictionary or file

The methods to_dict/from_dict and to_pandas/from_pandas allow transforming a Particles object into a dictionary or a pandas dataframe and back. By default the particles coordinates are transferred to CPU when using to_dict or to_pandas.

Such methods can be used to save or load particles coordinated to/from file as shown by the following examples:

Save and load from dictionary

import numpy as np

import xobjects as xo
import xpart as xp

# Create a Particles on your selected context (default is CPU)
context = xo.ContextCupy()
part = xp.Particles(_context=context, x=[1,2,3])

# Save particles to dict 
dct = part.to_dict()

# Load particles from dict 
part_from_dict = xp.Particles.from_dict(dct, _context=context)

# Complete source: xpart/examples/save_and_load/000_to_from_dict.py

Save and load from json file

import numpy as np

import xobjects as xo
import xpart as xp

# Create a Particles on your selected context (default is CPU)
context = xo.ContextCupy()
part = xp.Particles(_context=context, x=[1,2,3])

# Save particles to json
import json
with open('part.json', 'w') as fid:
    json.dump(part.to_dict(), fid, cls=xo.JEncoder)

# Load particles from json file to selected context
with open('part.json', 'r') as fid:
    part_from_json= xp.Particles.from_dict(json.load(fid), _context=context)

# Complete source: xpart/examples/save_and_load/001_save_load_json.py

Save and load from pickle file

import numpy as np

import xobjects as xo
import xpart as xp

# Create a Particles on your selected context (default is CPU)
context = xo.ContextCupy()
part = xp.Particles(_context=context, x=[1,2,3])

##########
# PICKLE #
##########

# Save particles to pickle file
import pickle
with open('part.pkl', 'wb') as fid:
    pickle.dump(part.to_dict(), fid)

# Load particles from json to selected context
with open('part.pkl', 'rb') as fid:
    part_from_pkl= xp.Particles.from_dict(pickle.load(fid), _context=context)

# Complete source: xpart/examples/save_and_load/002_save_load_pickle.py

Save and load using pandas

import numpy as np

import xobjects as xo
import xpart as xp

# Create a Particles on your selected context (default is CPU)
context = xo.ContextCupy()
part = xp.Particles(_context=context, x=[1,2,3])

##############
# PANDAS/HDF #
##############

# Save particles to hdf file via pandas
import pandas as pd
df = part.to_pandas()
df.to_hdf('part.hdf', key='df', mode='w')

# Read particles from hdf file via pandas
part_from_pdhdf = xp.Particles.from_pandas(pd.read_hdf('part.hdf'))

# Complete source: xpart/examples/save_and_load/003_save_load_with_pandas.py

Merging and filtering Particles objects

Merging Particles objects

The merge method can be used to merge Particles objects as shown by the following example:

import xpart as xp

p1 = xp.Particles(x=[1,2,3])
p2 = xp.Particles(x=[4, 5])
p3 = xp.Particles(x=6)

particles = xp.Particles.merge([p1,p2,p3])

print(particles.x) # gives [1. 2. 3. 4. 5. 6.]

# Complete source: xpart/examples/merge_copy_filter/000_merge.py

Filtering a Particles object

The filter method can be used to select a subset of particles satisfying a logical condition defined by the user.

import xpart as xp

p1 = xp.Particles(x=[1,2,3], px=[10, 20, 30])

mask = p1.x > 1

p2 = p1.filter(mask)

print(p2.x) # gives [2. 3.]
print(p2.px) # gives [20. 30.]

# Complete source: xpart/examples/merge_copy_filter/002_filter.py

Accessing particles coordinates on GPU contexts

When working on a GPU context, the coordinate attributes of particle objects are not numpy arrays as on the CPU contexts, but specific array types associated with the specific context (e.g. cupy arrays for contexts of type ContextCupy). Although such arrays can be directly inspected to a large extent, several actions, notably plotting with matplotlib and saving to pickle or json files, are not possible without explicitly transferring the data to the CPU memory.

For this purpose we recommend to use the specific functions provided by the context in order to keep the code usable on different contexts. For example:

import xobjects as xo
import xtrack as xt

context = xo.ContextCupy()

particles = xt.Particles(_context=context, x=[1, 2, 3])

# Avoid the following (which does not work if a CPU context is chosen):
# x_cpu = particles.x.get()

# Instead use the following (which is guaranteed to work on all contexts):
x_cpu = context.nparray_from_context_array(particles.x)