xcoll 0.5.5__py3-none-any.whl → 0.5.6__py3-none-any.whl

xcoll/__init__.py

@@ -7,17 +7,26 @@ from .general import _pkg_root, __version__, citation
 
 from .beam_elements import BlackAbsorber, BlackCrystal, EverestBlock, EverestCollimator, EverestCrystal, \
                            BlowUp, EmittanceMonitor, collimator_classes, crystal_classes, element_classes
-from .install import install_elements
-from .line_tools import assign_optics_to_collimators, open_collimators, send_to_parking, enable_scattering, disable_scattering
 from .scattering_routines.everest import materials, Material, CrystalMaterial
 from .colldb import CollimatorDatabase
 from .interaction_record import InteractionRecord
 from .rf_sweep import RFSweep
-from .initial_distribution import *
 from .lossmap import LossMap
 
 # Deprecated
 from ._manager import CollimatorManager
+from .install import install_elements
+from .line_tools import assign_optics_to_collimators, open_collimators, send_to_parking, enable_scattering, disable_scattering
+def generate_pencil_on_collimator(line, name, *args, **kwargs):
+    from warnings import warn
+    warn("`xcoll.generate_pencil_on_collimator()` is deprecated and will be removed. Use "
+       + "`line[coll].generate_pencil()` instead.", FutureWarning)
+    return line[name].generate_pencil(*args, **kwargs)
+def generate_delta_from_dispersion(line, at_element, *args, **kwargs):
+    from warnings import warn
+    warn("`xcoll.generate_delta_from_dispersion()` is deprecated and will be removed. Use "
+       + "`line[at_element].generate_delta()` instead.", FutureWarning)
+    return line[name].generate_delta(*args, **kwargs)
 
 # print("If you use Xcoll in your simulations, please cite us :-)")
 # print(citation)

xcoll/beam_elements/__init__.py

@@ -21,5 +21,7 @@ crystal_classes = tuple(v for v in globals().values()
 
 element_classes = block_classes + (BlowUp, EmittanceMonitor)
 
+# For compatibility with _K2Collimators
+_all_block_classes      = block_classes
 _all_collimator_classes = collimator_classes
 _all_crystal_classes    = crystal_classes

xcoll/beam_elements/base.py

@@ -922,6 +922,31 @@ class BaseCollimator(BaseBlock):
     # Methods
     # =======
 
+    def generate_pencil(self, num_particles, *, side='+-', pencil_spread=1e-6,
+                        impact_parameter=0, sigma_z=7.61e-2, twiss=None, longitudinal=None,
+                        longitudinal_betatron_cut=None, tw=None, **kwargs):
+        if not hasattr(self, '_line') or not hasattr(self, '_name'):
+            raise ValueError("Collimator is missing a pointer to the line. Install collimators "
+                           + "with `line.collimators.install()` (or use "
+                           + "`xcoll.initial_distribution.generate_pencil_on_collimator()`).")
+        from xcoll.initial_distribution import generate_pencil_on_collimator
+        return generate_pencil_on_collimator(line=self._line, name=self._name, side=side,
+                        num_particles=num_particles, pencil_spread=pencil_spread, tw=tw,
+                        impact_parameter=impact_parameter, sigma_z=sigma_z, twiss=twiss,
+                        longitudinal=longitudinal, longitudinal_betatron_cut=longitudinal_betatron_cut,
+                        **kwargs)
+
+    def generate_delta(self, *, plane, position_mm, nemitt_x, nemitt_y, betatron_cut=0,
+                       match_at_front=True, twiss=None):
+        if not hasattr(self, '_line') or not hasattr(self, '_name'):
+            raise ValueError("Collimator is missing a pointer to the line. Install collimators "
+                           + "with `line.collimators.install()` (or use "
+                           + "`xcoll.initial_distribution.generate_delta_from_dispersion()`).")
+        from xcoll.initial_distribution import generate_delta_from_dispersion
+        return generate_delta_from_dispersion(line=self._line, at_element=self._name, plane=plane,
+                        position_mm=position_mm, nemitt_x=nemitt_x, nemitt_y=nemitt_y, twiss=twiss,
+                        betatron_cut=betatron_cut, match_at_front=match_at_front)
+
     def _verify_consistency(self):
         BaseBlock._verify_consistency(self)
         # Verify angles
@@ -1345,6 +1370,9 @@ class BaseCrystal(BaseBlock):
     # Methods
     # =======
 
+    def generate_pencil(self, **kwargs):
+        return BaseCollimator.generate_pencil(self, **kwargs)
+
     def _verify_consistency(self):
         BaseBlock._verify_consistency(self)
         # Verify angles

xcoll/colldb.py

@@ -14,7 +14,6 @@ import xtrack as xt
 
 from .beam_elements import BlackAbsorber, BlackCrystal, EverestCollimator, EverestCrystal, \
                            BaseCollimator, BaseCrystal, _all_collimator_classes
-from .install import install_elements
 from .scattering_routines.everest.materials import SixTrack_to_xcoll
 
 
@@ -572,7 +571,7 @@ class CollimatorDatabase:
             else:
                 self._create_crystal(line, BlackCrystal, name, verbose=verbose)
         elements = [self._elements[name] for name in names]
-        install_elements(line, names, elements, need_apertures=need_apertures)
+        line.collimators.install(names, elements, need_apertures=need_apertures)
 
     def install_everest_collimators(self, line, *, names=None, families=None, verbose=False, need_apertures=True):
         names = self._get_names_from_line(line, names, families)
@@ -586,7 +585,7 @@ class CollimatorDatabase:
                                      lattice=self[name]['crystal'], verbose=verbose,
                                      miscut=self[name]['miscut'])
         elements = [self._elements[name] for name in names]
-        install_elements(line, names, elements, need_apertures=need_apertures)
+        line.collimators.install(names, elements, need_apertures=need_apertures)
 
 
     # ==================================

xcoll/general.py

@@ -12,5 +12,5 @@ citation = "F.F. Van der Veken, et al.: Recent Developments with the New Tools f
 # ======================
 # Do not change
 # ======================
-__version__ = '0.5.5'
+__version__ = '0.5.6'
 # ======================

xcoll/initial_distribution.py

@@ -4,6 +4,7 @@
 # ######################################### #
 
 import numpy as np
+from warnings import warn
 
 import xtrack as xt
 import xobjects as xo
@@ -14,7 +15,7 @@ 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, twiss=None, longitudinal=None,
-                                  longitudinal_betatron_cut=None, tw=None):
+                                  longitudinal_betatron_cut=None, tw=None, **kwargs):
     """
     Generate a pencil beam on a collimator.
     """
@@ -49,7 +50,7 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
         raise NotImplementedError("Pencil beam on a skew collimator not yet supported!")
 
     if tw is not None:
-        print("The argument tw is deprecated. Please use twiss instead.")
+        warn("The argument tw is deprecated. Please use twiss instead.", FutureWarning)
         if twiss is None:
             twiss = tw
 
@@ -129,15 +130,13 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
         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=at_element, #match_at_s=match_at_s,
-                _context=coll._buffer.context
+                line=line, at_element=at_element, _context=coll._buffer.context, **kwargs
         )
     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=at_element, #match_at_s=match_at_s,
-                _context=coll._buffer.context
+                line=line, at_element=at_element, _context=coll._buffer.context, **kwargs
         )
 
     part._init_random_number_generator()

xcoll/install.py

@@ -3,177 +3,33 @@
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
-import numpy as np
-import xtrack as xt
+from warnings import warn
 
-from .beam_elements import element_classes
+# This file is deprecated and will be removed in the future.
 
-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 install_elements(line, names, elements, *, at_s=None, apertures=None, need_apertures=False, s_tol=1.e-6):
+    warn("The function xcoll.install_elements() is deprecated and will be removed in the future. "
+       + "Please use line.collimators.install() instead.", FutureWarning)
+    line.collimators.install(names=names, elements=elements, at_s=at_s, apertures=apertures,
+                             need_apertures=need_apertures, s_tol=s_tol)
 
 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.")
-
+    warn("The function xcoll.check_element_position() is deprecated and will be removed "
+       + "in the future. Please use line.collimators.check_position() instead.", FutureWarning)
+    return line.collimators.check_position(name, s_start, s_end, at_s, length=length, s_tol=s_tol)
 
 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.")
-
+    warn("The function xcoll.get_apertures_at_s() is deprecated and will be removed in the future. "
+       + "Please use line.collimators.get_apertures_at_s() instead.", FutureWarning)
+    return get_apertures_at_s(s, s_tol=s_tol)
 
 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()
-
+    warn("The function xcoll.get_aperture_for_element() is deprecated and will be removed in the future. "
+       + "Please use line.collimators.get_aperture() instead.", FutureWarning)
+    return line.get_aperture(name, s_start, s_end, aperture=aperture, tab=tab, s_tol=s_tol)
 
 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)
-
+    warn("The function xcoll.prepare_space_for_element() is deprecated and will be removed in the future. "
+       + "Please use line.collimators.prepare_space() instead.", FutureWarning)
+    return line.collimators.prepare_space(name, s_start, s_end, tab=tab, s_tol=s_tol)