From d3e22c69d9a049285a4972ded6a71e50131a54f6 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Mon, 11 Nov 2024 10:53:50 +0100 Subject: [PATCH 1/4] add contexts in base wake --- xwakes/basewake.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/xwakes/basewake.py b/xwakes/basewake.py index d8f9edbe..1ab6ffee 100644 --- a/xwakes/basewake.py +++ b/xwakes/basewake.py @@ -12,6 +12,7 @@ def configure_for_tracking(self, zeta_range: Tuple[float, float], filling_scheme=None, circumference=None, bunch_selection=None, + _context=None, **kwargs # for multibunch compatibility ) -> None: from xfields.beam_elements.waketracker import WakeTracker @@ -24,9 +25,10 @@ def configure_for_tracking(self, zeta_range: Tuple[float, float], filling_scheme=filling_scheme, circumference=circumference, bunch_selection=bunch_selection, + _context=_context, **kwargs) - def _reconfigure_for_parallel(self, n_procs, my_rank) -> None: + def _reconfigure_for_parallel(self, n_procs, my_rank, _context) -> None: filled_slots = self._wake_tracker.slicer.filled_slots scheme = np.zeros(np.max(filled_slots) + 1, @@ -44,6 +46,7 @@ def _reconfigure_for_parallel(self, n_procs, my_rank) -> None: bunch_selection=bunch_selection_rank[my_rank], num_turns=self._wake_tracker.num_turns, circumference=self._wake_tracker.circumference, + _context=_context ) def track(self, particles) -> None: From dbde4c37e40cc1e4051b058c3448ee6e6c2a3991 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Mon, 11 Nov 2024 10:54:28 +0100 Subject: [PATCH 2/4] add gpu example --- examples/001d_instability_wake_table_gpu.py | 146 ++++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 examples/001d_instability_wake_table_gpu.py diff --git a/examples/001d_instability_wake_table_gpu.py b/examples/001d_instability_wake_table_gpu.py new file mode 100644 index 00000000..2b7a85b7 --- /dev/null +++ b/examples/001d_instability_wake_table_gpu.py @@ -0,0 +1,146 @@ +import numpy as np +import matplotlib.pyplot as plt +from scipy.signal import hilbert +from scipy.stats import linregress + +import xtrack as xt +import xpart as xp +import xwakes as xw +import xfields as xf +import xobjects as xo + +context = xo.ContextCupy(device=3) +#context = xo.ContextCpu() + +# Simulation settings +n_turns = 1000 + +wake_table_filename = xf.general._pkg_root.joinpath( + '../test_data/HLLHC_wake.dat') +wake_file_columns = ['time', 'longitudinal', 'dipole_x', 'dipole_y', + 'quadrupole_x', 'quadrupole_y', 'dipole_xy', + 'quadrupole_xy', 'dipole_yx', 'quadrupole_yx', + 'constant_x', 'constant_y'] +mytable = xw.read_headtail_file( + wake_file=wake_table_filename, + wake_file_columns=wake_file_columns +) +wf = xw.WakeFromTable( + table=mytable, + columns=['dipole_x', 'dipole_y', 'quadrupole_x', 'quadrupole_y'], +) +wf.configure_for_tracking( + zeta_range=(-0.375, 0.375), + num_slices=100, + num_turns=1, + _context=context +) + +one_turn_map = xt.LineSegmentMap( + length=27e3, betx=70., bety=80., + qx=62.31, qy=60.32, + longitudinal_mode='linear_fixed_qs', + dqx=-10., dqy=-10., # <-- to see fast mode-0 instability + qs=2e-3, bets=731.27, + _context=context +) + +# Generate line +line = xt.Line(elements=[one_turn_map, wf], + element_names=['one_turn_map', 'wf']) + + +line.particle_ref = xt.Particles(p0c=7e12, _context=context) +line.build_tracker(_context=context) + +# Generate particles +particles = xp.generate_matched_gaussian_bunch(line=line, + num_particles=40_000_000, total_intensity_particles=2.3e11, + nemitt_x=2e-6, nemitt_y=2e-6, sigma_z=0.075, + _context=context) + +# Apply a distortion to the bunch to trigger an instability +amplitude = 1e-3 +particles.x += amplitude +particles.y += amplitude + +flag_plot = True + +mean_x_xt = np.zeros(n_turns) +mean_y_xt = np.zeros(n_turns) + +plt.ion() + +fig1 = plt.figure(figsize=(6.4*1.7, 4.8)) +ax_x = fig1.add_subplot(121) +line1_x, = ax_x.plot(mean_x_xt, 'r-', label='average x-position') +line2_x, = ax_x.plot(mean_x_xt, 'm-', label='exponential fit') +ax_x.set_ylim(-3.5, -1) +ax_x.set_xlim(0, n_turns) +ax_y = fig1.add_subplot(122, sharex=ax_x) +line1_y, = ax_y.plot(mean_y_xt, 'b-', label='average y-position') +line2_y, = ax_y.plot(mean_y_xt, 'c-', label='exponential fit') +ax_y.set_ylim(-3.5, -1) +ax_y.set_xlim(0, n_turns) + +plt.xlabel('turn') +plt.ylabel('log10(average x-position)') +plt.legend() + +turns = np.linspace(0, n_turns - 1, n_turns) + +import time + +line.track(particles, num_turns=1) + +start = time.time() + +for i_turn in range(n_turns): + line.track(particles, num_turns=1) + + if i_turn % 50 == 0: + print(f'Turn: {i_turn}') + + ''' + mean_x_xt[i_turn] = np.mean(particles.x) + mean_y_xt[i_turn] = np.mean(particles.y) + + if i_turn % 50 == 0 and i_turn > 1: + i_fit_end = np.argmax(mean_x_xt) # i_turn + i_fit_start = int(i_fit_end * 0.9) + + # compute x instability growth rate + ampls_x_xt = np.abs(hilbert(mean_x_xt)) + fit_x_xt = linregress(turns[i_fit_start: i_fit_end], + np.log(ampls_x_xt[i_fit_start: i_fit_end])) + + # compute y instability growth rate + + ampls_y_xt = np.abs(hilbert(mean_y_xt)) + fit_y_xt = linregress(turns[i_fit_start: i_fit_end], + np.log(ampls_y_xt[i_fit_start: i_fit_end])) + + line1_x.set_xdata(turns[:i_turn]) + line1_x.set_ydata(np.log10(np.abs(mean_x_xt[:i_turn]))) + line2_x.set_xdata(turns[:i_turn]) + line2_x.set_ydata(np.log10(np.exp(fit_x_xt.intercept + + fit_x_xt.slope*turns[:i_turn]))) + + line1_y.set_xdata(turns[:i_turn]) + line1_y.set_ydata(np.log10(np.abs(mean_y_xt[:i_turn]))) + line2_y.set_xdata(turns[:i_turn]) + line2_y.set_ydata(np.log10(np.exp(fit_y_xt.intercept + + fit_y_xt.slope*turns[:i_turn]))) + print(f'xtrack h growth rate: {fit_x_xt.slope}') + print(f'xtrack v growth rate: {fit_y_xt.slope}') + + fig1.canvas.draw() + fig1.canvas.flush_events() + + out_folder = '.' + np.save(f'{out_folder}/mean_x.npy', mean_x_xt) + np.save(f'{out_folder}/mean_y.npy', mean_y_xt) + ''' + +end = time.time() +print(f'Time: {end - start}') \ No newline at end of file From d9a5cb68e8723194ae7423f2648840afc5913de4 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 21 Nov 2024 11:23:51 +0100 Subject: [PATCH 3/4] fix some tests --- tests/test_xwakes_kick_vs_pyheadtail.py | 65 ++++++---- tests/test_xwakes_tune_shift.py | 68 ++++++---- tests/test_xwakes_wake_kicks.py | 166 +++++++++++++++--------- 3 files changed, 192 insertions(+), 107 deletions(-) diff --git a/tests/test_xwakes_kick_vs_pyheadtail.py b/tests/test_xwakes_kick_vs_pyheadtail.py index 989c064c..d463b266 100644 --- a/tests/test_xwakes_kick_vs_pyheadtail.py +++ b/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' diff --git a/tests/test_xwakes_tune_shift.py b/tests/test_xwakes_tune_shift.py index be884238..57b7a55e 100644 --- a/tests/test_xwakes_tune_shift.py +++ b/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}') diff --git a/tests/test_xwakes_wake_kicks.py b/tests/test_xwakes_wake_kicks.py index c150596d..9240dc97 100644 --- a/tests/test_xwakes_wake_kicks.py +++ b/tests/test_xwakes_wake_kicks.py @@ -13,7 +13,7 @@ test_data_folder = pathlib.Path(__file__).parent.joinpath( '../test_data').absolute() -exclude_contexts = ['ContextPyopencl', 'ContextCupy'] +exclude_contexts = ['ContextPyopencl', 'ContextCpu'] kind_to_parameters = { 'longitudinal': {'plane': 'z', 'source_exponents': (0, 0), 'test_exponents': (0, 0)}, @@ -79,12 +79,13 @@ def test_wake_kick_single_bunch(test_context, kind, wake_type): zeta_range=zeta_range, num_slices=num_slices, bunch_spacing_zeta=circumference/h_bunch, - circumference=circumference + circumference=circumference, + _context=test_context ) line = xt.Line(elements=[wf]) - line.build_tracker() - line.particle_ref=xp.Particles(p0c=p0c) + line.build_tracker(_context=test_context) + line.particle_ref = xp.Particles(p0c=p0c) particles = xp.Particles(p0c=p0c, zeta=wf.slicer.zeta_centers.flatten(), @@ -130,14 +131,18 @@ def test_wake_kick_single_bunch(test_context, kind, wake_type): assert comp.test_exponents == kind_to_parameters[kk]['test_exponents'] expected = np.zeros_like(particles.zeta) - - for i_test, z_test in enumerate(particles.zeta): - expected[i_test] += (particles.x[i_test]**comp.test_exponents[0] * - particles.y[i_test]**comp.test_exponents[1] * - np.dot(particles.x**comp.source_exponents[0] * - particles.y**comp.source_exponents[1], - comp.function_vs_zeta(z_test - particles.zeta, - beta0=particles.beta0[0], + x_cpu = test_context.nparray_from_context_array(particles.x) + y_cpu = test_context.nparray_from_context_array(particles.y) + zeta_cpu = test_context.nparray_from_context_array(particles.zeta) + beta0_cpu = test_context.nparray_from_context_array(particles.beta0) + + for i_test, z_test in enumerate(zeta_cpu): + expected[i_test] += (x_cpu[i_test]**comp.test_exponents[0] * + y_cpu[i_test]**comp.test_exponents[1] * + np.dot(x_cpu**comp.source_exponents[0] * + y_cpu**comp.source_exponents[1], + comp.function_vs_zeta(z_test - zeta_cpu, + beta0=beta0_cpu[0], dzeta=1e-12)) * scale) xo.assert_allclose(getattr(particles, dict_p_bef[kk][0]) - dict_p_bef[kk][1], @@ -180,12 +185,13 @@ def test_wake_kick_multi_bunch(test_context, kind): filling_scheme=filling_scheme, bunch_spacing_zeta=circumference/h_bunch, bunch_selection=bunch_selection, - circumference=circumference + circumference=circumference, + _context=test_context ) line = xt.Line(elements=[wf]) - line.build_tracker() - line.particle_ref=xp.Particles(p0c=p0c) + line.build_tracker(_context=test_context) + line.particle_ref = xp.Particles(p0c=p0c) particles = xp.Particles(p0c=p0c, zeta=wf.slicer.zeta_centers.flatten(), @@ -219,6 +225,11 @@ def test_wake_kick_multi_bunch(test_context, kind): assert len(wf.components) == len(kind) + if hasattr(particles.zeta, 'get'): + particles_zeta = particles.zeta.get() + else: + particles_zeta = particles.zeta + for comp, kk in zip(wf.components, kind): if comp.plane == 'z': scale = -particles.q0**2 * qe**2 / ( @@ -230,15 +241,20 @@ def test_wake_kick_multi_bunch(test_context, kind): assert comp.source_exponents == kind_to_parameters[kk]['source_exponents'] assert comp.test_exponents == kind_to_parameters[kk]['test_exponents'] - expected = np.zeros_like(particles.zeta) + expected = np.zeros_like(particles_zeta) + + x_cpu = test_context.nparray_from_context_array(particles.x) + y_cpu = test_context.nparray_from_context_array(particles.y) + zeta_cpu = test_context.nparray_from_context_array(particles.zeta) + beta0_cpu = test_context.nparray_from_context_array(particles.beta0) - for i_test, z_test in enumerate(particles.zeta): - expected[i_test] += (particles.x[i_test]**comp.test_exponents[0] * - particles.y[i_test]**comp.test_exponents[1] * - np.dot(particles.x**comp.source_exponents[0] * - particles.y**comp.source_exponents[1], - comp.function_vs_zeta(z_test - particles.zeta, - beta0=particles.beta0[0], + for i_test, z_test in enumerate(zeta_cpu): + expected[i_test] += (x_cpu[i_test]**comp.test_exponents[0] * + y_cpu[i_test]**comp.test_exponents[1] * + np.dot(x_cpu**comp.source_exponents[0] * + y_cpu**comp.source_exponents[1], + comp.function_vs_zeta(z_test - zeta_cpu, + beta0=beta0_cpu[0], dzeta=1e-12)) * scale) xo.assert_allclose(getattr(particles, dict_p_bef[kk][0]) - dict_p_bef[kk][1], @@ -274,11 +290,13 @@ def test_beam_loading_theorem(test_context): delta_bef = particles.delta.copy() wf = xw.WakeResonator(kind='longitudinal', r=1e8, q=1e7, f_r=1e3) - wf.configure_for_tracking(zeta_range=zeta_range, num_slices=n_slices) + wf.configure_for_tracking(zeta_range=zeta_range, num_slices=n_slices, + _context=test_context) line = xt.Line(elements=[wf], element_names=['wf']) - line.build_tracker() + line.build_tracker(_context=test_context) + line.particle_ref = xp.Particles(p0c=p0c, _context=test_context) line.track(particles, num_turns=1) scale = -particles.q0**2 * qe**2 / (qe * particles.p0c[0] * particles.beta0[0] @@ -323,7 +341,9 @@ def test_wake_kick_multiturn(test_context, kind): wf = xw.WakeResonator(kind=kind, r=1e8, q=1e7, f_r=1e3) wf.configure_for_tracking(zeta_range=zeta_range, num_slices=n_slices, - num_turns=num_turns_wake, circumference=circumference) + num_turns=num_turns_wake, + circumference=circumference, + _context=test_context) i_source = -10 @@ -356,46 +376,53 @@ def test_wake_kick_multiturn(test_context, kind): line = xt.Line(elements=[wf], element_names=['wf']) - line.build_tracker() + line.build_tracker(_context=test_context) + line.particle_ref = xp.Particles(p0c=p0c, _context=test_context) line.track(particles, num_turns=3) scale = particles.q0**2 * qe**2 / ( - particles.p0c[0] * particles.beta0[0]* qe) * particles.weight[0] + p0c * particles.beta0[0]* qe) * particles.weight[0] + + x_cpu = test_context.nparray_from_context_array(particles.x) + y_cpu = test_context.nparray_from_context_array(particles.y) + zeta_cpu = test_context.nparray_from_context_array(particles.zeta) + beta0_cpu = test_context.nparray_from_context_array(particles.beta0) + weight_cpu = test_context.nparray_from_context_array(particles.weight) wake_kwargs = { - 'beta0': particles.beta0[0], + 'beta0': beta0_cpu[0], 'dzeta': 1e-20, } for comp, kk in zip(wf.components, kind): if comp.plane == 'z': scale = -particles.q0**2 * qe**2 / ( - particles.p0c[0] * particles.beta0[0]* qe) * particles.weight[0] + p0c * beta0_cpu[0]* qe) * weight_cpu[0] else: scale = particles.q0**2 * qe**2 / ( - particles.p0c[0] * particles.beta0[0]* qe) * particles.weight[0] + p0c * beta0_cpu[0]* qe) * weight_cpu[0] assert comp.plane == kind_to_parameters[kk]['plane'] assert comp.source_exponents == kind_to_parameters[kk]['source_exponents'] assert comp.test_exponents == kind_to_parameters[kk]['test_exponents'] - expected = np.zeros_like(particles.zeta) + expected = np.zeros_like(zeta_cpu) - for i_test, z_test in enumerate(particles.zeta): + for i_test, z_test in enumerate(zeta_cpu): expected[i_test] += ( - particles.x[i_test]**comp.test_exponents[0] * - particles.y[i_test]**comp.test_exponents[1] * - np.dot(particles.x**comp.source_exponents[0] * - particles.y**comp.source_exponents[1], - comp.function_vs_zeta(z_test - particles.zeta, + x_cpu[i_test]**comp.test_exponents[0] * + y_cpu[i_test]**comp.test_exponents[1] * + np.dot(x_cpu**comp.source_exponents[0] * + y_cpu**comp.source_exponents[1], + comp.function_vs_zeta(z_test - zeta_cpu, **wake_kwargs)) * scale * (num_turns_wake + 1)) expected[i_test] += ( - particles.x[i_test]**comp.test_exponents[0] * - particles.y[i_test]**comp.test_exponents[1] * - np.dot(particles.x**comp.source_exponents[0] * - particles.y**comp.source_exponents[1], - comp.function_vs_zeta(z_test - particles.zeta - + x_cpu[i_test]**comp.test_exponents[0] * + y_cpu[i_test]**comp.test_exponents[1] * + np.dot(x_cpu**comp.source_exponents[0] * + y_cpu**comp.source_exponents[1], + comp.function_vs_zeta(z_test - zeta_cpu - circumference, **wake_kwargs)) * scale * num_turns_wake) @@ -428,7 +455,9 @@ def test_wake_kick_multiturn_reset(test_context, kind): wf = xw.WakeResonator(kind=kind, r=1e8, q=1e7, f_r=1e3) wf.configure_for_tracking(zeta_range=zeta_range, num_slices=n_slices, - num_turns=num_turns_wake, circumference=circumference) + num_turns=num_turns_wake, + circumference=circumference, + _context=test_context) i_source = -10 @@ -449,7 +478,8 @@ def test_wake_kick_multiturn_reset(test_context, kind): line = xt.Line(elements=[wf], element_names=['wf']) - line.build_tracker() + line.build_tracker(_context=test_context) + line.particle_ref = xp.Particles(p0c=p0c, _context=test_context) line.track(particles, num_turns=num_turns_wake) particles.weight *= 0 @@ -573,7 +603,8 @@ def test_wake_kick_multibunch_pipeline(test_context, kind): filling_scheme=filling_scheme, bunch_selection=bunch_selection_0, num_turns=n_turns_wake, - circumference=circumference + circumference=circumference, + _context=test_context ) wf_0._wake_tracker.init_pipeline(pipeline_manager=pipeline_manager, element_name='wake', @@ -586,11 +617,12 @@ def test_wake_kick_multibunch_pipeline(test_context, kind): filling_scheme=filling_scheme, bunch_selection=bunch_selection_1, num_turns=n_turns_wake, - circumference=circumference + circumference=circumference, + _context=test_context ) wf_1._wake_tracker.init_pipeline(pipeline_manager=pipeline_manager, - element_name='wake', - partner_names=['b0']) + element_name='wake', + partner_names=['b0']) dict_p_bef_0 = {} dict_p_bef_1 = {} @@ -612,8 +644,10 @@ def test_wake_kick_multibunch_pipeline(test_context, kind): line_0 = xt.Line(elements=[wf_0]) line_1 = xt.Line(elements=[wf_1]) print('Initialising multitracker') - line_0.build_tracker() - line_1.build_tracker() + line_0.build_tracker(_context=test_context) + line_1.build_tracker(_context=test_context) + line_0.particle_ref = xp.Particles(p0c=p0c, _context=test_context) + line_1.particle_ref = xp.Particles(p0c=p0c, _context=test_context) multitracker = xt.PipelineMultiTracker( branches=[xt.PipelineBranch(line=line_0, particles=particles_0), xt.PipelineBranch(line=line_1, particles=particles_1)]) @@ -623,6 +657,14 @@ def test_wake_kick_multibunch_pipeline(test_context, kind): print('loading test data') particles_tot = xt.Particles.merge([particles_0, particles_1]) + x_cpu = test_context.nparray_from_context_array(particles.x) + y_cpu = test_context.nparray_from_context_array(particles.y) + zeta_cpu = test_context.nparray_from_context_array(particles.zeta) + x_tot_cpu = test_context.nparray_from_context_array(particles_tot.x) + y_tot_cpu = test_context.nparray_from_context_array(particles_tot.y) + zeta_tot_cpu = test_context.nparray_from_context_array(particles.zeta) + beta0_cpu = test_context.nparray_from_context_array(particles.beta0) + for particles, dict_p_bef, wf in [(particles_0, dict_p_bef_0, wf_0), (particles_1, dict_p_bef_1, wf_1)]: @@ -631,24 +673,24 @@ def test_wake_kick_multibunch_pipeline(test_context, kind): for comp, kk in zip(wf.components, kind): if comp.plane == 'z': scale = -particles.q0**2 * qe**2 / ( - particles.p0c[0] * particles.beta0[0]* qe) * particles.weight[0] + p0c * beta0_cpu * qe) * particles.weight[0] else: scale = particles.q0**2 * qe**2 / ( - particles.p0c[0] * particles.beta0[0]* qe) * particles.weight[0] + p0c * beta0_cpu * qe) * particles.weight[0] assert comp.plane == kind_to_parameters[kk]['plane'] assert comp.source_exponents == kind_to_parameters[kk]['source_exponents'] assert comp.test_exponents == kind_to_parameters[kk]['test_exponents'] - expected = np.zeros_like(particles.zeta) + expected = np.zeros_like(zeta_cpu) - for i_test, z_test in enumerate(particles.zeta): - expected[i_test] += (particles.x[i_test]**comp.test_exponents[0] * - particles.y[i_test]**comp.test_exponents[1] * - np.dot(particles_tot.x**comp.source_exponents[0] * - particles_tot.y**comp.source_exponents[1], - comp.function_vs_zeta(z_test - particles_tot.zeta, - beta0=particles_tot.beta0[0], - dzeta=1e-12)) * scale) + for i_test, z_test in enumerate(zeta_cpu): + expected[i_test] += (x_cpu[i_test]**comp.test_exponents[0] * + y_cpu[i_test]**comp.test_exponents[1] * + np.dot(x_tot_cpu**comp.source_exponents[0] * + y_tot_cpu**comp.source_exponents[1], + comp.function_vs_zeta(z_test - zeta_tot_cpu, + beta0=beta0_cpu[0], + dzeta=1e-12)) * scale) xo.assert_allclose(getattr(particles, dict_p_bef[kk][0]) - dict_p_bef[kk][1], expected, rtol=1e-4, atol=1e-20) From 101345da16d53b5f20d6b4e1c2adb1eef5091451 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Wed, 15 Jan 2025 10:27:12 +0100 Subject: [PATCH 4/4] add context to transverse damper --- xwakes/beam_elements/transverse_damper.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/xwakes/beam_elements/transverse_damper.py b/xwakes/beam_elements/transverse_damper.py index 476bc710..e763f0c7 100644 --- a/xwakes/beam_elements/transverse_damper.py +++ b/xwakes/beam_elements/transverse_damper.py @@ -50,13 +50,16 @@ def __init__(self, gain_x, gain_y, zeta_range, num_slices, 'py': gain_y, } + self.xoinitialize(**kwargs) + self.slicer = xf.UniformBinSlicer( filling_scheme=filling_scheme, bunch_selection=bunch_selection, zeta_range=zeta_range, num_slices=num_slices, bunch_spacing_zeta=bunch_spacing_zeta, - moments=['px', 'py'] + moments=['px', 'py'], + _context=self.context ) if filling_scheme is not None: @@ -74,7 +77,8 @@ def __init__(self, gain_x, gain_y, zeta_range, num_slices, bunch_spacing_zeta=bunch_spacing_zeta, num_periods=num_periods, num_turns=1, - circumference=circumference + circumference=circumference, + _context=self.context ) def _reconfigure_for_parallel(self, n_procs, my_rank):