From 3905b920e129ce57babe4bce15066923fca86754 Mon Sep 17 00:00:00 2001 From: Michi Date: Fri, 10 Jul 2026 15:17:48 +0200 Subject: [PATCH 1/2] get_line_with_second_order_maps(): proper treatment of thick split elements --- tests/test_second_order_taylor_map.py | 146 ++++++++++++++++++++++++++ xtrack/line.py | 28 ++++- 2 files changed, 170 insertions(+), 4 deletions(-) diff --git a/tests/test_second_order_taylor_map.py b/tests/test_second_order_taylor_map.py index 34a8fbf48..93eb28f6d 100644 --- a/tests/test_second_order_taylor_map.py +++ b/tests/test_second_order_taylor_map.py @@ -6,6 +6,7 @@ import pathlib import numpy as np +import pytest from cpymad.madx import Madx import xobjects as xo @@ -59,6 +60,151 @@ def test_line_with_second_order_maps(test_context): xo.assert_allclose(tw_map.line_length, tw.line_length, rtol=0, atol=5e-9) +@for_all_test_contexts +def test_line_with_second_order_maps_split_at_octupoles(test_context): + + line = xt.load(test_data_folder / + 'hllhc15_thick/lhc_thick_with_knobs.json') + line.build_tracker(_context=test_context) + line.vars['vrf400'] = 16 + line.vars['lagrf400.b1'] = 0.5 + + # power the Landau octupoles (strong amplitude detuning) + for vv in line.vars.keys(): + if vv.startswith('kof.') or vv.startswith('kod.'): + line.vars[vv] = 40. + + mo_names = [nn for nn in line.element_names if nn.startswith('mo.')] + assert len(mo_names) == 168 + + ele_cut = ['ip1', 'ip2', 'ip3', 'ip4', 'ip5', 'ip6', 'ip7'] + + # splitting also at the octupoles keeps them as exact elements between + # the maps (split elements are excluded from the maps) + line_maps_keep = line.get_line_with_second_order_maps( + split_at=ele_cut + mo_names) + line_maps = line.get_line_with_second_order_maps(split_at=ele_cut) + for lm in (line_maps_keep, line_maps): + lm.build_tracker(_context=test_context) + + # the kept octupoles are present as exact elements + assert isinstance(line_maps_keep[mo_names[0]], xt.Octupole) + assert mo_names[0] not in line_maps.element_names + + tw = line.twiss() + tw_map = line_maps_keep.twiss() + + xo.assert_allclose(tw_map.rows[ele_cut].s, tw.rows[ele_cut].s, rtol=0, atol=1e-12) + xo.assert_allclose(tw_map.rows[ele_cut].x, tw.rows[ele_cut].x, rtol=0, atol=1e-12) + xo.assert_allclose(tw_map.rows[ele_cut].px, tw.rows[ele_cut].px, rtol=0, atol=1e-12) + xo.assert_allclose(tw_map.rows[ele_cut].y, tw.rows[ele_cut].y, rtol=0, atol=1e-12) + xo.assert_allclose(tw_map.rows[ele_cut].py, tw.rows[ele_cut].py, rtol=0, atol=1e-12) + xo.assert_allclose(tw_map.rows[ele_cut].betx, tw.rows[ele_cut].betx, rtol=1e-5, atol=0) + xo.assert_allclose(tw_map.rows[ele_cut].bety, tw.rows[ele_cut].bety, rtol=1e-5, atol=0) + xo.assert_allclose(np.mod(tw_map.qx, 1), np.mod(tw.qx, 1), rtol=0, atol=1e-7) + xo.assert_allclose(np.mod(tw_map.qy, 1), np.mod(tw.qy, 1), rtol=0, atol=1e-7) + xo.assert_allclose(tw_map.dqx, tw.dqx, rtol=0, atol=5e-2) + xo.assert_allclose(tw_map.dqy, tw.dqy, rtol=0, atol=5e-2) + + # amplitude detuning is preserved by the kept octupoles (the small + # residual difference is the sextupole-driven detuning, which second + # order maps cannot carry), while it is absent from the plain map line: + # launch test particles at small and high amplitude and measure their + # tunes from the turn-by-turn data + def _tune_fft(xs): + """Fractional tune from turn-by-turn data (Hann window + parabolic + interpolation of the FFT peak).""" + xs = xs - np.mean(xs) + ff = np.abs(np.fft.rfft(xs * np.hanning(len(xs)))) + ii = np.argmax(ff[1:-1]) + 1 + dd = 0.5 * (ff[ii - 1] - ff[ii + 1]) / (ff[ii - 1] - 2 * ff[ii] + ff[ii + 1]) + return (ii + dd) / len(xs) + + dq = {} + for label, ll in (('ref', line), ('keep', line_maps_keep), + ('maps', line_maps)): + pp = ll.build_particles(x_norm=[0.5, 5., 0.5], y_norm=[0.5, 0.5, 5.], + nemitt_x=2.5e-6, nemitt_y=2.5e-6) + ll.track(pp, num_turns=512, turn_by_turn_monitor=True, + freeze_longitudinal=True) + mon = ll.record_last_track + ctx2np = test_context.nparray_from_context_array + qx = [_tune_fft(ctx2np(mon.x)[jj, :]) for jj in range(3)] + qy = [_tune_fft(ctx2np(mon.y)[jj, :]) for jj in range(3)] + dq[label] = (qx[1] - qx[0], qy[2] - qy[0]) # amplitude detuning + + assert dq['ref'][0] > 5e-3 # measured: ~1e-2 + assert dq['ref'][1] > 5e-3 + xo.assert_allclose(dq['keep'], dq['ref'], rtol=0, atol=1e-3) + xo.assert_allclose(dq['maps'], (0, 0), rtol=0, atol=5e-4) + + +@for_all_test_contexts +def test_second_order_maps_split_at_thick_elements(test_context): + + # FODO ring made only of exactly-linear elements (quadrupoles and + # expanded drifts) plus two thick octupoles kept exact in the map line + # by splitting at them: tracking through the map line must reproduce + # the full line exactly. + env = xt.Environment() + env.particle_ref = xt.Particles(p0c=10e9) + components = [] + for cc in range(8): + components += [ + env.new(f'qf{cc}', xt.Quadrupole, k1=0.12, length=0.5), + env.new(f'd1{cc}', xt.Drift, length=2.), + env.new(f'qd{cc}', xt.Quadrupole, k1=-0.12, length=0.5), + env.new(f'd2{cc}', xt.Drift, length=2.), + ] + # two octupoles back to back (empty span between kept elements) and a + # marker directly downstream (empty span between keep and split_at) + components[8:8] = [env.new('mo1', xt.Octupole, k3=3000., length=0.3), + env.new('mo2', xt.Octupole, k3=-2000., length=0.3), + env.new('m1', xt.Marker)] + line = env.new_line(components=components) + line.twiss_default['method'] = '4d' + line.build_tracker(_context=test_context) + + line_maps_keep = line.get_line_with_second_order_maps( + split_at=['m1', 'mo1', 'mo2']) + line_maps = line.get_line_with_second_order_maps(split_at=['m1']) + + # no maps are generated for the empty spans between adjacent cuts + names = list(line_maps_keep.element_names) + assert names[names.index('mo1') + 1] == 'mo2' + assert names[names.index('mo2') + 1] == 'm1' + + for lm in (line_maps_keep, line_maps): + lm.twiss_default['method'] = '4d' + lm.build_tracker(_context=test_context) + + tw = line.twiss() + tw_keep = line_maps_keep.twiss() + xo.assert_allclose(tw_keep.qx, tw.qx, rtol=0, atol=1e-8) + xo.assert_allclose(tw_keep.qy, tw.qy, rtol=0, atol=1e-8) + + p_test = dict(x=2e-3, px=1e-5, y=-1.5e-3, py=2e-5) + res = {} + for label, ll in (('keep', line_maps_keep), ('maps', line_maps), + ('ref', line)): + pp = ll.build_particles(**p_test) + ll.track(pp, num_turns=50) + res[label] = np.array([getattr(pp, cc)[0] for cc in ['x', 'px', 'y', 'py']]) + + # kept octupoles -> exact (everything else is linear) + xo.assert_allclose(res['keep'], res['ref'], rtol=0, atol=1e-12) + # octupoles inside the maps -> their nonlinearity is lost + assert np.max(np.abs(res['maps'] - res['ref'])) > 1e-5 + + # the original line is not affected by the map lines (shared elements + # must not be moved out of its buffer) + pp = line.build_particles(**p_test) + line.track(pp, num_turns=1) + + with pytest.raises(ValueError, match='not present in the line'): + line.get_line_with_second_order_maps(split_at=['does_not_exist']) + + @for_all_test_contexts def test_second_order_maps_against_madx(test_context, sandbox_cwd): diff --git a/xtrack/line.py b/xtrack/line.py index de6e14e4c..3267a57bf 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -5368,13 +5368,18 @@ def merge_consecutive_multipoles(self, inplace=True, keep=None): def get_line_with_second_order_maps(self, split_at): ''' - Return a new lines with segments definded by the elements in `split_at` + Return a new line with segments defined by the elements in `split_at` replaced by second order maps. Parameters ---------- split_at : list of str - Names of elements at which to split the line. + Names of elements at which to split the line. These elements are + kept as they are in the new line and are excluded from the maps: + each map spans from the exit of one split element to the + entrance of the next. Hence also thick and/or nonlinear elements + can be preserved exactly by splitting at them (e.g. octupoles, + to retain their amplitude detuning). Returns ------- @@ -5383,17 +5388,24 @@ def get_line_with_second_order_maps(self, split_at): ''' self._method_incompatible_with_compose() + missing = set(split_at) - set(self.element_names) + if missing: + raise ValueError(f'Elements {sorted(missing)} are not present in the line') + ele_cut_ext = split_at.copy() if self.element_names[0] not in ele_cut_ext: ele_cut_ext.insert(0, self.element_names[0]) if self.element_names[-1] not in ele_cut_ext: ele_cut_ext.append(self.element_names[-1]) + ele_cut_set = set(ele_cut_ext) ele_cut_sorted = [] for ee in self.element_names: - if ee in ele_cut_ext: + if ee in ele_cut_set: ele_cut_sorted.append(ee) + ele_idx = {nn: ii for ii, nn in enumerate(self.element_names)} + elements_map_line = [] names_map_line = [] tw = self.twiss() @@ -5402,8 +5414,16 @@ def get_line_with_second_order_maps(self, split_at): names_map_line.append(ele_cut_sorted[ii]) elements_map_line.append(self.get(ele_cut_sorted[ii])) + # the split element is placed in the new line as it is, hence it + # is excluded from the map: the map starts at its exit, i.e. at + # the entrance of the following element (relevant for thick + # split elements) + map_start = self.element_names[ele_idx[ele_cut_sorted[ii]] + 1] + if map_start == ele_cut_sorted[ii+1]: + continue # nothing between this element and the next cut + smap = xt.SecondOrderTaylorMap.from_line( - self, start=ele_cut_sorted[ii], + self, start=map_start, end=ele_cut_sorted[ii+1], twiss_table=tw, _buffer=self._buffer) From 6f8b5b4c7a9b509d1bb948b4a4b13daf97f6c062 Mon Sep 17 00:00:00 2001 From: Michi Date: Fri, 10 Jul 2026 17:02:59 +0200 Subject: [PATCH 2/2] fix edge case for lines with repeated element names (e.g. after slicing following element insertion) --- tests/test_second_order_taylor_map.py | 120 ++++++++++++++++++++++++++ xtrack/line.py | 38 +++++--- 2 files changed, 146 insertions(+), 12 deletions(-) diff --git a/tests/test_second_order_taylor_map.py b/tests/test_second_order_taylor_map.py index 93eb28f6d..56fbed3a7 100644 --- a/tests/test_second_order_taylor_map.py +++ b/tests/test_second_order_taylor_map.py @@ -205,6 +205,126 @@ def test_second_order_maps_split_at_thick_elements(test_context): line.get_line_with_second_order_maps(split_at=['does_not_exist']) +@for_all_test_contexts +def test_second_order_maps_split_at_repeated_names(test_context): + + # ring with REPEATED element names (same elements placed several times, + # as it happens e.g. for the drift pieces generated by element + # insertions): element handling refers to the disambiguated names + # 'name::N' used in the line/twiss tables. As above, all elements are + # exactly linear except one thick octupole, so the map line with the + # octupole split at must reproduce the full line exactly. + elements = { + 'qf': xt.Quadrupole(k1=0.12, length=0.5), + 'qd': xt.Quadrupole(k1=-0.12, length=0.5), + 'dd': xt.Drift(length=2.), + 'mo': xt.Octupole(k3=3000., length=0.3), + 'm1': xt.Marker(), + } + element_names = [] + for cc in range(8): + element_names += ['qf', 'dd', 'qd', 'dd'] + # octupole surrounded by repeated-name drifts (the map following it + # starts at a repeated element), marker in the second cell + element_names[2:2] = ['mo'] + element_names[8:8] = ['m1'] + line = xt.Line(elements=elements, element_names=element_names) + line.particle_ref = xt.Particles(p0c=10e9) + line.twiss_default['method'] = '4d' + line.build_tracker(_context=test_context) + + # split at the octupole (unique name) and at one particular occurrence + # of the repeated drift (disambiguated name) + line_maps_keep = line.get_line_with_second_order_maps( + split_at=['m1', 'mo', 'dd::5']) + line_maps = line.get_line_with_second_order_maps(split_at=['m1']) + + # the split elements appear in the new line under the same names, + # including the repeated first element of the line + for nn in ('qf::0', 'mo', 'm1', 'dd::5'): + assert nn in line_maps_keep.element_names + assert isinstance(line_maps_keep['dd::5'], xt.Drift) + + for lm in (line_maps_keep, line_maps): + lm.twiss_default['method'] = '4d' + lm.build_tracker(_context=test_context) + + tw = line.twiss() + tw_keep = line_maps_keep.twiss() + xo.assert_allclose(tw_keep.qx, tw.qx, rtol=0, atol=1e-8) + xo.assert_allclose(tw_keep.qy, tw.qy, rtol=0, atol=1e-8) + + p_test = dict(x=2e-3, px=1e-5, y=-1.5e-3, py=2e-5) + res = {} + for label, ll in (('keep', line_maps_keep), ('maps', line_maps), + ('ref', line)): + pp = ll.build_particles(**p_test) + ll.track(pp, num_turns=50) + res[label] = np.array([getattr(pp, cc)[0] for cc in ['x', 'px', 'y', 'py']]) + + # split octupole -> exact (everything else is linear) + xo.assert_allclose(res['keep'], res['ref'], rtol=0, atol=1e-12) + # octupole inside the maps -> its nonlinearity is lost + assert np.max(np.abs(res['maps'] - res['ref'])) > 1e-5 + + # a repeated plain name is ambiguous and not accepted + with pytest.raises(ValueError, match='not present in the line'): + line.get_line_with_second_order_maps(split_at=['dd']) + + +@for_all_test_contexts +def test_second_order_maps_split_at_sliced_line(test_context): + + # sliced lattice (thin slices and slice drifts, with parent/replica + # bookkeeping): the map line built on it must reproduce the sliced line + # exactly when the only nonlinear element (a thick octupole, excluded + # from the slicing) is split at. + env = xt.Environment() + env.particle_ref = xt.Particles(p0c=10e9) + components = [] + for cc in range(8): + components += [ + env.new(f'qf{cc}', xt.Quadrupole, k1=0.12, length=0.5), + env.new(f'd1{cc}', xt.Drift, length=2.), + env.new(f'qd{cc}', xt.Quadrupole, k1=-0.12, length=0.5), + env.new(f'd2{cc}', xt.Drift, length=2.), + ] + components[8:8] = [env.new('mo1', xt.Octupole, k3=3000., length=0.3), + env.new('m1', xt.Marker)] + line = env.new_line(components=components) + line.twiss_default['method'] = '4d' + line.slice_thick_elements(slicing_strategies=[ + xt.Strategy(slicing=None), # default: don't slice + xt.Strategy(slicing=xt.Teapot(4), element_type=xt.Quadrupole), + ]) + line.build_tracker(_context=test_context) + + line_maps_keep = line.get_line_with_second_order_maps( + split_at=['m1', 'mo1']) + line_maps = line.get_line_with_second_order_maps(split_at=['m1']) + for lm in (line_maps_keep, line_maps): + lm.twiss_default['method'] = '4d' + lm.build_tracker(_context=test_context) + + tw = line.twiss() + tw_keep = line_maps_keep.twiss() + xo.assert_allclose(tw_keep.qx, tw.qx, rtol=0, atol=1e-8) + xo.assert_allclose(tw_keep.qy, tw.qy, rtol=0, atol=1e-8) + + p_test = dict(x=2e-3, px=1e-5, y=-1.5e-3, py=2e-5) + res = {} + for label, ll in (('keep', line_maps_keep), ('maps', line_maps), + ('ref', line)): + pp = ll.build_particles(**p_test) + ll.track(pp, num_turns=50) + res[label] = np.array([getattr(pp, cc)[0] for cc in ['x', 'px', 'y', 'py']]) + + # split octupole -> exact (thin quad slices and drifts are linear) + xo.assert_allclose(res['keep'], res['ref'], rtol=0, atol=1e-12) + # octupole inside the maps -> its nonlinearity is lost + assert np.max(np.abs(res['maps'] - res['ref'])) > 1e-5 + + @for_all_test_contexts def test_second_order_maps_against_madx(test_context, sandbox_cwd): diff --git a/xtrack/line.py b/xtrack/line.py index 3267a57bf..5285928fd 100644 --- a/xtrack/line.py +++ b/xtrack/line.py @@ -5379,7 +5379,10 @@ def get_line_with_second_order_maps(self, split_at): each map spans from the exit of one split element to the entrance of the next. Hence also thick and/or nonlinear elements can be preserved exactly by splitting at them (e.g. octupoles, - to retain their amplitude detuning). + to retain their amplitude detuning). Repeated elements are + referred to by their disambiguated name 'name::N' (as shown in + the line table and in the twiss table); the same names are used + in the returned line. Returns ------- @@ -5388,37 +5391,48 @@ def get_line_with_second_order_maps(self, split_at): ''' self._method_incompatible_with_compose() - missing = set(split_at) - set(self.element_names) + if not self._has_valid_tracker(): + self.build_tracker() + + # element names disambiguated for repeated elements ('name::N', as + # in the line table and in the twiss table); for non-repeated + # elements they coincide with the plain element names + ele_names = self._element_names_unique + + missing = set(split_at) - set(ele_names) if missing: raise ValueError(f'Elements {sorted(missing)} are not present in the line') + ele_idx = {nn: ii for ii, nn in enumerate(ele_names)} + ele_cut_ext = split_at.copy() - if self.element_names[0] not in ele_cut_ext: - ele_cut_ext.insert(0, self.element_names[0]) - if self.element_names[-1] not in ele_cut_ext: - ele_cut_ext.append(self.element_names[-1]) + if ele_names[0] not in ele_cut_ext: + ele_cut_ext.insert(0, ele_names[0]) + if ele_names[-1] not in ele_cut_ext: + ele_cut_ext.append(ele_names[-1]) ele_cut_set = set(ele_cut_ext) ele_cut_sorted = [] - for ee in self.element_names: + for ee in ele_names: if ee in ele_cut_set: ele_cut_sorted.append(ee) - ele_idx = {nn: ii for ii, nn in enumerate(self.element_names)} - elements_map_line = [] names_map_line = [] tw = self.twiss() for ii in range(len(ele_cut_sorted)-1): names_map_line.append(ele_cut_sorted[ii]) - elements_map_line.append(self.get(ele_cut_sorted[ii])) + # element object by its unique name: `element_names` and + # `self.element_names` are index-aligned (all occurrences of a + # repeated element share the same object) + elements_map_line.append(self.get(self.element_names[ele_idx[ele_cut_sorted[ii]]])) # the split element is placed in the new line as it is, hence it # is excluded from the map: the map starts at its exit, i.e. at # the entrance of the following element (relevant for thick # split elements) - map_start = self.element_names[ele_idx[ele_cut_sorted[ii]] + 1] + map_start = ele_names[ele_idx[ele_cut_sorted[ii]] + 1] if map_start == ele_cut_sorted[ii+1]: continue # nothing between this element and the next cut @@ -5431,7 +5445,7 @@ def get_line_with_second_order_maps(self, split_at): elements_map_line.append(smap) names_map_line.append(ele_cut_sorted[-1]) - elements_map_line.append(self.get(ele_cut_sorted[-1])) + elements_map_line.append(self.get(self.element_names[ele_idx[ele_cut_sorted[-1]]])) line_maps = Line(elements=elements_map_line, element_names=names_map_line) line_maps.particle_ref = self.particle_ref.copy()