Feature/wakes on gpu

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_

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()

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()

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)