From 06a98cd5b1ccfd8b3841ecf9174cd95140aaf6ae Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Tue, 5 Nov 2024 14:33:45 +0100 Subject: [PATCH 01/13] move compressed profile data on gpu --- xfields/slicers/compressed_profile.py | 43 +++++++++++++++++---------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/xfields/slicers/compressed_profile.py b/xfields/slicers/compressed_profile.py index d3ef3188..a909794a 100644 --- a/xfields/slicers/compressed_profile.py +++ b/xfields/slicers/compressed_profile.py @@ -37,6 +37,7 @@ class CompressedProfile(xt.BeamElement): '_N_aux': xo.Int64, '_N_S': xo.Int64, 'num_turns': xo.Int64, + 'data': xo.Float64[:,:,:], } pkg_root = xf.general._pkg_root @@ -65,10 +66,13 @@ def __init__(self, num_turns=1, num_targets=None, num_slices_target=None, - circumference=None + circumference=None, + **kwargs, ): - self.xoinitialize() + if '_xobject' in kwargs.keys(): + self.xoinitialize(**kwargs) + return if num_turns > 1: assert circumference is not None, ( @@ -84,7 +88,7 @@ def __init__(self, self._N_1 = num_slices # N_1 in the self._z_P = bunch_spacing_zeta # P in the paper - self._N_S = num_periods # N_S in the paper + _N_S = num_periods # N_S in the paper if num_slices_target is not None: self._N_2 = num_slices_target @@ -94,24 +98,27 @@ def __init__(self, if num_targets is not None: self._N_T = num_targets else: - self._N_T = self._N_S - - self.num_turns = num_turns + self._N_T = _N_S self._BB = 1 # B in the paper # (for now we assume that B=0 is the first bunch in time # and the last one in zeta) - self._AA = self._BB - self._N_S + self._AA = self._BB - _N_S - self._N_aux = self._N_1 + self._N_2 # n_aux in the paper + _N_aux = self._N_1 + self._N_2 # n_aux in the paper # Compute m_aux - self._M_aux = (self._N_S + - self._N_T - 1) * self._N_aux # m_aux in the paper + self._M_aux = (_N_S + + self._N_T - 1) * _N_aux # m_aux in the paper self.moments_names = moments - self.data = np.zeros( - (len(moments), self.num_turns, self._M_aux), dtype=np.float64) + + data = np.zeros( + (len(moments), num_turns, self._M_aux), dtype=np.float64) + + self.xoinitialize(_N_S=_N_S, _N_aux=_N_aux, num_turns=num_turns, + data=data, **kwargs) + def __getitem__(self, key): assert isinstance(key, str), 'other modes not supported yet' @@ -168,8 +175,12 @@ def set_moments(self, i_source, i_turn, moments): i_start_in_moments_data = (self._N_S - i_source - 1) * self._N_aux i_end_in_moments_data = i_start_in_moments_data + self._N_1 + #if hasattr(vv, 'get'): + # self.data[i_moment, i_turn, + # i_start_in_moments_data:i_end_in_moments_data] = vv.get() + #else: self.data[i_moment, i_turn, - i_start_in_moments_data:i_end_in_moments_data] = vv + i_start_in_moments_data:i_end_in_moments_data] = vv def get_moment_profile(self, moment_name, i_turn): """ @@ -190,8 +201,8 @@ def get_moment_profile(self, moment_name, i_turn): The moment profile """ - z_out = np.zeros(self._N_S * self._N_1) - moment_out = np.zeros(self._N_S * self._N_1) + z_out = self._arr2ctx(np.zeros(self._N_S * self._N_1)) + moment_out = self._arr2ctx(np.zeros(self._N_S * self._N_1)) i_moment = self.moments_names.index(moment_name) _z_P = self._z_P or 0 for i_source in range(self._N_S): @@ -199,7 +210,7 @@ def get_moment_profile(self, moment_name, i_turn): i_end_out = i_start_out + self._N_1 z_out[i_start_out:i_end_out] = ( self._z_a + self.dz / 2 - - i_source * _z_P + self.dz * np.arange(self._N_1)) + - i_source * _z_P + self.dz * self._arr2ctx(np.arange(self._N_1))) i_start_in_moments_data = (self._N_S - i_source - 1) * self._N_aux i_end_in_moments_data = i_start_in_moments_data + self._N_1 From f7c29f88a81c3bad24d16b8ed791df7f265f6c33 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Tue, 5 Nov 2024 14:34:56 +0100 Subject: [PATCH 02/13] add gpu slicer examples --- examples/slicer/002b_example_slicer_gpu.py | 145 +++++++++++++++++++++ examples/slicer/005b_moment_profile_gpu.py | 89 +++++++++++++ 2 files changed, 234 insertions(+) create mode 100644 examples/slicer/002b_example_slicer_gpu.py create mode 100644 examples/slicer/005b_moment_profile_gpu.py diff --git a/examples/slicer/002b_example_slicer_gpu.py b/examples/slicer/002b_example_slicer_gpu.py new file mode 100644 index 00000000..499ce56b --- /dev/null +++ b/examples/slicer/002b_example_slicer_gpu.py @@ -0,0 +1,145 @@ +import xtrack as xt +import xpart as xp +import xfields as xf + +import numpy as np + +import xobjects as xo +context = xo.ContextCupy(device=3) +#context = xo.ContextCpu() + +line = xt.Line.from_json( + '../../../xtrack/test_data/hllhc15_thick/lhc_thick_with_knobs.json') +line.build_tracker(_context=context) + +line.vars['vrf400'] = 16 + +p = xp.generate_matched_gaussian_bunch( + line=line, + num_particles=int(1e5), + nemitt_x=2e-6, + nemitt_y=2.5e-6, + sigma_z=0.07) + +slicer = xf.UniformBinSlicer(zeta_range=(-999, +999), num_slices=1, _context=context) + +slicer.slice(p) + +cov_matrix = np.zeros((6, 6)) + +for ii, vii in enumerate(['x', 'px', 'y', 'py', 'zeta', 'delta']): + for jj, vjj in enumerate(['x', 'px', 'y', 'py', 'zeta', 'delta']): + if ii <= jj: + cov_matrix[ii, jj] = slicer.cov(vii, vjj)[0] + else: + cov_matrix[ii, jj] = cov_matrix[jj, ii] + +assert np.allclose(np.sqrt(cov_matrix[0, 0]), np.std(p.x), atol=0, rtol=1e-6) + +tw = line.twiss() +Sig = cov_matrix + +S = xt.linear_normal_form.S + +eival, eivec = np.linalg.eig(Sig @ S) + +# Keep only one from each complex conjugate pair +eival_list = [] +eivec_list = [] +for ii in range(6): + if ii == 0: + eival_list.append(eival[ii]) + eivec_list.append(eivec[:, ii]) + continue + found_conj = False + for jj in range(len(eival_list)): + if np.allclose(eival[ii], np.conj(eival_list[jj]), rtol=0, atol=1e-14): + found_conj = True + break + if not found_conj: + eival_list.append(eival[ii]) + eivec_list.append(eivec[:, ii]) + +assert len(eival_list) == 3 + +# Find longitudinal mode +norm = np.linalg.norm +i_long = 0 +if norm(eivec_list[1][5:6]) > norm(eivec_list[i_long][5:6]): + i_long = 1 +if norm(eivec_list[2][5:6]) > norm(eivec_list[i_long][5:6]): + i_long = 2 + +eival_zeta = eival_list[i_long] +eivec_zeta = eivec_list[i_long] + +# Find vertical mode +eival_list.pop(i_long) +eivec_list.pop(i_long) +if norm(eivec_list[0][3:4]) > norm(eivec_list[1][3:4]): + i_vert = 0 +else: + i_vert = 1 + +eival_y = eival_list[i_vert] +eivec_y = eivec_list[i_vert] + +# Find horizontal mode +eival_list.pop(i_vert) +eivec_list.pop(i_vert) +eival_x = eival_list[0] +eivec_x = eivec_list[0] + +nemitt_x = eival_x.imag * tw.gamma0 * tw.beta0 +nemitt_y = eival_y.imag * tw.gamma0 * tw.beta0 +nemitt_zeta = eival_zeta.imag * tw.gamma0 * tw.beta0 + + +print(f'{nemitt_x=} {nemitt_y=} {nemitt_zeta=}') +print('\n') + +dummy_lam = np.diag([ + np.exp(-1j*np.pi/3), np.exp(+1j*np.pi/3), + np.exp(-1j*np.pi/4), np.exp(+1j*np.pi/4), + np.exp(-1j*np.pi/5), np.exp(+1j*np.pi/5), +]) + +dummy_R = eivec @ dummy_lam @ np.linalg.inv(eivec) + +dummy_line = xt.Line(elements=[xt.Drift(length=1e-12)]) +p_dummy = line.build_particles(x=0) + +tw_from_sigmas = dummy_line.twiss( + particle_on_co=p_dummy, + R_matrix=dummy_R, + compute_chromatic_properties=False) + +print('betx/bety') +print(f'betx (from line) = {tw.betx[0]}') +print(f'betx (from sigmas) = {tw_from_sigmas.betx[0]}') +print(f'bety (from line) = {tw.bety[0]}') +print(f'bety (from sigmas) = {tw_from_sigmas.bety[0]}') +print() +print('alfx/alfy') +print(f'alfx (from line) = {tw.alfx[0]}') +print(f'alfx (from sigmas) = {tw_from_sigmas.alfx[0]}') +print(f'alfy (from line) = {tw.alfy[0]}') +print(f'alfy (from sigmas) = {tw_from_sigmas.alfy[0]}') +print() +print('dx/dy') +print(f'dx (from line) = {tw.dx[0]}') +print(f'dx (from sigmas) = {tw_from_sigmas.dx[0]}') +print(f'dy (from line) = {tw.dy[0]}') +print(f'dy (from sigmas) = {tw_from_sigmas.dy[0]}') +print() +print('coupled betas') +print(f'betx2 (from line) = {tw.betx2[0]}') +print(f'betx2 (from sigmas) = {tw_from_sigmas.betx2[0]}') +print(f'bety2 (from line) = {tw.bety2[0]}') +print(f'bety2 (from sigmas) = {tw_from_sigmas.bety2[0]}') + + +# - make it work with one slice +# - populate a Sigma matrix +# - p.get_sigma_matrix() +# - p.get_statistical_ \ No newline at end of file diff --git a/examples/slicer/005b_moment_profile_gpu.py b/examples/slicer/005b_moment_profile_gpu.py new file mode 100644 index 00000000..50ebd868 --- /dev/null +++ b/examples/slicer/005b_moment_profile_gpu.py @@ -0,0 +1,89 @@ +import xfields as xf +import xtrack as xt +import xobjects as xo +import xwakes as xw +from xfields.slicers.compressed_profile import CompressedProfile + +import numpy as np +import matplotlib.pyplot as plt + +import xobjects as xo +context = xo.ContextCupy(device=3) + +# base coordinates +zeta_coord = np.random.normal(loc=5e-3, scale=5e-2, + size=100_000) +bunch_spacing_zeta = 10 +num_bunches = 3 + +# n_bunches bunches with n_macroparticles each with different coordinates +particles = xt.Particles( + zeta=np.concatenate([zeta_coord*(bid + 1) - bunch_spacing_zeta*bid + for bid in range(num_bunches)]), + _context=context +) + +# dummy filling scheme +filling_scheme = np.ones(num_bunches, dtype=int) +bunch_selection = np.arange(num_bunches, dtype=int) + +zeta_range = (-1, 1) + +# we create a compressed profile object and a slicer object. The slicer object +# is used to compute the moments of the beam at the slices, which are then +# inserted in the profile +moments = ['x', 'y', 'px', 'py', 'num_particles'] + +compressed_profile = CompressedProfile( + moments=moments, + zeta_range=zeta_range, + num_slices=20, + bunch_spacing_zeta=bunch_spacing_zeta, + num_periods=num_bunches, + num_turns=1, + circumference=bunch_spacing_zeta*len(filling_scheme), + _context=context) + + +slicer = xf.UniformBinSlicer( + zeta_range=zeta_range, + num_slices=compressed_profile.num_slices, + filling_scheme=filling_scheme, + bunch_selection=bunch_selection, + bunch_spacing_zeta=bunch_spacing_zeta, + moments='all', + _context=context +) + +print('a') + +slicer.track(particles) + +print('b') + +for i_bunch in range(num_bunches): + dict_moments = { + 'num_particles': slicer.num_particles[i_bunch, :], + } + + for moment in moments: + if moment == 'num_particles': + continue + dict_moments[moment] = slicer.mean(moment)[i_bunch, :] + + compressed_profile.set_moments(i_turn=0, i_source=bunch_selection[i_bunch], + moments=dict_moments) + +print('c') + +long_profile = compressed_profile.get_moment_profile( + moment_name='num_particles', i_turn=0) + +long_profile_x = long_profile[0].get() +long_profile_y = long_profile[1].get() + +plt.figure(321) +plt.plot(long_profile_x, long_profile_y/np.sum(long_profile_y)) +plt.xlabel('zeta [m]') +plt.ylabel('longitudinal profile [a.u.]') +plt.show() From f7fc6432b0ca0de868ff1ab1583913032ff9727b Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Fri, 8 Nov 2024 15:44:21 +0100 Subject: [PATCH 03/13] more work --- ...ample_slicer.py => 002a_example_slicer.py} | 0 ...ment_profile.py => 005a_moment_profile.py} | 0 ...ke.py => 005c_moment_profile_from_wake.py} | 0 xfields/beam_elements/element_with_slicer.py | 13 +++- .../beam_elements/waketracker/convolution.py | 73 ++++++++++++++----- .../beam_elements/waketracker/waketracker.py | 12 ++- xfields/headers/compressed_profile.h | 2 + xfields/slicers/compressed_profile.py | 6 +- 8 files changed, 78 insertions(+), 28 deletions(-) rename examples/slicer/{002_example_slicer.py => 002a_example_slicer.py} (100%) rename examples/slicer/{005_moment_profile.py => 005a_moment_profile.py} (100%) rename examples/slicer/{005_moment_profile_from_wake.py => 005c_moment_profile_from_wake.py} (100%) diff --git a/examples/slicer/002_example_slicer.py b/examples/slicer/002a_example_slicer.py similarity index 100% rename from examples/slicer/002_example_slicer.py rename to examples/slicer/002a_example_slicer.py diff --git a/examples/slicer/005_moment_profile.py b/examples/slicer/005a_moment_profile.py similarity index 100% rename from examples/slicer/005_moment_profile.py rename to examples/slicer/005a_moment_profile.py diff --git a/examples/slicer/005_moment_profile_from_wake.py b/examples/slicer/005c_moment_profile_from_wake.py similarity index 100% rename from examples/slicer/005_moment_profile_from_wake.py rename to examples/slicer/005c_moment_profile_from_wake.py diff --git a/xfields/beam_elements/element_with_slicer.py b/xfields/beam_elements/element_with_slicer.py index 9943347c..e3106780 100644 --- a/xfields/beam_elements/element_with_slicer.py +++ b/xfields/beam_elements/element_with_slicer.py @@ -5,7 +5,7 @@ from xfields.slicers.compressed_profile import CompressedProfile -class ElementWithSlicer: +class ElementWithSlicer(xt.BeamElement): """ Base class for elements with a slicer. @@ -47,7 +47,10 @@ def __init__(self, bunch_selection=None, num_turns=1, circumference=None, - with_compressed_profile=False): + with_compressed_profile=False, + **kwargs): + + self.xoinitialize(**kwargs) self.with_compressed_profile = with_compressed_profile self.pipeline_manager = None @@ -88,7 +91,8 @@ def init_slicer(self, zeta_range, num_slices, filling_scheme, filling_scheme=filling_scheme, bunch_selection=bunch_selection, bunch_spacing_zeta=bunch_spacing_zeta, - moments=slicer_moments + moments=slicer_moments, + _context=self._context ) else: self.zeta_range = None @@ -115,7 +119,8 @@ def _initialize_moments( bunch_spacing_zeta=bunch_spacing_zeta, num_periods=num_periods, num_turns=num_turns, - circumference=circumference) + circumference=circumference, + _context=self.context) def init_pipeline(self, pipeline_manager, element_name, partner_names): self.pipeline_manager = pipeline_manager diff --git a/xfields/beam_elements/waketracker/convolution.py b/xfields/beam_elements/waketracker/convolution.py index c9a3b412..7075313f 100644 --- a/xfields/beam_elements/waketracker/convolution.py +++ b/xfields/beam_elements/waketracker/convolution.py @@ -1,12 +1,28 @@ import numpy as np +import xtrack as xt +import xobjects as xo + +try: + import cupy + cupy_imported = True +except ImportError: + cupy_imported = False + from scipy.constants import e as qe from typing import Tuple -class _ConvData: +class _ConvData(xt.BeamElement): + def __init__(self, component, waketracker=None, _flatten=False, log_moments=None, + **kwargs): - def __init__(self, component, waketracker=None, _flatten=False, log_moments=None): + self.xoinitialize(**kwargs) + + # for now we do not support Pyopencl to avoid complications with the + # rfft and irfft below + if type(self.context) == xo.ContextPyopencl: + raise NotImplementedError('Pyopencl not implemented yet') self._flatten = _flatten self.component = component @@ -34,6 +50,28 @@ def __init__(self, component, waketracker=None, _flatten=False, log_moments=None if source_exponents[1] > 1: raise NotImplementedError('Higher order moments not implemented yet') + def my_rfft(self, data, **kwargs): + if type(self.context) == xo.ContextCpu: + return np.fft.rfft(data, **kwargs) + elif type(self.context) == xo.ContextCupy: + if cupy_imported: + return cupy.fft.rfft(data, **kwargs) + else: + raise ImportError('Cupy not installed') + else: + raise NotImplementedError('Waketacker implemented only for CPU and Cupy') + + def my_irfft(self, data, **kwargs): + if type(self.context) == xo.ContextCpu: + return np.fft.irfft(data, **kwargs) + elif type(self.context) == xo.ContextCupy: + if cupy_imported: + return cupy.fft.irfft(data, **kwargs) + else: + raise ImportError('Cupy not installed') + else: + raise NotImplementedError('Waketacker implemented only for CPU and Cupy') + def _initialize_conv_data(self, _flatten=False, moments_data=None, beta0=None): assert moments_data is not None if not _flatten: @@ -50,22 +88,23 @@ def _initialize_conv_data(self, _flatten=False, moments_data=None, beta0=None): self._DD = self._BB # Build wake matrix - self.z_wake = _build_z_wake(moments_data._z_a, moments_data._z_b, + self.z_wake = self._arr2ctx(_build_z_wake(moments_data._z_a, moments_data._z_b, moments_data.num_turns, moments_data._N_aux, moments_data._M_aux, moments_data.circumference, moments_data.dz, self._AA, self._BB, self._CC, self._DD, - moments_data._z_P) + moments_data._z_P)) assert beta0 is not None - - self.G_aux = self.component.function_vs_zeta( - zeta=self.z_wake, beta0=beta0, dzeta=moments_data.dz) + # here below I had to add float() to beta0 because when using Cupy + # context particles.beta0[0] turns out to be a 0d array. To be checked + self.G_aux = self._arr2ctx(self.component.function_vs_zeta( + zeta=self.z_wake, beta0=float(beta0), dzeta=moments_data.dz)) # only positive frequencies because we are using rfft - phase_term = np.exp( + phase_term = self._arr2ctx(np.exp( 1j * 2 * np.pi * np.arange(self._M_aux//2 + 1) * - ((self._N_S - 1) * self._N_aux + self._N_1) / self._M_aux) + ((self._N_S - 1) * self._N_aux + self._N_1) / self._M_aux)) else: raise NotImplementedError('Flattened wakes are not implemented yet') @@ -100,8 +139,8 @@ def _initialize_conv_data(self, _flatten=False, moments_data=None, beta0=None): ((self._N_S_flatten - 1) * self._N_aux + self._N_1) / self._M_aux_flatten) - self._G_hat_dephased = phase_term * np.fft.rfft(self.G_aux, axis=1) - self._G_aux_shifted = np.fft.irfft(self._G_hat_dephased, axis=1) + self._G_hat_dephased = phase_term * self.my_rfft(self.G_aux, axis=1) + self._G_aux_shifted = self.my_irfft(self._G_hat_dephased, axis=1) def track(self, particles, i_slot_particles, i_slice_particles, moments_data): @@ -144,15 +183,15 @@ def _compute_convolution(self, moment_names, moments_data): if isinstance(moment_names, str): moment_names = [moment_names] - rho_aux = np.ones(shape=moments_data['result'].shape, - dtype=np.float64) + rho_aux = self._arr2ctx(np.ones(shape=moments_data['result'].shape, + dtype=np.float64)) for nn in moment_names: rho_aux *= moments_data[nn] if not self._flatten: - rho_hat = np.fft.rfft(rho_aux, axis=1) - res = np.fft.irfft(rho_hat * self._G_hat_dephased, axis=1) + rho_hat = self.my_rfft(rho_aux, axis=1) + res = self.my_irfft(rho_hat * self._G_hat_dephased, axis=1) else: rho_aux_flatten = np.zeros((1, self._M_aux_flatten), dtype=np.float64) @@ -162,8 +201,8 @@ def _compute_convolution(self, moment_names, moments_data): 0, tt * _N_aux_turn: (tt + 1) * _N_aux_turn] = \ rho_aux[tt, :_N_aux_turn] - rho_hat_flatten = np.fft.rfft(rho_aux_flatten, axis=1) - res_flatten = np.fft.irfft( + rho_hat_flatten = self.my_rfft(rho_aux_flatten, axis=1) + res_flatten = self.my_irfft( rho_hat_flatten * self._G_hat_dephased, axis=1).real self._res_flatten_fft = res_flatten # for debugging diff --git a/xfields/beam_elements/waketracker/waketracker.py b/xfields/beam_elements/waketracker/waketracker.py index da413a35..e21a6688 100644 --- a/xfields/beam_elements/waketracker/waketracker.py +++ b/xfields/beam_elements/waketracker/waketracker.py @@ -53,7 +53,10 @@ def __init__(self, components, num_turns=1, circumference=None, log_moments=None, - _flatten=False): + _flatten=False, + **kwargs): + + self.xoinitialize(**kwargs) self.components = components self.pipeline_manager = None @@ -75,8 +78,8 @@ def __init__(self, components, bunch_selection=bunch_selection, num_turns=num_turns, circumference=circumference, - with_compressed_profile=True - ) + with_compressed_profile=True, + _context=self.context) self._initialize_moments( zeta_range=zeta_range, # These are [a, b] in the paper @@ -114,7 +117,8 @@ def track(self, particles): continue cc._conv_data = _ConvData(component=cc, waketracker=self, - _flatten=self._flatten) + _flatten=self._flatten, + _context=self._context) cc._conv_data._initialize_conv_data(_flatten=self._flatten, moments_data=self.moments_data, beta0=beta0) diff --git a/xfields/headers/compressed_profile.h b/xfields/headers/compressed_profile.h index b2a7f05f..86cea534 100644 --- a/xfields/headers/compressed_profile.h +++ b/xfields/headers/compressed_profile.h @@ -1,11 +1,13 @@ #ifndef XFIELDS_COPRESSED_PROFILE_H #define XFIELDS_COPRESSED_PROFILE_H +/*gpufun*/ void CompressedProfile_track_local_particle(CompressedProfileData el, LocalParticle *part0){ } +/*gpufun*/ void CompressedProfile_interp_result( CompressedProfileData el, LocalParticle *part0, int64_t data_shape_0, diff --git a/xfields/slicers/compressed_profile.py b/xfields/slicers/compressed_profile.py index a909794a..d38079d4 100644 --- a/xfields/slicers/compressed_profile.py +++ b/xfields/slicers/compressed_profile.py @@ -156,9 +156,9 @@ def set_moments(self, i_source, i_turn, moments): A dictionary of the form {moment_name: moment_value} """ - - assert np.isscalar(i_source) - assert np.isscalar(i_turn) + #breakpoint() + #assert np.isscalar(i_source) + #assert np.isscalar(i_turn) assert i_source < self._N_S assert i_source >= 0 From 46d68ba28d219db71dbfe97e85b7dbbb9804e3cd Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 14 Nov 2024 11:42:20 +0100 Subject: [PATCH 04/13] initialize z_wake on cpu --- xfields/beam_elements/waketracker/convolution.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xfields/beam_elements/waketracker/convolution.py b/xfields/beam_elements/waketracker/convolution.py index 7075313f..c12b8272 100644 --- a/xfields/beam_elements/waketracker/convolution.py +++ b/xfields/beam_elements/waketracker/convolution.py @@ -88,13 +88,13 @@ def _initialize_conv_data(self, _flatten=False, moments_data=None, beta0=None): self._DD = self._BB # Build wake matrix - self.z_wake = self._arr2ctx(_build_z_wake(moments_data._z_a, moments_data._z_b, + self.z_wake = _build_z_wake(moments_data._z_a, moments_data._z_b, moments_data.num_turns, moments_data._N_aux, moments_data._M_aux, moments_data.circumference, moments_data.dz, self._AA, self._BB, self._CC, self._DD, - moments_data._z_P)) + moments_data._z_P) assert beta0 is not None # here below I had to add float() to beta0 because when using Cupy # context particles.beta0[0] turns out to be a 0d array. To be checked From 52f95c502bc3571040454dfae430a2c691de36c0 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 14 Nov 2024 11:49:12 +0100 Subject: [PATCH 05/13] fix issue with array access --- xfields/beam_elements/element_with_slicer.py | 4 ++-- xfields/slicers/compressed_profile.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/xfields/beam_elements/element_with_slicer.py b/xfields/beam_elements/element_with_slicer.py index e3106780..b6ac30a3 100644 --- a/xfields/beam_elements/element_with_slicer.py +++ b/xfields/beam_elements/element_with_slicer.py @@ -170,7 +170,7 @@ def _add_slicer_moments_to_moments_data(self, slicer): continue means[mm] = slicer.mean(mm) - for i_bunch_in_slicer, bunch_number in enumerate(slicer.bunch_selection): + for i_bunch_in_slicer, bunch_number in enumerate(slicer._xobject.bunch_selection): moments_bunch = {} for nn in means.keys(): moments_bunch[nn] = np.atleast_2d(means[nn])[i_bunch_in_slicer, :] @@ -180,7 +180,7 @@ def _add_slicer_moments_to_moments_data(self, slicer): self.moments_data.set_moments( moments=moments_bunch, i_turn=0, - i_source=slicer.filled_slots[bunch_number]) + i_source=slicer._xobject.filled_slots[bunch_number]) def _update_moments_for_new_turn(self, particles, _slice_result=None): if self.with_compressed_profile: diff --git a/xfields/slicers/compressed_profile.py b/xfields/slicers/compressed_profile.py index d38079d4..a909794a 100644 --- a/xfields/slicers/compressed_profile.py +++ b/xfields/slicers/compressed_profile.py @@ -156,9 +156,9 @@ def set_moments(self, i_source, i_turn, moments): A dictionary of the form {moment_name: moment_value} """ - #breakpoint() - #assert np.isscalar(i_source) - #assert np.isscalar(i_turn) + + assert np.isscalar(i_source) + assert np.isscalar(i_turn) assert i_source < self._N_S assert i_source >= 0 From 35ffa4127bd9140e9bcdc68b4bb318adffabaafe Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 14 Nov 2024 11:51:36 +0100 Subject: [PATCH 06/13] smart init of compressed profile data --- xfields/slicers/compressed_profile.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/xfields/slicers/compressed_profile.py b/xfields/slicers/compressed_profile.py index a909794a..fe76f646 100644 --- a/xfields/slicers/compressed_profile.py +++ b/xfields/slicers/compressed_profile.py @@ -113,11 +113,9 @@ def __init__(self, self.moments_names = moments - data = np.zeros( - (len(moments), num_turns, self._M_aux), dtype=np.float64) - self.xoinitialize(_N_S=_N_S, _N_aux=_N_aux, num_turns=num_turns, - data=data, **kwargs) + data=(len(moments), num_turns, self._M_aux), + **kwargs) def __getitem__(self, key): From afc28954de7b90b6a8c5db94448cf8045ce7eed6 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 14 Nov 2024 11:52:28 +0100 Subject: [PATCH 07/13] remove commented code --- xfields/slicers/compressed_profile.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/xfields/slicers/compressed_profile.py b/xfields/slicers/compressed_profile.py index fe76f646..ce814e7b 100644 --- a/xfields/slicers/compressed_profile.py +++ b/xfields/slicers/compressed_profile.py @@ -173,10 +173,6 @@ def set_moments(self, i_source, i_turn, moments): i_start_in_moments_data = (self._N_S - i_source - 1) * self._N_aux i_end_in_moments_data = i_start_in_moments_data + self._N_1 - #if hasattr(vv, 'get'): - # self.data[i_moment, i_turn, - # i_start_in_moments_data:i_end_in_moments_data] = vv.get() - #else: self.data[i_moment, i_turn, i_start_in_moments_data:i_end_in_moments_data] = vv From e87d11d31f2af730ff9f8f243c5b5aea9f138661 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 14 Nov 2024 12:00:29 +0100 Subject: [PATCH 08/13] fix init of rho_aux --- xfields/beam_elements/waketracker/convolution.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xfields/beam_elements/waketracker/convolution.py b/xfields/beam_elements/waketracker/convolution.py index c12b8272..d6868f9d 100644 --- a/xfields/beam_elements/waketracker/convolution.py +++ b/xfields/beam_elements/waketracker/convolution.py @@ -183,8 +183,9 @@ def _compute_convolution(self, moment_names, moments_data): if isinstance(moment_names, str): moment_names = [moment_names] - rho_aux = self._arr2ctx(np.ones(shape=moments_data['result'].shape, - dtype=np.float64)) + # we do zeros + 1 because there is no ones method in context + rho_aux = self.context.zeros(shape=moments_data['result'].shape, + dtype=np.float64) + 1 for nn in moment_names: rho_aux *= moments_data[nn] From 82c22654cc6919a80745466e7cc1d9def4ef5c66 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 14 Nov 2024 12:03:35 +0100 Subject: [PATCH 09/13] add another gpu example --- .../005d_moment_profile_from_wake_gpu.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 examples/slicer/005d_moment_profile_from_wake_gpu.py diff --git a/examples/slicer/005d_moment_profile_from_wake_gpu.py b/examples/slicer/005d_moment_profile_from_wake_gpu.py new file mode 100644 index 00000000..09ea20b5 --- /dev/null +++ b/examples/slicer/005d_moment_profile_from_wake_gpu.py @@ -0,0 +1,59 @@ +import xfields as xf +import xtrack as xt +import xobjects as xo +import xwakes as xw + +import numpy as np +import matplotlib.pyplot as plt + +import xobjects as xo +context = xo.ContextCupy(device=3) + +# base coordinates +zeta_coord = np.random.normal(loc=5e-3, scale=5e-2, + size=100_000) +bunch_spacing_zeta = 10 +num_bunches = 3 + +# n_bunches bunches with n_macroparticles each with different coordinates +particles = xt.Particles( + zeta=np.concatenate([zeta_coord*(bid + 1) - bunch_spacing_zeta*bid + for bid in range(num_bunches)]), + _context=context, +) +# dummy filling scheme +filling_scheme = np.ones(num_bunches, dtype=int) +bunch_selection = np.arange(num_bunches, dtype=int) + +# wake from which the compressed profile is taken +wf = xw.WakeResonator( + r=3e8, + q=1e7, + f_r=1e3, + kind='dipolar_x', +) + +wf.configure_for_tracking( + zeta_range=(-1, 1), + num_slices=20, + bunch_spacing_zeta=bunch_spacing_zeta, + filling_scheme=filling_scheme, + bunch_selection=bunch_selection, + num_turns=1, + circumference=bunch_spacing_zeta*len(filling_scheme), + _context=context, +) + +wf.track(particles) + +long_profile = wf._wake_tracker.moments_data.get_moment_profile( + moment_name='num_particles', i_turn=0) + +long_profile_x = long_profile[0].get() +long_profile_y = long_profile[1].get() + +plt.figure(321) +plt.plot(long_profile_x, long_profile_y/np.sum(long_profile_y)) +plt.xlabel('zeta [m]') +plt.ylabel('longitudinal profile [a.u.]') +plt.show() From d613030dbd07cce2c6ffc286675c5c920e72742c Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 14 Nov 2024 20:53:47 +0100 Subject: [PATCH 10/13] use nplike_lib --- xfields/beam_elements/waketracker/convolution.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/xfields/beam_elements/waketracker/convolution.py b/xfields/beam_elements/waketracker/convolution.py index d6868f9d..0209158d 100644 --- a/xfields/beam_elements/waketracker/convolution.py +++ b/xfields/beam_elements/waketracker/convolution.py @@ -183,9 +183,8 @@ def _compute_convolution(self, moment_names, moments_data): if isinstance(moment_names, str): moment_names = [moment_names] - # we do zeros + 1 because there is no ones method in context - rho_aux = self.context.zeros(shape=moments_data['result'].shape, - dtype=np.float64) + 1 + rho_aux = self.context.nplike_lib.zeros( + shape=moments_data['result'].shape, dtype=np.float64) for nn in moment_names: rho_aux *= moments_data[nn] From 7bdc77d06828e9afc93fa835f41f6e762c951ac3 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Thu, 14 Nov 2024 20:54:19 +0100 Subject: [PATCH 11/13] improve compressed profile test to work on GPU --- .../test_compressed_profile_interp_result.py | 40 ++++++++++--------- 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/tests/test_compressed_profile_interp_result.py b/tests/test_compressed_profile_interp_result.py index 5101960b..76234f7f 100644 --- a/tests/test_compressed_profile_interp_result.py +++ b/tests/test_compressed_profile_interp_result.py @@ -6,9 +6,10 @@ from xfields.slicers import CompressedProfile from xobjects.test_helpers import for_all_test_contexts import xtrack as xt +import xobjects as xo - -def test_compressed_profile_interp_result(): +@for_all_test_contexts +def test_compressed_profile_interp_result(test_context): # Machine settings moments = ['result'] zeta_range = (0, 1) @@ -25,30 +26,32 @@ def test_compressed_profile_interp_result(): bunch_spacing_zeta=bunch_spacing_zeta, num_periods=num_bunches, num_turns=num_turns, - circumference=circumference + circumference=circumference, + _context=test_context ) num_parts = 100 - interpolated_result = np.zeros(num_parts, dtype=float) - i_slice_particles = np.linspace(0, num_slices - 1, num_parts, dtype=int) + interpolated_result = test_context.zeros(num_parts, dtype=float) + i_slice_particles = test_context.nplike_lib.linspace(0, num_slices - 1, num_parts, dtype=int) - result_parts = np.zeros(num_parts) + result_parts = test_context.zeros(num_parts) def func(z, i): return z + circumference * i for i_turn in range(num_turns): moments = { - 'result': func(np.linspace(zeta_range[0] + dz / 2, + 'result': comp_prof._arr2ctx( + func(np.linspace(zeta_range[0] + dz / 2, zeta_range[1] - dz / 2, - num_slices), i_turn) + num_slices), i_turn)) } comp_prof.set_moments(moments=moments, i_turn=i_turn, i_source=0) - result_parts += func(i_slice_particles * dz + dz / 2, i_turn) + result_parts += comp_prof._arr2ctx(func(i_slice_particles * dz + dz / 2, i_turn)) particles = xt.Particles( mass0=xt.PROTON_MASS_EV, @@ -60,16 +63,17 @@ def func(z, i): zeta=np.zeros(num_parts), delta=np.zeros(num_parts), weight=np.ones(num_parts), + _context=test_context ) comp_prof._interp_result(particles=particles, - data_shape_0=comp_prof.data.shape[0], - data_shape_1=comp_prof.data.shape[1], - data_shape_2=comp_prof.data.shape[2], - data=comp_prof.data, - i_slot_particles=np.zeros(num_parts, dtype=int), - i_slice_particles=i_slice_particles, - out=interpolated_result - ) + data_shape_0=comp_prof.data.shape[0], + data_shape_1=comp_prof.data.shape[1], + data_shape_2=comp_prof.data.shape[2], + data=comp_prof.data, + i_slot_particles=test_context.nplike_lib.zeros(num_parts, dtype=int), + i_slice_particles=i_slice_particles, + out=interpolated_result + ) - assert np.allclose(result_parts, interpolated_result) + xo.assert_allclose(result_parts, interpolated_result) From 334752184328845eeffa7a897ae5f186393078c3 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Fri, 15 Nov 2024 10:49:49 +0100 Subject: [PATCH 12/13] small fix --- xfields/beam_elements/waketracker/convolution.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xfields/beam_elements/waketracker/convolution.py b/xfields/beam_elements/waketracker/convolution.py index 0209158d..49a67b01 100644 --- a/xfields/beam_elements/waketracker/convolution.py +++ b/xfields/beam_elements/waketracker/convolution.py @@ -183,7 +183,7 @@ def _compute_convolution(self, moment_names, moments_data): if isinstance(moment_names, str): moment_names = [moment_names] - rho_aux = self.context.nplike_lib.zeros( + rho_aux = self.context.nplike_lib.ones( shape=moments_data['result'].shape, dtype=np.float64) for nn in moment_names: From bdb9b2147d7610a240ff962b01403c719ca57451 Mon Sep 17 00:00:00 2001 From: Lorenzo Giacomel Date: Wed, 4 Dec 2024 10:36:29 +0100 Subject: [PATCH 13/13] improve doc --- .../beam_elements/waketracker/convolution.py | 21 +++++++++++++++++++ xfields/slicers/compressed_profile.py | 9 +++++++- xfields/slicers/uniform.py | 2 +- 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/xfields/beam_elements/waketracker/convolution.py b/xfields/beam_elements/waketracker/convolution.py index 49a67b01..998dd5c8 100644 --- a/xfields/beam_elements/waketracker/convolution.py +++ b/xfields/beam_elements/waketracker/convolution.py @@ -14,8 +14,29 @@ class _ConvData(xt.BeamElement): + ''' + This class is used to perform the convolution of the wakefields with the + beam. It is not meant to be used directly by the user, but only by the + WakeTracker class. + + The convolution algorithm is optimized to avoid the gaps between bunches, + using the algorithm described in the Xsuite physics manual (section 10.5), + which was originally developed for PyHEADTAIL by J. Komppula + (https://indico.cern.ch/event/735184/contributions/3032237/attachments/1668727/2676169/Multibunch_pyheadtail_algorithms.pdf) + and N. Mounet (https://indico.cern.ch/event/735184/contributions/3032242/attachments/1668613/2676354/20180615_PyHEADTAIL_convolution_algorithm.pdf). + ''' def __init__(self, component, waketracker=None, _flatten=False, log_moments=None, **kwargs): + ''' + Parameters: + ----------- + component: xfields.beam_elements.waketracker.WakeField + The wakefield component to be convolved with the beam + waketracker: xfields.beam_elements.waketracker.WakeTracker + The WakeTracker object that will use this ConvData object + log_moments: list of str + List of moments to be logged in the moments_data object + ''' self.xoinitialize(**kwargs) diff --git a/xfields/slicers/compressed_profile.py b/xfields/slicers/compressed_profile.py index ce814e7b..1895690e 100644 --- a/xfields/slicers/compressed_profile.py +++ b/xfields/slicers/compressed_profile.py @@ -6,7 +6,14 @@ class CompressedProfile(xt.BeamElement): """ - An object holding a compressed version of the beam data. + An object holding a compressed version of the beam data. This allows to + store the moments of the beam in a compressed way, i.e. avoiding to store + the moments in the empty slices between bunches. + The CompressedProfile is used for example in the _ConvData class for the + computation of the wake kicks. + The way this is handled is based on the algorithm devised by J. Komppula + (https://indico.cern.ch/event/735184/contributions/3032237/attachments/1668727/2676169/Multibunch_pyheadtail_algorithms.pdf) + and N. Mounet (https://indico.cern.ch/event/735184/contributions/3032242/attachments/1668613/2676354/20180615_PyHEADTAIL_convolution_algorithm.pdf). Parameters ----------. diff --git a/xfields/slicers/uniform.py b/xfields/slicers/uniform.py index 336737f8..8cebebd8 100644 --- a/xfields/slicers/uniform.py +++ b/xfields/slicers/uniform.py @@ -351,7 +351,7 @@ def _reshape_for_multibunch(self, data): return data.reshape(self.num_bunches, self.num_slices) def _to_npbuffer(self): - assert isinstance(self._context, xo.ContextCpu) + #assert isinstance(self._context, xo.ContextCpu) assert self._buffer.buffer.dtype == np.int8 return self._buffer.buffer[self._offset: self._offset + self._xobject._size]