diff --git a/tests/test_elements_thick.py b/tests/test_elements_thick.py index 47bf2b70f..966ed626b 100644 --- a/tests/test_elements_thick.py +++ b/tests/test_elements_thick.py @@ -1156,6 +1156,130 @@ def test_solenoid_with_mult_kicks(test_context, backtrack): xo.assert_allclose(p_test.pzeta, p_ref.pzeta, rtol=0, atol=1e-13) +@for_all_test_contexts +@pytest.mark.parametrize( + 'radiation_mode,config', + [ + ('mean', {}), + ('mean', {'XTRACK_SYNRAD_KICK_SAME_AS_FIRST': True}), + ('mean', {'XTRACK_SYNRAD_SCALE_SAME_AS_FIRST': True}), + ('quantum', {}), + ], +) +def test_drift_like_solenoid_with_kicks_radiation(test_context, radiation_mode, config): + config['XTRACK_USE_EXACT_DRIFTS'] = True + knl = [0.1, 0.4, 0.5] + ksl = [0.2, 0.3, 0.6] + + line_test = xt.Line(elements=[ + xt.Drift(length=0.5), + xt.Multipole(knl=knl, ksl=ksl), + xt.Drift(length=0.5), + ]) + + line_ref = xt.Line(elements=[ + xt.Solenoid(ks=0, length=1, knl=knl, ksl=ksl, num_multipole_kicks=1) + ]) + + coords = np.linspace(-0.05, 0.05, 10) + coords_6d = np.array(list(itertools.product(*(coords,) * 6))).T + + p0 = xt.Particles( + _context=test_context, + x=coords_6d[0], + px=coords_6d[1], + y=coords_6d[2], + py=coords_6d[3], + zeta=coords_6d[4], + delta=coords_6d[5], + ) + p0._init_random_number_generator(seeds=np.arange(len(coords_6d[0]), dtype=int)) + p_ref = p0.copy() + p_test = p0.copy() + + line_ref.build_tracker(_context=test_context) + line_ref.configure_radiation(model=radiation_mode) + line_ref.config.update(config) + line_ref.track(p_ref, num_turns=1) + + line_test.build_tracker(_context=test_context) + line_test.configure_radiation(model=radiation_mode) + line_test.config.update(config) + line_test.track(p_test, num_turns=1) + + xo.assert_allclose(p_test.x, p_ref.x, rtol=0, atol=1e-13) + xo.assert_allclose(p_test.px, p_ref.px, rtol=0, atol=1e-13) + xo.assert_allclose(p_test.y, p_ref.y, rtol=0, atol=1e-13) + xo.assert_allclose(p_test.py, p_ref.py, rtol=0, atol=1e-13) + xo.assert_allclose(p_test.delta, p_ref.delta, rtol=0, atol=1e-13) + xo.assert_allclose(p_test.pzeta, p_ref.pzeta, rtol=0, atol=1e-13) + + +@for_all_test_contexts +@pytest.mark.parametrize( + 'radiation_mode,config', + [ + ('mean', {}), + ('mean', {'XTRACK_SYNRAD_KICK_SAME_AS_FIRST': True}), + ('mean', {'XTRACK_SYNRAD_SCALE_SAME_AS_FIRST': True}), + ('quantum', {}), + ], +) +def test_solenoid_with_kicks_radiation(test_context, radiation_mode, config): + config['XTRACK_USE_EXACT_DRIFTS'] = True + + ks = 0.4 + l = 1.1 + knl = [0.1, 0.4, 0.5] + ksl = [0.2, 0.3, 0.6] + + sol_ref = xt.Solenoid(ks=ks, length=l) + sol_1 = xt.Solenoid(ks=ks, length=l, knl=knl, ksl=ksl, num_multipole_kicks=1) + sol_3 = xt.Solenoid(ks=ks, length=l, knl=knl, ksl=ksl, num_multipole_kicks=3) + + line_ref = xt.Line(elements=[sol_ref]) + line_1 = xt.Line(elements=[sol_1]) + line_3 = xt.Line(elements=[sol_3]) + + coords = np.linspace(-0.05, 0.05, 10) + coords_6d = np.array(list(itertools.product(*(coords,) * 6))).T + + p0 = xt.Particles( + _context=test_context, + x=coords_6d[0], + px=coords_6d[1], + y=coords_6d[2], + py=coords_6d[3], + zeta=coords_6d[4], + delta=coords_6d[5], + ) + p0._init_random_number_generator(seeds=np.arange(len(coords_6d[0]), dtype=int)) + p_ref = p0.copy() + p_1 = p0.copy() + p_3 = p0.copy() + + line_ref.build_tracker(_context=test_context) + line_ref.track(p_ref, num_turns=1) + + line_1.build_tracker(_context=test_context) + line_1.configure_radiation(model=radiation_mode) + line_1.config.update(config) + line_1.track(p_1, num_turns=1) + + line_3.build_tracker(_context=test_context) + line_3.configure_radiation(model=radiation_mode) + line_3.config.update(config) + line_3.track(p_3, num_turns=1) + + d_delta_ref = p_ref.delta - p0.delta + d_delta_1 = p_1.delta - p0.delta + d_delta_3 = p_3.delta - p0.delta + + xo.assert_allclose( + d_delta_1 - d_delta_ref, d_delta_3 - d_delta_ref, rtol=0.01, atol=1e-15 + ) + + @for_all_test_contexts def test_skew_quadrupole(test_context): k1 = 1.0 diff --git a/tests/test_slice_elements.py b/tests/test_slice_elements.py index 4b63d1c71..6bba4c370 100644 --- a/tests/test_slice_elements.py +++ b/tests/test_slice_elements.py @@ -1,8 +1,11 @@ -import xtrack as xt +import itertools + +import numpy as np import pytest -from xobjects.test_helpers import for_all_test_contexts import xobjects as xo +import xtrack as xt +from xobjects.test_helpers import for_all_test_contexts assert_allclose= xo.assert_allclose @@ -630,9 +633,9 @@ def test_thick_slice_octupole(test_context): @for_all_test_contexts def test_thick_slice_solenoid(test_context): - oct = xt.Solenoid(ks=0.1, length=1) + sol = xt.Solenoid(ks=0.1, length=1) - line = xt.Line(elements=[oct]) + line = xt.Line(elements=[sol]) line.slice_thick_elements( slicing_strategies=[xt.Strategy(xt.Teapot(10, mode='thick'))]) @@ -646,7 +649,7 @@ def test_thick_slice_solenoid(test_context): p_slice = p0.copy() line.track(p_slice) - oct.track(p_ref) + sol.track(p_ref) assert_allclose(p_slice.x, p_ref.x, rtol=0, atol=1e-10) assert_allclose(p_slice.px, p_ref.px, rtol=0, atol=1e-10) diff --git a/tests/test_slicing.py b/tests/test_slicing.py index 4a403960c..ac5d2a812 100644 --- a/tests/test_slicing.py +++ b/tests/test_slicing.py @@ -280,7 +280,7 @@ def test_slicing_strategy_matching(): if name == 'keep_this': assert isinstance(element, xt.Quadrupole) elif name == 'keep_drifts': - assert (element, xt.Drift) + assert isinstance(element, xt.Drift) elif name.startswith('drift_'): assert 'DriftSlice' in type(element).__name__ elif name == 'keep_thin': diff --git a/tests/test_spacecharge_in_ring.py b/tests/test_spacecharge_in_ring.py index a03632298..ce70a6834 100644 --- a/tests/test_spacecharge_in_ring.py +++ b/tests/test_spacecharge_in_ring.py @@ -157,7 +157,7 @@ def test_ring_with_spacecharge(test_context, mode): else: assert isinstance(line[0], xf.SpaceChargeBiGaussian) - if mode is not 'frozen': + if mode != 'frozen': assert line.iscollective else: assert not line.iscollective diff --git a/xtrack/beam_elements/elements_src/solenoid.h b/xtrack/beam_elements/elements_src/solenoid.h index 8c61e3649..0bc2d8467 100644 --- a/xtrack/beam_elements/elements_src/solenoid.h +++ b/xtrack/beam_elements/elements_src/solenoid.h @@ -1,6 +1,6 @@ // copyright ############################### // // This file is part of the Xtrack Package. // -// Copyright (c) CERN, 2023. // +// Copyright (c) CERN, 2024. // // ######################################### // #ifndef XTRACK_SOLENOID_H @@ -20,13 +20,13 @@ void Solenoid_track_local_particle(SolenoidData el, LocalParticle* part0) { factor_knl_ksl = -1; #endif - #ifndef XTRACK_SOLENOID_NO_SYNRAD // does this do the right thing? - double dp_record_entry = 0.; - double dpx_record_entry = 0.; - double dpy_record_entry = 0.; - double dp_record_exit = 0.; - double dpx_record_exit = 0.; - double dpy_record_exit = 0.; + #ifndef XTRACK_SOLENOID_NO_SYNRAD + double dp_record_entry = 0.; + double dpx_record_entry = 0.; + double dpy_record_entry = 0.; + double dp_record_exit = 0.; + double dpx_record_exit = 0.; + double dpy_record_exit = 0.; #endif int64_t num_multipole_kicks = SolenoidData_get_num_multipole_kicks(el); @@ -37,22 +37,34 @@ void Solenoid_track_local_particle(SolenoidData el, LocalParticle* part0) { const double slice_length = length / (num_multipole_kicks + 1); const double kick_weight = 1. / num_multipole_kicks; + //start_per_particle_block (part0->part) + #ifndef XTRACK_SOLENOID_NO_SYNRAD + double const old_px = LocalParticle_get_px(part); + double const old_py = LocalParticle_get_py(part); + double const old_ax = LocalParticle_get_ax(part); + double const old_ay = LocalParticle_get_ay(part); + double const old_zeta = LocalParticle_get_zeta(part); + #endif + for (int ii = 0; ii < num_multipole_kicks; ii++) { - //start_per_particle_block (part0->part) - Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag, - &dp_record_entry, &dpx_record_entry, &dpy_record_entry, - &dp_record_exit, &dpx_record_exit, &dpy_record_exit); + Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag); track_multipolar_kick_bend( part, order, inv_factorial_order, knl, ksl, factor_knl_ksl, kick_weight, 0, 0, 0, 0); - //end_per_particle_block } - //start_per_particle_block (part0->part) - Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag, - &dp_record_entry, &dpx_record_entry, &dpy_record_entry, - &dp_record_exit, &dpx_record_exit, &dpy_record_exit); + Solenoid_thick_track_single_particle(part, slice_length, ks, radiation_flag); + + #ifndef XTRACK_SOLENOID_NO_SYNRAD + if (radiation_flag > 0 && length > 0){ + Solenoid_apply_radiation_single_particle( + part, length, radiation_flag, + old_px, old_py, old_ax, old_ay, old_zeta, + &dp_record_entry, &dpx_record_entry, &dpy_record_entry + ); + } + #endif //end_per_particle_block } diff --git a/xtrack/beam_elements/elements_src/thick_slice_bend.h b/xtrack/beam_elements/elements_src/thick_slice_bend.h index fbc62953e..8c74ada48 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_bend.h +++ b/xtrack/beam_elements/elements_src/thick_slice_bend.h @@ -19,8 +19,8 @@ void ThickSliceBend_track_local_particle( const double h = ThickSliceBendData_get__parent_h(el); const double order = ThickSliceBendData_get__parent_order(el); const double inv_factorial_order = ThickSliceBendData_get__parent_inv_factorial_order(el); - const double* knl = ThickSliceBendData_getp1__parent_knl(el, 0); - const double* ksl = ThickSliceBendData_getp1__parent_ksl(el, 0); + /*gpuglmem*/ const double* knl = ThickSliceBendData_getp1__parent_knl(el, 0); + /*gpuglmem*/ const double* ksl = ThickSliceBendData_getp1__parent_ksl(el, 0); const int64_t model = ThickSliceBendData_get__parent_model(el); const int64_t num_multipole_kicks_parent = ThickSliceBendData_get__parent_num_multipole_kicks(el); diff --git a/xtrack/beam_elements/elements_src/thick_slice_quadrupole.h b/xtrack/beam_elements/elements_src/thick_slice_quadrupole.h index d213d864e..9301771a3 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_quadrupole.h +++ b/xtrack/beam_elements/elements_src/thick_slice_quadrupole.h @@ -19,8 +19,8 @@ void ThickSliceQuadrupole_track_local_particle( const int64_t num_multipole_kicks_parent = ThickSliceQuadrupoleData_get__parent_num_multipole_kicks(el); const double order = ThickSliceQuadrupoleData_get__parent_order(el); const double inv_factorial_order = ThickSliceQuadrupoleData_get__parent_inv_factorial_order(el); - const double* knl = ThickSliceQuadrupoleData_getp1__parent_knl(el, 0); - const double* ksl = ThickSliceQuadrupoleData_getp1__parent_ksl(el, 0); + /*gpuglmem*/ const double* knl = ThickSliceQuadrupoleData_getp1__parent_knl(el, 0); + /*gpuglmem*/ const double* ksl = ThickSliceQuadrupoleData_getp1__parent_ksl(el, 0); #ifndef XSUITE_BACKTRACK double const length = weight * ThickSliceQuadrupoleData_get__parent_length(el); // m diff --git a/xtrack/beam_elements/elements_src/thick_slice_solenoid.h b/xtrack/beam_elements/elements_src/thick_slice_solenoid.h index af6151aa2..7c5ef5430 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_solenoid.h +++ b/xtrack/beam_elements/elements_src/thick_slice_solenoid.h @@ -24,10 +24,9 @@ void ThickSliceSolenoid_track_local_particle( //start_per_particle_block (part0->part) - Solenoid_thick_track_single_particle( + Solenoid_thick_with_radiation_track_single_particle( part, length, ks, 0, // radiation flag, not supported for now - NULL, NULL, NULL, NULL, NULL, NULL); //end_per_particle_block diff --git a/xtrack/beam_elements/elements_src/track_bend.h b/xtrack/beam_elements/elements_src/track_bend.h index 3d77668d2..3636c6060 100644 --- a/xtrack/beam_elements/elements_src/track_bend.h +++ b/xtrack/beam_elements/elements_src/track_bend.h @@ -19,7 +19,7 @@ void Bend_track_local_particle_from_params(LocalParticle* part0, double length, double k0, double k1, double h, int64_t num_multipole_kicks, int64_t model, - double const* knl, double const* ksl, + /*gpuglmem*/ double const* knl, /*gpuglmem*/ double const* ksl, int64_t order, double inv_factorial_order, double factor_knl_ksl){ diff --git a/xtrack/beam_elements/elements_src/track_multipolar_components.h b/xtrack/beam_elements/elements_src/track_multipolar_components.h index fb6a79437..249e95634 100644 --- a/xtrack/beam_elements/elements_src/track_multipolar_components.h +++ b/xtrack/beam_elements/elements_src/track_multipolar_components.h @@ -9,8 +9,8 @@ /*gpufun*/ void track_multipolar_kick_bend( LocalParticle* part, int64_t order, double inv_factorial_order, - const double* knl, - const double* ksl, + /*gpuglmem*/ const double* knl, + /*gpuglmem*/ const double* ksl, double const factor_knl_ksl, double kick_weight, double k0, double k1, double h, double length){ diff --git a/xtrack/beam_elements/elements_src/track_quadrupole.h b/xtrack/beam_elements/elements_src/track_quadrupole.h index 8967c8723..4ce23a007 100644 --- a/xtrack/beam_elements/elements_src/track_quadrupole.h +++ b/xtrack/beam_elements/elements_src/track_quadrupole.h @@ -35,7 +35,7 @@ void normal_quad_with_rotation_track( void Quadrupole_from_params_track_local_particle( double length, double k1, double k1s, int64_t num_multipole_kicks, - double const* knl, double const* ksl, + /*gpuglmem*/ double const* knl, /*gpuglmem*/ double const* ksl, int64_t order, double inv_factorial_order, double factor_knl_ksl, LocalParticle* part0 diff --git a/xtrack/beam_elements/elements_src/track_solenoid.h b/xtrack/beam_elements/elements_src/track_solenoid.h index f7f083f31..020f28d1e 100644 --- a/xtrack/beam_elements/elements_src/track_solenoid.h +++ b/xtrack/beam_elements/elements_src/track_solenoid.h @@ -13,9 +13,7 @@ void Solenoid_thick_track_single_particle( LocalParticle* part, double length, double ks, - int64_t radiation_flag, - double* dp_record_entry, double* dpx_record_entry, double* dpy_record_entry, - double* dp_record_exit, double* dpx_record_exit, double* dpy_record_exit + int64_t radiation_flag ) { const double sk = ks / 2; @@ -74,66 +72,83 @@ void Solenoid_thick_track_single_particle( double const new_ax = -0.5 * Bz * new_y * q0 * QELEM / P0_J; double const new_ay = 0.5 * Bz * new_x * q0 * QELEM / P0_J; - #ifndef XTRACK_SOLENOID_NO_SYNRAD - double l_path, curv; - if (radiation_flag > 0 && length > 0){ + LocalParticle_set_x(part, new_x); + LocalParticle_set_px(part, new_px); + LocalParticle_set_y(part, new_y); + LocalParticle_set_py(part, new_py); + LocalParticle_add_to_zeta(part, add_to_zeta); + LocalParticle_add_to_s(part, length); + LocalParticle_set_ax(part, new_ax); + LocalParticle_set_ay(part, new_ay); +} - double const old_ax = LocalParticle_get_ax(part); - double const old_ay = LocalParticle_get_ay(part); +/*gpufun*/ +void Solenoid_apply_radiation_single_particle( + /*gpuglmem*/ LocalParticle* part, + const double length, + const int64_t radiation_flag, + const double old_px, const double old_py, + const double old_ax, const double old_ay, + const double old_zeta, + double* dp_record_exit, double* dpx_record_exit, double* dpy_record_exit +) { + double const rvv = LocalParticle_get_rvv(part); + double const new_ax = LocalParticle_get_ax(part); + double const new_ay = LocalParticle_get_ay(part); - double const old_px_mech = px - old_ax; - double const old_py_mech = py - old_ay; + double const old_px_mech = old_px - old_ax; + double const old_py_mech = old_py - old_ay; - double const new_px_mech = new_px - new_ax; - double const new_py_mech = new_py - new_ay; + double const new_px_mech = LocalParticle_get_px(part) - new_ax; + double const new_py_mech = LocalParticle_get_py(part) - new_ay; - double const dpx = new_px_mech - old_px_mech; - double const dpy = new_py_mech - old_py_mech; + double const dpx = new_px_mech - old_px_mech; + double const dpy = new_py_mech - old_py_mech; - curv = sqrt(dpx*dpx + dpy*dpy) / length; + double const curv = sqrt(dpx * dpx + dpy * dpy) / length; - // Path length for radiation - double const dzeta = add_to_zeta; - double const rvv = LocalParticle_get_rvv(part); - l_path = rvv * (length - dzeta); + // Path length for radiation + double const dzeta = LocalParticle_get_zeta(part) - old_zeta; + double const l_path = rvv * (length - dzeta); - // LocalParticle_add_to_px(part, -old_ax); - // LocalParticle_add_to_py(part, -old_ay); - // if (radiation_flag == 1){ - // synrad_average_kick(part, curv, l_path / 2, - // dp_record_entry, dpx_record_entry, dpy_record_entry); - // } - // else if (radiation_flag == 2){ - // synrad_emit_photons(part, curv, l_path / 2, NULL, NULL); - // } - // LocalParticle_add_to_px(part, old_ax); - // LocalParticle_add_to_py(part, old_ay); - } - #endif + LocalParticle_add_to_px(part, -new_ax); + LocalParticle_add_to_py(part, -new_ay); - LocalParticle_set_x(part, new_x); - LocalParticle_add_to_px(part, new_px - px); - LocalParticle_set_y(part, new_y); - LocalParticle_add_to_py(part, new_py - py); - LocalParticle_add_to_zeta(part, add_to_zeta); - LocalParticle_add_to_s(part, length); - LocalParticle_set_ax(part, new_ax); - LocalParticle_set_ay(part, new_ay); + if (radiation_flag == 1){ + synrad_average_kick(part, curv, l_path, + dp_record_exit, dpx_record_exit, dpy_record_exit); + } + else if (radiation_flag == 2){ + synrad_emit_photons(part, curv, l_path, NULL, NULL); + } + + LocalParticle_add_to_px(part, new_ax); + LocalParticle_add_to_py(part, new_ay); +} +/*gpufun*/ +void Solenoid_thick_with_radiation_track_single_particle( + LocalParticle* part, + double length, + double ks, + int64_t radiation_flag, + double* dp_record_exit, double* dpx_record_exit, double* dpy_record_exit +) { #ifndef XTRACK_SOLENOID_NO_SYNRAD + double const old_px = LocalParticle_get_px(part); + double const old_py = LocalParticle_get_py(part); + double const old_ax = LocalParticle_get_ax(part); + double const old_ay = LocalParticle_get_ay(part); + double const old_zeta = LocalParticle_get_zeta(part); + #endif + Solenoid_thick_track_single_particle(part, length, ks, radiation_flag); + + #ifndef XTRACK_SOLENOID_NO_SYNRAD if (radiation_flag > 0 && length > 0){ - LocalParticle_add_to_px(part, -new_ax); - LocalParticle_add_to_py(part, -new_ay); - if (radiation_flag == 1){ - synrad_average_kick(part, curv, l_path, - dp_record_exit, dpx_record_exit, dpy_record_exit); - } - else if (radiation_flag == 2){ - synrad_emit_photons(part, curv, l_path, NULL, NULL); - } - LocalParticle_add_to_px(part, new_ax); - LocalParticle_add_to_py(part, new_ay); + Solenoid_apply_radiation_single_particle(part, length, radiation_flag, + old_px, old_py, old_ax, old_ay, old_zeta, + dp_record_exit, dpx_record_exit, dpy_record_exit); } #endif }