get_line_with_second_order_maps(): split at thick elements

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,271 @@ 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_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):
 

xtrack/line.py

@@ -5368,13 +5368,21 @@ 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). 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
         -------
@@ -5383,15 +5391,30 @@ def get_line_with_second_order_maps(self, split_at):
         '''
         self._method_incompatible_with_compose()
 
+        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:
-            if ee in ele_cut_ext:
+        for ee in ele_names:
+            if ee in ele_cut_set:
                 ele_cut_sorted.append(ee)
 
         elements_map_line = []
@@ -5400,18 +5423,29 @@ def get_line_with_second_order_maps(self, split_at):
 
         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 = 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
 
             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)
             names_map_line.append(f'map_{ii}')
             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()