Fix bounds calculations

Author: szymonlopaciuk

examples/aperture_model/005_interpolate_on_curve.py

@@ -58,9 +58,10 @@ def arc_matrix(length: float, angle: float, tilt: float) -> np.ndarray:
 length = 3.2
 radius = 0.6
 
-bend_name = env.new('bend', xt.Bend, length=length, angle=angle, k0=0)
-anti_bend_name = env.new('anti_bend', xt.Bend, length=length, angle=-angle, k0=0)
-line = env.new_line(name='line', components=[bend_name, anti_bend_name])
+bend = env.new('bend', xt.Bend, length=length, angle=angle, k0=0)
+drift = env.new('drift', xt.Drift, length=length)
+anti_bend = env.new('anti_bend', xt.Bend, length=length, angle=-angle, k0=0)
+line = env.new_line(name='line', components=[bend, drift, anti_bend])
 sv = line.survey()
 
 shape = Circle(radius=radius)
@@ -87,19 +88,34 @@ def arc_matrix(length: float, angle: float, tilt: float) -> np.ndarray:
             survey_index=1,
             transformation=transform_matrix(),
         ),
+        TypePosition(
+            type_index=2,
+            survey_reference_name=sv.name[2],
+            survey_index=2,
+            transformation=transform_matrix(),
+        ),
     ],
     types=[
         ApertureType(curvature=angle / length, positions=profile_positions),
+        ApertureType(curvature=0, positions=profile_positions),
         ApertureType(curvature=-angle / length, positions=profile_positions),
     ],
     profiles=profiles,
-    type_names=['type0', 'type1'],
+    type_names=['type0', 'type1', 'type2'],
     profile_names=['circ0'],
 )
 
 ap = Aperture(line=line, model=model, num_profile_points=256)
 
-s_samples = np.linspace(0.0, 2 * length, 33, dtype=np.float32)
+bounds_table = ap.get_bounds_table()
+bounds_s = [0, length, length, 2 * length, 2 * length, 3 * length]
+xo.assert_allclose(bounds_table.s, bounds_s, atol=1e-6, rtol=1e-6)
+xo.assert_allclose(bounds_table.s_start, bounds_s, atol=1e-6, rtol=1e-6)
+xo.assert_allclose(bounds_table.s_end, bounds_s, atol=1e-6, rtol=1e-6)
+assert all(bounds_table.type_name == ['type0', 'type0', 'type1', 'type1', 'type2', 'type2'])
+assert all(bounds_table.profile_name == ['circ0'])
+
+s_samples = np.linspace(0, 3 * length, 51, dtype=np.float32)
 sections, poses = ap.cross_sections_at_s(s_samples)
 
 ax = plt.figure().add_subplot(projection="3d")

xtrack/aperture/aperture.py

@@ -743,6 +743,7 @@ def get_bounds_table(self):
 
         table = Table(
             data={
+                'name': np.array([f'{pn}_in_{tn}' for pn, tn in zip(profile_names, type_names)], dtype=np.str_),
                 'type_name': np.array(type_names, dtype=np.str_),
                 'profile_name': np.array(profile_names, dtype=np.str_),
                 's': np.array(s_positions, dtype=np.float32),
@@ -751,7 +752,7 @@ def get_bounds_table(self):
                 'shape': np.array(shapes, dtype=object),
                 'shape_param': np.array(shape_params, dtype=object),
             },
-            index='type_name',
+            index='name',
         )
         return table
 

xtrack/aperture/headers/path3d.h

@@ -287,7 +287,12 @@ float_type arc_segment_plane_intersect(const ArcSegment3D segment, const Point3D
 }
 
 
-inline float_type segment_plane_intersect(const Segment3D segment, const Point3D plane_point, const Point3D normal)
+inline float_type segment3d_plane_intersect(const Segment3D segment, const Point3D plane_point, const Point3D normal)
+/*
+    Given a `segment` (line or arc) and a plane defined with a point and a normal, get a parameter `t` along
+    the length of the `segment` at which the plane intersects the segment. If `t` is not in [0, 1] the plane
+    does not intersect the segment.
+*/
 {
     switch (segment.type) {
         case SEGMENT3D_LINE:
@@ -298,11 +303,13 @@ inline float_type segment_plane_intersect(const Segment3D segment, const Point3D
 }
 
 
-inline float_type closest_t_on_segment(const Point3D p, const Point3D a, const Point3D b)
+inline float_type closest_t_on_line_segment(const Point3D p, const LineSegment3D segment)
 /*
     Closest point parameter t on segment [a,b] to point p.
 */
 {
+    const Point3D a = segment.start;
+    const Point3D b = segment.end;
     const float_type eps = APER_PRECISION;
     const Point3D ab = point3d_sub(b, a);
     const Point3D ap = point3d_sub(p, a);
@@ -315,6 +322,99 @@ inline float_type closest_t_on_segment(const Point3D p, const Point3D a, const P
 }
 
 
+inline float_type closest_t_on_arc_segment(const Point3D p, const ArcSegment3D segment)
+/*
+    Closest point parameter t on an ArcSegment3D to point p.
+
+    The returned value is the normalized arc parameter, such that:
+        arc_segment_point_at(segment, t)
+    is the closest point on the supporting arc.
+
+    Like closest_t_on_segment(), this returns the unconstrained parameter.
+    Clamp to [0, 1] at the call site if you need the closest point strictly
+    on the finite arc segment.
+*/
+{
+    const float_type eps = APER_PRECISION;
+
+    if (segment.length < eps) return 0.f;
+
+    /*
+        Transform p into the local frame of segment.start.
+
+        Since segment.start is a rigid transform, its inverse is:
+            local = R^T * (p - T)
+    */
+    const float_type px = p.x - segment.start.mat[0][3];
+    const float_type py = p.y - segment.start.mat[1][3];
+    const float_type pz = p.z - segment.start.mat[2][3];
+
+    const float_type qx =
+        segment.start.mat[0][0] * px +
+        segment.start.mat[1][0] * py +
+        segment.start.mat[2][0] * pz;
+
+    const float_type qy =
+        segment.start.mat[0][1] * px +
+        segment.start.mat[1][1] * py +
+        segment.start.mat[2][1] * pz;
+
+    const float_type qz =
+        segment.start.mat[0][2] * px +
+        segment.start.mat[1][2] * py +
+        segment.start.mat[2][2] * pz;
+
+    /*
+        Undo the arc roll. In this frame the arc lies in the x-z plane:
+            x = (cos(theta) - 1) / k
+            z =  sin(theta) / k
+        with theta = k * s.
+    */
+    const float_type c_roll = cos(segment.roll);
+    const float_type s_roll = sin(segment.roll);
+
+    const float_type x =  c_roll * qx + s_roll * qy;
+    const float_type z =  qz;
+
+    const float_type k = segment.curvature;
+
+    /* Straight-line limit */
+    if (fabs(k) < eps) {
+        return z / segment.length;
+    }
+
+    /*
+        The arc is part of the circle:
+            (x + 1/k)^2 + z^2 = (1/k)^2
+
+        If C = (-1/k, 0) is the circle center in the local x-z plane, then
+        the closest point on the supporting circle is obtained by projecting
+        (x, z) radially from C.
+
+        For a point exactly on the arc:
+            x = (cos(theta) - 1) / k
+            z =  sin(theta) / k
+
+        so:
+            atan2(k*z, 1 + k*x) = theta
+    */
+    const float_type theta = atan2(k * z, 1.f + k * x);
+
+    return theta / (k * segment.length);
+}
+
+
+inline float_type closest_t_on_segment(const Point3D p, const Segment3D segment)
+{
+    switch (segment.type) {
+        case SEGMENT3D_LINE:
+            return closest_t_on_line_segment(p, segment.line);
+        case SEGMENT3D_ARC:
+            return closest_t_on_arc_segment(p, segment.arc);
+    }
+}
+
+
 inline Pose arc_matrix(const float_type length, const float_type angle, const float_type tilt)
 /*
     Get a transformation to the point at `length` along an arc of `angle`.

xtrack/aperture/headers/polygons.h

@@ -22,10 +22,11 @@ void cross_sections_at_s(
 void build_polygon_for_profile(float_type *const, const uint32_t, const Profile);
 uint32_t find_aperture_info_for_s(const ApertureBounds, const float_type s, const uint32_t lower_bound);
 
-static inline float_type survey_s_for_aperture(const TypePosition, const ProfilePosition, const SurveyData, uint32_t*);
+static inline float_type survey_s_for_aperture(const TypePosition, const ProfilePosition, const float_type curvature, const SurveyData, uint32_t*);
 static inline void bounds_on_s_for_aperture(
     const TypePosition,
     const ProfilePosition,
+    const float_type curvature,
     const SurveyData,
     const Point2D* const,
     const uint32_t num_poly_points,
@@ -97,6 +98,7 @@ void build_profile_polygons(
         const TypePosition type_pos = ApertureModel_getp1_type_positions(model, type_pos_idx);
         const uint32_t type_idx = TypePosition_get_type_index(type_pos);
         const ApertureType aper_type = ApertureModel_getp1_types(model, type_idx);
+        const float_type curvature = ApertureType_get_curvature(aper_type);
 
         /* Get the profile position, and the polygon */
         const ProfilePosition profile_pos = ApertureType_getp1_positions(aper_type, profile_pos_idx);
@@ -107,13 +109,13 @@ void build_profile_polygons(
 
         /* Get the survey s where the aperture actually sits */
         uint32_t installed_survey_index;
-        const float_type found_s = survey_s_for_aperture(type_pos, profile_pos, survey, &installed_survey_index);
+        const float_type found_s = survey_s_for_aperture(type_pos, profile_pos, curvature, survey, &installed_survey_index);
         ApertureBounds_set_s_positions(aperture_bounds, idx, found_s);
 
         /* Get the bounds in s that the aperture spans */
         float_type min_s, max_s;
         const Point2D* const profile_points = (Point2D*)poly;
-        bounds_on_s_for_aperture(type_pos, profile_pos, survey, profile_points, num_points, installed_survey_index, &min_s, &max_s);
+        bounds_on_s_for_aperture(type_pos, profile_pos, curvature, survey, profile_points, num_points, installed_survey_index, &min_s, &max_s);
         ApertureBounds_set_s_start(aperture_bounds, idx, min_s);
         ApertureBounds_set_s_end(aperture_bounds, idx, max_s);
     }
@@ -399,13 +401,13 @@ void cross_sections_at_s(
             We want to move to the type frame, so calculate the transformations
         */
         const Pose plane_in_world = pose_matrix_from_survey(survey_at_s, i);
-        // const Pose type_in_world = aperture_type_pose_in_world(type_pos, survey);
-        Pose type_in_world;
-        for (int i = 0; i < 4; i++) {
-            for (int j = 0; j < 4; j++) {
-                type_in_world.mat[i][j] = TypePosition_get_transformation(type_pos, i, j);
-            }
-        }
+        const Pose type_in_world = aperture_type_pose_in_world(type_pos, survey);
+//        Pose type_in_world;
+//        for (int i = 0; i < 4; i++) {
+//            for (int j = 0; j < 4; j++) {
+//                type_in_world.mat[i][j] = TypePosition_get_transformation(type_pos, i, j);
+//            }
+//        }
         const Pose world_in_type = pose_inverse_rigid(type_in_world);
         const Pose plane_in_type = matrix_multiply(world_in_type, plane_in_world);
         const Pose type_in_plane = pose_inverse_rigid(plane_in_type);
@@ -571,6 +573,7 @@ uint32_t find_aperture_info_for_s(
 static inline float_type survey_s_for_aperture(
     const TypePosition type_pos,
     const ProfilePosition profile_pos,
+    const float_type curvature,
     const SurveyData survey,
     uint32_t* found_survey_index
 )
@@ -588,9 +591,7 @@ static inline float_type survey_s_for_aperture(
 
     // Transformation from plane (s = 0) -> world
     Pose plane_in_world;
-    // TODO: Include curvature and correct frame calculation
-    const Pose world_in_world = identity();
-    aperture_profile_pose_frame(type_pos, profile_pos, 0, survey, world_in_world, &plane_in_world);
+    aperture_profile_pose_frame(type_pos, profile_pos, curvature, survey, identity(), &plane_in_world);
 
     /*
         For data from MAD-X etc. it's likely that it's the type's reference survey point where the profile
@@ -601,16 +602,16 @@ static inline float_type survey_s_for_aperture(
     ZigZagIterator it = zigzag_iterator_new(survey_idx, 0, num_survey_entries - 1);
     do
     {
-        LineSegment3D segment = survey_segment(survey, it.index);
+        Segment3D segment = survey_segment(survey, it.index);
         const Point3D plane_point = plane_initial_point(plane_in_world);
         const Point3D normal = plane_normal_vector(plane_in_world);
-        const float_type t = line_segment_plane_intersect(segment, plane_point, normal);
+        const float_type t = segment3d_plane_intersect(segment, plane_point, normal);
         const float_type type_s = SurveyData_get_s(survey, it.index);
 
         if (-eps <= t && t <= 1 + eps)
         {
             /* Current survey segment is intersected by the installed profile plane: compute the position. */
-            const float_type dist = t * segment_get_length(segment);
+            const float_type dist = t * segment3d_get_length(segment);
             found_s = type_s + dist;
             *found_survey_index = it.index;
             break;
@@ -624,6 +625,7 @@ static inline float_type survey_s_for_aperture(
 static inline void bounds_on_s_for_aperture(
     const TypePosition type_pos,
     const ProfilePosition profile_pos,
+    const float_type curvature,
     const SurveyData survey,
     const Point2D* const profile_points,
     const uint32_t num_poly_points,
@@ -642,18 +644,14 @@ static inline void bounds_on_s_for_aperture(
     segment (as opposed to being clamped to the edge points) we have found the right segment.
     This is a fair assumption as the diameter of a profile << radius of curvature of the survey,
     but if that is not the case, the bounds will not be correct.
-
-    TODO: Handling of arcs, for now the algorithm deals with straight segments only.
 */
 {
     const float_type eps = APER_PRECISION;
     const uint32_t num_survey_entries = SurveyData_len_s(survey);
 
     // Transformation profile local -> world frame
     Pose profile_in_world;
-    // TODO: Include curvature and correct frame calculation
-    const Pose world_in_world = identity();
-    aperture_profile_pose_frame(type_pos, profile_pos, 0, survey, world_in_world, &profile_in_world);
+    aperture_profile_pose_frame(type_pos, profile_pos, curvature, survey, identity(), &profile_in_world);
 
     float_type out_min = INFINITY;
     float_type out_max = -INFINITY;
@@ -675,17 +673,13 @@ static inline void bounds_on_s_for_aperture(
         ZigZagIterator it = zigzag_iterator_new(installed_survey_index, 0, num_survey_entries - 1);
         do
         {
-            const LineSegment3D seg = survey_segment(survey, it.index);
-
-            const Point3D a = seg.start;
-            const Point3D b = seg.end;
-
-            const float_type t = closest_t_on_segment(pt_in_world, a, b);
+            const Segment3D seg = survey_segment(survey, it.index);
+            const float_type t = closest_t_on_segment(pt_in_world, seg);
 
             if (-eps < t && t < 1 + eps) {
                 /* Candidate s on this segment */
                 const float_type seg_s_start = SurveyData_get_s(survey, it.index);
-                const float_type seg_len = segment_get_length(seg);
+                const float_type seg_len = segment3d_get_length(seg);
                 closest_s = seg_s_start + t * seg_len;
                 break;
             }

xtrack/aperture/headers/survey_tools.h

@@ -45,16 +45,52 @@ inline Point3D survey_point(SurveyData survey, uint32_t idx) {
 }
 
 
-inline LineSegment3D survey_segment(SurveyData survey, uint32_t idx) {
-    Point3D entry = survey_point(survey, idx);
-    Point3D exit = survey_point(survey, idx + 1);
+inline LineSegment3D survey_line_segment(SurveyData survey, uint32_t idx) {
+    const Point3D entry = survey_point(survey, idx);
+    const Point3D exit = survey_point(survey, idx + 1);
     return (LineSegment3D) {
         .start = entry,
         .end = exit
     };
 }
 
 
+inline ArcSegment3D survey_arc_segment(SurveyData survey, uint32_t idx) {
+    const Pose start = pose_matrix_from_survey(survey, idx);
+    const float_type length = SurveyData_get_length(survey, idx);
+    const float_type angle = SurveyData_get_angle(survey, idx);
+    const float_type roll = SurveyData_get_tilt(survey, idx);
+
+    return (ArcSegment3D) {
+        .start = start,
+        .length = length,
+        .curvature = angle / length,
+        .roll = roll
+    };
+}
+
+
+inline Segment3D survey_segment(SurveyData survey, uint32_t idx) {
+    const float_type angle = SurveyData_get_angle(survey, idx);
+    const float_type length = SurveyData_get_length(survey, idx);
+
+    if (fabs(angle) < APER_PRECISION || fabs(length) < APER_PRECISION)
+    {
+        return (Segment3D) {
+            .type = SEGMENT3D_LINE,
+            .line = survey_line_segment(survey, idx)
+        };
+    }
+    else
+    {
+        return (Segment3D) {
+            .type = SEGMENT3D_ARC,
+            .arc = survey_arc_segment(survey, idx)
+        };
+    }
+}
+
+
 SurveyEntry_s interpolate_survey_table_entry(
     const SurveyData survey,
     const float_type s_target,