Merge pull request #154 from xsuite/release/v0.6.4

Release 0.6.4

Author: freddieknets

examples/lhc_run3_lossmap.py

@@ -94,5 +94,5 @@
 
 print(f"Total calculation time {time.time()-start_time}s")
 
-ThisLM.plot()
+ThisLM.plot(savefig=Path(path_out, f'lossmap_B{beam}{plane}.pdf'))
 plt.show()

examples/lhc_run3_lossmap_with_blowup.py

@@ -115,7 +115,7 @@
 
 print(f"Total calculation time {time.time()-start_time}s")
 
-ThisLM.plot()
+ThisLM.plot(savefig=Path(path_out, f'lossmap_B{beam}{plane}.pdf'))
 plt.show()
 
 

examples/lhc_run3_lossmap_with_crystals.py

@@ -141,5 +141,5 @@
 
 print(f"Total calculation time {time.time()-start_time}s")
 
-ThisLM.plot()
+ThisLM.plot(savefig=Path(path_out, f'lossmap_B{beam}{plane}.pdf'))
 plt.show()

examples/lhc_run3_offmomentum_lossmap.py

@@ -140,5 +140,5 @@
 
 print(f"Total calculation time {time.time()-start_time}s")
 
-ThisLM.plot()
+ThisLM.plot(savefig=Path(path_out, f'lossmap_B{beam}{plane}.pdf'))
 plt.show()

examples/scatter.py

@@ -0,0 +1,148 @@
+import time
+import numpy as np
+import pandas as pd
+from pathlib import Path
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+
+import xobjects as xo
+import xpart as xp
+import xtrack as xt
+import xcoll as xc
+
+
+num_part = 100000
+_capacity = 2*num_part
+particle_ref = xp.Particles.reference_from_pdg_id(pdg_id='proton', p0c=7e12)
+
+everest_coll = xc.EverestCollimator(length=1.0, material=xc.materials.MolybdenumGraphite, jaw=[0.01, -0.01])
+
+fluka_coll = xc.FlukaCollimator(length=1.0, jaw=[0.01, -0.01])
+xc.FlukaEngine.start(elements=fluka_coll, names='tcp.c6l7.b1', debug_level=1, _capacity=_capacity, cwd='temp_scratch')
+xc.FlukaEngine.set_particle_ref(particle_ref=particle_ref)
+
+# g4_coll = xc.Geant4Collimator(length=1.0, material='mogr', jaw=[0.1, -0.1], geant4_id='g4_coll')
+
+
+# Plotting function
+# -----------------
+
+dri = xt.Drift(length=everest_coll.length)
+def plot_scatters(part_init, part, outfile, show_colorbar=True):
+    devmin = 1e-6
+    devmax = 1
+    part_dri = part_init.copy()
+    dri.track(part_dri)
+    part_dri.sort(interleave_lost_particles=True)
+    part.sort(interleave_lost_particles=True)
+    mask = part.state > 0
+    dev = 1.e3*np.sqrt((part.x[mask] - part_dri.x[part.parent_particle_id[mask]])**2
+                     + (part.y[mask] - part_dri.y[part.parent_particle_id[mask]])**2)
+    dev[dev<devmin] = devmin
+    if len(dev[dev>devmax]) > 0:
+        print(f'Warning: {len(dev[dev>devmax])} particles have deviation > {devmax} mm')
+
+    if show_colorbar:
+        fig = plt.figure(figsize=(7.11, 3.5))
+    else:
+        fig = plt.figure(figsize=(6, 3.5))
+    plt.scatter(1.e3*part_init.x[part.parent_particle_id[mask]],
+                1.e3*part_init.y[part.parent_particle_id[mask]],
+                c=dev, s=0.25, cmap='viridis',norm=mcolors.LogNorm(vmin=devmin, vmax=devmax))
+    plt.vlines(10, -5, 5, colors='r')
+    plt.xlim(8.5, 11.5)
+    plt.ylim(-5, 5)
+    plt.xlabel('x [mm]')
+    plt.ylabel('y [mm]')
+    if show_colorbar:
+        plt.colorbar(label='Deviation from drift [mm]')
+    plt.tight_layout()
+    plt.savefig(outfile, dpi=300)
+
+
+# Generate particles
+# ------------------
+
+y = np.random.normal(loc=0, scale=1e-3, size=num_part)
+py = np.random.normal(loc=0, scale=1e-3, size=num_part)
+delta = np.random.normal(loc=0, scale=3e-4, size=num_part)
+
+part_init = []
+x0 = 9.25e-3
+x = np.random.normal(loc=x0, scale=1e-4, size=num_part)
+part_init.append(xp.build_particles(particle_ref=particle_ref, x=x, px=0, y=y, py=py, delta=delta, _capacity=_capacity))
+
+x0 = 0.01
+x = np.random.normal(loc=x0, scale=1e-4, size=num_part)
+part_init.append(xp.build_particles(particle_ref=particle_ref, x=x, px=0, y=y, py=py, delta=delta, _capacity=_capacity))
+
+x0 = 0.01075
+x = np.random.normal(loc=x0, scale=1e-4, size=num_part)
+part_init.append(xp.build_particles(particle_ref=particle_ref, x=x, px=0, y=y, py=py, delta=delta, _capacity=_capacity))
+
+
+# Fluka
+# -----
+
+part_fluka = []
+for this_part_init in part_init:
+    this_part = this_part_init.copy()
+    t_start = time.time()
+    fluka_coll.track(this_part)
+    t_end = time.time()
+
+    print(f"FLUKA execution time: {t_end - t_start} s")
+    print(f"Survived in FLUKA: {len(np.unique(this_part.parent_particle_id[this_part.state>0]))}/{num_part}")
+    this_part.sort(interleave_lost_particles=True)
+    part_fluka.append(this_part)
+
+xc.FlukaEngine.stop()
+
+for idx in range(3):
+    plot_scatters(part_init[idx], part_fluka[idx], f'fluka_collimator_{idx}.png', show_colorbar=True)
+
+
+# Everest
+# -------
+
+# Pre-run for compilation
+this_part = part_init[0].copy()
+everest_coll.track(this_part)
+
+part_everest = []
+for this_part_init in part_init:
+    this_part = this_part_init.copy()
+    t_start = time.time()
+    everest_coll.track(this_part)
+    t_end = time.time()
+
+    print(f"Everest execution time: {t_end - t_start} s")
+    print(f"Survived in Everest: {len(np.unique(this_part.parent_particle_id[this_part.state>0]))}/{num_part}")
+    this_part.sort(interleave_lost_particles=True)
+    part_everest.append(this_part)
+
+for idx in range(3):
+    plot_scatters(part_init[idx], part_everest[idx], f'everest_collimator_{idx}.png', show_colorbar=True)
+
+
+# # Geant4
+# # ------
+
+# xc.Geant4Engine.start(elements=[g4_coll], seed=1993, bdsim_config_file='settings_protons.gmad', particle_ref=particle_ref)
+
+# part_g4 = []
+# for this_part_init in part_init:
+#     this_part = this_part_init.copy()
+#     t_start = time.time()
+#     g4_coll.track(this_part)
+#     t_end = time.time()
+
+#     print(f"Geant4 execution time: {t_end - t_start} s")
+#     print(f"Survived in Geant4: {len(np.unique(this_part.parent_particle_id[this_part.state>0]))}/{num_part}")
+#     this_part.sort(interleave_lost_particles=True)
+#     part_g4.append(this_part)
+
+# xc.Geant4Engine.stop()
+
+# for idx in range(3):
+#     plot_scatters(part_init[idx], part_g4[idx], f'g4_collimator_{idx}.png', show_colorbar=True)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xcoll"
-version = "0.6.3"
+version = "0.6.4"
 description = "Xsuite collimation package"
 authors = [
     {name="Frederik F. Van der Veken", email="[email protected]"},
@@ -25,7 +25,7 @@ dependencies = [
     "xobjects>=0.5.2",
     "xdeps>=0.10.6",
     "xpart>=0.23.1",
-    "xtrack>=0.89.3",
+    "xtrack>=0.90.1",
 ]
 
 [project.optional-dependencies]
@@ -46,7 +46,7 @@ exclude = ["xcoll/lib", "xcoll/config"]
 
 [poetry.group.dev.dependencies]
 pytest = ">=7.3"
-xaux = ">=0.3.5"
+xaux = ">=0.3.6"
 
 [build-system]
 # Needed for pip install -e (BTW: need pip version 22)

tests/test_rf_sweep.py

@@ -3,15 +3,14 @@
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
-from pathlib import Path
 import numpy as np
 import pytest
 
 import xtrack as xt
 import xcoll as xc
 from xobjects.test_helpers import for_all_test_contexts
 
-path = Path(__file__).parent / 'data'
+path = xc._pkg_root.parent / 'tests' / 'data'
 
 @for_all_test_contexts
 @pytest.mark.parametrize("sweep, beam", [[-300, 1], [300, 2], [3500, 3]],
@@ -20,27 +19,52 @@ def test_rf_sweep(sweep, beam, test_context):
     num_turns = 6000
     num_particles = 5
     if beam == 3:
+        bh = 3.5e-3
+        sweep *= 25
         line = xt.load(path / f'sequence_sps_q20_inj.json')['sps']
-        line.insert('markkk', xt.Marker(), at=2822)
+        line.insert('cavity_mid', xt.Marker(), at=2822) # Inside thick cavity as to slice it
         tt_c = line.get_table().rows['ac.*']
         assert 'ThickSliceCavity' in tt_c.element_type
     else:
+        bh = 3.e-4
         line = xt.load(path / f'sequence_lhc_run3_b{beam}.json')
 
-    line.build_tracker()
+    line.build_tracker(_context=test_context)
 
-    part = line.build_particles(delta=np.linspace(-2e-4, 2e-4, num_particles),
-                                x_norm=0, px_norm=0, y_norm=0, py_norm=0)
+    part = line.build_particles(delta=np.linspace(-0.7*bh, 0.7*bh, num_particles),
+                                x_norm=0, px_norm=0, y_norm=0, py_norm=0, _context=test_context)
 
     rf_sweep = xc.RFSweep(line)
     rf_sweep.prepare(sweep_per_turn=sweep/num_turns)
     rf_sweep.info()
-    # This sweep is 3.5 buckets, so check that all particles are at least 3 buckets away
     line.track(particles=part, num_turns=num_turns)
 
+    # This sweep is 3.5 buckets, so check that all particles are at least 3 buckets away
     # negative sweep => positive off-momentum etc
     if sweep < 0:
-        assert np.all(part.delta > 1.5e-3)
+        assert np.all(part.delta > 3*bh)
     else:
-        assert np.all(part.delta < -1.5e-3)
+        assert np.all(part.delta < -3*bh)
+
+
+def test_rf_sweep_old_style():
+    num_turns = 6000
+    num_particles = 5
+    line = xt.load(path / f'sequence_lhc_run3_b2.json')
 
+    line.build_tracker()
+
+    part = line.build_particles(delta=np.linspace(-2e-4, 2e-4, num_particles),
+                                x_norm=0, px_norm=0, y_norm=0, py_norm=0)
+
+    sweep = -300
+    rf_sweep = xc.RFSweep(line)
+    with pytest.warns(DeprecationWarning, match="The `sweep` argument is deprecated."):
+        with pytest.warns(DeprecationWarning, match="The `num_turns` argument is deprecated."):
+            rf_sweep.info(sweep=sweep, num_turns=num_turns)
+    with pytest.warns(DeprecationWarning):
+        rf_sweep.track(sweep=sweep, num_turns=num_turns, particles=part)
+
+    # This sweep is 3.5 buckets, so check that all particles are at least 3 buckets away
+    # negative sweep => positive off-momentum etc
+    assert np.all(part.delta > 1.5e-3)

tests/test_version.py

@@ -6,4 +6,4 @@
 from xcoll import __version__
 
 def test_version():
-    assert __version__ == '0.6.3'
+    assert __version__ == '0.6.4'

xcoll/general.py

@@ -12,5 +12,5 @@
 # ======================
 # Do not change
 # ======================
-__version__ = '0.6.3'
+__version__ = '0.6.4'
 # ======================