Merge pull request #23 from xsuite/investigate_omp

v0.1.0

Author: giadarol

examples/apertures/000_rect_aperture_standalone.py

@@ -26,8 +26,7 @@
         sigma = np.zeros(n_part),
         delta = np.zeros(n_part))
 
-part_dict = xt.pyparticles_to_xtrack_dict(pyst_part)
-particles = xt.Particles(_context=context, **part_dict)
+particles = xt.Particles(_context=context, **pyst_part.to_dict())
 
 aper_pyst = xline.elements.LimitRect(min_x=x_aper_min,
                                           max_x=x_aper_max,

examples/apertures/001_rect_aperture_withinline.py

@@ -3,7 +3,8 @@
 import xobjects as xo
 import xtrack as xt
 
-import xline
+import xline as xl
+import xpart as xp
 
 context = xo.ContextCpu()
 context = xo.ContextCupy()
@@ -17,7 +18,7 @@
 part_gen_range = 0.35
 n_part=10000
 
-pyst_part = xline.Particles(
+xparticles = xp.Particles(
         p0c=6500e9,
         x=np.random.uniform(-part_gen_range, part_gen_range, n_part),
         px = np.zeros(n_part),
@@ -26,27 +27,26 @@
         sigma = np.zeros(n_part),
         delta = np.zeros(n_part))
 
-part_dict = xt.pyparticles_to_xtrack_dict(pyst_part)
-particles = xt.Particles(_context=context, **part_dict)
+particles = xt.Particles(_context=context, **xparticles.to_dict())
 
-aper_pyst = xline.elements.LimitRect(min_x=x_aper_min,
+aper_xline = xl.elements.LimitRect(min_x=x_aper_min,
                                           max_x=x_aper_max,
                                           min_y=y_aper_min,
                                           max_y=y_aper_max)
 
 aper = xt.LimitRect(_context=context,
-                    **aper_pyst.to_dict())
+                    **aper_xline.to_dict())
 
-aper_pyst.track(pyst_part)
+aper_xline.track(xparticles)
 
 # Build a small test line
-pyst_line = xline.Line(elements=[
-                xline.elements.Drift(length=5.),
-                aper_pyst,
-                xline.elements.Drift(length=5.)],
+line = xl.Line(elements=[
+                xl.elements.Drift(length=5.),
+                aper_xline,
+                xl.elements.Drift(length=5.)],
                 element_names=['drift0', 'aper', 'drift1'])
 
-tracker = xt.Tracker(_context=context, sequence=pyst_line)
+tracker = xt.Tracker(_context=context, sequence=line)
 
 tracker.track(particles)
 
@@ -58,7 +58,7 @@
 
 id_alive = part_id[part_state>0]
 
-assert np.allclose(np.sort(pyst_part.particle_id), np.sort(id_alive))
+assert np.allclose(np.sort(xparticles.particle_id), np.sort(id_alive))
 assert np.allclose(part_s[part_state>0], 10.)
 assert np.allclose(part_s[part_state<1], 5.)
 

examples/full_ring/005_one_turn_check.py

@@ -111,8 +111,7 @@
 problem_found = False
 for ii, (eexl, nn) in enumerate(zip(sequence.elements, sequence.element_names)):
     vars_before = {vv :getattr(xl_part, vv) for vv in vars_to_check}
-    pp_dict = xt.pyparticles_to_xtrack_dict(xl_part)
-    particles.set_particle(ip_check, **pp_dict)
+    particles.set_particle(ip_check, **xl_part.to_dict())
 
     tracker.track(particles, ele_start=ii, num_elements=1)
 

examples/full_ring/006_timing.py

@@ -40,8 +40,9 @@
 # Choose a context #
 ####################
 
-n_part = 200
-context = xo.ContextCpu()
+n_part = 8000
+context = xo.ContextCpu(omp_num_threads=8)
+#context = xo.ContextCpu(omp_num_threads=0)
 
 #n_part = 20000
 #context = xo.ContextCupy()
@@ -75,6 +76,21 @@
 
 print('Import sequence')
 sequence = xl.Line.from_dict(input_data['line'])
+
+
+#fodo = [xl.Drift(length=1.), xl.Multipole(knl=[0, 1.]),
+#        xl.Drift(length=1.), xl.Multipole(knl=[0, -1.])]
+#
+#elelist = 4000*fodo
+#elenames = [f'e{ii}' for ii in range(len(elelist))]
+#
+#sequence = xl.Line(elements=elelist, element_names=elenames)
+
+
+
+
+
+
 if short_test:
     sequence = make_short_line(sequence)
 
@@ -84,7 +100,8 @@
 
 print('Build tracker')
 tracker = xt.Tracker(_context=context,
-                     sequence=sequence)
+                     sequence=sequence,
+                     save_source_as='source.c')
 
 ######################
 # Get some particles #