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

xcoll/__init__.py

@@ -1,15 +1,23 @@
 # copyright ############################### #
 # This file is part of the Xcoll Package.   #
-# Copyright (c) CERN, 2023.                 #
+# Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
 from .general import _pkg_root, __version__, citation
 
-from .beam_elements import BlackAbsorber, EverestBlock, EverestCollimator, EverestCrystal, element_classes
+from .beam_elements import BlackAbsorber, BlackCrystal, EverestBlock, EverestCollimator, EverestCrystal, 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 .manager import CollimatorManager
-from .colldb import CollimatorDatabase, load_SixTrack_colldb
+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
 
 # print("If you use Xcoll in your simulations, please cite us :-)")
 # print(citation)
+

xcoll/beam_elements/__init__.py

@@ -3,14 +3,16 @@
 # Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
-from .base import BaseBlock, BaseCollimator
-from .absorber import BlackAbsorber
+from .base import BaseBlock, BaseCollimator, BaseCrystal
+from .absorber import BlackAbsorber, BlackCrystal
 from .everest import EverestBlock, EverestCollimator, EverestCrystal
 
-_all_collimator_types = {BlackAbsorber, EverestCollimator, EverestCrystal}
-
 block_classes = tuple(v for v in globals().values()
                       if isinstance(v, type) and issubclass(v, BaseBlock) and v != BaseBlock)
+# Includes crystals
 collimator_classes = tuple(v for v in globals().values()
-                           if isinstance(v, type) and issubclass(v, BaseCollimator) and v != BaseCollimator)
+                           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 + collimator_classes

xcoll/beam_elements/absorber.py

@@ -1,11 +1,12 @@
 # copyright ############################### #
 # This file is part of the Xcoll Package.   #
-# Copyright (c) CERN, 2023.                 #
+# Copyright (c) CERN, 2024.                 #
 # ######################################### #
 
 import xtrack as xt
 import xobjects as xo
-from .base import BaseCollimator, InvalidXcoll
+from .base import BaseCollimator, BaseCrystal, InvalidXcoll
+from ..scattering_routines.geometry import XcollGeometry
 from ..general import _pkg_root
 
 
@@ -14,16 +15,16 @@ class BlackAbsorber(BaseCollimator):
         '_tracking':        xo.Int8
     }
 
-    _extra_c_sources = [
-        _pkg_root.joinpath('beam_elements','collimators_src','absorber.h')
-    ]
-
     isthick = True
     allow_track = True
     behaves_like_drift = True
     skip_in_loss_location_refinement = True
 
-    _depends_on = [BaseCollimator, xt.Drift, xt.SRotation, xt.XYShift]
+    _depends_on = [BaseCollimator, XcollGeometry]
+
+    _extra_c_sources = [
+        _pkg_root.joinpath('beam_elements','collimators_src','black_absorber.h')
+    ]
 
     _skip_in_to_dict       = BaseCollimator._skip_in_to_dict
     _store_in_to_dict      = BaseCollimator._store_in_to_dict
@@ -33,6 +34,39 @@ class BlackAbsorber(BaseCollimator):
         if '_xobject' not in kwargs:
             kwargs.setdefault('_tracking', True)
         super().__init__(**kwargs)
+        if not isinstance(self._context, xo.ContextCpu):
+            raise ValueError('BlackAbsorber is currently not supported on GPU.')
+
+    def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
+        return InvalidXcoll(length=-self.length, _context=_context, _buffer=_buffer, _offset=_offset)
+
+
+class BlackCrystal(BaseCrystal):
+    _xofields = { **BaseCrystal._xofields,
+        '_tracking':        xo.Int8
+    }
+
+    isthick = True
+    allow_track = True
+    behaves_like_drift = True
+    skip_in_loss_location_refinement = True
+
+    _depends_on = [BaseCrystal, XcollGeometry]
+
+    _extra_c_sources = [
+        _pkg_root.joinpath('beam_elements','collimators_src','black_crystal.h')
+    ]
+
+    _skip_in_to_dict       = BaseCrystal._skip_in_to_dict
+    _store_in_to_dict      = BaseCrystal._store_in_to_dict
+    _internal_record_class = BaseCrystal._internal_record_class
+
+    def __init__(self, **kwargs):
+        if '_xobject' not in kwargs:
+            kwargs.setdefault('_tracking', True)
+        super().__init__(**kwargs)
+        if not isinstance(self._context, xo.ContextCpu):
+            raise ValueError('BlackCrystal is currently not supported on GPU.')
 
     def get_backtrack_element(self, _context=None, _buffer=None, _offset=None):
         return InvalidXcoll(length=-self.length, _context=_context, _buffer=_buffer, _offset=_offset)