diff --git a/pyproject.toml b/pyproject.toml index 9773999bb..ba2bbf672 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,7 @@ tests = [ "PyHEADTAIL", "pytest", "pytest-mock", - "pymadng==0.9.3", + "pymadng==0.9.6", "requests-mock", "tfs-pandas", ] diff --git a/tests/test_madnginterface.py b/tests/test_madnginterface.py index 0120ce9f0..525538c37 100644 --- a/tests/test_madnginterface.py +++ b/tests/test_madnginterface.py @@ -216,6 +216,57 @@ def test_madng_conversion_drift_slice(): xo.assert_allclose(tw_ng.beta11_ng, tw.betx, rtol=1e-8) xo.assert_allclose(tw_ng.beta22_ng, tw.bety, rtol=1e-8) + +def test_madng_track_single_rbend(): + line = xt.Line( + elements=[xt.RBend(length_straight=2, angle=0.2)], + element_names=['rbend'], + particle_ref=xt.Particles(p0c=1e9), + ) + line.configure_bend_model( + core="bend-kick-bend", + integrator="uniform", + num_multipole_kicks=1, + edge="full", # MAD-NG has this by default, so we need to match it + ) + mng = line.to_madng(sequence_name='seq') + + X0 = { + 'x': 3e-5, + 'px': -1e-7, + 'y': -3e-5, + 'py': 1e-7, + 't': 0.0, + 'pt': 0.0, + } + mng.send( + """ + trk = MAD.track{ + sequence=MADX.seq, + nturn=1, + observe=0, + save='atentry', + X0=py:recv(), + } + """ + ).send(list(X0.values())) + madng_track = mng.trk.to_df().iloc[-1] + particles = xt.Particles( + p0c=1e9, + x=X0['x'], + px=X0['px'], + y=X0['y'], + py=X0['py'], + ) + line.track(particles) + + # This rtol below is concerning -> should be investigated, could be a bug in MAD-NG or in Xsuite. + xo.assert_allclose(madng_track.x, particles.x[0], rtol=7e-11, atol=0) #rtol is much worse -> so we use atol=0 + xo.assert_allclose(madng_track.px, particles.px[0], rtol=1.4e-10, atol=0) + xo.assert_allclose(madng_track.y, particles.y[0], rtol=1e-16, atol=0) + xo.assert_allclose(madng_track.py, particles.py[0], rtol=2e-16, atol=0) + xo.assert_allclose(madng_track.s, line.get_length(), rtol=1e-16, atol=0) + def test_madng_interface_with_slicing(): line = xt.load(test_data_folder / 'hllhc15_thick/lhc_thick_with_knobs.json') @@ -657,3 +708,69 @@ def test_madng_tpsa_optics_with_nonzero_initial_orbit(): tw_sol = line.twiss(**init) for q in quants: xo.assert_allclose(tw_sol[q, 'end'], target[q], rtol=2e-4, atol=1e-5) + + +def test_madng_tpsa_rmatrix_match_and_jacobian(): + """R-matrix matching through ``Line.match`` with TPSA derivatives. + + The Jacobian rows are checked against central finite differences of Xsuite's + own ``get_R_matrix``, and the matched terms against the same function after + solving, so the numbers are tested and not merely that the machinery runs. + Two ranges are used, because MAD-NG indexes the stored transfer maps by + range and an off-by-one there is invisible with a single one. Only + transverse terms are compared, as Xsuite and MAD-NG do not share + longitudinal coordinates. The finite difference step is small against the + knobs, which are of order 1e-2, yet far enough above the twiss round-off + floor: at 1e-8 the comparison degrades to 7e-5. + + Going through ``Line.match`` rather than through the action covers a lookup + that used to be fatal: matching labels the targets with ``rtag`` and the + MAD-NG interface stores the terms under that label, but the target read them + back under ``tag``. The miss fell through to ``get_R_matrix`` on the MAD-NG + table, which raised ``AttributeError`` on the absent ``values_at`` column. + """ + line = xt.load(test_data_folder / 'hllhc15_thick/lhc_thick_with_knobs.json') + line['on_disp'] = 0 + + specs = [('r12', 's.ds.l8.b1', 'ip8', 0, 1), + ('r34', 's.ds.l8.b1', 'e.ds.r8.b1', 2, 3)] + knobs = ['kq6.l8b1', 'kq7.l8b1'] + step = 1e-6 + + def terms_of(twiss): + return np.array([twiss.get_R_matrix(start=start, end=end)[ii, jj] + for _, start, end, ii, jj in specs]) + + initial_terms = terms_of(line.twiss4d()) + + expected_jac = np.zeros((len(specs), len(knobs))) + for col, knob in enumerate(knobs): + base = line[knob] + line[knob] = base + step + forward = terms_of(line.twiss4d()) + line[knob] = base - step + backward = terms_of(line.twiss4d()) + line[knob] = base + expected_jac[:, col] = (forward - backward) / (2 * step) + + target_terms = initial_terms * [0.9, 1.0] + opt = line.match( + use_tpsa=True, + solve=False, + vary=[xt.VaryList(knobs, step=1e-8)], + targets=[xt.TargetRmatrixTerm(term, value=value, start=start, end=end) + for (term, start, end, _, _), value in zip(specs, target_terms)], + ) + + action = opt.targets[0].action + got_terms = np.array([action.run()._data.attrs[t.rtag] for t in opt.targets]) + xo.assert_allclose(got_terms, initial_terms, rtol=1e-9, atol=1e-11) + + jac = np.array(action.acquire_jacobian()) + assert jac.shape == (len(specs), len(knobs)) + xo.assert_allclose(jac, expected_jac, rtol=1e-6, atol=1e-8) + + opt.solve() + assert opt.log()['penalty'][-1] < 1e-8 + xo.assert_allclose(terms_of(line.twiss4d()), target_terms, + rtol=1e-9, atol=1e-11) diff --git a/xtrack/madng_interface.py b/xtrack/madng_interface.py index 3d826b270..9ce007f63 100644 --- a/xtrack/madng_interface.py +++ b/xtrack/madng_interface.py @@ -1,13 +1,22 @@ -import numpy as np +from __future__ import annotations -from .match import Action -import os import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np import xtrack as xt +from xtrack.particles.particles import dptau2ddelta, ptau2delta +from xtrack.survey import SurveyTable -from xtrack.particles.particles import ptau2delta, dptau2ddelta +from .match import Action +# Xsuite name -> MAD-NG name. Not injective: ``delta`` and ``ptau`` are both +# ``pt`` in MAD-NG, which is why the reverse map is spelled out separately +# rather than derived from this one. NG_XS_MAP = { 'beta11': 'betx', 'beta22': 'bety', @@ -35,64 +44,130 @@ 'ptau': 'pt', } -BETA0_COLUMNS = ['x', 'px', 'y', 'py', 't', 'pt', - 'dx', 'dy', 'dpx', 'dpy', 'ddx', 'ddpx', 'ddy', 'ddpy', 'wx', 'phix', - 'wy', 'phiy', 'mu1', 'mu2', 'mu3', 'dmu1', 'dmu2', 'dmu3', 'r11', - 'r12', 'r21', 'r22', 'alfa11', 'alfa12', 'alfa13', 'alfa21', - 'alfa22', 'alfa23', 'alfa31', 'alfa32', 'alfa33', 'beta11', - 'beta12', 'beta13', 'beta21', 'beta22', 'beta23', 'beta31', - 'beta32', 'beta33', 'gama11', 'gama12', 'gama13', 'gama21', - 'gama22', 'gama23', 'gama31', 'gama32', 'gama33'] +# fmt: off +BETA0_COLUMNS = [ + 'x', 'px', 'y', 'py', 't', 'pt', + 'dx', 'dy', 'dpx', 'dpy', 'ddx', 'ddpx', 'ddy', 'ddpy', + 'wx', 'wxp', 'wy', 'wyp', + 'mu1', 'mu2', 'mu3', 'dmu1', 'dmu2', 'dmu3', + 'r11', 'r12', 'r21', 'r22', + 'alfa11', 'alfa12', 'alfa13', 'alfa21', 'alfa22', 'alfa23', 'alfa31', 'alfa32', 'alfa33', + 'beta11', 'beta12', 'beta13', 'beta21', 'beta22', 'beta23', 'beta31', 'beta32', 'beta33', + 'gama11', 'gama12', 'gama13', 'gama21', 'gama22', 'gama23', 'gama31', 'gama32', 'gama33', +] + +TW_BASE_COLUMNS = [ + 's', + 'beta11', 'beta22', 'beta33', + 'alfa11', 'alfa22', 'alfa33', + 'gama11', 'gama22', 'gama33', + 'x', 'px', 'y', 'py', 't', 'pt', + 'dx', 'dy', 'dpx', 'dpy', + 'mu1', 'mu2', 'mu3', +] + +OPTFUN_QUANTITIES = [ + 'beta11', 'beta22', 'alfa11', 'alfa22', 'gama11', 'gama22', 'mu1', 'mu2', + 'dx', 'dy', 'dpx', 'dpy', +] + +CHROM_COLUMNS = [ + 'dmu1', 'dmu2', 'dmu3', + 'Dx', 'Dpx', 'Dy', 'Dpy', + 'ddx', 'ddpx', 'ddy', 'ddpy', + 'wx', 'wy', 'wxp', 'wyp', +] + +COUPLING_COLUMNS = [ + 'alfa12', 'alfa13', 'alfa21', 'alfa23', 'alfa31', 'alfa32', + 'beta12', 'beta13', 'beta21', 'beta23', 'beta31', 'beta32', + 'gama12', 'gama13', 'gama21', 'gama23', 'gama31', 'gama32', + 'f1001', 'f1010', 'r11', 'r12', 'r21', 'r22', +] + +PART_COORDS = ['x', 'px', 'y', 'py', 't', 'pt'] + +BETA0_QUANTITIES = [ + 'beta11', 'beta22', 'alfa11', 'alfa22', 'dx', 'dpx', 'dy', 'dpy', +] + +TPSA_ALLOWED_TARGETS = { + 'beta11', 'beta22', 'alfa11', 'alfa22', 'mu1', 'mu2', + 'dx', 'dpx', 'dy', 'dpy', + 'x', 'px', 'y', 'py', 't', 'pt', +} +# fmt: on + +# No leading underscore: pymadng refuses to retrieve names it deems private +XSUITE_MADNG_ENV_NAME = 'xsuite_matching_env' -TW_BASE_COLUMNS = ['s', 'beta11', 'beta22', 'beta33', 'alfa11', 'alfa22', 'alfa33', - 'gama11', 'gama22', 'gama33', 'x', 'px', 'y', 'py', 't', 'pt', - 'dx', 'dy', 'dpx', 'dpy', 'mu1', 'mu2', 'mu3'] +# Prelude for MAD-NG commands, giving access to the table in which xsuite keeps +# its data (the sequence, the damaps, ...) under the local name ``env``. +MNG_ENV_PRELUDE = f'local env = {XSUITE_MADNG_ENV_NAME}\n' -OPTFUN_QUANTITIES = ['beta11', 'beta22', 'alfa11', 'alfa22', 'gama11', 'gama22', - 'dx', 'dy', 'dpx', 'dpy', 'mu1', 'mu2'] -CHROM_COLUMNS = ['dmu1', 'dmu2', 'dmu3', 'Dx', 'Dpx', 'Dy', - 'Dpy', 'ddx', 'ddpx', 'ddy', 'ddpy', 'wx', 'wy', 'phix', 'phiy'] +def _ng_run(mng: Any, script: str, *payload: Any) -> Any: + """Run ``script`` in MAD-NG with the xsuite environment bound to ``env``. -COUPLING_COLUMNS = ['alfa12', 'alfa13', 'alfa21', 'alfa23', 'alfa31', 'alfa32', - 'beta12', 'beta13', 'beta21', 'beta23', 'beta31', 'beta32', - 'gama12', 'gama13', 'gama21', 'gama23', 'gama31', 'gama32', - 'f1001', 'f1010', 'r11', 'r12', 'r21', 'r22'] + The values in ``payload`` are sent on pymadng's data channel in the order + given, to be read by the ``py:recv()`` calls of the script. Passing them + this way rather than interpolating them keeps names and floating point + values off the MAD-NG source, where they would have to be quoted and + rounded. + """ + mng.send(MNG_ENV_PRELUDE + script) + for value in payload: + mng.send(value) + return mng -TPSA_ALLOWED_TARGETS = { 'beta11', 'beta22', 'alfa11', 'alfa22', 'dx', 'dpx', 'dy', 'dpy', - 'mu1', 'mu2', 'x', 'px', 'y', 'py', 't', 'pt' } -XSUITE_MADNG_ENV_NAME = "_xsuite_matching_env" +def _ensure_madng_model(line: xt.Line) -> Any: + """Return the MAD-NG model attached to ``line``, creating it if needed.""" + if not hasattr(line.tracker, '_madng'): + line.build_madng_model() + return line.tracker._madng + + +def _normal_form_columns(values: Sequence[Any]) -> dict[str, Any]: + """Map the normal-form response array to its public column names.""" + #fmt: off + names = ( + 'q1','q2', 'dq1', 'dq2', + 'd2q1', 'd2q2', 'd3q1', 'd3q2', 'd4q1', 'd4q2', 'd5q1', 'd5q2', + 'dqxdjx', 'dqydjy', 'dqxdjy', 'dqydjx', + ) + #fmt: on + result = dict(zip(names, values)) + for name in ('dqxdjx', 'dqydjy', 'dqxdjy', 'dqydjx'): + result[name] *= 2.0 + return result + +# A variable being matched has been turned into a (c)tpsa, of which only the +# constant part may be updated, or the TPSA would be destroyed. +_LUA_SET_VAR = """ + local name, value = py:recv(), py:recv() + local var = MADX[name] + if MAD.typeid.is_tpsa(var) or MAD.typeid.is_ctpsa(var) then + var:set0(value) + else + MADX[name] = value + end + """ -def _lua_list(strings): - """Convert Python list → Lua { 'a', 'b', 'c' }.""" - return "{ " + ", ".join(f"'{s}'" for s in strings) + " }" class MadngVars: + """Expose MAD-NG variables through Xsuite's variable-update mechanism.""" - def __init__(self, mad): + def __init__(self, mad: Any) -> None: self.mad = mad - def __setitem__(self, key, value): - # Check for key if it's a ctpsa or tpsa - var = f"MADX['{key.replace('.', '_')}']" - is_tpsa = self.mad.send(f"py:send(MAD.typeid.is_tpsa({var}) or MAD.typeid.is_ctpsa({var}))").recv() - if is_tpsa: - self.mad.send(f"{var}:set0(py:recv())").send(value) - else: - self.mad[var] = value - + def __setitem__(self, key: str, value: Any) -> None: + # Only values are propagated; deferred expressions would need MADX's + # own environment, opened with ``MADX:open_env()``. + _ng_run(self.mad, _LUA_SET_VAR, key.replace('.', '_'), value) - #Expressions still to be handled, could use the following: - # mng.send( - # MADX:open_env() - # a = 3 - # b =\ 3 * a - # c =\ 4 * a - # MADX:close_env() - # ''') -def build_madng_model(line, sequence_name='seq', **kwargs): +def build_madng_model(line: xt.Line, sequence_name: str = 'seq', **kwargs: Any) -> Any: """ Build and attach the MAD-NG model associated with this line. @@ -108,7 +183,11 @@ def build_madng_model(line, sequence_name='seq', **kwargs): model : object Built MAD-NG model. """ - print('Building MAD-NG model for line', line.name, 'with sequence name', sequence_name) + # Printed rather than logged, as the rest of xtrack reports progress this way + print( + f'Building MAD-NG model for line {line.name} ' + f'with sequence name {sequence_name}' + ) if line.tracker is None: line.build_tracker() mng = line.to_madng(sequence_name=sequence_name, **kwargs) @@ -118,86 +197,228 @@ def build_madng_model(line, sequence_name='seq', **kwargs): line.vars.vars_to_update.add(line.tracker._madng_vars) return mng -def discard_madng_model(line): + +def discard_madng_model(line: xt.Line) -> None: + """Remove the MAD-NG model and variable hook attached to ``line``.""" + line.tracker._madng = None + line.vars.vars_to_update.remove(line.tracker._madng_vars) + + +def regen_madng_model(line: xt.Line) -> None: + """Discard and rebuild the MAD-NG model associated with ``line``.""" + discard_madng_model(line) + build_madng_model(line) + + +def to_ng_name(name: str) -> str: + """Return the MAD-NG name of the quantity ``name``. + + A trailing ``_ng`` marks a name that is a MAD-NG one already. """ - Discard the attached MAD-NG model for this line. + return name[:-3] if name.endswith('_ng') else XS_NG_MAP[name] - Returns - ------- - None - Removes the current MAD-NG model association. + +def to_ng_target(xs_qty: str) -> str: + """Return the MAD-NG name of a target quantity, rejecting unsupported ones.""" + qty = to_ng_name(xs_qty) + if qty not in TPSA_ALLOWED_TARGETS: + raise ValueError( + f"Target quantity '{xs_qty}' not allowed with TPSA matching." + ) + return qty + + +@dataclass +class _TwissRange: + """The element range of a MAD-NG Twiss, and how to trim its result. + + MAD-NG brackets the requested range with marker rows, and how many it adds + depends on how the range is expressed, so the trimming belongs with the + range itself rather than at the call site. """ - line.tracker._madng = None - line.tracker.vars_to_update.remove(line.tracker._madng_vars) - return -def regen_madng_model(line): + start: str | None + end: str | None + i_start: int + i_end: int + element_names: tuple[str, ...] + # Index at which a wrap-around range folds back onto the start of the line + wrap_idx: int | None + + @classmethod + def from_line( + cls, line: xt.Line, start: str | None, end: str | None + ) -> _TwissRange: + names = line._element_names_unique + i_start = names.index(start) if start is not None else 0 + i_end = names.index(end) if end is not None else len(names) - 1 + wrap_idx = None + if i_start > i_end > 1: + wrap_idx = len(line.element_names) - list(line.element_names).index(start) + return cls(start, end, i_start, i_end, names, wrap_idx) + + @property + def is_partial(self) -> bool: + """Whether a sub-range of the line was requested.""" + return self.start is not None and self.end is not None + + @property + def marker_nums(self) -> int: + """Number of extra marker rows MAD-NG adds for a wrap-around range.""" + return 2 if self.i_start > self.i_end else 0 + + def selected_names(self) -> np.ndarray: + """Element names covered by the range, with Xsuite's end-point marker.""" + names = self.element_names + if self.i_start > self.i_end: + selected = names[self.i_start :] + names[: self.i_end + 1] + else: + selected = names[self.i_start : self.i_end + 1] + return np.array(selected + ('_end_point',)) + + def trim(self, data: Any) -> np.ndarray: + """Drop the MAD-NG marker rows from one returned column.""" + data = np.atleast_1d(np.squeeze(data)) + if not self.is_partial: + return data[:-1] + if self.wrap_idx is not None: + return np.concatenate( + (data[0:1], data[0 : self.wrap_idx], data[self.wrap_idx + 2 :]) + ) + if self.marker_nums: + return np.concatenate((data[0:1], data[: -self.marker_nums])) + return np.concatenate((data[0:1], data)) + + +_LUA_TWISS = """ + local config, columns = py:recv(), py:recv() + config.sequence = env.sequence + + -- Initial conditions given as beta0 values rather than as a map. The key + -- is cleared again, as twiss only accepts the options it knows. + if config.beta0 then + config.X0 = MAD.beta0(config.beta0) + config.beta0 = nil + end + if config.trkrdt then + -- The map below fixes the order, so any map definition must give way + config.mapdef = nil + config.X0 = MAD.damap {nv=6, mo=4} + config.info = 2 + config.saverdt = true + config.coupling = true + config.chrom = true + end + + local mtbl = twiss(config) + for _, column in ipairs(columns) do py:send(mtbl[column], true) end """ - Regenerate the MAD-NG model associated with this line. - Returns - ------- - None - Rebuilds the MAD-NG model association. +_LUA_NORMAL_FORM = """ + local config = py:recv() + config.sequence = env.sequence + local _, mytrkflow = MAD.track(config) + + local normal in MAD.gphys -- like "from MAD.gphys import normal" + -- anh stands for anharmonicity + local nf = normal(mytrkflow[1]):analyse('anh') + + last_nf = nf + normal_forms_to_send = { + nf:q1{1}, -- qx from the normal form (fractional part) + nf:q2{1}, -- qy + nf:dq1{1}, -- dqx / d delta + nf:dq2{1}, -- dqy / d delta + nf:dq1{2}, -- d2 qx / d delta2 + nf:dq2{2}, -- d2 qy / d delta2 + nf:dq1{3}, -- d3 qx / d delta3 + nf:dq2{3}, -- d3 qy / d delta3 + nf:dq1{4}, -- d4 qx / d delta4 + nf:dq2{4}, -- d4 qy / d delta4 + nf:dq1{5}, -- d5 qx / d delta5 + nf:dq2{5}, -- d5 qy / d delta5 + nf:anhx{1, 0}, -- dqx / d(2 jx) + nf:anhy{0, 1}, -- dqy / d(2 jy) + nf:anhx{0, 1}, -- dqx / d(2 jy) + nf:anhy{1, 0}, -- dqy / d(2 jx) + } + py:send(normal_forms_to_send) """ - discard_madng_model(line) - build_madng_model(line) - return - -def _build_column_send_script(columns): - assert len(columns) > 0 - mng_columns_to_send = ["mtbl." + col for col in columns] - send_cmd = f''' - -- send columns to Python - columns = {{{", ".join(mng_columns_to_send)}}} - py:send(columns, true) - ''' - return send_cmd - -def _build_rdt_script(mng_sequence_name, rdts, columns): - assert len(rdts) > 0 - rdt_cmd = 'local rdts = {"' + '", "'.join(rdts) + '"}' - send_cmd = _build_column_send_script(columns) - # Create damap and X0, then twiss with rdts - script = f''' - local damap in MAD - {rdt_cmd} - - -- create phase-space damap at 4th order - local X0 = damap {{nv=6, mo=4}} - - -- twiss with RDTs - local mtbl = twiss {{ sequence={mng_sequence_name}, X0=X0, trkrdt=rdts, info=2, saverdt=true, coupling=true, chrom=true }} - - {send_cmd} - ''' - return script - -def _build_beta0_block_string(tw_kwargs): - flag_init = False - beta0_dict = {} - for k in tw_kwargs.keys(): - if k in BETA0_COLUMNS: - beta0_dict[k] = tw_kwargs[k] - flag_init = True - elif k in XS_NG_MAP: - beta0_dict[XS_NG_MAP[k]] = tw_kwargs[k] - flag_init = True - - if flag_init: - # Construct beta0 string - beta0_str = 'X0 = beta0 {' - for k, v in beta0_dict.items(): - beta0_str += f'{k} = {v}, ' - beta0_str = beta0_str[:-2] + '}, ' - else: - beta0_str = '' - return beta0_str -def _tw_ng(line, rdts=(), normal_form=False, - mapdef_twiss=2, mapdef_normal_form=4, - nslice=3, xsuite_tw=True, X0=None, compute_chromatic_properties=False, - coupling_edw_teng=False, method=4, **tw_kwargs): + +def _twiss_config( + method: int, + nslice: int, + mapdef: int, + coupling: bool, + chromatic: bool, + rng: _TwissRange, + X0: Any, + beta0_data: Mapping[str, Any] | None, + rdts: Sequence[str], +) -> dict[str, Any]: + """Build the option table for the MAD-NG ``twiss`` command.""" + config: dict[str, Any] = { + 'method': method, + 'implicit': True, + 'nslice': nslice, + 'misalign': True, + 'coupling': coupling, + 'chrom': chromatic, + } + if X0 is not None: + # A reference to a map that already lives in MAD-NG, whose order is + # fixed, so the map definition would have nothing left to say + config['X0'] = X0 + else: + # MAD-NG builds the map itself here, and the order is ours to choose + config['mapdef'] = mapdef + if beta0_data is not None: + # Turned into a map by MAD-NG, which knows the beta0 constructor + config['beta0'] = beta0_data + if rdts: + # An empty table would be truthy in MAD-NG, so only set it when needed + config['trkrdt'] = list(rdts) + if rng.is_partial: + config['range'] = f'{rng.start}/{rng.end}' + return config + + +def _add_chromatic_columns(tw: xt.TwissTable) -> None: + """Replace MAD-NG's chromatic amplitude and phase by their components.""" + for plane in ('x', 'y'): + wave = tw[f'w{plane}_ng'] * np.exp(1j * 2 * np.pi * tw[f'w{plane}p_ng']) + tw[f'a{plane}_ng'] = np.imag(wave) + tw[f'b{plane}_ng'] = np.real(wave) + del tw[f'w{plane}p_ng'] + + +def _add_normal_form_columns( + mng: Any, tw: xt.TwissTable, method: int, mapdef: int, nslice: int +) -> None: + """Compute the normal-form quantities and add them to ``tw``.""" + _ng_run( + mng, _LUA_NORMAL_FORM, + {'method': method, 'mapdef': mapdef, 'nslice': nslice}, + ) + for nn, val in _normal_form_columns(mng.recv('normal_forms_to_send')).items(): + tw[f'{nn}_nf_ng'] = val + + +def _tw_ng( + line: xt.Line, + rdts: Sequence[str] = (), + normal_form: bool = False, + mapdef_twiss: int = 2, + mapdef_normal_form: int = 4, + nslice: int = 3, + xsuite_tw: bool = True, + X0: Any = None, + compute_chromatic_properties: bool = False, + coupling_edw_teng: bool = False, + method: int = 4, + **tw_kwargs: Any, +) -> xt.TwissTable: """ Run a Twiss calculation using the MAD-NG model. @@ -210,7 +431,8 @@ def _tw_ng(line, rdts=(), normal_form=False, normal_form : bool, optional If ``True``, also compute normal-form quantities. mapdef_twiss : int, optional - Map order used for the MAD-NG Twiss computation. + Map order used for the MAD-NG Twiss computation. Only has an effect + when ``X0`` is not given, as a map carries its own order. mapdef_normal_form : int, optional Map order used for the MAD-NG normal-form computation. nslice : int, optional @@ -230,219 +452,121 @@ def _tw_ng(line, rdts=(), normal_form=False, Twiss table with MAD-NG columns. """ - _action = ActionTwissMadng(line, { - "rdts": rdts, - "normal_form": normal_form, - "mapdef_twiss": mapdef_twiss, - "mapdef_normal_form": mapdef_normal_form, - "nslice": nslice, - **tw_kwargs - }) + _action = ActionTwissMadng( + line, + { + 'rdts': rdts, + 'normal_form': normal_form, + 'mapdef_twiss': mapdef_twiss, + 'mapdef_normal_form': mapdef_normal_form, + 'nslice': nslice, + **tw_kwargs, + }, + ) - if not hasattr(line.tracker, '_madng'): - line.build_madng_model() - mng = line.tracker._madng + mng = _ensure_madng_model(line) start = tw_kwargs.get('start', None) end = tw_kwargs.get('end', None) init = tw_kwargs.get('init', None) + beta0_data = None if X0 is None: if init is not None and isinstance(init, xt.TwissTable): raise NotImplementedError('TwissTable as init not implemented.') - X0_str = _build_beta0_block_string(tw_kwargs) - else: - X0_str = f'X0 = {X0}, ' + beta0_data = { + ng_key: value + for key, value in tw_kwargs.items() + if (ng_key := XS_NG_MAP.get(key, key)) in BETA0_COLUMNS + } or None if (start is None) != (end is None): raise ValueError('Start and end must be specified together.') - if start is not None and end is not None and not X0_str: - raise ValueError('Initial conditions must be specified when start and end are given.') + rng = _TwissRange.from_line(line, start, end) - full_twiss_str = '' + if rng.is_partial and X0 is None and beta0_data is None: + raise ValueError( + 'Initial conditions must be specified when start and end are given.' + ) tw_columns = TW_BASE_COLUMNS.copy() - - full_twiss_str = f"implicit=true, nslice={nslice}, misalign=true, coupling={str(coupling_edw_teng).lower()}, chrom={str(compute_chromatic_properties).lower()}" - if coupling_edw_teng: tw_columns += COUPLING_COLUMNS if compute_chromatic_properties: tw_columns += CHROM_COLUMNS columns = tw_columns + list(rdts) - send_cmd = _build_column_send_script(columns) - - if len(rdts) > 0: - mng_script = _build_rdt_script(mng._sequence_name, rdts, columns) - else: - range_str = '' - - if start is not None and end is not None: - normal_form = False - # Range Twiss - range_str = f"range = '{start}/{end}', " - - mng_script = (''' - -- twiss without RDTs - local mtbl = twiss {sequence=''' f'{mng._sequence_name}, method={method},\ - {X0_str} {range_str} {full_twiss_str}' - '''} - ''' - + send_cmd) + if rng.is_partial and not rdts: + normal_form = False - mng.send(mng_script) + config = _twiss_config( + method=method, + nslice=nslice, + mapdef=mapdef_twiss, + coupling=coupling_edw_teng, + chromatic=compute_chromatic_properties, + rng=rng, + X0=X0, + beta0_data=beta0_data, + rdts=rdts, + ) - out = mng.recv('columns') - out_dct = {k: v for k, v in zip(columns, out)} - - # Add to table - names = line._element_names_unique - i_start = names.index(start) if start is not None else 0 - i_end = names.index(end) if end is not None else len(names) - 1 - marker_nums = 2 if i_start > i_end else 0 # MAD-NG wrap-around markers + _ng_run(mng, _LUA_TWISS, config, list(columns)) + out_dct = {c: mng.recv() for c in columns} if xsuite_tw: - xs_tw_kwargs = { - NG_XS_MAP.get(k, k): v for k, v in tw_kwargs.items() - } - + xs_tw_kwargs = {NG_XS_MAP.get(k, k): v for k, v in tw_kwargs.items()} tw = line.twiss(method='4d', reverse=False, **xs_tw_kwargs) - - if not xsuite_tw: - # Handle wrap-around range - if i_start > i_end: - name_co = np.array(names[i_start:] + names[:i_end + 1] + ('_end_point',)) - else: - name_co = np.array(names[i_start:i_end + 1] + ('_end_point',)) - - tw = xt.TwissTable({"name": name_co}) + else: + tw = xt.TwissTable({'name': rng.selected_names()}) tw._action = _action - # Consistency check - if start is None and end is None: - assert len(out[0]) == len(tw) + 1 + first_col = np.atleast_1d(np.squeeze(out_dct[columns[0]])) + if not rng.is_partial: + assert len(first_col) == len(tw) + 1 else: - assert len(out[0]) == len(tw) + marker_nums - 1 - - if start is None and end is None: - mode = "full" - elif i_start > i_end and i_end > 1: - mode = "wrap" - end_idx = len(line.element_names) - list(line.element_names).index(start) - elif marker_nums > 0: - mode = "marker" - else: - mode = "range" - - def _process_data(data): - data = np.atleast_1d(np.squeeze(data)) - if mode == "full": - return data[:-1] - elif mode == "wrap": - return np.concatenate((data[0:1], data[0:end_idx], data[end_idx + 2:])) - elif mode == "marker": - return np.concatenate((data[0:1], data[:-marker_nums])) - elif mode == "range": - return np.concatenate((data[0:1], data)) - else: - raise ValueError(f"Unexpected mode: {mode}") + assert len(first_col) == len(tw) + rng.marker_nums - 1 - # enforce marker for nn in tw_columns: - tw[f"{nn}_ng"] = _process_data(out_dct[nn]) + tw[f'{nn}_ng'] = rng.trim(out_dct[nn]) for nn in rdts: tw[nn] = np.atleast_1d(np.squeeze(out_dct[nn]))[:-1] if compute_chromatic_properties: - temp_x = tw.wx_ng * np.exp(1j*2*np.pi*tw.phix_ng) - tw['ax_ng'] = np.imag(temp_x) - tw['bx_ng'] = np.real(temp_x) - temp_y = tw.wy_ng * np.exp(1j*2*np.pi*tw.phiy_ng) - tw['ay_ng'] = np.imag(temp_y) - tw['by_ng'] = np.real(temp_y) - del tw['phix_ng'] - del tw['phiy_ng'] + _add_chromatic_columns(tw) if normal_form: - mng_script_nf = ( - ''' - local track in MAD -- like "from MAD import track" - local mytrktable, mytrkflow = MAD.track{sequence=''' - f'{mng._sequence_name}, method={method},mapdef={mapdef_normal_form}, nslice={nslice}' - '''} - - local normal in MAD.gphys -- like "from MAD.gphys import normal" - local my_norm_for = normal(mytrkflow[1]):analyse('anh') -- anh stands for anharmonicity - - local nf = my_norm_for - last_nf = my_norm_for - normal_forms_to_send = { - nf:q1{1}, -- qx from the normal form (fractional part) - nf:q2{1}, -- qy - nf:dq1{1}, -- dqx / d delta - nf:dq2{1}, -- dqy / d delta - nf:dq1{2}, -- d2 qx / d delta2 - nf:dq2{2}, -- d2 qy / d delta2 - nf:dq1{3}, -- d3 qx / d delta3 - nf:dq2{3}, -- d3 qy / d delta3 - nf:dq1{4}, -- d4 qx / d delta4 - nf:dq2{4}, -- d4 qy / d delta4 - nf:dq1{5}, -- d5 qx / d delta5 - nf:dq2{5}, -- d5 qy / d delta5 - nf:anhx{1, 0}, -- dqx / d(2 jx) - nf:anhy{0, 1}, -- dqy / d(2 jy) - nf:anhx{0, 1}, -- dqx / d(2 jy) - nf:anhy{1, 0}, -- dqy / d(2 jx) - } - py:send(normal_forms_to_send) - ''') - mng.send(mng_script_nf) - out_nf = mng.recv('normal_forms_to_send') - - dct_nf = dict( - q1 = out_nf[0], - q2 = out_nf[1], - dq1 = out_nf[2], - dq2 = out_nf[3], - d2q1 = out_nf[4], - d2q2 = out_nf[5], - d3q1 = out_nf[6], - d3q2 = out_nf[7], - d4q1 = out_nf[8], - d4q2 = out_nf[9], - d5q1 = out_nf[10], - d5q2 = out_nf[11], - dqxdjx = out_nf[12]*2., - dqydjy = out_nf[13]*2., - dqxdjy = out_nf[14]*2., - dqydjx = out_nf[15]*2., - ) - for nn in dct_nf: - tw[nn+'_nf_ng'] = dct_nf[nn] + _add_normal_form_columns(mng, tw, method, mapdef_normal_form, nslice) return tw -def madng_get_init(line, at): - if not hasattr(line.tracker, '_madng'): - line.build_madng_model() - mng = line.tracker._madng - if at == xt.START: - at = "1" - else: - at = f"'{at}'" - mng.send(f""" - local observed in MAD.element.flags - {mng._sequence_name}:select(observed, {{list = {{{at}}}}}) - twpart, mf = twiss {{sequence = {mng._sequence_name}, observe = 1, savemap = true, info = 2}} - {XSUITE_MADNG_ENV_NAME}.X0 = twpart[{at}].__map - """) - return f"{XSUITE_MADNG_ENV_NAME}.X0" +_LUA_GET_INIT = """ + local at = py:recv() + env.sequence:select(MAD.element.flags.observed, {list = {at}}) + local twpart = twiss { + sequence = env.sequence, observe = 1, savemap = true, info = 2 + } + env.X0 = twpart[at].__map + """ + -def _survey_ng(line): +def madng_get_init(line: xt.Line, at: Any) -> Any: + """Return a reference to the MAD-NG map at ``at``, computed with a Twiss. + + ``at`` is an element name, or ``xt.START`` for the start of the line. + """ + mng = _ensure_madng_model(line) + + # The location travels on the data channel: an element name needs no + # quoting there, and the start of the line is simply the index 1. + _ng_run(mng, _LUA_GET_INIT, 1 if at == xt.START else at) + return mng._env.X0 + + +def _survey_ng(line: xt.Line) -> SurveyTable: """ Run a survey using the MAD-NG model. @@ -450,69 +574,73 @@ def _survey_ng(line): Returns ------- - survey : xtrack.survey.SurveyTable + survey : SurveyTable Survey result produced by MAD-NG. """ - if not hasattr(line.tracker, '_madng'): - line.build_madng_model() - mng = line.tracker._madng - mng['srv'] = mng.survey(sequence=mng._sequence_name) + mng = _ensure_madng_model(line) + mng['srv'] = mng.survey(sequence=mng._sequence) survey_tab_keys = { - "x": "X", - "y": "Y", - "z": "Z", - "l": "length", - "kind": "element_type" + 'x': 'X', + 'y': 'Y', + 'z': 'Z', + 'l': 'length', + 'kind': 'element_type', } element_types = { - "drift": "Drift", - "sbend": "Bend", - "rbend": "RBend", - "quadrupole": "Quadrupole", - "sextupole": "Sextupole", - "octupole": "Octupole", - "multipole": "Multipole", - "kicker": "Kicker", # no coloring in survey plot - "rfcavity": "Cavity", - "marker": "Marker" + 'drift': 'Drift', + 'sbend': 'Bend', + 'rbend': 'RBend', + 'quadrupole': 'Quadrupole', + 'sextupole': 'Sextupole', + 'octupole': 'Octupole', + 'multipole': 'Multipole', + 'kicker': 'Kicker', # no coloring in survey plot + 'rfcavity': 'Cavity', + 'marker': 'Marker', } - # create SurveyTable from DataFrame survey_df = mng['srv'][0].to_df() survey_dict = survey_df.to_dict(orient='list') survey_dict = {k: np.array(v) for k, v in survey_dict.items()} - for k in survey_tab_keys.keys(): + for k, v in survey_tab_keys.items(): if k in survey_dict: - # Rename keys to match SurveyTable - survey_dict[survey_tab_keys[k]] = survey_dict[k] + survey_dict[v] = survey_dict[k] del survey_dict[k] - survey_dict['element_type'] = np.array([element_types.get(et, et) for et in survey_dict['element_type']]) + survey_dict['element_type'] = np.array( + [element_types.get(et, et) for et in survey_dict['element_type']] + ) - for i in survey_dict.keys(): + for i in survey_dict: # Interpretation of survey is shifted by 1 in MAD-NG vs. Xsuite if i in ['name', 'length', 'kind', 'element_type', 'angle', 'tilt']: survey_dict[i] = survey_dict[i][1:] else: survey_dict[i] = survey_dict[i][:-1] - survey_tab = xt.survey.SurveyTable(survey_dict) - return survey_tab + return SurveyTable(survey_dict) class ActionTwissMadng(Action): - def __init__(self, line, tw_kwargs={}, **kwargs): + """Matching action that evaluates Twiss parameters with MAD-NG.""" + + def __init__( + self, + line: xt.Line, + tw_kwargs: Mapping[str, Any] | None = None, + **kwargs: Any, + ) -> None: self.line = line - self.tw_kwargs = tw_kwargs + self.tw_kwargs = {} if tw_kwargs is None else dict(tw_kwargs) self.tw_kwargs.update(kwargs) - self._alredy_prepared = False + self._already_prepared = False self.X0 = None - def prepare(self, force=False): - - if self._alredy_prepared and not force: + def prepare(self, force: bool = False) -> None: + """Prepare the initial MAD-NG map used by the action.""" + if self._already_prepared and not force: return init = self.tw_kwargs.get('init', None) @@ -526,296 +654,462 @@ def prepare(self, force=False): assert isinstance(init, xt.TwissTable) self.X0 = madng_get_init(self.line, at=xt.START) - self._alredy_prepared = True - - def run(self): - return self.line.madng_twiss(xsuite_tw = False, X0=self.X0, **self.tw_kwargs) - -class ActionTwissMadngTPSA(Action): - def __init__(self, line, vary_names, targets = [], tw_kwargs={}, sum_rmat_tar=0, **kwargs): - self.line = line - self.vary_names = vary_names - self.targets = targets - self.optics_target_locations = None - self.optics_target_quantities = None - self.tw_kwargs = tw_kwargs - self.tw_kwargs.update(kwargs) - self.twiss_flag = None - self._already_prepared = False - self.match_rmat = False - self.match_opt = False - self.sum_rmat_tar = sum_rmat_tar - self.rmat_start_end_list = None - self.rmat_tags = None - self._last_res = None - self._needs_zeta_scale = [] - self._needs_delta_scale = [] - - def prepare(self, force=False): - """ - Prepare the MAD-NG TPSA matching environment. - This method sets up the MAD-NG environment for TPSA matching by - configuring the initial conditions, setting target locations, and quantities - based on the provided targets. - To achieve that, arrays and maps are created within MAD-NG to keep track of - the target locations, quantities and differential algebraic maps. - - Parameters - ---------- - force : bool, optional - If True, forces re-preparation even if already prepared. Default is False. - - Raises - ------ - ValueError - If the target quantity is not allowed with TPSA matching - or if start and end are provided without initial conditions. - """ - - if self._already_prepared and not force: - return - - # Collect initial conditions - init = self.tw_kwargs.get('init', None) - - if init is None: - init = self.line.madng_twiss(**self.tw_kwargs) - self.tw_kwargs.update({'init': init}) - - assert isinstance(init, xt.TwissTable) - - if not hasattr(self.line.tracker, '_madng'): - self.line.build_madng_model() - - self.mng = self.line.tracker._madng - - self.twiss_flag = any(isinstance(tar, (xt.TargetRelPhaseAdvance)) for tar in self.targets) - - # Collect initial coordinates - coord_assign_str = self._initialize_coordinates_str(init) - - # Process targets - - targets_map_str, xs_ng_target_map = self._process_targets(init) - - - # Lua script assembly - observables = [loc for loc in self.optics_target_locations] - - param_list_str = _lua_list(self.vary_names) - observables_str = _lua_list(observables) - optics_qty_str = _lua_list(list(self.optics_target_quantities)) - - init_cond_str = self._beta_block_str(init) - - rmat_array_str = '' - if self.match_rmat: - rmat_array_str = f"{XSUITE_MADNG_ENV_NAME}.rmat_map_arr = table.new({self.sum_rmat_tar}, 0)\n" + self._already_prepared = True - mng_init_str = r''' - ''' + XSUITE_MADNG_ENV_NAME + r''' = {} -- to avoid variable name clashes - local obs_flag = MAD.element.flags.observed + def run(self, allow_failure: bool = False) -> xt.TwissTable: + """Evaluate and return the requested MAD-NG Twiss table.""" + return self.line.madng_twiss(xsuite_tw=False, X0=self.X0, **self.tw_kwargs) + + +# The matrices are sent one by one: a Lua table would arrive as a reference +# rather than as its contents. +_LUA_RMATRICES = """ + for i, options in ipairs(py:recv()) do + options.sequence = env.sequence + options.X0 = env.empty_X0 + local _, flow = MAD.track(options) + env.rmat_map_arr[i] = flow[1] + py:send(flow[1]:get1()) + end + """ - local pts=''' + observables_str + r''' +# One row per target, one column per matching variable. Both counts come from +# the environment: the initial map carries one parameter per matching variable. +_LUA_JACOBIAN = """ + local varylen = env.init_X0_map:np() + local nv = env.init_X0_map:nv() + env.jac = MAD.matrix(#env.targets_arr, varylen) + + for i, target in ipairs(env.targets_arr) do + local map = nil + if target.optfun or target.orbit then + map = env.target_loc_map[target.loc] + elseif target.rmat then + map = env.rmat_map_arr[target.tag] + end - ''' + self.mng._sequence_name + r''':select(obs_flag, {list=pts}) + local monom = MAD.monomial(nv + varylen) -- BUILD MONOMIAL + for j = 1, map:np(), 1 do + local jac_idx = (i-1)*varylen + j - local params = ''' + param_list_str + r''' + -- Quantity which can be calculated with optfun + if target.optfun then + -- If loc_start (phase advance) is defined, we provide initial map + local a0 = target.loc_start + and env.target_loc_map[target.loc_start] + env.jac[jac_idx] = MAD.gphys.optfun( + map, target.qty .. "_", j, 1, a0) - local X0 = MAD.damap { - nv=6, -- number of variables - mo=2, -- max order of variables - np=#params, -- number of parameters - po=1, -- max order of parameters - pn=params, -- parameter names - } + -- Orbit Quantity + elseif target.orbit then + monom[nv + j] = 1 + env.jac[jac_idx] = map[target.orbit]:get(monom) + monom[nv + j] = 0 - -- Converting to TPSA (mutating type) - for _, v in ipairs(params) do - MADX[v] = MADX[v] + X0[v] + elseif target.rmat then + -- rmatrix terms are extracted directly from the damap + local ind_1 = tonumber(target.qty:sub(2,2)) + local ind_2 = tonumber(target.qty:sub(3,3)) + monom[nv + j] = 1 + monom[ind_2] = 1 + env.jac[jac_idx] = map[ind_1]:get(monom) + monom[nv + j] = 0 + monom[ind_2] = 0 end + end + end - ''' + init_cond_str + r''' - - local map1 = MAD.gphys.bet2map(B0, X0:copy()) - - ''' + coord_assign_str + r''' + py:send(env.jac) + """ - -- Maps target locations to damaps - -- e.g. { 'BPM1' = damap1, 'BPM2' = damap2, ... } - ''' + XSUITE_MADNG_ENV_NAME + r'''.target_loc_map = table.new(0, ''' + str(len(self.optics_target_locations)) + r''') +_LUA_CLEANUP = """ + for _, var_name in ipairs(py:recv()) do + MADX[var_name] = MADX[var_name]:get0() + end + """ - -- Maps rmat tags to rmat damaps - -- e.g. {rmat_damap1, rmat_damap2, ... } - ''' + rmat_array_str + r''' +_LUA_TRACK = """ + local operation, options = py:recv(), py:recv() + options.sequence = env.sequence + options.X0 = env.init_X0_map + env.trk = MAD[operation](options) + """ - -- Array of targets with additional info (location, quantity, orbit/optical function) - -- e.g. { { loc = 'BPM1', qty = 'beta11', optfun = true }, { loc = 'BPM2', qty = 'x', orbit = 1 }, ... } - ''' + XSUITE_MADNG_ENV_NAME + r'''.targets_arr = table.new(''' + str(len(self.targets)) + r''', 0) +# After a Track calculation the optical functions are not present in the table, +# so they are added as derived columns, named as the user (Xsuite) defined the +# targets. A Twiss already carries every quantity TPSA matching allows, so the +# test below skips them all and no flag is needed to tell the two apart. +_LUA_RESULT = """ + local trk = env.trk + + for _, tar in ipairs(env.tar_optics_qtys) do + if not trk[env.xs_ng_target_map[tar]] then + trk:addcol(tar, \\ri -> MAD.gphys.optfun( + trk[ri].__map, env.xs_ng_target_map[tar] .. '_')) + end + end - -- List (Array) of target optics/orbit quantities (suitable for MAD-NG), - -- e.g. {'beta11', 'alfa11', ...} - ''' + XSUITE_MADNG_ENV_NAME + r'''.tar_optics_qtys = ''' + optics_qty_str + r''' + -- Save damaps + for _, location in ipairs(py:recv()) do + env.target_loc_map[location] = trk[location].__map + end + py:send(trk) + """ - -- Initial map for tracking/twiss - ''' + XSUITE_MADNG_ENV_NAME + r'''.init_X0_map = map1 +# The values are received in the order in which they are sent, and are only +# kept in a local when used more than once. +_LUA_PREPARE_TPSA = """ + local locations = py:recv() -- target locations + env.sequence:select(MAD.element.flags.observed, {list = locations}) - -- Identity map - ''' + XSUITE_MADNG_ENV_NAME + r'''.empty_X0 = X0 + local params = py:recv() -- names of the matching variables - -- Defining targets array - ''' + targets_map_str + r''' + local X0 = MAD.damap { + nv=6, -- number of variables + mo=2, -- max order of variables + np=#params, -- number of parameters + po=1, -- max order of parameters + pn=params, -- parameter names + } - -- Mapping from xsuite quantity names to madng quantity names - ''' + xs_ng_target_map + r''' - ''' + -- Converting to TPSA (mutating type) + for _, v in ipairs(params) do + MADX[v] = MADX[v] + X0[v] + end + local map1 = MAD.gphys.bet2map(MAD.beta0(py:recv()), X0:copy()) - self.mng.send(mng_init_str) + -- Initial orbit, as {name, value} pairs + for _, coordinate in ipairs(py:recv()) do + map1[coordinate[1]]:set0(coordinate[2]) + end - self._already_prepared = True + -- Maps target locations to damaps + -- e.g. { 'BPM1' = damap1, 'BPM2' = damap2, ... } + env.target_loc_map = table.new(0, #locations) - def _process_targets(self, init): - self.optics_target_locations = set() - self.optics_target_quantities = set() - self.rmat_start_end_list = [None] * self.sum_rmat_tar - self.rmat_tags = [] + -- Maps rmat tags to rmat damaps + -- e.g. {rmat_damap1, rmat_damap2, ... } + env.rmat_map_arr = table.new(py:recv(), 0) - start = self.tw_kwargs.get('start', None) - end = self.tw_kwargs.get('end', None) + -- Initial map for tracking/twiss + env.init_X0_map = map1 - targets_map_str = '' - xs_ng_target_map = XSUITE_MADNG_ENV_NAME + '.xs_ng_target_map = {}\n' + -- Identity map + env.empty_X0 = X0 + """ - for i, target in enumerate(self.targets): - if isinstance(target.tar, tuple): - self.match_opt = True - qty_orig = target.tar[0] - loc = target.tar[1] - qty = qty_orig[:-3] if qty_orig.endswith('_ng') else XS_NG_MAP[qty_orig] - assert qty in TPSA_ALLOWED_TARGETS, f"Target quantity '{qty_orig}' not allowed with TPSA matching." +@dataclass +class _Target: + """One matching target, in the form the MAD-NG environment expects. - self.optics_target_locations.add(loc) + Exactly one of ``optfun``, ``orbit`` and ``rmat`` says how the quantity is + evaluated: from an optical function, from a coordinate of the damap, or + from a transfer matrix term. + """ - aux = '' - if qty in OPTFUN_QUANTITIES: - aux = 'optfun = true' - elif qty in (part_coords := ['x', 'px', 'y', 'py', 't', 'pt']): - aux = f'orbit = {part_coords.index(qty) + 1}' + loc: str + qty: str # MAD-NG name of the quantity + xs_qty: str # Name the user asked for, which may be an Xsuite one + loc_start: str | None = None + optfun: bool = False + orbit: int | None = None # 1-based index of the coordinate in the damap + rmat: bool = False + tag: int | None = None + rtag: str | None = None # Xsuite's '_r' label for the term + + def as_ng(self) -> dict[str, Any]: + """Return the Lua table for this target; omitted keys read as nil.""" + data: dict[str, Any] = {'loc': self.loc, 'qty': self.qty} + if self.loc_start is not None: + data['loc_start'] = self.loc_start + if self.optfun: + data['optfun'] = True + if self.orbit is not None: + data['orbit'] = self.orbit + if self.rmat: + data['rmat'] = True + # Xsuite numbers the transfer maps from zero, MAD-NG arrays from one + data['tag'] = self.tag + 1 if self.tag is not None else None + return data + + +def _target_range( + target: Any, init: xt.TwissTable, start: str | None, end: str | None +) -> tuple[str, str]: + """Resolve the endpoints of a target, expanding Xsuite's placeholders.""" + if target.start != '__ele_start__': + loc_start = target.start + elif start is not None: + loc_start = start + else: + loc_start = init.name[0] - if qty_orig == 'zeta': - self._needs_zeta_scale.append(i) - elif qty_orig == 'delta': - self._needs_delta_scale.append(i) + if target.end != '__ele_stop__': + loc_end = target.end + elif end is not None: + loc_end = end + else: + loc_end = init.name[-2] + + return loc_start, loc_end + + +def _build_targets( + targets: Sequence[Any], + init: xt.TwissTable, + start: str | None, + end: str | None, +) -> list[_Target]: + """Translate the Xsuite targets into their MAD-NG counterparts.""" + built = [] + for target in targets: + if isinstance(target.tar, tuple): + xs_qty, loc = target.tar + qty = to_ng_target(xs_qty) + built.append(_Target( + loc=loc, + qty=qty, + xs_qty=xs_qty, + optfun=qty in OPTFUN_QUANTITIES, + orbit=( + PART_COORDS.index(qty) + 1 + if qty not in OPTFUN_QUANTITIES and qty in PART_COORDS + else None + ), + )) + + elif hasattr(target, 'start') and hasattr(target, 'end'): + loc_start, loc_end = _target_range(target, init, start, end) + + if isinstance(target, xt.TargetRelPhaseAdvance): + built.append(_Target( + loc=loc_end, + qty=to_ng_target(target.var), + xs_qty=target.var, + loc_start=loc_start, + optfun=True, + )) + + elif isinstance(target, xt.TargetRmatrixTerm): + built.append(_Target( + loc=loc_end, + qty=target.term, + xs_qty=target.term, + loc_start=loc_start, + rmat=True, + tag=int(target.rtag.split('_')[0]), + rtag=target.rtag, + )) - # set string for quantity mapping + loc to save in madng - targets_map_str += f"{XSUITE_MADNG_ENV_NAME}.targets_arr[{i+1}] = {{ loc = '{loc}', qty = '{qty}', {aux} }}\n" - xs_ng_target_map += f"{XSUITE_MADNG_ENV_NAME}.xs_ng_target_map['{qty_orig}'] = '{qty}'\n" + else: + raise NotImplementedError( + f'Target of type {type(target)} not implemented for ' + 'MAD-NG TPSA matching.' + ) - self.optics_target_quantities.add(qty_orig) + else: + raise NotImplementedError( + f'Target of type {type(target)} not implemented for ' + 'MAD-NG TPSA matching.' + ) + return built - elif hasattr(target, "start") and hasattr(target, "end"): - if target.start != "__ele_start__": - loc_start = target.start - elif start is not None: - loc_start = start - else: - loc_start = init.name[0] - if target.end != "__ele_stop__": - loc_end = target.end - elif end is not None: - loc_end = end - else: - loc_end = init.name[-2] - if isinstance(target, xt.TargetRelPhaseAdvance): - self.match_opt = True - qty_orig = target.var - qty = qty_orig[:-3] if qty_orig.endswith('_ng') else XS_NG_MAP[qty_orig] +class ActionTwissMadngTPSA(Action): + """Matching action using MAD-NG TPSA maps for optics sensitivities.""" + + def __init__( + self, + line: xt.Line, + vary_names: Sequence[str], + targets: Sequence[Any] | None = None, + tw_kwargs: Mapping[str, Any] | None = None, + sum_rmat_tar: int = 0, + **kwargs: Any, + ) -> None: + self.line = line + self.vary_names = vary_names + self.targets = [] if targets is None else targets + self.mng: Any = None + self.optics_target_locations: list[str] = [] + self.optics_target_quantities: set[str] = set() + self.tw_kwargs = {} if tw_kwargs is None else dict(tw_kwargs) + self.tw_kwargs.update(kwargs) + self._already_prepared = False + self.sum_rmat_tar = sum_rmat_tar + self.rmat_start_end_list: list[tuple[str, str]] = [] + self.rmat_tags: list[str] = [] + self._last_res: Any = None + self._needs_zeta_scale: list[int] = [] + self._needs_delta_scale: list[int] = [] + + @property + def twiss_flag(self) -> bool: + """Whether a Twiss is needed, as a Track gives no phase advance.""" + return any( + isinstance(tar, xt.TargetRelPhaseAdvance) for tar in self.targets + ) - assert qty in TPSA_ALLOWED_TARGETS, f"Target quantity '{target.var}' not allowed with TPSA matching." + @property + def match_rmat(self) -> bool: + """Whether any target is a transfer matrix term.""" + return any( + isinstance(tar, xt.TargetRmatrixTerm) for tar in self.targets + ) - self.optics_target_locations.add(loc_start) - self.optics_target_locations.add(loc_end) - self.optics_target_quantities.add(qty_orig) + def prepare(self, force: bool = False) -> None: + """ + Prepare the MAD-NG TPSA matching environment. + This method sets up the MAD-NG environment for TPSA matching by + configuring the initial conditions, setting target locations, and quantities + based on the provided targets. + To achieve that, arrays and maps are created within MAD-NG to keep track of + the target locations, quantities and differential algebraic maps. - targets_map_str += f"{XSUITE_MADNG_ENV_NAME}.targets_arr[{i+1}] = {{ loc = '{loc_end}', qty = '{qty}', loc_start = '{loc_start}', optfun = true }}\n" - xs_ng_target_map += f"{XSUITE_MADNG_ENV_NAME}.xs_ng_target_map['{target.var}'] = '{qty}'\n" + Parameters + ---------- + force : bool, optional + If True, forces re-preparation even if already prepared. Default is False. - elif isinstance(target, xt.TargetRmatrixTerm): - self.match_rmat = True - qty = target.term - tag = target.rtag.split('_')[0] - idx = int(tag) + Raises + ------ + ValueError + If the target quantity is not allowed with TPSA matching + or if start and end are provided without initial conditions. + """ - self.rmat_start_end_list[idx] = (loc_start, loc_end) - self.rmat_tags.append(target.rtag) + if self._already_prepared and not force: + return - targets_map_str += f"{XSUITE_MADNG_ENV_NAME}.targets_arr[{i+1}] = {{ loc = '{loc_end}', qty = '{qty}', loc_start = '{loc_start}', rmat = true, tag = {tag} }}\n" - xs_ng_target_map += f"{XSUITE_MADNG_ENV_NAME}.xs_ng_target_map['{target.term}'] = '{qty}'\n" + init = self.tw_kwargs.get('init', None) - else: - raise NotImplementedError(f"Target of type {type(target)} not implemented for MAD-NG TPSA matching.") + if init is None: + init = self.line.madng_twiss(**self.tw_kwargs) + self.tw_kwargs.update({'init': init}) - self.optics_target_locations = list(self.optics_target_locations) + assert isinstance(init, xt.TwissTable) + self.mng = _ensure_madng_model(self.line) + + # Keep dynamic values in pymadng's native data channel. The MAD-NG + # command below is deliberately independent of target names and values. + targets = self._process_targets(init) + beta0_data, coordinates = self._initial_conditions(init) + self.mng._env.xs_ng_target_map = {t.xs_qty: t.qty for t in targets} + self.mng._env.targets_arr = [t.as_ng() for t in targets] + self.mng._env.tar_optics_qtys = list(self.optics_target_quantities) + + _ng_run( + self.mng, + _LUA_PREPARE_TPSA, + self.optics_target_locations, + list(self.vary_names), + beta0_data, + coordinates, + self.sum_rmat_tar, + ) - return targets_map_str, xs_ng_target_map + self._already_prepared = True - def _beta_block_str(self, init): - madng_init_flag = "x_ng" in init.cols - quantity_appendix = "_ng" if "x_ng" in init.cols else "" - start = self.tw_kwargs.get('start', None) - start_loc = 0 if start is None else start - - init_cond_str = f"""local B0 = MAD.beta0 {{ - beta11 = {init['beta11'+quantity_appendix,start_loc] if madng_init_flag else init['betx',start_loc]}, - beta22 = {init['beta22'+quantity_appendix,start_loc] if madng_init_flag else init['bety',start_loc]}, - alfa11 = {init['alfa11'+quantity_appendix,start_loc] if madng_init_flag else init['alfx',start_loc]}, - alfa22 = {init['alfa22'+quantity_appendix,start_loc] if madng_init_flag else init['alfy',start_loc]}, - dx = {init['dx'+quantity_appendix,start_loc] if madng_init_flag else init['dx',start_loc]}, - dpx = {init['dpx'+quantity_appendix,start_loc] if madng_init_flag else init['dpx',start_loc]}, - dy = {init['dy'+quantity_appendix,start_loc] if madng_init_flag else init['dy',start_loc]}, - dpy = {init['dpy'+quantity_appendix,start_loc] if madng_init_flag else init['dpy',start_loc]}, - }}""" - - return init_cond_str - - def _initialize_coordinates_str(self, init): - madng_init_flag = "x_ng" in init.cols - quantity_appendix = "_ng" if madng_init_flag else "" - beta0 = self.line.particle_ref.beta0[0] - start = self.tw_kwargs.get('start', None) - start_loc = 0 if start is None else start + def _process_targets(self, init: xt.TwissTable) -> list[_Target]: + """Translate the Xsuite targets and derive the state they imply.""" + targets = _build_targets( + self.targets, + init, + start=self.tw_kwargs.get('start', None), + end=self.tw_kwargs.get('end', None), + ) - init_coord = np.zeros(6) - init_coord[0] = init['x' + quantity_appendix, start_loc] - init_coord[1] = init['px' + quantity_appendix, start_loc] - init_coord[2] = init['y' + quantity_appendix, start_loc] - init_coord[3] = init['py' + quantity_appendix, start_loc] + locations: set[str] = set() + for target in targets: + if target.rmat: + # Transfer matrices are tracked separately, over their own range + continue + locations.add(target.loc) + if target.loc_start is not None: + locations.add(target.loc_start) + self.optics_target_locations = list(locations) + + self.optics_target_quantities = { + target.xs_qty for target in targets if not target.rmat + } + self.rmat_tags = [t.rtag for t in targets if t.rtag is not None] + + # Ordered by tag, as MAD-NG indexes the transfer maps by their position. + # The start is never absent, ``_target_range`` having resolved it. + rmat_ranges = { + t.tag: (t.loc_start, t.loc) + for t in targets + if t.rmat and t.loc_start is not None + } + self.rmat_start_end_list = [ + rmat_ranges[tag] for tag in range(self.sum_rmat_tar) + ] + + # Only orbit targets are expressed in Xsuite units that MAD-NG does not + # share, so only their Jacobian rows need rescaling. + self._needs_zeta_scale = [ + i for i, t in enumerate(targets) + if t.orbit is not None and t.xs_qty == 'zeta' + ] + self._needs_delta_scale = [ + i for i, t in enumerate(targets) + if t.orbit is not None and t.xs_qty == 'delta' + ] + + return targets + + def _initial_conditions( + self, init: xt.TwissTable + ) -> tuple[dict[str, Any], list[list[Any]]]: + """Extract the beta0 values and the closed orbit at the start of the range. + + ``init`` may come either from a MAD-NG Twiss (columns suffixed with + ``_ng``) or from an Xsuite Twiss (columns with Xsuite names). Both are + read here with MAD-NG naming conventions. + """ + loc = self.tw_kwargs.get('start', None) or 0 - if madng_init_flag: - init_coord[4] = init['t_ng', start_loc] - init_coord[5] = init['pt_ng', start_loc] + if 'x_ng' in init.cols: + values = {nn: init[f'{nn}_ng', loc] for nn in BETA0_QUANTITIES + PART_COORDS} else: - init_coord[4] = init['zeta', start_loc] / beta0 - init_coord[5] = init['ptau', start_loc] # ptau corresponds to pt + beta0 = self.line.particle_ref.beta0[0] + values = { + nn: init[NG_XS_MAP.get(nn, nn), loc] + for nn in BETA0_QUANTITIES + PART_COORDS[:4] + } + values['t'] = init['zeta', loc] / beta0 + values['pt'] = init['ptau', loc] # ptau corresponds to pt - # Build small Lua snippet setting the initial orbit, e.g. - # ``map1.x:set0(val) ...``. We must use ``:set0`` (set the 0th-order / + beta0_data = {nn: values[nn] for nn in BETA0_QUANTITIES} + + # The orbit is applied in MAD-NG with ``:set0`` (set the 0th-order / # constant part) and NOT ``map1.x = val``: to preserve TPSA. # Otherwise TPSA is replaced which corrupts the A-matrix row and # the optical functions for non-zero orbit. + coordinates = [ + [nn, values[nn]] for nn in PART_COORDS if abs(values[nn]) > 1e-12 + ] - coord_assign = " ".join( - f"map1.{p}:set0({v})" - for p, v in zip(['x','px','y','py','t','pt'], init_coord) - if abs(v) > 1e-12 - ) - return coord_assign + return beta0_data, coordinates + + def _track_options(self, start: str | None, end: str | None) -> dict[str, Any]: + """Return the MAD-NG track/twiss options for the given range. + + The sequence is passed by reference, so that its name never has to be + interpolated into a MAD-NG command. + """ + options = { + 'savemap': True, + 'observe': 1, + } + if start is not None and end is not None: + options['range'] = f'{start}/{end}' + return options - def run(self): + def run(self, allow_failure: bool = False) -> xt.TwissTable: """ Execute the MAD-NG TPSA matching action. This method performs either a Twiss or Track operation in MAD-NG @@ -833,69 +1127,31 @@ def run(self): if self._already_prepared is False: self.prepare() - start = self.tw_kwargs.get('start', None) - end = self.tw_kwargs.get('end', None) - - operation = "twiss" if self.twiss_flag else "track" - range_str = f"range='{start}/{end}', " if (start and end) else "" - - mng_track_str = ( - f"local trk, mflw = MAD.{operation}{{\n" - f" sequence={self.mng._sequence_name},\n" - f" X0={XSUITE_MADNG_ENV_NAME}.init_X0_map,\n" - f" savemap=true,\n" - f" observe=1,\n" - f" {range_str}\n" - f"}}\n" - f"{XSUITE_MADNG_ENV_NAME}.trk = trk\n" + _ng_run( + self.mng, + _LUA_TRACK, + 'twiss' if self.twiss_flag else 'track', + self._track_options( + self.tw_kwargs.get('start'), self.tw_kwargs.get('end') + ), + ) + _ng_run(self.mng, _LUA_RESULT, self.optics_target_locations) + res = xt.TwissTable( + self.mng.recv(f'{XSUITE_MADNG_ENV_NAME}.trk').to_df() ) - self.mng.send(mng_track_str) - - loc_map_str = '' - loc_map_str = "\n".join(f"{XSUITE_MADNG_ENV_NAME}.target_loc_map['{loc}'] = {XSUITE_MADNG_ENV_NAME}.trk['{loc}'].__map" for loc in self.optics_target_locations) - - # Twiss if self.twiss_flag: - mng_table_str = r''' - local trk = ''' + XSUITE_MADNG_ENV_NAME + r'''.trk - ''' + loc_map_str + r''' - py:send(trk) - ''' - - res = self.mng.send(mng_table_str).recv(XSUITE_MADNG_ENV_NAME + '.trk').to_df() - res = xt.TwissTable(res) - - # Add quantities which are not present yet with the name corresponding to the target quantity + # Alias the target quantities to the names used by the user (Xsuite) for qty in self.optics_target_quantities: if qty not in res.cols: - res[qty] = res[qty[:-3] if qty.endswith('_ng') else XS_NG_MAP[qty]] - - # Track - else: - mng_table_str = r''' - local trk = ''' + XSUITE_MADNG_ENV_NAME + r'''.trk - -- Add derived columns which are not present due to Track calculation - -- and use target names as defined from the user (Xsuite) - for _, tar in ipairs( ''' + XSUITE_MADNG_ENV_NAME + r'''.tar_optics_qtys ) do - if not trk[''' + XSUITE_MADNG_ENV_NAME + r'''.xs_ng_target_map[tar]] then - trk:addcol(tar, \ri -> MAD.gphys.optfun(trk[ri].__map, ''' + XSUITE_MADNG_ENV_NAME + r'''.xs_ng_target_map[tar] .. '_')) - end - end - - -- Save damaps - ''' + loc_map_str + r''' - py:send(trk) - ''' - - res_ng = self.mng.send(mng_table_str).recv(XSUITE_MADNG_ENV_NAME + '.trk') - res_tab = res_ng.to_df() - res = xt.TwissTable(res_tab) + res[qty] = res[to_ng_name(qty)] if 'zeta' in self.optics_target_quantities: res._data.loc[:, 'zeta'] = res['t'] * self.line.particle_ref.beta0[0] if 'delta' in self.optics_target_quantities: - res._data.loc[:, 'delta'] = ptau2delta(res['pt'], self.line.particle_ref.beta0[0]) + res._data.loc[:, 'delta'] = ptau2delta( + res['pt'], self.line.particle_ref.beta0[0] + ) if self.match_rmat: res = self.handle_rmatrices(res) @@ -903,46 +1159,29 @@ def run(self): self._last_res = res return res - def handle_rmatrices(self, res): - rmat_str = '' - rmatrices = [] - for i in range(self.sum_rmat_tar): - start_rmat = self.rmat_start_end_list[i][0] - end_rmat = self.rmat_start_end_list[i][1] + def handle_rmatrices(self, res: xt.TwissTable) -> xt.TwissTable: + """Evaluate requested transfer-matrix terms and attach them to ``res``.""" + options_arr = [] + for start_rmat, end_rmat in self.rmat_start_end_list: if start_rmat == '__ele_start__': start_rmat = self.tw_kwargs.get('start', None) if end_rmat == '__ele_stop__': end_rmat = self.tw_kwargs.get('end', None) + options_arr.append(self._track_options(start_rmat, end_rmat)) - range_str = '' - if start_rmat is not None and end_rmat is not None: - range_str = f"range = '{start_rmat}/{end_rmat}', " - rmat_str = ( - f"local trkid, mflwid = MAD.track{{\n" - f" sequence={self.mng._sequence_name},\n" - f" X0={XSUITE_MADNG_ENV_NAME}.empty_X0,\n" - f" savemap=true,\n" - f" observe=1,\n" - f" {range_str}\n" - f"}}\n" - f"{XSUITE_MADNG_ENV_NAME}.rmat_map_arr[{i}] = mflwid[1]\n" - f"local rmat = mflwid[1]:get1()\n" - f"py:send(rmat)\n" - ) - - rmat_res = self.mng.send(rmat_str).recv('rmat') - rmatrices.append(rmat_res) + _ng_run(self.mng, _LUA_RMATRICES, options_arr) + rmatrices = [self.mng.recv() for _ in options_arr] for tag in self.rmat_tags: - t0, term = tag.split("_") + t0, term = tag.split('_') ii = int(term[1]) - 1 jj = int(term[2]) - 1 res._data.attrs[tag] = rmatrices[int(t0)][ii, jj] return res - def acquire_jacobian(self): - ''' + def acquire_jacobian(self) -> np.ndarray: + """ Acquire the Jacobian matrix for the TPSA matching targets and variables. This method computes the Jacobian matrix for the specified targets and variables using MAD-NG's TPSA capabilities. It constructs @@ -955,114 +1194,101 @@ def acquire_jacobian(self): np.ndarray A 2D NumPy array representing the Jacobian matrix, where each row corresponds to a target and each column corresponds to a variable. - ''' - - tar_len_str = f"local tarlen = {len(self.targets)}\n" - vary_len_str = f"local varylen = {len(self.vary_names)}\n" - jac_decl_str = f"{XSUITE_MADNG_ENV_NAME}.jac = MAD.matrix(tarlen, varylen)\n" - - mng_str = tar_len_str + vary_len_str + jac_decl_str + r''' - -- Compute Jacobian - for i, target in ipairs( ''' + XSUITE_MADNG_ENV_NAME + r'''.targets_arr ) do - local map = nil - local nv = ''' + XSUITE_MADNG_ENV_NAME + r'''.init_X0_map:nv() - if target.optfun or target.orbit then - map = ''' + XSUITE_MADNG_ENV_NAME + r'''.target_loc_map[target.loc] - elseif target.rmat then - map = ''' + XSUITE_MADNG_ENV_NAME + r'''.rmat_map_arr[target.tag] - end - - local monom = MAD.monomial(nv + varylen) -- BUILD MONOMIAL - for j = 1, map.np(map), 1 do - -- Quantity which can be calculated with optfun - if target.optfun then - -- If loc_start (phase advance) is defined, we provide initial map - if target.loc_start then - local a0 = ''' + XSUITE_MADNG_ENV_NAME + r'''.target_loc_map[target.loc_start] - ''' + XSUITE_MADNG_ENV_NAME + r'''.jac[(i-1)*varylen + j] = MAD.gphys.optfun(map, target.qty .. "_", j, 1, a0) - else - ''' + XSUITE_MADNG_ENV_NAME + r'''.jac[(i-1)*varylen + j] = MAD.gphys.optfun(map, target.qty .. "_", j, 1) - end - - -- Orbit Quantity - elseif target.orbit then - monom[nv + j] = 1 - ''' + XSUITE_MADNG_ENV_NAME + r'''.jac[(i-1)*varylen + j] = map[target.orbit]:get(monom) - monom[nv + j] = 0 - - elseif target.rmat then - -- rmatrix terms are extracted directly from the damap - local ind_1 = tonumber(target.qty:sub(2,2)) - local ind_2 = tonumber(target.qty:sub(3,3)) - monom[nv + j] = 1 - monom[ind_2] = 1 - ''' + XSUITE_MADNG_ENV_NAME + r'''.jac[(i-1)*varylen + j] = map[ind_1]:get(monom) - monom[nv + j] = 0 - monom[ind_2] = 0 - end - end - end - - py:send(''' + XSUITE_MADNG_ENV_NAME + r'''.jac) - ''' - - self.mng.send(mng_str) + """ + _ng_run(self.mng, _LUA_JACOBIAN) jac = np.array(self.mng.recv()) for i in self._needs_zeta_scale: jac[i, :] *= self.line.particle_ref.beta0[0] for i in self._needs_delta_scale: - jac[i, :] *= dptau2ddelta(self._last_res['delta', self.targets[i].tar[1]], self.line.particle_ref.beta0[0]) + jac[i, :] *= dptau2ddelta( + self._last_res['delta', self.targets[i].tar[1]], + self.line.particle_ref.beta0[0], + ) return jac - def cleanup(self): - # Need to reconvert TPSAs to normal values + def cleanup(self) -> None: + """Restore matched MAD-NG variables and release the TPSA environment.""" if self._already_prepared is True: - mng_str = '' - for var_name in self.vary_names: - mng_str += f"MADX['{var_name}'] = MADX['{var_name}']:get0()\n" - mng_str += f"{XSUITE_MADNG_ENV_NAME}.X0 = nil\n" - self.mng.send(mng_str) + _ng_run(self.mng, _LUA_CLEANUP, list(self.vary_names)) + self.mng._env.X0 = None self._already_prepared = False -def line_to_madng(line, sequence_name='seq', temp_fname=None, keep_files=False, - **kwargs): + +def line_to_madng( + line: xt.Line, + sequence_name: str = 'seq', + temp_fname: str | None = None, + keep_files: bool = False, + **kwargs: Any, +) -> Any: + """Serialize an Xsuite line and load it into a new MAD-NG session. + + Parameters + ---------- + line : xtrack.Line + Line to convert into a MAD-NG sequence. + sequence_name : str, optional + Name assigned to the generated MAD-NG sequence. + temp_fname : str, optional + Prefix for the temporary MAD-NG input file. + keep_files : bool, optional + Keep the generated input file after loading it. + **kwargs + Additional options passed to :class:`pymadng.MAD`. + + Returns + ------- + object + The initialized ``pymadng.MAD`` session. + """ try: - _ge = xt.elements._get_expr if temp_fname is None: - temp_fname = 'temp_madng_' + str(uuid.uuid4()) + temp_fname = f'temp_madng_{uuid.uuid4()}' from .mad_writer import to_madng_sequence + madx_seq = to_madng_sequence(line, name=sequence_name) - with open(f'{temp_fname}.mad', 'w') as fid: - fid.write(madx_seq) + Path(f'{temp_fname}.mad').write_text(madx_seq) from pymadng import MAD - nocharge = str(kwargs.pop('nocharge', True)).lower() - - mng = MAD(**kwargs) - mng.send(f""" - local mad_func = loadfile('{temp_fname}.mad', nil, MADX) - assert(mad_func) - mad_func() - MAD.option.nocharge = {nocharge} - MADX.option.rbarc = true - {XSUITE_MADNG_ENV_NAME} = {{}} -- to avoid variable name clashes - """) + nocharge = kwargs.pop('nocharge', True) + + # Typed as Any: xtrack stashes its own state on the MAD-NG session + mng: Any = MAD(**kwargs) + mng.MAD.option.nocharge = nocharge + mng.MADX.option.rbarc = True + mng.send('assert(loadfile(py:recv(), nil, MADX))()').send( + str(Path(f'{temp_fname}.mad').resolve()) + ) + mng[XSUITE_MADNG_ENV_NAME] = [] # Create an empty table in MAD-NG to store xsuite data + mng._env = mng[XSUITE_MADNG_ENV_NAME] mng._init_madx_data = madx_seq - mng[sequence_name] = mng.MADX[sequence_name] # this ensures that the file has been read - mng[sequence_name].beam = mng.beam(particle="'custom'", - mass=line.particle_ref.mass0 / 1e9, # xsuite mass eV -> ng mass GeV. - charge=line.particle_ref.q0, - betgam=line.particle_ref.beta0[0] * line.particle_ref.gamma0[0]) + # A variable that does not exist in the MAD-X environment defaults to 0 + sequence = mng.MADX[sequence_name] + if sequence == 0: + raise ValueError( + f"Sequence '{sequence_name}' not found in MAD-NG model. " + "Check the generated MAD-NG input file for errors." + ) + # The sequence is kept in the xsuite environment, so that commands can + # refer to it without its name ever being interpolated into them. + mng._sequence = sequence + mng._env.sequence = sequence + mng[sequence_name] = sequence + mng[sequence_name].beam = mng.beam( + particle="'custom'", + mass=line.particle_ref.mass0 / 1e9, # xsuite mass eV -> ng mass GeV. + charge=line.particle_ref.q0, + betgam=line.particle_ref.beta0[0] * line.particle_ref.gamma0[0], + ) finally: if not keep_files: - for nn in [temp_fname + '.madx', temp_fname + '.mad']: - if os.path.isfile(nn): - os.remove(nn) + for nn in [f'{temp_fname}.madx', f'{temp_fname}.mad']: + Path(nn).unlink(missing_ok=True) return mng diff --git a/xtrack/match.py b/xtrack/match.py index 928e872c7..7f05fd20d 100644 --- a/xtrack/match.py +++ b/xtrack/match.py @@ -608,8 +608,11 @@ def compute(self, tw): 'Only terms of the R-matrix in the form "r11", "r12", "r21", "r22", etc' ' are supported') - if hasattr(tw._data, 'attrs') and self.tag in tw._data.attrs: - return tw._data.attrs[self.tag] + # rtag identifies both the range and the term; it is set by OptimizeLine + # and is the key under which the MAD-NG interface stores the value + rmat_key = getattr(self, 'rtag', None) or self.tag + if hasattr(tw._data, 'attrs') and rmat_key in tw._data.attrs: + return tw._data.attrs[rmat_key] if self.start is xt.START: self.start = tw.name[0]