Combined CPU - GPU simulations

When performing simulations on GPU contexts, it is possible to execute on CPU the computation for specific beam elements, for example PyHEADTAIL elements that do not support the Xsuite GPU contexts. For that purpose, one needs to set a `needs_cpu` flag on the concerned elements before inserting them in an Xtrack line. This instructs the Xtrack tracker to handle the required data transfers.

Presently PyHEADTAIL elements are not able to skip lost particles in the Xtrack Particles objects. This can be handled by setting the flag `needs_hidden_lost_particles`, which instructs the Xtrack tracker to temporarily hide the lost particles for the concerned elements.

The following example shows a simulation including the SPS lattice elements, space-charge elements and wakefields. The lattice tracking and the spacecharge calculations are performed on GPU with Xtrack/Xfields while the wakefield computation is performed on CPU with PyHEADTAIL.

import numpy as np

from cpymad.madx import Madx

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

###############################
# Enable pyheadtail interface #
###############################

xt.enable_pyheadtail_interface()

############
# Settings #
############

gpu_device = 0
seq_name = "sps"
qx0,qy0 = 20.13, 20.18
p0c = 26e9
cavity_name = "actcse.31632"
cavity_lag = 180
frequency = 200e6
rf_voltage = 4e6
use_wakes = True
n_slices_wakes = 100
limit_z = 0.7
bunch_intensity = 1e11
sigma_z = 22.5e-2
nemitt_x=2.5e-6
nemitt_y=2.5e-6
n_part=int(1e4)
num_turns=2
num_spacecharge_interactions = 540

# mode = 'frozen'
mode = 'pic'

#######################################################
# Use cpymad to load the sequence and match the tunes #
#######################################################

mad = Madx()
mad.call("./madx_sps/sps.seq")
mad.call("./madx_sps/lhc_q20.str")
mad.call("./madx_sps/macro.madx")
mad.beam()
mad.use(seq_name)

mad.twiss()
tw_thick = mad.table.twiss.dframe()
summ_thick = mad.table.summ.dframe()

mad.input("""select, flag=makethin, slice=1, thick=false;
            makethin, sequence=sps, style=teapot, makedipedge=false;""")
mad.use(seq_name)
mad.exec(f"sps_match_tunes({qx0},{qy0});")

##################
# Switch context #
##################

if gpu_device is None:
   context = xo.ContextCpu()
else:
   context = xo.ContextCupy(device=gpu_device)

#################################
# Build line from MAD-X lattice #
#################################

line = xt.Line.from_madx_sequence(sequence=mad.sequence[seq_name],
           deferred_expressions=True, install_apertures=True,
           apply_madx_errors=False)
# Define reference particle
line.particle_ref = xt.Particles(p0c=p0c,mass0=xt.PROTON_MASS_EV)

################
# Switch on RF #
################

line[cavity_name].voltage = rf_voltage
line[cavity_name].lag = cavity_lag
line[cavity_name].frequency = frequency

############################################
# Twiss (check that the lattice is stable) #
############################################

line.build_tracker()
tw = line.twiss()

# We unfreeze the line so that we can still insert elements in the
# original one (would be prevented by the existence of a tracker linked to the
# line).
line.unfreeze()

#####################################
# Install spacecharge interactions) #
#####################################

# Install frozen space-charge lenses

lprofile = xf.LongitudinalProfileQGaussian(
       number_of_particles=bunch_intensity,
       sigma_z=sigma_z,
       z0=0.,
       q_parameter=1.)

xf.install_spacecharge_frozen(line=line,
                  particle_ref=line.particle_ref,
                  longitudinal_profile=lprofile,
                  nemitt_x=nemitt_x, nemitt_y=nemitt_y,
                  sigma_z=sigma_z,
                  num_spacecharge_interactions=num_spacecharge_interactions,
                  )

# Switch to PIC or quasi-frozen
if mode == 'frozen':
   pass # Already configured in line
elif mode == 'quasi-frozen':
   xf.replace_spacecharge_with_quasi_frozen(
                                   line,
                                   update_mean_x_on_track=True,
                                   update_mean_y_on_track=True)
elif mode == 'pic':
   pic_collection, all_pics = xf.replace_spacecharge_with_PIC(
       _context=context, line=line,
       n_sigmas_range_pic_x=8,
       n_sigmas_range_pic_y=8,
       nx_grid=256, ny_grid=256, nz_grid=100,
       n_lims_x=7, n_lims_y=3,
       z_range=(-3*sigma_z, 3*sigma_z))
else:
   raise ValueError(f'Invalid mode: {mode}')

######################
# Install wakefields #
######################

if use_wakes:

   # We rescale the waketable to take into account for the difference between
   # the average beta functions (for which the table is calculated) and the beta
   # functions at the location in the lattice at which the wake element is
   # installed.
   wakefields = np.genfromtxt('wakes/kickerSPSwake_2020_oldMKP.wake')
   wakefields[:,1] *= 54.65/tw['betx'][0]
   wakefields[:,2] *= 54.51/tw['bety'][0]
   wakefields[:,3] *= 54.65/tw['betx'][0]
   wakefields[:,4] *= 54.65/tw['betx'][0]
   wakefile = 'wakes/wakefields.wake'
   np.savetxt(wakefile,wakefields,delimiter='\t')

   # Build PyHEADTAIL wakefield element
   from PyHEADTAIL.particles.slicing import UniformBinSlicer
   from PyHEADTAIL.impedances.wakes import WakeTable, WakeField
   n_slices_wakes = 500
   slicer_for_wakefields = UniformBinSlicer(n_slices_wakes, z_cuts=(-limit_z, limit_z))
   waketable = WakeTable(wakefile, ['time', 'dipole_x', 'dipole_y',
                  'quadrupole_x', 'quadrupole_y'])
   wakefield = WakeField(slicer_for_wakefields, waketable)

   # Specity that the wakefield element needs to run on CPU and that lost
   # particles need to be hidden for this element (required by PyHEADTAIL)
   wakefield.needs_cpu = True
   wakefield.needs_hidden_lost_particles = True

   # Insert element in the line
   line.insert_element(element=wakefield,name="wakefield", at_s=0)

#################
# Build Tracker #
#################

line.build_tracker(_context=context)

######################
# Generate particles #
######################

# (we choose to match the distribution without accounting for spacecharge)
line_sc_off = line.filter_elements(exclude_types_starting_with='SpaceCh')

particles = xp.generate_matched_gaussian_bunch(_context=context,
        num_particles=n_part, total_intensity_particles=bunch_intensity,
        nemitt_x=nemitt_x, nemitt_y=nemitt_y, sigma_z=sigma_z,
        particle_ref=line.particle_ref, line=line_sc_off)
particles.circumference = line.get_length() # Needed by pyheadtail

###########################################################
# We use a phase monitor to measure the tune turn by turn #
###########################################################

phasem = xp.PhaseMonitor(line,
                 num_particles=n_part, twiss=line_sc_off.twiss())

#########
# Track #
#########

for turn in range(num_turns):
   phasem.measure(particles)
   #import pdb; pdb.set_trace()
   line.track(particles)

##################
# Plot footprint #
##################

import matplotlib.pyplot as plt
plt.close('all')
f,ax = plt.subplots()
ax.plot(phasem.qx, phasem.qy,'b.',ms=1)
ax.set_xlim(0, 0.5)
ax.set_ylim(0, 0.5)
plt.show()

# Complete source: xtrack/examples/pyheadtail_interface/004_imped_spacech_cpu_gpu.py