xcoll 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

xcoll/__init__.py.orig

@@ -0,0 +1,26 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+from .general import _pkg_root, __version__, citation
+
+<<<<<<< HEAD
+from .beam_elements import BlackAbsorber, BlackCrystal, EverestBlock, EverestCollimator, EverestCrystal, \
+                           Geant4Collimator, BlowUp, EmittanceMonitor, collimator_classes, crystal_classes, element_classes
+=======
+from .beam_elements import BlackAbsorber, BlackCrystal, TransparentCollimator, TransparentCrystal, \
+                           EverestBlock, EverestCollimator, EverestCrystal, BlowUp, EmittanceMonitor, \
+                           collimator_classes, crystal_classes, element_classes
+>>>>>>> main
+from .scattering_routines.everest import materials, Material, CrystalMaterial
+from .scattering_routines.geant4 import Geant4Engine
+from .colldb import CollimatorDatabase
+from .interaction_record import InteractionRecord
+from .rf_sweep import RFSweep
+from .lossmap import LossMap, MultiLossMap
+
+
+# print("If you use Xcoll in your simulations, please cite us :-)")
+# print(citation)
+

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.6.0'
+__version__ = '0.6.1'
 # ======================

xcoll/general.py.orig

@@ -0,0 +1,20 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2024.                 #
+# ######################################### #
+
+from pathlib import Path
+
+_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
+# ======================
+<<<<<<< HEAD
+__version__ = '0.5.11.dev0+geant4'
+=======
+__version__ = '0.6.0'
+>>>>>>> main
+# ======================

xcoll/initial_distribution.py.orig

@@ -0,0 +1,256 @@
+# copyright ############################### #
+# This file is part of the Xcoll package.   #
+# Copyright (c) CERN, 2025.                 #
+# ######################################### #
+
+import numpy as np
+from warnings import warn
+
+import xtrack as xt
+import xobjects as xo
+import xpart as xp
+
+<<<<<<< HEAD
+from .beam_elements import collimator_classes, EverestCrystal, Geant4Collimator
+=======
+from .beam_elements import collimator_classes, BaseCrystal
+>>>>>>> main
+
+
+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, _capacity=None, tw=None,
+                                  _longitudinal_coords=None, **kwargs):
+    """Generate a pencil beam on a collimator."""
+
+    # Do some general checks
+    if not line._has_valid_tracker():
+        raise ValueError("Please build tracker before generating pencil distribution!")
+    coll = line[name]
+    if not isinstance(coll, tuple(collimator_classes)):
+        raise ValueError("Need to provide a valid collimator!")
+    if coll.optics is None:
+        raise ValueError("Need to assign optics to collimators before generating pencil distribution!")
+    num_particles = int(num_particles)
+    if len(line.get_elements_of_type(Geant4Collimator)[0]) > 0:
+        kwargs.setdefault('_capacity', 2*num_particles)
+
+    # 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 coll.side == 'left':
+        if side == '-':
+            raise ValueError("Cannot generate a pencil on the right jaw of a left-"
+                             "sided collimator!")
+        side = '+'
+    if coll.side == 'right':
+        if side == '+':
+            raise ValueError("Cannot generate a pencil on the left jaw of a right-"
+                             "sided collimator!")
+        side = '-'
+
+    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:
+        twiss = line.twiss()
+
+    # Longitudinal plane
+    if _longitudinal_coords:
+        zeta = _longitudinal_coords[0]
+        delta = _longitudinal_coords[1]
+
+    else:
+        zeta, delta = _generate_longitudinal_dist(line, num_particles, sigma_z, longitudinal)
+
+        # 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)
+            zeta_plus = zeta[:num_plus] if hasattr(zeta, '__iter__') else zeta
+            zeta_min  = zeta[num_plus:] if hasattr(zeta, '__iter__') else zeta
+            delta_plus = delta[:num_plus] if hasattr(delta, '__iter__') else delta
+            delta_min  = delta[num_plus:] if hasattr(delta, '__iter__') else delta
+            if _capacity:
+                _capacity_plus = int(_capacity/2)
+                _capacity_min  = int(_capacity - _capacity_plus)
+            else:
+                _capacity_plus = None
+                _capacity_min  = None
+            part_plus = generate_pencil_on_collimator(line=line, name=name, num_particles=num_plus,
+                                impact_parameter=impact_parameter, _capacity=_capacity_plus,
+                                side='+', pencil_spread=pencil_spread, twiss=twiss, 
+                                _longitudinal_coords=[zeta_plus, delta_plus],
+                                **kwargs)
+            part_min = generate_pencil_on_collimator(line=line, name=name, num_particles=num_min,
+                                impact_parameter=impact_parameter, _capacity=_capacity_min,
+                                side='-', pencil_spread=pencil_spread, twiss=twiss, 
+                                _longitudinal_coords=[zeta_min, delta_min],
+                                **kwargs)
+
+            part = xt.Particles.merge([part_plus, part_min])
+            part.start_tracking_at_element = part_plus.start_tracking_at_element
+            assert part.start_tracking_at_element == part_min.start_tracking_at_element
+            return part
+
+    pencil, p_pencil, transverse_norm, p_transverse_norm, is_converging, at_element = \
+                    _generate_4D_pencil_one_jaw(line, name, num_particles, plane, side,
+                                                impact_parameter, pencil_spread, twiss, **kwargs)
+
+    # 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=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, _context=coll._buffer.context, **kwargs
+        )
+
+    part._init_random_number_generator()
+
+    if not is_converging:
+        dri = xt.Drift(length=-coll.length)
+        dri.track(part)
+        part.start_tracking_at_element -= 1
+        part.at_element -= 1
+
+    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,
+                                pencil_spread, twiss=None, _capacity=None, **kwargs):
+    coll = line[name]
+    beam_sizes = twiss.get_beam_covariance(nemitt_x=coll.nemitt_x, nemitt_y=coll.nemitt_y)
+
+    # Is it converging or diverging?
+    # TODO: dispersion might change this
+    # TODO: skew collimators
+    tolerance_tilt = 1e-12 # 0.1 urad tolerance on jaw tilt  =>  we prioritise converging
+    divergence = coll.divergence
+    if side == '+':
+        if isinstance(coll, BaseCrystal):
+            # A pencil on the crystal should always be upstream
+            is_converging = True
+            pencil_pos = coll.jaw_U + impact_parameter
+        else:
+            betatron_angle = coll.gap_L * divergence
+            is_converging = coll.tilt_L + tolerance_tilt >= betatron_angle
+            print(f"Left jaw of collimator {name} is {'con' if is_converging else 'di'}verging.")
+            if is_converging:
+                pencil_pos = coll.jaw_LU + impact_parameter
+            else:
+                pencil_pos = coll.jaw_LD + impact_parameter
+    elif side == '-':
+        if isinstance(coll, BaseCrystal):
+            # A pencil on the crystal should always be upstream
+            is_converging = True
+            pencil_pos = coll.jaw_U - impact_parameter
+        else:
+            betatron_angle = coll.gap_R * divergence
+            is_converging = coll.tilt_R - tolerance_tilt <= betatron_angle
+            print(f"Right jaw of collimator {name} is {'con' if is_converging else 'di'}verging.")
+            if is_converging:
+                pencil_pos = coll.jaw_RU - impact_parameter
+            else:
+                pencil_pos = coll.jaw_RD - impact_parameter
+    else:
+        raise ValueError(f"Sinde {side} not supported in  _generate_4D_pencil_one_jaw!")
+
+    if is_converging:
+        # pencil at front of jaw
+        sigma = beam_sizes.rows[name:f'{name}>>1'][f'sigma_{plane}'][0]
+        tw_at_s = twiss.rows[name]
+        at_element = name
+    else:
+        # pencil at back of jaw
+        sigma = beam_sizes.rows[name:f'{name}>>1'][f'sigma_{plane}'][1]
+        tw_at_s = twiss.rows[f'{name}>>1']
+        at_element = line.element_names[line.element_names.index(name)+1]
+
+    dr_sigmas = pencil_spread/sigma
+
+    # 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=at_element, side=side, twiss=tw_at_s, **kwargs
+                       )
+
+    # 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, is_converging, at_element
+
+
+def _generate_longitudinal_dist(line, num_particles, sigma_z, longitudinal):
+    # TODO: make this more general, make this better
+    if longitudinal is None:
+        return 0, 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':
+        return 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):
+        return longitudinal['zeta'], longitudinal['delta']
+    else:
+        return longitudinal[0], longitudinal[1]

xcoll/plot.py

@@ -5,7 +5,6 @@
 
 import numpy as np
 from numbers import Number
-import matplotlib.pyplot as plt
 
 
 _NORMS = ["total", "coll_max", "max", "none", "raw"]
@@ -13,6 +12,7 @@ _NORMS = ["total", "coll_max", "max", "none", "raw"]
 
 def _plot_lossmap_base(lossmap: dict, *, norm="total", ax=None, xlim=None, ylim=None,
                        legend=True, grid=True, energy=False, show=True, savefig=None):
+    import matplotlib.pyplot as plt
     if norm in _NORMS:
         norm = norm.lower()
     elif not isinstance(norm, Number):

{xcoll-0.6.0.dist-info → xcoll-0.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xcoll
-Version: 0.6.0
+Version: 0.6.1
 Summary: Xsuite collimation package
 Home-page: https://github.com/xsuite/xcoll
 License: Apache 2.0
@@ -21,7 +21,7 @@ Requires-Dist: ruamel-yaml (>=0.17.31,<0.18.0) ; extra == "tests"
 Requires-Dist: xdeps (>=0.10.5)
 Requires-Dist: xobjects (>=0.5.0)
 Requires-Dist: xpart (>=0.23.0)
-Requires-Dist: xtrack (>=0.84.5)
+Requires-Dist: xtrack (>=0.84.6)
 Project-URL: Repository, https://github.com/xsuite/xcoll
 Description-Content-Type: text/markdown
 

{xcoll-0.6.0.dist-info → xcoll-0.6.1.dist-info}/RECORD

@@ -1,6 +1,7 @@
 LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
 NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
 xcoll/__init__.py,sha256=wEk3szUGeeN2Bg7_nx2D2pCgzRR0zcbEcWCP33p-e3I,854
+xcoll/__init__.py.orig,sha256=h9YXS-ZlmC-ZWLoLoAe8V175f2xc6E9gfoe_VjRYwBw,1172
 xcoll/beam_elements/__init__.py,sha256=qu2-LefWfEKZHfALGqZo97DTrwdZVh7Sp97V-3XokJU,1189
 xcoll/beam_elements/absorber.py,sha256=Gw1l1ng1MqF0dRTk7YhkiPAoinS5xsTASMxMDWGsmTw,2842
 xcoll/beam_elements/base.py,sha256=mjnzLbZ-GC56N56Eq5r5oQB8cLmeaqR5oEH1-FnS_H4,57816
@@ -18,17 +19,19 @@ xcoll/beam_elements/everest.py,sha256=iJva_S0OQnmG-qj7rIAiUCssu5WL_Qs5P_n6Qx-IWz
 xcoll/beam_elements/monitor.py,sha256=baMYT_PNya8KX35ReeMwk5Cm04ePUqiDF_s69L5GTX0,16968
 xcoll/beam_elements/transparent.py,sha256=AZOWfH3dvsetL6DjQGejUgX7T18e_3omIIffG_1gwc0,2894
 xcoll/colldb.py,sha256=WaJC6gkDNVvxo65OhucTM6Fwo6pgRFHNnBnHSkiX7AQ,31054
-xcoll/general.py,sha256=kSLikai9J-AWrL9oj9LXhvIKN6bMW7-RpWd4LZCH5A0,534
+xcoll/general.py,sha256=mXL1LaNUxbgXm4AthRJ_7j5oA48cgPDtksQwYmPzLXY,534
+xcoll/general.py.orig,sha256=AvtA-RygKcKCRLDkldzzwG7Q7GlBI49df3aNWe-rr_Y,603
 xcoll/headers/checks.h,sha256=qdXsOTBOK1MwW6bdFF93j4yE648mcDtEv5rGN1w9sfk,1582
 xcoll/headers/particle_states.py,sha256=N8ehLtg99358I4RGCJ3sjuKCk2LVdGqgVx2xGoHrRdM,1959
 xcoll/initial_distribution.py,sha256=CHhqeCVjVOyaBWhfy2sDgLwIgtXZoqcTW-su-M2YByE,10989
+xcoll/initial_distribution.py.orig,sha256=PRZwBtQ2qY_2g6f879IEhUxmNu1tpTArt-qKYNqpjSI,11223
 xcoll/interaction_record/__init__.py,sha256=UFoLiKa-z2oX7YoszP-7Vgdt1nM6kT382v1CaIu8_u0,50
 xcoll/interaction_record/interaction_record.py,sha256=cCjfVCYUJaV0Y-pCljA2_u8GLyhHt7zNjbSXkE7cRO4,13299
 xcoll/interaction_record/interaction_record_src/interaction_record.h,sha256=0rNagnfSGc2i1jauOMIcDbj9QFic9dV_MOyqVx1kw5Q,6067
 xcoll/interaction_record/interaction_types.py,sha256=XTeg2dcImMpqhTqzkJKh6JYam2v9gUMhED-dDKLsilk,2894
 xcoll/line_tools.py,sha256=JBySkKKZ-RGumqAZF_8PSQ3OV0M87P78NeQgbXAdFc4,16699
 xcoll/lossmap.py,sha256=HLdhej9tUjk36T5BulVWurJtfiNa61L_1X7-8kT6JoQ,28407
-xcoll/plot.py,sha256=uy19w0aLufYKpk6RuL265Y8Skl4NgzEuJFF8ur0ezYw,4067
+xcoll/plot.py,sha256=27fYdXstPg3o8XyeBlPuD2yz0GR3dKkNWy4kO2BD9WA,4071
 xcoll/rf_sweep.py,sha256=hMgW15uTp7qPprdlHCZBjOOe8Emdm-R0TfxPDoRXnJU,9021
 xcoll/scattering_routines/engine.py,sha256=1VCBEZLMd5orFo3sZqK51NrqbqELBVrbsz6bB2SXsFQ,22962
 xcoll/scattering_routines/environment.py,sha256=xZ8A7XfvSfpe8vYbM5oD5qbcp8Tuag5qWAI6VqgGUlA,12033
@@ -128,8 +131,8 @@ xcoll/scattering_routines/geometry/rotation.h,sha256=lO3RaQBys9r0ROMjR8T8Rr7UsIE
 xcoll/scattering_routines/geometry/segments.h,sha256=7nKnnin2ByxkKyaYwGvFaqgLQg5uBba4CdLHL7L3iQs,7667
 xcoll/scattering_routines/geometry/sort.h,sha256=b1MkFO2ddzv1fWGeQzsLuz46qo2pKyRSXHjoAEVU7Ts,5763
 xcoll/xaux.py,sha256=mh4E_jhaUFm2LOKawJ0b6X4zuz8JiJhx6w7GcnRgIXQ,2291
-xcoll-0.6.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-xcoll-0.6.0.dist-info/METADATA,sha256=pvIhoCK7JIHzHWv-Alykk78jooTeySkRYg0cAofiHtk,2676
-xcoll-0.6.0.dist-info/NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
-xcoll-0.6.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-xcoll-0.6.0.dist-info/RECORD,,
+xcoll-0.6.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+xcoll-0.6.1.dist-info/METADATA,sha256=qZ6WjBbpMtmyC5jeyM39MTNXw779decuHpGDHOShgpY,2676
+xcoll-0.6.1.dist-info/NOTICE,sha256=6DO_E7WCdRKc42vUoVVBPGttvQi4mRt9fAcxj9u8zy8,74
+xcoll-0.6.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+xcoll-0.6.1.dist-info/RECORD,,