WIP Enable inverse tracking for dipole fringe edges

tests/test_elements_thick.py

@@ -1138,6 +1138,37 @@ def test_fringe_implementations(test_context):
     xo.assert_allclose(np.linalg.det(R_ptc), 1, rtol=0, atol=1e-8) # Symplecticity check
 
 
+@for_all_test_contexts
+def test_full_dipole_edge_backtrack(test_context):
+
+    skip_if_forbid_compile()
+
+    fringe = xt.DipoleEdge(
+        k=0.12, fint=100, hgap=0.035, model='full')
+
+    line = xt.Line(elements=[fringe])
+    line.particle_ref = xp.Particles(
+        mass0=xp.PROTON_MASS_EV, beta0=0.5)
+    line.reset_s_at_end_turn = False
+    line.build_tracker(_context=test_context)
+
+    p0 = line.build_particles(px=0.5, py=0.001, y=0.01, delta=0.1)
+    p_test = p0.copy(_context=test_context)
+
+    line.track(p_test)
+    assert np.all(p_test.state == 1)
+    line.track(p_test, backtrack=True)
+
+    p0.move(_context=xo.context_default)
+    p_test.move(_context=xo.context_default)
+
+    assert np.all(p_test.state == 1)
+    for coordinate in ['x', 'px', 'y', 'py', 'zeta', 'delta', 's']:
+        xo.assert_allclose(
+            getattr(p_test, coordinate), getattr(p0, coordinate),
+            rtol=0, atol=1e-12)
+
+
 @for_all_test_contexts
 def test_backtrack_with_bend_quadrupole_and_cfm(test_context):
 
@@ -1188,8 +1219,8 @@ def test_backtrack_with_bend_quadrupole_and_cfm(test_context):
     xo.assert_allclose(p2.zeta, p0.zeta, atol=1e-15, rtol=0)
     xo.assert_allclose(p2.delta, p0.delta, atol=1e-15, rtol=0)
 
-    # Same for dipole edge
-    de = xt.DipoleEdge(e1=0.1, k=3, fint=0.3)
+    # Same for a full dipole edge
+    de = xt.DipoleEdge(e1=0.1, k=3, fint=0.3, model='full')
     line = xt.Line(elements=[de])
     line.particle_ref = xp.Particles(mass0=xp.PROTON_MASS_EV, beta0=0.5)
     line.reset_s_at_end_turn = False
@@ -1198,16 +1229,27 @@ def test_backtrack_with_bend_quadrupole_and_cfm(test_context):
                                 zeta=0.05, delta=0.01)
     p1 = p0.copy(_context=test_context)
     line.track(p1)
-    p1.move(_context=xo.context_default)
     assert np.all(p1.state == 1)
-    line.configure_bend_model(edge='full')
     p2 = p1.copy(_context=test_context)
     line.track(p2, backtrack=True)
-    p2.move(_context=xo.context_default)
-    assert np.all(p2.state == -32)
 
-    # Same for combined function magnet
-    cfm = xt.Bend(length=1.0, k1=0.2, angle=0.1)
+    p0.move(_context=xo.context_default)
+    p2.move(_context=xo.context_default)
+    assert np.all(p2.state == 1)
+    xo.assert_allclose(p2.s, p0.s, atol=1e-14, rtol=0)
+    xo.assert_allclose(p2.x, p0.x, atol=1e-14, rtol=0)
+    xo.assert_allclose(p2.px, p0.px, atol=1e-14, rtol=0)
+    xo.assert_allclose(p2.y, p0.y, atol=1e-14, rtol=0)
+    xo.assert_allclose(p2.py, p0.py, atol=1e-14, rtol=0)
+    xo.assert_allclose(p2.zeta, p0.zeta, atol=1e-14, rtol=0)
+    xo.assert_allclose(p2.delta, p0.delta, atol=1e-14, rtol=0)
+
+    # Same for a combined-function magnet with dipole-only full edges
+    cfm = xt.Bend(
+        length=1.0, k1=0.2, angle=0.1,
+        edge_entry_model='dipole-only', edge_exit_model='dipole-only',
+        edge_entry_fint=0.5, edge_exit_fint=0.5,
+        edge_entry_hgap=0.03, edge_exit_hgap=0.03)
     line = xt.Line(elements=[cfm])
     line.particle_ref = xp.Particles(mass0=xp.PROTON_MASS_EV, beta0=0.5)
     line.reset_s_at_end_turn = False

xtrack/beam_elements/elements_src/dipoleedge.h

@@ -52,22 +52,26 @@ void DipoleEdge_track_local_particle(DipoleEdgeData el, LocalParticle* part0){
     }
     else if (model == 1){
 
-        if (LocalParticle_check_track_flag(part0, XS_FLAG_BACKTRACK)){
-            START_PER_PARTICLE_BLOCK(part0, part);
-                LocalParticle_kill_particle(part, -32);
-            END_PER_PARTICLE_BLOCK;
-            return;
-        }
-
         double const e1 = DipoleEdgeData_get_e1(el);
         double const fint = DipoleEdgeData_get_fint(el);
         double const hgap = DipoleEdgeData_get_hgap(el);
         double const k = DipoleEdgeData_get_k(el);
         int64_t const side = DipoleEdgeData_get_side(el);
+        int64_t const backtrack =
+            LocalParticle_check_track_flag(part0, XS_FLAG_BACKTRACK);
 
-        START_PER_PARTICLE_BLOCK(part0, part);
-            DipoleEdgeNonLinear_single_particle(part, k, e1, fint, hgap, side);
-        END_PER_PARTICLE_BLOCK;
+        if (backtrack){
+            START_PER_PARTICLE_BLOCK(part0, part);
+                DipoleEdgeNonLinear_single_particle_backtrack(
+                    part, k, e1, fint, hgap, side);
+            END_PER_PARTICLE_BLOCK;
+        }
+        else{
+            START_PER_PARTICLE_BLOCK(part0, part);
+                DipoleEdgeNonLinear_single_particle(
+                    part, k, e1, fint, hgap, side);
+            END_PER_PARTICLE_BLOCK;
+        }
 
     }
 

xtrack/beam_elements/elements_src/magnet_edge.h

@@ -38,6 +38,7 @@ void MagnetEdge_track_local_particle(MagnetEdgeData el, LocalParticle* part0)
         part0,
         model,
         is_exit,
+        is_exit, // physical_is_exit
         half_gap,
         knorm,
         kskew,

xtrack/beam_elements/elements_src/track_dipole_edge_nonlinear.h

@@ -43,4 +43,39 @@ void DipoleEdgeNonLinear_single_particle(LocalParticle* part,
     }
 }
 
+
+GPUFUN
+void DipoleEdgeNonLinear_single_particle_backtrack(LocalParticle* part,
+            double const k, double const e1, double const fint, double const hgap,
+            int64_t const side
+){
+
+    double sin_, cos_, tan_;
+    if (fabs(e1) < 10e-10) {
+        sin_ = -999.0; cos_ = -999.0; tan_ = -999.0;
+    }
+    else{
+        sin_ = sin(e1); cos_ = cos(e1); tan_ = tan(e1);
+    }
+
+    if (side == 0){ // entry: invert wedge, fringe, rotation
+        if (sin_ > -99.){
+            Wedge_single_particle(part, e1, k);
+        }
+        DipoleFringe_single_particle_backtrack(part, fint, hgap, k);
+        if (sin_ > -99.){
+            YRotation_single_particle(part, sin_, cos_, tan_);
+        }
+    }
+    else if (side == 1){ // exit: invert rotation, fringe, wedge
+        if (sin_ > -99.){
+            YRotation_single_particle(part, sin_, cos_, tan_);
+        }
+        DipoleFringe_single_particle_backtrack(part, fint, hgap, -k);
+        if (sin_ > -99.){
+            Wedge_single_particle(part, e1, k);
+        }
+    }
+}
+
 #endif