Adapt examples

giadarol · giadarol · commit 900f609a987d · 2025-11-02T23:23:09.000+01:00
diff --git a/examples/particles_generation/001a_build_particles_set.py b/examples/particles_generation/001a_build_particles_set.py
@@ -9,11 +9,8 @@
 # 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])
+particles = xp.build_particles(particle_ref=p0, y=[1,2,3])
 
 # Inspect
 print(particles.p0c[1]) # gives 7e12
diff --git a/examples/particles_generation/001at_build_particles_set_with_tracker.py b/examples/particles_generation/001at_build_particles_set_with_tracker.py
@@ -2,22 +2,16 @@
 # This file is part of the Xpart Package.   #
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
-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 line
-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'])
-line.build_tracker(_context=ctx)
+# Get a line
+line = xt.load('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
 
 # Attach a reference particle to the line
-line.particle_ref = xp.Particles(p0c=7e12, mass0=xp.PROTON_MASS_EV, q0=1, x =1 , y=3)
+line.set_particle_ref('proton', p0c=7e12, x=1, y=3)
 
 # Built a set of three particles with different y coordinates
 # (context and particle_ref are taken from the line)
diff --git a/examples/particles_generation/001b_build_particles_shift.py b/examples/particles_generation/001b_build_particles_shift.py
@@ -9,12 +9,8 @@
 # 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(mode='shift', particle_ref=p0, y=[1,2,3],
-                               _context=ctx)
+particles = xp.build_particles(mode='shift', particle_ref=p0, y=[1,2,3])
 
 # Inspect
 print(particles.p0c[1]) # gives 7e12
diff --git a/examples/particles_generation/001bt_build_particles_shift_with_tracker.py b/examples/particles_generation/001bt_build_particles_shift_with_tracker.py
@@ -2,22 +2,16 @@
 # This file is part of the Xpart Package.   #
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
-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 line
-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'])
-line.build_tracker(_context=ctx)
+# Load machine model
+line = xt.load('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
 
 # Attach a reference particle to the line
-line.particle_ref = xp.Particles(p0c=7e12, mass0=xp.PROTON_MASS_EV, q0=1, x=1 , y=3)
+line.set_particle_ref('proton', p0c=7e12, x=1, y=3)
 
 # Built a set of three particles with different y coordinates
 # (context and particle_ref are taken from the line)
diff --git a/examples/particles_generation/001c_build_particles_normalized.py b/examples/particles_generation/001c_build_particles_normalized.py
@@ -3,24 +3,14 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-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'])
-line.build_tracker(_context=ctx)
+line = xt.load('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
 
 # Attach a reference particle to the line
-line.particle_ref = xp.Particles(mass0=xp.PROTON_MASS_EV, q0=1, p0c=7e12, x=1, y=3)
+line.set_particle_ref('proton', p0c=7e12, x=1, y=3)
 
 # Built a set of three particles with different x coordinates
 particles = line.build_particles(
diff --git a/examples/particles_generation/002_halo.py b/examples/particles_generation/002_halo.py
@@ -3,7 +3,6 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-import json
 import numpy as np
 
 import xpart as xp
@@ -13,13 +12,9 @@
 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)
-line = xt.Line.from_dict(input_data['line'])
-line.particle_ref = xp.Particles.from_dict(input_data['particle'])
-line.build_tracker()
+line = xt.load('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
+
+line.set_particle_ref('proton', p0c=7e12)
 
 # Horizontal plane: generate cut halo distribution
 (x_in_sigmas, px_in_sigmas, r_points, theta_points
diff --git a/examples/particles_generation/003_pencil.py b/examples/particles_generation/003_pencil.py
@@ -3,7 +3,6 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-import json
 import numpy as np
 
 import xpart as xp
@@ -13,13 +12,9 @@
 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)
-line=xt.Line.from_dict(input_data['line'])
-line.particle_ref = xp.Particles.from_dict(input_data['particle'])
-line.build_tracker()
+line = xt.load('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
+
+line.set_particle_ref('proton', p0c=7e12)
 
 # Horizontal plane: generate gaussian distribution in normalized coordinates
 x_in_sigmas, px_in_sigmas = xp.generate_2D_gaussian(num_particles)
diff --git a/examples/particles_generation/003a_pencil_on_collimator.py b/examples/particles_generation/003a_pencil_on_collimator.py
@@ -3,7 +3,6 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-import json
 import numpy as np
 
 import xpart as xp
@@ -13,17 +12,13 @@
 nemitt_x = 2.5e-6
 nemitt_y = 3e-6
 
-# Load machine model
-filename = ('../../../xtrack/test_data/hllhc15_noerrors_nobb/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'])
-line.build_tracker()
+line = xt.load('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
+
+line.set_particle_ref('proton', p0c=7e12)
 
 # Location of the collimator
 at_element = 'tcp.d6l7.b1'
-at_s = line.line.get_s_position(at_element) + 1.
+at_s = line.get_s_position(at_element) + 1.
 y_cut = 3e-3 # position of the jaw
 pencil_dr_sigmas = 3 # width of the pencil
 side = '+' # side of the pencil
diff --git a/examples/particles_generation/003b_pencil_on_collimator_dispersion.py b/examples/particles_generation/003b_pencil_on_collimator_dispersion.py
@@ -3,7 +3,6 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-import json
 import numpy as np
 
 import xpart as xp
@@ -13,14 +12,9 @@
 nemitt_x = 2.5e-6
 nemitt_y = 3e-6
 
-# Load machine model
-filename = ('../../../xtrack/test_data/hllhc15_noerrors_nobb/line_and_particle.json')
-line = xt.Line.from_json(filename)
-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'])
-line.build_tracker()
+line = xt.load('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
+
+line.set_particle_ref('proton', p0c=7e12)
 
 # Location of the collimator
 at_element = 'tcp.6l3.b1' # High dispersion
@@ -54,7 +48,7 @@
                 at_element=at_element, match_at_s=at_s)
 
 # Drift to at_s position for checking
-drift_to_at_s = xt.Drift(length=at_s-line.line.get_s_position(at_element))
+drift_to_at_s = xt.Drift(length=at_s - line.get_s_position(at_element))
 drift_to_at_s.track(particles)
 
 # Checks and plots
diff --git a/examples/particles_generation/004_generate_gaussian.py b/examples/particles_generation/004_generate_gaussian.py
@@ -3,8 +3,6 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-import json
-
 import numpy as np
 from scipy.constants import e as qe
 from scipy.constants import m_p
@@ -18,14 +16,8 @@
 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'])
-
-line.build_tracker()
+line = xt.load('../../../xtrack/test_data/sps_w_spacecharge/line_no_spacecharge_and_particle.json')
+line.set_particle_ref('proton', p0c=26e9)
 
 particles = xp.generate_matched_gaussian_bunch(
          num_particles=n_part, total_intensity_particles=bunch_intensity,
diff --git a/examples/particles_generation/004b_generate_gaussian_long_bunch.py b/examples/particles_generation/004b_generate_gaussian_long_bunch.py
@@ -3,8 +3,6 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-import json
-
 import numpy as np
 from scipy.constants import e as qe
 from scipy.constants import m_p
@@ -18,11 +16,7 @@
 nemitt_x = 2e-6
 nemitt_y = 2.5e-6
 
-filename = ('../../../xtrack/test_data/sps_ions/line_and_particle.json')
-with open(filename, 'r') as fid:
-    ddd = json.load(fid)
-line = xt.Line.from_dict(ddd)
-line.build_tracker()
+line = xt.load('../../../xtrack/test_data/sps_ions/line_and_particle.json')
 
 particles = xp.generate_matched_gaussian_bunch(
          num_particles=n_part, total_intensity_particles=bunch_intensity,
diff --git a/examples/particles_generation/006_match_at_element.py b/examples/particles_generation/006_match_at_element.py
@@ -3,29 +3,22 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-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'])
-line.build_tracker()
+line = xt.load('../../../xtrack/test_data/lhc_no_bb/line_and_particle.json')
+line.set_particle_ref('proton', p0c=7e12)
 
 # Match distribution at a given element
-particles = tracker.build_particles(x_norm=[0,1,2], px_norm=[0,0,0], # in sigmas
+particles = line.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
+particles = line.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
+                   match_at_s=line.get_s_position('ip6') + 100
                    )
 
 
diff --git a/examples/particles_generation/007a_build_multi_bunch_beam.py b/examples/particles_generation/007a_build_multi_bunch_beam.py
diff --git a/xpart/longitudinal/generate_longitudinal.py b/xpart/longitudinal/generate_longitudinal.py
@@ -185,6 +185,8 @@ def generate_longitudinal_coordinates(
         line = tracker.line
 
     if line is not None:
+        if not line._has_valid_tracker():
+            line.build_tracker()
         if particle_ref is None:
             particle_ref = line.particle_ref
         assert particle_ref is not None
diff --git a/xpart/transverse_generators/pencil.py b/xpart/transverse_generators/pencil.py
@@ -130,7 +130,8 @@ def generate_2D_pencil_with_absolute_cut(num_particles,
         line = tracker.line
 
     if line is not None:
-        assert line.tracker is not None
+        if not line._has_valid_tracker():
+            line.build_tracker()
 
     # kwargs are passed to line.twiss
 
