diff --git a/examples/element_internal_record/000_internal_record.py b/examples/element_internal_record/000_internal_record.py index 71c272137..0c6c5719f 100644 --- a/examples/element_internal_record/000_internal_record.py +++ b/examples/element_internal_record/000_internal_record.py @@ -36,20 +36,22 @@ class TestElementRecord(xo.HybridClass): # element number. track_method_source = r''' -/*gpufun*/ -void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ +#include "xtrack/headers/track.h" +GPUFUN +void TestElement_track_local_particle(TestElementData el, LocalParticle* part0) +{ // Extract the record and record_index TestElementRecordData record = TestElementData_getp_internal_record(el, part0); RecordIndex record_index = NULL; - if (record){ + if (record) { record_index = TestElementRecordData_getp__index(record); } int64_t n_kicks = TestElementData_get_n_kicks(el); printf("n_kicks %d\n", (int)n_kicks); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); for (int64_t i = 0; i < n_kicks; i++) { double rr = 1e-6 * RandomUniform_generate(part); @@ -72,8 +74,7 @@ class TestElementRecord(xo.HybridClass): } } - - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''' diff --git a/examples/element_internal_record/001_multirecord.py b/examples/element_internal_record/001_multirecord.py index 2d081187c..d4f21212c 100644 --- a/examples/element_internal_record/001_multirecord.py +++ b/examples/element_internal_record/001_multirecord.py @@ -41,7 +41,9 @@ class TestElementRecord(xo.HybridClass): # code of the beam element. TestElement_track_method_source = r''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ // Extract the record and record_index @@ -60,7 +62,7 @@ class TestElementRecord(xo.HybridClass): int64_t n_kicks = TestElementData_get_n_kicks(el); printf("n_kicks %d\n", (int)n_kicks); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); // Record in table1 info about the ingoing particle if (record){ @@ -106,7 +108,7 @@ class TestElementRecord(xo.HybridClass): } } - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''' diff --git a/examples/element_internal_record/002_record_in_individual_element.py b/examples/element_internal_record/002_record_in_individual_element.py index 57cc4581f..a144876b9 100644 --- a/examples/element_internal_record/002_record_in_individual_element.py +++ b/examples/element_internal_record/002_record_in_individual_element.py @@ -40,7 +40,9 @@ class TestElementRecord(xo.HybridClass): # code of the beam element. TestElement_track_method_source = r''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ // Extract the record and record_index @@ -59,7 +61,7 @@ class TestElementRecord(xo.HybridClass): int64_t n_kicks = TestElementData_get_n_kicks(el); // printf("n_kicks %d\n", (int)n_kicks); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); // Record in table1 info about the ingoing particle if (record){ @@ -105,7 +107,7 @@ class TestElementRecord(xo.HybridClass): } } - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''' diff --git a/examples/random_number_generator/001_usage_in_beam_element.py b/examples/random_number_generator/001_usage_in_beam_element.py index 32cc71e71..e91908af8 100644 --- a/examples/random_number_generator/001_usage_in_beam_element.py +++ b/examples/random_number_generator/001_usage_in_beam_element.py @@ -27,12 +27,15 @@ class TestElement(xt.BeamElement): _extra_c_sources = [ ''' - /*gpufun*/ - void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ - //start_per_particle_block (part0->part) + #include "xtrack/headers/track.h" + + GPUFUN + void TestElement_track_local_particle(TestElementData el, LocalParticle* part0) + { + START_PER_PARTICLE_BLOCK(part0, part); double rr = !!GENERATOR!!_generate(part); LocalParticle_set_x(part, rr); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''.replace('!!GENERATOR!!', generator_to_test) ] diff --git a/examples/spin_lep/chirp_kicker.py b/examples/spin_lep/chirp_kicker.py index e0a0b377d..de8079743 100644 --- a/examples/spin_lep/chirp_kicker.py +++ b/examples/spin_lep/chirp_kicker.py @@ -19,22 +19,22 @@ class VerticalChirpKicker(xt.BeamElement): _extra_c_sources =[''' - #include - #include + #include "xtrack/headers/track.h" + #include "xtrack/beam_elements/elements_src/track_magnet.h" - /*gpufun*/ + GPUFUN void VerticalChirpKicker_track_local_particle( - VerticalChirpKickerData el, LocalParticle* part0){ - + VerticalChirpKickerData el, LocalParticle* part0) + { double const k0sl = VerticalChirpKickerData_get_k0sl(el); double const q_start = VerticalChirpKickerData_get_q_start(el); double const q_end = q_start + VerticalChirpKickerData_get_q_span(el); double const num_turns = VerticalChirpKickerData_get_num_turns(el); double const length = VerticalChirpKickerData_get_length(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); double const at_turn = LocalParticle_get_at_turn(part); - if (at_turn < num_turns){ + if (at_turn < num_turns) { // integrating to get the instantaneous phase double const phi = 2 * PI * q_start * at_turn + PI * (q_end - q_start) / ((double) num_turns) * ((double) at_turn * at_turn); @@ -60,9 +60,9 @@ class VerticalChirpKicker(xt.BeamElement): 0, // frame curvature length, length // lpath - same for a thin element - ); - #endif + ); + #endif } - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] diff --git a/examples/true_rbend/000a_check_edge_full.py b/examples/true_rbend/000a_check_edge_full.py index 84f5bc326..774cf4493 100644 --- a/examples/true_rbend/000a_check_edge_full.py +++ b/examples/true_rbend/000a_check_edge_full.py @@ -12,7 +12,7 @@ l_ref = xt.Line([b_ref]) l_ref.append('end', xt.Marker()) l_ref.particle_ref = xt.Particles(p0c=10e9) -tw_ref0 = l_ref.twiss(betx=1, bety=1) +tw_ref0 = l_ref.twiss(betx=1, bety=1, strengths=True) tw_ref = l_ref.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) b_test = xt.RBend( @@ -20,6 +20,7 @@ b_test.rbend_model = 'straight-body' b_test.model = 'bend-kick-bend' b_test.num_multipole_kicks = 100 +b_test.rbend_compensate_sagitta = True b_test.edge_entry_model = edge_model b_test.edge_exit_model = edge_model l_test = xt.Line([b_test]) diff --git a/examples/true_rbend/000c_check_edge_full_angle._diff.py b/examples/true_rbend/000c_check_edge_full_angle._diff.py new file mode 100644 index 000000000..4ca1ba7dc --- /dev/null +++ b/examples/true_rbend/000c_check_edge_full_angle._diff.py @@ -0,0 +1,72 @@ +import xtrack as xt +import numpy as np +import xobjects as xo + +edge_model = 'full' + +b_ref = xt.RBend(angle=0.1, k0_from_h=True, length_straight=3., rbend_angle_diff=0.1) +b_ref.edge_entry_model = edge_model +b_ref.edge_exit_model = edge_model +b_ref.model = 'rot-kick-rot' +b_ref.num_multipole_kicks = 100 +l_ref = xt.Line([b_ref]) +l_ref.append('end', xt.Marker()) +l_ref.particle_ref = xt.Particles(p0c=10e9) +tw_ref0 = l_ref.twiss(betx=1, bety=1, strengths=True) +tw_ref = l_ref.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + +b_test = xt.RBend( + angle=0.1, k0_from_h=True, length_straight=3, rbend_angle_diff=0.1, + rbend_shift=0, rbend_compensate_sagitta=False) +b_test.rbend_model = 'straight-body' +b_test.model = 'bend-kick-bend' +b_test.num_multipole_kicks = 100 +b_test.edge_entry_model = edge_model +b_test.edge_exit_model = edge_model + +b_test.rbend_shift += b_test._x0_in + +l_test = xt.Line([b_test]) +l_test.append('end', xt.Marker()) +l_test.particle_ref = xt.Particles(p0c=10e9) +tw_test0 = l_test.twiss(betx=1, bety=1) +tw_test = l_test.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + +xo.assert_allclose(tw_ref0.betx, tw_test0.betx, rtol=1e-9, atol=0.0) +xo.assert_allclose(tw_ref0.bety, tw_test0.bety, rtol=1e-9, atol=0.0) +xo.assert_allclose(tw_ref0.x, tw_test0.x, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref0.y, tw_test0.y, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref0.s, tw_test0.s, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref0.zeta, tw_test0.zeta, rtol=0, atol=1e-11) +xo.assert_allclose(tw_ref0.px, tw_test0.px, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref0.py, tw_test0.py, rtol=0, atol=1e-12) + +xo.assert_allclose(tw_ref.betx, tw_test.betx, rtol=5e-9, atol=0.0) +xo.assert_allclose(tw_ref.bety, tw_test.bety, rtol=5e-9, atol=0.0) +xo.assert_allclose(tw_ref.x, tw_test.x, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref.y, tw_test.y, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref.s, tw_test.s, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref.zeta, tw_test.zeta, rtol=0, atol=1e-1) +xo.assert_allclose(tw_ref.px, tw_test.px, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref.py, tw_test.py, rtol=0, atol=1e-12) + +l_sliced = l_test.copy(shallow=True) +l_sliced.cut_at_s(np.linspace(0, l_test.get_length(), 100)) +tw_test_sliced0 = l_sliced.twiss(betx=1, bety=1) + +xo.assert_allclose(tw_test_sliced0.betx[-1], tw_test0.betx[-1], rtol=1e-9, atol=0.0) +xo.assert_allclose(tw_test_sliced0.bety[-1], tw_test0.bety[-1], rtol=1e-9, atol=0.0) +xo.assert_allclose(tw_test_sliced0.x[-1], tw_test0.x[-1], rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced0.y[-1], tw_test0.y[-1], rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced0.s[-1], tw_test0.s[-1], rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced0.zeta[-1], tw_test0.zeta[-1], rtol=0, atol=1e-11) +xo.assert_allclose(tw_test_sliced0.px[-1], tw_test0.px[-1], rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced0.py[-1], tw_test0.py[-1], rtol=0, atol=1e-12) + +import matplotlib.pyplot as plt +plt.close('all') + +tw_test_sliced0.plot('x') +plt.xlim(-0.1, 3.1) + +plt.show() \ No newline at end of file diff --git a/examples/true_rbend/000d_check_edge_linear_angle_diff.py b/examples/true_rbend/000d_check_edge_linear_angle_diff.py new file mode 100644 index 000000000..eacac1a65 --- /dev/null +++ b/examples/true_rbend/000d_check_edge_linear_angle_diff.py @@ -0,0 +1,96 @@ +import xtrack as xt +import numpy as np +import xobjects as xo + +edge_model = 'linear' + +b_ref = xt.RBend(angle=0.1, k0_from_h=True, length_straight=3., rbend_angle_diff=0.1) +b_ref.edge_entry_model = edge_model +b_ref.edge_exit_model = edge_model +b_ref.model = 'rot-kick-rot' +b_ref.num_multipole_kicks = 100 +l_ref = xt.Line([b_ref]) +l_ref.append('end', xt.Marker()) +l_ref.particle_ref = xt.Particles(p0c=10e9) +tw_ref0 = l_ref.twiss(betx=1, bety=1) +tw_ref = l_ref.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + +b_test = xt.RBend( + angle=0.1, k0_from_h=True, length_straight=3, rbend_angle_diff=0.1, + rbend_shift=0, rbend_compensate_sagitta=False) +b_test.rbend_model = 'straight-body' +b_test.model = 'bend-kick-bend' +b_test.num_multipole_kicks = 100 +b_test.edge_entry_model = edge_model +b_test.edge_exit_model = edge_model + +b_test.rbend_shift += b_test._x0_in + +l_test = xt.Line([b_test]) +l_test.append('end', xt.Marker()) +l_test.particle_ref = xt.Particles(p0c=10e9) +tw_test0 = l_test.twiss(betx=1, bety=1) +tw_test = l_test.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + +xo.assert_allclose(tw_ref0.betx, tw_test0.betx, rtol=1e-9, atol=0.0) +xo.assert_allclose(tw_ref0.bety, tw_test0.bety, rtol=1e-9, atol=0.0) +xo.assert_allclose(tw_ref0.x, tw_test0.x, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref0.y, tw_test0.y, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref0.s, tw_test0.s, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref0.zeta, tw_test0.zeta, rtol=0, atol=1e-11) +xo.assert_allclose(tw_ref0.px, tw_test0.px, rtol=0, atol=1e-12) +xo.assert_allclose(tw_ref0.py, tw_test0.py, rtol=0, atol=1e-12) + +xo.assert_allclose(tw_ref.betx, tw_test.betx, rtol=5e-6, atol=0.0) +xo.assert_allclose(tw_ref.bety, tw_test.bety, rtol=5e-6, atol=0.0) +xo.assert_allclose(tw_ref.x, tw_test.x, rtol=0, atol=2e-8) +xo.assert_allclose(tw_ref.y, tw_test.y, rtol=0, atol=2e-8) +xo.assert_allclose(tw_ref.s, tw_test.s, rtol=0, atol=2e-12) +xo.assert_allclose(tw_ref.zeta, tw_test.zeta, rtol=0, atol=1e-1) +xo.assert_allclose(tw_ref.px, tw_test.px, rtol=0, atol=5e-9) +xo.assert_allclose(tw_ref.py, tw_test.py, rtol=0, atol=5e-9) + +tw_back = l_test.twiss(init=tw_test, init_at='end') + +assert tw_back.orientation == 'backward' +xo.assert_allclose(tw_back.betx, tw_test.betx, rtol=5e-6, atol=0.0) +xo.assert_allclose(tw_back.bety, tw_test.bety, rtol=5e-6, atol=0.0) +xo.assert_allclose(tw_back.x, tw_test.x, rtol=0, atol=1e-8) +xo.assert_allclose(tw_back.y, tw_test.y, rtol=0, atol=1e-8) +xo.assert_allclose(tw_back.s, tw_test.s, rtol=0, atol=1e-12) +xo.assert_allclose(tw_back.zeta, tw_test.zeta, rtol=0, atol=1e-1) +xo.assert_allclose(tw_back.px, tw_test.px, rtol=0, atol=1e-9) +xo.assert_allclose(tw_back.py, tw_test.py, rtol=0, atol=1e-9) + +l_sliced = l_test.copy(shallow=True) +l_sliced.cut_at_s(np.linspace(0, l_test.get_length(), 100)) +tw_test_sliced0 = l_sliced.twiss(betx=1, bety=1) + +xo.assert_allclose(tw_test_sliced0.betx[-1], tw_test0.betx[-1], rtol=1e-9, atol=0.0) +xo.assert_allclose(tw_test_sliced0.bety[-1], tw_test0.bety[-1], rtol=1e-9, atol=0.0) +xo.assert_allclose(tw_test_sliced0.x[-1], tw_test0.x[-1], rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced0.y[-1], tw_test0.y[-1], rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced0.s[-1], tw_test0.s[-1], rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced0.zeta[-1], tw_test0.zeta[-1], rtol=0, atol=1e-11) +xo.assert_allclose(tw_test_sliced0.px[-1], tw_test0.px[-1], rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced0.py[-1], tw_test0.py[-1], rtol=0, atol=1e-12) + +tw_test_sliced_back = l_sliced.twiss(init=tw_test_sliced0, init_at='end') + +assert tw_test_sliced_back.orientation == 'backward' +xo.assert_allclose(tw_test_sliced_back.betx, tw_test_sliced0.betx, rtol=5e-9, atol=0.0) +xo.assert_allclose(tw_test_sliced_back.bety, tw_test_sliced0.bety, rtol=5e-9, atol=0.0) +xo.assert_allclose(tw_test_sliced_back.x, tw_test_sliced0.x, rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced_back.y, tw_test_sliced0.y, rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced_back.s, tw_test_sliced0.s, rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced_back.zeta, tw_test_sliced0.zeta, rtol=0, atol=1e-11) +xo.assert_allclose(tw_test_sliced_back.px, tw_test_sliced0.px, rtol=0, atol=1e-12) +xo.assert_allclose(tw_test_sliced_back.py, tw_test_sliced0.py, rtol=0, atol=1e-12) + +import matplotlib.pyplot as plt +plt.close('all') + +tw_test_sliced0.plot('x') +plt.xlim(-0.1, 3.1) + +plt.show() \ No newline at end of file diff --git a/examples/true_rbend/002_survey_hbend.py b/examples/true_rbend/002_survey_hbend.py index a8c7a7ad8..c858ff893 100644 --- a/examples/true_rbend/002_survey_hbend.py +++ b/examples/true_rbend/002_survey_hbend.py @@ -127,19 +127,18 @@ # _end_point 5 0 assert np.all(sv_straight['name'] == [ - 'start', 'drift_1..0', 'drift_1..1', 'mb_entry', 'mb..entry_map', + 'start', '||drift_3::0', '||drift_4', 'mb_entry', 'mb..entry_map', 'mb..0', 'mb..1', 'mb..2', 'mb..3', 'mid', 'mb..4', 'mb..5', - 'mb..6', 'mb..7', 'mb..exit_map', 'mb_exit', 'drift_2..0', - 'drift_2..1', 'end', '_end_point' -]) + 'mb..6', 'mb..7', 'mb..exit_map', 'mb_exit', '||drift_5', + '||drift_3::1', 'end', '_end_point']) # Assert entire columns using np.all -assert np.all(sv_straight['element_type'] == ['Marker', 'DriftSlice', 'DriftSlice', 'Marker', - 'ThinSliceRBendEntry', 'ThickSliceRBend', 'ThickSliceRBend', - 'ThickSliceRBend', 'ThickSliceRBend', 'Marker', 'ThickSliceRBend', +assert np.all(sv_straight['element_type'] == [ + 'Marker', 'Drift', 'Drift', 'Marker', 'ThinSliceRBendEntry', 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', - 'ThinSliceRBendExit', 'Marker', 'DriftSlice', 'DriftSlice', - 'Marker', '']) + 'ThickSliceRBend', 'Marker', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThinSliceRBendExit', + 'Marker', 'Drift', 'Drift', 'Marker', '']) xo.assert_allclose( sv_straight['angle'], @@ -257,19 +256,17 @@ # _end_point 5 0 assert np.all(sv_curved['name'] == [ - 'start', 'drift_1..0', 'drift_1..1', 'mb_entry', 'mb..entry_map', - 'mb..0', 'mb..1', 'mb..2', 'mb..3', 'mid', 'mb..4', 'mb..5', - 'mb..6', 'mb..7', 'mb..exit_map', 'mb_exit', 'drift_2..0', - 'drift_2..1', 'end', '_end_point' -]) + 'start', '||drift_3::0', '||drift_4', 'mb_entry', 'mb..entry_map', + 'mb..0', 'mb..1', 'mb..2', 'mb..3', 'mid', 'mb..4', 'mb..5', + 'mb..6', 'mb..7', 'mb..exit_map', 'mb_exit', '||drift_5', + '||drift_3::1', 'end', '_end_point']) assert np.all(sv_curved['element_type'] == [ - 'Marker', 'DriftSlice', 'DriftSlice', 'Marker', - 'ThinSliceRBendEntry', 'ThickSliceRBend', 'ThickSliceRBend', - 'ThickSliceRBend', 'ThickSliceRBend', 'Marker', 'ThickSliceRBend', - 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', - 'ThinSliceRBendExit', 'Marker', 'DriftSlice', 'DriftSlice', - 'Marker', '']) + 'Marker', 'Drift', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'Marker', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThinSliceRBendExit', + 'Marker', 'Drift', 'Drift', 'Marker', '']) xo.assert_allclose( sv_curved['angle'], diff --git a/examples/true_rbend/002a_survey_hbend_angle_diff.py b/examples/true_rbend/002a_survey_hbend_angle_diff.py new file mode 100644 index 000000000..f27261eeb --- /dev/null +++ b/examples/true_rbend/002a_survey_hbend_angle_diff.py @@ -0,0 +1,310 @@ +import xtrack as xt +import numpy as np +import xobjects as xo + +env = xt.Environment(particle_ref=xt.Particles(p0c=10e9)) + +edge_model = 'full' + +line = env.new_line(length=5, components=[ + env.new('mb', 'RBend', angle=0.3, k0_from_h=True, length_straight=3, + rbend_angle_diff=0.3, + model='bend-kick-bend', + rbend_model='straight-body', edge_entry_model=edge_model, edge_exit_model=edge_model, + anchor='start', at=1.)]) +line['mb'].rbend_compensate_sagitta = False +line['mb'].rbend_shift = line['mb']._x0_in +line.insert('start', xt.Marker(), at=0) +line.append('end', xt.Marker()) + +line_no_slice = line.copy(shallow=True) + +line.slice_thick_elements( + slicing_strategies=[ + xt.Strategy(slicing=xt.Uniform(6, mode='thick'))]) + +line['mb'].rbend_model = 'straight-body' +sv_straight = line.survey() +tt_straight = line.get_table(attr=True) +tw_straight = line.twiss(betx=1, bety=1) +p_straight = (sv_straight.p0 + tw_straight.x[:, None] * sv_straight['ex'] + + tw_straight.y[:, None] * sv_straight['ey']) +tw_straight['X'] = p_straight[:, 0] +tw_straight['Y'] = p_straight[:, 1] +tw_straight['Z'] = p_straight[:, 2] + +sv_straight_start = line.survey(element0='start', + X0=sv_straight['X', 'start'], + Y0=sv_straight['Y', 'start'], + Z0=sv_straight['Z', 'start'], + theta0=sv_straight['theta', 'start'], + phi0=sv_straight['phi', 'start'], + psi0=sv_straight['psi', 'start']) +sv_straight_end = line.survey(element0='end', + X0=sv_straight['X', 'end'], + Y0=sv_straight['Y', 'end'], + Z0=sv_straight['Z', 'end'], + theta0=sv_straight['theta', 'end'], + phi0=sv_straight['phi', 'end'], + psi0=sv_straight['psi', 'end']) + +sv_no_slice_start = line_no_slice.survey(element0='start', + X0=sv_straight['X', 'start'], + Y0=sv_straight['Y', 'start'], + Z0=sv_straight['Z', 'start'], + theta0=sv_straight['theta', 'start'], + phi0=sv_straight['phi', 'start'], + psi0=sv_straight['psi', 'start']) +sv_no_slice_end = line_no_slice.survey(element0='end', + X0=sv_straight['X', 'end'], + Y0=sv_straight['Y', 'end'], + Z0=sv_straight['Z', 'end'], + theta0=sv_straight['theta', 'end'], + phi0=sv_straight['phi', 'end'], + psi0=sv_straight['psi', 'end']) +tw_no_slice_straight = line_no_slice.twiss(betx=1, bety=1) + +line['mb'].rbend_model = 'curved-body' +sv_curved = line.survey() +tt_curved = line.get_table(attr=True) +tw_curved = line.twiss(betx=1, bety=1) +p_curved = (sv_curved.p0 + tw_curved.x[:, None] * sv_curved['ex'] + + tw_curved.y[:, None] * sv_curved['ey']) +tw_curved['X'] = p_curved[:, 0] +tw_curved['Y'] = p_curved[:, 1] +tw_curved['Z'] = p_curved[:, 2] + +sv_curved_start = line.survey(element0='start', + X0=sv_curved['X', 'start'], + Y0=sv_curved['Y', 'start'], + Z0=sv_curved['Z', 'start'], + theta0=sv_curved['theta', 'start'], + phi0=sv_curved['phi', 'start'], + psi0=sv_curved['psi', 'start']) +sv_curved_end = line.survey(element0='end', + X0=sv_curved['X', 'end'], + Y0=sv_curved['Y', 'end'], + Z0=sv_curved['Z', 'end'], + theta0=sv_curved['theta', 'end'], + phi0=sv_curved['phi', 'end'], + psi0=sv_curved['psi', 'end']) +sv_no_slice_curved_start = line_no_slice.survey(element0='start', + X0=sv_curved['X', 'start'], + Y0=sv_curved['Y', 'start'], + Z0=sv_curved['Z', 'start'], + theta0=sv_curved['theta', 'start'], + phi0=sv_curved['phi', 'start'], + psi0=sv_curved['psi', 'start']) +sv_no_slice_curved_end = line_no_slice.survey(element0='end', + X0=sv_curved['X', 'end'], + Y0=sv_curved['Y', 'end'], + Z0=sv_curved['Z', 'end'], + theta0=sv_curved['theta', 'end'], + phi0=sv_curved['phi', 'end'], + psi0=sv_curved['psi', 'end']) +tw_no_slice_curved = line_no_slice.twiss(betx=1, bety=1) + + +assert np.all(sv_straight['name'] == [ + 'start', '||drift_1', 'mb_entry', 'mb..entry_map', 'mb..0', + 'mb..1', 'mb..2', 'mb..3', 'mb..4', 'mb..5', 'mb..exit_map', + 'mb_exit', '||drift_2', 'end', '_end_point']) + +# Assert entire columns using np.all +assert np.all(sv_straight['element_type'] == [ + 'Marker', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThinSliceRBendExit', 'Marker', 'Drift', 'Marker', '']) + +xo.assert_allclose( + sv_straight['angle'], + np.array([ + 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0.3, 0. , 0. , + 0. , 0. ]), + atol=1e-12 +) + +xo.assert_allclose(sv_straight['s'], np.array([ + 0. , 0. , 1. , 1. , 1. , + 1.5075795 , 2.01515901, 2.52273851, 3.03031802, 3.53789752, + 4.04547703, 4.04547703, 4.04547703, 5. , 5. ] +), atol=1e-5) + +xo.assert_allclose( + sv_straight['rot_s_rad'], 0, atol=1e-14) + +sv_straight.cols['X Y Z'] +xo.assert_allclose(sv_straight['Y'], 0, atol=1e-14) +xo.assert_allclose(sv_straight['Z'], np.array( + [0. , 0. , 1. , 1. , 1. , + 1.5 , 2. , 2.5 , 3. , 3.5 , + 4. , 4. , 4. , 4.91189063, 4.91189063]), + atol=1e-8) +xo.assert_allclose(sv_straight['X'], np.array([ + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + -1.38777878e-17, -1.38777878e-17, -1.38777878e-17, -1.38777878e-17, + -1.38777878e-17, -1.38777878e-17, -1.38777878e-17, -4.53405654e-01, + -4.53405654e-01, -7.35486481e-01, -7.35486481e-01]), + atol=1e-8) + + +sv_straight.cols['theta phi psi'] +xo.assert_allclose(sv_straight['phi'], 0, atol=1e-14) +xo.assert_allclose(sv_straight['psi'], 0, atol=1e-14) +xo.assert_allclose(sv_straight['theta'], np.array([ + 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , + -0.3, -0.3, -0.3, -0.3])) + + +sv_curved.cols['s element_type angle'] +assert np.all(sv_curved['name'] == [ + 'start', '||drift_1', 'mb_entry', 'mb..entry_map', 'mb..0', + 'mb..1', 'mb..2', 'mb..3', 'mb..4', 'mb..5', 'mb..exit_map', + 'mb_exit', '||drift_2', 'end', '_end_point']) + +assert np.all(sv_curved['element_type'] == [ + 'Marker', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThinSliceRBendExit', 'Marker', 'Drift', 'Marker', '']) + +xo.assert_allclose( + sv_curved['angle'], + np.array([ + 0. , 0. , 0. , 0. , 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0. , + 0. , 0. , 0. , 0. ]), + atol=1e-8 +) + +xo.assert_allclose(sv_curved['s'], np.array([ + 0. , 0. , 1. , 1. , 1. , + 1.5075795 , 2.01515901, 2.52273851, 3.03031802, 3.53789752, + 4.04547703, 4.04547703, 4.04547703, 5. , 5. ] +), atol=1e-5) + +xo.assert_allclose( + sv_curved['rot_s_rad'], 0, atol=1e-14) + +sv_curved.cols['X Y Z'] +xo.assert_allclose(sv_curved['Y'], 0, atol=1e-14) +xo.assert_allclose(sv_curved['Z'], np.array([ + 0. , 0. , 1. , 1. , 1. , 1.507368, + 2.013468, 2.517035, 3.01681 , 3.511544, 4. , 4. , + 4. , 4.911891, 4.911891]), atol=1e-5) +xo.assert_allclose(sv_curved['X'], np.array( + [ 0. , 0. , 0. , 0. , 0. , + -0.01268684, -0.05071567, -0.11399141, -0.20235593, -0.31558835, + -0.45340565, -0.45340565, -0.45340565, -0.73548648, -0.73548648]), + atol=1e-8) + +sv_curved.cols['theta phi psi'] +xo.assert_allclose(sv_curved['phi'], 0, atol=1e-14) +xo.assert_allclose(sv_curved['psi'], 0, atol=1e-14) +xo.assert_allclose(sv_curved['theta'], np.array([ + 0. , 0. , 0. , 0. , 0. , -0.05, -0.1 , -0.15, -0.2 , + -0.25, -0.3 , -0.3 , -0.3 , -0.3 , -0.3 ], + ), atol=1e-8) + +for nn in ['start', 'end']: + xo.assert_allclose(sv_straight['X', nn], sv_curved['X', nn], atol=1e-14) + xo.assert_allclose(sv_straight['Y', nn], sv_curved['Y', nn], atol=1e-14) + xo.assert_allclose(sv_straight['Z', nn], sv_curved['Z', nn], atol=1e-14) + xo.assert_allclose(sv_straight['theta', nn], sv_curved['theta', nn], atol=1e-14) + xo.assert_allclose(sv_straight['phi', nn], sv_curved['phi', nn], atol=1e-14) + xo.assert_allclose(sv_straight['psi', nn], sv_curved['psi', nn], atol=1e-14) + +xo.assert_allclose(sv_straight_start['X'], sv_straight['X'], atol=1e-14) +xo.assert_allclose(sv_straight_start['Y'], sv_straight['Y'], atol=1e-14) +xo.assert_allclose(sv_straight_start['Z'], sv_straight['Z'], atol=1e-14) +xo.assert_allclose(sv_straight_start['theta'], sv_straight['theta'], atol=1e-14) +xo.assert_allclose(sv_straight_start['phi'], sv_straight['phi'], atol=1e-14) +xo.assert_allclose(sv_straight_start['psi'], sv_straight['psi'], atol=1e-14) + +xo.assert_allclose(sv_straight_end['X'], sv_straight['X'], atol=1e-14) +xo.assert_allclose(sv_straight_end['Y'], sv_straight['Y'], atol=1e-14) +xo.assert_allclose(sv_straight_end['Z'], sv_straight['Z'], atol=1e-14) +xo.assert_allclose(sv_straight_end['theta'], sv_straight['theta'], atol=1e-14) +xo.assert_allclose(sv_straight_end['phi'], sv_straight['phi'], atol=1e-14) +xo.assert_allclose(sv_straight_end['psi'], sv_straight['psi'], atol=1e-14) + +xo.assert_allclose(sv_curved_start['X'], sv_curved['X'], atol=1e-14) +xo.assert_allclose(sv_curved_start['Y'], sv_curved['Y'], atol=1e-14) +xo.assert_allclose(sv_curved_start['Z'], sv_curved['Z'], atol=1e-14) +xo.assert_allclose(sv_curved_start['theta'], sv_curved['theta'], atol=1e-14) +xo.assert_allclose(sv_curved_start['phi'], sv_curved['phi'], atol=1e-14) +xo.assert_allclose(sv_curved_start['psi'], sv_curved['psi'], atol=1e-14) + +xo.assert_allclose(sv_curved_end['X'], sv_curved['X'], atol=1e-14) +xo.assert_allclose(sv_curved_end['Y'], sv_curved['Y'], atol=1e-14) +xo.assert_allclose(sv_curved_end['Z'], sv_curved['Z'], atol=1e-14) +xo.assert_allclose(sv_curved_end['theta'], sv_curved['theta'], atol=1e-14) +xo.assert_allclose(sv_curved_end['phi'], sv_curved['phi'], atol=1e-14) +xo.assert_allclose(sv_curved_end['psi'], sv_curved['psi'], atol=1e-14) +xo.assert_allclose(tw_straight['X', 'start'], 0, atol=1e-14) +xo.assert_allclose(tw_straight['Y', 'start'], 0, atol=1e-14) +xo.assert_allclose(tw_straight['Z', 'start'], 0, atol=1e-14) +xo.assert_allclose(tw_curved['X', 'start'], 0, atol=1e-14) +xo.assert_allclose(tw_curved['Y', 'start'], 0, atol=1e-14) +xo.assert_allclose(tw_curved['Z', 'start'], 0, atol=1e-14) +xo.assert_allclose(tw_straight['X', 'mb_entry'], tw_curved['X', 'mb_entry'], atol=1e-14) +xo.assert_allclose(tw_straight['Y', 'mb_entry'], tw_curved['Y', 'mb_entry'], atol=1e-14) +xo.assert_allclose(tw_straight['Z', 'mb_entry'], tw_curved['Z', 'mb_entry'], atol=1e-14) +xo.assert_allclose(tw_straight['X', 'mb_exit'], tw_curved['X', 'mb_exit'], atol=1e-14) +xo.assert_allclose(tw_straight['Y', 'mb_exit'], tw_curved['Y', 'mb_exit'], atol=1e-14) +xo.assert_allclose(tw_straight['Z', 'mb_exit'], tw_curved['Z', 'mb_exit'], atol=1e-14) + +xo.assert_allclose(tw_straight['x', 'mb_entry'], 0, atol=1e-14) +xo.assert_allclose(tw_straight['y', 'mb_entry'], 0, atol=1e-14) +xo.assert_allclose(tw_straight['x', 'mb_exit'], 0, atol=1e-14) +xo.assert_allclose(tw_straight['y', 'mb_exit'], 0, atol=1e-14) +xo.assert_allclose(tw_curved['x', 'mb_entry'], 0, atol=1e-14) +xo.assert_allclose(tw_curved['y', 'mb_entry'], 0, atol=1e-14) +xo.assert_allclose(tw_curved['x', 'mb_exit'], 0, atol=1e-14) +xo.assert_allclose(tw_curved['y', 'mb_exit'], 0, atol=1e-14) + +xo.assert_allclose(tw_no_slice_curved['x', 'start'], tw_curved['x', 'start'], + atol=1e-14) +xo.assert_allclose(tw_no_slice_curved['y', 'start'], tw_curved['y', 'start'], + atol=1e-14) +xo.assert_allclose(tw_no_slice_curved['x', 'end'], tw_curved['x', 'end'], + atol=1e-14) +xo.assert_allclose(tw_no_slice_curved['y', 'end'], tw_curved['y', 'end'], + atol=1e-14) +xo.assert_allclose(tw_no_slice_straight['x', 'start'], tw_curved['x', 'start'], + atol=1e-14) +xo.assert_allclose(tw_no_slice_straight['y', 'start'], tw_curved['y', 'start'], + atol=1e-14) +xo.assert_allclose(tw_no_slice_straight['x', 'end'], tw_curved['x', 'end'], + atol=1e-14) +xo.assert_allclose(tw_no_slice_straight['y', 'end'], tw_curved['y', 'end'], + atol=1e-14) + +for nn in ['start', 'end']: + # Compare no_slice survey vs curved survey + xo.assert_allclose(sv_no_slice_curved_start['X', nn], sv_curved['X', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['Y', nn], sv_curved['Y', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['Z', nn], sv_curved['Z', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['theta', nn], sv_curved['theta', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['phi', nn], sv_curved['phi', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['psi', nn], sv_curved['psi', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['s', nn], sv_curved['s', nn], atol=1e-14) + + xo.assert_allclose(sv_no_slice_curved_end['X', nn], sv_curved['X', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['Y', nn], sv_curved['Y', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['Z', nn], sv_curved['Z', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['theta', nn], sv_curved['theta', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['phi', nn], sv_curved['phi', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['psi', nn], sv_curved['psi', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['s', nn], sv_curved['s', nn], atol=1e-14) + +import matplotlib.pyplot as plt +plt.close('all') +sv_straight.plot(projection='ZX') +plt.plot(sv_curved.Z, sv_curved.X, '.-', color='r', alpha=0.7) +plt.suptitle('Straight body') + +sv_curved.plot(projection='ZX') +plt.suptitle('Curved body') + +plt.show() \ No newline at end of file diff --git a/examples/true_rbend/008_d1_d2_example.py b/examples/true_rbend/008_d1_d2_example.py new file mode 100644 index 000000000..e0d4d936d --- /dev/null +++ b/examples/true_rbend/008_d1_d2_example.py @@ -0,0 +1,352 @@ +import xtrack as xt +import numpy as np +import xobjects as xo + +edge_model = 'full' # linear of full + +# TODO check angle column in survey + +env = xt.Environment() +env.vars.default_to_zero = True +line = env.new_line(compose=True) +line.new('start', 'Marker', at=0.) +line.new('d1a', 'RBend', length_straight=1.0, k0='k0d1a', anchor='start', at='dz_d1a') +line.new('d1b', 'RBend', length_straight=1.0, k0='k0d1b', anchor='start', at='dz_d1b') +line.new('d2', 'RBend', length_straight=1.0, k0='k0d2', anchor='start', at='dz_d2') +line.new('end', 'Marker', at='dz_end') + + # ------ measure geometry in the straight reference frame ------ + +# Positions in the straight reference frame +env['dz_d1a'] = 1. +env['dz_d1b'] = 3. +env['dz_d2'] = 8. +env['dz_end'] = 10. + +line.end_compose() +line.set_particle_ref('proton', p0c=1e9) +line.configure_drift_model('exact') +line.set(env.elements.get_table().rows.match(element_type='RBend'), + model='bend-kick-bend', edge_entry_model=edge_model, + edge_exit_model=edge_model) + +env['k0d1a'] = 'k0d1' +env['k0d1b'] = 'k0d1' + +opt = line.match( + solve=False, + betx=1, bety=1, + vary=[xt.VaryList(['k0d1', 'k0d2'], step=1e-5)], + targets=xt.TargetSet(x=1., px=0.0, at='end'), +) +opt.solve() + +# ---- build geometry with curved reference frame ---- + +# Twiss in the straight reference system +tw0 = line.twiss(betx=1, bety=1, strengths=True) + +if edge_model == 'linear': + # Set fdown angles to match the trajectory (used only for linear edges) + for nn in ['d1a', 'd1b', 'd2']: + line[nn].edge_entry_angle_fdown = np.arcsin(tw0['px', nn]) + line[nn].edge_exit_angle_fdown = -np.arcsin(tw0['px', nn + '>>1']) + tw0 = line.twiss(betx=1, bety=1, strengths=True) + for nn in ['d1a', 'd1b', 'd2']: + line[nn].edge_entry_angle_fdown = 0 + line[nn].edge_exit_angle_fdown = 0 + +line.regenerate_from_composer() + +# Update positions according to path length +env['dz_d1a'] = tw0['s', 'd1a'] - tw0['zeta', 'd1a'] +env['dz_d1b'] = tw0['s', 'd1b'] - tw0['zeta', 'd1b'] +env['dz_d2'] = tw0['s', 'd2'] - tw0['zeta', 'd2'] +env['dz_end'] = tw0['s', 'end'] - tw0['zeta', 'end'] + +# Introduce magnet curvatures +for nn in ['d1a', 'd1b', 'd2']: + line[nn].k0 = 0 + line[nn].k0_from_h = True + line[nn].rbend_compensate_sagitta = False + line[nn].rbend_model = 'straight-body' + +d1a_angle_in = np.arcsin(tw0['px', 'd1a']) +d1b_angle_in = np.arcsin(tw0['px', 'd1b']) +d2_angle_in = np.arcsin(tw0['px', 'd2']) +d1a_angle_out = -d1b_angle_in +d1b_angle_out = -d2_angle_in +d2_angle_out = -np.arcsin(tw0['px', 'end']) + +line['d1a'].angle = d1a_angle_in + d1a_angle_out +line['d1b'].angle = d1b_angle_in + d1b_angle_out +line['d2'].angle = d2_angle_in + d2_angle_out + +line['d1a'].rbend_angle_diff = d1a_angle_out - d1a_angle_in +line['d1b'].rbend_angle_diff = d1b_angle_out - d1b_angle_in +line['d2'].rbend_angle_diff = d2_angle_out - d2_angle_in + +# Set rbend shifts +line['d1a'].rbend_shift += line['d1a']._x0_in - tw0['x', 'd1a'] +line['d1b'].rbend_shift += line['d1b']._x0_in - tw0['x', 'd1b'] +line['d2'].rbend_shift += line['d2']._x0_out - tw0['x', 'end'] # to illustrate that out can be set as well + +line.end_compose() + +sv = line.survey() +tw = line.twiss(betx=1, bety=1) +sv_back = line.survey(element0='end', X0=sv.X[-1], Y0=sv.Y[-1], Z0=sv.Z[-1], + theta0=sv.theta[-1], phi0=sv.phi[-1], psi0=sv.psi[-1]) +if edge_model == 'linear': + tw_back = line.twiss(init_at='end', init=tw.get_twiss_init('end')) + +# slice for plot +l_sliced =line.copy(shallow=True) +l_sliced.slice_thick_elements( + slicing_strategies=[ + xt.Strategy(slicing=xt.Uniform(3, mode='thick')) + ]) + +sv_sliced = l_sliced.survey() +tw_sliced = l_sliced.twiss(betx=1, bety=1) +sv_sliced_back = l_sliced.survey(element0='end', + X0=sv_sliced.X[-1], Y0=sv_sliced.Y[-1], Z0=sv_sliced.Z[-1], + theta0=sv_sliced.theta[-1], phi0=sv_sliced.phi[-1], + psi0=sv_sliced.psi[-1]) +if edge_model == 'linear': + tw_sliced_back = l_sliced.twiss(init_at='end', + init=tw_sliced.get_twiss_init('end')) + +# Combine twiss and survey to get actual trajectory +trajectory = sv_sliced.p0 + tw_sliced.x[:, None] * sv_sliced.ex + tw_sliced.y[:, None] * sv_sliced.ey + +tw0['path_length'] = tw0.s - tw0.zeta +tw0['diff_path_length'] = np.diff(tw0.path_length, append=tw0.path_length[-1]) + +xo.assert_allclose(tw0.path_length, tw.s, atol=1e-14) + +xo.assert_allclose(tw0['diff_path_length', 'd1a'], line['d1a'].length, atol=1e-14) +xo.assert_allclose(tw0['diff_path_length', 'd1b'], line['d1b'].length, atol=1e-14) +xo.assert_allclose(tw0['diff_path_length', 'd2'], line['d2'].length, atol=1e-14) + +xo.assert_allclose(tw0['px', 'd1a'], 0, atol=1e-14) +xo.assert_allclose(tw0['px', 'd1b'], np.sin(line['d1b']._angle_in), atol=1e-14) +xo.assert_allclose(tw0['px', 'd2'], np.sin(line['d2']._angle_in), atol=1e-14) + +xo.assert_allclose(tw0['px', 'd1b'], -np.sin(line['d1a']._angle_out), atol=1e-14) +xo.assert_allclose(tw0['px', 'd2'], -np.sin(line['d1b']._angle_out), atol=1e-14) +xo.assert_allclose(tw0['px', 'end'], -np.sin(line['d2']._angle_out), atol=1e-14) + +xo.assert_allclose(tw0['x', 'd1a'], line['d1a']._x0_in, atol=1e-14) +xo.assert_allclose(tw0['x', 'd1b'], line['d1b']._x0_in, atol=1e-14) +xo.assert_allclose(tw0['x', 'd2'], line['d2']._x0_in, atol=1e-14) + +xo.assert_allclose(tw0['x', 'd1a>>1'], line['d1a']._x0_out, atol=1e-14) +xo.assert_allclose(tw0['x', 'd1b>>1'], line['d1b']._x0_out, atol=1e-14) +xo.assert_allclose(tw0['x', 'd2>>1'], line['d2']._x0_out, atol=1e-14) + +assert np.all(sv.element_type == + ['Marker', 'Drift', 'RBend', 'Drift', 'RBend', 'Drift', 'RBend', + 'Drift', 'Marker', '']) +xo.assert_allclose(sv.angle, np.array([ + 0. , 0. , -0.08249992, 0. , -0.08306823, + 0. , 0.16556815, 0. , 0. , 0. ]), + rtol=1e-7) + +xo.assert_allclose(sv.Z, tw0.s, atol=0, rtol=5e-9) +xo.assert_allclose(sv.X, tw0.x, atol=0, rtol=3e-8) +xo.assert_allclose(sv.Y, tw0.y, atol=1e-14) +xo.assert_allclose(sv.theta, np.arcsin(tw0.px), atol=1e-14) +xo.assert_allclose(sv.psi, 0., atol=1e-14) +xo.assert_allclose(sv.phi, 0., atol=1e-14) + +xo.assert_allclose(tw.x, 0, atol=1e-14) +xo.assert_allclose(tw.zeta, 0, atol=1e-14) +xo.assert_allclose(tw.y, 0, atol=1e-14) + +xo.assert_allclose(tw.betx[-1], tw0.betx[-1], rtol=1e-10) +xo.assert_allclose(tw.bety[-1], tw0.bety[-1], rtol=1e-10) +xo.assert_allclose(tw.alfx[-1], tw0.alfx[-1], rtol=1e-10) +xo.assert_allclose(tw.alfy[-1], tw0.alfy[-1], rtol=1e-10) +xo.assert_allclose(tw.dx[-1], tw0.dx[-1], rtol=1e-10) +xo.assert_allclose(tw.dpx[-1], tw0.dpx[-1], atol=1e-10) + +if edge_model == 'linear': + xo.assert_allclose(tw_back.s, tw.s, atol=1e-14) + xo.assert_allclose(tw_back.x, tw.x, atol=1e-14) + xo.assert_allclose(tw_back.y, tw.y, atol=1e-14) + xo.assert_allclose(tw_back.betx, tw.betx, rtol=1e-10) + xo.assert_allclose(tw_back.bety, tw.bety, rtol=1e-10) + xo.assert_allclose(tw_back.alfx, tw.alfx, atol=1e-8) + xo.assert_allclose(tw_back.alfy, tw.alfy, atol=1e-8) + xo.assert_allclose(tw_back.dx, tw.dx, atol=1e-9) + xo.assert_allclose(tw_back.dpx, tw.dpx, atol=1e-9) + +xo.assert_allclose(sv_back.s, sv.s, atol=1e-14) +xo.assert_allclose(sv_back.X, sv.X, atol=1e-14) +xo.assert_allclose(sv_back.Y, sv.Y, atol=1e-14) +xo.assert_allclose(sv_back.Z, sv.Z, atol=1e-14) +xo.assert_allclose(sv_back.theta, sv.theta, atol=1e-14) +xo.assert_allclose(sv_back.phi, sv.phi, atol=1e-14) +xo.assert_allclose(sv_back.psi, sv.psi, atol=1e-14) +xo.assert_allclose(sv.angle, sv.angle, atol=1e-14) + +sv_sliced.cols['s angle theta X'].show() +# name s angle theta X +# start 0 0 0 0 +# ||drift_1 0 0 0 0 +# d1a_entry 1 0 0 0 +# d1a..entry_map 1 0 0 0 +# d1a..0 1 0 0 0 +# d1a..1 1.33371 0 0 0 +# d1a..2 1.66742 0 0 0 +# d1a..exit_map 2.00114 -0.0824999 0 0 +# d1a_exit 2.00114 0 0.0824999 0.0412734 +# ||drift_3 2.00114 0 0.0824999 0.0412734 +# d1b_entry 3.00455 0 0.0824999 0.123961 +# d1b..entry_map 3.00455 0.0824999 0.0824999 0.123961 +# d1b..0 3.00455 0 -1.14732e-17 -1.38778e-17 +# d1b..1 3.34056 0 -1.14732e-17 -1.77022e-17 +# d1b..2 3.67657 0 -1.14732e-17 -2.15266e-17 +# d1b..exit_map 4.01258 -0.165568 -1.14732e-17 -2.5351e-17 +# d1b_exit 4.01258 0 0.165568 0.248635 +# ||drift_4 4.01258 0 0.165568 0.248635 +# d2_entry 8.06804 0 0.165568 0.917026 +# d2..entry_map 8.06804 0.165568 0.165568 0.917026 +# d2..0 8.06804 0 -9.64632e-18 1.11022e-16 +# d2..1 8.4029 0 -9.64632e-18 1.07807e-16 +# d2..2 8.73776 0 -9.64632e-18 1.04591e-16 +# d2..exit_map 9.07262 1.38778e-17 -9.64632e-18 1.01376e-16 +# d2_exit 9.07262 0 -9.64632e-18 1 +# ||drift_5 9.07262 0 -9.64632e-18 1 +# end 10.0726 0 -9.64632e-18 1 +# _end_point 10.0726 0 -9.64632e-18 1 + +xo.assert_allclose(sv_sliced.s[-1], tw0.path_length[-1], atol=0, rtol=1e-14) +xo.assert_allclose(sv_sliced.X[-1], tw0.x[-1], atol=0, rtol=1e-14) +xo.assert_allclose(sv_sliced.Y, 0, atol=1e-14) + +assert np.all(sv_sliced.element_type == + np.array(['Marker', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThinSliceRBendExit', 'Marker', 'Drift', 'Marker', + 'ThinSliceRBendEntry', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThinSliceRBendExit', 'Marker', 'Drift', + 'Marker', 'ThinSliceRBendEntry', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThinSliceRBendExit', + 'Marker', 'Drift', 'Marker', ''])) + +xo.assert_allclose(sv_sliced.angle, np.array([ + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, -8.24999219e-02, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 8.24999219e-02, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, -1.65568148e-01, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.65568148e-01, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]), + rtol=1e-8, atol=1e-14) + +xo.assert_allclose(sv_sliced['s', 'd1a..entry_map'], sv['s', 'd1a'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd1b..entry_map'], sv['s', 'd1b'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd2..entry_map'], sv['s', 'd2'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd1a..exit_map>>1'], sv['s', 'd1a>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd1b..exit_map>>1'], sv['s', 'd1b>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd2..exit_map>>1'], sv['s', 'd2>>1'], atol=1e-14) + +xo.assert_allclose(sv_sliced['theta', 'd1a..entry_map'], sv['theta', 'd1a'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1b..entry_map'], sv['theta', 'd1b'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd2..entry_map'], sv['theta', 'd2'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1a..exit_map>>1'], sv['theta', 'd1a>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1b..exit_map>>1'], sv['theta', 'd1b>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd2..exit_map>>1'], sv['theta', 'd2>>1'], atol=1e-14) + +xo.assert_allclose(sv_sliced['X', 'd1a..entry_map'], sv['X', 'd1a'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1b..entry_map'], sv['X', 'd1b'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd2..entry_map'], sv['X', 'd2'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1a..exit_map>>1'], sv['X', 'd1a>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1b..exit_map>>1'], sv['X', 'd1b>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd2..exit_map>>1'], sv['X', 'd2>>1'], atol=1e-14) + +xo.assert_allclose(sv_sliced['Z', 'd1a..entry_map'], sv['Z', 'd1a'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd1b..entry_map'], sv['Z', 'd1b'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd2..entry_map'], sv['Z', 'd2'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd1a..exit_map>>1'], sv['Z', 'd1a>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd1b..exit_map>>1'], sv['Z', 'd1b>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd2..exit_map>>1'], sv['Z', 'd2>>1'], atol=1e-14) + +xo.assert_allclose(sv_sliced['theta', 'd1a..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1b..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd2..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1a..exit_map'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1b..exit_map'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd2..exit_map'], 0., atol=1e-14) + +xo.assert_allclose(sv_sliced['X', 'd1a..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1b..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd2..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1a..exit_map'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1b..exit_map'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd2..exit_map'], 0., atol=1e-14) + +xo.assert_allclose(sv_sliced_back.s, sv_sliced.s, atol=1e-14) +xo.assert_allclose(sv_sliced_back.X, sv_sliced.X, atol=1e-14) +xo.assert_allclose(sv_sliced_back.Y, sv_sliced.Y, atol=1e-14) +xo.assert_allclose(sv_sliced_back.Z, sv_sliced.Z, atol=1e-14) +xo.assert_allclose(sv_sliced_back.theta, sv_sliced.theta, atol=1e-14) +xo.assert_allclose(sv_sliced_back.phi, sv_sliced.phi, atol=1e-14) +xo.assert_allclose(sv_sliced_back.psi, sv_sliced.psi, atol=1e-14) +xo.assert_allclose(sv_sliced_back.angle, sv_sliced.angle, atol=1e-14) + +# Twiss checks + +xo.assert_allclose(tw_sliced.s, sv_sliced.s, atol=0, rtol=1e-14) + +xo.assert_allclose(tw_sliced['x', 'd1a..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1b..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd2..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1a..exit_map>>1'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1b..exit_map>>1'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd2..exit_map>>1'], 0, atol=1e-14) + +xo.assert_allclose(tw_sliced['px', 'd1a..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd1b..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd2..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd1a..exit_map>>1'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd1b..exit_map>>1'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd2..exit_map>>1'], 0, atol=1e-14) + +xo.assert_allclose(tw_sliced['x', 'd1a..entry_map>>1'], tw0['x','d1a'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1b..entry_map>>1'], tw0['x','d1b'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd2..entry_map>>1'], tw0['x','d2'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1a..exit_map'], tw0['x','d1a>>1'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1b..exit_map'], tw0['x','d1b>>1'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd2..exit_map'], tw0['x','d2>>1'], atol=1e-14) + +xo.assert_allclose(tw_sliced.betx[-1], tw0.betx[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.bety[-1], tw0.bety[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.alfx[-1], tw0.alfx[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.alfy[-1], tw0.alfy[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.dx[-1], tw0.dx[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.dpx[-1], tw0.dpx[-1], atol=1e-10) + +if edge_model == 'linear': + xo.assert_allclose(tw_sliced_back.s, tw_sliced.s, atol=1e-14) + xo.assert_allclose(tw_sliced_back.x, tw_sliced.x, atol=1e-14) + xo.assert_allclose(tw_sliced_back.y, tw_sliced.y, atol=1e-14) + xo.assert_allclose(tw_sliced_back.betx, tw_sliced.betx, rtol=1e-8) + xo.assert_allclose(tw_sliced_back.bety, tw_sliced.bety, rtol=1e-8) + xo.assert_allclose(tw_sliced_back.alfx, tw_sliced.alfx, atol=1e-8) + xo.assert_allclose(tw_sliced_back.alfy, tw_sliced.alfy, atol=1e-8) + xo.assert_allclose(tw_sliced_back.dx, tw_sliced.dx, atol=1e-9) + xo.assert_allclose(tw_sliced_back.dpx, tw_sliced.dpx, atol=1e-9) + + + +import matplotlib.pyplot as plt +plt.close('all') +tw0.plot('x') +sv_sliced.plot(element_width=4.) +plt.plot(trajectory[:, 2], trajectory[:, 0], color='C1', linestyle='--') +plt.plot(tw0.s, tw0.x, ':', color='C2') + +plt.show() diff --git a/examples/true_rbend/008a_d1_d2_example_vertical.py b/examples/true_rbend/008a_d1_d2_example_vertical.py new file mode 100644 index 000000000..8aceb1634 --- /dev/null +++ b/examples/true_rbend/008a_d1_d2_example_vertical.py @@ -0,0 +1,322 @@ +import xtrack as xt +import numpy as np +import xobjects as xo +import numpy as np + +edge_model = 'full' # linear of full + +env = xt.Environment() +env.vars.default_to_zero = True +line = env.new_line(compose=True) +line.new('start', 'Marker', at=0.) +line.new('d1a', 'RBend', rot_s_rad=np.pi/2, length_straight=1.0, k0='k0d1a', anchor='start', at='dz_d1a') +line.new('d1b', 'RBend', rot_s_rad=np.pi/2, length_straight=1.0, k0='k0d1b', anchor='start', at='dz_d1b') +line.new('d2', 'RBend', rot_s_rad=np.pi/2, length_straight=1.0, k0='k0d2', anchor='start', at='dz_d2') +line.new('end', 'Marker', at='dz_end') + + # ------ measure geometry in the straight reference frame ------ + +# Positions in the straight reference frame +env['dz_d1a'] = 1. +env['dz_d1b'] = 3. +env['dz_d2'] = 8. +env['dz_end'] = 10. + +line.end_compose() +line.set_particle_ref('proton', p0c=1e9) +line.configure_drift_model('exact') +line.set(env.elements.get_table().rows.match(element_type='RBend'), + model='bend-kick-bend', edge_entry_model=edge_model, + edge_exit_model=edge_model) + +env['k0d1a'] = 'k0d1' +env['k0d1b'] = 'k0d1' + +opt = line.match( + solve=False, + betx=1, bety=1, + vary=[xt.VaryList(['k0d1', 'k0d2'], step=1e-5)], + targets=xt.TargetSet(y=1., py=0.0, at='end'), +) +opt.solve() + +# ---- build geometry with curved reference frame ---- + +# Twiss in the straight reference system +tw0 = line.twiss(betx=1, bety=1, strengths=True) + +if edge_model == 'linear': + # Set fdown angles to match the trajectory (used only for linear edges) + for nn in ['d1a', 'd1b', 'd2']: + line[nn].edge_entry_angle_fdown = np.arcsin(tw0['py', nn]) + line[nn].edge_exit_angle_fdown = -np.arcsin(tw0['py', nn + '>>1']) + tw0 = line.twiss(betx=1, bety=1, strengths=True) + for nn in ['d1a', 'd1b', 'd2']: + line[nn].edge_entry_angle_fdown = 0 + line[nn].edge_exit_angle_fdown = 0 + +line.regenerate_from_composer() + +# Update positions according to path length +env['dz_d1a'] = tw0['s', 'd1a'] - tw0['zeta', 'd1a'] +env['dz_d1b'] = tw0['s', 'd1b'] - tw0['zeta', 'd1b'] +env['dz_d2'] = tw0['s', 'd2'] - tw0['zeta', 'd2'] +env['dz_end'] = tw0['s', 'end'] - tw0['zeta', 'end'] + +# Introduce magnet curvatures +for nn in ['d1a', 'd1b', 'd2']: + line[nn].k0 = 0 + line[nn].k0_from_h = True + line[nn].rbend_compensate_sagitta = False + line[nn].rbend_model = 'straight-body' + +d1a_angle_in = np.arcsin(tw0['py', 'd1a']) +d1b_angle_in = np.arcsin(tw0['py', 'd1b']) +d2_angle_in = np.arcsin(tw0['py', 'd2']) +d1a_angle_out = -d1b_angle_in +d1b_angle_out = -d2_angle_in +d2_angle_out = -np.arcsin(tw0['py', 'end']) + +line['d1a'].angle = d1a_angle_in + d1a_angle_out +line['d1b'].angle = d1b_angle_in + d1b_angle_out +line['d2'].angle = d2_angle_in + d2_angle_out + +line['d1a'].rbend_angle_diff = d1a_angle_out - d1a_angle_in +line['d1b'].rbend_angle_diff = d1b_angle_out - d1b_angle_in +line['d2'].rbend_angle_diff = d2_angle_out - d2_angle_in + +# Set rbend shifts +line['d1a'].rbend_shift += line['d1a']._x0_in - tw0['y', 'd1a'] +line['d1b'].rbend_shift += line['d1b']._x0_in - tw0['y', 'd1b'] +line['d2'].rbend_shift += line['d2']._x0_out - tw0['y', 'end'] # to illustrate that out can be set as well + +line.end_compose() + +sv = line.survey() +tw = line.twiss(betx=1, bety=1) +sv_back = line.survey(element0='end', X0=sv.X[-1], Y0=sv.Y[-1], Z0=sv.Z[-1], + theta0=sv.theta[-1], phi0=sv.phi[-1], psi0=sv.psi[-1]) +if edge_model == 'linear': + tw_back = line.twiss(init_at='end', init=tw.get_twiss_init('end')) + +# slice for plot +l_sliced =line.copy(shallow=True) +l_sliced.slice_thick_elements( + slicing_strategies=[ + xt.Strategy(slicing=xt.Uniform(3, mode='thick')) + ]) + +sv_sliced = l_sliced.survey() +tw_sliced = l_sliced.twiss(betx=1, bety=1) +sv_sliced_back = l_sliced.survey(element0='end', + X0=sv_sliced.X[-1], Y0=sv_sliced.Y[-1], Z0=sv_sliced.Z[-1], + theta0=sv_sliced.theta[-1], phi0=sv_sliced.phi[-1], + psi0=sv_sliced.psi[-1]) +if edge_model == 'linear': + tw_sliced_back = l_sliced.twiss(init_at='end', + init=tw_sliced.get_twiss_init('end')) + +# Combine twiss and survey to get actual trajectory +trajectory = sv_sliced.p0 + tw_sliced.x[:, None] * sv_sliced.ex + tw_sliced.y[:, None] * sv_sliced.ey + +tw0['path_length'] = tw0.s - tw0.zeta +tw0['diff_path_length'] = np.diff(tw0.path_length, append=tw0.path_length[-1]) + +xo.assert_allclose(tw0.path_length, tw.s, atol=1e-14) + +xo.assert_allclose(tw0['diff_path_length', 'd1a'], line['d1a'].length, atol=1e-14) +xo.assert_allclose(tw0['diff_path_length', 'd1b'], line['d1b'].length, atol=1e-14) +xo.assert_allclose(tw0['diff_path_length', 'd2'], line['d2'].length, atol=1e-14) + +xo.assert_allclose(tw0['py', 'd1a'], 0, atol=1e-14) +xo.assert_allclose(tw0['py', 'd1b'], np.sin(line['d1b']._angle_in), atol=1e-14) +xo.assert_allclose(tw0['py', 'd2'], np.sin(line['d2']._angle_in), atol=1e-14) + +xo.assert_allclose(tw0['py', 'd1b'], -np.sin(line['d1a']._angle_out), atol=1e-14) +xo.assert_allclose(tw0['py', 'd2'], -np.sin(line['d1b']._angle_out), atol=1e-14) +xo.assert_allclose(tw0['py', 'end'], -np.sin(line['d2']._angle_out), atol=1e-14) + +xo.assert_allclose(tw0['y', 'd1a'], line['d1a']._x0_in, atol=1e-14) +xo.assert_allclose(tw0['y', 'd1b'], line['d1b']._x0_in, atol=1e-14) +xo.assert_allclose(tw0['y', 'd2'], line['d2']._x0_in, atol=1e-14) + +xo.assert_allclose(tw0['y', 'd1a>>1'], line['d1a']._x0_out, atol=1e-14) +xo.assert_allclose(tw0['y', 'd1b>>1'], line['d1b']._x0_out, atol=1e-14) +xo.assert_allclose(tw0['y', 'd2>>1'], line['d2']._x0_out, atol=1e-14) + +assert np.all(sv.element_type == + ['Marker', 'Drift', 'RBend', 'Drift', 'RBend', 'Drift', 'RBend', + 'Drift', 'Marker', '']) +xo.assert_allclose(sv.angle, np.array([ + 0. , 0. , -0.08249992, 0. , -0.08306823, + 0. , 0.16556815, 0. , 0. , 0. ]), + rtol=1e-7) + +xo.assert_allclose(sv.Z, tw0.s, atol=0, rtol=5e-9) +xo.assert_allclose(sv.X, tw0.x, atol=0, rtol=3e-8) +xo.assert_allclose(sv.Y, tw0.y, atol=1e-14) +xo.assert_allclose(sv.theta, 0., atol=1e-14) +xo.assert_allclose(sv.psi, 0., atol=1e-14) +xo.assert_allclose(sv.phi, np.arcsin(tw0.py), atol=1e-14) + +xo.assert_allclose(tw.x, 0, atol=1e-14) +xo.assert_allclose(tw.zeta, 0, atol=1e-14) +xo.assert_allclose(tw.y, 0, atol=1e-14) + +xo.assert_allclose(tw.betx[-1], tw0.betx[-1], rtol=1e-10) +xo.assert_allclose(tw.bety[-1], tw0.bety[-1], rtol=1e-10) +xo.assert_allclose(tw.alfx[-1], tw0.alfx[-1], rtol=1e-10) +xo.assert_allclose(tw.alfy[-1], tw0.alfy[-1], rtol=1e-10) +xo.assert_allclose(tw.dx[-1], tw0.dx[-1], rtol=1e-10) +xo.assert_allclose(tw.dpx[-1], tw0.dpx[-1], atol=1e-10) + +if edge_model == 'linear': + xo.assert_allclose(tw_back.s, tw.s, atol=1e-14) + xo.assert_allclose(tw_back.x, tw.x, atol=1e-14) + xo.assert_allclose(tw_back.y, tw.y, atol=1e-14) + xo.assert_allclose(tw_back.betx, tw.betx, rtol=1e-10) + xo.assert_allclose(tw_back.bety, tw.bety, rtol=1e-10) + xo.assert_allclose(tw_back.alfx, tw.alfx, atol=1e-8) + xo.assert_allclose(tw_back.alfy, tw.alfy, atol=1e-8) + xo.assert_allclose(tw_back.dx, tw.dx, atol=1e-9) + xo.assert_allclose(tw_back.dpx, tw.dpx, atol=1e-9) + +xo.assert_allclose(sv_back.s, sv.s, atol=1e-14) +xo.assert_allclose(sv_back.X, sv.X, atol=1e-14) +xo.assert_allclose(sv_back.Y, sv.Y, atol=1e-14) +xo.assert_allclose(sv_back.Z, sv.Z, atol=1e-14) +xo.assert_allclose(sv_back.theta, sv.theta, atol=1e-14) +xo.assert_allclose(sv_back.phi, sv.phi, atol=1e-14) +xo.assert_allclose(sv_back.psi, sv.psi, atol=1e-14) +xo.assert_allclose(sv.angle, sv.angle, atol=1e-14) + +sv_sliced.cols['s angle theta X'].show() + +xo.assert_allclose(sv_sliced.s[-1], tw0.path_length[-1], atol=0, rtol=1e-14) +xo.assert_allclose(sv_sliced.Y[-1], tw0.y[-1], atol=0, rtol=1e-14) +xo.assert_allclose(sv_sliced.X, 0, atol=1e-14) + +assert np.all(sv_sliced.element_type == + np.array(['Marker', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThinSliceRBendExit', 'Marker', 'Drift', 'Marker', + 'ThinSliceRBendEntry', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThinSliceRBendExit', 'Marker', 'Drift', + 'Marker', 'ThinSliceRBendEntry', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThinSliceRBendExit', + 'Marker', 'Drift', 'Marker', ''])) + +xo.assert_allclose(sv_sliced.angle, np.array([ + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, -8.24999219e-02, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 8.24999219e-02, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, -1.65568148e-01, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.65568148e-01, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]), + rtol=1e-8, atol=1e-14) + +xo.assert_allclose(sv_sliced['s', 'd1a..entry_map'], sv['s', 'd1a'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd1b..entry_map'], sv['s', 'd1b'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd2..entry_map'], sv['s', 'd2'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd1a..exit_map>>1'], sv['s', 'd1a>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd1b..exit_map>>1'], sv['s', 'd1b>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['s', 'd2..exit_map>>1'], sv['s', 'd2>>1'], atol=1e-14) + +xo.assert_allclose(sv_sliced['theta', 'd1a..entry_map'], sv['theta', 'd1a'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1b..entry_map'], sv['theta', 'd1b'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd2..entry_map'], sv['theta', 'd2'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1a..exit_map>>1'], sv['theta', 'd1a>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1b..exit_map>>1'], sv['theta', 'd1b>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd2..exit_map>>1'], sv['theta', 'd2>>1'], atol=1e-14) + +xo.assert_allclose(sv_sliced['X', 'd1a..entry_map'], sv['X', 'd1a'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1b..entry_map'], sv['X', 'd1b'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd2..entry_map'], sv['X', 'd2'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1a..exit_map>>1'], sv['X', 'd1a>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1b..exit_map>>1'], sv['X', 'd1b>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd2..exit_map>>1'], sv['X', 'd2>>1'], atol=1e-14) + +xo.assert_allclose(sv_sliced['Z', 'd1a..entry_map'], sv['Z', 'd1a'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd1b..entry_map'], sv['Z', 'd1b'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd2..entry_map'], sv['Z', 'd2'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd1a..exit_map>>1'], sv['Z', 'd1a>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd1b..exit_map>>1'], sv['Z', 'd1b>>1'], atol=1e-14) +xo.assert_allclose(sv_sliced['Z', 'd2..exit_map>>1'], sv['Z', 'd2>>1'], atol=1e-14) + +xo.assert_allclose(sv_sliced['theta', 'd1a..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1b..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd2..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1a..exit_map'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd1b..exit_map'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['theta', 'd2..exit_map'], 0., atol=1e-14) + +xo.assert_allclose(sv_sliced['X', 'd1a..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1b..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd2..entry_map>>1'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1a..exit_map'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd1b..exit_map'], 0., atol=1e-14) +xo.assert_allclose(sv_sliced['X', 'd2..exit_map'], 0., atol=1e-14) + +xo.assert_allclose(sv_sliced_back.s, sv_sliced.s, atol=1e-14) +xo.assert_allclose(sv_sliced_back.X, sv_sliced.X, atol=1e-14) +xo.assert_allclose(sv_sliced_back.Y, sv_sliced.Y, atol=1e-14) +xo.assert_allclose(sv_sliced_back.Z, sv_sliced.Z, atol=1e-14) +xo.assert_allclose(sv_sliced_back.theta, sv_sliced.theta, atol=1e-14) +xo.assert_allclose(sv_sliced_back.phi, sv_sliced.phi, atol=1e-14) +xo.assert_allclose(sv_sliced_back.psi, sv_sliced.psi, atol=1e-14) +xo.assert_allclose(sv_sliced_back.angle, sv_sliced.angle, atol=1e-14) + +# Twiss checks + +xo.assert_allclose(tw_sliced.s, sv_sliced.s, atol=0, rtol=1e-14) + +xo.assert_allclose(tw_sliced['x', 'd1a..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1b..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd2..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1a..exit_map>>1'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1b..exit_map>>1'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd2..exit_map>>1'], 0, atol=1e-14) + +xo.assert_allclose(tw_sliced['px', 'd1a..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd1b..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd2..entry_map'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd1a..exit_map>>1'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd1b..exit_map>>1'], 0, atol=1e-14) +xo.assert_allclose(tw_sliced['px', 'd2..exit_map>>1'], 0, atol=1e-14) + +xo.assert_allclose(tw_sliced['x', 'd1a..entry_map>>1'], tw0['x','d1a'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1b..entry_map>>1'], tw0['x','d1b'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd2..entry_map>>1'], tw0['x','d2'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1a..exit_map'], tw0['x','d1a>>1'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd1b..exit_map'], tw0['x','d1b>>1'], atol=1e-14) +xo.assert_allclose(tw_sliced['x', 'd2..exit_map'], tw0['x','d2>>1'], atol=1e-14) + +xo.assert_allclose(tw_sliced.betx[-1], tw0.betx[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.bety[-1], tw0.bety[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.alfx[-1], tw0.alfx[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.alfy[-1], tw0.alfy[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.dx[-1], tw0.dx[-1], rtol=1e-10) +xo.assert_allclose(tw_sliced.dpx[-1], tw0.dpx[-1], atol=1e-10) + +if edge_model == 'linear': + xo.assert_allclose(tw_sliced_back.s, tw_sliced.s, atol=1e-14) + xo.assert_allclose(tw_sliced_back.x, tw_sliced.x, atol=1e-14) + xo.assert_allclose(tw_sliced_back.y, tw_sliced.y, atol=1e-14) + xo.assert_allclose(tw_sliced_back.betx, tw_sliced.betx, rtol=1e-8) + xo.assert_allclose(tw_sliced_back.bety, tw_sliced.bety, rtol=1e-8) + xo.assert_allclose(tw_sliced_back.alfx, tw_sliced.alfx, atol=1e-8) + xo.assert_allclose(tw_sliced_back.alfy, tw_sliced.alfy, atol=1e-8) + xo.assert_allclose(tw_sliced_back.dx, tw_sliced.dx, atol=1e-9) + xo.assert_allclose(tw_sliced_back.dpx, tw_sliced.dpx, atol=1e-9) + + + +import matplotlib.pyplot as plt +plt.close('all') +tw0.plot('y') +sv_sliced.plot(element_width=4., projection='ZY') +plt.plot(trajectory[:, 2], trajectory[:, 1], color='C1', linestyle='--') +plt.plot(tw0.s, tw0.y, ':', color='C2') + +plt.show() diff --git a/examples/true_rbend/t000d_backtrack.py b/examples/true_rbend/t000d_backtrack.py new file mode 100644 index 000000000..b7345bb2d --- /dev/null +++ b/examples/true_rbend/t000d_backtrack.py @@ -0,0 +1,38 @@ +import xtrack as xt +import numpy as np +import xobjects as xo + +edge_model = 'linear' + +# b_ref = xt.RBend(angle=0.1, k0_from_h=True, length_straight=3., rbend_angle_diff=0.1/2) +# b_ref.edge_entry_model = edge_model +# b_ref.edge_exit_model = edge_model +# b_ref.model = 'rot-kick-rot' +# b_ref.num_multipole_kicks = 100 +# l_ref = xt.Line([b_ref]) +# l_ref.append('end', xt.Marker()) +# l_ref.particle_ref = xt.Particles(p0c=10e9) +# tw_ref0 = l_ref.twiss(betx=1, bety=1) +# tw_ref = l_ref.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + +b_test = xt.RBend( + angle=0.1, k0_from_h=True, length_straight=3, rbend_angle_diff=0.1) +b_test.rbend_model = 'straight-body' +b_test.model = 'bend-kick-bend' +b_test.num_multipole_kicks = 100 +b_test.edge_entry_model = edge_model +b_test.edge_exit_model = edge_model +l_test = xt.Line([b_test]) +l_test.append('end', xt.Marker()) +l_test.particle_ref = xt.Particles(p0c=10e9) + +p = l_test.particle_ref.copy() +p.x = 2e-3 +p.px = 1e-3 +p.y = 2e-3 +p.py = 2e-3 +p.delta = 1e-3 + +l_test.track(p) +print("\n\n") +l_test.track(p, backtrack=True) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 966ba860e..72e1fab9e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,67 @@ [build-system] build-backend = 'setuptools.build_meta' requires = [ - 'setuptools >= 43.0.0', - 'numpy', + 'setuptools>=77.0.0', +] + +[project] +name = "xtrack" +dynamic = ["version"] +description = "Tracking library for particle accelerators" +readme = "README.md" +requires-python = ">=3.9" +license = "Apache-2.0" +authors = [ + { name = "G. Iadarola et al." } +] +dependencies = [ + "numpy>=1.0", + "pandas>=2.0", + "scipy", + "tqdm", + "requests", + "xobjects", + "xdeps", +] + +[project.urls] +"Homepage" = "https://xsuite.readthedocs.io/" +"Bug Tracker" = "https://github.com/xsuite/xsuite/issues" +"Documentation" = "https://xsuite.readthedocs.io/" +"Source Code" = "https://github.com/xsuite/xtrack" +"Download" = "https://pypi.python.org/pypi/xtrack" + +[project.optional-dependencies] +tests = [ + "cpymad", + "nafflib", + "PyHEADTAIL", + "pytest", + "pytest-mock", + "pymadng", + "requests-mock", + "tfs-pandas", +] +notebooks = [ + "jupyter", + "ipympl", + "xplt" +] + +[tool.setuptools.packages.find] +where = ["."] +include = ["xtrack", "ducktrack"] + +[tool.setuptools.dynamic] +version = { attr = "xtrack._version.__version__" } + +[tool.setuptools] +include-package-data = true + +[project.entry-points.xobjects] +include = "xtrack" + +[pytest] +markers = [ + "context_dependent: marks test as one that depends on the execution context", ] diff --git a/setup.py b/setup.py deleted file mode 100644 index ff4151360..000000000 --- a/setup.py +++ /dev/null @@ -1,56 +0,0 @@ -# copyright ############################### # -# This file is part of the Xtrack Package. # -# Copyright (c) CERN, 2021. # -# ######################################### # - -from setuptools import setup, find_packages, Extension -from pathlib import Path - -####################################### -# Prepare list of compiled extensions # -####################################### - -extensions = [] - -######### -# Setup # -######### - -version_file = Path(__file__).parent / 'xtrack/_version.py' -dd = {} -with open(version_file.absolute(), 'r') as fp: - exec(fp.read(), dd) -__version__ = dd['__version__'] - -setup( - name='xtrack', - version=__version__, - description='Tracking library for particle accelerators', - long_description='Tracking library for particle accelerators', - url='https://xsuite.readthedocs.io/', - author='G. Iadarola et al.', - license='Apache 2.0', - download_url="https://pypi.python.org/pypi/xtrack", - project_urls={ - "Bug Tracker": "https://github.com/xsuite/xsuite/issues", - "Documentation": 'https://xsuite.readthedocs.io/', - "Source Code": "https://github.com/xsuite/xtrack", - }, - packages=find_packages(), - ext_modules = extensions, - include_package_data=True, - install_requires=[ - 'numpy>=1.0', - "pandas>=2.0", - 'scipy', - 'tqdm', - 'requests', - 'xobjects', - 'xdeps' - ], - extras_require={ - 'tests': ['cpymad', 'nafflib', 'PyHEADTAIL', 'pytest', 'pytest-mock', - 'pymadng', 'requests-mock', 'tfs-pandas'], - 'notebooks': ['jupyter', 'ipympl', 'xplt'], - }, -) diff --git a/tests/test_aperture_table.py b/tests/test_aperture_table.py index 2a26192d1..a2ef4f9dd 100644 --- a/tests/test_aperture_table.py +++ b/tests/test_aperture_table.py @@ -61,9 +61,9 @@ def test_aperture_table(): aper_check = aper.rows['veba.20250.a_aper' : 'vebb.20270.b_aper'] assert np.all(aper_check.name == np.array( - ['veba.20250.a_aper', '||drift_979', 'mba.20250', '||drift_991', - 'veba.20250.b_aper', '||drift_671', 'vebb.20270.a_aper', - '||drift_992', 'mbb.20270', '||drift_993', 'vebb.20270.b_aper'])) + ['veba.20250.a_aper', '||drift_976', 'mba.20250', '||drift_988', + 'veba.20250.b_aper', '||drift_668', 'vebb.20270.a_aper', + '||drift_989', 'mbb.20270', '||drift_990', 'vebb.20270.b_aper'])) xo.assert_allclose(aper_check.s, np.array([ 1225.8247 , 1225.8247 , 1226.0176907, 1232.2777093, 1232.4827 , 1232.4827 , 1232.4847 , 1232.4847 , diff --git a/tests/test_atomic_add.py b/tests/test_atomic_add.py deleted file mode 100644 index 97a5afbca..000000000 --- a/tests/test_atomic_add.py +++ /dev/null @@ -1,90 +0,0 @@ -# copyright ############################### # -# This file is part of the Xtrack Package. # -# Copyright (c) CERN, 2025. # -# ######################################### # - -import pytest -import numpy as np - -import xtrack as xt -import xobjects as xo -from xobjects.test_helpers import for_all_test_contexts - - -@for_all_test_contexts -@pytest.mark.parametrize("overload", [True, False], ids=['overload', 'no_overload']) -@pytest.mark.parametrize("ctype", [xo.Int8, xo.Int16, xo.Int32, xo.Int64, - xo.UInt8, xo.UInt16, xo.UInt32, - xo.UInt64, xo.Float32, xo.Float64]) -def test_atomic(overload, ctype, test_context): - if overload: - func_name = 'atomicAdd' - else: - func_name = f'atomicAdd_{ctype.__name__.lower()[0]}{ctype.__name__.split("t")[1]}' - class TestAtomic(xt.BeamElement): - _xofields = {f'val': ctype} - allow_track = False - _extra_c_sources = [f''' - #include - #include - - GPUKERN - void run_atomic_test(TestAtomicData el, GPUGLMEM {ctype._c_type}* increments, - GPUGLMEM {ctype._c_type}* retvals, int length) {{ - VECTORIZE_OVER(ii, length); - GPUGLMEM {ctype._c_type}* val = TestAtomicData_getp_val(el); - {ctype._c_type} ret = {func_name}(val, increments[ii]); - retvals[ii] = ret; - END_VECTORIZE; - }} - '''] - - _kernels = { - 'run_atomic_test': xo.Kernel( - c_name='run_atomic_test', - args=[xo.Arg(xo.ThisClass, name='el'), - xo.Arg(ctype, pointer=True, name='increments'), - xo.Arg(ctype, pointer=True, name='retvals'), - xo.Arg(xo.Int32, name='length')], - n_threads='length') - } - - atomic = TestAtomic(_context=test_context, val=0) - assert atomic.val == 0 - - # Test with all increments = 1, so we can check the return values easily. - num_steps = 10000 - if ctype.__name__.startswith('Int') or ctype.__name__.startswith('Uint'): - # Less steps to avoid overflow - num_steps = min(num_steps, 2**(8*ctype._size - 1) - 1) - increments = test_context.zeros(shape=(num_steps,), dtype=ctype._dtype) + 1 - retvals = test_context.zeros(shape=(num_steps,), dtype=ctype._dtype) - atomic.run_atomic_test(increments=increments, retvals=retvals, length=num_steps) - assert atomic.val == num_steps - retvals = np.sort(test_context.nparray_from_context_array(retvals)) - assert np.allclose(retvals, np.arange(num_steps, dtype=ctype._dtype), atol=1.e-15, rtol=1.e-15) - - # Test with random increments, where we now only can check the total sum - # (retvals can be anything). Watch out: overflow is undefined behaviour, - # except for unsigned integers, so we skip this test for signed integers. - atomic.val = 0 - retvals = test_context.zeros(shape=(num_steps,), dtype=ctype._dtype) - if ctype.__name__.startswith('Uint'): - low = 0 - high = 2**(8*ctype._size) - 1 - increments = np.random.randint(low, high+1, size=num_steps, dtype=ctype._dtype) - increments = test_context.nparray_to_context_array(increments) - atomic.run_atomic_test(increments=increments, retvals=retvals, length=num_steps) - increments = test_context.nparray_from_context_array(increments) - assert atomic.val == (np.sum(increments).item() % (2**(8*ctype._size))) - - elif ctype.__name__.startswith('Float'): - increments = np.zeros(shape=(num_steps,), dtype=ctype._dtype) - increments += np.random.uniform(0, 10, size=num_steps) - increments = test_context.nparray_to_context_array(increments) - atomic.run_atomic_test(increments=increments, retvals=retvals, length=num_steps) - increments = test_context.nparray_from_context_array(increments) - if ctype == xo.Float32: - assert np.isclose(atomic.val, np.sum(increments), atol=10., rtol=1.e-4) - else: - assert np.isclose(atomic.val, np.sum(increments), atol=1.e-6, rtol=1.e-12) diff --git a/tests/test_element_internal_record.py b/tests/test_element_internal_record.py index c61a6d555..1034872ea 100644 --- a/tests/test_element_internal_record.py +++ b/tests/test_element_internal_record.py @@ -26,7 +26,9 @@ class TestElementRecord(xo.HybridClass): extra_src = [] extra_src.append(r''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ // Extract the record and record_index @@ -39,8 +41,7 @@ class TestElementRecord(xo.HybridClass): int64_t n_kicks = TestElementData_get_n_kicks(el); //printf("n_kicks %d\n", (int)n_kicks); - //start_per_particle_block (part0->part) - + START_PER_PARTICLE_BLOCK(part0, part); for (int64_t i = 0; i < n_kicks; i++) { double rr = 1e-6 * RandomUniform_generate(part); LocalParticle_add_to_px(part, rr); @@ -61,9 +62,7 @@ class TestElementRecord(xo.HybridClass): } } } - - - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''') @@ -205,7 +204,9 @@ class TestElementRecord(xo.HybridClass): extra_src = [] extra_src.append(r''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ // Extract the record and record_index @@ -218,8 +219,7 @@ class TestElementRecord(xo.HybridClass): int64_t n_kicks = TestElementData_get_n_kicks(el); //printf("n_kicks %d\n", (int)n_kicks); - //start_per_particle_block (part0->part) - + START_PER_PARTICLE_BLOCK(part0, part); for (int64_t i = 0; i < n_kicks; i++) { // We don't apply the kick, otherwise the twiss fails double rr = 1e-6 * RandomUniform_generate(part); @@ -240,9 +240,7 @@ class TestElementRecord(xo.HybridClass): } } } - - - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''') @@ -308,7 +306,9 @@ class TestElementRecord(xo.HybridClass): extra_src = [] extra_src.append(r''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ // Extract the record and record_index @@ -327,8 +327,7 @@ class TestElementRecord(xo.HybridClass): int64_t n_kicks = TestElementData_get_n_kicks(el); // printf("n_kicks %d\n", (int)n_kicks); - //start_per_particle_block (part0->part) - + START_PER_PARTICLE_BLOCK(part0, part); // Record in table1 info about the ingoing particle if (record){ // Get a slot in table1 @@ -372,8 +371,7 @@ class TestElementRecord(xo.HybridClass): } } } - - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''') @@ -565,7 +563,9 @@ class TestElementRecord(xo.HybridClass): extra_src = [] extra_src.append(r''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ // Extract the record and record_index @@ -584,8 +584,7 @@ class TestElementRecord(xo.HybridClass): int64_t n_kicks = TestElementData_get_n_kicks(el); // printf("n_kicks %d\n", (int)n_kicks); - //start_per_particle_block (part0->part) - + START_PER_PARTICLE_BLOCK(part0, part); // Record in table1 info about the ingoing particle if (record){ // Get a slot in table1 @@ -629,8 +628,7 @@ class TestElementRecord(xo.HybridClass): } } } - - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''') diff --git a/tests/test_elements.py b/tests/test_elements.py index 96da7e068..60b2f1f1b 100644 --- a/tests/test_elements.py +++ b/tests/test_elements.py @@ -542,34 +542,36 @@ def test_exciter(test_context): test_source = r""" -/*gpufun*/ +#include "xtrack/headers/track.h" + +GPUFUN void test_function(TestElementData el, LocalParticle* part0, - /*gpuglmem*/ double* b){ + GPUGLMEM double* b){ double const a = TestElementData_get_a(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); const int64_t ipart = part->ipart; double const val = b[ipart]; LocalParticle_add_to_s(part, val + a); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } -/*gpufun*/ +GPUFUN void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ double const a = TestElementData_get_a(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_set_s(part, a); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } """ diff --git a/tests/test_environment.py b/tests/test_environment.py index d5ccf429e..ea947242b 100644 --- a/tests/test_environment.py +++ b/tests/test_environment.py @@ -482,7 +482,7 @@ def test_assemble_ring(): == np.array([nn+'.cell.2' for nn in tt_cell.name[:-1] if not nn.startswith('||drift')])) assert np.all(tt_cell2.s == tt_cell.s) - tt_cell2_nodrift = tt_cell2.rows[~tt_cell2.rows.mask['\|\|drift.*']] + tt_cell2_nodrift = tt_cell2.rows[~tt_cell2.rows.mask[r'\|\|drift.*']] assert tt_cell2_nodrift.isreplica[:-1].all() assert tt_cell2['parent_name', 'mq.d.l.cell.2'] == 'mq.d.l' assert tt_cell2['parent_name', 'mq.f.l.cell.2'] == 'mq.f.l' @@ -907,7 +907,7 @@ def test_assemble_ring_builders(): assert np.all([nn for nn in tt_cell2.name[:-1] if not nn.startswith('||drift')] == np.array([nn + '.cell.2' for nn in tt_cell.name[:-1] if not nn.startswith('||drift')])) assert np.all(tt_cell2.s == tt_cell.s) - tt_cell2_nodrift = tt_cell2.rows[~tt_cell2.rows.mask['\|\|drift.*']] + tt_cell2_nodrift = tt_cell2.rows[~tt_cell2.rows.mask[r'\|\|drift.*']] assert tt_cell2_nodrift.isreplica[:-1].all() assert tt_cell2['parent_name', 'mq.d.l.cell.2'] == 'mq.d.l' assert tt_cell2['parent_name', 'mq.f.l.cell.2'] == 'mq.f.l' diff --git a/tests/test_native_madloader_sps.py b/tests/test_native_madloader_sps.py index 4021557ba..c32a34b93 100644 --- a/tests/test_native_madloader_sps.py +++ b/tests/test_native_madloader_sps.py @@ -115,5 +115,5 @@ def test_native_madloader_sps(): xo.assert_allclose(twtest.wy_chrom[-1], twref.wy_chrom[-1], rtol=1e-4, atol=0) xo.assert_allclose(twtest.ax_chrom[-1], twref.ax_chrom[-1], rtol=1e-4, atol=0) xo.assert_allclose(twtest.ay_chrom[-1], twref.ay_chrom[-1], rtol=1e-4, atol=0) - xo.assert_allclose(twtest.dqx, twref.dqx, rtol=1e-4, atol=0) - xo.assert_allclose(twtest.dqy, twref.dqy, rtol=1e-4, atol=0) + xo.assert_allclose(twtest.dqx, twref.dqx, rtol=2e-4, atol=0) + xo.assert_allclose(twtest.dqy, twref.dqy, rtol=2e-4, atol=0) diff --git a/tests/test_particles.py b/tests/test_particles.py index 5827b04de..372fd2939 100644 --- a/tests/test_particles.py +++ b/tests/test_particles.py @@ -20,18 +20,19 @@ class TestElement(xt.BeamElement): } _extra_c_sources = [""" + #include "xtrack/headers/track.h" #define XT_OMP_SKIP_REORGANIZE - /*gpufun*/ + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0 ) { - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); int64_t state = check_is_active(part); int64_t id = LocalParticle_get_particle_id(part); TestElementData_set_states(el, id, state); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } """] @@ -96,16 +97,18 @@ class TestElement(xt.BeamElement): } _extra_c_sources = [""" - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0 ) { - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); int64_t state = check_is_active(part); int64_t id = LocalParticle_get_particle_id(part); TestElementData_set_states(el, id, state); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } """] diff --git a/tests/test_particles_basics.py b/tests/test_particles_basics.py index 79e537ad3..4217937c1 100644 --- a/tests/test_particles_basics.py +++ b/tests/test_particles_basics.py @@ -345,14 +345,16 @@ class TestElement(xt.BeamElement): 'pz_only': xo.Int64, } _extra_c_sources = [''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ double const value = TestElementData_get_value(el); int const pz_only = (int) TestElementData_get_pz_only(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_add_to_energy(part, value, pz_only); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] @@ -419,13 +421,15 @@ class TestElement(xt.BeamElement): } _extra_c_sources =[''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ double const value = TestElementData_get_value(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_update_delta(part, value); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] @@ -460,13 +464,15 @@ class TestElement(xt.BeamElement): } _extra_c_sources = [''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ double const value = TestElementData_get_value(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_update_ptau(part, value); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] @@ -500,14 +506,16 @@ class TestElement(xt.BeamElement): 'value': xo.Float64, } _extra_c_sources = [''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ double const value = TestElementData_get_value(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); double const pzeta = LocalParticle_get_pzeta(part); LocalParticle_update_pzeta(part, pzeta+value); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] @@ -543,13 +551,15 @@ class TestElement(xt.BeamElement): 'value': xo.Float64, } _extra_c_sources = [''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ double const value = TestElementData_get_value(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_update_p0c(part, value); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] @@ -590,18 +600,20 @@ class ScaleAng(xt.BeamElement): 'scale_y2': xo.Float64, } _extra_c_sources = [''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void ScaleAng_track_local_particle( ScaleAngData el, LocalParticle* part0){ double const scale_x = ScaleAngData_get_scale_x(el); double const scale_y = ScaleAngData_get_scale_y(el); double const scale_x2 = ScaleAngData_get_scale_x2(el); double const scale_y2 = ScaleAngData_get_scale_y2(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_scale_xp(part, scale_x); LocalParticle_scale_yp(part, scale_y); LocalParticle_scale_xp_yp(part, scale_x2, scale_y2); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] @@ -613,18 +625,20 @@ class KickAng(xt.BeamElement): 'kick_y2': xo.Float64, } _extra_c_sources = [''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void KickAng_track_local_particle( KickAngData el, LocalParticle* part0){ double const kick_x = KickAngData_get_kick_x(el); double const kick_y = KickAngData_get_kick_y(el); double const kick_x2 = KickAngData_get_kick_x2(el); double const kick_y2 = KickAngData_get_kick_y2(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_add_to_xp(part, kick_x); LocalParticle_add_to_yp(part, kick_y); LocalParticle_add_to_xp_yp(part, kick_x2, kick_y2); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] @@ -648,18 +662,20 @@ class ScaleAngExact(xt.BeamElement): 'scale_y2': xo.Float64, } _extra_c_sources = [''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void ScaleAngExact_track_local_particle( ScaleAngExactData el, LocalParticle* part0){ double const scale_x = ScaleAngExactData_get_scale_x(el); double const scale_y = ScaleAngExactData_get_scale_y(el); double const scale_x2 = ScaleAngExactData_get_scale_x2(el); double const scale_y2 = ScaleAngExactData_get_scale_y2(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_scale_exact_xp(part, scale_x); LocalParticle_scale_exact_yp(part, scale_y); LocalParticle_scale_exact_xp_yp(part, scale_x2, scale_y2); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] @@ -671,18 +687,20 @@ class KickAngExact(xt.BeamElement): 'kick_y2': xo.Float64, } _extra_c_sources = [''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void KickAngExact_track_local_particle( KickAngExactData el, LocalParticle* part0){ double const kick_x = KickAngExactData_get_kick_x(el); double const kick_y = KickAngExactData_get_kick_y(el); double const kick_x2 = KickAngExactData_get_kick_x2(el); double const kick_y2 = KickAngExactData_get_kick_y2(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); LocalParticle_add_to_exact_xp(part, kick_x); LocalParticle_add_to_exact_yp(part, kick_y); LocalParticle_add_to_exact_xp_yp(part, kick_x2, kick_y2); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''] diff --git a/tests/test_radial_steering.py b/tests/test_radial_steering.py index eff775ae8..480fde4f2 100644 --- a/tests/test_radial_steering.py +++ b/tests/test_radial_steering.py @@ -35,7 +35,7 @@ def test_radial_steering(test_context): dzeta = tw.circumference * df_hz / f_rf - line.unfreeze() + line.discard_tracker() line.append_element(element=xt.ZetaShift(dzeta=dzeta), name='zeta_shift') line.build_tracker() diff --git a/tests/test_random_gen.py b/tests/test_random_gen.py index b3a9d3f2f..378a8ce07 100644 --- a/tests/test_random_gen.py +++ b/tests/test_random_gen.py @@ -31,13 +31,15 @@ class TestElement(xt.BeamElement): _extra_c_sources = [ ''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); double rr = !!GENERATOR!!_generate(part); LocalParticle_set_x(part, rr); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } '''.replace('!!GENERATOR!!', generator) ] diff --git a/tests/test_random_gen_exp.py b/tests/test_random_gen_exp.py index ea6b3ec01..3614fbd1d 100644 --- a/tests/test_random_gen_exp.py +++ b/tests/test_random_gen_exp.py @@ -29,13 +29,15 @@ class TestElement(xt.BeamElement): _extra_c_sources = [ ''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); double rr = RandomExponential_generate(part); LocalParticle_set_x(part, rr); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''' ] diff --git a/tests/test_random_gen_gauss.py b/tests/test_random_gen_gauss.py index 1e92441cf..5579124d9 100644 --- a/tests/test_random_gen_gauss.py +++ b/tests/test_random_gen_gauss.py @@ -28,13 +28,15 @@ class TestElement(xt.BeamElement): _extra_c_sources = [ ''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); double rr = RandomNormal_generate(part); LocalParticle_set_x(part, rr); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''' ] diff --git a/tests/test_random_gen_ruth.py b/tests/test_random_gen_ruth.py index 7139444d1..bd624190b 100644 --- a/tests/test_random_gen_ruth.py +++ b/tests/test_random_gen_ruth.py @@ -36,14 +36,16 @@ class TestElement(xt.BeamElement): _extra_c_sources = [ ''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0){ RandomRutherfordData rng = TestElementData_getp_rng(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); double rr = RandomRutherford_generate(rng, part); LocalParticle_set_x(part, rr); - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''' ] diff --git a/tests/test_rbend_straight_body.py b/tests/test_rbend_straight_body.py index 9cb44afa8..59da69fc5 100644 --- a/tests/test_rbend_straight_body.py +++ b/tests/test_rbend_straight_body.py @@ -3,6 +3,7 @@ import xobjects as xo from cpymad.madx import Madx from xobjects.test_helpers import for_all_test_contexts +import pytest import pathlib @@ -157,6 +158,144 @@ def test_rbend_straight_body_edge_linear(test_context): xo.assert_allclose(tw_test_sliced_back.px, tw_test_sliced0.px, rtol=0, atol=1e-12) xo.assert_allclose(tw_test_sliced_back.py, tw_test_sliced0.py, rtol=0, atol=1e-12) +def test_rbend_straight_body_edge_full_angle_diff(): + + edge_model = 'full' + + b_ref = xt.RBend(angle=0.1, k0_from_h=True, length_straight=3., rbend_angle_diff=0.1) + b_ref.edge_entry_model = edge_model + b_ref.edge_exit_model = edge_model + b_ref.model = 'rot-kick-rot' + b_ref.num_multipole_kicks = 100 + l_ref = xt.Line([b_ref]) + l_ref.append('end', xt.Marker()) + l_ref.particle_ref = xt.Particles(p0c=10e9) + tw_ref0 = l_ref.twiss(betx=1, bety=1, strengths=True) + tw_ref = l_ref.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + + b_test = xt.RBend( + angle=0.1, k0_from_h=True, length_straight=3, rbend_angle_diff=0.1, + rbend_shift=0, rbend_compensate_sagitta=False) + b_test.rbend_model = 'straight-body' + b_test.model = 'bend-kick-bend' + b_test.num_multipole_kicks = 100 + b_test.edge_entry_model = edge_model + b_test.edge_exit_model = edge_model + + b_test.rbend_shift += b_test._x0_in + + l_test = xt.Line([b_test]) + l_test.append('end', xt.Marker()) + l_test.particle_ref = xt.Particles(p0c=10e9) + tw_test0 = l_test.twiss(betx=1, bety=1) + tw_test = l_test.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + + xo.assert_allclose(tw_ref0.betx, tw_test0.betx, rtol=1e-9, atol=0.0) + xo.assert_allclose(tw_ref0.bety, tw_test0.bety, rtol=1e-9, atol=0.0) + xo.assert_allclose(tw_ref0.x, tw_test0.x, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref0.y, tw_test0.y, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref0.s, tw_test0.s, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref0.zeta, tw_test0.zeta, rtol=0, atol=1e-11) + xo.assert_allclose(tw_ref0.px, tw_test0.px, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref0.py, tw_test0.py, rtol=0, atol=1e-12) + + xo.assert_allclose(tw_ref.betx, tw_test.betx, rtol=5e-9, atol=0.0) + xo.assert_allclose(tw_ref.bety, tw_test.bety, rtol=5e-9, atol=0.0) + xo.assert_allclose(tw_ref.x, tw_test.x, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref.y, tw_test.y, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref.s, tw_test.s, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref.zeta, tw_test.zeta, rtol=0, atol=1e-1) + xo.assert_allclose(tw_ref.px, tw_test.px, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref.py, tw_test.py, rtol=0, atol=1e-12) + +def test_rbend_straight_body_edge_linear_angle_diff(): + + edge_model = 'linear' + + b_ref = xt.RBend(angle=0.1, k0_from_h=True, length_straight=3., rbend_angle_diff=0.1) + b_ref.edge_entry_model = edge_model + b_ref.edge_exit_model = edge_model + b_ref.model = 'rot-kick-rot' + b_ref.num_multipole_kicks = 100 + l_ref = xt.Line([b_ref]) + l_ref.append('end', xt.Marker()) + l_ref.particle_ref = xt.Particles(p0c=10e9) + tw_ref0 = l_ref.twiss(betx=1, bety=1) + tw_ref = l_ref.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + + b_test = xt.RBend( + angle=0.1, k0_from_h=True, length_straight=3, rbend_angle_diff=0.1, + rbend_shift=0, rbend_compensate_sagitta=False) + b_test.rbend_model = 'straight-body' + b_test.model = 'bend-kick-bend' + b_test.num_multipole_kicks = 100 + b_test.edge_entry_model = edge_model + b_test.edge_exit_model = edge_model + + b_test.rbend_shift += b_test._x0_in + + l_test = xt.Line([b_test]) + l_test.append('end', xt.Marker()) + l_test.particle_ref = xt.Particles(p0c=10e9) + tw_test0 = l_test.twiss(betx=1, bety=1) + tw_test = l_test.twiss(betx=1, bety=1, x=2e-3, px=1e-3, y=2e-3, py=2e-3, delta=1e-3) + + xo.assert_allclose(tw_ref0.betx, tw_test0.betx, rtol=1e-9, atol=0.0) + xo.assert_allclose(tw_ref0.bety, tw_test0.bety, rtol=1e-9, atol=0.0) + xo.assert_allclose(tw_ref0.x, tw_test0.x, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref0.y, tw_test0.y, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref0.s, tw_test0.s, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref0.zeta, tw_test0.zeta, rtol=0, atol=1e-11) + xo.assert_allclose(tw_ref0.px, tw_test0.px, rtol=0, atol=1e-12) + xo.assert_allclose(tw_ref0.py, tw_test0.py, rtol=0, atol=1e-12) + + xo.assert_allclose(tw_ref.betx, tw_test.betx, rtol=5e-6, atol=0.0) + xo.assert_allclose(tw_ref.bety, tw_test.bety, rtol=5e-6, atol=0.0) + xo.assert_allclose(tw_ref.x, tw_test.x, rtol=0, atol=2e-8) + xo.assert_allclose(tw_ref.y, tw_test.y, rtol=0, atol=2e-8) + xo.assert_allclose(tw_ref.s, tw_test.s, rtol=0, atol=2e-12) + xo.assert_allclose(tw_ref.zeta, tw_test.zeta, rtol=0, atol=1e-1) + xo.assert_allclose(tw_ref.px, tw_test.px, rtol=0, atol=5e-9) + xo.assert_allclose(tw_ref.py, tw_test.py, rtol=0, atol=5e-9) + + tw_back = l_test.twiss(init=tw_test, init_at='end') + + assert tw_back.orientation == 'backward' + xo.assert_allclose(tw_back.betx, tw_test.betx, rtol=5e-6, atol=0.0) + xo.assert_allclose(tw_back.bety, tw_test.bety, rtol=5e-6, atol=0.0) + xo.assert_allclose(tw_back.x, tw_test.x, rtol=0, atol=1e-8) + xo.assert_allclose(tw_back.y, tw_test.y, rtol=0, atol=1e-8) + xo.assert_allclose(tw_back.s, tw_test.s, rtol=0, atol=1e-12) + xo.assert_allclose(tw_back.zeta, tw_test.zeta, rtol=0, atol=1e-1) + xo.assert_allclose(tw_back.px, tw_test.px, rtol=0, atol=1e-9) + xo.assert_allclose(tw_back.py, tw_test.py, rtol=0, atol=1e-9) + + l_sliced = l_test.copy(shallow=True) + l_sliced.cut_at_s(np.linspace(0, l_test.get_length(), 100)) + tw_test_sliced0 = l_sliced.twiss(betx=1, bety=1) + + xo.assert_allclose(tw_test_sliced0.betx[-1], tw_test0.betx[-1], rtol=1e-9, atol=0.0) + xo.assert_allclose(tw_test_sliced0.bety[-1], tw_test0.bety[-1], rtol=1e-9, atol=0.0) + xo.assert_allclose(tw_test_sliced0.x[-1], tw_test0.x[-1], rtol=0, atol=1e-12) + xo.assert_allclose(tw_test_sliced0.y[-1], tw_test0.y[-1], rtol=0, atol=1e-12) + xo.assert_allclose(tw_test_sliced0.s[-1], tw_test0.s[-1], rtol=0, atol=1e-12) + xo.assert_allclose(tw_test_sliced0.zeta[-1], tw_test0.zeta[-1], rtol=0, atol=1e-11) + xo.assert_allclose(tw_test_sliced0.px[-1], tw_test0.px[-1], rtol=0, atol=1e-12) + xo.assert_allclose(tw_test_sliced0.py[-1], tw_test0.py[-1], rtol=0, atol=1e-12) + + tw_test_sliced_back = l_sliced.twiss(init=tw_test_sliced0, init_at='end') + + assert tw_test_sliced_back.orientation == 'backward' + xo.assert_allclose(tw_test_sliced_back.betx, tw_test_sliced0.betx, rtol=5e-9, atol=0.0) + xo.assert_allclose(tw_test_sliced_back.bety, tw_test_sliced0.bety, rtol=5e-9, atol=0.0) + xo.assert_allclose(tw_test_sliced_back.x, tw_test_sliced0.x, rtol=0, atol=1e-12) + xo.assert_allclose(tw_test_sliced_back.y, tw_test_sliced0.y, rtol=0, atol=1e-12) + xo.assert_allclose(tw_test_sliced_back.s, tw_test_sliced0.s, rtol=0, atol=1e-12) + xo.assert_allclose(tw_test_sliced_back.zeta, tw_test_sliced0.zeta, rtol=0, atol=1e-11) + xo.assert_allclose(tw_test_sliced_back.px, tw_test_sliced0.px, rtol=0, atol=1e-12) + xo.assert_allclose(tw_test_sliced_back.py, tw_test_sliced0.py, rtol=0, atol=1e-12) + + def test_rbend_straight_sps(): env = xt.load([test_data_folder / 'sps_thick/sps.seq', @@ -703,6 +842,304 @@ def test_rbend_straight_body_survey_h(): xo.assert_allclose(sv_no_slice_straight_end['psi', nn], sv_straight['psi', nn], atol=5e-14) xo.assert_allclose(sv_no_slice_straight_end['s', nn], sv_straight['s', nn], atol=5e-14) +def test_rbend_straight_survey_h_angle_diff(): + + env = xt.Environment(particle_ref=xt.Particles(p0c=10e9)) + + edge_model = 'full' + + line = env.new_line(length=5, components=[ + env.new('mb', 'RBend', angle=0.3, k0_from_h=True, length_straight=3, + rbend_angle_diff=0.3, + model='bend-kick-bend', + rbend_model='straight-body', edge_entry_model=edge_model, edge_exit_model=edge_model, + anchor='start', at=1.)]) + line['mb'].rbend_compensate_sagitta = False + line['mb'].rbend_shift = line['mb']._x0_in + line.insert('start', xt.Marker(), at=0) + line.append('end', xt.Marker()) + + line_no_slice = line.copy(shallow=True) + + line.slice_thick_elements( + slicing_strategies=[ + xt.Strategy(slicing=xt.Uniform(6, mode='thick'))]) + + line['mb'].rbend_model = 'straight-body' + sv_straight = line.survey() + tt_straight = line.get_table(attr=True) + tw_straight = line.twiss(betx=1, bety=1) + p_straight = (sv_straight.p0 + tw_straight.x[:, None] * sv_straight['ex'] + + tw_straight.y[:, None] * sv_straight['ey']) + tw_straight['X'] = p_straight[:, 0] + tw_straight['Y'] = p_straight[:, 1] + tw_straight['Z'] = p_straight[:, 2] + + sv_straight_start = line.survey(element0='start', + X0=sv_straight['X', 'start'], + Y0=sv_straight['Y', 'start'], + Z0=sv_straight['Z', 'start'], + theta0=sv_straight['theta', 'start'], + phi0=sv_straight['phi', 'start'], + psi0=sv_straight['psi', 'start']) + sv_straight_end = line.survey(element0='end', + X0=sv_straight['X', 'end'], + Y0=sv_straight['Y', 'end'], + Z0=sv_straight['Z', 'end'], + theta0=sv_straight['theta', 'end'], + phi0=sv_straight['phi', 'end'], + psi0=sv_straight['psi', 'end']) + + sv_no_slice_start = line_no_slice.survey(element0='start', + X0=sv_straight['X', 'start'], + Y0=sv_straight['Y', 'start'], + Z0=sv_straight['Z', 'start'], + theta0=sv_straight['theta', 'start'], + phi0=sv_straight['phi', 'start'], + psi0=sv_straight['psi', 'start']) + sv_no_slice_end = line_no_slice.survey(element0='end', + X0=sv_straight['X', 'end'], + Y0=sv_straight['Y', 'end'], + Z0=sv_straight['Z', 'end'], + theta0=sv_straight['theta', 'end'], + phi0=sv_straight['phi', 'end'], + psi0=sv_straight['psi', 'end']) + tw_no_slice_straight = line_no_slice.twiss(betx=1, bety=1) + + line['mb'].rbend_model = 'curved-body' + sv_curved = line.survey() + tt_curved = line.get_table(attr=True) + tw_curved = line.twiss(betx=1, bety=1) + p_curved = (sv_curved.p0 + tw_curved.x[:, None] * sv_curved['ex'] + + tw_curved.y[:, None] * sv_curved['ey']) + tw_curved['X'] = p_curved[:, 0] + tw_curved['Y'] = p_curved[:, 1] + tw_curved['Z'] = p_curved[:, 2] + + sv_curved_start = line.survey(element0='start', + X0=sv_curved['X', 'start'], + Y0=sv_curved['Y', 'start'], + Z0=sv_curved['Z', 'start'], + theta0=sv_curved['theta', 'start'], + phi0=sv_curved['phi', 'start'], + psi0=sv_curved['psi', 'start']) + sv_curved_end = line.survey(element0='end', + X0=sv_curved['X', 'end'], + Y0=sv_curved['Y', 'end'], + Z0=sv_curved['Z', 'end'], + theta0=sv_curved['theta', 'end'], + phi0=sv_curved['phi', 'end'], + psi0=sv_curved['psi', 'end']) + sv_no_slice_curved_start = line_no_slice.survey(element0='start', + X0=sv_curved['X', 'start'], + Y0=sv_curved['Y', 'start'], + Z0=sv_curved['Z', 'start'], + theta0=sv_curved['theta', 'start'], + phi0=sv_curved['phi', 'start'], + psi0=sv_curved['psi', 'start']) + sv_no_slice_curved_end = line_no_slice.survey(element0='end', + X0=sv_curved['X', 'end'], + Y0=sv_curved['Y', 'end'], + Z0=sv_curved['Z', 'end'], + theta0=sv_curved['theta', 'end'], + phi0=sv_curved['phi', 'end'], + psi0=sv_curved['psi', 'end']) + tw_no_slice_curved = line_no_slice.twiss(betx=1, bety=1) + + + assert np.all(sv_straight['name'] == [ + 'start', '||drift_1', 'mb_entry', 'mb..entry_map', 'mb..0', + 'mb..1', 'mb..2', 'mb..3', 'mb..4', 'mb..5', 'mb..exit_map', + 'mb_exit', '||drift_2', 'end', '_end_point']) + + # Assert entire columns using np.all + assert np.all(sv_straight['element_type'] == [ + 'Marker', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThinSliceRBendExit', 'Marker', 'Drift', 'Marker', '']) + + xo.assert_allclose( + sv_straight['angle'], + np.array([ + 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0.3, 0. , 0. , + 0. , 0. ]), + atol=1e-12 + ) + + xo.assert_allclose(sv_straight['s'], np.array([ + 0. , 0. , 1. , 1. , 1. , + 1.5075795 , 2.01515901, 2.52273851, 3.03031802, 3.53789752, + 4.04547703, 4.04547703, 4.04547703, 5. , 5. ] + ), atol=1e-5) + + xo.assert_allclose( + sv_straight['rot_s_rad'], 0, atol=1e-14) + + sv_straight.cols['X Y Z'] + xo.assert_allclose(sv_straight['Y'], 0, atol=1e-14) + xo.assert_allclose(sv_straight['Z'], np.array( + [0. , 0. , 1. , 1. , 1. , + 1.5 , 2. , 2.5 , 3. , 3.5 , + 4. , 4. , 4. , 4.91189063, 4.91189063]), + atol=1e-8) + xo.assert_allclose(sv_straight['X'], np.array([ + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + -1.38777878e-17, -1.38777878e-17, -1.38777878e-17, -1.38777878e-17, + -1.38777878e-17, -1.38777878e-17, -1.38777878e-17, -4.53405654e-01, + -4.53405654e-01, -7.35486481e-01, -7.35486481e-01]), + atol=1e-8) + + + sv_straight.cols['theta phi psi'] + xo.assert_allclose(sv_straight['phi'], 0, atol=1e-14) + xo.assert_allclose(sv_straight['psi'], 0, atol=1e-14) + xo.assert_allclose(sv_straight['theta'], np.array([ + 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , + -0.3, -0.3, -0.3, -0.3])) + + + sv_curved.cols['s element_type angle'] + assert np.all(sv_curved['name'] == [ + 'start', '||drift_1', 'mb_entry', 'mb..entry_map', 'mb..0', + 'mb..1', 'mb..2', 'mb..3', 'mb..4', 'mb..5', 'mb..exit_map', + 'mb_exit', '||drift_2', 'end', '_end_point']) + + assert np.all(sv_curved['element_type'] == [ + 'Marker', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThinSliceRBendExit', 'Marker', 'Drift', 'Marker', '']) + + xo.assert_allclose( + sv_curved['angle'], + np.array([ + 0. , 0. , 0. , 0. , 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0. , + 0. , 0. , 0. , 0. ]), + atol=1e-8 + ) + + xo.assert_allclose(sv_curved['s'], np.array([ + 0. , 0. , 1. , 1. , 1. , + 1.5075795 , 2.01515901, 2.52273851, 3.03031802, 3.53789752, + 4.04547703, 4.04547703, 4.04547703, 5. , 5. ] + ), atol=1e-5) + + xo.assert_allclose( + sv_curved['rot_s_rad'], 0, atol=1e-14) + + sv_curved.cols['X Y Z'] + xo.assert_allclose(sv_curved['Y'], 0, atol=1e-14) + xo.assert_allclose(sv_curved['Z'], np.array([ + 0. , 0. , 1. , 1. , 1. , 1.507368, + 2.013468, 2.517035, 3.01681 , 3.511544, 4. , 4. , + 4. , 4.911891, 4.911891]), atol=1e-5) + xo.assert_allclose(sv_curved['X'], np.array( + [ 0. , 0. , 0. , 0. , 0. , + -0.01268684, -0.05071567, -0.11399141, -0.20235593, -0.31558835, + -0.45340565, -0.45340565, -0.45340565, -0.73548648, -0.73548648]), + atol=1e-8) + + sv_curved.cols['theta phi psi'] + xo.assert_allclose(sv_curved['phi'], 0, atol=1e-14) + xo.assert_allclose(sv_curved['psi'], 0, atol=1e-14) + xo.assert_allclose(sv_curved['theta'], np.array([ + 0. , 0. , 0. , 0. , 0. , -0.05, -0.1 , -0.15, -0.2 , + -0.25, -0.3 , -0.3 , -0.3 , -0.3 , -0.3 ], + ), atol=1e-8) + + for nn in ['start', 'end']: + xo.assert_allclose(sv_straight['X', nn], sv_curved['X', nn], atol=1e-14) + xo.assert_allclose(sv_straight['Y', nn], sv_curved['Y', nn], atol=1e-14) + xo.assert_allclose(sv_straight['Z', nn], sv_curved['Z', nn], atol=1e-14) + xo.assert_allclose(sv_straight['theta', nn], sv_curved['theta', nn], atol=1e-14) + xo.assert_allclose(sv_straight['phi', nn], sv_curved['phi', nn], atol=1e-14) + xo.assert_allclose(sv_straight['psi', nn], sv_curved['psi', nn], atol=1e-14) + + xo.assert_allclose(sv_straight_start['X'], sv_straight['X'], atol=1e-14) + xo.assert_allclose(sv_straight_start['Y'], sv_straight['Y'], atol=1e-14) + xo.assert_allclose(sv_straight_start['Z'], sv_straight['Z'], atol=1e-14) + xo.assert_allclose(sv_straight_start['theta'], sv_straight['theta'], atol=1e-14) + xo.assert_allclose(sv_straight_start['phi'], sv_straight['phi'], atol=1e-14) + xo.assert_allclose(sv_straight_start['psi'], sv_straight['psi'], atol=1e-14) + + xo.assert_allclose(sv_straight_end['X'], sv_straight['X'], atol=1e-14) + xo.assert_allclose(sv_straight_end['Y'], sv_straight['Y'], atol=1e-14) + xo.assert_allclose(sv_straight_end['Z'], sv_straight['Z'], atol=1e-14) + xo.assert_allclose(sv_straight_end['theta'], sv_straight['theta'], atol=1e-14) + xo.assert_allclose(sv_straight_end['phi'], sv_straight['phi'], atol=1e-14) + xo.assert_allclose(sv_straight_end['psi'], sv_straight['psi'], atol=1e-14) + + xo.assert_allclose(sv_curved_start['X'], sv_curved['X'], atol=1e-14) + xo.assert_allclose(sv_curved_start['Y'], sv_curved['Y'], atol=1e-14) + xo.assert_allclose(sv_curved_start['Z'], sv_curved['Z'], atol=1e-14) + xo.assert_allclose(sv_curved_start['theta'], sv_curved['theta'], atol=1e-14) + xo.assert_allclose(sv_curved_start['phi'], sv_curved['phi'], atol=1e-14) + xo.assert_allclose(sv_curved_start['psi'], sv_curved['psi'], atol=1e-14) + + xo.assert_allclose(sv_curved_end['X'], sv_curved['X'], atol=1e-14) + xo.assert_allclose(sv_curved_end['Y'], sv_curved['Y'], atol=1e-14) + xo.assert_allclose(sv_curved_end['Z'], sv_curved['Z'], atol=1e-14) + xo.assert_allclose(sv_curved_end['theta'], sv_curved['theta'], atol=1e-14) + xo.assert_allclose(sv_curved_end['phi'], sv_curved['phi'], atol=1e-14) + xo.assert_allclose(sv_curved_end['psi'], sv_curved['psi'], atol=1e-14) + xo.assert_allclose(tw_straight['X', 'start'], 0, atol=1e-14) + xo.assert_allclose(tw_straight['Y', 'start'], 0, atol=1e-14) + xo.assert_allclose(tw_straight['Z', 'start'], 0, atol=1e-14) + xo.assert_allclose(tw_curved['X', 'start'], 0, atol=1e-14) + xo.assert_allclose(tw_curved['Y', 'start'], 0, atol=1e-14) + xo.assert_allclose(tw_curved['Z', 'start'], 0, atol=1e-14) + xo.assert_allclose(tw_straight['X', 'mb_entry'], tw_curved['X', 'mb_entry'], atol=1e-14) + xo.assert_allclose(tw_straight['Y', 'mb_entry'], tw_curved['Y', 'mb_entry'], atol=1e-14) + xo.assert_allclose(tw_straight['Z', 'mb_entry'], tw_curved['Z', 'mb_entry'], atol=1e-14) + xo.assert_allclose(tw_straight['X', 'mb_exit'], tw_curved['X', 'mb_exit'], atol=1e-14) + xo.assert_allclose(tw_straight['Y', 'mb_exit'], tw_curved['Y', 'mb_exit'], atol=1e-14) + xo.assert_allclose(tw_straight['Z', 'mb_exit'], tw_curved['Z', 'mb_exit'], atol=1e-14) + + xo.assert_allclose(tw_straight['x', 'mb_entry'], 0, atol=1e-14) + xo.assert_allclose(tw_straight['y', 'mb_entry'], 0, atol=1e-14) + xo.assert_allclose(tw_straight['x', 'mb_exit'], 0, atol=1e-14) + xo.assert_allclose(tw_straight['y', 'mb_exit'], 0, atol=1e-14) + xo.assert_allclose(tw_curved['x', 'mb_entry'], 0, atol=1e-14) + xo.assert_allclose(tw_curved['y', 'mb_entry'], 0, atol=1e-14) + xo.assert_allclose(tw_curved['x', 'mb_exit'], 0, atol=1e-14) + xo.assert_allclose(tw_curved['y', 'mb_exit'], 0, atol=1e-14) + + xo.assert_allclose(tw_no_slice_curved['x', 'start'], tw_curved['x', 'start'], + atol=1e-14) + xo.assert_allclose(tw_no_slice_curved['y', 'start'], tw_curved['y', 'start'], + atol=1e-14) + xo.assert_allclose(tw_no_slice_curved['x', 'end'], tw_curved['x', 'end'], + atol=1e-14) + xo.assert_allclose(tw_no_slice_curved['y', 'end'], tw_curved['y', 'end'], + atol=1e-14) + xo.assert_allclose(tw_no_slice_straight['x', 'start'], tw_curved['x', 'start'], + atol=1e-14) + xo.assert_allclose(tw_no_slice_straight['y', 'start'], tw_curved['y', 'start'], + atol=1e-14) + xo.assert_allclose(tw_no_slice_straight['x', 'end'], tw_curved['x', 'end'], + atol=1e-14) + xo.assert_allclose(tw_no_slice_straight['y', 'end'], tw_curved['y', 'end'], + atol=1e-14) + + for nn in ['start', 'end']: + # Compare no_slice survey vs curved survey + xo.assert_allclose(sv_no_slice_curved_start['X', nn], sv_curved['X', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['Y', nn], sv_curved['Y', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['Z', nn], sv_curved['Z', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['theta', nn], sv_curved['theta', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['phi', nn], sv_curved['phi', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['psi', nn], sv_curved['psi', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_start['s', nn], sv_curved['s', nn], atol=1e-14) + + xo.assert_allclose(sv_no_slice_curved_end['X', nn], sv_curved['X', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['Y', nn], sv_curved['Y', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['Z', nn], sv_curved['Z', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['theta', nn], sv_curved['theta', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['phi', nn], sv_curved['phi', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['psi', nn], sv_curved['psi', nn], atol=1e-14) + xo.assert_allclose(sv_no_slice_curved_end['s', nn], sv_curved['s', nn], atol=1e-14) + def test_rbend_straight_body_survey_v(): env = xt.Environment(particle_ref=xt.Particles(p0c=10e9)) @@ -1189,7 +1626,6 @@ def test_rbend_straight_body_survey_v(): xo.assert_allclose(sv_no_slice_straight_end['psi', nn], sv_straight['psi', nn], atol=5e-14) xo.assert_allclose(sv_no_slice_straight_end['s', nn], sv_straight['s', nn], atol=5e-14) - def test_rbend_straight_body_thin_slices_coarse(): env = xt.Environment(particle_ref=xt.Particles(p0c=10e9)) @@ -1843,3 +2279,648 @@ def test_rbend_loading_writing_in_interfaces(): # Check madng sv_ng = line.madng_survey() xo.assert_allclose(sv_ng['s', 'b'], 2.5 - l_arc / 2, rtol=1e-12) # b si at entrance + +@pytest.mark.parametrize('edge_model', ['linear', 'full']) +def test_rbend_straight_body_chicane_h(edge_model): + + env = xt.Environment() + env.vars.default_to_zero = True + line = env.new_line(compose=True) + line.new('start', 'Marker', at=0.) + line.new('d1a', 'RBend', length_straight=1.0, k0='k0d1a', anchor='start', at='dz_d1a') + line.new('d1b', 'RBend', length_straight=1.0, k0='k0d1b', anchor='start', at='dz_d1b') + line.new('d2', 'RBend', length_straight=1.0, k0='k0d2', anchor='start', at='dz_d2') + line.new('end', 'Marker', at='dz_end') + + # ------ measure geometry in the straight reference frame ------ + + # Positions in the straight reference frame + env['dz_d1a'] = 1. + env['dz_d1b'] = 3. + env['dz_d2'] = 8. + env['dz_end'] = 10. + + line.end_compose() + line.set_particle_ref('proton', p0c=1e9) + line.configure_drift_model('exact') + line.set(env.elements.get_table().rows.match(element_type='RBend'), + model='bend-kick-bend', edge_entry_model=edge_model, + edge_exit_model=edge_model) + + env['k0d1a'] = 'k0d1' + env['k0d1b'] = 'k0d1' + + opt = line.match( + solve=False, + betx=1, bety=1, + vary=[xt.VaryList(['k0d1', 'k0d2'], step=1e-5)], + targets=xt.TargetSet(x=1., px=0.0, at='end'), + ) + opt.solve() + + # ---- build geometry with curved reference frame ---- + + # Twiss in the straight reference system + tw0 = line.twiss(betx=1, bety=1, strengths=True) + + if edge_model == 'linear': + # Set fdown angles to match the trajectory (used only for linear edges) + for nn in ['d1a', 'd1b', 'd2']: + line[nn].edge_entry_angle_fdown = np.arcsin(tw0['px', nn]) + line[nn].edge_exit_angle_fdown = -np.arcsin(tw0['px', nn + '>>1']) + tw0 = line.twiss(betx=1, bety=1, strengths=True) + for nn in ['d1a', 'd1b', 'd2']: + line[nn].edge_entry_angle_fdown = 0 + line[nn].edge_exit_angle_fdown = 0 + + line.regenerate_from_composer() + + # Update positions according to path length + env['dz_d1a'] = tw0['s', 'd1a'] - tw0['zeta', 'd1a'] + env['dz_d1b'] = tw0['s', 'd1b'] - tw0['zeta', 'd1b'] + env['dz_d2'] = tw0['s', 'd2'] - tw0['zeta', 'd2'] + env['dz_end'] = tw0['s', 'end'] - tw0['zeta', 'end'] + + # Introduce magnet curvatures + for nn in ['d1a', 'd1b', 'd2']: + line[nn].k0 = 0 + line[nn].k0_from_h = True + line[nn].rbend_compensate_sagitta = False + line[nn].rbend_model = 'straight-body' + + d1a_angle_in = np.arcsin(tw0['px', 'd1a']) + d1b_angle_in = np.arcsin(tw0['px', 'd1b']) + d2_angle_in = np.arcsin(tw0['px', 'd2']) + d1a_angle_out = -d1b_angle_in + d1b_angle_out = -d2_angle_in + d2_angle_out = -np.arcsin(tw0['px', 'end']) + + line['d1a'].angle = d1a_angle_in + d1a_angle_out + line['d1b'].angle = d1b_angle_in + d1b_angle_out + line['d2'].angle = d2_angle_in + d2_angle_out + + line['d1a'].rbend_angle_diff = d1a_angle_out - d1a_angle_in + line['d1b'].rbend_angle_diff = d1b_angle_out - d1b_angle_in + line['d2'].rbend_angle_diff = d2_angle_out - d2_angle_in + + # Set rbend shifts + line['d1a'].rbend_shift += line['d1a']._x0_in - tw0['x', 'd1a'] + line['d1b'].rbend_shift += line['d1b']._x0_in - tw0['x', 'd1b'] + line['d2'].rbend_shift += line['d2']._x0_out - tw0['x', 'end'] # to illustrate that out can be set as well + + line.end_compose() + + sv = line.survey() + tw = line.twiss(betx=1, bety=1) + sv_back = line.survey(element0='end', X0=sv.X[-1], Y0=sv.Y[-1], Z0=sv.Z[-1], + theta0=sv.theta[-1], phi0=sv.phi[-1], psi0=sv.psi[-1]) + if edge_model == 'linear': + tw_back = line.twiss(init_at='end', init=tw.get_twiss_init('end')) + + # slice for plot + l_sliced =line.copy(shallow=True) + l_sliced.slice_thick_elements( + slicing_strategies=[ + xt.Strategy(slicing=xt.Uniform(3, mode='thick')) + ]) + + sv_sliced = l_sliced.survey() + tw_sliced = l_sliced.twiss(betx=1, bety=1) + sv_sliced_back = l_sliced.survey(element0='end', + X0=sv_sliced.X[-1], Y0=sv_sliced.Y[-1], Z0=sv_sliced.Z[-1], + theta0=sv_sliced.theta[-1], phi0=sv_sliced.phi[-1], + psi0=sv_sliced.psi[-1]) + if edge_model == 'linear': + tw_sliced_back = l_sliced.twiss(init_at='end', + init=tw_sliced.get_twiss_init('end')) + + # Combine twiss and survey to get actual trajectory + trajectory = sv_sliced.p0 + tw_sliced.x[:, None] * sv_sliced.ex + tw_sliced.y[:, None] * sv_sliced.ey + + tw0['path_length'] = tw0.s - tw0.zeta + tw0['diff_path_length'] = np.diff(tw0.path_length, append=tw0.path_length[-1]) + + xo.assert_allclose(tw0.path_length, tw.s, atol=1e-14) + + xo.assert_allclose(tw0['diff_path_length', 'd1a'], line['d1a'].length, atol=1e-14) + xo.assert_allclose(tw0['diff_path_length', 'd1b'], line['d1b'].length, atol=1e-14) + xo.assert_allclose(tw0['diff_path_length', 'd2'], line['d2'].length, atol=1e-14) + + xo.assert_allclose(tw0['px', 'd1a'], 0, atol=1e-14) + xo.assert_allclose(tw0['px', 'd1b'], np.sin(line['d1b']._angle_in), atol=1e-14) + xo.assert_allclose(tw0['px', 'd2'], np.sin(line['d2']._angle_in), atol=1e-14) + + xo.assert_allclose(tw0['px', 'd1b'], -np.sin(line['d1a']._angle_out), atol=1e-14) + xo.assert_allclose(tw0['px', 'd2'], -np.sin(line['d1b']._angle_out), atol=1e-14) + xo.assert_allclose(tw0['px', 'end'], -np.sin(line['d2']._angle_out), atol=1e-14) + + xo.assert_allclose(tw0['x', 'd1a'], line['d1a']._x0_in, atol=1e-14) + xo.assert_allclose(tw0['x', 'd1b'], line['d1b']._x0_in, atol=1e-14) + xo.assert_allclose(tw0['x', 'd2'], line['d2']._x0_in, atol=1e-14) + + xo.assert_allclose(tw0['x', 'd1a>>1'], line['d1a']._x0_out, atol=1e-14) + xo.assert_allclose(tw0['x', 'd1b>>1'], line['d1b']._x0_out, atol=1e-14) + xo.assert_allclose(tw0['x', 'd2>>1'], line['d2']._x0_out, atol=1e-14) + + assert np.all(sv.element_type == + ['Marker', 'Drift', 'RBend', 'Drift', 'RBend', 'Drift', 'RBend', + 'Drift', 'Marker', '']) + xo.assert_allclose(sv.angle, np.array([ + 0. , 0. , -0.08249992, 0. , -0.08306823, + 0. , 0.16556815, 0. , 0. , 0. ]), + rtol=1e-7) + + xo.assert_allclose(sv.Z, tw0.s, atol=0, rtol=5e-9) + xo.assert_allclose(sv.X, tw0.x, atol=0, rtol=3e-8) + xo.assert_allclose(sv.Y, tw0.y, atol=1e-14) + xo.assert_allclose(sv.theta, np.arcsin(tw0.px), atol=1e-14) + xo.assert_allclose(sv.psi, 0., atol=1e-14) + xo.assert_allclose(sv.phi, 0., atol=1e-14) + + xo.assert_allclose(tw.x, 0, atol=1e-14) + xo.assert_allclose(tw.zeta, 0, atol=1e-14) + xo.assert_allclose(tw.y, 0, atol=1e-14) + + xo.assert_allclose(tw.betx[-1], tw0.betx[-1], rtol=1e-10) + xo.assert_allclose(tw.bety[-1], tw0.bety[-1], rtol=1e-10) + xo.assert_allclose(tw.alfx[-1], tw0.alfx[-1], rtol=1e-10) + xo.assert_allclose(tw.alfy[-1], tw0.alfy[-1], rtol=1e-10) + xo.assert_allclose(tw.dx[-1], tw0.dx[-1], rtol=1e-10) + xo.assert_allclose(tw.dpx[-1], tw0.dpx[-1], atol=1e-10) + + if edge_model == 'linear': + xo.assert_allclose(tw_back.s, tw.s, atol=1e-14) + xo.assert_allclose(tw_back.x, tw.x, atol=1e-14) + xo.assert_allclose(tw_back.y, tw.y, atol=1e-14) + xo.assert_allclose(tw_back.betx, tw.betx, rtol=1e-10) + xo.assert_allclose(tw_back.bety, tw.bety, rtol=1e-10) + xo.assert_allclose(tw_back.alfx, tw.alfx, atol=1e-8) + xo.assert_allclose(tw_back.alfy, tw.alfy, atol=1e-8) + xo.assert_allclose(tw_back.dx, tw.dx, atol=1e-9) + xo.assert_allclose(tw_back.dpx, tw.dpx, atol=1e-9) + + xo.assert_allclose(sv_back.s, sv.s, atol=1e-14) + xo.assert_allclose(sv_back.X, sv.X, atol=1e-14) + xo.assert_allclose(sv_back.Y, sv.Y, atol=1e-14) + xo.assert_allclose(sv_back.Z, sv.Z, atol=1e-14) + xo.assert_allclose(sv_back.theta, sv.theta, atol=1e-14) + xo.assert_allclose(sv_back.phi, sv.phi, atol=1e-14) + xo.assert_allclose(sv_back.psi, sv.psi, atol=1e-14) + xo.assert_allclose(sv.angle, sv.angle, atol=1e-14) + + sv_sliced.cols['s angle theta X'].show() + # name s angle theta X + # start 0 0 0 0 + # ||drift_1 0 0 0 0 + # d1a_entry 1 0 0 0 + # d1a..entry_map 1 0 0 0 + # d1a..0 1 0 0 0 + # d1a..1 1.33371 0 0 0 + # d1a..2 1.66742 0 0 0 + # d1a..exit_map 2.00114 -0.0824999 0 0 + # d1a_exit 2.00114 0 0.0824999 0.0412734 + # ||drift_3 2.00114 0 0.0824999 0.0412734 + # d1b_entry 3.00455 0 0.0824999 0.123961 + # d1b..entry_map 3.00455 0.0824999 0.0824999 0.123961 + # d1b..0 3.00455 0 -1.14732e-17 -1.38778e-17 + # d1b..1 3.34056 0 -1.14732e-17 -1.77022e-17 + # d1b..2 3.67657 0 -1.14732e-17 -2.15266e-17 + # d1b..exit_map 4.01258 -0.165568 -1.14732e-17 -2.5351e-17 + # d1b_exit 4.01258 0 0.165568 0.248635 + # ||drift_4 4.01258 0 0.165568 0.248635 + # d2_entry 8.06804 0 0.165568 0.917026 + # d2..entry_map 8.06804 0.165568 0.165568 0.917026 + # d2..0 8.06804 0 -9.64632e-18 1.11022e-16 + # d2..1 8.4029 0 -9.64632e-18 1.07807e-16 + # d2..2 8.73776 0 -9.64632e-18 1.04591e-16 + # d2..exit_map 9.07262 1.38778e-17 -9.64632e-18 1.01376e-16 + # d2_exit 9.07262 0 -9.64632e-18 1 + # ||drift_5 9.07262 0 -9.64632e-18 1 + # end 10.0726 0 -9.64632e-18 1 + # _end_point 10.0726 0 -9.64632e-18 1 + + xo.assert_allclose(sv_sliced.s[-1], tw0.path_length[-1], atol=0, rtol=1e-14) + xo.assert_allclose(sv_sliced.X[-1], tw0.x[-1], atol=0, rtol=1e-14) + xo.assert_allclose(sv_sliced.Y, 0, atol=1e-14) + + assert np.all(sv_sliced.element_type == + np.array(['Marker', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThinSliceRBendExit', 'Marker', 'Drift', 'Marker', + 'ThinSliceRBendEntry', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThinSliceRBendExit', 'Marker', 'Drift', + 'Marker', 'ThinSliceRBendEntry', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThinSliceRBendExit', + 'Marker', 'Drift', 'Marker', ''])) + + xo.assert_allclose(sv_sliced.angle, np.array([ + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, -8.24999219e-02, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 8.24999219e-02, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, -1.65568148e-01, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.65568148e-01, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]), + rtol=1e-8, atol=1e-14) + + xo.assert_allclose(sv_sliced['s', 'd1a..entry_map'], sv['s', 'd1a'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd1b..entry_map'], sv['s', 'd1b'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd2..entry_map'], sv['s', 'd2'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd1a..exit_map>>1'], sv['s', 'd1a>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd1b..exit_map>>1'], sv['s', 'd1b>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd2..exit_map>>1'], sv['s', 'd2>>1'], atol=1e-14) + + xo.assert_allclose(sv_sliced['theta', 'd1a..entry_map'], sv['theta', 'd1a'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1b..entry_map'], sv['theta', 'd1b'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd2..entry_map'], sv['theta', 'd2'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1a..exit_map>>1'], sv['theta', 'd1a>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1b..exit_map>>1'], sv['theta', 'd1b>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd2..exit_map>>1'], sv['theta', 'd2>>1'], atol=1e-14) + + xo.assert_allclose(sv_sliced['X', 'd1a..entry_map'], sv['X', 'd1a'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1b..entry_map'], sv['X', 'd1b'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd2..entry_map'], sv['X', 'd2'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1a..exit_map>>1'], sv['X', 'd1a>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1b..exit_map>>1'], sv['X', 'd1b>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd2..exit_map>>1'], sv['X', 'd2>>1'], atol=1e-14) + + xo.assert_allclose(sv_sliced['Z', 'd1a..entry_map'], sv['Z', 'd1a'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd1b..entry_map'], sv['Z', 'd1b'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd2..entry_map'], sv['Z', 'd2'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd1a..exit_map>>1'], sv['Z', 'd1a>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd1b..exit_map>>1'], sv['Z', 'd1b>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd2..exit_map>>1'], sv['Z', 'd2>>1'], atol=1e-14) + + xo.assert_allclose(sv_sliced['theta', 'd1a..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1b..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd2..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1a..exit_map'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1b..exit_map'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd2..exit_map'], 0., atol=1e-14) + + xo.assert_allclose(sv_sliced['X', 'd1a..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1b..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd2..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1a..exit_map'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1b..exit_map'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd2..exit_map'], 0., atol=1e-14) + + xo.assert_allclose(sv_sliced_back.s, sv_sliced.s, atol=1e-14) + xo.assert_allclose(sv_sliced_back.X, sv_sliced.X, atol=1e-14) + xo.assert_allclose(sv_sliced_back.Y, sv_sliced.Y, atol=1e-14) + xo.assert_allclose(sv_sliced_back.Z, sv_sliced.Z, atol=1e-14) + xo.assert_allclose(sv_sliced_back.theta, sv_sliced.theta, atol=1e-14) + xo.assert_allclose(sv_sliced_back.phi, sv_sliced.phi, atol=1e-14) + xo.assert_allclose(sv_sliced_back.psi, sv_sliced.psi, atol=1e-14) + xo.assert_allclose(sv_sliced_back.angle, sv_sliced.angle, atol=1e-14) + + # Twiss checks + + xo.assert_allclose(tw_sliced.s, sv_sliced.s, atol=0, rtol=1e-14) + + xo.assert_allclose(tw_sliced['x', 'd1a..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1b..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd2..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1a..exit_map>>1'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1b..exit_map>>1'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd2..exit_map>>1'], 0, atol=1e-14) + + xo.assert_allclose(tw_sliced['px', 'd1a..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd1b..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd2..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd1a..exit_map>>1'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd1b..exit_map>>1'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd2..exit_map>>1'], 0, atol=1e-14) + + xo.assert_allclose(tw_sliced['x', 'd1a..entry_map>>1'], tw0['x','d1a'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1b..entry_map>>1'], tw0['x','d1b'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd2..entry_map>>1'], tw0['x','d2'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1a..exit_map'], tw0['x','d1a>>1'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1b..exit_map'], tw0['x','d1b>>1'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd2..exit_map'], tw0['x','d2>>1'], atol=1e-14) + + xo.assert_allclose(tw_sliced.betx[-1], tw0.betx[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.bety[-1], tw0.bety[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.alfx[-1], tw0.alfx[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.alfy[-1], tw0.alfy[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.dx[-1], tw0.dx[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.dpx[-1], tw0.dpx[-1], atol=1e-10) + + if edge_model == 'linear': + xo.assert_allclose(tw_sliced_back.s, tw_sliced.s, atol=1e-14) + xo.assert_allclose(tw_sliced_back.x, tw_sliced.x, atol=1e-14) + xo.assert_allclose(tw_sliced_back.y, tw_sliced.y, atol=1e-14) + xo.assert_allclose(tw_sliced_back.betx, tw_sliced.betx, rtol=1e-8) + xo.assert_allclose(tw_sliced_back.bety, tw_sliced.bety, rtol=1e-8) + xo.assert_allclose(tw_sliced_back.alfx, tw_sliced.alfx, atol=1e-8) + xo.assert_allclose(tw_sliced_back.alfy, tw_sliced.alfy, atol=1e-8) + xo.assert_allclose(tw_sliced_back.dx, tw_sliced.dx, atol=1e-9) + xo.assert_allclose(tw_sliced_back.dpx, tw_sliced.dpx, atol=1e-9) + +@pytest.mark.parametrize('edge_model', ['linear', 'full']) +def test_rbend_straight_body_chicane_v(edge_model): + + env = xt.Environment() + env.vars.default_to_zero = True + line = env.new_line(compose=True) + line.new('start', 'Marker', at=0.) + line.new('d1a', 'RBend', rot_s_rad=np.pi/2, length_straight=1.0, k0='k0d1a', anchor='start', at='dz_d1a') + line.new('d1b', 'RBend', rot_s_rad=np.pi/2, length_straight=1.0, k0='k0d1b', anchor='start', at='dz_d1b') + line.new('d2', 'RBend', rot_s_rad=np.pi/2, length_straight=1.0, k0='k0d2', anchor='start', at='dz_d2') + line.new('end', 'Marker', at='dz_end') + + # ------ measure geometry in the straight reference frame ------ + + # Positions in the straight reference frame + env['dz_d1a'] = 1. + env['dz_d1b'] = 3. + env['dz_d2'] = 8. + env['dz_end'] = 10. + + line.end_compose() + line.set_particle_ref('proton', p0c=1e9) + line.configure_drift_model('exact') + line.set(env.elements.get_table().rows.match(element_type='RBend'), + model='bend-kick-bend', edge_entry_model=edge_model, + edge_exit_model=edge_model) + + env['k0d1a'] = 'k0d1' + env['k0d1b'] = 'k0d1' + + opt = line.match( + solve=False, + betx=1, bety=1, + vary=[xt.VaryList(['k0d1', 'k0d2'], step=1e-5)], + targets=xt.TargetSet(y=1., py=0.0, at='end'), + ) + opt.solve() + + # ---- build geometry with curved reference frame ---- + + # Twiss in the straight reference system + tw0 = line.twiss(betx=1, bety=1, strengths=True) + + if edge_model == 'linear': + # Set fdown angles to match the trajectory (used only for linear edges) + for nn in ['d1a', 'd1b', 'd2']: + line[nn].edge_entry_angle_fdown = np.arcsin(tw0['py', nn]) + line[nn].edge_exit_angle_fdown = -np.arcsin(tw0['py', nn + '>>1']) + tw0 = line.twiss(betx=1, bety=1, strengths=True) + for nn in ['d1a', 'd1b', 'd2']: + line[nn].edge_entry_angle_fdown = 0 + line[nn].edge_exit_angle_fdown = 0 + + line.regenerate_from_composer() + + # Update positions according to path length + env['dz_d1a'] = tw0['s', 'd1a'] - tw0['zeta', 'd1a'] + env['dz_d1b'] = tw0['s', 'd1b'] - tw0['zeta', 'd1b'] + env['dz_d2'] = tw0['s', 'd2'] - tw0['zeta', 'd2'] + env['dz_end'] = tw0['s', 'end'] - tw0['zeta', 'end'] + + # Introduce magnet curvatures + for nn in ['d1a', 'd1b', 'd2']: + line[nn].k0 = 0 + line[nn].k0_from_h = True + line[nn].rbend_compensate_sagitta = False + line[nn].rbend_model = 'straight-body' + + d1a_angle_in = np.arcsin(tw0['py', 'd1a']) + d1b_angle_in = np.arcsin(tw0['py', 'd1b']) + d2_angle_in = np.arcsin(tw0['py', 'd2']) + d1a_angle_out = -d1b_angle_in + d1b_angle_out = -d2_angle_in + d2_angle_out = -np.arcsin(tw0['py', 'end']) + + line['d1a'].angle = d1a_angle_in + d1a_angle_out + line['d1b'].angle = d1b_angle_in + d1b_angle_out + line['d2'].angle = d2_angle_in + d2_angle_out + + line['d1a'].rbend_angle_diff = d1a_angle_out - d1a_angle_in + line['d1b'].rbend_angle_diff = d1b_angle_out - d1b_angle_in + line['d2'].rbend_angle_diff = d2_angle_out - d2_angle_in + + # Set rbend shifts + line['d1a'].rbend_shift += line['d1a']._x0_in - tw0['y', 'd1a'] + line['d1b'].rbend_shift += line['d1b']._x0_in - tw0['y', 'd1b'] + line['d2'].rbend_shift += line['d2']._x0_out - tw0['y', 'end'] # to illustrate that out can be set as well + + line.end_compose() + + sv = line.survey() + tw = line.twiss(betx=1, bety=1) + sv_back = line.survey(element0='end', X0=sv.X[-1], Y0=sv.Y[-1], Z0=sv.Z[-1], + theta0=sv.theta[-1], phi0=sv.phi[-1], psi0=sv.psi[-1]) + if edge_model == 'linear': + tw_back = line.twiss(init_at='end', init=tw.get_twiss_init('end')) + + # slice for plot + l_sliced =line.copy(shallow=True) + l_sliced.slice_thick_elements( + slicing_strategies=[ + xt.Strategy(slicing=xt.Uniform(3, mode='thick')) + ]) + + sv_sliced = l_sliced.survey() + tw_sliced = l_sliced.twiss(betx=1, bety=1) + sv_sliced_back = l_sliced.survey(element0='end', + X0=sv_sliced.X[-1], Y0=sv_sliced.Y[-1], Z0=sv_sliced.Z[-1], + theta0=sv_sliced.theta[-1], phi0=sv_sliced.phi[-1], + psi0=sv_sliced.psi[-1]) + if edge_model == 'linear': + tw_sliced_back = l_sliced.twiss(init_at='end', + init=tw_sliced.get_twiss_init('end')) + + # Combine twiss and survey to get actual trajectory + trajectory = sv_sliced.p0 + tw_sliced.x[:, None] * sv_sliced.ex + tw_sliced.y[:, None] * sv_sliced.ey + + tw0['path_length'] = tw0.s - tw0.zeta + tw0['diff_path_length'] = np.diff(tw0.path_length, append=tw0.path_length[-1]) + + xo.assert_allclose(tw0.path_length, tw.s, atol=1e-14) + + xo.assert_allclose(tw0['diff_path_length', 'd1a'], line['d1a'].length, atol=1e-14) + xo.assert_allclose(tw0['diff_path_length', 'd1b'], line['d1b'].length, atol=1e-14) + xo.assert_allclose(tw0['diff_path_length', 'd2'], line['d2'].length, atol=1e-14) + + xo.assert_allclose(tw0['py', 'd1a'], 0, atol=1e-14) + xo.assert_allclose(tw0['py', 'd1b'], np.sin(line['d1b']._angle_in), atol=1e-14) + xo.assert_allclose(tw0['py', 'd2'], np.sin(line['d2']._angle_in), atol=1e-14) + + xo.assert_allclose(tw0['py', 'd1b'], -np.sin(line['d1a']._angle_out), atol=1e-14) + xo.assert_allclose(tw0['py', 'd2'], -np.sin(line['d1b']._angle_out), atol=1e-14) + xo.assert_allclose(tw0['py', 'end'], -np.sin(line['d2']._angle_out), atol=1e-14) + + xo.assert_allclose(tw0['y', 'd1a'], line['d1a']._x0_in, atol=1e-14) + xo.assert_allclose(tw0['y', 'd1b'], line['d1b']._x0_in, atol=1e-14) + xo.assert_allclose(tw0['y', 'd2'], line['d2']._x0_in, atol=1e-14) + + xo.assert_allclose(tw0['y', 'd1a>>1'], line['d1a']._x0_out, atol=1e-14) + xo.assert_allclose(tw0['y', 'd1b>>1'], line['d1b']._x0_out, atol=1e-14) + xo.assert_allclose(tw0['y', 'd2>>1'], line['d2']._x0_out, atol=1e-14) + + assert np.all(sv.element_type == + ['Marker', 'Drift', 'RBend', 'Drift', 'RBend', 'Drift', 'RBend', + 'Drift', 'Marker', '']) + xo.assert_allclose(sv.angle, np.array([ + 0. , 0. , -0.08249992, 0. , -0.08306823, + 0. , 0.16556815, 0. , 0. , 0. ]), + rtol=1e-7) + + xo.assert_allclose(sv.Z, tw0.s, atol=0, rtol=5e-9) + xo.assert_allclose(sv.X, tw0.x, atol=0, rtol=3e-8) + xo.assert_allclose(sv.Y, tw0.y, atol=1e-14) + xo.assert_allclose(sv.theta, 0., atol=1e-14) + xo.assert_allclose(sv.psi, 0., atol=1e-14) + xo.assert_allclose(sv.phi, np.arcsin(tw0.py), atol=1e-14) + + xo.assert_allclose(tw.x, 0, atol=1e-14) + xo.assert_allclose(tw.zeta, 0, atol=1e-14) + xo.assert_allclose(tw.y, 0, atol=1e-14) + + xo.assert_allclose(tw.betx[-1], tw0.betx[-1], rtol=1e-10) + xo.assert_allclose(tw.bety[-1], tw0.bety[-1], rtol=1e-10) + xo.assert_allclose(tw.alfx[-1], tw0.alfx[-1], rtol=1e-10) + xo.assert_allclose(tw.alfy[-1], tw0.alfy[-1], rtol=1e-10) + xo.assert_allclose(tw.dx[-1], tw0.dx[-1], rtol=1e-10) + xo.assert_allclose(tw.dpx[-1], tw0.dpx[-1], atol=1e-10) + + if edge_model == 'linear': + xo.assert_allclose(tw_back.s, tw.s, atol=1e-14) + xo.assert_allclose(tw_back.x, tw.x, atol=1e-14) + xo.assert_allclose(tw_back.y, tw.y, atol=1e-14) + xo.assert_allclose(tw_back.betx, tw.betx, rtol=1e-10) + xo.assert_allclose(tw_back.bety, tw.bety, rtol=1e-10) + xo.assert_allclose(tw_back.alfx, tw.alfx, atol=1e-8) + xo.assert_allclose(tw_back.alfy, tw.alfy, atol=1e-8) + xo.assert_allclose(tw_back.dx, tw.dx, atol=1e-9) + xo.assert_allclose(tw_back.dpx, tw.dpx, atol=1e-9) + + xo.assert_allclose(sv_back.s, sv.s, atol=1e-14) + xo.assert_allclose(sv_back.X, sv.X, atol=1e-14) + xo.assert_allclose(sv_back.Y, sv.Y, atol=1e-14) + xo.assert_allclose(sv_back.Z, sv.Z, atol=1e-14) + xo.assert_allclose(sv_back.theta, sv.theta, atol=1e-14) + xo.assert_allclose(sv_back.phi, sv.phi, atol=1e-14) + xo.assert_allclose(sv_back.psi, sv.psi, atol=1e-14) + xo.assert_allclose(sv.angle, sv.angle, atol=1e-14) + + sv_sliced.cols['s angle theta X'].show() + + xo.assert_allclose(sv_sliced.s[-1], tw0.path_length[-1], atol=0, rtol=1e-14) + xo.assert_allclose(sv_sliced.Y[-1], tw0.y[-1], atol=0, rtol=1e-14) + xo.assert_allclose(sv_sliced.X, 0, atol=1e-14) + + assert np.all(sv_sliced.element_type == + np.array(['Marker', 'Drift', 'Marker', 'ThinSliceRBendEntry', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThinSliceRBendExit', 'Marker', 'Drift', 'Marker', + 'ThinSliceRBendEntry', 'ThickSliceRBend', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThinSliceRBendExit', 'Marker', 'Drift', + 'Marker', 'ThinSliceRBendEntry', 'ThickSliceRBend', + 'ThickSliceRBend', 'ThickSliceRBend', 'ThinSliceRBendExit', + 'Marker', 'Drift', 'Marker', ''])) + + xo.assert_allclose(sv_sliced.angle, np.array([ + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, -8.24999219e-02, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 8.24999219e-02, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, -1.65568148e-01, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.65568148e-01, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]), + rtol=1e-8, atol=1e-14) + + xo.assert_allclose(sv_sliced['s', 'd1a..entry_map'], sv['s', 'd1a'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd1b..entry_map'], sv['s', 'd1b'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd2..entry_map'], sv['s', 'd2'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd1a..exit_map>>1'], sv['s', 'd1a>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd1b..exit_map>>1'], sv['s', 'd1b>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['s', 'd2..exit_map>>1'], sv['s', 'd2>>1'], atol=1e-14) + + xo.assert_allclose(sv_sliced['theta', 'd1a..entry_map'], sv['theta', 'd1a'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1b..entry_map'], sv['theta', 'd1b'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd2..entry_map'], sv['theta', 'd2'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1a..exit_map>>1'], sv['theta', 'd1a>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1b..exit_map>>1'], sv['theta', 'd1b>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd2..exit_map>>1'], sv['theta', 'd2>>1'], atol=1e-14) + + xo.assert_allclose(sv_sliced['X', 'd1a..entry_map'], sv['X', 'd1a'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1b..entry_map'], sv['X', 'd1b'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd2..entry_map'], sv['X', 'd2'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1a..exit_map>>1'], sv['X', 'd1a>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1b..exit_map>>1'], sv['X', 'd1b>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd2..exit_map>>1'], sv['X', 'd2>>1'], atol=1e-14) + + xo.assert_allclose(sv_sliced['Z', 'd1a..entry_map'], sv['Z', 'd1a'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd1b..entry_map'], sv['Z', 'd1b'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd2..entry_map'], sv['Z', 'd2'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd1a..exit_map>>1'], sv['Z', 'd1a>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd1b..exit_map>>1'], sv['Z', 'd1b>>1'], atol=1e-14) + xo.assert_allclose(sv_sliced['Z', 'd2..exit_map>>1'], sv['Z', 'd2>>1'], atol=1e-14) + + xo.assert_allclose(sv_sliced['theta', 'd1a..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1b..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd2..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1a..exit_map'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd1b..exit_map'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['theta', 'd2..exit_map'], 0., atol=1e-14) + + xo.assert_allclose(sv_sliced['X', 'd1a..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1b..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd2..entry_map>>1'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1a..exit_map'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd1b..exit_map'], 0., atol=1e-14) + xo.assert_allclose(sv_sliced['X', 'd2..exit_map'], 0., atol=1e-14) + + xo.assert_allclose(sv_sliced_back.s, sv_sliced.s, atol=1e-14) + xo.assert_allclose(sv_sliced_back.X, sv_sliced.X, atol=1e-14) + xo.assert_allclose(sv_sliced_back.Y, sv_sliced.Y, atol=1e-14) + xo.assert_allclose(sv_sliced_back.Z, sv_sliced.Z, atol=1e-14) + xo.assert_allclose(sv_sliced_back.theta, sv_sliced.theta, atol=1e-14) + xo.assert_allclose(sv_sliced_back.phi, sv_sliced.phi, atol=1e-14) + xo.assert_allclose(sv_sliced_back.psi, sv_sliced.psi, atol=1e-14) + xo.assert_allclose(sv_sliced_back.angle, sv_sliced.angle, atol=1e-14) + + # Twiss checks + + xo.assert_allclose(tw_sliced.s, sv_sliced.s, atol=0, rtol=1e-14) + + xo.assert_allclose(tw_sliced['x', 'd1a..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1b..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd2..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1a..exit_map>>1'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1b..exit_map>>1'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd2..exit_map>>1'], 0, atol=1e-14) + + xo.assert_allclose(tw_sliced['px', 'd1a..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd1b..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd2..entry_map'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd1a..exit_map>>1'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd1b..exit_map>>1'], 0, atol=1e-14) + xo.assert_allclose(tw_sliced['px', 'd2..exit_map>>1'], 0, atol=1e-14) + + xo.assert_allclose(tw_sliced['x', 'd1a..entry_map>>1'], tw0['x','d1a'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1b..entry_map>>1'], tw0['x','d1b'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd2..entry_map>>1'], tw0['x','d2'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1a..exit_map'], tw0['x','d1a>>1'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd1b..exit_map'], tw0['x','d1b>>1'], atol=1e-14) + xo.assert_allclose(tw_sliced['x', 'd2..exit_map'], tw0['x','d2>>1'], atol=1e-14) + + xo.assert_allclose(tw_sliced.betx[-1], tw0.betx[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.bety[-1], tw0.bety[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.alfx[-1], tw0.alfx[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.alfy[-1], tw0.alfy[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.dx[-1], tw0.dx[-1], rtol=1e-10) + xo.assert_allclose(tw_sliced.dpx[-1], tw0.dpx[-1], atol=1e-10) + + if edge_model == 'linear': + xo.assert_allclose(tw_sliced_back.s, tw_sliced.s, atol=1e-14) + xo.assert_allclose(tw_sliced_back.x, tw_sliced.x, atol=1e-14) + xo.assert_allclose(tw_sliced_back.y, tw_sliced.y, atol=1e-14) + xo.assert_allclose(tw_sliced_back.betx, tw_sliced.betx, rtol=1e-8) + xo.assert_allclose(tw_sliced_back.bety, tw_sliced.bety, rtol=1e-8) + xo.assert_allclose(tw_sliced_back.alfx, tw_sliced.alfx, atol=1e-8) + xo.assert_allclose(tw_sliced_back.alfy, tw_sliced.alfy, atol=1e-8) + xo.assert_allclose(tw_sliced_back.dx, tw_sliced.dx, atol=1e-9) + xo.assert_allclose(tw_sliced_back.dpx, tw_sliced.dpx, atol=1e-9) diff --git a/tests/test_tracker.py b/tests/test_tracker.py index 4aff5d2c9..b016206b8 100644 --- a/tests/test_tracker.py +++ b/tests/test_tracker.py @@ -381,12 +381,14 @@ class TestElement(xt.BeamElement): 'dummy': xo.Float64, } _extra_c_sources = [""" - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle( TestElementData el, LocalParticle* part0) { - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); #if TEST_FLAG == 2 LocalParticle_set_x(part, 7); @@ -396,7 +398,7 @@ class TestElement(xt.BeamElement): LocalParticle_set_y(part, 42); #endif - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } """] @@ -728,7 +730,9 @@ class TestElement(xt.BeamElement): _extra_c_sources = [ r''' - /*gpufun*/ + #include "xtrack/headers/track.h" + + GPUFUN void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){ // Extract the record and record_index TestElementRecordData record = TestElementData_getp_internal_record(el, part0); @@ -739,7 +743,7 @@ class TestElement(xt.BeamElement): int64_t elem_field = TestElementData_get_element_field(el); - //start_per_particle_block (part0->part) + START_PER_PARTICLE_BLOCK(part0, part); if (record) { // Record exists // Get a slot in the record (this is thread safe) int64_t i_slot = RecordIndex_get_slot(record_index); @@ -757,7 +761,7 @@ class TestElement(xt.BeamElement): ); } } - //end_per_particle_block + END_PER_PARTICLE_BLOCK; } ''' ] diff --git a/xtrack/base_element.py b/xtrack/base_element.py index c9e6ce046..e392971ed 100644 --- a/xtrack/base_element.py +++ b/xtrack/base_element.py @@ -43,7 +43,7 @@ } """ -def _handle_per_particle_blocks(sources, local_particle_src): +def _handle_per_particle_blocks(sources): if isinstance(sources, str): sources = (sources, ) @@ -59,8 +59,6 @@ def _handle_per_particle_blocks(sources, local_particle_src): else: strss = ss - strss = strss.replace('/*placeholder_for_local_particle_src*/', - local_particle_src) if '//start_per_particle_block' in strss: lines = strss.splitlines() @@ -73,7 +71,7 @@ def _handle_per_particle_blocks(sources, local_particle_src): # TODO: this is very dirty, just for check!!!!! out.append('\n'.join(lines)) else: - out.append(ss) + out.append(strss) if wasstring: @@ -441,9 +439,7 @@ def init_pipeline(self, pipeline_manager, name, partners_names=[]): def compile_kernels(self, extra_classes=(), *args, **kwargs): if 'apply_to_source' not in kwargs.keys(): kwargs['apply_to_source'] = [] - kwargs['apply_to_source'].append( - partial(_handle_per_particle_blocks, - local_particle_src=Particles.gen_local_particle_api())) + kwargs['apply_to_source'].append(_handle_per_particle_blocks) context = self._context cls = self.__class__ diff --git a/xtrack/beam_elements/acdipole.py b/xtrack/beam_elements/acdipole.py index 7dd3a02eb..bf5322c1b 100644 --- a/xtrack/beam_elements/acdipole.py +++ b/xtrack/beam_elements/acdipole.py @@ -8,7 +8,8 @@ class ACDipole(xt.BeamElement): """ ACDipole is a thin element that applies an oscillating kick to the beam in the x or y direction. It is used for beam excitations in circular machines for optics measurements. The kick is - defined by the voltages, frequencies,1 and phase lags. + a sinusoidal function defined by a voltage amplitude (with ramped up and down), a fixed frequency + (small compared to the revolution frequency) and fixed phase lag. If the dipole is not in twiss mode (typically used for tracking simulations): The kick is applied as a function of the turn number, and it can be ramped up and down @@ -34,6 +35,7 @@ class ACDipole(xt.BeamElement): freq : float | None The driven frequency of the AC dipole, in units of 2π per turn. This is equivalent to the fractional driven tune. If `None`, freq is set to zero. + Note that freq must be small compared to the revolution frequency. _This is the only parameter that is used in _both_ tracking and twiss modes._ lag : float | None The phase lag of the AC dipole, in units of radians. This is only used in diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 5b4670c9b..958a34d2c 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -1747,7 +1747,9 @@ class RBend(_BendCommon, BeamElement): **_BendCommon._common_xofields, 'length_straight': xo.Float64, 'rbend_model': xo.Int64, + 'rbend_compensate_sagitta': xo.Field(xo.Int64, default=True), 'rbend_shift': xo.Float64, + 'rbend_angle_diff': xo.Float64, } allow_loss_refinement = True @@ -1755,7 +1757,9 @@ class RBend(_BendCommon, BeamElement): _rename = { **_BendCommon._common_rename, 'length_straight': '_length_straight', - 'rbend_model': '_rbend_model' + 'rbend_model': '_rbend_model', + 'rbend_angle_diff': '_rbend_angle_diff', + 'rbend_compensate_sagitta': '_rbend_compensate_sagitta', } _depends_on = [RandomUniformAccurate, RandomExponential] @@ -1786,6 +1790,7 @@ def __init__(self, **kwargs): kwargs.get('length_straight'), kwargs.get('h'), kwargs.get('angle'), + rbend_angle_diff=kwargs.get('rbend_angle_diff', 0.0), ) if self.k0_from_h: @@ -1801,13 +1806,24 @@ def __init__(self, **kwargs): if rbend_model is not None: self.rbend_model = rbend_model + @property + def rbend_angle_diff(self): + return self._rbend_angle_diff + + @rbend_angle_diff.setter + def rbend_angle_diff(self, value): + self._rbend_angle_diff = value + self.set_bend_params(length_straight=self._length_straight, angle=self._angle, + rbend_angle_diff=value) + @property def length(self): return self._length @length.setter def length(self, value): - self.set_bend_params(length=value, angle=self.angle) + self.set_bend_params(length=value, angle=self.angle, + rbend_angle_diff=self._rbend_angle_diff) @property def h(self): @@ -1815,7 +1831,8 @@ def h(self): @h.setter def h(self, value): - self.set_bend_params(h=value, length_straight=self.length_straight) + self.set_bend_params(h=value, length_straight=self._length_straight, + rbend_angle_diff=self._rbend_angle_diff) @property def angle(self): @@ -1823,7 +1840,8 @@ def angle(self): @angle.setter def angle(self, value): - self.set_bend_params(angle=value, length_straight=self.length_straight) + self.set_bend_params(angle=value, length_straight=self.length_straight, + rbend_angle_diff=self._rbend_angle_diff) @property def length_straight(self): @@ -1831,7 +1849,8 @@ def length_straight(self): @length_straight.setter def length_straight(self, value): - self.set_bend_params(length_straight=value, angle=self.angle) + self.set_bend_params(length_straight=value, angle=self.angle, + rbend_angle_diff=self._rbend_angle_diff) @property def rbend_model(self): @@ -1844,12 +1863,21 @@ def rbend_model(self, value): except KeyError: raise ValueError(f'Invalid rbend_model: {value}') - def set_bend_params(self, length=None, length_straight=None, h=None, angle=None): + @property + def rbend_compensate_sagitta(self): + return bool(self._rbend_compensate_sagitta) + + @rbend_compensate_sagitta.setter + def rbend_compensate_sagitta(self, value): + self._rbend_compensate_sagitta = int(bool(value)) + + def set_bend_params(self, length=None, length_straight=None, h=None, angle=None, + rbend_angle_diff=None): ( length, length_straight, h, angle - ) = self.compute_bend_params(length, length_straight, h, angle) + ) = self.compute_bend_params(length, length_straight, h, angle, + rbend_angle_diff=rbend_angle_diff) - # None becomes NaN in numpy buffers if length is not None: self._length = length if length_straight is not None: @@ -1858,12 +1886,12 @@ def set_bend_params(self, length=None, length_straight=None, h=None, angle=None) self._h = h if angle is not None: self._angle = angle - if self.k0_from_h: self._k0 = self.h @staticmethod - def compute_bend_params(length=None, length_straight=None, h=None, angle=None): + def compute_bend_params(length=None, length_straight=None, h=None, angle=None, + rbend_angle_diff=None): """Compute the bend parameters (length, h) from the given arguments. The arguments are checked for consistency and the missing ones are @@ -1873,6 +1901,7 @@ def compute_bend_params(length=None, length_straight=None, h=None, angle=None): `angle`. If both are missing for each pair, zero will be assumed for the length and for `h`. """ + assert rbend_angle_diff is not None kwargs = locals() given = { arg_name: arg_value @@ -1905,18 +1934,47 @@ def compute_bend_params(length=None, length_straight=None, h=None, angle=None): else: length = length or length_straight length_straight = length - elif 'length' in given and 'h' in given: + elif 'length' in given and 'h' in given: # case 3 angle = length * h - length_straight = length * np.sinc(0.5 * angle / np.pi) - elif 'length' in given and 'angle' in given: + theta_in = 0.5 * angle - rbend_angle_diff / 2 + theta_out = 0.5 * angle + rbend_angle_diff / 2 + if abs(angle) < 1e-10: + length_straight = length + else: + length_straight = (1/h) * (np.sin(theta_in) + np.sin(theta_out)) + elif 'length' in given and 'angle' in given: # case 2 h = angle / length - length_straight = length * np.sinc(0.5 * angle / np.pi) - elif 'length_straight' in given and 'h' in given: - angle = 2 * np.arcsin(length_straight * h / 2) - length = length_straight / np.sinc(0.5 * angle / np.pi) - elif 'length_straight' in given and 'angle' in given: - h = 2 * np.sin(angle / 2) / length_straight - length = length_straight / np.sinc(0.5 * angle / np.pi) + theta_in = 0.5 * angle - rbend_angle_diff / 2 + theta_out = 0.5 * angle + rbend_angle_diff / 2 + if abs(angle) < 1e-10: + length_straight = length + else: + length_straight = (1/h) * (np.sin(theta_in) + np.sin(theta_out)) + elif 'length_straight' in given and 'h' in given: # case 4 + angle = 2 * np.arcsin(length_straight * h / 2 / np.cos(rbend_angle_diff / 2)) + if abs(angle) < 1e-10: + length = length_straight + else: + length = angle / h + elif 'length_straight' in given and 'angle' in given: # case 1 + theta_in = 0.5 * angle - rbend_angle_diff / 2 + theta_out = 0.5 * angle + rbend_angle_diff / 2 + if abs(angle) < 1e-10: + length = length_straight + h = 0 + else: + h = (np.sin(theta_in) + np.sin(theta_out)) / length_straight + length = angle / h + elif 'h' in given and 'angle' in given: + theta_in = 0.5 * angle - rbend_angle_diff / 2 + theta_out = 0.5 * angle + rbend_angle_diff / 2 + if abs(angle) < 1e-10: + length = length_straight = 0 + else: + length = angle / h + length_straight = (1/h) * (np.sin(theta_in) + np.sin(theta_out)) + else: + raise RuntimeError("Internal error in RBend parameter computation.") # Verify consistency errors = [] @@ -1937,15 +1995,52 @@ def compute_bend_params(length=None, length_straight=None, h=None, angle=None): # By this point the given values are proved to be consistent, prefer # the values given by the user, not to accumulate numerical noise. return ( - length, + given.get('length', length), given.get('length_straight', length_straight), - h, + given.get('h', h), given.get('angle', angle), ) @property def hxl(self): return self.h * self.length + @property + def _angle_in(self): + return 0.5 * self.angle - self._rbend_angle_diff / 2 + + @property + def _angle_out(self): + return 0.5 * self.angle + self._rbend_angle_diff / 2 + + @property + def _x0_mid(self): + out = -self.rbend_shift + if abs(self.angle) > 1e-10 and self.rbend_compensate_sagitta: + out += 0.5 / self.h * (1 - np.cos(self.angle / 2)) + return out + + @property + def _x0_in(self): + out = self._x0_mid + if abs(self.angle) > 1e-10: + px0_in = np.sin(self._angle_in) + px0_mid = px0_in - self.h * self.length_straight / 2 + sqrt_mid = np.sqrt(1 - px0_mid * px0_mid) + cos_theta_in = np.cos(self._angle_in) + out -= 1 / self.h * (sqrt_mid - cos_theta_in) + return out + + @property + def _x0_out(self): + out = self._x0_mid + if abs(self.angle) > 1e-10: + px0_out = np.sin(self._angle_out) + px0_mid = px0_out - self.h * self.length_straight / 2 + sqrt_mid = np.sqrt(1 - px0_mid * px0_mid) + cos_theta_out = np.cos(self._angle_out) + out += 1 / self.h * (cos_theta_out - sqrt_mid) + return out + @property def radiation_flag(self): return 0.0 diff --git a/xtrack/beam_elements/elements_src/bend.h b/xtrack/beam_elements/elements_src/bend.h index 4d85d1e38..084a30d85 100644 --- a/xtrack/beam_elements/elements_src/bend.h +++ b/xtrack/beam_elements/elements_src/bend.h @@ -48,7 +48,10 @@ void Bend_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ BendData_get_edge_entry_active(el), /*edge_exit_active*/ BendData_get_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/magnet.h b/xtrack/beam_elements/elements_src/magnet.h index 40d581ebb..7c0f0bbf8 100644 --- a/xtrack/beam_elements/elements_src/magnet.h +++ b/xtrack/beam_elements/elements_src/magnet.h @@ -49,7 +49,10 @@ void Magnet_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ MagnetData_get_edge_entry_active(el), /*edge_exit_active*/ MagnetData_get_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/multipole.h b/xtrack/beam_elements/elements_src/multipole.h index 47699832d..0c33ca020 100644 --- a/xtrack/beam_elements/elements_src/multipole.h +++ b/xtrack/beam_elements/elements_src/multipole.h @@ -47,7 +47,10 @@ void Multipole_track_local_particle(MultipoleData el, LocalParticle* part0){ /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/octupole.h b/xtrack/beam_elements/elements_src/octupole.h index adcefb12a..ac0ca8402 100644 --- a/xtrack/beam_elements/elements_src/octupole.h +++ b/xtrack/beam_elements/elements_src/octupole.h @@ -48,7 +48,10 @@ void Octupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ OctupoleData_get_edge_entry_active(el), /*edge_exit_active*/ OctupoleData_get_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/quadrupole.h b/xtrack/beam_elements/elements_src/quadrupole.h index 63dfb38a4..59756ce6a 100644 --- a/xtrack/beam_elements/elements_src/quadrupole.h +++ b/xtrack/beam_elements/elements_src/quadrupole.h @@ -48,7 +48,10 @@ void Quadrupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ QuadrupoleData_get_edge_entry_active(el), /*edge_exit_active*/ QuadrupoleData_get_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/rbend.h b/xtrack/beam_elements/elements_src/rbend.h index 65e80ae57..030439603 100644 --- a/xtrack/beam_elements/elements_src/rbend.h +++ b/xtrack/beam_elements/elements_src/rbend.h @@ -48,7 +48,10 @@ void RBend_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ RBendData_get_rbend_model(el), + /*rbend_compensate_sagitta*/ RBendData_get_rbend_compensate_sagitta(el), /*rbend_shift*/ RBendData_get_rbend_shift(el), + /*rbend_angle_diff*/ RBendData_get_rbend_angle_diff(el), + /*length_straight*/ RBendData_get_length_straight(el), /*body_active*/ 1, /*edge_entry_active*/ RBendData_get_edge_entry_active(el), /*edge_exit_active*/ RBendData_get_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/sextupole.h b/xtrack/beam_elements/elements_src/sextupole.h index d5d4c2e60..e7fef72ae 100644 --- a/xtrack/beam_elements/elements_src/sextupole.h +++ b/xtrack/beam_elements/elements_src/sextupole.h @@ -48,7 +48,10 @@ void Sextupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ SextupoleData_get_edge_entry_active(el), /*edge_exit_active*/ SextupoleData_get_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/slnd.h b/xtrack/beam_elements/elements_src/slnd.h index ce868967c..41c852c8e 100644 --- a/xtrack/beam_elements/elements_src/slnd.h +++ b/xtrack/beam_elements/elements_src/slnd.h @@ -48,7 +48,10 @@ void UniformSolenoid_track_local_particle( /*x0_solenoid*/ UniformSolenoidData_get_x0(el), /*y0_solenoid*/ UniformSolenoidData_get_y0(el), /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ UniformSolenoidData_get_edge_entry_active(el), /*edge_exit_active*/ UniformSolenoidData_get_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/thick_slice_bend.h b/xtrack/beam_elements/elements_src/thick_slice_bend.h index 847d2e9ce..11bac1206 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_bend.h +++ b/xtrack/beam_elements/elements_src/thick_slice_bend.h @@ -51,7 +51,10 @@ void ThickSliceBend_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thick_slice_multipole.h b/xtrack/beam_elements/elements_src/thick_slice_multipole.h index 917249de9..954734296 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_multipole.h +++ b/xtrack/beam_elements/elements_src/thick_slice_multipole.h @@ -50,7 +50,10 @@ void ThickSliceMultipole_track_local_particle(ThickSliceMultipoleData el, LocalP /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thick_slice_octupole.h b/xtrack/beam_elements/elements_src/thick_slice_octupole.h index aaaecad09..2aba5ed95 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_octupole.h +++ b/xtrack/beam_elements/elements_src/thick_slice_octupole.h @@ -51,7 +51,10 @@ void ThickSliceOctupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thick_slice_quadrupole.h b/xtrack/beam_elements/elements_src/thick_slice_quadrupole.h index 6f44e74be..895853c52 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_quadrupole.h +++ b/xtrack/beam_elements/elements_src/thick_slice_quadrupole.h @@ -51,7 +51,10 @@ void ThickSliceQuadrupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thick_slice_rbend.h b/xtrack/beam_elements/elements_src/thick_slice_rbend.h index b4c29e059..330e50d28 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_rbend.h +++ b/xtrack/beam_elements/elements_src/thick_slice_rbend.h @@ -51,7 +51,10 @@ void ThickSliceRBend_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ ThickSliceRBendData_get__parent_rbend_model(el), + /*rbend_compensate_sagitta*/ ThickSliceRBendData_get__parent_rbend_compensate_sagitta(el), /*rbend_shift*/ ThickSliceRBendData_get__parent_rbend_shift(el), + /*rbend_angle_diff*/ ThickSliceRBendData_get__parent_rbend_angle_diff(el), + /*length_straight*/ ThickSliceRBendData_get__parent_length_straight(el), /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thick_slice_sextupole.h b/xtrack/beam_elements/elements_src/thick_slice_sextupole.h index 0ee1a9586..2cbc62d46 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_sextupole.h +++ b/xtrack/beam_elements/elements_src/thick_slice_sextupole.h @@ -51,7 +51,10 @@ void ThickSliceSextupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thick_slice_uniform_solenoid.h b/xtrack/beam_elements/elements_src/thick_slice_uniform_solenoid.h index 644f04fd2..af4688fb7 100644 --- a/xtrack/beam_elements/elements_src/thick_slice_uniform_solenoid.h +++ b/xtrack/beam_elements/elements_src/thick_slice_uniform_solenoid.h @@ -51,7 +51,10 @@ void ThickSliceUniformSolenoid_track_local_particle( /*x0_solenoid*/ ThickSliceUniformSolenoidData_get__parent_x0(el), /*y0_solenoid*/ ThickSliceUniformSolenoidData_get__parent_y0(el), /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_bend.h b/xtrack/beam_elements/elements_src/thin_slice_bend.h index d95a9a20f..dff2ad132 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_bend.h +++ b/xtrack/beam_elements/elements_src/thin_slice_bend.h @@ -51,7 +51,10 @@ void ThinSliceBend_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_bend_entry.h b/xtrack/beam_elements/elements_src/thin_slice_bend_entry.h index e615d2a39..c532f5d55 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_bend_entry.h +++ b/xtrack/beam_elements/elements_src/thin_slice_bend_entry.h @@ -51,7 +51,10 @@ void ThinSliceBendEntry_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ ThinSliceBendEntryData_get__parent_edge_entry_active(el), /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_bend_exit.h b/xtrack/beam_elements/elements_src/thin_slice_bend_exit.h index 10b9143e6..7e5126a42 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_bend_exit.h +++ b/xtrack/beam_elements/elements_src/thin_slice_bend_exit.h @@ -51,7 +51,10 @@ void ThinSliceBendExit_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ 0, /*edge_exit_active*/ ThinSliceBendExitData_get__parent_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/thin_slice_multipole.h b/xtrack/beam_elements/elements_src/thin_slice_multipole.h index 4664f9e2e..9c5ddab77 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_multipole.h +++ b/xtrack/beam_elements/elements_src/thin_slice_multipole.h @@ -50,7 +50,10 @@ void ThinSliceMultipole_track_local_particle(ThinSliceMultipoleData el, LocalPar /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_octupole.h b/xtrack/beam_elements/elements_src/thin_slice_octupole.h index 79223fcb6..f3595a648 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_octupole.h +++ b/xtrack/beam_elements/elements_src/thin_slice_octupole.h @@ -51,7 +51,10 @@ void ThinSliceOctupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_octupole_entry.h b/xtrack/beam_elements/elements_src/thin_slice_octupole_entry.h index 0db9bd677..32e874535 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_octupole_entry.h +++ b/xtrack/beam_elements/elements_src/thin_slice_octupole_entry.h @@ -51,7 +51,10 @@ void ThinSliceOctupoleEntry_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ ThinSliceOctupoleEntryData_get__parent_edge_entry_active(el), /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_octupole_exit.h b/xtrack/beam_elements/elements_src/thin_slice_octupole_exit.h index f937ba744..51ee618b7 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_octupole_exit.h +++ b/xtrack/beam_elements/elements_src/thin_slice_octupole_exit.h @@ -51,7 +51,10 @@ void ThinSliceOctupoleExit_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ 0, /*edge_exit_active*/ ThinSliceOctupoleExitData_get__parent_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/thin_slice_quadrupole.h b/xtrack/beam_elements/elements_src/thin_slice_quadrupole.h index 0a18140d3..b35dbb274 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_quadrupole.h +++ b/xtrack/beam_elements/elements_src/thin_slice_quadrupole.h @@ -51,7 +51,10 @@ void ThinSliceQuadrupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_quadrupole_entry.h b/xtrack/beam_elements/elements_src/thin_slice_quadrupole_entry.h index 834cefbcf..98087c1af 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_quadrupole_entry.h +++ b/xtrack/beam_elements/elements_src/thin_slice_quadrupole_entry.h @@ -51,7 +51,10 @@ void ThinSliceQuadrupoleEntry_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ ThinSliceQuadrupoleEntryData_get__parent_edge_entry_active(el), /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_quadrupole_exit.h b/xtrack/beam_elements/elements_src/thin_slice_quadrupole_exit.h index 67d5326ba..d0ace46c1 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_quadrupole_exit.h +++ b/xtrack/beam_elements/elements_src/thin_slice_quadrupole_exit.h @@ -51,7 +51,10 @@ void ThinSliceQuadrupoleExit_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ 0, /*edge_exit_active*/ ThinSliceQuadrupoleExitData_get__parent_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/thin_slice_rbend.h b/xtrack/beam_elements/elements_src/thin_slice_rbend.h index 4c34395a1..c86a9c790 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_rbend.h +++ b/xtrack/beam_elements/elements_src/thin_slice_rbend.h @@ -51,7 +51,10 @@ void ThinSliceRBend_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ ThinSliceRBendData_get__parent_rbend_model(el), + /*rbend_compensate_sagitta*/ ThinSliceRBendData_get__parent_rbend_compensate_sagitta(el), /*rbend_shift*/ ThinSliceRBendData_get__parent_rbend_shift(el), + /*rbend_angle_diff*/ ThinSliceRBendData_get__parent_rbend_angle_diff(el), + /*length_straight*/ ThinSliceRBendData_get__parent_length_straight(el), /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_rbend_entry.h b/xtrack/beam_elements/elements_src/thin_slice_rbend_entry.h index 7302c169d..21585f99e 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_rbend_entry.h +++ b/xtrack/beam_elements/elements_src/thin_slice_rbend_entry.h @@ -51,7 +51,10 @@ void ThinSliceRBendEntry_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ ThinSliceRBendEntryData_get__parent_rbend_model(el), + /*rbend_compensate_sagitta*/ ThinSliceRBendEntryData_get__parent_rbend_compensate_sagitta(el), /*rbend_shift*/ ThinSliceRBendEntryData_get__parent_rbend_shift(el), + /*rbend_angle_diff*/ ThinSliceRBendEntryData_get__parent_rbend_angle_diff(el), + /*length_straight*/ ThinSliceRBendEntryData_get__parent_length_straight(el), /*body_active*/ 0, // disabled /*edge_entry_active*/ ThinSliceRBendEntryData_get__parent_edge_entry_active(el), /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_rbend_exit.h b/xtrack/beam_elements/elements_src/thin_slice_rbend_exit.h index 539b7035f..e9fcc8846 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_rbend_exit.h +++ b/xtrack/beam_elements/elements_src/thin_slice_rbend_exit.h @@ -51,7 +51,10 @@ void ThinSliceRBendExit_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ ThinSliceRBendExitData_get__parent_rbend_model(el), + /*rbend_compensate_sagitta*/ ThinSliceRBendExitData_get__parent_rbend_compensate_sagitta(el), /*rbend_shift*/ ThinSliceRBendExitData_get__parent_rbend_shift(el), + /*rbend_angle_diff*/ ThinSliceRBendExitData_get__parent_rbend_angle_diff(el), + /*length_straight*/ ThinSliceRBendExitData_get__parent_length_straight(el), /*body_active*/ 0, // disabled /*edge_entry_active*/ 0, /*edge_exit_active*/ ThinSliceRBendExitData_get__parent_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/thin_slice_sextupole.h b/xtrack/beam_elements/elements_src/thin_slice_sextupole.h index 926fe4780..94d16f1bc 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_sextupole.h +++ b/xtrack/beam_elements/elements_src/thin_slice_sextupole.h @@ -51,7 +51,10 @@ void ThinSliceSextupole_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_sextupole_entry.h b/xtrack/beam_elements/elements_src/thin_slice_sextupole_entry.h index e558960d7..98c5c57ef 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_sextupole_entry.h +++ b/xtrack/beam_elements/elements_src/thin_slice_sextupole_entry.h @@ -51,7 +51,10 @@ void ThinSliceSextupoleEntry_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ ThinSliceSextupoleEntryData_get__parent_edge_entry_active(el), /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_sextupole_exit.h b/xtrack/beam_elements/elements_src/thin_slice_sextupole_exit.h index 0fae887e9..1ab80d415 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_sextupole_exit.h +++ b/xtrack/beam_elements/elements_src/thin_slice_sextupole_exit.h @@ -51,7 +51,10 @@ void ThinSliceSextupoleExit_track_local_particle( /*x0_solenoid*/ 0., /*y0_solenoid*/ 0., /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ 0, /*edge_exit_active*/ ThinSliceSextupoleExitData_get__parent_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/thin_slice_uniform_solenoid_entry.h b/xtrack/beam_elements/elements_src/thin_slice_uniform_solenoid_entry.h index 5a8b126fb..4e0b5c5d4 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_uniform_solenoid_entry.h +++ b/xtrack/beam_elements/elements_src/thin_slice_uniform_solenoid_entry.h @@ -51,7 +51,10 @@ void ThinSliceUniformSolenoidEntry_track_local_particle( /*x0_solenoid*/ ThinSliceUniformSolenoidEntryData_get__parent_x0(el), /*y0_solenoid*/ ThinSliceUniformSolenoidEntryData_get__parent_y0(el), /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ ThinSliceUniformSolenoidEntryData_get__parent_edge_entry_active(el), /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/elements_src/thin_slice_uniform_solenoid_exit.h b/xtrack/beam_elements/elements_src/thin_slice_uniform_solenoid_exit.h index b6da5814b..302f85760 100644 --- a/xtrack/beam_elements/elements_src/thin_slice_uniform_solenoid_exit.h +++ b/xtrack/beam_elements/elements_src/thin_slice_uniform_solenoid_exit.h @@ -51,7 +51,10 @@ void ThinSliceUniformSolenoidExit_track_local_particle( /*x0_solenoid*/ ThinSliceUniformSolenoidExitData_get__parent_x0(el), /*y0_solenoid*/ ThinSliceUniformSolenoidExitData_get__parent_y0(el), /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 0, // disabled /*edge_entry_active*/ 0, /*edge_exit_active*/ ThinSliceUniformSolenoidExitData_get__parent_edge_exit_active(el), diff --git a/xtrack/beam_elements/elements_src/track_magnet.h b/xtrack/beam_elements/elements_src/track_magnet.h index c71751ad8..299a4799a 100644 --- a/xtrack/beam_elements/elements_src/track_magnet.h +++ b/xtrack/beam_elements/elements_src/track_magnet.h @@ -173,6 +173,7 @@ void track_magnet_body_single_particle( } else{ + // START GENERATED INTEGRATION CODE if (integrator == 1){ // TEAPOT @@ -271,6 +272,7 @@ void track_magnet_body_single_particle( } + GPUFUN void track_magnet_particles( double const weight, @@ -304,7 +306,10 @@ void track_magnet_particles( double x0_solenoid, double y0_solenoid, int64_t rbend_model, // -1: not used, 0: auto, 1: curved body, 2: straight body + int64_t rbend_compensate_sagitta, double rbend_shift, + double rbend_angle_diff, + double length_straight, int64_t body_active, int64_t edge_entry_active, int64_t edge_exit_active, @@ -323,48 +328,67 @@ void track_magnet_particles( double factor_knl_ksl = 1.0; // Used only for rbend-straight-body - double rbend_half_angle = 0.; - double cos_rbha = 0.; - double sin_rbha = 0.; + double theta_in = 0.; + double theta_out = 0.; + double cos_theta_in = 1.; + double sin_theta_in = 0.; + double cos_theta_out = 1.; + double sin_theta_out = 0.; double length_curved = 0.; - double dx_rb = 0.; - + double x0_mid = 0.; + double x0_in = 0.; + double x0_out = 0.; if (rbend_model == 0){ // auto mode, curved body rbend_model = 1; } + double const angle = h * length; if (rbend_model == 1){ // curved body - double const angle = h * length; - edge_entry_angle += angle / 2.0; - edge_exit_angle += angle / 2.0; + edge_entry_angle += (angle - rbend_angle_diff) / 2.0; + edge_exit_angle += (angle + rbend_angle_diff) / 2.0; } else if (rbend_model == 2){ // straight body - double sinc_rbha; - rbend_half_angle = h * length / 2; - if (fabs(rbend_half_angle) > 1e-10){ - sin_rbha = sin(rbend_half_angle); - cos_rbha = cos(rbend_half_angle); - sinc_rbha = sin_rbha / rbend_half_angle; + theta_in = (angle - rbend_angle_diff) / 2.0; + if (fabs(theta_in) > 1e-10){ + sin_theta_in = sin(theta_in); + cos_theta_in = cos(theta_in); } - else { - sinc_rbha = 1.; - cos_rbha = 1.; + theta_out = (angle + rbend_angle_diff) / 2.0; + if (fabs(theta_out) > 1e-10){ + sin_theta_out = sin(theta_out); + cos_theta_out = cos(theta_out); } + length_curved = length; - length = length_curved * sinc_rbha; - if (fabs(rbend_half_angle) > 1e-10){ + length = length_straight; + + x0_mid -= rbend_shift; + + if (rbend_compensate_sagitta && fabs(angle) > 1e-10){ // shift by half the sagitta - dx_rb = 0.5 / h * (1 - cos_rbha) + rbend_shift; - }; + double cos_rbha = cos(angle / 2.); + x0_mid += 0.5 / h * (1 - cos_rbha); + } + + x0_in = x0_mid; + x0_out = x0_mid; + if (fabs(angle) > 1e-10){ + double const px0_in = sin(theta_in); + double const px0_mid = px0_in - h * length_straight / 2; + double const sqrt_mid = sqrt(1 - px0_mid * px0_mid); + x0_in -= 1/h *(sqrt_mid - cos_theta_in); + x0_out += 1/h * (cos_theta_out - sqrt_mid); + } + ; h = 0; // treat magnet as straight // We are entering the fringe with an angle, the linear fringe // needs to come from the expansion around the right angle - edge_entry_angle_fdown += rbend_half_angle; - edge_exit_angle_fdown += rbend_half_angle; + edge_entry_angle_fdown += theta_in; + edge_exit_angle_fdown += theta_out; } double core_length, core_length_curved,factor_knl_ksl_body, @@ -383,8 +407,14 @@ void track_magnet_particles( VSWAP(edge_entry_angle_fdown, edge_exit_angle_fdown); VSWAP(edge_entry_fint, edge_exit_fint); VSWAP(edge_entry_hgap, edge_exit_hgap); - rbend_half_angle = -rbend_half_angle; - sin_rbha = -sin_rbha; + VSWAP(theta_in, theta_out); + VSWAP(cos_theta_in, cos_theta_out); + VSWAP(sin_theta_in, sin_theta_out); + theta_in = -theta_in; + theta_out = -theta_out; + sin_theta_in = -sin_theta_in; + sin_theta_out = -sin_theta_out; + VSWAP(x0_in, x0_out); } else { core_length = length * weight; @@ -430,8 +460,9 @@ void track_magnet_particles( if (rbend_model == 2){ // straight body --> curvature in the edges START_PER_PARTICLE_BLOCK(part0, part); - YRotation_single_particle(part, -sin_rbha, cos_rbha, -sin_rbha/cos_rbha); - LocalParticle_add_to_x(part, -dx_rb); + YRotation_single_particle(part, -sin_theta_in, cos_theta_in, + -sin_theta_in/cos_theta_in); + LocalParticle_add_to_x(part, x0_in); END_PER_PARTICLE_BLOCK; } @@ -581,8 +612,9 @@ void track_magnet_particles( if (rbend_model == 2){ // straight body --> curvature in the edges START_PER_PARTICLE_BLOCK(part0, part); - LocalParticle_add_to_x(part, dx_rb); // shift by half sagitta - YRotation_single_particle(part, -sin_rbha, cos_rbha, -sin_rbha/cos_rbha); + LocalParticle_add_to_x(part, -x0_out); // shift by half sagitta + YRotation_single_particle(part, -sin_theta_out, cos_theta_out, + -sin_theta_out/cos_theta_out); END_PER_PARTICLE_BLOCK; } } diff --git a/xtrack/beam_elements/elements_src/track_magnet.template.h b/xtrack/beam_elements/elements_src/track_magnet.template.h index 2380f4390..8ea6f2bfd 100644 --- a/xtrack/beam_elements/elements_src/track_magnet.template.h +++ b/xtrack/beam_elements/elements_src/track_magnet.template.h @@ -191,6 +191,7 @@ void track_magnet_body_single_particle( } + GPUFUN void track_magnet_particles( double const weight, @@ -224,7 +225,10 @@ void track_magnet_particles( double x0_solenoid, double y0_solenoid, int64_t rbend_model, // -1: not used, 0: auto, 1: curved body, 2: straight body + int64_t rbend_compensate_sagitta, double rbend_shift, + double rbend_angle_diff, + double length_straight, int64_t body_active, int64_t edge_entry_active, int64_t edge_exit_active, @@ -243,48 +247,67 @@ void track_magnet_particles( double factor_knl_ksl = 1.0; // Used only for rbend-straight-body - double rbend_half_angle = 0.; - double cos_rbha = 0.; - double sin_rbha = 0.; + double theta_in = 0.; + double theta_out = 0.; + double cos_theta_in = 1.; + double sin_theta_in = 0.; + double cos_theta_out = 1.; + double sin_theta_out = 0.; double length_curved = 0.; - double dx_rb = 0.; - + double x0_mid = 0.; + double x0_in = 0.; + double x0_out = 0.; if (rbend_model == 0){ // auto mode, curved body rbend_model = 1; } + double const angle = h * length; if (rbend_model == 1){ // curved body - double const angle = h * length; - edge_entry_angle += angle / 2.0; - edge_exit_angle += angle / 2.0; + edge_entry_angle += (angle - rbend_angle_diff) / 2.0; + edge_exit_angle += (angle + rbend_angle_diff) / 2.0; } else if (rbend_model == 2){ // straight body - double sinc_rbha; - rbend_half_angle = h * length / 2; - if (fabs(rbend_half_angle) > 1e-10){ - sin_rbha = sin(rbend_half_angle); - cos_rbha = cos(rbend_half_angle); - sinc_rbha = sin_rbha / rbend_half_angle; + theta_in = (angle - rbend_angle_diff) / 2.0; + if (fabs(theta_in) > 1e-10){ + sin_theta_in = sin(theta_in); + cos_theta_in = cos(theta_in); } - else { - sinc_rbha = 1.; - cos_rbha = 1.; + theta_out = (angle + rbend_angle_diff) / 2.0; + if (fabs(theta_out) > 1e-10){ + sin_theta_out = sin(theta_out); + cos_theta_out = cos(theta_out); } + length_curved = length; - length = length_curved * sinc_rbha; - if (fabs(rbend_half_angle) > 1e-10){ + length = length_straight; + + x0_mid -= rbend_shift; + + if (rbend_compensate_sagitta && fabs(angle) > 1e-10){ // shift by half the sagitta - dx_rb = 0.5 / h * (1 - cos_rbha) + rbend_shift; - }; + double cos_rbha = cos(angle / 2.); + x0_mid += 0.5 / h * (1 - cos_rbha); + } + + x0_in = x0_mid; + x0_out = x0_mid; + if (fabs(angle) > 1e-10){ + double const px0_in = sin(theta_in); + double const px0_mid = px0_in - h * length_straight / 2; + double const sqrt_mid = sqrt(1 - px0_mid * px0_mid); + x0_in -= 1/h *(sqrt_mid - cos_theta_in); + x0_out += 1/h * (cos_theta_out - sqrt_mid); + } + ; h = 0; // treat magnet as straight // We are entering the fringe with an angle, the linear fringe // needs to come from the expansion around the right angle - edge_entry_angle_fdown += rbend_half_angle; - edge_exit_angle_fdown += rbend_half_angle; + edge_entry_angle_fdown += theta_in; + edge_exit_angle_fdown += theta_out; } double core_length, core_length_curved,factor_knl_ksl_body, @@ -303,8 +326,14 @@ void track_magnet_particles( VSWAP(edge_entry_angle_fdown, edge_exit_angle_fdown); VSWAP(edge_entry_fint, edge_exit_fint); VSWAP(edge_entry_hgap, edge_exit_hgap); - rbend_half_angle = -rbend_half_angle; - sin_rbha = -sin_rbha; + VSWAP(theta_in, theta_out); + VSWAP(cos_theta_in, cos_theta_out); + VSWAP(sin_theta_in, sin_theta_out); + theta_in = -theta_in; + theta_out = -theta_out; + sin_theta_in = -sin_theta_in; + sin_theta_out = -sin_theta_out; + VSWAP(x0_in, x0_out); } else { core_length = length * weight; @@ -350,8 +379,9 @@ void track_magnet_particles( if (rbend_model == 2){ // straight body --> curvature in the edges START_PER_PARTICLE_BLOCK(part0, part); - YRotation_single_particle(part, -sin_rbha, cos_rbha, -sin_rbha/cos_rbha); - LocalParticle_add_to_x(part, -dx_rb); + YRotation_single_particle(part, -sin_theta_in, cos_theta_in, + -sin_theta_in/cos_theta_in); + LocalParticle_add_to_x(part, x0_in); END_PER_PARTICLE_BLOCK; } @@ -501,8 +531,9 @@ void track_magnet_particles( if (rbend_model == 2){ // straight body --> curvature in the edges START_PER_PARTICLE_BLOCK(part0, part); - LocalParticle_add_to_x(part, dx_rb); // shift by half sagitta - YRotation_single_particle(part, -sin_rbha, cos_rbha, -sin_rbha/cos_rbha); + LocalParticle_add_to_x(part, -x0_out); // shift by half sagitta + YRotation_single_particle(part, -sin_theta_out, cos_theta_out, + -sin_theta_out/cos_theta_out); END_PER_PARTICLE_BLOCK; } } diff --git a/xtrack/beam_elements/elements_src/variable_solenoid.h b/xtrack/beam_elements/elements_src/variable_solenoid.h index 3332fff8a..e2cc9f75a 100644 --- a/xtrack/beam_elements/elements_src/variable_solenoid.h +++ b/xtrack/beam_elements/elements_src/variable_solenoid.h @@ -75,7 +75,10 @@ void VariableSolenoid_track_local_particle( /*x0_solenoid*/ x0, /*y0_solenoid*/ y0, /*rbend_model*/ -1, // not rbend - /*rbend_shift*/ 0., + /*rbend_compensate_sagitta*/ 0, // not rbend + /*rbend_shift*/ 0., // not rbend + /*rbend_angle_diff*/ 0., // not rbend + /*length_straight*/ 0., // not rbend /*body_active*/ 1, /*edge_entry_active*/ 0, /*edge_exit_active*/ 0, diff --git a/xtrack/beam_elements/slice_elements_edge.py b/xtrack/beam_elements/slice_elements_edge.py index 7572c5d4c..281a6ceac 100644 --- a/xtrack/beam_elements/slice_elements_edge.py +++ b/xtrack/beam_elements/slice_elements_edge.py @@ -273,7 +273,6 @@ def get_equivalent_element(self): def _propagate_survey(self, v, w, backtrack): if self._parent.rbend_model == "straight-body": - rbend_shift_tot = self._parent.sagitta / 2 + self._parent.rbend_shift if backtrack: if abs(self._parent.angle) > 1e-10: # avoid numerical issues v, w = survey_advance_element( @@ -282,8 +281,8 @@ def _propagate_survey(self, v, w, backtrack): length = 0, angle = 0, tilt = 0, - ref_shift_x = -rbend_shift_tot * np.cos(self._parent.rot_s_rad), - ref_shift_y = -rbend_shift_tot * np.sin(self._parent.rot_s_rad), + ref_shift_x = self._parent._x0_in * np.cos(self._parent.rot_s_rad), + ref_shift_y = self._parent._x0_in * np.sin(self._parent.rot_s_rad), ref_rot_x_rad = 0, ref_rot_y_rad = 0, ref_rot_s_rad = 0, @@ -292,7 +291,7 @@ def _propagate_survey(self, v, w, backtrack): v = v, w = w, length = 0, - angle = -self._parent.angle / 2., + angle = -self._parent._angle_in, tilt = self._parent.rot_s_rad, ref_shift_x = 0, ref_shift_y = 0, @@ -305,7 +304,7 @@ def _propagate_survey(self, v, w, backtrack): v = v, w = w, length = 0, - angle = self._parent.angle / 2., + angle = self._parent._angle_in, tilt = self._parent.rot_s_rad, ref_shift_x = 0, ref_shift_y = 0, @@ -320,8 +319,8 @@ def _propagate_survey(self, v, w, backtrack): length = 0, angle = 0, tilt = 0, - ref_shift_x = rbend_shift_tot * np.cos(self._parent.rot_s_rad), - ref_shift_y = rbend_shift_tot * np.sin(self._parent.rot_s_rad), + ref_shift_x = -self._parent._x0_in * np.cos(self._parent.rot_s_rad), + ref_shift_y = -self._parent._x0_in * np.sin(self._parent.rot_s_rad), ref_rot_x_rad = 0, ref_rot_y_rad = 0, ref_rot_s_rad = 0, @@ -359,13 +358,12 @@ def get_equivalent_element(self): def _propagate_survey(self, v, w, backtrack): if self._parent.rbend_model == "straight-body": - rbend_shift_tot = self._parent.sagitta / 2 + self._parent.rbend_shift if backtrack: v, w = survey_advance_element( v = v, w = w, length = 0, - angle = -self._parent.angle / 2., + angle = -self._parent._angle_out, tilt = self._parent.rot_s_rad, ref_shift_x = 0, ref_shift_y = 0, @@ -380,8 +378,8 @@ def _propagate_survey(self, v, w, backtrack): length = 0, angle = 0, tilt = 0, - ref_shift_x = rbend_shift_tot * np.cos(self._parent.rot_s_rad), - ref_shift_y = rbend_shift_tot * np.sin(self._parent.rot_s_rad), + ref_shift_x = -self._parent._x0_out * np.cos(self._parent.rot_s_rad), + ref_shift_y = -self._parent._x0_out * np.sin(self._parent.rot_s_rad), ref_rot_x_rad = 0, ref_rot_y_rad = 0, ref_rot_s_rad = 0, @@ -394,8 +392,8 @@ def _propagate_survey(self, v, w, backtrack): length = 0, angle = 0, tilt = 0, - ref_shift_x = -rbend_shift_tot * np.cos(self._parent.rot_s_rad), - ref_shift_y = -rbend_shift_tot * np.sin(self._parent.rot_s_rad), + ref_shift_x = self._parent._x0_out * np.cos(self._parent.rot_s_rad), + ref_shift_y = self._parent._x0_out * np.sin(self._parent.rot_s_rad), ref_rot_x_rad = 0, ref_rot_y_rad = 0, ref_rot_s_rad = 0, @@ -404,7 +402,7 @@ def _propagate_survey(self, v, w, backtrack): v = v, w = w, length = 0, - angle = self._parent.angle / 2., + angle = self._parent._angle_out, tilt = self._parent.rot_s_rad, ref_shift_x = 0, ref_shift_y = 0, diff --git a/xtrack/headers/atomicadd.h b/xtrack/headers/atomicadd.h index 360dda64d..0a44925b3 100644 --- a/xtrack/headers/atomicadd.h +++ b/xtrack/headers/atomicadd.h @@ -3,348 +3,11 @@ // Copyright (c) CERN, 2025. // // ########################################### // -#ifndef _ATOMICADD_H_ -#define _ATOMICADD_H_ +#ifndef XTRACK_ATOMICADD_H_DEPRECATED +#define XTRACK_ATOMICADD_H_DEPRECATED -#include +// The xtrack/headers/atomicadd.h header is deprecated. Use xobjects/headers/atomicadd.h instead. +#include "xobjects/headers/atomicadd.h" -// Ensure no conflicting atomicAdd macro is defined earlier for CPU or OpenCL -#ifndef XO_CONTEXT_CUDA -#ifdef atomicAdd - #warning "Warning: atomicAdd macro already defined, undefining it." - #undef atomicAdd -#endif // atomicAdd -#endif // XO_CONTEXT_CUDA - - -// ################################# // -// ### CPU Contexts ### // -// ################################# // - -#ifdef XO_CONTEXT_CPU -#include - -#ifdef XO_CONTEXT_CPU_OPENMP - #ifndef _OPENMP - #error "XO_CONTEXT_CPU_OPENMP set, but compiled without -fopenmp" - #endif - // OpenMP atomic capture gives us read+write atomically - #define OMP_ATOMIC_CAPTURE _Pragma("omp atomic capture") -#else - #define OMP_ATOMIC_CAPTURE // No OpenMP: non-atomic fallback -#endif // XO_CONTEXT_CPU_OPENMP - -// Macro to define atomicAdd for different types, will be overloaded via _Generic. -#define DEF_ATOMIC_ADD(T, SUF) \ -GPUFUN T atomicAdd_##SUF(GPUGLMEM T *addr, T val) { \ - T old; \ - const T inc = (T)val; \ - OMP_ATOMIC_CAPTURE \ - { old = *addr; *addr = *addr + inc; } \ - return old; \ -} - -DEF_ATOMIC_ADD(int8_t , i8) -DEF_ATOMIC_ADD(int16_t, i16) -DEF_ATOMIC_ADD(int32_t, i32) -DEF_ATOMIC_ADD(int64_t, i64) -DEF_ATOMIC_ADD(uint8_t , u8) -DEF_ATOMIC_ADD(uint16_t, u16) -DEF_ATOMIC_ADD(uint32_t, u32) -DEF_ATOMIC_ADD(uint64_t, u64) -DEF_ATOMIC_ADD(float , f32) -DEF_ATOMIC_ADD(double , f64) - -// Using _Generic to select the right function based on type (since C11). -// See https://en.cppreference.com/w/c/language/generic.html -#define atomicAdd(addr, val) _Generic((addr), \ - int8_t*: atomicAdd_i8, \ - int16_t*: atomicAdd_i16, \ - int32_t*: atomicAdd_i32, \ - int64_t*: atomicAdd_i64, \ - uint8_t*: atomicAdd_u8, \ - uint16_t*: atomicAdd_u16, \ - uint32_t*: atomicAdd_u32, \ - uint64_t*: atomicAdd_u64, \ - float*: atomicAdd_f32, \ - double*: atomicAdd_f64 \ -)(addr, (val)) -#endif // XO_CONTEXT_CPU - - -// ################################# // -// ### GPU Contexts ### // -// ################################# // - -/* -------------- Design notes for GPU atomics --------------------------------- -* We will go for a unified approach for CUDA and OpenCL, with macros where needed. -* We implement 8/16-bit atomic add by atomically CAS'ing the *containing* - * 32-bit word. Only the target byte/halfword lane is modified; the neighbor - * lane is preserved. This is linearisable: each successful CAS is one atomic - * RMW on the 32-bit word. - * Assumptions: little-endian lane layout (true on NVIDIA) and natural alignment - * of the 16-bit addresses (addr % 2 == 0). 8-bit has no alignment requirement. - * Return value matches CUDA semantics: the **old** value at *addr* (fetch-add). - * ---------------------------------------------------------------------------*/ - -// Info: a CAS function (Compare-And-Swap) takes three arguments: a pointer to -// the value to be modified, the expected old value, and the new value. The -// second argument helps to recognise if another thread has modified the value -// in the meantime. The function compares the value at the address with the -// expected values, and if they are the same, it writes the new value at the -// address. In any case, the function returns the actual old value at the -// address. If the returned value is different from the expected value, it -// means that another thread has modified the value in the meantime. - -// Define types and macros for CUDA and OpenCL -// ------------------------------------------- -#if defined(XO_CONTEXT_CUDA) - // CUDA compiler may not have , so define the types if needed. - #ifdef __CUDACC_RTC__ - // NVRTC (CuPy RawModule default) can’t see , so detect it via __CUDACC_RTC__ - typedef signed char int8_t; - typedef short int16_t; - typedef int int32_t; - typedef long long int64_t; - typedef unsigned char uint8_t; - typedef unsigned short uint16_t; - typedef unsigned int uint32_t; - typedef unsigned long long uint64_t; - #else - // Alternatively, NVCC path is fine with host headers - #include - #endif // __CUDACC_RTC__ - #define GPUVOLATILE GPUGLMEM - #define __XT_CAS_U32(ptr, exp, val) atomicCAS((GPUVOLATILE uint32_t*)(ptr), (uint32_t)(exp), (uint32_t)(val)) - #define __XT_CAS_U64(ptr, exp, val) atomicCAS((GPUVOLATILE uint64_t*)(ptr), (uint64_t)(exp), (uint64_t)(val)) - #define __XT_AS_U32_FROM_F32(x) __float_as_uint(x) - #define __XT_AS_F32_FROM_U32(x) __uint_as_float(x) - #define __XT_AS_U64_FROM_F64(x) __double_as_longlong(x) - #define __XT_AS_F64_FROM_U64(x) __longlong_as_double(x) -#elif defined(XO_CONTEXT_CL) - // It seems OpenCL already has the types from defined. - // typedef char int8_t; - // typedef short int16_t; - // typedef int int32_t; - // typedef long int64_t; - // typedef unsigned char uint8_t; - // typedef unsigned short uint16_t; - // typedef unsigned int uint32_t; - // typedef unsigned long uint64_t; - #define GPUVOLATILE GPUGLMEM volatile - #if __OPENCL_C_VERSION__ < 110 - // Map 1.0 "atom_*" names to 1.1+ "atomic_*" - #pragma OPENCL EXTENSION cl_khr_global_int32_base_atomics : enable - #define atomic_add atom_add - #define atomic_cmpxchg atom_cmpxchg - #endif - #pragma OPENCL EXTENSION cl_khr_int64_base_atomics : enable - #pragma OPENCL EXTENSION cl_khr_fp64 : enable - #define __XT_CAS_U32(ptr, exp, val) atomic_cmpxchg((GPUVOLATILE uint32_t*)(ptr), (uint32_t)(exp), (uint32_t)(val)) - #define __XT_CAS_U64(ptr, exp, val) atom_cmpxchg((GPUVOLATILE uint64_t*)(ptr), (uint64_t)(exp), (uint64_t)(val)) - #define __XT_AS_U32_FROM_F32(x) as_uint(x) - #define __XT_AS_F32_FROM_U32(x) as_float(x) - #define __XT_AS_U64_FROM_F64(x) as_ulong(x) - #define __XT_AS_F64_FROM_U64(x) as_double(x) -#endif // XO_CONTEXT_CUDA / XO_CONTEXT_CL - -// Define atomic functions per type -// -------------------------------- -#if defined(XO_CONTEXT_CUDA) || defined(XO_CONTEXT_CL) - // Helper: compute (base 32-bit word pointer, shift, mask) for a byte in that word. - GPUFUN inline void __xt_lane8(GPUVOLATILE void* addr, GPUVOLATILE uint32_t** w, uint32_t* sh, uint32_t* mask){ - size_t a = (size_t)addr; - *w = (GPUVOLATILE uint32_t*)(a & ~(size_t)3); // word: align down to 4-byte boundary - *sh = (uint32_t)((a & (size_t)3) * 8U); // shift 0,8,16,24 depending on byte lane - *mask = (uint32_t)(0xFFu << *sh); - } - // Helper: same for a halfword (16-bit) in the containing 32-bit word. - GPUFUN inline void __xt_lane16(GPUVOLATILE void* addr, GPUVOLATILE uint32_t** w, uint32_t* sh, uint32_t* mask){ - size_t a = (size_t)addr; - *w = (GPUVOLATILE uint32_t*)(a & ~(size_t)3); // word: align down to 4-byte boundary - *sh = (uint32_t)((a & (size_t)2) ? 16U : 0U); // shift0 or 16 depending on halfword - *mask = (uint32_t)(0xFFFFU << *sh); - } - - // ---------------- 8-bit: int8_t / uint8_t (CAS on 32-bit word) -------------- - GPUFUN int8_t atomicAdd_i8(GPUVOLATILE int8_t* addr, int8_t val){ - GPUVOLATILE uint32_t *w; - uint32_t sh, mask; - __xt_lane8(addr, &w, &sh, &mask); - uint32_t old = *w, assumed, b, nb, nw; // byte, newbyte, newword - do { - assumed = old; - b = (assumed & mask) >> sh; // Extract current 8-bit lane - nb = (uint32_t)((uint8_t)b + (uint8_t)val); // Add in modulo-256 (two's complement) - nw = (assumed & ~mask) | ((nb & 0xFFU) << sh); // Merge back updated lane; leave neighbor lanes intact - // Try to publish; if someone raced us, retry with their value - old = __XT_CAS_U32(w, assumed, nw); - } while (old != assumed); - return (int8_t)((assumed & mask) >> sh); - } - GPUFUN uint8_t atomicAdd_u8(GPUVOLATILE uint8_t* addr, uint8_t val){ - GPUVOLATILE uint32_t *w; - uint32_t sh, mask; - __xt_lane8(addr, &w, &sh, &mask); - uint32_t old = *w, assumed, b, nb, nw; - do { - assumed = old; - b = (assumed & mask) >> sh; - nb = (uint32_t)(b + val); /* mod 256 */ - nw = (assumed & ~mask) | ((nb & 0xFFU) << sh); - old = __XT_CAS_U32(w, assumed, nw); - } while (old != assumed); - return (uint8_t)((assumed & mask) >> sh); - } - - // ---------------- 16-bit: int16_t / uint16_t (CAS on 32-bit word) ----------- - GPUFUN int16_t atomicAdd_i16(GPUVOLATILE int16_t* addr, int16_t val){ - GPUVOLATILE uint32_t* w; - uint32_t sh, mask; - __xt_lane16(addr, &w, &sh, &mask); - uint32_t old = *w, assumed, b, nb, nw; - do { - assumed = old; - b = (assumed & mask) >> sh; - nb = (uint32_t)((uint16_t)b + (uint16_t)val); - nw = (assumed & ~mask) | ((nb & 0xFFFFU) << sh); - old = __XT_CAS_U32(w, assumed, nw); - } while (old != assumed); - return (int16_t)((assumed & mask) >> sh); - } - GPUFUN uint16_t atomicAdd_u16(GPUVOLATILE uint16_t* addr, uint16_t val){ - GPUVOLATILE uint32_t* w; - uint32_t sh, mask; - __xt_lane16(addr, &w, &sh, &mask); - uint32_t old = *w, assumed, b, nb, nw; - do { - assumed = old; - b = (assumed & mask) >> sh; - nb = (uint32_t)(b + val); - nw = (assumed & ~mask) | ((nb & 0xFFFFU) << sh); - old = __XT_CAS_U32(w, assumed, nw); - } while (old != assumed); - return (uint16_t)((assumed & mask) >> sh); - } - - // ---------------- 32-bit: int32_t / uint32_t (built-in) ----------- - GPUFUN int32_t atomicAdd_i32(GPUVOLATILE int32_t* addr, int32_t val){ - #ifdef XO_CONTEXT_CUDA - return atomicAdd(addr, val); - #else // XO_CONTEXT_CL - return atomic_add(addr, val); - #endif // XO_CONTEXT_CUDA / XO_CONTEXT_CL - } - GPUFUN uint32_t atomicAdd_u32(GPUVOLATILE uint32_t* addr, uint32_t val){ - #ifdef XO_CONTEXT_CUDA - return atomicAdd(addr, val); - #else // XO_CONTEXT_CL - return atomic_add(addr, val); - #endif // XO_CONTEXT_CUDA / XO_CONTEXT_CL - } - - // ---------------- 64-bit: int64_t / uint64_t (built-in or CAS on 64-bit word) ----------- - GPUFUN int64_t atomicAdd_i64(GPUVOLATILE int64_t* addr, int64_t val){ - uint64_t old, nw; - do { - old = *addr; - nw = old + val; - } while (__XT_CAS_U64((GPUVOLATILE uint64_t*)addr, old, nw) != old); - return old; - } - GPUFUN uint64_t atomicAdd_u64(GPUVOLATILE uint64_t* addr, uint64_t val){ - #ifdef XO_CONTEXT_CUDA - return atomicAdd(addr, val); - #else // XO_CONTEXT_CL - return atom_add(addr, val); - #endif // XO_CONTEXT_CUDA / XO_CONTEXT_CL - } - - // ---------------- 32-bit: float (built-in or CAS on 32-bit word) ----------- - GPUFUN float atomicAdd_f32(GPUVOLATILE float* addr, float val){ - #ifdef XO_CONTEXT_CUDA - return atomicAdd(addr, val); - #else // XO_CONTEXT_CL - uint32_t old, nw; - do { - old = __XT_AS_U32_FROM_F32(*addr); - nw = __XT_AS_U32_FROM_F32(__XT_AS_F32_FROM_U32(old) + val); - } while (__XT_CAS_U32((GPUVOLATILE uint32_t*)addr, old, nw) != old); - return __XT_AS_F32_FROM_U32(old); - #endif // XO_CONTEXT_CUDA / XO_CONTEXT_CL - } - - // ---------------- 64-bit: float (built-in or CAS on 64-bit word) ----------- - GPUFUN double atomicAdd_f64(GPUVOLATILE double* addr, double val){ - #if __CUDA_ARCH__ >= 600 - return atomicAdd(addr, val); - #else // XO_CONTEXT_CL || __CUDA_ARCH__ < 600 - uint64_t old, nw; - do { - old = __XT_AS_U64_FROM_F64(*addr); - nw = __XT_AS_U64_FROM_F64(__XT_AS_F64_FROM_U64(old) + val); - } while (__XT_CAS_U64((GPUVOLATILE uint64_t*)addr, old, nw) != old); - return __XT_AS_F64_FROM_U64(old); - #endif // __CUDA_ARCH__ >= 600 / XO_CONTEXT_CL - } -#endif // defined(XO_CONTEXT_CUDA) || defined(XO_CONTEXT_CL) - -// Define the overloaded function -// ------------------------------ -#ifdef XO_CONTEXT_CUDA - // NVRTC (CuPy RawModule) usually compiles under extern "C". - // In C, function overloading is not possible, but we can cheat by doing it in - // C++ (with a different name to avoid clashes with the built-in atomicAdd). - // This function will then be remapped to atomicAdd() via a macro in C. - #ifdef __cplusplus - extern "C++" { - - GPUFUN int8_t xt_atomicAdd(GPUVOLATILE int8_t* p, int8_t v) { return atomicAdd_i8 (p, v); } - GPUFUN uint8_t xt_atomicAdd(GPUVOLATILE uint8_t* p, uint8_t v) { return atomicAdd_u8 (p, v); } - GPUFUN int16_t xt_atomicAdd(GPUVOLATILE int16_t* p, int16_t v) { return atomicAdd_i16(p, v); } - GPUFUN uint16_t xt_atomicAdd(GPUVOLATILE uint16_t*p, uint16_t v) { return atomicAdd_u16(p, v); } - GPUFUN int64_t xt_atomicAdd(GPUVOLATILE int64_t*p, int64_t v) { return atomicAdd_i64(p, v); } - - // Existing type definitions: forward to CUDA built-ins - GPUFUN int32_t xt_atomicAdd(GPUVOLATILE int32_t* p, int32_t v) { return ::atomicAdd(p, v); } - GPUFUN uint32_t xt_atomicAdd(GPUVOLATILE uint32_t* p, uint32_t v) { return ::atomicAdd(p, v); } - GPUFUN uint64_t xt_atomicAdd(GPUVOLATILE uint64_t* p, uint64_t v) { return ::atomicAdd(p, v); } - GPUFUN float xt_atomicAdd(GPUVOLATILE float* p, float v) { return ::atomicAdd(p, v); } - #if __CUDA_ARCH__ >= 600 - GPUFUN double xt_atomicAdd(GPUVOLATILE double* p, double v) { return ::atomicAdd(p, v); } - #else - GPUFUN double xt_atomicAdd(GPUVOLATILE double* p, double v) { return atomicAdd_f64(p, v); } - #endif - - } - #endif // __cplusplus - - // ---------- Global remap of the public name on device code ---------- - // Define AFTER the wrappers so we don't macro-rewrite our own calls. - #ifdef atomicAdd - #undef atomicAdd - #endif - #define atomicAdd(ptr, val) xt_atomicAdd((ptr), (val)) -#endif /* XO_CONTEXT_CUDA */ - -#ifdef XO_CONTEXT_CL - #if !__has_attribute(overloadable) - #error "The current OpenCL compiler/architecture does not support __attribute__((overloadable))" - #endif - #define OCL_OVERLOAD __attribute__((overloadable)) - GPUFUN int8_t OCL_OVERLOAD atomicAdd(GPUVOLATILE int8_t* p, int8_t v) { return atomicAdd_i8 (p, v); } - GPUFUN uint8_t OCL_OVERLOAD atomicAdd(GPUVOLATILE uint8_t* p, uint8_t v) { return atomicAdd_u8 (p, v); } - GPUFUN int16_t OCL_OVERLOAD atomicAdd(GPUVOLATILE int16_t* p, int16_t v) { return atomicAdd_i16(p, v); } - GPUFUN uint16_t OCL_OVERLOAD atomicAdd(GPUVOLATILE uint16_t*p, uint16_t v) { return atomicAdd_u16(p, v); } - GPUFUN int64_t OCL_OVERLOAD atomicAdd(GPUVOLATILE int64_t*p, int64_t v) { return atomicAdd_i64(p, v); } - GPUFUN float OCL_OVERLOAD atomicAdd(GPUVOLATILE float* p, float v) { return atomicAdd_f32(p, v); } - GPUFUN double OCL_OVERLOAD atomicAdd(GPUVOLATILE double* p, double v) { return atomicAdd_f64(p, v); } - - // Existing type definitions: forward to OpenCL built-ins - GPUFUN int32_t OCL_OVERLOAD atomicAdd(GPUVOLATILE int32_t* p, int32_t v) { return atomic_add(p, v); } - GPUFUN uint32_t OCL_OVERLOAD atomicAdd(GPUVOLATILE uint32_t* p, uint32_t v) { return atomic_add(p, v); } - GPUFUN uint64_t OCL_OVERLOAD atomicAdd(GPUVOLATILE uint64_t* p, uint64_t v) { return atom_add(p, v); } -#endif // XO_CONTEXT_CL - -#endif //_ATOMICADD_H_ +#endif // XTRACK_ATOMICADD_H_DEPRECATED diff --git a/xtrack/headers/constants.h b/xtrack/headers/constants.h index 9159b38ff..d876cc4df 100644 --- a/xtrack/headers/constants.h +++ b/xtrack/headers/constants.h @@ -58,4 +58,20 @@ #define MASS_ELECTRON (9.1093837e-31) #endif /* !defined( MASS_ELECTRON ) */ +#if !defined( SQRT3 ) + #define SQRT3 1.732050807568877 +#endif /* !defined( SQRT3 ) */ + +#if !defined( SQRT2 ) + #define SQRT2 (1.414213562373095048801688724209698078569671875376948073176679738) +#endif /* !defined( SQRT2 ) */ + +#if !defined( TWO_OVER_SQRT_PI ) + #define TWO_OVER_SQRT_PI (1.128379167095512573896158903121545171688101258657997713688171443418) +#endif /* !defined( TWO_OVER_SQRT_PI ) */ + +#if !defined( ALPHA_EM ) + #define ALPHA_EM 0.0072973525693 +#endif /* !defined( ALPHA_EM ) */ + #endif /* XTRACK_CONSTANTS_H */ diff --git a/xtrack/headers/track.h b/xtrack/headers/track.h index d803bf46b..b55b4b6bc 100644 --- a/xtrack/headers/track.h +++ b/xtrack/headers/track.h @@ -1,7 +1,14 @@ +// copyright ############################### // +// This file is part of the Xtrack Package. // +// Copyright (c) CERN, 2025. // +// ######################################### // + #ifndef XTRACK_TRACK_H #define XTRACK_TRACK_H -#include +#include "xobjects/headers/common.h" +#include "xobjects/headers/atomicadd.h" +#include "xtrack/headers/constants.h" /* The particle tracking "decorators" for all the contexts. @@ -22,12 +29,6 @@ #define END_PER_PARTICLE_BLOCK \ } \ } - - #define VECTORIZE_OVER(INDEX_NAME, COUNT) \ - for (int INDEX_NAME = 0; INDEX_NAME < (COUNT); INDEX_NAME++) { - - #define END_VECTORIZE \ - } #endif // XO_CONTEXT_CPU_SERIAL #ifdef XO_CONTEXT_CPU_OPENMP @@ -48,14 +49,6 @@ } \ } \ } - - #define VECTORIZE_OVER(INDEX_NAME, COUNT) \ - _Pragma("omp parallel for") \ - for (int INDEX_NAME = 0; INDEX_NAME < (COUNT); INDEX_NAME++) { - - #define END_VECTORIZE \ - } - #endif // XO_CONTEXT_CPU_OPENMP @@ -67,14 +60,6 @@ #define END_PER_PARTICLE_BLOCK \ } - - #define VECTORIZE_OVER(INDEX_NAME, COUNT) { \ - int INDEX_NAME = blockDim.x * blockIdx.x + threadIdx.x; \ - if (INDEX_NAME < (COUNT)) { - - #define END_VECTORIZE \ - } \ - } #endif // XO_CONTEXT_CUDA @@ -86,54 +71,9 @@ #define END_PER_PARTICLE_BLOCK \ } - - #define VECTORIZE_OVER(INDEX_NAME, COUNT) \ - { \ - int INDEX_NAME = get_global_id(0); - - #define END_VECTORIZE \ - } #endif // XO_CONTEXT_CL -/* - Qualifier keywords for GPU and optimisation -*/ - -#ifdef XO_CONTEXT_CPU // for both serial and OpenMP - #define GPUKERN - #define GPUFUN static inline - #define GPUGLMEM - #define RESTRICT restrict -#endif - - -#ifdef XO_CONTEXT_CUDA - #define GPUKERN __global__ - #define GPUFUN __device__ - #define GPUGLMEM - #define RESTRICT -#endif // XO_CONTEXT_CUDA - - -#ifdef XO_CONTEXT_CL - #define GPUKERN __kernel - #define GPUFUN - #define GPUGLMEM __global - #define RESTRICT -#endif // XO_CONTEXT_CL - - -/* - Common maths-related macros -*/ - -#define POW2(X) ((X)*(X)) -#define POW3(X) ((X)*(X)*(X)) -#define POW4(X) ((X)*(X)*(X)*(X)) -#define NONZERO(X) ((X) != 0.0) - - #ifndef START_PER_PARTICLE_BLOCK #error "Unknown context, or the expected context (XO_CONTEXT_*) flag undefined. Try updating Xobjects?" #endif diff --git a/xtrack/internal_record.py b/xtrack/internal_record.py index 4aa1d78b6..fe4f1560f 100644 --- a/xtrack/internal_record.py +++ b/xtrack/internal_record.py @@ -33,8 +33,8 @@ ''' _RecordIndex_get_slot_source = r''' -#include -#include +#include "xobjects/headers/common.h" +#include "xobjects/headers/atomicadd.h" GPUFUN int64_t RecordIndex_get_slot(RecordIndex record_index){ diff --git a/xtrack/line.py b/xtrack/line.py index 421830ab9..64b4e1293 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -1369,7 +1369,8 @@ def build_particles( particle_ref : Particle object Reference particle defining the reference quantities (mass0, q0, p0c, gamma0, etc.). Its coordinates (x, py, y, py, zeta, delta) are ignored - unless `mode`='shift' is selected. + unless `mode`='shift' is selected. If this is None (default), the + reference particle associated with this line is used. num_particles : int Number of particles to be generated (used if provided coordinates are all scalar). @@ -3076,7 +3077,7 @@ def cycle(self, index_first_element=None, name_first_element=None, track_kernel = None if inplace: - self.unfreeze() + self.discard_tracker() self.element_names = new_element_names new_line = self else: @@ -4294,10 +4295,11 @@ def _freeze(self): self.element_names = tuple(self.element_names) def unfreeze(self): - - # Unfreeze the line. This is useful if you want to modify the line - # after it has been frozen (most likely by calling `build_tracker`). - + """See `Line.discard_tracker()`. This function is deprecated.""" + _print( + '`Line.unfreeze()` is deprecated and will be removed in future ' + 'versions. Please use `Line.discard_tracker()` instead.' + ) self.discard_tracker() def _frozen_check(self): @@ -4947,6 +4949,7 @@ def _get_attr_cache(self): '_parent_h': (('_parent', 'h'), None), '_parent_hxl': (('_parent', 'hxl'), None), '_parent_rbend_model': (('_parent', 'rbend_model'), None), + '_parent_rbend_angle_diff': (('_parent', 'rbend_angle_diff'), None), '_parent_voltage': (('_parent', 'voltage'), None), '_parent_lag': (('_parent', 'lag'), None), @@ -5833,18 +5836,30 @@ def _angle_rbend_correction_from_attr(attr): # Retrieve element_type from tracker cache (remove _end_point) element_type = attr.line.tracker._tracker_data_base._line_table.element_type[:-1] - mask_rbend_edges = ((element_type == 'ThinSliceRBendEntry') - | (element_type == 'ThinSliceRBendExit')) + mask_rbend_edge_entry = (element_type == 'ThinSliceRBendEntry') + mask_rbend_edge_exit = (element_type == 'ThinSliceRBendExit') + mask_rbend_body_slices = ((element_type == 'ThinSliceRBend') | (element_type == 'ThickSliceRBend')) mask_parent_is_rbend_straigth_body = (attr['_parent_rbend_model'] == 2) - mask_rbend_edges_straight_body = (mask_rbend_edges - & mask_parent_is_rbend_straigth_body) + mask_rbend_edges_entry_straight_body = (mask_rbend_edge_entry + & mask_parent_is_rbend_straigth_body) + mask_rbend_edges_exit_straight_body = (mask_rbend_edge_exit + & mask_parent_is_rbend_straigth_body) angle[mask_parent_is_rbend_straigth_body & mask_rbend_body_slices] = 0 - angle[mask_rbend_edges_straight_body] = 0.5 * ( - attr['_parent_h'][mask_rbend_edges_straight_body] - * attr['_parent_length'][mask_rbend_edges_straight_body]) + + # angle_in + angle[mask_rbend_edges_entry_straight_body] = 0.5 * (( + attr['_parent_h'][mask_rbend_edges_entry_straight_body] + * attr['_parent_length'][mask_rbend_edges_entry_straight_body]) + - attr['_parent_rbend_angle_diff'][mask_rbend_edges_entry_straight_body]) + + # angle_out + angle[mask_rbend_edges_exit_straight_body] = 0.5 * (( + attr['_parent_h'][mask_rbend_edges_exit_straight_body] + * attr['_parent_length'][mask_rbend_edges_exit_straight_body]) + + attr['_parent_rbend_angle_diff'][mask_rbend_edges_exit_straight_body]) return angle diff --git a/xtrack/monitors/beam_position_monitor.h b/xtrack/monitors/beam_position_monitor.h index 81e50569e..f8821eb6c 100644 --- a/xtrack/monitors/beam_position_monitor.h +++ b/xtrack/monitors/beam_position_monitor.h @@ -9,7 +9,7 @@ #ifndef XTRACK_BEAM_POSITION_MONITOR_H #define XTRACK_BEAM_POSITION_MONITOR_H -#include +#include "xtrack/headers/track.h" GPUFUN diff --git a/xtrack/monitors/beam_position_monitor.py b/xtrack/monitors/beam_position_monitor.py index 65b3e8ba9..3780261ab 100644 --- a/xtrack/monitors/beam_position_monitor.py +++ b/xtrack/monitors/beam_position_monitor.py @@ -12,7 +12,6 @@ from ..base_element import BeamElement from ..beam_elements import Marker from ..internal_record import RecordIndex -from ..general import _pkg_root class BeamPositionMonitorRecord(xo.Struct): @@ -40,8 +39,7 @@ class BeamPositionMonitor(BeamElement): properties = [field.name for field in BeamPositionMonitorRecord._fields] _extra_c_sources = [ - _pkg_root.joinpath('headers/atomicadd.h'), - _pkg_root.joinpath('monitors/beam_position_monitor.h') + '#include "xtrack/monitors/beam_position_monitor.h"', ] def __init__(self, *, particle_id_range=None, particle_id_start=None, num_particles=None, diff --git a/xtrack/monitors/beam_profile_monitor.h b/xtrack/monitors/beam_profile_monitor.h index 37daff113..dd720c7e8 100644 --- a/xtrack/monitors/beam_profile_monitor.h +++ b/xtrack/monitors/beam_profile_monitor.h @@ -9,7 +9,7 @@ #ifndef XTRACK_BEAM_PROFILE_MONITOR_H #define XTRACK_BEAM_PROFILE_MONITOR_H -#include +#include "xtrack/headers/track.h" GPUFUN diff --git a/xtrack/monitors/beam_profile_monitor.py b/xtrack/monitors/beam_profile_monitor.py index 24c7853ed..d23727010 100644 --- a/xtrack/monitors/beam_profile_monitor.py +++ b/xtrack/monitors/beam_profile_monitor.py @@ -12,7 +12,6 @@ from ..base_element import BeamElement from ..beam_elements import Marker from ..internal_record import RecordIndex -from ..general import _pkg_root class BeamProfileMonitorRecord(xo.Struct): counts_x = xo.Float64[:] @@ -45,8 +44,7 @@ class BeamProfileMonitor(BeamElement): properties = [field.name for field in BeamProfileMonitorRecord._fields] _extra_c_sources = [ - _pkg_root.joinpath('headers/atomicadd.h'), - _pkg_root.joinpath('monitors/beam_profile_monitor.h') + '#include "xtrack/monitors/beam_profile_monitor.h"', ] diff --git a/xtrack/monitors/beam_size_monitor.h b/xtrack/monitors/beam_size_monitor.h index 66bfa23cf..cb4102bad 100644 --- a/xtrack/monitors/beam_size_monitor.h +++ b/xtrack/monitors/beam_size_monitor.h @@ -13,8 +13,8 @@ GPUFUN -void BeamSizeMonitor_track_local_particle(BeamSizeMonitorData el, LocalParticle* part0){ - +void BeamSizeMonitor_track_local_particle(BeamSizeMonitorData el, LocalParticle* part0) +{ // get parameters int64_t const start_at_turn = BeamSizeMonitorData_get_start_at_turn(el); int64_t particle_id_start = BeamSizeMonitorData_get_particle_id_start(el); @@ -44,19 +44,19 @@ void BeamSizeMonitor_track_local_particle(BeamSizeMonitorData el, LocalParticle* double x = LocalParticle_get_x(part); double y = LocalParticle_get_y(part); - /*gpuglmem*/ double* count = BeamSizeMonitorRecord_getp1_count(record, slot); + GPUGLMEM double* count = BeamSizeMonitorRecord_getp1_count(record, slot); atomicAdd(count, 1); - /*gpuglmem*/ double * x_sum = BeamSizeMonitorRecord_getp1_x_sum(record, slot); + GPUGLMEM double * x_sum = BeamSizeMonitorRecord_getp1_x_sum(record, slot); atomicAdd(x_sum, x); - /*gpuglmem*/ double * y_sum = BeamSizeMonitorRecord_getp1_y_sum(record, slot); + GPUGLMEM double * y_sum = BeamSizeMonitorRecord_getp1_y_sum(record, slot); atomicAdd(y_sum, y); - /*gpuglmem*/ double * x2_sum = BeamSizeMonitorRecord_getp1_x2_sum(record, slot); + GPUGLMEM double * x2_sum = BeamSizeMonitorRecord_getp1_x2_sum(record, slot); atomicAdd(x2_sum, x*x); - /*gpuglmem*/ double * y2_sum = BeamSizeMonitorRecord_getp1_y2_sum(record, slot); + GPUGLMEM double * y2_sum = BeamSizeMonitorRecord_getp1_y2_sum(record, slot); atomicAdd(y2_sum, y*y); } } diff --git a/xtrack/monitors/beam_size_monitor.py b/xtrack/monitors/beam_size_monitor.py index dff81bbb7..be66c4010 100644 --- a/xtrack/monitors/beam_size_monitor.py +++ b/xtrack/monitors/beam_size_monitor.py @@ -8,11 +8,10 @@ import numpy as np import xobjects as xo - from ..base_element import BeamElement from ..beam_elements import Marker from ..internal_record import RecordIndex -from ..general import _pkg_root + class BeamSizeMonitorRecord(xo.Struct): count = xo.Float64[:] @@ -40,8 +39,7 @@ class BeamSizeMonitor(BeamElement): properties = [field.name for field in BeamSizeMonitorRecord._fields] _extra_c_sources = [ - _pkg_root.joinpath('headers/atomicadd.h'), - _pkg_root.joinpath('monitors/beam_size_monitor.h') + '#include "xtrack/monitors/beam_size_monitor.h"', ] def __init__(self, *, particle_id_range=None, particle_id_start=None, num_particles=None, diff --git a/xtrack/monitors/particles_monitor.h b/xtrack/monitors/particles_monitor.h index d10faba96..5c889def4 100644 --- a/xtrack/monitors/particles_monitor.h +++ b/xtrack/monitors/particles_monitor.h @@ -48,8 +48,11 @@ void ParticlesMonitor_track_local_particle(ParticlesMonitorData el, } else if (n_repetitions > 1){ if (at_turn < start_at_turn){ - return; //only_for_context cuda opencl - break; //only_for_context cpu_serial cpu_openmp + #ifdef XO_CONTEXT_CPU + break; + #else + return; + #endif } int64_t const i_frame = (at_turn - start_at_turn) / repetition_period; if (i_frame < n_repetitions diff --git a/xtrack/multiline_legacy/multiline_legacy.py b/xtrack/multiline_legacy/multiline_legacy.py index 8f0f7a9bb..987dfd12c 100644 --- a/xtrack/multiline_legacy/multiline_legacy.py +++ b/xtrack/multiline_legacy/multiline_legacy.py @@ -479,7 +479,7 @@ def install_beambeam_interactions(self, clockwise_line, anticlockwise_line, # Trackers need to be invalidated to add elements for nn, ll in self.lines.items(): - ll.unfreeze() + ll.discard_tracker() if clockwise_line is not None and anticlockwise_line is not None: circumference_cw = self.lines[clockwise_line].get_length() diff --git a/xtrack/multisetter/multisetter.h b/xtrack/multisetter/multisetter.h new file mode 100644 index 000000000..ce15d8bf0 --- /dev/null +++ b/xtrack/multisetter/multisetter.h @@ -0,0 +1,76 @@ +// copyright ############################### // +// This file is part of the Xtrack Package. // +// Copyright (c) CERN, 2025. // +// ######################################### // + +#ifndef XTRACK_MULTISETTER_H +#define XTRACK_MULTISETTER_H + +#include "xobjects/headers/common.h" + + +GPUKERN +void get_values_at_offsets_float64( + MultiSetterData data, + GPUGLMEM int8_t* buffer, + GPUGLMEM double* out){ + + int64_t num_offsets = MultiSetterData_len_offsets(data); + + VECTORIZE_OVER(ii, num_offsets); + int64_t offs = MultiSetterData_get_offsets(data, ii); + + double val = *((GPUGLMEM double*)(buffer + offs)); + out[ii] = val; + END_VECTORIZE; +} + +GPUKERN +void get_values_at_offsets_int64( + MultiSetterData data, + GPUGLMEM int8_t* buffer, + GPUGLMEM int64_t* out){ + + int64_t num_offsets = MultiSetterData_len_offsets(data); + + VECTORIZE_OVER(ii, num_offsets); + int64_t offs = MultiSetterData_get_offsets(data, ii); + + int64_t val = *((GPUGLMEM int64_t*)(buffer + offs)); + out[ii] = val; + END_VECTORIZE; +} + +GPUKERN +void set_values_at_offsets_float64( + MultiSetterData data, + GPUGLMEM int8_t* buffer, + GPUGLMEM double* input){ + + int64_t num_offsets = MultiSetterData_len_offsets(data); + + VECTORIZE_OVER(ii, num_offsets); + int64_t offs = MultiSetterData_get_offsets(data, ii); + + double val = input[ii]; + *((GPUGLMEM double*)(buffer + offs)) = val; + END_VECTORIZE; +} + +GPUKERN +void set_values_at_offsets_int64( + MultiSetterData data, + GPUGLMEM int8_t* buffer, + GPUGLMEM int64_t* input){ + + int64_t num_offsets = MultiSetterData_len_offsets(data); + + VECTORIZE_OVER(ii, num_offsets); + int64_t offs = MultiSetterData_get_offsets(data, ii); + + int64_t val = input[ii]; + *((GPUGLMEM int64_t*)(buffer + offs)) = val; + END_VECTORIZE; +} + +#endif /* XTRACK_MULTISETTER_H */ \ No newline at end of file diff --git a/xtrack/multisetter/multisetter.py b/xtrack/multisetter/multisetter.py index e963eece6..b3cdb1b1c 100644 --- a/xtrack/multisetter/multisetter.py +++ b/xtrack/multisetter/multisetter.py @@ -3,75 +3,6 @@ import xobjects as xo import xtrack as xt -source = """ - -/*gpukern*/ -void get_values_at_offsets_float64( - MultiSetterData data, - /*gpuglmem*/ int8_t* buffer, - /*gpuglmem*/ double* out){ - - int64_t num_offsets = MultiSetterData_len_offsets(data); - - for (int64_t ii = 0; ii < num_offsets; ii++) { //vectorize_over ii num_offsets - int64_t offs = MultiSetterData_get_offsets(data, ii); - - double val = *((/*gpuglmem*/ double*)(buffer + offs)); - out[ii] = val; - } //end_vectorize -} - -/*gpukern*/ -void get_values_at_offsets_int64( - MultiSetterData data, - /*gpuglmem*/ int8_t* buffer, - /*gpuglmem*/ int64_t* out){ - - int64_t num_offsets = MultiSetterData_len_offsets(data); - - for (int64_t ii = 0; ii < num_offsets; ii++) { //vectorize_over ii num_offsets - int64_t offs = MultiSetterData_get_offsets(data, ii); - - int64_t val = *((/*gpuglmem*/ int64_t*)(buffer + offs)); - out[ii] = val; - } //end_vectorize -} - -/*gpukern*/ -void set_values_at_offsets_float64( - MultiSetterData data, - /*gpuglmem*/ int8_t* buffer, - /*gpuglmem*/ double* input){ - - int64_t num_offsets = MultiSetterData_len_offsets(data); - - for (int64_t ii = 0; ii < num_offsets; ii++) { //vectorize_over ii num_offsets - int64_t offs = MultiSetterData_get_offsets(data, ii); - - double val = input[ii]; - *((/*gpuglmem*/ double*)(buffer + offs)) = val; - } //end_vectorize -} - -/*gpukern*/ -void set_values_at_offsets_int64( - MultiSetterData data, - /*gpuglmem*/ int8_t* buffer, - /*gpuglmem*/ int64_t* input){ - - int64_t num_offsets = MultiSetterData_len_offsets(data); - - for (int64_t ii = 0; ii < num_offsets; ii++) { //vectorize_over ii num_offsets - int64_t offs = MultiSetterData_get_offsets(data, ii); - - int64_t val = input[ii]; - *((/*gpuglmem*/ int64_t*)(buffer + offs)) = val; - } //end_vectorize -} - - - -""" class MultiSetter(xo.HybridClass): _xofields = { @@ -79,7 +10,7 @@ class MultiSetter(xo.HybridClass): } _extra_c_sources = [ - source, + '#include "xtrack/multisetter/multisetter.h"', ] _kernels = { diff --git a/xtrack/particles/local_particle_custom_api.h b/xtrack/particles/local_particle_custom_api.h new file mode 100644 index 000000000..19b9ce078 --- /dev/null +++ b/xtrack/particles/local_particle_custom_api.h @@ -0,0 +1,289 @@ +// copyright ############################### // +// This file is part of the Xtrack Package. // +// Copyright (c) CERN, 2025. // +// ######################################### // + +#ifndef XTRACK_LOCAL_PARTICLE_CUSTOM_API_H +#define XTRACK_LOCAL_PARTICLE_CUSTOM_API_H + +#include "xtrack/headers/track.h" + + +GPUFUN +double LocalParticle_get_energy0(LocalParticle* part) { + double const p0c = LocalParticle_get_p0c(part); + double const m0 = LocalParticle_get_mass0(part); + return sqrt(p0c * p0c + m0 * m0); +} + + +GPUFUN +void LocalParticle_update_ptau(LocalParticle* part, double new_ptau_value) { + double const beta0 = LocalParticle_get_beta0(part); + double const ptau = new_ptau_value; + double const irpp = sqrt(ptau * ptau + 2 * ptau / beta0 + 1); + double const new_rpp = 1. / irpp; + double const new_rvv = irpp / (1 + beta0 * ptau); + + LocalParticle_set_delta(part, irpp - 1); + LocalParticle_set_rvv(part, new_rvv); + LocalParticle_set_ptau(part, ptau); + LocalParticle_set_rpp(part, new_rpp); +} + + +GPUFUN +void LocalParticle_update_delta(LocalParticle* part, double new_delta_value) { + double const beta0 = LocalParticle_get_beta0(part); + double const delta_beta0 = new_delta_value * beta0; + double const ptau_beta0 = sqrt(delta_beta0 * delta_beta0 + 2 * delta_beta0 * beta0 + 1) - 1; + double const one_plus_delta = 1 + new_delta_value; + double const rvv = (one_plus_delta) / (1 + ptau_beta0); + double const rpp = 1 / one_plus_delta; + double const ptau = ptau_beta0 / beta0; + + LocalParticle_set_delta(part, new_delta_value); + LocalParticle_set_rvv(part, rvv); + LocalParticle_set_rpp(part, rpp); + LocalParticle_set_ptau(part, ptau); +} + + +GPUFUN +double LocalParticle_get_pzeta(LocalParticle* part) { + double const ptau = LocalParticle_get_ptau(part); + double const beta0 = LocalParticle_get_beta0(part); + return ptau / beta0; +} + + +GPUFUN +void LocalParticle_update_pzeta(LocalParticle* part, double new_pzeta_value) { + double const beta0 = LocalParticle_get_beta0(part); + LocalParticle_update_ptau(part, beta0 * new_pzeta_value); +} + + +GPUFUN +void increment_at_element(LocalParticle* part0, int64_t const increment) { + START_PER_PARTICLE_BLOCK(part0, part); + LocalParticle_add_to_at_element(part, increment); + END_PER_PARTICLE_BLOCK; +} + + +GPUFUN +void increment_at_turn(LocalParticle* part0, int flag_reset_s){ + START_PER_PARTICLE_BLOCK(part0, part); + LocalParticle_add_to_at_turn(part, 1); + LocalParticle_set_at_element(part, 0); + if (flag_reset_s > 0) { + LocalParticle_set_s(part, 0.); + } + END_PER_PARTICLE_BLOCK; +} + + +GPUFUN +void increment_at_turn_backtrack( + LocalParticle* part0, + int flag_reset_s, + double const line_length, + int64_t const num_elements +) { + START_PER_PARTICLE_BLOCK(part0, part); + LocalParticle_add_to_at_turn(part, -1); + LocalParticle_set_at_element(part, num_elements); + if (flag_reset_s > 0) { + LocalParticle_set_s(part, line_length); + } + END_PER_PARTICLE_BLOCK; +} + + +/* check_is_active has different implementation on CPU and GPU */ +#if defined(XO_CONTEXT_CPU_SERIAL) + + GPUFUN + int64_t check_is_active(LocalParticle* part) { + int64_t ipart = 0; + while (ipart < part->_num_active_particles){ + #ifdef XSUITE_RESTORE_LOSS + ipart++; + #else + if (part->state[ipart] < 1) { + LocalParticle_exchange(part, ipart, part->_num_active_particles - 1); + part->_num_active_particles--; + part->_num_lost_particles++; + } else { + ipart++; + } + #endif + } + + if (part->_num_active_particles == 0){ + return 0; //All particles lost + } else { + return 1; //Some stable particles are still present + } + } + +#elif defined(XO_CONTEXT_CPU_OPENMP) + + GPUFUN + int64_t check_is_active(LocalParticle* part) { + #ifndef SKIP_SWAPS + int64_t ipart = part->ipart; + int64_t endpart = part->endpart; + int64_t left = ipart; + int64_t right = endpart - 1; + int64_t swap_made = 0; + int64_t has_alive = 0; + + if (left == right) + return part->state[left] > 0; + + while (left < right) { + if (part->state[left] > 0) { + left++; + has_alive = 1; + } + else if (part->state[right] <= 0) + right--; + else { + LocalParticle_exchange(part, left, right); + left++; + right--; + swap_made = 1; + } + } + return swap_made || has_alive; + #else + return 1; + #endif + } + + + GPUFUN + void count_reorganized_particles(LocalParticle* part) { + int64_t num_active = 0; + int64_t num_lost = 0; + + for (int64_t i = part->ipart; i < part->endpart; i++) { + if (part->state[i] <= -999999999) + break; + else if (part->state[i] > 0) + num_active++; + else + num_lost++; + } + + part->_num_active_particles = num_active; + part->_num_lost_particles = num_lost; + } + +#else // not XO_CONTEXT_CPU_SERIAL and not XO_CONTEXT_CPU_OPENMP + + GPUFUN + int64_t check_is_active(LocalParticle* part) { + return LocalParticle_get_state(part) > 0; + }; + +#endif + + +GPUFUN +void LocalParticle_add_to_energy(LocalParticle* part, double delta_energy, int pz_only ) +{ + double ptau = LocalParticle_get_ptau(part); + double const p0c = LocalParticle_get_p0c(part); + double const charge_ratio = LocalParticle_get_charge_ratio(part); + double const chi = LocalParticle_get_chi(part); + double const mass_ratio = charge_ratio / chi; + + ptau += delta_energy / p0c / mass_ratio; + + double const old_rpp = LocalParticle_get_rpp(part); + + LocalParticle_update_ptau(part, ptau); + + if (!pz_only) { + double const new_rpp = LocalParticle_get_rpp(part); + double const f = old_rpp / new_rpp; + LocalParticle_scale_px(part, f); + LocalParticle_scale_py(part, f); + } +} + + +GPUFUN +void LocalParticle_update_p0c(LocalParticle* part, double new_p0c_value) +{ + double const mass0 = LocalParticle_get_mass0(part); + double const old_p0c = LocalParticle_get_p0c(part); + double const old_delta = LocalParticle_get_delta(part); + double const old_beta0 = LocalParticle_get_beta0(part); + + double const ppc = old_p0c * old_delta + old_p0c; + double const new_delta = (ppc - new_p0c_value) / new_p0c_value; + + double const new_energy0 = sqrt(new_p0c_value * new_p0c_value + mass0 * mass0); + double const new_beta0 = new_p0c_value / new_energy0; + double const new_gamma0 = new_energy0 / mass0; + + LocalParticle_set_p0c(part, new_p0c_value); + LocalParticle_set_gamma0(part, new_gamma0); + LocalParticle_set_beta0(part, new_beta0); + + LocalParticle_update_delta(part, new_delta); + + LocalParticle_scale_px(part, old_p0c / new_p0c_value); + LocalParticle_scale_py(part, old_p0c / new_p0c_value); + + LocalParticle_scale_zeta(part, new_beta0 / old_beta0); + +} + +GPUFUN +void LocalParticle_kill_particle(LocalParticle* part, int64_t kill_state) { + LocalParticle_set_x(part, 1e30); + LocalParticle_set_px(part, 1e30); + LocalParticle_set_y(part, 1e30); + LocalParticle_set_py(part, 1e30); + LocalParticle_set_zeta(part, 1e30); + LocalParticle_update_delta(part, -1); // zero energy + LocalParticle_set_state(part, kill_state); +} + + +#ifdef XTRACK_GLOBAL_XY_LIMIT + + GPUFUN + void global_aperture_check(LocalParticle* part0) + { + if (LocalParticle_check_track_flag(part0, XS_FLAG_IGNORE_GLOBAL_APERTURE)) + { + return; + } + + START_PER_PARTICLE_BLOCK(part0, part); + double const x = LocalParticle_get_x(part); + double const y = LocalParticle_get_y(part); + + int64_t const is_alive = (int64_t)( + (x >= -XTRACK_GLOBAL_XY_LIMIT) && + (x <= XTRACK_GLOBAL_XY_LIMIT) && + (y >= -XTRACK_GLOBAL_XY_LIMIT) && + (y <= XTRACK_GLOBAL_XY_LIMIT) + ); + + // I assume that if I am in the function is because + if (!is_alive) { + LocalParticle_set_state(part, -1); + } + END_PER_PARTICLE_BLOCK; + } + +#endif + +#endif /* XTRACK_LOCAL_PARTICLE_CUSTOM_API_H */ diff --git a/xtrack/particles/particles.py b/xtrack/particles/particles.py index eca5b324a..b7a5438b9 100644 --- a/xtrack/particles/particles.py +++ b/xtrack/particles/particles.py @@ -23,69 +23,398 @@ LAST_INVALID_STATE = -999999999 +size_vars = ( + (xo.Int64, '_capacity'), + (xo.Int64, '_num_active_particles'), + (xo.Int64, '_num_lost_particles'), + (xo.Int64, 'start_tracking_at_element'), +) +# Capacity is always kept up to date +# the other two are placeholders to be used if needed +# i.e. on ContextCpu + +scalar_vars = ( + (xo.Float64, 'q0'), + (xo.Float64, 'mass0'), + (xo.Float64, 't_sim'), +) + +part_energy_vars = ( + (xo.Float64, 'ptau'), + (xo.Float64, 'delta'), + (xo.Float64, 'rpp'), + (xo.Float64, 'rvv'), +) + +per_particle_vars = ( + ( + (xo.Float64, 'p0c'), + (xo.Float64, 'gamma0'), + (xo.Float64, 'beta0'), + (xo.Float64, 's'), + (xo.Float64, 'zeta'), + (xo.Float64, 'x'), + (xo.Float64, 'y'), + (xo.Float64, 'px'), + (xo.Float64, 'py'), + ) + + part_energy_vars + + ( + (xo.Float64, 'chi'), + (xo.Float64, 'charge_ratio'), + (xo.Float64, 'weight'), + (xo.Float64, 'ax'), + (xo.Float64, 'ay'), + (xo.Float64, 'spin_x'), + (xo.Float64, 'spin_y'), + (xo.Float64, 'spin_z'), + (xo.Float64, 'anomalous_magnetic_moment'), + (xo.Int64, 'pdg_id'), + (xo.Int64, 'particle_id'), + (xo.Int64, 'at_element'), + (xo.Int64, 'at_turn'), + (xo.Int64, 'state'), + (xo.Int64, 'parent_particle_id'), + (xo.UInt32, '_rng_s1'), + (xo.UInt32, '_rng_s2'), + (xo.UInt32, '_rng_s3'), + (xo.UInt32, '_rng_s4') + ) +) -class Particles(xo.HybridClass): - _cname = 'ParticlesData' - size_vars = ( - (xo.Int64, '_capacity'), - (xo.Int64, '_num_active_particles'), - (xo.Int64, '_num_lost_particles'), - (xo.Int64, 'start_tracking_at_element'), +def gen_local_particle_api(): + src_lines = [ + '#include "xobjects/headers/common.h"', + '#include "xtrack/particles/rng_src/base_rng.h"', + '#include "xtrack/particles/rng_src/particles_rng.h"', + ] + for name, mass in mass__dict__.items(): + if name.endswith('_MASS_EV'): + src_lines.append(f'#define {name} {mass}') + + src_lines.append('typedef struct {') + + for tt, vv in size_vars + scalar_vars: + src_lines.append(' ' + tt._c_type + ' ' + vv + ';') + + for tt, vv in per_particle_vars: + src_lines.append(' GPUGLMEM ' + tt._c_type + '* ' + vv + ';') + + src_lines.append(' int64_t ipart;') + src_lines.append(' int64_t endpart;') + src_lines.append(' uint64_t track_flags;') + src_lines.append(' GPUGLMEM int8_t* io_buffer;') + src_lines.append('} LocalParticle;') + src_typedef = '\n'.join(src_lines) + + # Get io buffer + src_lines = [] + src_lines.append( + ''' + GPUFUN + GPUGLMEM int8_t* LocalParticle_get_io_buffer(LocalParticle* part){ + return part->io_buffer; + } + + ''' ) - # Capacity is always kept up to date - # the other two are placeholders to be used if needed - # i.e. on ContextCpu - - scalar_vars = ( - (xo.Float64, 'q0'), - (xo.Float64, 'mass0'), - (xo.Float64, 't_sim'), + + # Get track flag + src_lines.append( + ''' + GPUFUN + uint64_t LocalParticle_check_track_flag(LocalParticle* part, uint8_t index){ + return (part->track_flags >> index) & 1; + } + ''' ) - part_energy_vars = ( - (xo.Float64, 'ptau'), - (xo.Float64, 'delta'), - (xo.Float64, 'rpp'), - (xo.Float64, 'rvv'), + # Particles_to_LocalParticle + src_lines.append( + ''' + GPUFUN + void Particles_to_LocalParticle(ParticlesData source, + LocalParticle* dest, + int64_t id, + int64_t eid){''' + ) + for _, vv in size_vars + scalar_vars: + src_lines.append( + f' dest->{vv} = ParticlesData_get_' + vv + '(source);' + ) + + for _, vv in per_particle_vars: + src_lines.append( + f' dest->{vv} = ParticlesData_getp1_' + vv + '(source, 0);' + ) + + src_lines.append(' dest->ipart = id;') + src_lines.append(' dest->endpart = eid;') + src_lines.append('}') + src_particles_to_local = '\n'.join(src_lines) + + # LocalParticle_to_Particles + src_lines = [] + src_lines.append( + ''' + GPUFUN + void LocalParticle_to_Particles(LocalParticle* source, + ParticlesData dest, + int64_t id, + int64_t set_scalar){''' ) + src_lines.append('if (set_scalar){') + for _, vv in size_vars + scalar_vars: + src_lines.append( + ' ParticlesData_set_' + vv + '(dest,' + f' LocalParticle_get_{vv}(source));' + ) + src_lines.append('}') + + for _, vv in per_particle_vars: + src_lines.append( + ' ParticlesData_set_' + vv + '(dest, id, ' + f' LocalParticle_get_{vv}(source));' + ) + src_lines.append('}') + src_local_to_particles = '\n'.join(src_lines) + + # Adders + src_lines = [] + for tt, vv in per_particle_vars: + src_lines.append( + ''' + GPUFUN + void LocalParticle_add_to_''' + vv + f'(LocalParticle* part, {tt._c_type} value)' + + '{' + ) + src_lines.append(f'#ifndef FREEZE_VAR_{vv}') + src_lines.append(f' part->{vv}[part->ipart] += value;') + src_lines.append('#endif') + src_lines.append('}\n') + src_adders = '\n'.join(src_lines) + + # Scalers + src_lines = [] + for tt, vv in per_particle_vars: + src_lines.append( + ''' + GPUFUN + void LocalParticle_scale_''' + vv + f'(LocalParticle* part, {tt._c_type} value)' + + '{' + ) + src_lines.append(f'#ifndef FREEZE_VAR_{vv}') + src_lines.append(f' part->{vv}[part->ipart] *= value;') + src_lines.append('#endif') + src_lines.append('}\n') + src_scalers = '\n'.join(src_lines) + + # Setters + src_lines = [] + for tt, vv in per_particle_vars: + src_lines.append( + ''' + GPUFUN + void LocalParticle_set_''' + vv + f'(LocalParticle* part, {tt._c_type} value)' + + '{' + ) + src_lines.append(f'#ifndef FREEZE_VAR_{vv}') + src_lines.append(f' part->{vv}[part->ipart] = value;') + src_lines.append('#endif') + src_lines.append('}') + src_setters = '\n'.join(src_lines) + + # Getters + src_lines = [] - per_particle_vars = ( - ( - (xo.Float64, 'p0c'), - (xo.Float64, 'gamma0'), - (xo.Float64, 'beta0'), - (xo.Float64, 's'), - (xo.Float64, 'zeta'), - (xo.Float64, 'x'), - (xo.Float64, 'y'), - (xo.Float64, 'px'), - (xo.Float64, 'py'), + for tt, vv in size_vars + scalar_vars: + src_lines.append('GPUFUN') + src_lines.append( + f'{tt._c_type} LocalParticle_get_' + vv + + '(LocalParticle* part)' + + '{' + ) + src_lines.append(f' return part->{vv};') + src_lines.append('}') + + for tt, vv in per_particle_vars: + src_lines.append('GPUFUN') + src_lines.append( + f'{tt._c_type} LocalParticle_get_' + vv + + '(LocalParticle* part)' + + '{' + ) + src_lines.append(f' return part->{vv}[part->ipart];') + src_lines.append('}') + + src_getters = '\n'.join(src_lines) + + # Angles + src_angles_lines = [] + for exact in ['', 'exact_']: + # xp (py as transverse) and vice versa + for xx, yy in [['x', 'y'], ['y', 'x']]: + # Getter + src_angles_lines.append('GPUFUN') + src_angles_lines.append( + f'double LocalParticle_get_{exact}{xx}p(LocalParticle* part){{' + ) + src_angles_lines.append( + f' double const p{xx} = LocalParticle_get_p{xx}(part);' + ) + if exact == 'exact_': + src_angles_lines.append( + f' double const p{yy} = LocalParticle_get_p{yy}(part);' + ) + src_angles_lines.append( + ' double const one_plus_delta = 1. + LocalParticle_get_delta(part);' + ) + src_angles_lines.append( + ' double const rpp = 1./sqrt(one_plus_delta*one_plus_delta - px*px - py*py);' + ) + else: + src_angles_lines.append( + ' double const rpp = LocalParticle_get_rpp(part);' + ) + src_angles_lines.append( + ' // INFO: this is not the angle, but sin(angle)' + ) + src_angles_lines.append(f' return p{xx}*rpp;') + src_angles_lines.append('}') + src_angles_lines.append('') + + for xx, yy in [['x', 'y'], ['y', 'x']]: + # Setter + src_angles_lines.append('GPUFUN') + src_angles_lines.append( + f'void LocalParticle_set_{exact}{xx}p(LocalParticle* part, double {xx}p){{' + ) + src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}') + src_angles_lines.append( + ' double rpp = LocalParticle_get_rpp(part);' + ) + if exact == 'exact_': + src_angles_lines.append( + f' // Careful! If {yy}p also changes, use LocalParticle_set_{exact}xp_yp!' + ) + src_angles_lines.append( + f' double const {yy}p = LocalParticle_get_{exact}{yy}p(part);' + ) + src_angles_lines.append(' rpp *= sqrt(1 + xp*xp + yp*yp);') + src_angles_lines.append( + f' // INFO: {xx}p is not the angle, but sin(angle)' + ) + src_angles_lines.append( + f' LocalParticle_set_p{xx}(part, {xx}p/rpp);' + ) + src_angles_lines.append('#endif') + src_angles_lines.append('}') + src_angles_lines.append('') + + for xx, yy in [['x', 'y'], ['y', 'x']]: + # Adder + src_angles_lines.append('GPUFUN') + src_angles_lines.append( + f'void LocalParticle_add_to_{exact}{xx}p(LocalParticle* part, double {xx}p){{' + ) + src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}') + src_angles_lines.append( + f' LocalParticle_set_{exact}{xx}p(part, ' + + f'LocalParticle_get_{exact}{xx}p(part) + {xx}p);' + ) + src_angles_lines.append('#endif') + src_angles_lines.append('}') + src_angles_lines.append('') + # Scaler + src_angles_lines.append('GPUFUN') + src_angles_lines.append( + f'void LocalParticle_scale_{exact}{xx}p(LocalParticle* part, double value){{' + ) + src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}') + src_angles_lines.append( + f' LocalParticle_set_{exact}{xx}p(part, ' + + f'LocalParticle_get_{exact}{xx}p(part) * value);' + ) + src_angles_lines.append('#endif') + src_angles_lines.append('}') + src_angles_lines.append('') + # Double setter, adder, scaler + src_angles_lines.append('GPUFUN') + src_angles_lines.append( + f'void LocalParticle_set_{exact}xp_yp(LocalParticle* part, double xp, double yp){{' + ) + src_angles_lines.append(' double rpp = LocalParticle_get_rpp(part);') + if exact == 'exact_': + src_angles_lines.append(' rpp *= sqrt(1 + xp*xp + yp*yp);') + for xx in ['x', 'y']: + src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}') + src_angles_lines.append( + f' LocalParticle_set_p{xx}(part, {xx}p/rpp);' + ) + src_angles_lines.append('#endif') + src_angles_lines.append('}') + src_angles_lines.append('') + src_angles_lines.append('GPUFUN') + src_angles_lines.append( + f'void LocalParticle_add_to_{exact}xp_yp(LocalParticle* part, double xp, double yp){{' + ) + src_angles_lines.append( + f' LocalParticle_set_{exact}xp_yp(part, ' + + f'LocalParticle_get_{exact}xp(part) + xp, ' + + f'LocalParticle_get_{exact}yp(part) + yp);' + ) + src_angles_lines.append('}') + src_angles_lines.append('') + src_angles_lines.append('GPUFUN') + src_angles_lines.append( + f'void LocalParticle_scale_{exact}xp_yp(LocalParticle* part, double value_x, double value_y){{' ) - + part_energy_vars + - ( - (xo.Float64, 'chi'), - (xo.Float64, 'charge_ratio'), - (xo.Float64, 'weight'), - (xo.Float64, 'ax'), - (xo.Float64, 'ay'), - (xo.Float64, 'spin_x'), - (xo.Float64, 'spin_y'), - (xo.Float64, 'spin_z'), - (xo.Float64, 'anomalous_magnetic_moment'), - (xo.Int64, 'pdg_id'), - (xo.Int64, 'particle_id'), - (xo.Int64, 'at_element'), - (xo.Int64, 'at_turn'), - (xo.Int64, 'state'), - (xo.Int64, 'parent_particle_id'), - (xo.UInt32, '_rng_s1'), - (xo.UInt32, '_rng_s2'), - (xo.UInt32, '_rng_s3'), - (xo.UInt32, '_rng_s4') + src_angles_lines.append( + f' LocalParticle_set_{exact}xp_yp(part, ' + + f'LocalParticle_get_{exact}xp(part) * value_x, ' + + f'LocalParticle_get_{exact}yp(part) * value_y);' ) + src_angles_lines.append('}') + src_angles = '\n'.join(src_angles_lines) + + # Particle exchangers + src_exchange = ''' + GPUFUN + void LocalParticle_exchange(LocalParticle* part, int64_t i1, int64_t i2){ + ''' + for tt, vv in per_particle_vars: + src_exchange += '\n'.join( + [ + '\n {', + f' {tt._c_type} temp = part->{vv}[i2];', + f' part->{vv}[i2] = part->{vv}[i1];', + f' part->{vv}[i1] = temp;', + ' }'] + ) + src_exchange += '}\n' + + source = '\n\n'.join( + [src_typedef, src_adders, src_getters, + src_setters, src_scalers, src_exchange, + src_particles_to_local, src_local_to_particles, + src_angles] ) + source += """ + #include "xtrack/particles/local_particle_custom_api.h" + """ + return source + + +class Particles(xo.HybridClass): + _cname = 'ParticlesData' + + size_vars = size_vars + scalar_vars = scalar_vars + part_energy_vars = part_energy_vars + per_particle_vars = per_particle_vars + _xofields = { **{nn: tt for tt, nn in size_vars + scalar_vars}, **{nn: tt[:] for tt, nn in per_particle_vars}, @@ -95,9 +424,7 @@ class Particles(xo.HybridClass): if (not vv[1].startswith('_') and vv[1] != 't_sim')] + ['t_sim'] _extra_c_sources = [ - Path(__file__).parent.joinpath('rng_src', 'base_rng.h'), - Path(__file__).parent.joinpath('rng_src', 'particles_rng.h'), - '\n /*placeholder_for_local_particle_src*/ \n' + gen_local_particle_api() ] _rename = { @@ -1471,551 +1798,6 @@ def add_to_energy(self, delta_energy): def set_particle(self, index, set_scalar_vars=False, **kwargs): raise NotImplementedError('This functionality has been removed') - @classmethod - def gen_local_particle_api(cls, mode='no_local_copy'): - if mode != 'no_local_copy': - raise NotImplementedError - - src_lines = [''] - for name, mass in mass__dict__.items(): - if name.endswith('_MASS_EV'): - src_lines.append(f'#define {name} {mass}') - - src_lines.append('typedef struct {') - - for tt, vv in cls.size_vars + cls.scalar_vars: - src_lines.append(' ' + tt._c_type + ' ' + vv + ';') - - for tt, vv in cls.per_particle_vars: - src_lines.append(' /*gpuglmem*/ ' + tt._c_type + '* ' + vv + ';') - - src_lines.append(' int64_t ipart;') - src_lines.append(' int64_t endpart;') - src_lines.append(' uint64_t track_flags;') - src_lines.append(' /*gpuglmem*/ int8_t* io_buffer;') - src_lines.append('} LocalParticle;') - src_typedef = '\n'.join(src_lines) - - # Get io buffer - src_lines = [] - src_lines.append(''' - /*gpufun*/ - /*gpuglmem*/ int8_t* LocalParticle_get_io_buffer(LocalParticle* part){ - return part->io_buffer; - } - - ''') - - # Get track flag - src_lines.append(''' - /*gpufun*/ - uint64_t LocalParticle_check_track_flag(LocalParticle* part, uint8_t index){ - return (part->track_flags >> index) & 1; - } - ''') - - # Particles_to_LocalParticle - src_lines.append(''' - /*gpufun*/ - void Particles_to_LocalParticle(ParticlesData source, - LocalParticle* dest, - int64_t id, - int64_t eid){''') - for _, vv in cls.size_vars + cls.scalar_vars: - src_lines.append( - f' dest->{vv} = ParticlesData_get_' + vv + '(source);') - - for _, vv in cls.per_particle_vars: - src_lines.append( - f' dest->{vv} = ParticlesData_getp1_' + vv + '(source, 0);') - - src_lines.append(' dest->ipart = id;') - src_lines.append(' dest->endpart = eid;') - src_lines.append('}') - src_particles_to_local = '\n'.join(src_lines) - - # LocalParticle_to_Particles - src_lines = [] - src_lines.append(''' - /*gpufun*/ - void LocalParticle_to_Particles(LocalParticle* source, - ParticlesData dest, - int64_t id, - int64_t set_scalar){''') - src_lines.append('if (set_scalar){') - for _, vv in cls.size_vars + cls.scalar_vars: - src_lines.append( - ' ParticlesData_set_' + vv + '(dest,' - f' LocalParticle_get_{vv}(source));') - src_lines.append('}') - - for _, vv in cls.per_particle_vars: - src_lines.append( - ' ParticlesData_set_' + vv + '(dest, id, ' - f' LocalParticle_get_{vv}(source));') - src_lines.append('}') - src_local_to_particles = '\n'.join(src_lines) - - # Adders - src_lines = [] - for tt, vv in cls.per_particle_vars: - src_lines.append(''' - /*gpufun*/ - void LocalParticle_add_to_''' + vv + f'(LocalParticle* part, {tt._c_type} value)' - + '{') - src_lines.append(f'#ifndef FREEZE_VAR_{vv}') - src_lines.append(f' part->{vv}[part->ipart] += value;') - src_lines.append('#endif') - src_lines.append('}\n') - src_adders = '\n'.join(src_lines) - - # Scalers - src_lines = [] - for tt, vv in cls.per_particle_vars: - src_lines.append(''' - /*gpufun*/ - void LocalParticle_scale_''' + vv + f'(LocalParticle* part, {tt._c_type} value)' - + '{') - src_lines.append(f'#ifndef FREEZE_VAR_{vv}') - src_lines.append(f' part->{vv}[part->ipart] *= value;') - src_lines.append('#endif') - src_lines.append('}\n') - src_scalers = '\n'.join(src_lines) - - # Setters - src_lines = [] - for tt, vv in cls.per_particle_vars: - src_lines.append(''' - /*gpufun*/ - void LocalParticle_set_''' + vv + f'(LocalParticle* part, {tt._c_type} value)' - + '{') - src_lines.append(f'#ifndef FREEZE_VAR_{vv}') - src_lines.append(f' part->{vv}[part->ipart] = value;') - src_lines.append('#endif') - src_lines.append('}') - src_setters = '\n'.join(src_lines) - - # Getters - src_lines = [] - - for tt, vv in cls.size_vars + cls.scalar_vars: - src_lines.append('/*gpufun*/') - src_lines.append(f'{tt._c_type} LocalParticle_get_' + vv - + '(LocalParticle* part)' - + '{') - src_lines.append(f' return part->{vv};') - src_lines.append('}') - - for tt, vv in cls.per_particle_vars: - src_lines.append('/*gpufun*/') - src_lines.append(f'{tt._c_type} LocalParticle_get_' + vv - + '(LocalParticle* part)' - + '{') - src_lines.append(f' return part->{vv}[part->ipart];') - src_lines.append('}') - - src_getters = '\n'.join(src_lines) - - # Angles - src_angles_lines = [] - for exact in ['', 'exact_']: - # xp (py as transverse) and vice versa - for xx, yy in [['x', 'y'], ['y', 'x']]: - # Getter - src_angles_lines.append('/*gpufun*/') - src_angles_lines.append(f'double LocalParticle_get_{exact}{xx}p(LocalParticle* part){{') - src_angles_lines.append(f' double const p{xx} = LocalParticle_get_p{xx}(part);') - if exact == 'exact_': - src_angles_lines.append(f' double const p{yy} = LocalParticle_get_p{yy}(part);') - src_angles_lines.append(' double const one_plus_delta = 1. + LocalParticle_get_delta(part);') - src_angles_lines.append( - ' double const rpp = 1./sqrt(one_plus_delta*one_plus_delta - px*px - py*py);') - else: - src_angles_lines.append(' double const rpp = LocalParticle_get_rpp(part);') - src_angles_lines.append(' // INFO: this is not the angle, but sin(angle)') - src_angles_lines.append(f' return p{xx}*rpp;') - src_angles_lines.append('}') - src_angles_lines.append('') - - for xx, yy in [['x', 'y'], ['y', 'x']]: - # Setter - src_angles_lines.append('/*gpufun*/') - src_angles_lines.append(f'void LocalParticle_set_{exact}{xx}p(LocalParticle* part, double {xx}p){{') - src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}') - src_angles_lines.append(' double rpp = LocalParticle_get_rpp(part);') - if exact == 'exact_': - src_angles_lines.append( - f' // Careful! If {yy}p also changes, use LocalParticle_set_{exact}xp_yp!') - src_angles_lines.append(f' double const {yy}p = LocalParticle_get_{exact}{yy}p(part);') - src_angles_lines.append(' rpp *= sqrt(1 + xp*xp + yp*yp);') - src_angles_lines.append(f' // INFO: {xx}p is not the angle, but sin(angle)') - src_angles_lines.append(f' LocalParticle_set_p{xx}(part, {xx}p/rpp);') - src_angles_lines.append('#endif') - src_angles_lines.append('}') - src_angles_lines.append('') - - for xx, yy in [['x', 'y'], ['y', 'x']]: - # Adder - src_angles_lines.append('/*gpufun*/') - src_angles_lines.append(f'void LocalParticle_add_to_{exact}{xx}p(LocalParticle* part, double {xx}p){{') - src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}') - src_angles_lines.append(f' LocalParticle_set_{exact}{xx}p(part, ' - + f'LocalParticle_get_{exact}{xx}p(part) + {xx}p);') - src_angles_lines.append('#endif') - src_angles_lines.append('}') - src_angles_lines.append('') - # Scaler - src_angles_lines.append('/*gpufun*/') - src_angles_lines.append(f'void LocalParticle_scale_{exact}{xx}p(LocalParticle* part, double value){{') - src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}') - src_angles_lines.append(f' LocalParticle_set_{exact}{xx}p(part, ' - + f'LocalParticle_get_{exact}{xx}p(part) * value);') - src_angles_lines.append('#endif') - src_angles_lines.append('}') - src_angles_lines.append('') - # Double setter, adder, scaler - src_angles_lines.append('/*gpufun*/') - src_angles_lines.append(f'void LocalParticle_set_{exact}xp_yp(LocalParticle* part, double xp, double yp){{') - src_angles_lines.append(' double rpp = LocalParticle_get_rpp(part);') - if exact == 'exact_': - src_angles_lines.append(' rpp *= sqrt(1 + xp*xp + yp*yp);') - for xx in ['x', 'y']: - src_angles_lines.append(f'#ifndef FREEZE_VAR_p{xx}') - src_angles_lines.append(f' LocalParticle_set_p{xx}(part, {xx}p/rpp);') - src_angles_lines.append('#endif') - src_angles_lines.append('}') - src_angles_lines.append('') - src_angles_lines.append('/*gpufun*/') - src_angles_lines.append( - f'void LocalParticle_add_to_{exact}xp_yp(LocalParticle* part, double xp, double yp){{') - src_angles_lines.append(f' LocalParticle_set_{exact}xp_yp(part, ' - + f'LocalParticle_get_{exact}xp(part) + xp, ' - + f'LocalParticle_get_{exact}yp(part) + yp);') - src_angles_lines.append('}') - src_angles_lines.append('') - src_angles_lines.append('/*gpufun*/') - src_angles_lines.append( - f'void LocalParticle_scale_{exact}xp_yp(LocalParticle* part, double value_x, double value_y){{') - src_angles_lines.append(f' LocalParticle_set_{exact}xp_yp(part, ' - + f'LocalParticle_get_{exact}xp(part) * value_x, ' - + f'LocalParticle_get_{exact}yp(part) * value_y);') - src_angles_lines.append('}') - src_angles = '\n'.join(src_angles_lines) - - # Particle exchangers - src_exchange = ''' - /*gpufun*/ - void LocalParticle_exchange(LocalParticle* part, int64_t i1, int64_t i2){ - ''' - for tt, vv in cls.per_particle_vars: - src_exchange += '\n'.join([ - '\n {', - f' {tt._c_type} temp = part->{vv}[i2];', - f' part->{vv}[i2] = part->{vv}[i1];', - f' part->{vv}[i1] = temp;', - ' }']) - src_exchange += '}\n' - - custom_source = ''' - /*gpufun*/ - double LocalParticle_get_energy0(LocalParticle* part){ - - double const p0c = LocalParticle_get_p0c(part); - double const m0 = LocalParticle_get_mass0(part); - - return sqrt( p0c * p0c + m0 * m0 ); - } - - /*gpufun*/ - void LocalParticle_update_ptau(LocalParticle* part, double new_ptau_value){ - - double const beta0 = LocalParticle_get_beta0(part); - - double const ptau = new_ptau_value; - - double const irpp = sqrt(ptau*ptau + 2*ptau/beta0 +1); - - double const new_rpp = 1./irpp; - LocalParticle_set_delta(part, irpp - 1.); - - double const new_rvv = irpp/(1 + beta0*ptau); - LocalParticle_set_rvv(part, new_rvv); - LocalParticle_set_ptau(part, ptau); - - LocalParticle_set_rpp(part, new_rpp ); - } - - /*gpufun*/ - void LocalParticle_update_delta(LocalParticle* part, double new_delta_value){ - double const beta0 = LocalParticle_get_beta0(part); - double const delta_beta0 = new_delta_value * beta0; - double const ptau_beta0 = sqrt( delta_beta0 * delta_beta0 + - 2. * delta_beta0 * beta0 + 1. ) - 1.; - - double const one_plus_delta = 1. + new_delta_value; - double const rvv = ( one_plus_delta ) / ( 1. + ptau_beta0 ); - double const rpp = 1. / one_plus_delta; - double const ptau = ptau_beta0 / beta0; - - LocalParticle_set_delta(part, new_delta_value); - - LocalParticle_set_rvv(part, rvv ); - LocalParticle_set_rpp(part, rpp ); - LocalParticle_set_ptau(part, ptau ); - - } - - /*gpufun*/ - double LocalParticle_get_pzeta(LocalParticle* part){ - - double const ptau = LocalParticle_get_ptau(part); - double const beta0 = LocalParticle_get_beta0(part); - - return ptau/beta0; - - } - - /*gpufun*/ - void LocalParticle_update_pzeta(LocalParticle* part, double new_pzeta_value){ - - double const beta0 = LocalParticle_get_beta0(part); - LocalParticle_update_ptau(part, beta0*new_pzeta_value); - - } - - /*gpufun*/ - void increment_at_element(LocalParticle* part0, int64_t const increment){ - - //start_per_particle_block (part0->part) - LocalParticle_add_to_at_element(part, increment); - //end_per_particle_block - - } - - /*gpufun*/ - void increment_at_turn(LocalParticle* part0, int flag_reset_s){ - - //start_per_particle_block (part0->part) - LocalParticle_add_to_at_turn(part, 1); - LocalParticle_set_at_element(part, 0); - if (flag_reset_s>0){ - LocalParticle_set_s(part, 0.); - } - //end_per_particle_block - } - - /*gpufun*/ - void increment_at_turn_backtrack(LocalParticle* part0, int flag_reset_s, - double const line_length, - int64_t const num_elements){ - - //start_per_particle_block (part0->part) - LocalParticle_add_to_at_turn(part, -1); - LocalParticle_set_at_element(part, num_elements); - if (flag_reset_s>0){ - LocalParticle_set_s(part, line_length); - } - //end_per_particle_block - } - - // check_is_active has different implementation on CPU and GPU - - #define CPU_SERIAL_IMPLEM //only_for_context cpu_serial - #define CPU_OMP_IMPLEM //only_for_context cpu_openmp - - #ifdef CPU_SERIAL_IMPLEM - - /*gpufun*/ - int64_t check_is_active(LocalParticle* part) { - int64_t ipart=0; - while (ipart < part->_num_active_particles){ - #ifdef XSUITE_RESTORE_LOSS - ipart++; - #else - if (part->state[ipart]<1){ - LocalParticle_exchange( - part, ipart, part->_num_active_particles-1); - part->_num_active_particles--; - part->_num_lost_particles++; - } - else{ - ipart++; - } - #endif - } - - if (part->_num_active_particles==0){ - return 0;//All particles lost - } else { - return 1; //Some stable particles are still present - } - } - - #else // not CPU_SERIAL_IMPLEM - #ifdef CPU_OMP_IMPLEM - - /*gpufun*/ - int64_t check_is_active(LocalParticle* part) { - #ifndef SKIP_SWAPS - int64_t ipart = part->ipart; - int64_t endpart = part->endpart; - - int64_t left = ipart; - int64_t right = endpart - 1; - int64_t swap_made = 0; - int64_t has_alive = 0; - - if (left == right) return part->state[left] > 0; - - while (left < right) { - if (part->state[left] > 0) { - left++; - has_alive = 1; - } - else if (part->state[right] <= 0) right--; - else { - LocalParticle_exchange(part, left, right); - left++; - right--; - swap_made = 1; - } - } - - return swap_made || has_alive; - #else - return 1; - #endif - } - - /*gpufun*/ - void count_reorganized_particles(LocalParticle* part) { - int64_t num_active = 0; - int64_t num_lost = 0; - - for (int64_t i = part->ipart; i < part->endpart; i++) { - if (part->state[i] <= -999999999) break; - else if (part->state[i] > 0) num_active++; - else num_lost++; - } - - part->_num_active_particles = num_active; - part->_num_lost_particles = num_lost; - } - - #else // not CPU_SERIAL_IMPLEM and not CPU_OMP_IMPLEM - - /*gpufun*/ - int64_t check_is_active(LocalParticle* part) { - return LocalParticle_get_state(part)>0; - }; - - #endif // CPU_OMP_IMPLEM - #endif // CPU_SERIAL_IMPLEM - - #undef CPU_SERIAL_IMPLEM //only_for_context cpu_serial - #undef CPU_OMP_IMPLEM //only_for_context cpu_openmp - - - ''' - - source = '\n\n'.join([src_typedef, src_adders, src_getters, - src_setters, src_scalers, src_exchange, - src_particles_to_local, src_local_to_particles, - src_angles, custom_source]) - - source += """ - #ifdef XTRACK_GLOBAL_XY_LIMIT - - /*gpufun*/ - void global_aperture_check(LocalParticle* part0) { - if (LocalParticle_check_track_flag( - part0, XS_FLAG_IGNORE_GLOBAL_APERTURE)){ - return; - } - - //start_per_particle_block (part0->part) - double const x = LocalParticle_get_x(part); - double const y = LocalParticle_get_y(part); - - int64_t const is_alive = (int64_t)( - (x >= -XTRACK_GLOBAL_XY_LIMIT) && - (x <= XTRACK_GLOBAL_XY_LIMIT) && - (y >= -XTRACK_GLOBAL_XY_LIMIT) && - (y <= XTRACK_GLOBAL_XY_LIMIT) ); - - // I assume that if I am in the function is because - if (!is_alive){ - LocalParticle_set_state(part, -1); - } - //end_per_particle_block - } - - #endif - - /*gpufun*/ - void LocalParticle_add_to_energy(LocalParticle* part, double delta_energy, int pz_only ){ - double ptau = LocalParticle_get_ptau(part); - double const p0c = LocalParticle_get_p0c(part); - double const charge_ratio = LocalParticle_get_charge_ratio(part); - double const chi = LocalParticle_get_chi(part); - double const mass_ratio = charge_ratio / chi; - - ptau += delta_energy/p0c / mass_ratio; - - double const old_rpp = LocalParticle_get_rpp(part); - - LocalParticle_update_ptau(part, ptau); - - if (!pz_only) { - double const new_rpp = LocalParticle_get_rpp(part); - double const f = old_rpp / new_rpp; - LocalParticle_scale_px(part, f); - LocalParticle_scale_py(part, f); - } - } - - - /*gpufun*/ - void LocalParticle_update_p0c(LocalParticle* part, double new_p0c_value){ - - double const mass0 = LocalParticle_get_mass0(part); - double const old_p0c = LocalParticle_get_p0c(part); - double const old_delta = LocalParticle_get_delta(part); - double const old_beta0 = LocalParticle_get_beta0(part); - - double const ppc = old_p0c * old_delta + old_p0c; - double const new_delta = (ppc - new_p0c_value)/new_p0c_value; - - double const new_energy0 = sqrt(new_p0c_value*new_p0c_value + mass0 * mass0); - double const new_beta0 = new_p0c_value / new_energy0; - double const new_gamma0 = new_energy0 / mass0; - - LocalParticle_set_p0c(part, new_p0c_value); - LocalParticle_set_gamma0(part, new_gamma0); - LocalParticle_set_beta0(part, new_beta0); - - LocalParticle_update_delta(part, new_delta); - - LocalParticle_scale_px(part, old_p0c/new_p0c_value); - LocalParticle_scale_py(part, old_p0c/new_p0c_value); - - LocalParticle_scale_zeta(part, new_beta0/old_beta0); - - } - - /*gpufun*/ - void LocalParticle_kill_particle(LocalParticle* part, int64_t kill_state) { - LocalParticle_set_x(part, 1e30); - LocalParticle_set_px(part, 1e30); - LocalParticle_set_y(part, 1e30); - LocalParticle_set_py(part, 1e30); - LocalParticle_set_zeta(part, 1e30); - LocalParticle_update_delta(part, -1); // zero energy - LocalParticle_set_state(part, kill_state); - } - """ - return source - @classmethod def part_energy_varnames(cls): return [vv for tt, vv in cls.part_energy_vars] diff --git a/xtrack/tracker.py b/xtrack/tracker.py index 27128d045..ba38bd178 100644 --- a/xtrack/tracker.py +++ b/xtrack/tracker.py @@ -47,7 +47,6 @@ def __init__( track_kernel=None, particles_monitor_class=None, extra_headers=(), - local_particle_src=None, _prebuilding_kernels=False, ): @@ -62,16 +61,12 @@ def __init__( raise ValueError("`enable_pipeline_hold` is not implemented in " "non-collective mode") - if local_particle_src is None: - local_particle_src = xt.Particles.gen_local_particle_api() - if not particles_monitor_class: particles_monitor_class = self._get_default_monitor_class() self.line = line self.particles_monitor_class = particles_monitor_class self.extra_headers = extra_headers - self.local_particle_src = local_particle_src self._enable_pipeline_hold = enable_pipeline_hold self.use_prebuilt_kernels = use_prebuilt_kernels self.track_flags = TrackFlags() @@ -684,9 +679,7 @@ def _build_kernel( kernel_descriptions=kernels, extra_headers=self._config_to_headers() + headers, extra_classes=kernel_element_classes + extra_classes, - apply_to_source=[ - partial(_handle_per_particle_blocks, - local_particle_src=self.local_particle_src)], + apply_to_source=[_handle_per_particle_blocks], specialize=True, compile=compile, save_source_as=f'{module_name}.c' if module_name else None, @@ -1128,8 +1121,6 @@ def _track_no_collective( self.track_flags.XS_FLAG_BACKTRACK = True return self._track_no_collective(**kwargs) - self.local_particle_src = particles.gen_local_particle_api() - if freeze_longitudinal: kwargs = locals().copy() kwargs.pop('self') diff --git a/xtrack/twiss.py b/xtrack/twiss.py index 77fd4d625..983a691b4 100644 --- a/xtrack/twiss.py +++ b/xtrack/twiss.py @@ -2837,7 +2837,6 @@ def _build_auxiliary_tracker_with_extra_markers(tracker, at_s, marker_prefix, io_buffer=tracker.io_buffer, line=auxline, particles_monitor_class=None, - local_particle_src=tracker.local_particle_src ) auxtracker.line.config = tracker.line.config.copy() auxtracker.line._extra_config = tracker.line._extra_config.copy()