From 5ab48ab89f62a3afc9bdd9cebd565aaf5a296d31 Mon Sep 17 00:00:00 2001 From: giadarol Date: Sat, 11 Feb 2023 10:16:38 +0100 Subject: [PATCH] Handle ```skip_in_loss_location_refinement``` for non-collective collimators --- .../loss_location_refinement.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/xtrack/loss_location_refinement/loss_location_refinement.py b/xtrack/loss_location_refinement/loss_location_refinement.py index b284f6ea8..7d913bfd6 100644 --- 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