xcoll 0.3.6__py3-none-any.whl → 0.4.0__py3-none-any.whl

xcoll/scattering_routines/geometry/collimator_geometry.h

@@ -0,0 +1,219 @@
+// copyright ############################### #
+// This file is part of the Xcoll Package.   #
+// Copyright (c) CERN, 2024.                 #
+// ######################################### #
+
+#ifndef XCOLL_COLL_GEOM_H
+#define XCOLL_COLL_GEOM_H
+#include <math.h>
+#include <stdio.h>
+
+typedef struct CollimatorGeometry_ {
+    // Collimator jaws (with tilts)
+    double jaw_LU;
+    double jaw_RU;
+    double length;
+    int8_t side;
+    // Jaw angles
+    double sin_zL;
+    double cos_zL;
+    double sin_zR;
+    double cos_zR;
+    double sin_zDiff;
+    double cos_zDiff;
+    int8_t jaws_parallel;
+    // Tilts
+    double sin_yL;
+    double cos_yL;
+    double sin_yR;
+    double cos_yR;
+    // Segments
+    Segment* segments_L;
+    Segment* segments_R;
+    // Impact table
+    InteractionRecordData record;
+    RecordIndex record_index;
+    int8_t record_touches;
+} CollimatorGeometry_;
+typedef CollimatorGeometry_* CollimatorGeometry;
+
+
+// This function checks if a particle hits a jaw (and which).
+// Return value: 0 (no hit), 1 (hit on left jaw), -1 (hit on right jaw).
+// Furthermore, the particle is moved to the location where it hits the jaw (drifted to the end if no hit),
+//              and transformed to the reference frame of that jaw.
+/*gpufun*/
+int8_t hit_jaws_check_and_transform(LocalParticle* part, CollimatorGeometry restrict cg){
+    double part_x, part_tan;
+    int8_t is_hit = 0;
+    double s_L = 1.e21;
+    double s_R = 1.e21;
+
+    // Find the first hit on the left jaw (if not a right-sided collimator)
+    if (cg->side != -1){
+        SRotation_single_particle(part, cg->sin_zL, cg->cos_zL);
+        part_x = LocalParticle_get_x(part);
+#ifdef XCOLL_USE_EXACT
+        part_tan = LocalParticle_get_exact_xp(part);
+#else
+        part_tan = LocalParticle_get_xp(part);
+#endif
+        s_L = get_s_of_first_crossing(part_x, part_tan, cg->segments_L, 3);
+        if (s_L < S_MAX){
+            is_hit = 1;
+        } else if (cg->side == 1){
+            // If left-sided and no hit, rotate back to lab frame
+            SRotation_single_particle(part, -cg->sin_zL, cg->cos_zL);
+        }
+    }
+
+    // if rightsided:            lab  frame
+    // if leftsided  and no hit: lab  frame
+    // if leftsided  and hit:    left frame
+    // if bothsided  and no hit: left frame
+    // if bothsided  and hit:    left frame
+
+    // Find the first hit on the right jaw (if not a left-sided collimator)
+    if (cg->side != 1){
+        if (cg->side == -1){
+            // We didn't rotate to the left frame earlier, so full rotation to the right frame now
+            SRotation_single_particle(part, cg->sin_zR, cg->cos_zR);
+        } else if (!cg->jaws_parallel){
+            // We rotated to the left frame before, so now rotate the difference
+            SRotation_single_particle(part, cg->sin_zDiff, cg->cos_zDiff);
+        }
+        part_x = LocalParticle_get_x(part);
+#ifdef XCOLL_USE_EXACT
+        part_tan = LocalParticle_get_exact_xp(part);
+#else
+        part_tan = LocalParticle_get_xp(part);
+#endif
+        s_R = get_s_of_first_crossing(part_x, part_tan, cg->segments_R, 3);
+        if (s_R < S_MAX && s_R < s_L){
+            is_hit = -1;
+        } else if (is_hit == 1){
+            if (!cg->jaws_parallel){
+                // Rotate back to left frame
+                SRotation_single_particle(part, -cg->sin_zDiff, cg->cos_zDiff);
+            }
+        } else {
+            // No hit, rotate back to lab frame
+            SRotation_single_particle(part, -cg->sin_zR, cg->cos_zR);
+        }
+    }
+
+    // if rightsided and no hit: lab   frame
+    // if rightsided and hit:    right frame
+    // if leftsided  and no hit: lab   frame
+    // if leftsided  and hit:    left  frame
+    // if bothsided  and no hit: lab  frame
+    // if bothsided  and hit:    hit   frame
+
+
+    // Drift to the impact position or end, and move to jaw frame if relevant
+    if (is_hit == 1){
+        // Move to the impact position
+#ifdef XCOLL_USE_EXACT
+        Drift_single_particle_exact(part, s_L);
+#else
+        Drift_single_particle_expanded(part, s_L);
+#endif
+        // Shift the reference frame to the jaw corner LU
+        XYShift_single_particle(part, cg->jaw_LU, 0);
+        LocalParticle_add_to_s(part, -cg->length/2*(1 - cg->cos_yL));
+        // Rotate the reference frame to tilt
+        double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), asin(cg->sin_yL));
+        LocalParticle_set_s(part, new_s);
+        if (cg->record_touches){
+            InteractionRecordData_log(cg->record, cg->record_index, part, XC_ENTER_JAW_L);
+        }
+
+    } else if (is_hit == -1){
+        // Move to the impact position
+#ifdef XCOLL_USE_EXACT
+        Drift_single_particle_exact(part, s_R);
+#else
+        Drift_single_particle_expanded(part, s_R);
+#endif
+        // Shift the reference frame to the jaw corner RU
+        XYShift_single_particle(part, cg->jaw_RU, 0);
+        LocalParticle_add_to_s(part, -cg->length/2*(1 - cg->cos_yR));
+        // Rotate the reference frame to tilt
+        double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), asin(cg->sin_yR));
+        LocalParticle_set_s(part, new_s);
+        // Mirror x
+        LocalParticle_scale_x(part, -1);
+#ifdef XCOLL_USE_EXACT
+        LocalParticle_scale_exact_xp(part, -1);
+#else
+        LocalParticle_scale_xp(part, -1);
+#endif
+        if (cg->record_touches){
+            InteractionRecordData_log(cg->record, cg->record_index, part, XC_ENTER_JAW_R);
+        }
+
+    } else {
+#ifdef XCOLL_USE_EXACT
+        Drift_single_particle_exact(part, cg->length);
+#else
+        Drift_single_particle_expanded(part, cg->length);
+#endif
+    }
+
+    return is_hit;
+}
+
+
+/*gpufun*/
+void hit_jaws_transform_back(int8_t is_hit, LocalParticle* part, CollimatorGeometry restrict cg){
+    if (is_hit == 1){
+        // Rotate back from tilt
+        double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), -asin(cg->sin_yL));
+        LocalParticle_set_s(part, new_s);
+        // Shift the reference frame back from jaw corner LU
+        XYShift_single_particle(part, -cg->jaw_LU, 0);
+        LocalParticle_add_to_s(part, cg->length/2*(1 - cg->cos_yL));
+        // If particle survived, drift to end of element
+        if (LocalParticle_get_state(part) > 0){
+            if (cg->record_touches){
+                InteractionRecordData_log(cg->record, cg->record_index, part, XC_EXIT_JAW);
+            }
+#ifdef XCOLL_USE_EXACT
+            Drift_single_particle_exact(part, cg->length - LocalParticle_get_s(part));
+#else
+            Drift_single_particle_expanded(part, cg->length - LocalParticle_get_s(part));
+#endif
+        }
+        SRotation_single_particle(part, -cg->sin_zL, cg->cos_zL);
+
+    } else if (is_hit == -1){
+        // Mirror back
+        LocalParticle_scale_x(part, -1);
+#ifdef XCOLL_USE_EXACT
+        LocalParticle_scale_exact_xp(part, -1);
+#else
+        LocalParticle_scale_xp(part, -1);
+#endif
+        // Rotate back from tilt
+        double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), -asin(cg->sin_yR));
+        LocalParticle_set_s(part, new_s);
+        // Shift the reference frame back from jaw corner RU
+        XYShift_single_particle(part, -cg->jaw_RU, 0);
+        LocalParticle_add_to_s(part, cg->length/2*(1 - cg->cos_yR));
+        // If particle survived, drift to end of element
+        if (LocalParticle_get_state(part) > 0){
+            if (cg->record_touches){
+                InteractionRecordData_log(cg->record, cg->record_index, part, XC_EXIT_JAW);
+            }
+#ifdef XCOLL_USE_EXACT
+            Drift_single_particle_exact(part, cg->length - LocalParticle_get_s(part));
+#else
+            Drift_single_particle_expanded(part, cg->length - LocalParticle_get_s(part));
+#endif
+        }
+        SRotation_single_particle(part, -cg->sin_zR, cg->cos_zR);
+    }
+}
+
+
+#endif /* XCOLL_COLL_GEOM_H */

xcoll/scattering_routines/geometry/crystal_geometry.h

@@ -0,0 +1,150 @@
+// copyright ############################### #
+// This file is part of the Xcoll Package.   #
+// Copyright (c) CERN, 2024.                 #
+// ######################################### #
+
+#ifndef XCOLL_CRY_GEOM_H
+#define XCOLL_CRY_GEOM_H
+#include <math.h>
+#include <stdio.h>
+
+typedef struct CrystalGeometry_ {
+    // Crystal inner upstream corner (with tilt)
+    double jaw_U;
+    double length;
+    int8_t side;
+    // Jaw angle
+    double sin_z;
+    double cos_z;
+    // Tilt
+    double sin_y;
+    double cos_y;
+    // Crystal geometry
+    double bending_radius;
+    double bending_angle;
+    double width;
+    double height;
+    double miscut_angle;
+    double s_B;    // Bend centre
+    double x_B;
+    double s_P;    // Miscut centre
+    double x_P;
+    double t_VImax;
+    // Segments
+    Segment* segments;
+    // Impact table
+    InteractionRecordData record;
+    RecordIndex record_index;
+    int8_t record_touches;
+} CrystalGeometry_;
+typedef CrystalGeometry_* CrystalGeometry;
+
+
+// This function checks if a particle hits a jaw (and which).
+// Return value: 0 (no hit), 1 (hit on left jaw), -1 (hit on right jaw).
+// Furthermore, the particle is moved to the location where it hits the jaw (drifted to the end if no hit),
+//              and transformed to the reference frame of that jaw.
+/*gpufun*/
+int8_t hit_crystal_check_and_transform(LocalParticle* part, CrystalGeometry restrict cg){
+    double part_x, part_tan_x, part_y, part_tan_y;
+    double s = 1.e21;
+
+    // Crystal should be single-sided
+    if (cg->side == 0){
+        LocalParticle_kill_particle(part, XC_ERR_NOT_IMPLEMENTED);
+    }
+
+    SRotation_single_particle(part, cg->sin_z, cg->cos_z);
+    part_x = LocalParticle_get_x(part);
+#ifdef XCOLL_USE_EXACT
+    part_tan_x = LocalParticle_get_exact_xp(part);
+#else
+    part_tan_x = LocalParticle_get_xp(part);
+#endif
+    part_y = LocalParticle_get_y(part);
+#ifdef XCOLL_USE_EXACT
+    part_tan_y = LocalParticle_get_exact_yp(part);
+#else
+    part_tan_y = LocalParticle_get_yp(part);
+#endif
+    s = get_s_of_first_crossing_with_vlimit(part_x, part_tan_x, part_y, part_tan_y, cg->segments, 4, -cg->height/2, cg->height/2);
+
+    if (s < S_MAX){
+        // Hit: Drift to the impact position, and move to jaw frame if relevant
+#ifdef XCOLL_USE_EXACT
+        Drift_single_particle_exact(part, s);
+#else
+        Drift_single_particle_expanded(part, s);
+#endif
+        // Shift the reference frame to the upstream jaw corner (for a crystal, this is always at s=0)
+        XYShift_single_particle(part, cg->jaw_U, 0);
+        // Rotate the reference frame to tilt
+        double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), asin(cg->sin_y));
+        LocalParticle_set_s(part, new_s);
+        if (cg->side == 1){
+            if (cg->record_touches){
+                InteractionRecordData_log(cg->record, cg->record_index, part, XC_ENTER_JAW_L);
+            }
+
+        } else {
+            // Mirror x
+            LocalParticle_scale_x(part, -1);
+#ifdef XCOLL_USE_EXACT
+            LocalParticle_scale_exact_xp(part, -1);
+#else
+            LocalParticle_scale_xp(part, -1);
+#endif
+            if (cg->record_touches){
+                InteractionRecordData_log(cg->record, cg->record_index, part, XC_ENTER_JAW_R);
+            }
+        }
+        return cg->side;
+
+    } else {
+        // No hit, rotate back to lab frame
+        SRotation_single_particle(part, -cg->sin_z, cg->cos_z);
+        // Drift to end
+#ifdef XCOLL_USE_EXACT
+        Drift_single_particle_exact(part, cg->length);
+#else
+        Drift_single_particle_expanded(part, cg->length);
+#endif
+        return 0;
+    }
+}
+
+
+/*gpufun*/
+void hit_crystal_transform_back(int8_t is_hit, LocalParticle* part, CrystalGeometry restrict cg){
+    if (is_hit != 0){
+        if (cg->side == -1){
+            // Mirror back
+            LocalParticle_scale_x(part, -1);
+#ifdef XCOLL_USE_EXACT
+            LocalParticle_scale_exact_xp(part, -1);
+#else
+            LocalParticle_scale_xp(part, -1);
+#endif
+        }
+        // Rotate back from tilt
+        double new_s = YRotation_single_particle_rotate_only(part, LocalParticle_get_s(part), -asin(cg->sin_y));
+        LocalParticle_set_s(part, new_s);
+        // Shift the reference frame back from jaw corner U
+        XYShift_single_particle(part, -cg->jaw_U, 0);
+        // If particle survived, drift to end of element
+        if (LocalParticle_get_state(part) > 0){
+            if (cg->record_touches){
+                InteractionRecordData_log(cg->record, cg->record_index, part, XC_EXIT_JAW);
+            }
+#ifdef XCOLL_USE_EXACT
+            Drift_single_particle_exact(part, cg->length - LocalParticle_get_s(part));
+#else
+            Drift_single_particle_expanded(part, cg->length - LocalParticle_get_s(part));
+#endif
+        }
+        SRotation_single_particle(part, -cg->sin_z, cg->cos_z);
+    }
+}
+
+
+#endif /* XCOLL_CRY_GEOM_H */

xcoll/scattering_routines/geometry/geometry.py

@@ -0,0 +1,26 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import xtrack as xt
+from ...interaction_record import InteractionRecord
+from ...general import _pkg_root
+
+class XcollGeometry(xt.BeamElement):
+    _xofields = {}
+
+    allow_track = False
+
+    _depends_on = [xt.Drift, xt.XYShift, xt.SRotation, xt.YRotation, InteractionRecord]
+
+    _extra_c_sources = [
+        _pkg_root.joinpath('scattering_routines','geometry','sort.h'),
+        _pkg_root.joinpath('scattering_routines','geometry','segments.h'),
+        _pkg_root.joinpath('scattering_routines','geometry','objects.h'),
+        _pkg_root.joinpath('scattering_routines','geometry','methods.h'),
+        _pkg_root.joinpath('scattering_routines','geometry','get_s.h'),
+        _pkg_root.joinpath('scattering_routines','geometry','rotation.h'),
+        _pkg_root.joinpath('scattering_routines','geometry','collimator_geometry.h'),
+        _pkg_root.joinpath('scattering_routines','geometry','crystal_geometry.h')
+    ]

xcoll/scattering_routines/geometry/get_s.h

@@ -0,0 +1,92 @@
+// copyright ############################### #
+// This file is part of the Xcoll Package.   #
+// Copyright (c) CERN, 2024.                 #
+// ######################################### #
+
+#ifndef XCOLL_GEOM_GET_S_H
+#define XCOLL_GEOM_GET_S_H
+#include <math.h>
+#include <stdio.h>
+#include <stdint.h>
+#include <stdlib.h>
+
+
+// IMPORTANT:
+// These functions assume that the particle moves towards positive s!
+// (hence no backscattering/backtracking is allowed)
+
+
+/*gpufun*/
+double get_s_of_first_crossing(double part_x, double part_tan, Segment* segments, \
+                               int8_t n_segments){
+    int8_t n_hit = 0;
+    double* s = (double*) malloc(XC_MAX_CROSS_PER_SEGMENT*n_segments*sizeof(double));
+    find_crossing(&n_hit, s, part_x, part_tan, segments, n_segments);
+    if (n_hit==0){
+        // No crossing
+        free(s);
+        return S_MAX;
+    }
+    double result = s[0];
+    free(s);
+    return result;
+}
+
+/*gpufun*/
+double get_s_of_crossing_after_s(double part_x, double part_tan, Segment* segments, \
+                                 int8_t n_segments, double current_s){
+    int8_t n_hit = 0;
+    double* s = (double*) malloc(XC_MAX_CROSS_PER_SEGMENT*n_segments*sizeof(double));
+    find_crossing(&n_hit, s, part_x, part_tan, segments, n_segments);
+    for (int8_t i=0; i<n_hit; i++){
+        if (s[i] >= current_s){
+            double result = s[i];
+            free(s);
+            return result;
+        }
+    }
+    // No crossing
+    free(s);
+    return S_MAX;
+}
+
+/*gpufun*/
+double get_s_of_first_crossing_with_vlimit(double part_x, double part_tan_x, \
+                                double part_y, double part_tan_y, Segment* segments, \
+                                int8_t n_segments, double y_min, double y_max){
+    int8_t n_hit = 0;
+    double* s = (double*) malloc(XC_MAX_CROSS_PER_SEGMENT*n_segments*sizeof(double));
+    find_crossing_with_vlimit(&n_hit, s, part_x, part_tan_x, part_y, part_tan_y, \
+                              segments, n_segments, y_min, y_max);
+    if (n_hit==0){
+        // No crossing
+        free(s);
+        return S_MAX;
+    }
+    double result = s[0];
+    free(s);
+    return result;
+}
+
+/*gpufun*/
+double get_s_of_crossing_after_s_with_vlimit(double part_x, double part_tan_x, \
+                                double part_y, double part_tan_y, Segment* segments, \
+                                int8_t n_segments, double y_min, double y_max, double current_s){
+    int8_t n_hit = 0;
+    double* s = (double*) malloc(XC_MAX_CROSS_PER_SEGMENT*n_segments*sizeof(double));
+    find_crossing_with_vlimit(&n_hit, s, part_x, part_tan_x, part_y, part_tan_y, \
+                              segments, n_segments, y_min, y_max);
+    for (int8_t i=0; i<n_hit; i++){
+        if (s[i] >= current_s){
+            double result = s[i];
+            free(s);
+            return result;
+        }
+    }
+    // No crossing
+    free(s);
+    return S_MAX;
+}
+
+
+#endif /* XCOLL_GEOM_GET_S_H */

xcoll/scattering_routines/geometry/methods.h

@@ -0,0 +1,111 @@
+// copyright ############################### #
+// This file is part of the Xcoll Package.   #
+// Copyright (c) CERN, 2024.                 #
+// ######################################### #
+
+#ifndef XCOLL_GEOM_METHODS_H
+#define XCOLL_GEOM_METHODS_H
+#include <math.h>
+#include <stdio.h>
+#include <stdint.h>
+#include <stdlib.h>
+
+#define S_MAX 1.e21
+
+
+/*gpufun*/
+void find_crossing(int8_t* n_hit, double* s, double part_x, double part_tan, \
+                       Segment* segments, int8_t n_segments){
+    for (int8_t i=0; i<n_segments;i++) {
+        segments[i]->crossing(n_hit, s, part_x, part_tan, segments[i]);
+    }
+    sort_array_of_double(s, (int64_t) *n_hit);
+}
+
+
+// IMPORTANT:
+// The array and interval are assumed to be sorted!
+// Furthermore, the array should have one extra slot allocated at the end, in case it needs to be expanded..
+// This is always true for the arrays created by get_s, as we create them with 2*n_segments slots.
+/*gpufun*/
+void calculate_overlap_array_interval(double* arr, int8_t* length, double* interval){
+    if (arr[0] > interval[1]){
+        // No overlap
+        *length = 0;
+    }
+    if ((*length)%2 == 1){
+        // Special case: last interval of array is open until infinity,
+        // so we add an extra point at the end to represent infinity.
+        arr[*length] = S_MAX*0.1;
+        (*length)++;
+    } else if (arr[*length-1] < interval[0]){
+        // No overlap
+        *length = 0;
+    }
+    int8_t i_start = 0;
+    // Find the start of overlap
+    for (int8_t i=0; i<*length; i++){
+        if (arr[i] >= interval[0]){
+            if (i%2 == 0){
+                // This is the first point of overlap
+                i_start = i;
+            } else {
+                // The vertical restriction is the first point of overlap
+                i_start = i-1;
+                arr[i_start] = interval[0];
+            }
+            break;
+        }
+    }
+    // Find the end of overlap
+    int8_t i_stop = *length - 1;
+    for (int8_t i=0; i<*length; i++){
+        if (arr[i] >= interval[1]){
+            if (i%2 == 0){
+                // The previous point is the last point of overlap
+                i_stop = i-1;
+            } else {
+                // The vertical restriction is the first point of overlap
+                i_stop = i;
+                arr[i_stop] = interval[1];
+            }
+            break;
+        }
+    }
+    *length = i_stop - i_start + 1;
+    if (i_start > 0){
+        for (int8_t i=0; i<*length; i++){
+            arr[i] = arr[i+i_start];
+        }
+    }
+}
+
+
+/*gpufun*/
+void find_crossing_with_vlimit(int8_t* n_hit, double* s, double part_x, double part_tan_x, \
+                               double part_y, double part_tan_y, Segment* segments, \
+                               int8_t n_segments, double y_min, double y_max){
+    if (fabs(part_tan_y) < 1.e-12){
+        // Trajectory parallel to s axis
+        if (part_y < y_min || part_y > y_max){
+            // No crossing
+            return;
+        } else {
+            // The particle is completely inside the vertical limits, so only check horizontal
+            find_crossing(n_hit, s, part_x, part_tan_x, segments, n_segments);
+            return;
+        }
+    } else {
+        find_crossing(n_hit, s, part_x, part_tan_x, segments, n_segments);
+        // restrict_s is the region [s0, s1] where the particle is inside the vertical limits
+        double* restrict_s = (double*) malloc(2*sizeof(double));
+        restrict_s[0] = (y_min - part_y)/part_tan_y;
+        restrict_s[1] = (y_max - part_y)/part_tan_y;
+        SWAP(restrict_s, 0, 1);   // To make sure these are sorted
+        calculate_overlap_array_interval(s, n_hit, restrict_s);
+        free(restrict_s);
+    }
+}
+
+
+#endif /* XCOLL_GEOM_METHODS_H */