Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
25 commits
Select commit Hold shift + click to select a range
3c05ce7
Expanded test to showcase other parameters
babreufig Oct 9, 2025
42902d1
More parameters
babreufig Oct 9, 2025
19e05f3
Add nocharge option with true as default
babreufig Oct 13, 2025
0bc57c5
Merge branch 'release/v0.92.1' into feature/madng_interface
babreufig Oct 13, 2025
df0e03c
Merge branch 'xsuite:main' into feature/madng_interface
babreufig Oct 23, 2025
902cbcf
Add MAD-NG Action for TPSA matching
babreufig Nov 5, 2025
8192b53
Added cleanup procedures to make it possible to twiss afterwards
babreufig Nov 5, 2025
e60591e
Nicer way of defining Damap and default cleanup to true in step
babreufig Nov 6, 2025
44cc2ef
Add chroma properties and coupling as flags for madng twiss, added su…
babreufig Nov 18, 2025
e8e0931
Unnecessary comment
babreufig Nov 18, 2025
8e5fce4
Allows TPSA with full twiss to be used
babreufig Nov 18, 2025
5d2c529
Support TPSA matching without initial conditions
babreufig Nov 19, 2025
482c04b
Fixed cleanup
babreufig Nov 19, 2025
0c904c7
Add tpsa file
babreufig Nov 19, 2025
5d97ef7
Add possibility of using xsuite and madng targets interchangeably
babreufig Nov 19, 2025
963d800
Added test for TPSA matching
babreufig Nov 19, 2025
0ed0027
Improved twiss for phase advances, now is much faster
babreufig Nov 21, 2025
17d2616
Made orbits work, added more comments
babreufig Nov 24, 2025
0538bf1
Added orbit bump test
babreufig Nov 25, 2025
6e528ba
Remove tpsa.py
babreufig Nov 25, 2025
815731b
Removed buggy import
babreufig Nov 25, 2025
41634c3
To avoid multiple cleanups
babreufig Nov 25, 2025
cabbcbd
Remove unnecessary property
babreufig Nov 26, 2025
656316b
Some simplifications
babreufig Nov 27, 2025
cfcef60
Sets rbarc to true when converting line to madng
babreufig Nov 27, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 230 additions & 13 deletions tests/test_madnginterface.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import xobjects as xo
import pathlib
import numpy as np
from xtrack._temp import lhc_match as lm

test_data_folder = pathlib.Path(
__file__).parent.joinpath('../test_data').absolute()
Expand Down Expand Up @@ -68,7 +69,7 @@ def test_madng_interface_with_multipole_errors_and_misalignments():
line[nn_quad].shift_y = sy * line.ref['on_error']
line[nn_quad].rot_s_rad = rr * line.ref['on_error']
line[nn_quad].knl[2] = kkk * line.ref['on_error']
tw = line.madng_twiss()
tw = line.madng_twiss(coupling_edw_teng=True, compute_chromatic_properties=True)

xo.assert_allclose(tw.x, tw.x_ng, atol=5e-4*tw.x.std(), rtol=0)
xo.assert_allclose(tw.y, tw.y_ng, atol=5e-4*tw.y.std(), rtol=0)
Expand All @@ -82,7 +83,7 @@ def test_madng_interface_with_multipole_errors_and_misalignments():
xo.assert_allclose(tw.by_chrom, tw.by_ng, atol=5e-3*tw.wy_chrom.max(), rtol=0)

line['on_error'] = 0
tw = line.madng_twiss()
tw = line.madng_twiss(coupling_edw_teng=True, compute_chromatic_properties=True)
xo.assert_allclose(tw.x, 0, atol=1e-10, rtol=0)
xo.assert_allclose(tw.y, 0, atol=1e-10, rtol=0)
xo.assert_allclose(tw.betx2, 0, atol=1e-10, rtol=0)
Expand Down Expand Up @@ -173,7 +174,7 @@ def test_madng_interface_with_slicing():
line.cut_at_s(np.arange(1000))

tw_xs = line.twiss4d()
tw = line.madng_twiss()
tw = line.madng_twiss(coupling_edw_teng=True, compute_chromatic_properties=True)

assert len(tw) == len(tw_xs)

Expand All @@ -196,25 +197,25 @@ def test_madng_twiss_with_initial_conditions():
line = xt.load(test_data_folder /
'hllhc15_thick/lhc_thick_with_knobs.json')
#pytest.set_trace()
tw_xs = line.twiss(betx=120, bety=150)
tw = line.madng_twiss(beta11=120, beta22=150)
tw_xs = line.twiss(betx=120, bety=150, alfx=5, alfy=5, dx=1e-4)
tw = line.madng_twiss(beta11=120, beta22=150, alfa11=5, alfa22=5, dx=1e-4)

assert len(tw) == len(tw_xs)
assert len(tw.betx) == len(tw.beta11_ng)

xo.assert_allclose(tw.betx, tw.beta11_ng, rtol=1e-7, atol=1e-6)
xo.assert_allclose(tw.bety, tw.beta22_ng, rtol=1e-7, atol=1e-6)
xo.assert_allclose(tw.alfx, tw.alfa11_ng, rtol=1e-7, atol=1e-6)
xo.assert_allclose(tw.alfy, tw.alfa22_ng, rtol=1e-7, atol=1e-6)
xo.assert_allclose(tw.betx, tw.beta11_ng, rtol=1e-6, atol=1e-6)
xo.assert_allclose(tw.bety, tw.beta22_ng, rtol=1e-6, atol=1e-6)
xo.assert_allclose(tw.alfx, tw.alfa11_ng, rtol=1e-6, atol=1e-6)
xo.assert_allclose(tw.alfy, tw.alfa22_ng, rtol=1e-6, atol=1e-6)
xo.assert_allclose(tw.dx, tw.dx_ng, rtol=1e-8, atol=1e-6)
xo.assert_allclose(tw.dy, tw.dy_ng, rtol=1e-8, atol=1e-6)
xo.assert_allclose(tw.dpx, tw.dpx_ng, rtol=1e-8, atol=1e-6)
xo.assert_allclose(tw.dpy, tw.dpy_ng, rtol=1e-8, atol=1e-6)
xo.assert_allclose(tw.x, tw.x_ng, rtol=1e-8, atol=1e-6)
xo.assert_allclose(tw.y, tw.y_ng, rtol=1e-8, atol=1e-6)

tw2_xs = line.twiss(start='s.ds.l8.b1', end='ip1', betx=100, bety=34)
tw2_xsng = line.madng_twiss(start='s.ds.l8.b1', end='ip1', beta11=100, beta22=34, xsuite_tw=False)
tw2_xs = line.twiss(start='s.ds.l8.b1', end='ip1', betx=100, bety=34, dx=1e-5)
tw2_xsng = line.madng_twiss(start='s.ds.l8.b1', end='ip1', beta11=100, beta22=34, dx=1e-5, xsuite_tw=False)

assert len(tw2_xs.betx) == len(tw2_xsng.beta11_ng)
xo.assert_allclose(tw2_xs.betx, tw2_xsng.beta11_ng, rtol=1e-8, atol=1e-6)
Expand Down Expand Up @@ -248,16 +249,34 @@ def test_madng_twiss_with_initial_conditions():
xo.assert_allclose(tw3_xsng.alfx, tw3_xsng.alfa11_ng, rtol=1e-8, atol=1e-6)
xo.assert_allclose(tw3_xsng.alfy, tw3_xsng.alfa22_ng, rtol=1e-8, atol=1e-6)

tw4_xs = line.twiss(start='ip3', end='ip4', betx=121.5668, bety=218.58374, alfx=2.295, alfy=-2.6429, dx=-0.51)
tw4_xsng = line.madng_twiss(start='ip3', end='ip4', beta11=121.5668, beta22=218.58374, alfa11=2.295,
alfa22=-2.6429, dx=-0.51, xsuite_tw=False)

assert len(tw4_xs.betx) == len(tw4_xsng.beta11_ng)
xo.assert_allclose(tw4_xs.betx, tw4_xsng.beta11_ng, rtol=1e-6, atol=1e-5)
xo.assert_allclose(tw4_xs.bety, tw4_xsng.beta22_ng, rtol=1e-6, atol=1e-5)
xo.assert_allclose(tw4_xs.alfx, tw4_xsng.alfa11_ng, rtol=1e-6, atol=1e-5)
xo.assert_allclose(tw4_xs.alfy, tw4_xsng.alfa22_ng, rtol=1e-6, atol=1e-5)
xo.assert_allclose(tw4_xs.dx, tw4_xsng.dx_ng, rtol=1e-7, atol=1e-8)
xo.assert_allclose(tw4_xs.dy, tw4_xsng.dy_ng, rtol=1e-7, atol=1e-8)
xo.assert_allclose(tw4_xs.x, tw4_xsng.x_ng, rtol=1e-8, atol=1e-10)
xo.assert_allclose(tw4_xs.y, tw4_xsng.y_ng, rtol=1e-8, atol=1e-10)
xo.assert_allclose(tw4_xs.px, tw4_xsng.px_ng, rtol=1e-8, atol=1e-10)
xo.assert_allclose(tw4_xs.py, tw4_xsng.py_ng, rtol=1e-8, atol=1e-10)
xo.assert_allclose(tw4_xs.mux, tw4_xsng.mu1_ng, rtol=1e-8, atol=1e-5)
xo.assert_allclose(tw4_xs.muy, tw4_xsng.mu2_ng, rtol=1e-8, atol=1e-5)

def test_madng_slices():
line = xt.load(test_data_folder /
'hllhc15_thick/lhc_thick_with_knobs.json')
tw = line.twiss4d()

twng = line.madng_twiss()
twng = line.madng_twiss(compute_chromatic_properties=True)

line.cut_at_s(np.linspace(0, line.get_length(), 5000))
tw_sliced = line.twiss4d()
twng_sliced = line.madng_twiss()
twng_sliced = line.madng_twiss(compute_chromatic_properties=True)
tt_sliced = line.get_table()

assert np.all(np.array(sorted(list(set(tt_sliced.element_type)))) ==
Expand Down Expand Up @@ -302,3 +321,201 @@ def test_madng_slices():
xo.assert_allclose(twng_ip.wy_ng, twng_ip_sliced.wy_ng, rtol=1e-3)
xo.assert_allclose(twng_ip.dx_ng, twng_ip_sliced.dx_ng, atol=1e-6)
xo.assert_allclose(twng_ip.dy_ng, twng_ip_sliced.dy_ng, atol=1e-6)

def test_madng_match_optics():
collider = xt.Environment.from_json(test_data_folder /
'hllhc15_thick/hllhc15_collider_thick.json')
collider.vars.load_madx(test_data_folder /
'hllhc15_thick/opt_round_150_1500.madx')

line = collider.lhcb1
tw0 = line.madng_twiss()

lm.set_var_limits_and_steps(collider)

# Match with Xsuite Targets
opt = line.match(
solve=False,
default_tol={None: 1e-8, 'betx': 1e-6, 'bety': 1e-6, 'alfx': 1e-6, 'alfy': 1e-6},
start='s.ds.l8.b1', end='ip1',
init=tw0, init_at=xt.START,
vary=[
# Only IR8 quadrupoles including DS
xt.VaryList(['kq6.l8b1', 'kq7.l8b1', 'kq8.l8b1', 'kq9.l8b1', 'kq10.l8b1',
'kqtl11.l8b1', 'kqt12.l8b1', 'kqt13.l8b1',
'kq4.l8b1', 'kq5.l8b1', 'kq4.r8b1', 'kq5.r8b1',
'kq6.r8b1', 'kq7.r8b1', 'kq8.r8b1', 'kq9.r8b1',
'kq10.r8b1', 'kqtl11.r8b1', 'kqt12.r8b1', 'kqt13.r8b1'])],
targets=[
xt.TargetSet(at='ip8', tars=('betx', 'bety', 'alfx', 'alfy', 'dx', 'dpx'), value=tw0, weight=1),
xt.TargetSet(at='ip1', betx=0.15, bety=0.1, alfx=0, alfy=0, dx=0, dpx=0, weight=1),
xt.TargetRelPhaseAdvance('mux', value = tw0['mux', 'ip1.l1'] - tw0['mux', 's.ds.l8.b1'], start='s.ds.l8.b1', end='ip1.l1', weight=1),
xt.TargetRelPhaseAdvance('muy', value = tw0['muy', 'ip1.l1'] - tw0['muy', 's.ds.l8.b1'], start='s.ds.l8.b1', end='ip1.l1', weight=1),
],
use_tpsa=True)

opt.step(30)

assert opt._err.call_counter < 20
assert len(opt.log()) < 10

tw = line.twiss(init=tw0, start='s.ds.l8.b1', end='ip1')

xo.assert_allclose(tw['betx', 'ip1'], 0.15, atol=1e-6, rtol=0)
xo.assert_allclose(tw['bety', 'ip1'], 0.1, atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfx', 'ip1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfy', 'ip1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['dx', 'ip1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['dy', 'ip1'], 0., atol=1e-6, rtol=0)

xo.assert_allclose(tw['betx', 'ip8'], tw0['betx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['bety', 'ip8'], tw0['bety', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfx', 'ip8'], tw0['alfx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfy', 'ip8'], tw0['alfy', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['dx', 'ip8'], tw0['dx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['dy', 'ip8'], tw0['dy', 'ip8'], atol=1e-6, rtol=0)

xo.assert_allclose(tw['mux', 'ip1.l1'] - tw['mux', 's.ds.l8.b1'], tw0['mux', 'ip1.l1'] - tw0['mux', 's.ds.l8.b1'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['muy', 'ip1.l1'] - tw['muy', 's.ds.l8.b1'], tw0['muy', 'ip1.l1'] - tw0['muy', 's.ds.l8.b1'], atol=1e-6, rtol=0)

opt.reload(0)
opt.actions[0].cleanup()

# Match with MAD-NG and Xsuite Targets mixed
opt = line.match(
solve=False,
default_tol={None: 1e-8, 'betx': 1e-6, 'bety': 1e-6, 'alfx': 1e-6, 'alfy': 1e-6},
start='s.ds.l8.b1', end='ip1',
init=tw0, init_at=xt.START,
vary=[
# Only IR8 quadrupoles including DS
xt.VaryList(['kq6.l8b1', 'kq7.l8b1', 'kq8.l8b1', 'kq9.l8b1', 'kq10.l8b1',
'kqtl11.l8b1', 'kqt12.l8b1', 'kqt13.l8b1',
'kq4.l8b1', 'kq5.l8b1', 'kq4.r8b1', 'kq5.r8b1',
'kq6.r8b1', 'kq7.r8b1', 'kq8.r8b1', 'kq9.r8b1',
'kq10.r8b1', 'kqtl11.r8b1', 'kqt12.r8b1', 'kqt13.r8b1'])],
targets=[
xt.TargetSet(at='ip8', tars=('beta11_ng', 'bety', 'alfa11_ng', 'alfy', 'dx_ng', 'dpx'), value=tw0, weight=1),
xt.TargetSet(at='ip1', betx=0.15, beta22_ng=0.1, alfx=0, alfa22_ng=0, dx=0, dpx_ng=0, weight=1),
xt.TargetRelPhaseAdvance('mux', value = tw0['mux', 'ip1.l1'] - tw0['mux', 's.ds.l8.b1'], start='s.ds.l8.b1', end='ip1.l1', weight=1),
xt.TargetRelPhaseAdvance('mu2_ng', value = tw0['mu2_ng', 'ip1.l1'] - tw0['mu2_ng', 's.ds.l8.b1'], start='s.ds.l8.b1', end='ip1.l1', weight=1),
],
use_tpsa=True)

opt.step(30)

assert opt._err.call_counter < 20
assert len(opt.log()) < 10

tw = line.twiss(init=tw0, start='s.ds.l8.b1', end='ip1')

xo.assert_allclose(tw['betx', 'ip1'], 0.15, atol=1e-6, rtol=0)
xo.assert_allclose(tw['bety', 'ip1'], 0.1, atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfx', 'ip1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfy', 'ip1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['dx', 'ip1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['dy', 'ip1'], 0., atol=1e-6, rtol=0)

xo.assert_allclose(tw['betx', 'ip8'], tw0['betx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['bety', 'ip8'], tw0['bety', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfx', 'ip8'], tw0['alfx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfy', 'ip8'], tw0['alfy', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['dx', 'ip8'], tw0['dx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['dy', 'ip8'], tw0['dy', 'ip8'], atol=1e-6, rtol=0)

xo.assert_allclose(tw['mux', 'ip1.l1'] - tw['mux', 's.ds.l8.b1'], tw0['mux', 'ip1.l1'] - tw0['mux', 's.ds.l8.b1'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['muy', 'ip1.l1'] - tw['muy', 's.ds.l8.b1'], tw0['muy', 'ip1.l1'] - tw0['muy', 's.ds.l8.b1'], atol=1e-6, rtol=0)

opt.reload(0)
opt.actions[0].cleanup()

# Match on full line without initial conditions
opt = line.match(
solve=False,
default_tol={None: 1e-8, 'betx': 1e-6, 'bety': 1e-6, 'alfx': 1e-6, 'alfy': 1e-6},
vary=[
# Only IR8 quadrupoles including DS
xt.VaryList(['kq6.l8b1', 'kq7.l8b1', 'kq8.l8b1', 'kq9.l8b1', 'kq10.l8b1',
'kqtl11.l8b1', 'kqt12.l8b1', 'kqt13.l8b1',
'kq4.l8b1', 'kq5.l8b1', 'kq4.r8b1', 'kq5.r8b1',
'kq6.r8b1', 'kq7.r8b1', 'kq8.r8b1', 'kq9.r8b1',
'kq10.r8b1', 'kqtl11.r8b1', 'kqt12.r8b1', 'kqt13.r8b1'])],
targets=[
xt.TargetSet(at='ip8', tars=('beta11_ng', 'beta22_ng', 'alfa11_ng', 'alfa22_ng', 'dx_ng', 'dpx_ng'), value=tw0, weight=1),
xt.TargetSet(at='ip1.l1', beta11_ng=0.15, beta22_ng=0.1, alfa11_ng=0, alfa22_ng=0, dx_ng=0, dpx_ng=0, weight=1),
xt.TargetRelPhaseAdvance('mu1_ng', value = tw0['mu1_ng', 'ip1.l1'] - tw0['mu1_ng', 's.ds.l8.b1'], start='s.ds.l8.b1', end='ip1.l1', weight=1),
xt.TargetRelPhaseAdvance('mu2_ng', value = tw0['mu2_ng', 'ip1.l1'] - tw0['mu2_ng', 's.ds.l8.b1'], start='s.ds.l8.b1', end='ip1.l1', weight=1),
],
use_tpsa=True)

opt.step(30)

assert opt._err.call_counter < 20
assert len(opt.log()) < 10

tw = line.twiss(init=tw0)

xo.assert_allclose(tw['betx', 'ip1.l1'], 0.15, atol=1e-6, rtol=0)
xo.assert_allclose(tw['bety', 'ip1.l1'], 0.1, atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfx', 'ip1.l1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfy', 'ip1.l1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['dx', 'ip1.l1'], 0., atol=1e-6, rtol=0)
xo.assert_allclose(tw['dy', 'ip1.l1'], 0., atol=1e-6, rtol=0)

xo.assert_allclose(tw['betx', 'ip8'], tw0['betx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['bety', 'ip8'], tw0['bety', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfx', 'ip8'], tw0['alfx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['alfy', 'ip8'], tw0['alfy', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['dx', 'ip8'], tw0['dx', 'ip8'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['dy', 'ip8'], tw0['dy', 'ip8'], atol=1e-6, rtol=0)

xo.assert_allclose(tw['mux', 'ip1.l1'] - tw['mux', 's.ds.l8.b1'], tw0['mux', 'ip1.l1'] - tw0['mux', 's.ds.l8.b1'], atol=1e-6, rtol=0)
xo.assert_allclose(tw['muy', 'ip1.l1'] - tw['muy', 's.ds.l8.b1'], tw0['muy', 'ip1.l1'] - tw0['muy', 's.ds.l8.b1'], atol=1e-6, rtol=0)

def test_madng_orbit_bump():
env = xt.Environment()
env.vars.default_to_zero = True
line = env.new_line(length=10, components=[
env.new('corr1', xt.Multipole, isthick=True,
knl=['kick_h_1'], ksl=['kick_v_1'], length=0.1, at=1),
env.new('corr2', xt.Multipole, isthick=True,
knl=['kick_h_2'], ksl=['kick_v_2'], length=0.1, at=2),
env.new('corr3', xt.Multipole, isthick=True,
knl=['kick_h_3'], ksl=['kick_v_3'], length=0.1, at=8),
env.new('corr4', xt.Multipole, isthick=True,
knl=['kick_h_4'], ksl=['kick_v_4'], length=0.1, at=9),
env.new('mid', xt.Marker, at=5),
env.new('end', xt.Marker, at=10)
])
line.set_particle_ref('proton', p0c=26e9)

opt = line.match(
solve=False,
betx=1, bety=1,
vary=xt.VaryList(['kick_h_1', 'kick_v_1',
'kick_h_2', 'kick_v_2',
'kick_h_3', 'kick_v_3',
'kick_h_4', 'kick_v_4']),
targets=[
xt.TargetSet(x=1e-3, y=-2e-3, px=0, py=0, at='mid'),
xt.TargetSet(x=0, y=0, px=0, py=0, at='end'),
],
use_tpsa=True
)

jac_ng = opt._err.get_jacobian(opt._err._get_x())

jac_opt = np.array([[-40, 0, -30, 0, 0, 0, 0, 0],
[-100, 0, -100, 0, 0, 0, 0, 0],
[0, 40, 0, 30, 0, 0, 0, 0],
[0, 100, 0, 100, 0, 0, 0, 0],
[-90, 0, -80, 0, -20, 0, -10, 0],
[-100, 0, -100, 0, -100, 0, -100, 0],
[0, 90, 0, 80, 0, 20, 0, 10],
[0, 100, 0, 100, 0, 100, 0, 100]])

xo.assert_allclose(jac_ng, jac_opt, rtol=1e-12, atol=1e-12)

opt.solve()

assert opt._err.call_counter < 7
Loading