fix some tests

Author: lgiacome

tests/test_xwakes_kick_vs_pyheadtail.py

@@ -7,15 +7,18 @@
 import xwakes as xw
 import xtrack as xt
 import xobjects as xo
+from xobjects.test_helpers import for_all_test_contexts
 
 test_data_folder = pathlib.Path(__file__).parent.joinpath(
     '../test_data').absolute()
 
-def test_xwakes_kick_vs_pyheadtail_table_dipolar():
+@for_all_test_contexts(excluding="ContextPyopencl")
+def test_xwakes_kick_vs_pyheadtail_table_dipolar(test_context):
 
     from xpart.pyheadtail_interface.pyhtxtparticles import PyHtXtParticles
 
-    p = xt.Particles(p0c=7e12, zeta=np.linspace(-1, 1, 100000))
+    p = xt.Particles(p0c=7e12, zeta=np.linspace(-1, 1, 100000),
+                    _context=test_context)
     p.x[p.zeta > 0] += 1e-3
     p.y[p.zeta > 0] += 1e-3
     p_table = p.copy()
@@ -31,7 +34,8 @@ def test_xwakes_kick_vs_pyheadtail_table_dipolar():
                         'quadrupolar_xy', 'dipolar_yx', 'quadrupolar_yx',
                         'constant_x', 'constant_y'])
     wake_from_table = xw.WakeFromTable(table, columns=['dipolar_x', 'dipolar_y'])
-    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=100)
+    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=100,
+                                           _context=test_context)
 
     # Zotter convention
     assert table['dipolar_x'].values[1] > 0
@@ -72,11 +76,13 @@ def test_xwakes_kick_vs_pyheadtail_table_dipolar():
     xo.assert_allclose(p_table.px, p_ref.px, atol=1e-30, rtol=2e-3)
     xo.assert_allclose(p_table.py, p_ref.py, atol=1e-30, rtol=2e-3)
 
-def test_xwakes_kick_vs_pyheadtail_table_quadrupolar():
+@for_all_test_contexts(excluding="ContextPyopencl")
+def test_xwakes_kick_vs_pyheadtail_table_quadrupolar(test_context):
 
     from xpart.pyheadtail_interface.pyhtxtparticles import PyHtXtParticles
 
-    p = xt.Particles(p0c=7e12, zeta=np.linspace(-1, 1, 100000))
+    p = xt.Particles(p0c=7e12, zeta=np.linspace(-1, 1, 100000),
+                     _context=test_context)
     p.x[p.zeta > 0] += 1e-3
     p.y[p.zeta > 0] += 1e-3
     p_table = p.copy()
@@ -93,7 +99,8 @@ def test_xwakes_kick_vs_pyheadtail_table_quadrupolar():
                         'constant_x', 'constant_y'])
 
     wake_from_table = xw.WakeFromTable(table, columns=['quadrupolar_x', 'quadrupolar_y'])
-    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=100)
+    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=100,
+                                           _context=test_context)
 
     # This is specific of this table
     assert table['quadrupolar_x'].values[1] < 0
@@ -136,13 +143,16 @@ def test_xwakes_kick_vs_pyheadtail_table_quadrupolar():
     xo.assert_allclose(p_table.py, p_ref.py, atol=1e-30, rtol=2e-3)
 
 
-def test_xwakes_kick_vs_pyheadtail_table_longitudinal():
+@for_all_test_contexts(excluding="ContextPyopencl")
+def test_xwakes_kick_vs_pyheadtail_table_longitudinal(test_context):
 
     from xpart.pyheadtail_interface.pyhtxtparticles import PyHtXtParticles
 
     p = xt.Particles.merge([
-        xt.Particles(p0c=7e12, zeta=np.linspace(-1e-3, 1e-3, 100000)),
-        xt.Particles(p0c=7e12, zeta=1e-6+np.zeros(100000))
+        xt.Particles(p0c=7e12, zeta=np.linspace(-1e-3, 1e-3, 100000),
+                     _context=test_context),
+        xt.Particles(p0c=7e12, zeta=1e-6+np.zeros(100000),
+                     _context=test_context),
     ])
 
     p_table = p.copy()
@@ -158,7 +168,9 @@ def test_xwakes_kick_vs_pyheadtail_table_longitudinal():
                         'quadrupolar_xy', 'dipolar_yx', 'quadrupolar_yx',
                         'constant_x', 'constant_y'])
     wake_from_table = xw.WakeFromTable(table, columns=['time', 'longitudinal'])
-    wake_from_table.configure_for_tracking(zeta_range=(-2e-3, 2e-3), num_slices=1000)
+    wake_from_table.configure_for_tracking(zeta_range=(-2e-3, 2e-3),
+                                           num_slices=1000,
+                                           _context=test_context)
 
     assert len(wake_from_table.components) == 1
     assert wake_from_table.components[0].plane == 'z'
@@ -191,12 +203,13 @@ def test_xwakes_kick_vs_pyheadtail_table_longitudinal():
     assert np.max(p_ref.delta) > 1e-12
     xo.assert_allclose(p_table.delta, p_ref.delta, atol=1e-14, rtol=0)
 
-def test_xwakes_kick_vs_pyheadtail_resonator_dipolar():
+@for_all_test_contexts(excluding="ContextPyopencl")
+def test_xwakes_kick_vs_pyheadtail_resonator_dipolar(test_context):
 
     from xpart.pyheadtail_interface.pyhtxtparticles import PyHtXtParticles
 
     p = xt.Particles(p0c=7e12, zeta=np.linspace(-1, 1, 100000),
-                    weight=1e14)
+                    weight=1e14, _context=test_context)
     p.x[p.zeta > 0] += 1e-3
     p.y[p.zeta > 0] += 1e-3
     p_table = p.copy()
@@ -207,7 +220,8 @@ def test_xwakes_kick_vs_pyheadtail_resonator_dipolar():
         r=1e8, q=1e7, f_r=1e9,
         kind=xw.Yokoya('circular'), # equivalent to: kind=['dipolar_x', 'dipolar_y'],
     )
-    wake.configure_for_tracking(zeta_range=(-1, 1), num_slices=50)
+    wake.configure_for_tracking(zeta_range=(-1, 1), num_slices=50,
+                                _context=test_context)
 
     assert len(wake.components) == 2
     assert wake.components[0].plane == 'x'
@@ -237,7 +251,8 @@ def test_xwakes_kick_vs_pyheadtail_resonator_dipolar():
     table = pd.DataFrame({'time': t_samples, 'dipolar_x': w_dipole_x_samples,
                             'dipolar_y': w_dipole_y_samples})
     wake_from_table = xw.WakeFromTable(table)
-    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=50)
+    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=50,
+                                           _context=test_context)
 
     assert len(wake_from_table.components) == 2
     assert wake_from_table.components[0].plane == 'x'
@@ -277,12 +292,13 @@ def test_xwakes_kick_vs_pyheadtail_resonator_dipolar():
     xo.assert_allclose(p_table.px, p_ref.px, rtol=0, atol=2e-3*np.max(np.abs(p_ref.px)))
     xo.assert_allclose(p_table.py, p_ref.py, rtol=0, atol=2e-3*np.max(np.abs(p_ref.py)))
 
-def test_xwakes_kick_vs_pyheadtail_resonator_quadrupolar():
+@for_all_test_contexts(excluding="ContextPyopencl")
+def test_xwakes_kick_vs_pyheadtail_resonator_quadrupolar(test_context):
 
     from xpart.pyheadtail_interface.pyhtxtparticles import PyHtXtParticles
 
     p = xt.Particles(p0c=7e12, zeta=np.linspace(-1, 1, 100000),
-                    weight=1e14)
+                    weight=1e14, _context=test_context)
     p.x[p.zeta > 0] += 1e-3
     p.y[p.zeta > 0] += 1e-3
     p_table = p.copy()
@@ -293,7 +309,8 @@ def test_xwakes_kick_vs_pyheadtail_resonator_quadrupolar():
         r=1e8, q=1e7, f_r=1e9,
         kind=['quadrupolar_x', 'quadrupolar_y'],
     )
-    wake.configure_for_tracking(zeta_range=(-1, 1), num_slices=50)
+    wake.configure_for_tracking(zeta_range=(-1, 1), num_slices=50,
+                                _context=test_context)
 
     assert len(wake.components) == 2
     assert wake.components[0].plane == 'x'
@@ -322,7 +339,8 @@ def test_xwakes_kick_vs_pyheadtail_resonator_quadrupolar():
     table = pd.DataFrame({'time': t_samples, 'quadrupolar_x': w_quadrupole_x_samples,
                             'quadrupolar_y': w_quadrupole_y_samples})
     wake_from_table = xw.WakeFromTable(table)
-    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=50)
+    wake_from_table.configure_for_tracking(zeta_range=(-1, 1), num_slices=50,
+                                           _context=test_context)
 
     assert len(wake_from_table.components) == 2
     assert wake_from_table.components[0].plane == 'x'
@@ -365,12 +383,13 @@ def test_xwakes_kick_vs_pyheadtail_resonator_quadrupolar():
     xo.assert_allclose(p_table.px, p_ref.px, rtol=0, atol=2e-3*np.max(np.abs(p_ref.px)))
     xo.assert_allclose(p_table.py, p_ref.py, rtol=0, atol=2e-3*np.max(np.abs(p_ref.py)))
 
-def test_xwakes_kick_vs_pyheadtail_resonator_longitudinal():
+@for_all_test_contexts(excluding="ContextPyopencl")
+def test_xwakes_kick_vs_pyheadtail_resonator_longitudinal(test_context):
 
     from xpart.pyheadtail_interface.pyhtxtparticles import PyHtXtParticles
 
     p = xt.Particles(p0c=7e12, zeta=np.linspace(-1, 1, 100000),
-                    weight=1e14)
+                     weight=1e14, _context=test_context)
     p.x[p.zeta > 0] += 1e-3
     p.y[p.zeta > 0] += 1e-3
     p_table = p.copy()
@@ -381,7 +400,8 @@ def test_xwakes_kick_vs_pyheadtail_resonator_longitudinal():
         r=1e8, q=1e7, f_r=1e9,
         kind='longitudinal'
     )
-    wake.configure_for_tracking(zeta_range=(-1.01, 1.01), num_slices=50)
+    wake.configure_for_tracking(zeta_range=(-1.01, 1.01), num_slices=50,
+                                _context=test_context)
 
     assert len(wake.components) == 1
     assert wake.components[0].plane == 'z'
@@ -402,7 +422,8 @@ def test_xwakes_kick_vs_pyheadtail_resonator_longitudinal():
     w_longitudinal_x_samples[0] *= 2 # Undo sampling weight
     table = pd.DataFrame({'time': t_samples, 'longitudinal': w_longitudinal_x_samples})
     wake_from_table = xw.WakeFromTable(table)
-    wake_from_table.configure_for_tracking(zeta_range=(-1.01, 1.01), num_slices=50)
+    wake_from_table.configure_for_tracking(zeta_range=(-1.01, 1.01), num_slices=50,
+                                           _context=test_context)
 
     assert len(wake_from_table.components) == 1
     assert wake_from_table.components[0].plane == 'z'

tests/test_xwakes_tune_shift.py

@@ -7,16 +7,19 @@
 import xpart as xp
 import xwakes as xw
 import xobjects as xo
+from xobjects.test_helpers import for_all_test_contexts
 
+@for_all_test_contexts(excluding="ContextPyopencl")
 @pytest.mark.parametrize('wake_type', ['dipolar', 'quadrupolar'])
 @pytest.mark.parametrize('plane', ['x', 'y'])
-def test_tune_shift_transverse(wake_type, plane):
+def test_tune_shift_transverse(wake_type, plane, test_context):
 
     table = pd.DataFrame({'time': [0, 10],
                         f'{wake_type}_{plane}': [6e15, 6e15]})
 
     wake = xw.WakeFromTable(table)
-    wake.configure_for_tracking(zeta_range=(-20e-2, 20e-2), num_slices=100)
+    wake.configure_for_tracking(zeta_range=(-20e-2, 20e-2), num_slices=100,
+                                _context=test_context)
 
     assert len(wake.components) == 1
     assert wake.components[0].plane == plane
@@ -49,15 +52,15 @@ def test_tune_shift_transverse(wake_type, plane):
     line_no_wake = xt.Line(elements=[one_turn_map])
     line_with_wake = xt.Line(elements=[one_turn_map, wake])
 
-    line_no_wake.particle_ref = xt.Particles(p0c=2e9)
-    line_with_wake.particle_ref = xt.Particles(p0c=2e9)
+    line_no_wake.particle_ref = xt.Particles(p0c=2e9, _context=test_context)
+    line_with_wake.particle_ref = xt.Particles(p0c=2e9, _context=test_context)
 
-    line_no_wake.build_tracker()
-    line_with_wake.build_tracker()
+    line_no_wake.build_tracker(_context=test_context)
+    line_with_wake.build_tracker(_context=test_context)
 
     p = xp.generate_matched_gaussian_bunch(line=line_no_wake,
                 num_particles=1000, nemitt_x=1e-6, nemitt_y=1e-6, sigma_z=0.07,
-                total_intensity_particles=1e11)
+                total_intensity_particles=1e11, _context=test_context)
     p.x += 1e-3
     p.y += 1e-3
 
@@ -67,8 +70,8 @@ def test_tune_shift_transverse(wake_type, plane):
                 mean_y=lambda line, part: part.y.mean())
 
     line_no_wake.track(p0, num_turns=100,
-                    log=mylog,
-                    with_progress=10)
+                       log=mylog,
+                       with_progress=10)
     log_no_wake = line_no_wake.log_last_track
 
     line_with_wake.track(p, num_turns=100,
@@ -78,8 +81,15 @@ def test_tune_shift_transverse(wake_type, plane):
 
     import nafflib as nl
 
-    tune_no_wake = nl.get_tune(log_no_wake[f'mean_{plane}'])
-    tune_with_wake = nl.get_tune(log_with_wake[f'mean_{plane}'])
+    mean_no_wake = test_context.nplike_lib.array(log_no_wake[f'mean_{plane}'])
+    mean_with_wake = test_context.nplike_lib.array(log_with_wake[f'mean_{plane}'])
+
+    if hasattr(mean_no_wake, 'get'):
+        tune_no_wake = nl.get_tune(mean_no_wake.get())
+        tune_with_wake = nl.get_tune(mean_with_wake.get())
+    else:
+        tune_no_wake = nl.get_tune(mean_no_wake)
+        tune_with_wake = nl.get_tune(mean_with_wake)
 
     print(f'Tune without wake: {tune_no_wake}')
     print(f'Tune with wake: {tune_with_wake}')
@@ -89,13 +99,15 @@ def test_tune_shift_transverse(wake_type, plane):
     xo.assert_allclose(tune_no_wake, {'x': 0.28, 'y': 0.31}[plane], atol=1e-6, rtol=0)
     xo.assert_allclose(tune_with_wake, tune_no_wake - 2e-3, atol=0.3e-3, rtol=0)
 
-def test_tune_shift_longitudinal():
+@for_all_test_contexts(excluding="ContextPyopencl")
+def test_tune_shift_longitudinal(test_context):
 
     table = pd.DataFrame({'time': [0, 10],
                         'longitudinal': [1e13, 1e13]})
 
     wake = xw.WakeFromTable(table)
-    wake.configure_for_tracking(zeta_range=(-20e-2, 20e-2), num_slices=100)
+    wake.configure_for_tracking(zeta_range=(-20e-2, 20e-2), num_slices=100,
+                                _context=test_context)
 
     assert len(wake.components) == 1
     assert wake.components[0].plane == 'z'
@@ -124,19 +136,20 @@ def test_tune_shift_longitudinal():
     line_no_wake = xt.Line(elements=[one_turn_map])
     line_with_wake = xt.Line(elements=[one_turn_map, wake])
 
-    line_no_wake.particle_ref = xt.Particles(p0c=2e9)
-    line_with_wake.particle_ref = xt.Particles(p0c=2e9)
+    line_no_wake.particle_ref = xt.Particles(p0c=2e9, _context=test_context)
+    line_with_wake.particle_ref = xt.Particles(p0c=2e9, _context=test_context)
 
-    line_no_wake.build_tracker()
-    line_with_wake.build_tracker()
+    line_no_wake.build_tracker(_context=test_context)
+    line_with_wake.build_tracker(_context=test_context)
 
     p = xp.generate_matched_gaussian_bunch(line=line_no_wake,
                 num_particles=1000, nemitt_x=1e-6, nemitt_y=1e-6, sigma_z=0.07,
-                total_intensity_particles=1e11)
+                total_intensity_particles=1e11,
+                _context=test_context)
     p.zeta += 5e-3
 
     p0 = p.copy()
-
+    
     mylog = xt.Log(mean_zeta=lambda line, part: part.zeta.mean(),
                 mean_delta=lambda line, part: part.delta.mean())
 
@@ -146,14 +159,23 @@ def test_tune_shift_longitudinal():
     log_no_wake = line_no_wake.log_last_track
 
     line_with_wake.track(p, num_turns=3000,
-                            log=mylog,
-                            with_progress=10)
+                         log=mylog,
+                         with_progress=10)
     log_with_wake = line_with_wake.log_last_track
 
     import nafflib as nl
 
-    tune_no_wake = nl.get_tune(log_no_wake[f'mean_zeta'])
-    tune_with_wake = nl.get_tune(log_with_wake[f'mean_zeta']-np.mean(log_with_wake[f'mean_zeta']))
+    mean_zeta_no_wake = test_context.nplike_lib.array(log_no_wake[f'mean_zeta'])
+    mean_zeta_with_wake = test_context.nplike_lib.array(log_with_wake[f'mean_zeta'])
+
+    if hasattr(mean_zeta_no_wake, 'get'):
+        tune_no_wake = nl.get_tune(mean_zeta_no_wake.get())
+        tune_with_wake = nl.get_tune(mean_zeta_with_wake.get() -
+                                     np.mean(mean_zeta_with_wake.mean().get()))
+    else:
+        tune_no_wake = nl.get_tune(mean_zeta_no_wake)
+        tune_with_wake = nl.get_tune(mean_zeta_with_wake -
+                                     np.mean(mean_zeta_with_wake))
 
     print(f'Tune without wake: {tune_no_wake}')
     print(f'Tune with wake: {tune_with_wake}')