PS ion sequence

examples/spacecharge/000_spacecharge_example_ps.py

@@ -0,0 +1,127 @@
+# copyright ############################### #
+# This file is part of the Xtrack Package.  #
+# Copyright (c) CERN, 2021.                 #
+# ######################################### #
+
+import json
+import numpy as np
+
+import xobjects as xo
+import xpart as xp
+import xtrack as xt
+import xfields as xf
+
+############
+# Settings #
+############
+
+# Load PS line from test data
+fname_line = ('../../test_data/ps_ion/PS_2022_Pb_nominal.json')
+
+# Nominal Pb ion parameters
+bunch_intensity = 6.0e8
+sigma_z = 6.0
+nemitt_x=0.75e-6
+nemitt_y=0.5e-6
+n_part=int(1e6)
+num_turns=32
+
+num_spacecharge_interactions = 540
+
+# Available modes: frozen/quasi-frozen/pic
+# mode = 'pic'
+#mode = 'quasi-frozen'
+mode = 'frozen'
+
+# Choose solver between `FFTSolver2p5DAveraged` and `FFTSolver2p5D`
+pic_solver = 'FFTSolver2p5DAveraged'
+
+####################
+# Choose a context #
+####################
+
+context = xo.ContextCpu()
+#context = xo.ContextCupy()
+#context = xo.ContextPyopencl('0.0')
+
+print(context)
+
+#############
+# Load line #
+#############
+
+line = xt.Line.from_json(fname_line)
+line.build_tracker(_context=context, compile=False)
+
+#############################################
+# Install spacecharge interactions (frozen) #
+#############################################
+
+lprofile = xf.LongitudinalProfileQGaussian(
+        number_of_particles=bunch_intensity,
+        sigma_z=sigma_z,
+        z0=0.,
+        q_parameter=1.)
+
+xf.install_spacecharge_frozen(line=line,
+                   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(
+        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),
+        solver=pic_solver)
+else:
+    raise ValueError(f'Invalid mode: {mode}')
+
+#################
+# Build Tracker #
+#################
+
+line.build_tracker(_context=context)
+
+######################
+# Generate particles #
+######################
+
+# Build a line without spacecharge (recycling the track kernel)
+line_sc_off = line.filter_elements(exclude_types_starting_with='SpaceCh')
+line_sc_off.build_tracker(track_kernel=line.tracker.track_kernel)
+
+# (we choose to match the distribution without accounting for spacecharge)
+particles = xp.generate_matched_gaussian_bunch(line=line_sc_off,
+         num_particles=n_part, total_intensity_particles=bunch_intensity,
+         nemitt_x=nemitt_x, nemitt_y=nemitt_y, sigma_z=sigma_z)
+
+####################################################
+# Check that everything is on the selected context #
+####################################################
+
+assert particles._context == context
+assert len(set([ee._buffer for ee in line.elements])) == 1 # All on same context
+assert line._context == context
+
+#########
+# Track #
+#########
+line.track(particles, num_turns=num_turns, with_progress=1)
+

test_data/ps_ion/000_make_ps_line.py

@@ -0,0 +1,66 @@
+# copyright ############################### #
+# This file is part of the Xtrack Package.  #
+# Copyright (c) CERN, 2021.                 #
+# ######################################### #
+
+import json
+import numpy as np
+
+import xtrack as xt
+import xpart as xp
+import xobjects as xo
+from cpymad.madx import Madx
+
+# Load matched MADX sequence with correct tunes and chromaticities
+mad = Madx()
+mad.call('PS_2022_Pb_ions_matched.seq')
+mad.use('ps')
+mad.twiss()
+
+twmad_4d = mad.table.twiss.dframe()
+summad_4d = mad.table.summ.dframe()
+
+# RF values as taken from LSA for Pb ions of PS
+# In the ps_ss.seq, the 10 MHz cavities are PR_ACC10 - there are 12 of them in the straight sections
+harmonic_nb = 16
+V_RF = 38.0958e3  # V (control room definition, energy gain per charge). 
+nn = 'pa.c10.11'  # for now use the first of the 10 MHz RF cavities 
+
+# We switch on one cavity 
+charge = mad.sequence.ps.beam.charge
+mad.sequence.ps.elements[nn].lag = 0 # 0 as we are below transition
+mad.sequence.ps.elements[nn].volt = V_RF/1e6*charge # different convention between madx and xsuite
+mad.sequence.ps.elements[nn].freq = mad.sequence['ps'].beam.freq0*harmonic_nb
+
+twmad_6d = mad.table.twiss.dframe()
+summad_6d = mad.table.summ.dframe()
+
+mad.emit()
+qs_mad = mad.table.emitsumm.qs[0]
+
+# Make xsuite line and tracker
+line = xt.Line.from_madx_sequence(mad.sequence.ps)
+line.particle_ref = xp.Particles(mass0=mad.sequence.ps.beam.mass*1e9,
+                                 q0=mad.sequence.ps.beam.charge,
+                                 gamma0=mad.sequence.ps.beam.gamma)
+tracker = line.build_tracker()
+
+# Activate the RF cavity in Xsuite
+line[nn].lag = 0  
+line[nn].voltage =  V_RF # In Xsuite for ions, do not multiply by charge as in MADX
+line[nn].frequency = mad.sequence['ps'].beam.freq0*1e6*harmonic_nb
+
+tw = tracker.twiss()
+
+# Relativistic beta factor needed when comparing MAD-X and Xsuite
+beta0 = mad.sequence['ps'].beam.beta 
+
+assert np.isclose(tw.qs, qs_mad, atol=1e-6)
+assert np.isclose(tw.qx, summad_4d.q1, atol=1e-5)
+assert np.isclose(tw.qy, summad_4d.q2, atol=1e-5)
+assert np.isclose(tw.dqx, summad_6d.dq1*beta0, atol=0.5)
+assert np.isclose(tw.dqy, summad_6d.dq2*beta0, atol=0.5)
+
+with open('ps_line_and_particle.json', 'w') as f:
+    json.dump(line.to_dict(), f, cls=xo.JEncoder, indent=4)
+