xcoll 0.5.5__tar.gz → 0.5.7__tar.gz

{xcoll-0.5.5 → xcoll-0.5.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xcoll
-Version: 0.5.5
+Version: 0.5.7
 Summary: Xsuite collimation package
 Home-page: https://github.com/xsuite/xcoll
 License: Apache 2.0

{xcoll-0.5.5 → xcoll-0.5.7}/pyproject.toml

@@ -1,17 +1,19 @@
 [tool.poetry]
 name = "xcoll"
-version = "0.5.5"
+version = "0.5.7"
 description = "Xsuite collimation package"
 homepage = "https://github.com/xsuite/xcoll"
 repository = "https://github.com/xsuite/xcoll"
 authors = [
            "Frederik F. Van der Veken <frederik@cern.ch>",
            "Simone O. Solstrand <simone.otelie.solstrand@cern.ch>",
-           "Bjorn Lindstrom",
-           "Andre Donadon Servelle",
+           "Björn Lindstrom <bjorn.lindstrom@cern.ch>",
+           "Giacomo Broggi <giacomo.broggi@cern.ch>",
+           "André Donadon Servelle <andre.donadon.servelle@cern.ch>",
+           "Dora E. Veres <dora.erzsebet.veres@cern.ch>",
            "Despina Demetriadou <despina.demetriadou@cern.ch>",
            "Andrey Abramov <andrey.abramov@cern.ch>",
-           "Giovanni Iadorala <giovanni.iadorala@cern.ch>"
+           "Giovanni Iadarola <giovanni.iadarola@cern.ch>"
 ]
 readme = "README.md"
 license = "Apache 2.0"
@@ -42,9 +44,7 @@ build-backend = "poetry.core.masonry.api"
 
 # pyproject.toml
 [tool.pytest.ini_options]
-addopts = "-ra --durations=10 --durations-min=1 --junitxml=pytest_results.xml"
-junit_suite_name = "Xcoll"
-junit_duration_report = "call"
+addopts = "-ra --durations=10 --durations-min=1"
 python_functions = ["test_"]
 testpaths = [
     "tests",

{xcoll-0.5.5 → xcoll-0.5.7}/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-0.5.5 → xcoll-0.5.7}/xcoll/beam_elements/__init__.py

@@ -20,6 +20,3 @@ crystal_classes = tuple(v for v in globals().values()
                         if isinstance(v, type) and issubclass(v, BaseCrystal) and v != BaseCrystal)
 
 element_classes = block_classes + (BlowUp, EmittanceMonitor)
-
-_all_collimator_classes = collimator_classes
-_all_crystal_classes    = crystal_classes

xcoll-0.5.7/xcoll/beam_elements/__init__.py.orig

@@ -0,0 +1,34 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+from .base import BaseBlock, BaseCollimator, BaseCrystal
+from .absorber import BlackAbsorber, BlackCrystal
+from .everest import EverestBlock, EverestCollimator, EverestCrystal
+from .blowup import BlowUp
+from .monitor import EmittanceMonitor
+
+block_classes = tuple(v for v in globals().values()
+                      if isinstance(v, type) and issubclass(v, BaseBlock) and v != BaseBlock
+                      and v != BaseCollimator and v != BaseCrystal)
+# Includes crystals
+collimator_classes = tuple(v for v in globals().values()
+                           if isinstance(v, type) and (issubclass(v, BaseCollimator) or issubclass(v, BaseCrystal))
+                           and v != BaseCollimator and v != BaseCrystal)
+crystal_classes = tuple(v for v in globals().values()
+                        if isinstance(v, type) and issubclass(v, BaseCrystal) and v != BaseCrystal)
+
+element_classes = block_classes + (BlowUp, EmittanceMonitor)
+
+<<<<<<< HEAD
+# These should not go into any of the classes lists
+from .k2 import _K2Collimator, _K2Crystal
+_all_collimator_classes = collimator_classes + (_K2Collimator, _K2Crystal)
+_all_crystal_classes    = crystal_classes + (_K2Crystal,)
+=======
+# For compatibility with _K2Collimators
+_all_block_classes      = block_classes
+_all_collimator_classes = collimator_classes
+_all_crystal_classes    = crystal_classes
+>>>>>>> release/v0.6.0

{xcoll-0.5.5 → xcoll-0.5.7}/xcoll/beam_elements/base.py

@@ -597,9 +597,6 @@ class BaseCollimator(BaseBlock):
 
     def assign_optics(self, *, nemitt_x=None, nemitt_y=None, beta_gamma_rel=None, name=None, twiss=None,
                       twiss_upstream=None, twiss_downstream=None):
-        from xcoll.beam_elements import _all_collimator_classes
-        if not isinstance(self, _all_collimator_classes):
-            raise ValueError("Please install collimator before assigning optics.")
         if nemitt_x is None:
             if self.nemitt_x is None:
                 raise ValueError("Need to provide `nemitt_x`.")
@@ -922,6 +919,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 +1367,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-0.5.5 → xcoll-0.5.7}/xcoll/colldb.py

@@ -13,8 +13,7 @@ from pathlib import Path
 import xtrack as xt
 
 from .beam_elements import BlackAbsorber, BlackCrystal, EverestCollimator, EverestCrystal, \
-                           BaseCollimator, BaseCrystal, _all_collimator_classes
-from .install import install_elements
+                           BaseCollimator, BaseCrystal, collimator_classes
 from .scattering_routines.everest.materials import SixTrack_to_xcoll
 
 
@@ -527,7 +526,7 @@ class CollimatorDatabase:
     def _check_installed(self, line, name, collimator_class):
             # Check that collimator is not installed as different type
             # TODO: automatically replace collimator type and print warning
-            if isinstance(line[name], _all_collimator_classes):
+            if isinstance(line[name], collimator_classes):
                 raise ValueError(f"Trying to install {name} as {collimator_class.__name__}, "
                                + f"but it is already installed as {line[name].__class__.__name__}!\n"
                                + f"Please reconstruct the line.")
@@ -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-0.5.5 → xcoll-0.5.7}/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.7'
 # ======================

{xcoll-0.5.5 → xcoll-0.5.7}/xcoll/initial_distribution.py

@@ -4,17 +4,18 @@
 # ######################################### #
 
 import numpy as np
+from warnings import warn
 
 import xtrack as xt
 import xobjects as xo
 import xpart as xp
 
-from .beam_elements import _all_collimator_classes, EverestCrystal
+from .beam_elements import 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.
     """
@@ -24,7 +25,7 @@ def generate_pencil_on_collimator(line, name, num_particles, *, side='+-', penci
 
     coll = line[name]
 
-    if not isinstance(coll, tuple(_all_collimator_classes)):
+    if not isinstance(coll, tuple(collimator_classes)):
         raise ValueError("Need to provide a valid collimator!")
 
     if coll.optics is None:
@@ -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-0.5.7/xcoll/install.py

@@ -0,0 +1,35 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+from warnings import warn
+
+# 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):
+    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):
+    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):
+    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):
+    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):
+    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)

xcoll-0.5.7/xcoll/line_tools.py

@@ -0,0 +1,311 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+import numpy as np
+from warnings import warn
+
+import xtrack as xt
+
+from .beam_elements import element_classes, collimator_classes, block_classes
+
+
+class XcollScatteringAPI:
+    def __init__(self, line):
+        self._line = line
+
+    @property
+    def line(self):
+        return self._line
+
+    def enable(self):
+        elements = self.line.get_elements_of_type(element_classes)[0]
+        if len(elements) == 0:
+            print("No xcoll elements found in line.")
+        else:
+            nemitt_x = None
+            nemitt_y = None
+            for el in elements:
+                if hasattr(el, 'optics') and el.optics is not None:
+                    if nemitt_x is None:
+                        nemitt_x = el.nemitt_x
+                    if nemitt_y is None:
+                        nemitt_y = el.nemitt_y
+                    if not np.isclose(el.nemitt_x, nemitt_x) \
+                    or not np.isclose(el.nemitt_x, nemitt_x):
+                        raise ValueError("Not all collimators have the same "
+                                    + "emittance. This is not supported.")
+                if hasattr(el, 'enable_scattering'):
+                    el.enable_scattering()
+
+    def disable(self):
+        elements = self.line.get_elements_of_type(element_classes)[0]
+        if len(elements) == 0:
+            print("No xcoll elements found in line.")
+        else:
+            for el in elements:
+                if hasattr(el, 'disable_scattering'):
+                    el.disable_scattering()
+
+
+class XcollCollimatorAPI:
+    def __init__(self, line):
+        self._line = line
+
+    @property
+    def line(self):
+        return self._line
+
+    def install(self, names, elements, *, at_s=None, apertures=None, need_apertures=False, s_tol=1.e-6):
+        if self.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:
+            print(el.__class__)
+            assert isinstance(el, block_classes)
+            el._tracking = False
+
+        # Get positions
+        tab = self.line.get_table()
+        tt = tab.rows[[name for name in names if name in self.line.element_names]]
+        s_start = []
+        for name, s, l in zip(names, at_s, length):
+            if s is None:
+                s_start.append(self._get_s_start(name, length=l, table=tt))
+            else:
+                s_start.append(s)
+        s_start = np.array(s_start)
+        s_end = s_start + length
+
+        # Check positions
+        l_line = self.line.get_length()
+        for s1, s2, name, s3 in zip(s_start, s_end, names, at_s):
+            self.check_position(name, s_start=s1, s_end=s2, at_s=s3, length=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 = self.get_aperture(name, s_start=s1, s_end=s2, aperture=aper, table=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, el in zip(s_start, s_end, names, elements):
+            self.prepare_space(name, s_start=s1, s_end=s2, table=tab, s_tol=s_tol)
+            el._line = self.line
+            el._name = name
+
+        # Install
+        self.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):
+                self.line.insert_element(element=aper1, name=f'{name}_aper_upstream', at=name, s_tol=s_tol)
+                idx = self.line.element_names.index(name) + 1
+                self.line.insert_element(element=aper2, name=f'{name}_aper_downstream', at=idx, s_tol=s_tol)
+
+    def check_position(self, name, *, s_start, s_end, at_s, length=None, s_tol=1.e-6):
+        if at_s is None:
+            if name not in self.line.element_names:
+                raise ValueError(f"Element {name} not found in line. Provide `at_s`.")
+        elif name in self.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 = self.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(self, s, *, table=None, s_tol=1.e-6):
+        if table is None:
+            table = self.line.get_table()
+        tab_s = table.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(self, name, *, s_start, s_end, aperture=None, table=None, s_tol=1.e-6):
+        if aperture is not None:
+            if isinstance(aperture, str):
+                aper1 = self.line[aperture]
+                aper2 = self.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 = self.line[aperture[0]]
+                if isinstance(aperture[1], str):
+                    aper2 = self.line[aperture[1]]
+            else:
+                aper1 = aperture
+                aper2 = aperture
+            if not xt.line._is_aperture(aper1, self.line):
+                raise ValueError(f"Not a valid aperture: {aper1}")
+            if not xt.line._is_aperture(aper2, self.line):
+                raise ValueError(f"Not a valid aperture: {aper2}")
+            return aper1.copy(), aper2.copy()
+        else:
+            if table is None:
+                table = self.line.get_table()
+            aper1 = self.get_apertures_at_s(s=s_start, table=table, s_tol=s_tol)
+            aper2 = self.get_apertures_at_s(s=s_end, table=table, 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 = self.get_apertures_at_s(s=(s_start+s_end)/2, table=table, s_tol=s_tol)
+                if aper_mid is None:
+                    raise ValueError(f"No aperture found for {name}! Please provide one.")
+                if self.line[aper_mid].allow_rot_and_shift \
+                and xt.base_element._tranformations_active(self.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 self.line[aper1].copy(), self.line[aper2].copy()
+
+    def prepare_space(self, name, *, s_start, s_end, table=None, s_tol=1.e-6):
+        if table is None:
+            table = self.line.get_table()
+        tt = table.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 = self.line[element_name].length if hasattr(self.line[element_name], 'length') else 0
+            self.line.element_dict[element_name] = xt.Drift(length=length)
+
+    def get_optics_at(self, names, *, twiss=None, tw=None):
+        if tw is not None:
+            warn("The argument tw is deprecated. Please use twiss instead.", FutureWarning)
+            if twiss is None:
+                twiss = tw
+        if twiss is None:
+            if not self.line._has_valid_tracker():
+                raise Exception("Please build the tracker before computing the optics for the openings!")
+            twiss = self.line.twiss()
+        if not hasattr(names, '__iter__') and not isinstance(names, str):
+            names = [names]
+        coll_entry_indices = twiss.rows.indices[names]
+        tw_entry = twiss.rows[coll_entry_indices]
+        tw_exit = twiss.rows[coll_entry_indices+1]
+        tw_exit.name = tw_entry.name
+        return tw_entry, tw_exit
+
+    def assign_optics(self, *, nemitt_x=None, nemitt_y=None, twiss=None, tw=None):
+        if tw is not None:
+            warn("The argument tw is deprecated. Please use twiss instead.", FutureWarning)
+            if twiss is None:
+                twiss = tw
+        if not self.line._has_valid_tracker():
+            raise Exception("Please build tracker before setting the openings!")
+        names = self.line.get_elements_of_type(collimator_classes)[1]
+        tw_upstream, tw_downstream = self.get_optics_at(names, twiss=twiss)
+        beta_gamma_rel = self.line.particle_ref._xobject.gamma0[0]*self.line.particle_ref._xobject.beta0[0]
+        for coll in names:
+            self.line[coll].assign_optics(name=coll, nemitt_x=nemitt_x, nemitt_y=nemitt_x, twiss_upstream=tw_upstream,
+                                    twiss_downstream=tw_downstream, beta_gamma_rel=beta_gamma_rel)
+
+    def open(self, names=None):
+        if names is None:
+            names = self.line.get_elements_of_type(collimator_classes)[1]
+        if len(names) == 0:
+            print("No collimators found in line.")
+        else:
+            for coll in names:
+                self.line[coll].open_jaws(keep_tilts=False)
+                self.line[coll].gap = None
+
+    def to_parking(self, names=None):
+        if names is None:
+            names = self.line.get_elements_of_type(collimator_classes)[1]
+        if len(names) == 0:
+            print("No collimators found in line.")
+        else:
+            raise NotImplementedError("Need to move this to new type manager or so.")
+
+    def _get_s_start(self, name, *, length, table=None):
+        if table is None:
+            table = self.line.get_table()
+        if name in self.line.element_names and hasattr(self.line[name], 'length'):
+            existing_length = self.line[name].length
+        else:
+            existing_length = 0
+        if name not in table.name:
+            raise ValueError(f"Element {name} not found in line. Need to manually provide `at_s`.")
+        return table.rows[name].s[0] + existing_length/2. - length/2
+
+
+
+# Deprecated; to be removed
+# -------------------------
+
+def assign_optics_to_collimators(line, nemitt_x=None, nemitt_y=None, twiss=None):
+    warn("The function xcoll.assign_optics_to_collimators() is deprecated and will be "
+       + "removed in the future. Please use line.scattering.assign_optics() instead.", FutureWarning)
+    line.collimators.assign_optics(nemitt_x=nemitt_x, nemitt_y=nemitt_y, twiss=twiss)
+
+def get_optics_at(names, *, twiss=None, line=None):
+    warn("The function xcoll.get_optics_at() is deprecated and will be "
+       + "removed in the future. Please use line.scattering.get_optics_at() instead.", FutureWarning)
+    return line.collimators.get_optics_at(names=names, twiss=twiss)
+
+def open_collimators(line, names=None):
+    warn("The function xcoll.open_collimators() is deprecated and will be "
+       + "removed in the future. Please use line.scattering.open_collimators() instead.", FutureWarning)
+    line.collimators.open(names=names)
+
+def send_to_parking(line, names=None):
+    warn("The function xcoll.send_to_parking() is deprecated and will be "
+       + "removed in the future. Please use line.scattering.send_to_parking() instead.", FutureWarning)
+    line.collimators.to_parking(names=names)
+
+def enable_scattering(line):
+    warn("The function xcoll.enable_scattering() is deprecated and will be "
+       + "removed in the future. Please use line.scattering.enable() instead.", FutureWarning)
+    line.scattering.enable()
+
+def disable_scattering(line):
+    warn("The function xcoll.disable_scattering() is deprecated and will be "
+       + "removed in the future. Please use line.scattering.disable() instead.", FutureWarning)
+    line.scattering.disable()

{xcoll-0.5.5 → xcoll-0.5.7}/xcoll/lossmap.py

@@ -12,7 +12,7 @@ import xtrack as xt
 import xpart as xp
 import xobjects as xo
 
-from .beam_elements import _all_collimator_classes, _all_crystal_classes
+from .beam_elements import collimator_classes, crystal_classes
 
 
 class LossMap:
@@ -114,7 +114,7 @@ class LossMap:
     def _correct_absorbed(self, verbose=True):
         # Correct particles that are at an aperture directly after a collimator
         part = self._part
-        coll_classes = list(set(_all_collimator_classes) - set(_all_crystal_classes))
+        coll_classes = list(set(collimator_classes) - set(crystal_classes))
         coll_elements = self._line.get_elements_of_type(coll_classes)[1]
         for idx, elem in enumerate(part.at_element):
             if part.state[idx] == 0:
@@ -156,7 +156,7 @@ class LossMap:
 
 
     def _make_coll_summary(self):
-        collimator_names = self._line.get_elements_of_type(_all_collimator_classes)[1]
+        collimator_names = self._line.get_elements_of_type(collimator_classes)[1]
         coll_mask     = (self._part.state <= -330) & (self._part.state >= -340)
         coll_losses   = np.array([self._line.element_names[i]
                                   for i in self._part.at_element[coll_mask]])