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

xcoll/general.py

@@ -1,6 +1,6 @@
 # copyright ############################### #
 # This file is part of the Xcoll Package.   #
-# Copyright (c) CERN, 2023.                 #
+# Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
 from pathlib import Path
@@ -9,8 +9,8 @@ _pkg_root = Path(__file__).parent.absolute()
 
 citation = "F.F. Van der Veken, et al.: Recent Developments with the New Tools for Collimation Simulations in Xsuite, Proceedings of HB2023, Geneva, Switzerland."
 
-# ===================
+# ======================
 # Do not change
-# ===================
-__version__ = '0.3.6'
-# ===================
+# ======================
+__version__ = '0.5.0'
+# ======================

xcoll/headers/checks.h

@@ -31,5 +31,5 @@ int8_t xcoll_check_particle_init(RandomRutherfordData rng, LocalParticle* part)
     }
     return is_tracking*rng_is_set*ruth_is_set;
 }
-    
+
 #endif /* XCOLL_CHECKS_H */

xcoll/headers/particle_states.h

@@ -1,6 +1,6 @@
 // copyright ############################### #
 // This file is part of the Xcoll Package.   #
-// Copyright (c) CERN, 2023.                 #
+// Copyright (c) CERN, 2024.                 #
 // ######################################### #
 
 #ifndef XCOLL_STATES_H
@@ -21,5 +21,5 @@
 #define  XC_ERR_NOT_IMPLEMENTED     -391
 #define  XC_ERR_INVALID_XOFIELD     -392
 #define  XC_ERR                     -399
-    
+
 #endif /* XCOLL_STATES_H */

xcoll/initial_distribution.py

@@ -0,0 +1,207 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import numpy as np
+
+import xtrack as xt
+import xobjects as xo
+import xpart as xp
+
+from .beam_elements import _all_collimator_classes, EverestCrystal
+
+
+def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', pencil_spread=1e-6,
+                                  impact_parameter=0, sigma_z=7.61e-2, tw=None, longitudinal=None,
+                                  longitudinal_betatron_cut=None):
+    """
+    Generate a pencil beam on a collimator.
+    """
+
+    if not line._has_valid_tracker():
+        raise Exception("Please build tracker before generating pencil distribution!")
+
+    coll = line[name]
+
+    if not isinstance(coll, tuple(_all_collimator_classes)):
+        raise ValueError("Need to provide a valid collimator!")
+
+    if coll.optics is None:
+        raise Exception("Need to assign optics to collimators before generating pencil distribution!")
+
+    num_particles = int(num_particles)
+
+    if coll.side == 'left':
+        side = '+'
+    if coll.side == 'right':
+        side = '-'
+
+    # Define the plane
+    angle = coll.angle
+    if abs(np.mod(angle-90,180)-90) < 1e-6:
+        plane = 'x'
+        transv_plane = 'y'
+    elif abs(np.mod(angle,180)-90) < 1e-6:
+        plane = 'y'
+        transv_plane = 'x'
+    else:
+        raise NotImplementedError("Pencil beam on a skew collimator not yet supported!")
+
+    if tw is None:
+        tw = line.twiss()    # TODO: can we do this smarter by caching?
+
+    # Is it converging or diverging?    # TODO: This might change with a tilt!!!!!!
+    s_front = line.get_s_position(name)
+    s_back  = s_front + coll.length
+    is_converging  = tw[f'alf{plane}', name] > 0
+    print(f"Collimator {name} is {'con' if is_converging else 'di'}verging.")
+
+    beam_sizes = tw.get_beam_covariance(nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y)
+    if is_converging:
+        # pencil at front of jaw
+        match_at_s = s_front
+        sigma = beam_sizes.rows[name:f'{name}%%1'][f'sigma_{plane}'][0]
+        sigma_transv = beam_sizes.rows[name:f'{name}%%1'][f'sigma_{transv_plane}'][0]
+    else:
+        # pencil at back of jaw
+        match_at_s = s_back
+        sigma = beam_sizes.rows[name:f'{name}%%1'][f'sigma_{plane}'][1]
+        sigma_transv = beam_sizes.rows[name:f'{name}%%1'][f'sigma_{transv_plane}'][1]
+    dr_sigmas = pencil_spread/sigma
+
+    # Generate 4D coordinates
+    # TODO: there is some looping in the calculation here and in xpart. Can it be improved?
+    if side == '+-':
+        num_plus = int(num_particles/2)
+        num_min  = int(num_particles - num_plus)
+        coords_plus = _generate_4D_pencil_one_jaw(line, name, num_plus, plane, '+', impact_parameter, dr_sigmas, match_at_s, is_converging)
+        coords_min  = _generate_4D_pencil_one_jaw(line, name, num_min,  plane, '-', impact_parameter, dr_sigmas, match_at_s, is_converging)
+        coords      = [ [*c_plus, *c_min] for c_plus, c_min in zip(coords_plus, coords_min)]
+    else:
+        coords      = _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter, dr_sigmas, match_at_s, is_converging)
+    pencil            = coords[0]
+    p_pencil          = coords[1]
+    transverse_norm   = coords[2]
+    p_transverse_norm = coords[3]
+
+    # Longitudinal plane
+    # TODO: make this more general, make this better
+    if longitudinal is None:
+        delta = 0
+        zeta  = 0
+    elif longitudinal == 'matched_dispersion':
+        raise NotImplementedError
+        # if longitudinal_betatron_cut is None:
+        #     cut = 0
+        # else:
+        #     cut = np.random.uniform(-longitudinal_betatron_cut, longitudinal_betatron_cut,
+        #                             num_particles)
+        # delta = generate_delta_from_dispersion(line, name, plane=plane, position_mm=pencil,
+        #                                        nemitt_x=nemitt_x, nemitt_y=nemitt_y, twiss=tw,
+        #                                        betatron_cut=cut, match_at_front=is_converging)
+        # zeta  = 0
+    elif longitudinal == 'bucket':
+        zeta, delta = xp.generate_longitudinal_coordinates(
+                num_particles=num_particles, distribution='gaussian', sigma_z=sigma_z, line=line
+        )
+    elif not hasattr(longitudinal, '__iter__'):
+        raise ValueError
+    elif len(longitudinal) != 2:
+        raise ValueError
+    elif isinstance(longitudinal, str):
+        raise ValueError
+    elif isinstance(longitudinal, dict):
+        zeta = longitudinal['zeta']
+        delta = longitudinal['delta']
+    else:
+        zeta = longitudinal[0]
+        delta = longitudinal[1]
+
+    # Build the particles
+    if plane == 'x':
+        part = xp.build_particles(
+                x=pencil, px=p_pencil, y_norm=transverse_norm, py_norm=p_transverse_norm,
+                zeta=zeta, delta=delta, nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y,
+                line=line, at_element=name, match_at_s=match_at_s,
+                _context=coll._buffer.context
+        )
+    else:
+        part = xp.build_particles(
+                x_norm=transverse_norm, px_norm=p_transverse_norm, y=pencil, py=p_pencil, 
+                zeta=zeta, delta=delta, nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y,
+                line=line, at_element=name, match_at_s=match_at_s,
+                _context=coll._buffer.context
+        )
+
+    part._init_random_number_generator()
+
+    return part
+
+
+def generate_delta_from_dispersion(line, at_element, *, plane, position_mm, nemitt_x, nemitt_y,
+                                   twiss=None, betatron_cut=0, match_at_front=True):
+    if line.tracker is None:
+        raise ValueError("Need to build tracker first!")
+    if not hasattr(betatron_cut, '__iter__'):
+        if hasattr(position_mm, '__iter__'):
+            betatron_cut = np.full_like(position_mm, betatron_cut)
+    elif not hasattr(position_mm, '__iter__'):
+        position_mm = np.full_like(betatron_cut, position_mm)
+    elif len(position_mm) != len(betatron_cut):
+        raise ValueError
+    if plane not in ['x', 'y']:
+        raise ValueError("The variable 'plane' needs to be either 'x' or 'y'!")
+
+    if twiss is None:
+        twiss = line.twiss()
+
+    beam_sizes = twiss.get_beam_covariance(nemitt_x=nemitt_x, nemitt_y=nemitt_y)
+    beam_sizes = beam_sizes.rows[at_element:f'{at_element}%%1'][f'sigma_{plane}']
+    sigma = beam_sizes[0] if match_at_front else beam_sizes[1]
+    delta  = (position_mm - betatron_cut*sigma - twiss.rows[at_element][plane])
+    delta /= twiss.rows[at_element][f'd{plane}']
+
+    return delta
+
+
+def _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side, impact_parameter,
+                                dr_sigmas, match_at_s, is_converging):
+    coll = line[name]
+
+    if side == '+':
+        if is_converging:
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_U + impact_parameter
+            else:
+                pencil_pos = coll.jaw_LU + impact_parameter
+        else:
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_D - impact_parameter
+            else:
+                pencil_pos = coll.jaw_LD + impact_parameter
+    elif side == '-':
+        if is_converging:
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_U - impact_parameter
+            else:
+                pencil_pos = coll.jaw_RU - impact_parameter
+        else:
+            if isinstance(coll, EverestCrystal):
+                pencil_pos = coll.jaw_D + impact_parameter
+            else:
+                pencil_pos = coll.jaw_RD - impact_parameter
+
+    # Collimator plane: generate pencil distribution
+    pencil, p_pencil = xp.generate_2D_pencil_with_absolute_cut(
+                        num_particles, plane=plane, absolute_cut=pencil_pos, line=line,
+                        dr_sigmas=dr_sigmas, nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y,
+                        at_element=name, side=side,match_at_s=match_at_s
+                       )
+
+    # Other plane: generate gaussian distribution in normalized coordinates
+    transverse_norm   = np.random.normal(size=num_particles)
+    p_transverse_norm = np.random.normal(size=num_particles)
+
+    return [pencil, p_pencil, transverse_norm, p_transverse_norm]
+

xcoll/install.py

@@ -0,0 +1,179 @@
+# copyright ############################### #
+# This file is part of the Xcoll Package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import numpy as np
+import xtrack as xt
+
+from .beam_elements import element_classes
+
+def install_elements(line, names, elements, *, at_s=None, apertures=None, need_apertures=False, s_tol=1.e-6):
+    if line._has_valid_tracker():
+        raise Exception("Tracker already built!\nPlease install collimators before building "
+                      + "tracker!")
+
+    if not hasattr(names, '__iter__') or isinstance(names, str):
+        names = [names]
+    if not hasattr(elements, '__iter__') or isinstance(elements, str):
+        elements = [elements]
+    names = np.array(names)
+    length = np.array([coll.length for coll in elements])
+    assert len(length) == len(names)
+    if not hasattr(at_s, '__iter__'):
+        at_s = [at_s for _ in range(len(names))]
+    assert len(at_s) == len(names)
+    if isinstance(apertures, str) or not hasattr(apertures, '__iter__'):
+        apertures = [apertures for _ in range(len(names))]
+    assert len(apertures) == len(names)
+
+    # Verify elements
+    for el in elements:
+        assert isinstance(el, element_classes)
+        el._tracking = False
+
+    # Get positions
+    tab = line.get_table()
+    tt = tab.rows[[name for name in names if name in line.element_names]]
+    s_start = []
+    for name, s, l in zip(names, at_s, length):
+        if s is None:
+            s_start.append(_get_s_start(line, name, l, tt))
+        else:
+            s_start.append(s)
+    s_start = np.array(s_start)
+    s_end = s_start + length
+
+    # Check positions
+    l_line = line.get_length() 
+    for s1, s2, name, s3 in zip(s_start, s_end, names, at_s):
+        check_element_position(line, name, s1, s2, s3, l_line, s_tol=s_tol)
+
+    # Look for apertures
+    aper_upstream   = []
+    aper_downstream = []
+    for s1, s2, name, aper in zip(s_start, s_end, names, apertures):
+        if not need_apertures:
+            aper_upstream.append(None)
+            aper_downstream.append(None)
+        else:
+            aper1, aper2 = get_aperture_for_element(line, name, s1, s2, aper, tab, s_tol=s_tol)
+            aper_upstream.append(aper1)
+            aper_downstream.append(aper2)
+
+    # Remove elements at location of collimator (by changing them into markers)
+    for s1, s2, name in zip(s_start, s_end, names):
+        prepare_space_for_element(line, name, s1, s2, tab=tab, s_tol=s_tol)
+
+    # Install
+    line._insert_thick_elements_at_s(element_names=list(names), elements=elements, at_s=s_start, s_tol=s_tol)
+
+    # Install apertures
+    if need_apertures:
+        for s1, name, aper1, aper2 in zip(s_start, names, aper_upstream, aper_downstream):
+            line.insert_element(element=aper1, name=f'{name}_aper_upstream', at=name, s_tol=s_tol)
+            idx = line.element_names.index(name) + 1
+            line.insert_element(element=aper2, name=f'{name}_aper_downstream', at=idx, s_tol=s_tol)
+
+
+def _get_s_start(line, name, length, tab=None):
+    if tab is None:
+        tab = line.get_table()
+    if name in line.element_names and hasattr(line[name], 'length'):
+        existing_length = line[name].length
+    else:
+        existing_length = 0
+    if name not in tab.name:
+        raise ValueError(f"Element {name} not found in line. Need to mmanually provide `at_s`.")
+    return tab.rows[name].s[0] + existing_length/2. - length/2
+
+
+def check_element_position(line, name, s_start, s_end, at_s, length=None, s_tol=1.e-6):
+    if at_s is None:
+        if name not in line.element_names:
+            raise ValueError(f"Element {name} not found in line. Provide `at_s`.")
+    elif name in line.element_names:
+        if at_s < s_start or at_s > s_end:
+            raise ValueError(f"Element {name} already exists in line at different "
+                           + f"location: at_s = {at_s}, exists at [{s_start}, {s_end}].")
+    if length is None:
+        length = line.get_length() 
+    if s_start <= s_tol:
+        raise ValueError(f"Position of {name} too close to start of line. Please cycle.")
+    if s_end >= length - s_tol:
+        raise ValueError(f"Position of {name} too close to end of line. Please cycle.")
+
+
+def get_apertures_at_s(tab, s, s_tol=1.e-6):
+    tab_s = tab.rows[s-s_tol:s+s_tol:'s']
+    aper = tab_s.rows[[cls.startswith('Limit') for cls in tab_s.element_type]]
+    if len(aper) == 0:
+        return None
+    elif len(aper) == 1:
+        return aper.name[0]
+    else:
+        raise ValueError(f"Multiple apertures found at location {s} with "
+                       + f"tolerance {s_tol}: {aper.name}. Not supported.")
+
+
+def get_aperture_for_element(line, name, s_start, s_end, aperture=None, tab=None, s_tol=1.e-6):
+    if aperture is not None:
+        if isinstance(aperture, str):
+            aper1 = line[aperture]
+            aper2 = line[aperture]
+        elif hasattr(aperture, '__iter__'):
+            if len(aperture) != 2:
+                raise ValueError(f"The value `aperture` should be None or a list "
+                               + f"[upstream, downstream].")
+            assert aperture[0] is not None and aperture[1] is not None
+            if isinstance(aperture[0], str):
+                aper1 = line[aperture[0]]
+            if isinstance(aperture[1], str):
+                aper2 = line[aperture[1]]
+        else:
+            aper1 = aperture
+            aper2 = aperture
+        if not xt.line._is_aperture(aper1, line):
+            raise ValueError(f"Not a valid aperture: {aper1}")
+        if not xt.line._is_aperture(aper2, line):
+            raise ValueError(f"Not a valid aperture: {aper2}")
+        return aper1.copy(), aper2.copy()
+    else:
+        if tab is None:
+            tab = line.get_table()
+        aper1 = get_apertures_at_s(tab, s_start, s_tol=s_tol)
+        aper2 = get_apertures_at_s(tab, s_end, s_tol=s_tol)
+        if aper1 is None and aper2 is not None:
+            aper1 = aper2
+            print(f"Warning: Could not find upstream aperture for {name}! "
+                + f"Used copy of downstream aperture. Proceed with caution.")
+        elif aper2 is None and aper1 is not None:
+            aper2 = aper1
+            print(f"Warning: Could not find downstream aperture for {name}! "
+                + f"Used copy of upstream aperture. Proceed with caution.")
+        elif aper1 is None and aper2 is None:
+            aper_mid = get_apertures_at_s(tab, (s_start+s_end)/2, s_tol=s_tol)
+            if aper_mid is None:
+                raise ValueError(f"No aperture found for {name}! Please provide one.")
+            if line[aper_mid].allow_rot_and_shift \
+            and xt.base_element._tranformations_active(line[aper_mid]):
+                print(f"Warning: Using the centre aperture for {name}, but "
+                    + f"transformations are present. Proceed with caution.")
+            aper1 = aper_mid
+            aper2 = aper_mid
+        return line[aper1].copy(), line[aper2].copy()
+
+
+def prepare_space_for_element(line, name, s_start, s_end, tab=None, s_tol=1.e-6):
+    if tab is None:
+        tab = line.get_table()
+    tt = tab.rows[s_start-s_tol:s_end+s_tol:'s']
+    for element_name, element_type in zip(tt.name[:-1], tt.element_type[:-1]):
+        if element_type == 'Marker' or element_type.startswith('Drift'):
+            continue
+        if not element_type.startswith('Limit'):
+            print(f"Warning: Removed active element {element_name} "
+                + f"at location inside collimator {name}!")
+        length = line[element_name].length if hasattr(line[element_name], 'length') else 0
+        line.element_dict[element_name] = xt.Drift(length=length)
+

xcoll/interaction_record/__init__.py

@@ -0,0 +1 @@
+from .interaction_record import InteractionRecord