Working with Particles objects¶
Table of Contents
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])
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 tracker.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
Equivalently one can use the tracker.build_particles
function (automatically
infers context and reference particle from the tracker):
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 tracker
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'])
tracker = line.build_tracker(_context=ctx)
# Attach a reference particle to the tracker
tracker.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 tracker)
particles = tracker.build_particles(y=[1,2,3])
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 x 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
Equivalently one can use the tracker.build_particles
function (automatically
infers context and reference particle from the tracker):
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 tracker
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'])
tracker = line.build_tracker(_context=ctx)
# Attach a reference particle to the tracker
tracker.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 tracker)
particles = tracker.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
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 tracker
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'])
tracker = line.build_tracker(_context=ctx)
# Attach a reference particle to the tracker
tracker.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 = tracker.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)
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)
tracker = xt.Tracker(line=xt.Line.from_dict(input_data['line']))
tracker.particle_ref = xp.Particles.from_dict(input_data['particle'])
# 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, tracker=tracker)
# Build particles:
# - scale with given emittances
# - transform to physical coordinates (using 1-turn matrix)
# - handle dispersion
# - center around the closed orbit
particles = tracker.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)
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)
tracker = xt.Tracker(line=xt.Line.from_dict(input_data['line']))
tracker.particle_ref = xp.Particles.from_dict(input_data['particle'])
# 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 = tracker.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)
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'])
tracker = 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,
tracker=tracker)
Matching distribution at custom location in the ring¶
The functions xtrack.Tracker.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'])
tracker = 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
)
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
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)
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)
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)
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'))
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.]
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.]
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 xpart as xp
context = xo.ContextCupy()
particles = xp.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)