Merge pull request #17 from xsuite/feature/avoid_multitracker

Interface with xtrack/xfields for tracking with wakes

.github/workflows/pr_workflow.yaml

@@ -10,15 +10,17 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.11
+      - name: Set up Python 3.10
         uses: actions/setup-python@v2
         with:
-          python-version: 3.11
+          python-version: "3.10"
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install .[tests]
+          pip install wheel
+          pip install numpy cython h5py
+          pip install xtrack xpart xfields nafflib
+          pip install --no-build-isolation .[tests]
       - name: Run tests
         run: |
-          cd tests
-          pytest -v
+          pytest -v tests

.gitlab-ci.yml

@@ -11,6 +11,7 @@ run_tests:
   before_script:
     - pip install .
     - pip install pytest
+    - pip install xtrack xfields xpart
   script:
     - pytest
 
@@ -27,4 +28,4 @@ tagging:
   rules:
     - if: $CI_COMMIT_TAG
       when: never
-    - when: always                     
+    - when: always

contributors.txt

@@ -0,0 +1,12 @@
+The following people contributed to the development of this package:
+
+CERN contributors:
+ - D. Amorim
+ - X. Buffat
+ - L. Giacomel
+ - G. Iadarola
+ - L. Mether
+ - N. Mounet
+ - M. Rognlien
+ - E. Vik
+

examples/000a_sps_tune_shift.py

@@ -0,0 +1,113 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pathlib
+import h5py
+
+import xwakes as xw
+import xtrack as xt
+import xfields as xf
+import xpart as xp
+import nafflib
+
+# a few machine parameters
+p0c = 26e9
+circumference = 6911.5662
+
+h_RF = [4620, 4*4620]
+V_RF = [4.5e6, 0.45e6]
+dphi_RF = [180, 0]
+f_rev = 43347.25733575443
+f_RF = [f_rev*h for h in h_RF]
+
+bucket_length = circumference / h_RF[0]
+
+Q_x = 20.13
+Q_y = 20.18
+xi = 0.05
+
+# initialize the one turn map
+one_turn_map = xt.LineSegmentMap(
+    length=circumference, betx=54.65, bety=54.51,
+    dnqx=np.array([Q_x, xi*Q_x]), dnqy=np.array([Q_y, xi*Q_y]),
+    longitudinal_mode='nonlinear',
+    voltage_rf=V_RF, frequency_rf=f_RF,
+    lag_rf=dphi_RF, momentum_compaction_factor=0.0031
+)
+
+# initialize the wakes
+test_data_folder = pathlib.Path(__file__).parent.joinpath('../test_data').absolute()
+wake_table_name_kickers = test_data_folder.joinpath('sps_wake_2022/MKP_wake.dat')
+wake_file_columns = ['time', 'dipole_x', 'dipole_y', 'quadrupole_x', 'quadrupole_y']
+components = ['dipole_x', 'dipole_y', 'quadrupole_x', 'quadrupole_y']
+wake_df_kickers = xw.read_headtail_file(wake_table_name_kickers,
+                                        wake_file_columns)
+wf_kickers = xw.WakeFromTable(table=wake_df_kickers, columns=components)
+
+wake_table_name_wall = test_data_folder.joinpath('sps_wake_2022/wall_wake.dat')
+wake_df_wall = xw.read_headtail_file(wake_table_name_wall,
+                                     wake_file_columns)
+wf_wall = xw.WakeFromTable(table=wake_df_wall, columns=components)
+
+# sum the wakes and prepare for tracking
+wf_sps = wf_kickers + wf_wall
+wf_sps.configure_for_tracking(
+    zeta_range=(-0.5*bucket_length, 0.5*bucket_length),
+    num_slices=100,
+    bunch_spacing_zeta=5*bucket_length,
+    num_turns=37,
+    circumference=circumference
+)
+
+# initialize a damper with 100 turns gain
+transverse_damper = xf.TransverseDamper(
+    gain_x=2/100, gain_y=2/100,
+    zeta_range=(-0.5*bucket_length, 0.5*bucket_length),
+    num_slices=100,
+    bunch_spacing_zeta=5*bucket_length,
+    circumference=circumference,
+)
+
+# initialize a monitor for the average transverse positions
+monitor = xf.CollectiveMonitor(
+    base_file_name=f'sps_tune_shift',
+    monitor_bunches=True,
+    monitor_slices=False,
+    monitor_particles=False,
+    flush_data_every=100,
+    stats_to_store=['mean_x', 'mean_y'],
+    backend='hdf5',
+    zeta_range=(-0.5*bucket_length, 0.5*bucket_length),
+    num_slices=100,
+    bunch_spacing_zeta=5*bucket_length,
+)
+
+elements = [one_turn_map, wf_sps, transverse_damper, monitor]
+element_names = ['one_turn_map', 'wake', 'transverse_damper', 'monitor']
+line = xt.Line(elements, element_names=element_names)
+
+line.particle_ref = xt.Particles(p0c=26e9)
+line.build_tracker()
+
+# initialize a matched gaussian bunch
+particles = xp.generate_matched_gaussian_bunch(
+    num_particles=100_000,
+    total_intensity_particles=2.3e11,
+    nemitt_x=2e-6, nemitt_y=2e-6,
+    sigma_z=0.21734953205,
+    line=line,
+)
+
+# apply a small kick to the particles
+particles.px += 1e-3
+particles.py += 1e-3
+
+# track the particles
+line.track(particles, num_turns=200, with_progress=1)
+
+# read mean positions from the monitor file
+with h5py.File(monitor.base_file_name + '_bunches.h5', 'r') as h5file:
+    mean_x = h5file['0']['mean_x'][:]
+    mean_y = h5file['0']['mean_y'][:]
+
+print('Q_x:', nafflib.tune(mean_x))
+print('Q_y:', nafflib.tune(mean_y))

examples/000b_sps_tune_shift_multibunch.py

@@ -0,0 +1,139 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pathlib
+import h5py
+
+import xwakes as xw
+import xtrack as xt
+import xfields as xf
+import xpart as xp
+import nafflib
+
+# a few machine parameters
+p0c = 26e9
+circumference = 6911.5662
+
+h_RF = [4620, 4*4620]
+V_RF = [4.5e6, 0.45e6]
+dphi_RF = [180, 0]
+f_rev = 43347.25733575443
+f_RF = [f_rev*h for h in h_RF]
+
+bucket_length = circumference / h_RF[0]
+bunch_spacing_buckets = 5
+
+Q_x = 20.13
+Q_y = 20.18
+xi = 0.05
+
+# prepare the filling scheme
+filling_scheme = np.zeros(int(h_RF[0]/bunch_spacing_buckets))
+n_bunches = 12
+filling_scheme[0:n_bunches] = 1
+
+# initialize the one turn map
+one_turn_map = xt.LineSegmentMap(
+    length=circumference, betx=54.65, bety=54.51,
+    dnqx=np.array([Q_x, xi*Q_x]), dnqy=np.array([Q_y, xi*Q_y]),
+    longitudinal_mode='nonlinear',
+    voltage_rf=V_RF, frequency_rf=f_RF,
+    lag_rf=dphi_RF, momentum_compaction_factor=0.0031
+)
+
+# initialize the wakes
+test_data_folder = pathlib.Path(__file__).parent.joinpath('../test_data').absolute()
+wake_table_name_kickers = test_data_folder.joinpath('sps_wake_2022/MKP_wake.dat')
+wake_file_columns = ['time', 'dipole_x', 'dipole_y', 'quadrupole_x', 'quadrupole_y']
+components = ['dipole_x', 'dipole_y', 'quadrupole_x', 'quadrupole_y']
+wake_df_kickers = xw.read_headtail_file(wake_table_name_kickers,
+                                        wake_file_columns)
+wf_kickers = xw.WakeFromTable(table=wake_df_kickers, columns=components)
+
+wake_table_name_wall = test_data_folder.joinpath('sps_wake_2022/wall_wake.dat')
+wake_df_wall = xw.read_headtail_file(wake_table_name_wall,
+                                     wake_file_columns)
+wf_wall = xw.WakeFromTable(table=wake_df_wall, columns=components)
+
+# sum the wakes and prepare for tracking
+wf_sps = wf_kickers + wf_wall
+wf_sps.configure_for_tracking(
+    zeta_range=(-0.5*bucket_length, 0.5*bucket_length),
+    num_slices=100,
+    bunch_spacing_zeta=5*bucket_length,
+    num_turns=37,
+    circumference=circumference,
+    filling_scheme=filling_scheme
+)
+
+# initialize a damper with 100 turns gain
+transverse_damper = xf.TransverseDamper(
+    gain_x=2/100, gain_y=2/100,
+    zeta_range=(-0.5*bucket_length, 0.5*bucket_length),
+    num_slices=100,
+    bunch_spacing_zeta=5*bucket_length,
+    circumference=circumference,
+    filling_scheme=filling_scheme
+)
+
+# initialize a monitor for the average transverse positions
+monitor = xf.CollectiveMonitor(
+    base_file_name=f'sps_tune_shift',
+    monitor_bunches=True,
+    monitor_slices=False,
+    monitor_particles=False,
+    flush_data_every=100,
+    stats_to_store=['mean_x', 'mean_y'],
+    backend='hdf5',
+    zeta_range=(-0.5*bucket_length, 0.5*bucket_length),
+    num_slices=100,
+    bunch_spacing_zeta=5*bucket_length,
+    filling_scheme=filling_scheme
+)
+
+elements = [one_turn_map, wf_sps, transverse_damper, monitor]
+element_names = ['one_turn_map', 'wake', 'transverse_damper', 'monitor']
+line = xt.Line(elements, element_names=element_names)
+
+line.particle_ref = xt.Particles(p0c=26e9)
+line.build_tracker()
+
+# initialize a matched gaussian bunch
+particles = xp.generate_matched_gaussian_multibunch_beam(
+    bunch_num_particles=100_000,
+    bunch_intensity_particles=2.3e11,
+    nemitt_x=2e-6, nemitt_y=2e-6,
+    sigma_z=0.21734953205,
+    filling_scheme=filling_scheme,
+    bucket_length=bucket_length,
+    bunch_spacing_buckets=bunch_spacing_buckets,
+    circumference=circumference,
+    line=line,
+)
+
+# apply a kick to the particles
+particles.px += 1e-3
+particles.py += 1e-3
+
+# track the particles
+line.track(particles, num_turns=101, with_progress=1)
+
+qx_bunch = []
+qy_bunch = []
+
+# read mean positions from the monitor file
+with h5py.File(monitor.base_file_name + '_bunches.h5', 'r') as h5file:
+    for bunch in h5file.keys():
+        mean_x = h5file[bunch]['mean_x'][:]
+        mean_y = h5file[bunch]['mean_y'][:]
+
+        qx_bunch.append(nafflib.tune(mean_x))
+        qy_bunch.append(nafflib.tune(mean_y))
+
+import matplotlib.pyplot as plt
+plt.close('all')
+plt.plot(qx_bunch, 'rx', label='Q_x')
+plt.plot(qy_bunch, 'bo', label='Q_y')
+plt.xlabel('Bunch number')
+plt.ylabel('Tune')
+plt.legend()
+plt.show()