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

xcoll/line_tools.py

@@ -4,12 +4,49 @@
 # ######################################### #
 
 import numpy as np
-import warnings
-warnings.simplefilter("always")
+from warnings import warn
 
 import xtrack as xt
 
-from .beam_elements import element_classes, _all_collimator_classes
+from .beam_elements import element_classes, _all_collimator_classes, _all_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:
@@ -20,8 +57,169 @@ class XcollCollimatorAPI:
     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, _all_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
 
-    def get_optics_at(self, names, *, twiss=None):
+        # 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!")
@@ -34,10 +232,14 @@ class XcollCollimatorAPI:
         tw_exit.name = tw_entry.name
         return tw_entry, tw_exit
 
-    def assign_optics(self, nemitt_x=None, nemitt_y=None, twiss=None):
+    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(_all_collimator_classes)[1]
+        names = self.line.get_elements_of_type(_all_collimator_classes, _all_block_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:
@@ -46,7 +248,7 @@ class XcollCollimatorAPI:
 
     def open(self, names=None):
         if names is None:
-            names = self.line.get_elements_of_type(_all_collimator_classes)[1]
+            names = self.line.get_elements_of_type(_all_collimator_classes, _all_block_classes)[1]
         if len(names) == 0:
             print("No collimators found in line.")
         else:
@@ -62,73 +264,48 @@ class XcollCollimatorAPI:
         else:
             raise NotImplementedError("Need to move this to new type manager or so.")
 
-
-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.")
+    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:
-            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()
+            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):
-    warnings.warn("The function xcoll.assign_optics_to_collimators() is deprecated and will be "
-                + "removed in the future. Please use line.scattering.assign_optics() instead.", DeprecationWarning)
+    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):
-    warnings.warn("The function xcoll.get_optics_at() is deprecated and will be "
-                + "removed in the future. Please use line.scattering.get_optics_at() instead.", DeprecationWarning)
+    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):
-    warnings.warn("The function xcoll.open_collimators() is deprecated and will be "
-                + "removed in the future. Please use line.scattering.open_collimators() instead.", DeprecationWarning)
+    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):
-    warnings.warn("The function xcoll.send_to_parking() is deprecated and will be "
-                + "removed in the future. Please use line.scattering.send_to_parking() instead.", DeprecationWarning)
+    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):
-    warnings.warn("The function xcoll.enable_scattering() is deprecated and will be "
-                + "removed in the future. Please use line.scattering.enable() instead.", DeprecationWarning)
+    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):
-    warnings.warn("The function xcoll.disable_scattering() is deprecated and will be "
-                + "removed in the future. Please use line.scattering.disable() instead.", DeprecationWarning)
+    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.dist-info → xcoll-0.5.6.dist-info}/METADATA

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

{xcoll-0.5.5.dist-info → xcoll-0.5.6.dist-info}/RECORD

@@ -1,10 +1,10 @@
 LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
-xcoll/__init__.py,sha256=e4yt0ZiXp-bjeCoRdcF1TQD60yP30tnCp7ipXWKlPtM,1022
+xcoll/__init__.py,sha256=b_61vh5irhf5fPmqTFJlyhNSt4rmftXg9uXPIEpgVB4,1612
 xcoll/_manager.py,sha256=9NQKaNxZR2I1ChMVBeKQc0A8h6W8gVgRRg72a5NgbXU,808
-xcoll/beam_elements/__init__.py,sha256=oB2FY3PmQ-OsvDSjZaBsbJksOMdFoy0my9_N_K1asOY,1210
+xcoll/beam_elements/__init__.py,sha256=GX3Gj3R1VpRxKft9TZIy0cw5mNG4o1KGNfE28-ExcGs,1290
 xcoll/beam_elements/absorber.py,sha256=efK6gyUgD4x_FnBLpMR7-5_HCdp_753nkYikcdn6ulw,2502
-xcoll/beam_elements/base.py,sha256=tgZwlgaKr34lebyGMIVHa8UqNpqJR4dTPzr-N9W_vxg,50922
+xcoll/beam_elements/base.py,sha256=xLeelb8Amro15kPbcSCBpsDbQ4SKG7qQb2tb5_1-ArQ,52933
 xcoll/beam_elements/blowup.py,sha256=gBXdlISvoDiMjXVpA77ls5QdAU3H9krwvFt2bSW_NII,8029
 xcoll/beam_elements/elements_src/black_absorber.h,sha256=jCyQoaZ7VsIu8coQ99J1dDtAGqZpFIxxFLBlfhO6Drw,5027
 xcoll/beam_elements/elements_src/black_crystal.h,sha256=x7oUwsrkUwUrFvhFsq8vtPkJJTSi1hYAhSt6fjzkJQo,4622
@@ -15,17 +15,17 @@ xcoll/beam_elements/elements_src/everest_collimator.h,sha256=DlWkb5RORnp5VwI7E31
 xcoll/beam_elements/elements_src/everest_crystal.h,sha256=w1wHWhrvodvcJ2Atr8LSaVYHfCU3Y63KAjZZFFv87OI,12123
 xcoll/beam_elements/everest.py,sha256=PA_VWpnPrIuO1xN__eKyT_ejbGZK7p93QHDVi3re7cM,8541
 xcoll/beam_elements/monitor.py,sha256=zzMdN3JMFSAs-30_ntRvd5qZGdsXfGtColhiFDuMcIk,16928
-xcoll/colldb.py,sha256=IyWozG_37gRf7LGYDz4LMB2HzjexMOapFgIqzeym31w,31027
-xcoll/general.py,sha256=7KCioNHNLS4wm9VU9iPFm0N2E6MlhkKCLpPULrQagAg,534
+xcoll/colldb.py,sha256=vp9m4A0_4y7RVkrpaowt9-08Umhb8AwAkXkiwFikuCc,30993
+xcoll/general.py,sha256=hqc9uwLEA8CVMamQTCgV-N-fnAbwU9vV-q8b1KYmz8U,534
 xcoll/headers/checks.h,sha256=qdXsOTBOK1MwW6bdFF93j4yE648mcDtEv5rGN1w9sfk,1582
 xcoll/headers/particle_states.h,sha256=DZa_ZSaJrjnA8aHFriZKfRCkArQ8nK1t445MRwevDtA,840
-xcoll/initial_distribution.py,sha256=0O4hmO69CuTFDD6QA5G7iEMNx47vaMHn0A9IPw1dYi0,8799
-xcoll/install.py,sha256=HumxIzaPR6wW_6HFtksoUnUlaS1jMp7RbV7FhmIBWmg,7685
+xcoll/initial_distribution.py,sha256=_pm93FNNTO1GROni5j5sCdw8McXQ7PUMyVSXLbDNCxI,8789
+xcoll/install.py,sha256=SxEFQnhWXlsXyPBIo847q6wPgId_f5ZtFRR1awGbkjc,2108
 xcoll/interaction_record/__init__.py,sha256=UFoLiKa-z2oX7YoszP-7Vgdt1nM6kT382v1CaIu8_u0,50
 xcoll/interaction_record/interaction_record.py,sha256=G2sFXgfmkgpDMaHWJaC80aAreHbY2rGhbkhaMhWKQBg,13192
 xcoll/interaction_record/interaction_record_src/interaction_record.h,sha256=0rNagnfSGc2i1jauOMIcDbj9QFic9dV_MOyqVx1kw5Q,6067
 xcoll/interaction_record/interaction_types.py,sha256=Vh6oFYKdRNOx9hc_E0iT8auNZVKC0qYVp_p7oyClZ8o,2894
-xcoll/line_tools.py,sha256=1yrw9b3j-yO3uyRpnKhlU0iHxDNeTj5M1jj9v3GqBj0,5634
+xcoll/line_tools.py,sha256=XAOHKLP1I0MuUWlTQvnIWy07MbaUxz64QeUtbMiNpo4,14396
 xcoll/lossmap.py,sha256=pV6uf-CKt6PROK_EZ2_DieJftRI2NcsE3CLG23Fth-I,8087
 xcoll/rf_sweep.py,sha256=P2X1S9pGi4OpNYnzYfQVyblFt2p8aw9EWHsKDkAuYt0,8936
 xcoll/scattering_routines/everest/__init__.py,sha256=7lkkeZ1liBjXVHCuRpgzZI6ohzHVMj5uJBO792147XY,286
@@ -406,9 +406,8 @@ xcoll/scattering_routines/geometry/objects.h,sha256=A5ktGvVlSkC4hUsI_PQFsE80CuDw
 xcoll/scattering_routines/geometry/rotation.h,sha256=lO3RaQBys9r0ROMjR8T8Rr7UsIEm-H9_C_70Nwz4MXY,701
 xcoll/scattering_routines/geometry/segments.h,sha256=7nKnnin2ByxkKyaYwGvFaqgLQg5uBba4CdLHL7L3iQs,7667
 xcoll/scattering_routines/geometry/sort.h,sha256=b1MkFO2ddzv1fWGeQzsLuz46qo2pKyRSXHjoAEVU7Ts,5763
-xcoll/scattering_routines/geometry/temp.c,sha256=uak0prUg6e441GuR8A9fvAUSfg-E_jaAfGw9fu2Ogms,32437
-xcoll-0.5.5.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-xcoll-0.5.5.dist-info/METADATA,sha256=uN-Ple230sy6CY2egyegJCG9mSGd3nO3C5aWDMWU42M,2709
-xcoll-0.5.5.dist-info/NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
-xcoll-0.5.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-xcoll-0.5.5.dist-info/RECORD,,
+xcoll-0.5.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+xcoll-0.5.6.dist-info/METADATA,sha256=5ySOFXnd3w9Lu6ZAzodeT5WGHSmud6esOCZh-L10mxQ,2709
+xcoll-0.5.6.dist-info/NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
+xcoll-0.5.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+xcoll-0.5.6.dist-info/RECORD,,