Aperture modelling tools for Xtrack

.gitignore

@@ -15,4 +15,4 @@ __pycache__
 xtrack/prebuilt_kernels
 !xtrack/prebuilt_kernels/*.py
 checkpoint_restart.dat
-
+.idea

examples/aperture_model/000_lhc_mqfa.py

@@ -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()

examples/aperture_model/000_lhc_whole.py

@@ -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()

examples/aperture_model/001_transform_and_interpolate.py

@@ -0,0 +1,156 @@
+import xtrack as xt
+import xobjects as xo
+import numpy as np
+import matplotlib.pyplot as plt
+from xtrack.aperture.aperture import Aperture
+from xtrack.aperture.transform import transform_matrix
+from xtrack.aperture.structures import ApertureModel, Pipe, Circle, Profile, ProfilePosition, Rectangle, PipePosition
+
+
+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),
+]
+
+pipes = [
+    Pipe(curvature=0., positions=[
+        ProfilePosition(profile_index=1, shift_s=0, rot_s_rad=np.deg2rad(15)),
+        ProfilePosition(profile_index=1, shift_s=5.5, rot_s_rad=np.deg2rad(90)),
+        ProfilePosition(profile_index=0, shift_s=11, rot_x_rad=np.deg2rad(10)),
+    ]),
+]
+
+pipe_positions = [
+    PipePosition(
+        pipe_index=0,
+        survey_reference_name='drift::0',
+        survey_index=sv.name.tolist().index('drift::0'),
+        transformation=transform_matrix(shift_x=-1.5),
+    ),
+]
+
+model = ApertureModel(
+    line_name='line',
+    pipe_positions=pipe_positions,
+    pipes=pipes,
+    profiles=profiles,
+    pipe_names=['type0'],
+    pipe_position_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.pipe_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.shift_s,
+            theta=profile_pos.rot_y_rad,
+            phi=profile_pos.rot_x_rad,
+            psi=profile_pos.rot_s_rad,
+        )
+
+        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()