Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
66 commits
Select commit Hold shift + click to select a range
8932519
Fixes for aperture loading in the madloaders
szymonlopaciuk Jan 29, 2026
62b50b1
load with limits or without
szymonlopaciuk Feb 16, 2026
8e855a1
First prototype of aperture modelling functionality
szymonlopaciuk Feb 19, 2026
b8630f5
Remove 'line_name' from the Aperture API
szymonlopaciuk Feb 24, 2026
d91145d
Don't keep rebuilding cross sections, default num points set for model
szymonlopaciuk Feb 24, 2026
9449ff7
Compute the s positions of installed profiles and their bounds
szymonlopaciuk Mar 2, 2026
61b29ea
Compute cross_sections only once at init
szymonlopaciuk Mar 2, 2026
38956d4
Don't generate a separate polygon for all CrossSections, store in Pro…
szymonlopaciuk Mar 2, 2026
38731f9
Implement interpolation between two profiles
szymonlopaciuk Mar 2, 2026
b2ed127
Import aperture_offsets from MAD-X-imported line metadata
szymonlopaciuk Mar 3, 2026
516c776
Fix for OpenMP
szymonlopaciuk Mar 3, 2026
fb661cf
Refactor common code in beam_aperture.h
szymonlopaciuk Mar 4, 2026
f95f481
Use proper interpolation routine in n1 calculations
szymonlopaciuk Mar 4, 2026
8438812
Add a test for aperture interpolations
szymonlopaciuk Mar 4, 2026
1a2fc19
Fix bug in interpolation logic
szymonlopaciuk Mar 4, 2026
f729b48
Generate zero length segment points correctly
szymonlopaciuk Mar 4, 2026
a2c95db
Fixes for interpolation continuity
szymonlopaciuk Mar 4, 2026
4b9249b
Add code to validate bounds
szymonlopaciuk Mar 4, 2026
38016ed
Codex: interpolate on both sides of the cross section plane
szymonlopaciuk Mar 4, 2026
35d9a14
Interpolate on both sides, simplify logic
szymonlopaciuk Mar 4, 2026
b6c854b
Document the new cross_sections_at_s
szymonlopaciuk Mar 4, 2026
d5368e7
Remove cgeom references
szymonlopaciuk Mar 4, 2026
4fdea0d
Unify matrix usages
szymonlopaciuk Mar 5, 2026
6072d41
Finally fix stupid MAD-X zero apertures
szymonlopaciuk Mar 5, 2026
7e70a82
Implement simpler beam envelope generation function and use interpola…
szymonlopaciuk Mar 5, 2026
c6b23f2
Checkpoint before we interpolate on curve
szymonlopaciuk Mar 6, 2026
8c59b04
Generate aperture poses on a curve
szymonlopaciuk Mar 6, 2026
e807a9b
Interpolate on a curve
szymonlopaciuk Mar 6, 2026
f0a8358
adding example aperture
rdemaria Mar 6, 2026
fcfd6eb
fix benchmark scripts
szymonlopaciuk Mar 6, 2026
062c042
slowly fix scripts
szymonlopaciuk Mar 6, 2026
35bad98
Add a table functionality
szymonlopaciuk Mar 6, 2026
0ab6950
FIXME: bad survey_ref in MAD-X offsets, wrong type position transform…
szymonlopaciuk Mar 10, 2026
1788264
Fix bounds calculations
szymonlopaciuk Mar 11, 2026
b85a7ee
Do not reinvent the wheel
szymonlopaciuk Mar 12, 2026
2f42dc2
Fixes for bend/unbend (ambiguous for >90°)
szymonlopaciuk Mar 12, 2026
389fb33
Readd validity checking
szymonlopaciuk Mar 12, 2026
9638153
Switch to analytic arc_segment_plane_intersect
szymonlopaciuk Mar 12, 2026
897629a
Remove the numerical implementation of `arc_segment_plane_intersect`
szymonlopaciuk Mar 12, 2026
44bb6ce
Add stability tests
szymonlopaciuk Mar 13, 2026
f0a7dba
Offsets almost working?
szymonlopaciuk Mar 13, 2026
631146c
Constant offsets work, non working example for dx/ddx
szymonlopaciuk Mar 16, 2026
6740204
Implement MAD-X aperture offsets
szymonlopaciuk Mar 16, 2026
e80a8c9
Fixes for tests
szymonlopaciuk Mar 16, 2026
4d882db
Include (switchable with a flag) aperture tols in beam size computation
szymonlopaciuk Mar 17, 2026
86079a7
More flexible ray method
szymonlopaciuk Mar 17, 2026
a6d848d
Tighten tolerances, remove needless _skip_validity_checks (now in f64)
szymonlopaciuk Mar 17, 2026
b9e14f8
Serialisation, progress tracking, validation, tests
szymonlopaciuk Mar 18, 2026
6c6c4ba
Zigzag wraps around the ring
szymonlopaciuk Mar 18, 2026
13b6e93
Fix aperture interpolation not taking the last point sometimes
szymonlopaciuk Mar 19, 2026
bbf0c90
Add proper tests for the zigzag
szymonlopaciuk Mar 19, 2026
8549c92
Add regression test for the unwrapped interpolation
szymonlopaciuk Mar 19, 2026
1ee1bd2
Add correlated dispersion to the closed orbit rather than the sigma c…
szymonlopaciuk Mar 20, 2026
846b7ac
Add a new exact computation method for n1
szymonlopaciuk Mar 23, 2026
fbe22bb
Fix precision problem in `convolve_poly_and_segment`
szymonlopaciuk Mar 23, 2026
281c234
Fix the `exact` method (wrong ellipse radius calculation) & add bench…
szymonlopaciuk Mar 23, 2026
e3f283b
wip codex: pull cross section generation into n1 calc loop
szymonlopaciuk Mar 23, 2026
004cf68
Generate cross sections within n1 loops
szymonlopaciuk Mar 24, 2026
6fb8bc9
In `cross_sections_at_s` output interpolated tols
szymonlopaciuk Mar 24, 2026
21b063b
Test aperture tolerance interpolation
szymonlopaciuk Mar 24, 2026
136534f
Change header names, move some C functions around
szymonlopaciuk Mar 24, 2026
6322caa
Renames
szymonlopaciuk Mar 24, 2026
8218959
Restore analytical tests working
szymonlopaciuk Mar 24, 2026
ff63fec
Fix OpenMP temporarily broken
szymonlopaciuk Mar 24, 2026
ab99ba4
Add plotting functions and optional envelope discretisation in n1 funcs
szymonlopaciuk Mar 25, 2026
b9bbfbd
More generic IR examples
szymonlopaciuk Mar 25, 2026
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
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -15,4 +15,4 @@ __pycache__
xtrack/prebuilt_kernels
!xtrack/prebuilt_kernels/*.py
checkpoint_restart.dat

.idea
93 changes: 93 additions & 0 deletions examples/aperture_model/000_lhc_mqfa.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
import matplotlib.pyplot as plt
import xobjects as xo
import xtrack as xt

from xtrack.aperture import Aperture

context = xo.ContextCpu(omp_num_threads='auto')

lhc_with_metadata = xt.load('./lhc_aperture.json')
b1 = lhc_with_metadata['b1']
lhc_length = b1.get_length()

aperture_model = Aperture.from_line_with_madx_metadata(b1, num_profile_points=100, context=context)

mqxfa_name = 'mqy.4r1.b1'

# Calculate n1 with the ``rays`` method
n1_rays, tw_rays = aperture_model.get_aperture_sigmas_at_element(
element_name=mqxfa_name,
resolution=0.1,
method='rays',
output_cross_sections=True,
)
sig_rays = n1_rays.n1
aper_rays = n1_rays.cross_section

sig_hvd_rays, _, _ = aperture_model.get_hvd_aperture_sigmas_at_element(
element_name=mqxfa_name,
resolution=0.1,
)

# Calculate n1's with the ``bisection`` method
n1_bisect, tw_bisect = aperture_model.get_aperture_sigmas_at_element(
element_name=mqxfa_name,
resolution=0.1,
method='bisection',
output_cross_sections=True,
output_max_envelopes=True,
)
sig_bisect = n1_bisect.n1
aper_bisect = n1_bisect.cross_section
max_envelope = n1_bisect.envelope

# Get envelope at arbitrary sigma
envelopes, tw_envel = aperture_model.get_envelope_at_element(
element_name=mqxfa_name,
resolution=0.1,
sigmas=1,
)

aper_table = aperture_model.cross_sections_at_element(
element_name=mqxfa_name,
resolution=0.1,
)
aper_envel = aper_table.cross_section

# PLOT envelope sigmas
plt.plot(tw_rays.s, sig_hvd_rays[:, 0], label=r'horizonal envelope [$\sigma$] (rays)')
plt.plot(tw_rays.s, sig_hvd_rays[:, 1], label=r'vertical envelope [$\sigma$] (rays)')
plt.plot(tw_rays.s, sig_hvd_rays[:, 2], label=r'diagonal envelope [$\sigma$] (rays)')
plt.plot(tw_rays.s, sig_rays, label=r'min envelope [$\sigma$] (rays)', linestyle=':')

plt.plot(tw_bisect.s, sig_bisect, label=r'max envelope [$\sigma$] (bisection)', linestyle='--')

plt.legend()
plt.show()

plt.scatter(tw_rays.x, tw_rays.y)

# PLOT max sigma beam in aperture
for pt in aper_rays:
plt.plot(pt[:, 0], pt[:, 1], c='k')

for pt in aper_bisect:
plt.plot(pt[:, 0], pt[:, 1], c='gray', linestyle='--')

for pt in max_envelope:
plt.plot(pt[:, 0], pt[:, 1])

plt.gca().set_aspect('equal')
plt.legend()
plt.show()

# PLOT arbitrary sigma beam in aperture
for pt in aper_envel:
plt.plot(pt[:, 0], pt[:, 1], c='k')

for pt in envelopes:
plt.plot(pt[:, 0], pt[:, 1])

plt.gca().set_aspect('equal')
plt.legend()
plt.show()
97 changes: 97 additions & 0 deletions examples/aperture_model/000_lhc_whole.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
import matplotlib.pyplot as plt
import numpy as np

import xobjects as xo
import xtrack as xt

from xtrack.aperture import Aperture

context = xo.ContextCpu(omp_num_threads='auto')

lhc_with_metadata = xt.load('./benchmark1/lhc_aperture.json')
b1 = lhc_with_metadata['b1']
lhc_length = b1.get_length()

aperture_model = Aperture.from_line_with_madx_metadata(b1, num_profile_points=100, context=context)

end_s = lhc_length / 8
s_positions = np.linspace(0, end_s, int(end_s))

# Calculate n1 with the ``rays`` method
n1_rays, tw_rays = aperture_model.get_aperture_sigmas_at_s(
s_positions=np.linspace(0, lhc_length, int(lhc_length)),
method='rays',
output_cross_sections=True,
)
sig_rays = n1_rays.n1
aper_rays = n1_rays.cross_section

# Calculate extents
envel, tw_envel = aperture_model.get_envelope_at_s(
s_positions=np.linspace(0, lhc_length, int(lhc_length)),
sigmas=5,
envelopes_num_points=12,
include_aper_tols=False,
)

cs_s = np.linspace(0, lhc_length, int(lhc_length))
cs_table = aperture_model.cross_sections_at_s(s_positions=cs_s)
cs = cs_table.cross_section

# Calculate n1's with the ``bisection`` method
# sig_bisect, tw_bisect, aper_bisect, max_envelope = aperture_model.get_aperture_sigmas_at_s(
# s_positions=np.linspace(0, lhc_length, int(lhc_length)),
# method='bisection',
# )

# PLOT envelope sigmas
# plt.plot(tw_rays.s, sig_rays, label=r'min envelope [$\sigma$] (rays)')

# plt.plot(tw_bisect.s, sig_bisect, label=r'max envelope [$\sigma$] (bisection)', linestyle='--')

# plt.plot(tw_rays.s, np.max(aper_rays[:, :, 0], axis=1), label=r'+ horizontal extent [mm]')
# plt.plot(tw_rays.s, np.min(aper_rays[:, :, 0], axis=1), label=r'- horizontal extent [mm]')
# plt.plot(tw_rays.s, np.max(aper_rays[:, :, 1], axis=1), label=r'+ vertical extent [mm]')
# plt.plot(tw_rays.s, np.min(aper_rays[:, :, 1], axis=1), label=r'- vertical extent [mm]')

# plt.legend()
# plt.show()
#
# plt.scatter(tw_rays.x, tw_rays.y)


fig, (ax, ay) = plt.subplots(2, sharex=True)
fig.suptitle(r'Interpolated apertures and beam at 3$\sigma$')

min_envel_x = np.min(envel[:, :, 0], axis=1) * 1000
max_envel_x = np.max(envel[:, :, 0], axis=1) * 1000
min_aper_x = np.min(cs[:, :, 0], axis=1) * 1000
max_aper_x = np.max(cs[:, :, 0], axis=1) * 1000

ax.fill_between(tw_envel.s, min_envel_x, max_envel_x, color='b', alpha=0.3)
ax.plot(cs_s, min_aper_x, color='k', marker='.')
ax.plot(cs_s, max_aper_x, color='k', marker='.')
ax.set_ylabel(r'horizontal aperture [mm]')
ax.set_ylim([-100, 100])

# ax_sig = ax.twinx()
# ax_sig.plot(tw_rays.s, sig_rays[:, 0], label=r'horizonal envelope [$\sigma$] (rays)', color='pink')
# ax_sig.set_ylabel(r'horizontal n1 [$sigma$]')

min_envel_y = np.min(envel[:, :, 1], axis=1) * 1000
max_envel_y = np.max(envel[:, :, 1], axis=1) * 1000
min_aper_y = np.min(cs[:, :, 1], axis=1) * 1000
max_aper_y = np.max(cs[:, :, 1], axis=1) * 1000
ay.set_ylabel(r'vertical aperture [mm]')
ay.set_ylim([-100, 100])

ay.fill_between(tw_envel.s, min_envel_y, max_envel_y, color='r', alpha=0.3)
ay.plot(cs_s, min_aper_y, color='k', marker='.')
ay.plot(cs_s, max_aper_y, color='k', marker='.')

# ay_sig = ay.twinx()
# ay_sig.plot(tw_rays.s, sig_rays[:, 1], label=r'vertical envelope [$\sigma$] (rays)', color='violet')
# ay_sig.set_ylabel(r'horizontal n1 [$sigma$]')

ay.set_xlabel(r's [m]')
plt.show()
154 changes: 154 additions & 0 deletions examples/aperture_model/001_transform_and_interpolate.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,154 @@
import xtrack as xt
import xobjects as xo
import numpy as np
import matplotlib.pyplot as plt
from xtrack.aperture.aperture import Aperture, transform_matrix
from xtrack.aperture.structures import ApertureModel, ApertureType, Circle, Profile, ProfilePosition, Rectangle, TypePosition


env = xt.Environment()

l = 1
dx = 1
angle = np.deg2rad(30)
l_straight = dx / np.sin(angle / 2)
rho = 0.5 * l_straight / np.sin(angle / 2)
l_curv = rho * angle

drift = env.new('drift', xt.Drift, length=l)
rot_plus = env.new('rot_plus', xt.Bend, length=l_curv, angle=angle, k0=0)
rot_minus = env.new('rot_minus', xt.Bend, length=l_curv, angle=-angle, k0=0)

line = env.new_line(
name='line',
components=[drift, rot_plus, drift, drift, rot_minus, drift],
)

sv = line.survey()

circle = Circle(radius=2)
rectangle = Rectangle(half_width=2, half_height=1)

profiles = [
Profile(shape=circle, tol_r=0, tol_x=0, tol_y=0),
Profile(shape=rectangle, tol_r=0, tol_x=0, tol_y=0),
]

types = [
ApertureType(curvature=0., positions=[
ProfilePosition(profile_index=1, s_position=0, rot_s=np.deg2rad(15)),
ProfilePosition(profile_index=1, s_position=5.5, rot_s=np.deg2rad(90)),
ProfilePosition(profile_index=0, s_position=11, rot_x=np.deg2rad(10)),
]),
]

type_positions = [
TypePosition(
type_index=0,
survey_reference_name='drift::0',
survey_index=sv.name.tolist().index('drift::0'),
transformation=transform_matrix(dx=-1.5),
),
]

model = ApertureModel(
line_name='line',
type_positions=type_positions,
types=types,
profiles=profiles,
type_names=['type0'],
profile_names=['circle', 'rectangle'],
)

line_sliced = line.copy()
line_sliced.cut_at_s(np.linspace(0, line_sliced.get_length(), 100))
sv_sliced = line_sliced.survey()
ax = plt.figure().add_subplot(projection='3d')
ax.plot(sv_sliced.Z, sv_sliced.X, sv_sliced.Y, c='b')
ax.set_xlabel('Z [m]')
ax.set_ylabel('X [m]')
ax.set_zlabel('Y [m]')

ax.auto_scale_xyz([0, 12], [-6, 6], [-6, 6])

aper = Aperture(line, model)


def matrix_from_survey_point(sv_row):
matrix = np.identity(4)
matrix[:3, 0] = sv_row.ex
matrix[:3, 1] = sv_row.ey
matrix[:3, 2] = sv_row.ez
matrix[:3, 3] = np.hstack([sv_row.X, sv_row.Y, sv_row.Z])
return matrix


def poly2d_to_hom(poly2d):
num_points = poly2d.shape[0]
poly_hom = np.column_stack((poly2d, np.zeros(num_points), np.ones(num_points))).T
return poly_hom


for type_pos in aper.model.type_positions:
aper_type = aper.model.type_for_position(type_pos)
sv_ref = sv.rows[type_pos.survey_index]

sv_ref_matrix = matrix_from_survey_point(sv_ref)
type_matrix = type_pos.transformation.to_nparray()

for profile_pos in aper_type.positions:
profile = aper.model.profile_for_position(profile_pos)

num_points = 100
poly = aper.polygon_for_profile(profile, num_points)
poly_hom = poly2d_to_hom(poly)

profile_position_matrix = transform_matrix(
dx=profile_pos.shift_x,
dy=profile_pos.shift_y,
ds=profile_pos.s_position,
theta=profile_pos.rot_y,
phi=profile_pos.rot_x,
psi=profile_pos.rot_s,
)

poly_in_sv_frame = sv_ref_matrix @ type_matrix @ profile_position_matrix @ poly_hom

xs, ys, zs = poly_in_sv_frame[:3]
ax.plot(zs, xs, ys, c='r')


def poses_at_s(line, s_positions):
"""Return a local coordinate system (each represented by a homogeneous matrix) at all ``s_positions``."""
poses = np.zeros(shape=(len(s_positions), 4, 4), dtype=np.float32)
line_sliced = line.copy()
line_sliced.cut_at_s(s_positions)
survey_sliced = line_sliced.survey()
sv_indices = np.searchsorted(survey_sliced.s, s_positions)

for idx, sv_idx in enumerate(sv_indices):
row = survey_sliced.rows[sv_idx]
poses[idx, :3, 0] = row.ex
poses[idx, :3, 1] = row.ey
poses[idx, :3, 2] = row.ez
poses[idx, :, 3] = np.hstack([row.X, row.Y, row.Z, 1])

return poses


s_for_cuts = np.linspace(1, 11, 20)
cs_table = aper.cross_sections_at_s(s_for_cuts)
profiles, poses = cs_table.cross_section, cs_table.pose

expected_poses = poses_at_s(line, s_for_cuts)
xo.assert_allclose(poses, expected_poses, atol=1e-6, rtol=1e-6)

for idx, s in enumerate(s_for_cuts):
profile = profiles[idx]
profile_hom = poly2d_to_hom(profile)
profile_in_sv_frame = poses[idx] @ profile_hom

xs, ys, zs = profile_in_sv_frame[:3]
ax.plot(zs, xs, ys, c='g')

plt.show()
Loading