Handle skip_in_loss_location_refinement for non-collective coll…

…imators
xsuite · Feb 11, 2023 · 5ab48ab · 5ab48ab
1 parent 76a5309
commit 5ab48ab
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/xtrack/loss_location_refinement/loss_location_refinement.py b/xtrack/loss_location_refinement/loss_location_refinement.py
@@ -346,7 +346,8 @@ def build_interp_tracker(_buffer, s0, s1, s_interp, aper_0, aper_1, aper_interp,
 
     for i_ele in range(i_start_thin_0+1, i_start_thin_1):
         ee = tracker.line.elements[i_ele]
-        if not ee.__class__.__name__.startswith('Drift'):
+        if (not ee.__class__.__name__.startswith('Drift') and not (
+                hasattr(ee, 'behaves_like_drift') and ee.behaves_like_drift)):
             assert not hasattr(ee, 'isthick') or not ee.isthick
             ss_ee = tracker._tracker_data.element_s_locations[i_ele]
             elements.append(ee.copy(_buffer=_buffer))
@@ -379,13 +380,18 @@ def build_interp_tracker(_buffer, s0, s1, s_interp, aper_0, aper_1, aper_interp,
 
 def find_previous_drift(tracker, i_aperture):
 
+    line = tracker.line
+
     ii=i_aperture
     found = False
     while not(found):
-        ccnn = tracker.line.elements[ii].__class__.__name__
+        ee = line.element_dict[line.element_names[ii]]
+        ccnn = ee.__class__.__name__
         #print(ccnn)
         if ccnn == 'Drift':
             found = True
+        elif hasattr(ee, 'behaves_like_drift') and ee.behaves_like_drift:
+            found = True
         else:
             ii -= 1
     i_start = ii + 1