xtrack 0.57.1__tar.gz → 0.57.2__tar.gz

{xtrack-0.57.1 → xtrack-0.57.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.57.1
+Version: 0.57.2
 Summary: Tracking library for particle accelerators
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xtrack

{xtrack-0.57.1 → xtrack-0.57.2}/tests/test_element_characterization_functions.py

@@ -41,7 +41,7 @@ def test_has_backtrack(test_context):
     line.build_tracker(_context=test_context)
 
     _has_backtrack = xt.line._has_backtrack
-    _allow_backtrack = xt.line._allow_backtrack
+    _allow_loss_refinement = xt.line._allow_loss_refinement
 
     assert _has_backtrack(line['e0'], line)
     assert _has_backtrack(line['e1'], line)
@@ -50,12 +50,12 @@ def test_has_backtrack(test_context):
     assert _has_backtrack(line['a1'], line)
     assert _has_backtrack(line['a2'], line)
 
-    assert not _allow_backtrack(line['e0'], line)
-    assert not _allow_backtrack(line['e1'], line)
-    assert not _allow_backtrack(line['e2'], line)
-    assert _allow_backtrack(line['a0'], line)
-    assert _allow_backtrack(line['a1'], line)
-    assert _allow_backtrack(line['a2'], line)
+    assert not _allow_loss_refinement(line['e0'], line)
+    assert not _allow_loss_refinement(line['e1'], line)
+    assert not _allow_loss_refinement(line['e2'], line)
+    assert _allow_loss_refinement(line['a0'], line)
+    assert _allow_loss_refinement(line['a1'], line)
+    assert _allow_loss_refinement(line['a2'], line)
 
 def test_is_drift_behaves_like_drift():
 

{xtrack-0.57.1 → xtrack-0.57.2}/tests/test_line.py

@@ -1006,4 +1006,30 @@ def test_insert_thick_element_reuse_marker_name():
 
     assert np.all(tt.name == ['d1..0', 'm1', 'd2..1', '_end_point'])
     assert np.all(tt.parent_name == ['d1', None, 'd2', None])
-    assert_allclose(tt.s, [0. , 0.5, 1.5, 2. ], rtol=0, atol=1e-14)
+    assert_allclose(tt.s, [0. , 0.5, 1.5, 2. ], rtol=0, atol=1e-14)
+
+def test_multiple_thick_elements():
+    line = xt.Line(
+        elements=[xt.Drift(length=1.0) for i in range(10)]
+    )
+
+    s_insert = np.array([2.5, 5.5, 7])
+    l_insert = np.array([1.0, 1.0, 1.0])
+    ele_insert = [xt.Sextupole(length=l) for l in l_insert]
+
+    line._insert_thick_elements_at_s(
+        element_names=[f'insertion_{i}' for i in range(len(s_insert))],
+        elements=ele_insert,
+        at_s=s_insert
+    )
+
+    tt = line.get_table()
+
+    assert np.all(tt.name == ['e0', 'e1', 'e2..0', 'insertion_0', 'e3..1', 'e4', 'e5..0',
+        'insertion_1', 'e6..1', 'insertion_2', 'e8', 'e9', '_end_point'])
+    xo.assert_allclose(tt.s, [ 0. ,  1. ,  2. ,  2.5,  3.5,  4. ,  5. ,
+                                    5.5,  6.5,  7. ,  8. , 9. , 10. ])
+
+    assert np.all(tt.element_type == ['Drift', 'Drift', 'DriftSlice', 'Sextupole', 'DriftSlice', 'Drift',
+        'DriftSlice', 'Sextupole', 'DriftSlice', 'Sextupole', 'Drift',
+        'Drift', ''])

{xtrack-0.57.1 → xtrack-0.57.2}/tests/test_tracker.py

@@ -709,7 +709,7 @@ def test_track_log_and_merit_function(pimms_mad, test_context):
     fit = np.polyfit(np.arange(num_turns), line.log_last_track['kqfa'], 1, full=True)
     (slope, _), residual, _, _, _ = fit
     assert slope > 0
-    assert residual < 1e-29
+    assert residual < 1e-28
     assert np.isclose(line.log_last_track['kqfa'][0], kqfa_before, atol=1e-14, rtol=0)
     assert np.isclose(line.log_last_track['kqfa'][-1], line.vv['kqfa'], atol=1e-14, rtol=0)
 

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/__init__.py

@@ -37,7 +37,7 @@ from .multisetter import MultiSetter
 from .footprint import Footprint, LinearRescale
 
 # Flags and test functions
-from .line import _is_drift, _behaves_like_drift, _is_aperture, _is_thick, _allow_backtrack
+from .line import _is_drift, _behaves_like_drift, _is_aperture, _is_thick, _allow_loss_refinement
 from .line import _lines_equal, _apertures_equal
 
 from .slicing import Strategy, Uniform, Teapot

xtrack-0.57.2/xtrack/_version.py

@@ -0,0 +1 @@
+__version__ = '0.57.2'

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/base_element.py

@@ -431,7 +431,7 @@ class BeamElement(xo.HybridClass, metaclass=MetaBeamElement):
     behaves_like_drift = False
     allow_track = True
     has_backtrack = False
-    allow_backtrack = False
+    allow_loss_refinement = False
     allow_rot_and_shift = True
     skip_in_loss_location_refinement = False
     needs_rng = False

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/beam_elements/elements.py

@@ -45,7 +45,7 @@ class Marker(BeamElement):
         '_dummy': xo.Int64}
 
     behaves_like_drift = True
-    allow_backtrack = True
+    allow_loss_refinement = True
     has_backtrack = True
     allow_rot_and_shift = False
 
@@ -72,7 +72,7 @@ class Drift(BeamElement):
     isthick = True
     behaves_like_drift = True
     has_backtrack = True
-    allow_backtrack = True
+    allow_loss_refinement = True
     allow_rot_and_shift = False
 
     _extra_c_sources = [
@@ -138,7 +138,7 @@ class XYShift(BeamElement):
         'dy': xo.Float64,
         }
 
-    allow_backtrack = True
+    allow_loss_refinement = True
     has_backtrack = True
     allow_rot_and_shift = False
 
@@ -279,7 +279,7 @@ class SRotation(BeamElement):
         'sin_z': xo.Float64,
         }
 
-    allow_backtrack = True
+    allow_loss_refinement = True
     has_backtrack = True
     allow_rot_and_shift = False
 
@@ -348,7 +348,7 @@ class XRotation(BeamElement):
         'tan_angle': xo.Float64,
         }
 
-    allow_backtrack = True
+    allow_loss_refinement = True
     has_backtrack = True
     allow_rot_and_shift = False
 
@@ -431,7 +431,7 @@ class YRotation(BeamElement):
     '''
 
     has_backtrack = True
-    allow_backtrack = True
+    allow_loss_refinement = True
     allow_rot_and_shift = False
 
     _xofields={

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/beam_elements/slice_elements_thick.py

@@ -203,6 +203,7 @@ class DriftSliceBend(BeamElement):
     rot_and_shift_from_parent = False
     _skip_in_to_dict = ['_parent']
     has_backtrack = True
+    allow_loss_refinement = True
     _force_moveable = True
     isthick = True
     _inherit_strengths = False
@@ -239,6 +240,7 @@ class DriftSliceQuadrupole(BeamElement):
     rot_and_shift_from_parent = False
     _skip_in_to_dict = ['_parent']
     has_backtrack = True
+    allow_loss_refinement = True
     _force_moveable = True
     isthick = True
     _inherit_strengths = False
@@ -276,6 +278,7 @@ class DriftSliceSextupole(BeamElement):
     rot_and_shift_from_parent = False
     _skip_in_to_dict = ['_parent']
     has_backtrack = True
+    allow_loss_refinement = True
     _force_moveable = True
     isthick = True
     _inherit_strengths = False
@@ -313,6 +316,7 @@ class DriftSliceOctupole(BeamElement):
     rot_and_shift_from_parent = False
     _skip_in_to_dict = ['_parent']
     has_backtrack = True
+    allow_loss_refinement = True
     _force_moveable = True
     isthick = True
     _inherit_strengths = False
@@ -349,6 +353,7 @@ class DriftSlice(BeamElement):
     rot_and_shift_from_parent = False
     _skip_in_to_dict = ['_parent']
     has_backtrack = True
+    allow_loss_refinement = True
     _force_moveable = True
     isthick = True
     _inherit_strengths = False

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/line.py

@@ -1855,7 +1855,7 @@ class Line:
     def _elements_intersecting_s(
             self,
             s: List[float],
-            tol=1e-16,
+            s_tol=1e-6,
     ) -> Dict[str, List[float]]:
         """Given a list of s positions, return a list of elements 'cut' by s.
 
@@ -1863,9 +1863,9 @@ class Line:
         ---------
         s
             A list of s positions.
-        tol
+        s_tol
             Tolerance used when checking if s falls inside an element, or
-            at its edge. Defaults to 1e-16.
+            at its edge. Defaults to 1e-6.
 
         Returns
         -------
@@ -1890,12 +1890,12 @@ class Line:
                     start, name = next(all_s_iter)
                     continue
 
-                if np.isclose(current_s, start, atol=tol, rtol=0):
+                if np.isclose(current_s, start, atol=s_tol, rtol=0):
                     current_s = next(current_s_iter)
                     continue
 
                 end = start + _length(element, self)
-                if np.isclose(current_s, end, atol=tol, rtol=0):
+                if np.isclose(current_s, end, atol=s_tol, rtol=0):
                     current_s = next(current_s_iter)
                     continue
 
@@ -1916,9 +1916,9 @@ class Line:
 
         return cuts_for_element
 
-    def cut_at_s(self, s: List[float]):
+    def cut_at_s(self, s: List[float], s_tol=1e-6):
         """Slice the line so that positions in s never fall inside an element."""
-        cuts_for_element = self._elements_intersecting_s(s)
+        cuts_for_element = self._elements_intersecting_s(s, s_tol=s_tol)
         strategies = [Strategy(None)]  # catch-all, ignore unaffected elements
 
         for name, cuts in cuts_for_element.items():
@@ -2005,11 +2005,12 @@ class Line:
 
         s_vect_upstream = np.array(self.get_s_position(mode='upstream'))
         if _is_thick(element, self) and _length(element, self) > 0:
-            s_vect_downstream = np.array(self.get_s_position(mode='downstream'))
-            i_first_drift_to_cut = np.where(s_vect_downstream > s_start_ele)[0][0]
-            i_last_drift_to_cut = np.where(s_vect_upstream < s_end_ele)[0][-1]
-            assert i_first_drift_to_cut <= i_last_drift_to_cut
-            self.element_names[i_first_drift_to_cut:i_last_drift_to_cut + 1] = [name]
+            i_first_removal = np.where(np.abs(s_vect_upstream - s_start_ele) < s_tol)[0][-1]
+            i_last_removal = np.where(np.abs(s_vect_upstream - s_end_ele) < s_tol)[0][0] - 1
+            xo.assert_allclose(s_vect_upstream[i_last_removal + 1]
+                              - s_vect_upstream[i_first_removal],
+                                _length(element, self), atol=2 * s_tol, rtol=0)
+            self.element_names[i_first_removal:i_last_removal + 1] = [name]
         else:
             i_closest = np.argmin(np.abs(s_vect_upstream - at_s))
             assert np.abs(s_vect_upstream[i_closest] - at_s) < s_tol
@@ -2891,7 +2892,7 @@ class Line:
             ee = self.element_dict[name]
             if isinstance(ee, xt.Replica):
                 ee = ee.resolve(self)
-            if _allow_backtrack(ee, self) and not name in needs_aperture:
+            if _allow_loss_refinement(ee, self) and not name in needs_aperture:
                 dont_need_aperture[name] = True
             if name.endswith('_entry') or name.endswith('_exit'):
                 dont_need_aperture[name] = True
@@ -3615,7 +3616,7 @@ class Line:
         )
         return cache
 
-    def _insert_thin_elements_at_s(self, elements_to_insert):
+    def _insert_thin_elements_at_s(self, elements_to_insert, s_tol=0.5e-6):
 
         '''
         Example:
@@ -3629,7 +3630,6 @@ class Line:
         '''
 
         self._frozen_check()
-        s_tol = 0.5e-6
 
         s_cuts = [ee[0] for ee in elements_to_insert]
         s_cuts = np.sort(s_cuts)
@@ -3673,6 +3673,69 @@ class Line:
                 if insert_at is not None:
                     insert_at += 1
 
+    def _insert_thick_elements_at_s(self, element_names, elements,
+                                    at_s, s_tol=1e-6):
+
+        assert isinstance(element_names, (list, tuple))
+        assert isinstance(elements, (list, tuple))
+        assert isinstance(at_s, (list, tuple, np.ndarray))
+        assert len(element_names) == len(elements) == len(at_s)
+
+        self._frozen_check()
+
+        s_insert = np.array(at_s)
+        l_insert = np.array([_length(ee, None) for ee in elements])
+        ele_insert = list(elements).copy()
+        name_insert = list(element_names).copy()
+
+        end_insert = np.array(s_insert) + np.array(l_insert)
+
+        self.cut_at_s(list(s_insert) + list(end_insert))
+
+        i_sorted = np.argsort(s_insert)
+        s_insert_sorted = s_insert[i_sorted]
+        ele_insert_sorted = [ele_insert[i] for i in i_sorted]
+        name_insert_sorted = [name_insert[i] for i in i_sorted]
+        end_insert_sorted = end_insert[i_sorted]
+
+        assert np.all(s_insert_sorted[:-1] < end_insert_sorted[1:]), (
+                    'Overlapping insertions')
+
+        old_element_names = self.element_names
+
+        s_tol = 1e-6
+
+        s_vect_upstream = np.array(self.get_s_position(mode='upstream'))
+
+        i_replace = np.zeros(len(s_vect_upstream), dtype=int)
+        mask_remove = np.zeros(len(s_vect_upstream), dtype=bool)
+
+        i_replace[:] = -1
+
+        for ii in range(len(s_insert_sorted)):
+            ss_start = s_insert_sorted[ii]
+            ss_end = end_insert_sorted[ii]
+
+            i_first_removal = np.where(np.abs(s_vect_upstream - ss_start) < s_tol)[0][-1]
+            i_last_removal = np.where(np.abs(s_vect_upstream - ss_end) < s_tol)[0][0] - 1
+
+            i_replace[i_first_removal] = ii
+            mask_remove[i_first_removal+1:i_last_removal+1] = True
+
+        new_element_names = []
+        for ii, nn in enumerate(old_element_names):
+            if mask_remove[ii]:
+                continue
+            if i_replace[ii] != -1:
+                new_element_names.append(name_insert_sorted[i_replace[ii]])
+            else:
+                new_element_names.append(nn)
+
+        for new_nn, new_ee in zip(name_insert_sorted, ele_insert_sorted):
+            self.element_dict[new_nn] = new_ee
+
+        self.element_names = new_element_names
+
     @property
     def _line_before_slicing(self):
         if self._element_names_before_slicing is None:
@@ -3871,10 +3934,10 @@ def _is_collective(element, line):
     return iscoll
 
 # whether backtrack in loss location refinement is allowed
-def _allow_backtrack(element, line):
+def _allow_loss_refinement(element, line):
     if isinstance(element, xt.Replica):
         element = element.resolve(line)
-    return hasattr(element, 'allow_backtrack') and element.allow_backtrack
+    return hasattr(element, 'allow_loss_refinement') and element.allow_loss_refinement
 
 # whether element has backtrack capability
 def _has_backtrack(element, line):

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/loss_location_refinement/loss_location_refinement.py

@@ -10,7 +10,7 @@ import xobjects as xo
 import xtrack as xt
 
 from ..beam_elements import LimitPolygon, XYShift, SRotation, Drift, Marker
-from ..line import (Line, _is_thick, _behaves_like_drift, _allow_backtrack,
+from ..line import (Line, _is_thick, _behaves_like_drift, _allow_loss_refinement,
                     _has_backtrack, _is_aperture)
 
 from ..general import _print
@@ -52,7 +52,7 @@ class LossLocationRefinement:
         If True, the lines used to refine the loss location are saved.
     allowed_backtrack_types : list
         List of element types through which the backtracking is allowed.
-        Elements exposing the attribute `allow_backtrack` are automatically
+        Elements exposing the attribute `allow_loss_refinement` are automatically
         added to the list.
 
     '''
@@ -261,7 +261,7 @@ def refine_loss_location_single_aperture(particles, i_aper_1, i_end_thin_0,
         can_backtrack = True
         if not _has_backtrack(ee, line):
             can_backtrack = False
-        elif not _allow_backtrack(ee, line):
+        elif not _allow_loss_refinement(ee, line):
             can_backtrack = False
 
             # Check for override

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/monitors/beam_position_monitor.py

@@ -35,7 +35,7 @@ class BeamPositionMonitor(BeamElement):
     }
 
     behaves_like_drift = True
-    allow_backtrack = True
+    allow_loss_refinement = True
 
     properties = [field.name for field in BeamPositionMonitorRecord._fields]
 

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/monitors/beam_profile_monitor.py

@@ -40,7 +40,7 @@ class BeamProfileMonitor(BeamElement):
     }
 
     behaves_like_drift = True
-    allow_backtrack = True
+    allow_loss_refinement = True
 
     properties = [field.name for field in BeamProfileMonitorRecord._fields]
 

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/monitors/beam_size_monitor.py

@@ -35,7 +35,7 @@ class BeamSizeMonitor(BeamElement):
     }
 
     behaves_like_drift = True
-    allow_backtrack = True
+    allow_loss_refinement = True
 
     properties = [field.name for field in BeamSizeMonitorRecord._fields]
 

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/monitors/last_turns_monitor.py

@@ -42,7 +42,7 @@ class LastTurnsMonitor(BeamElement):
     }
 
     behaves_like_drift = True
-    allow_backtrack = True
+    allow_loss_refinement = True
 
     # TODO: find a way to dynamically change what properties are being saved by this monitor
     properties = [field.name for field in LastTurnsData._fields if field.name != 'lost_at_offset']

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/monitors/particles_monitor.py

@@ -162,7 +162,7 @@ class ParticlesMonitor(BeamElement):
 
     behaves_like_drift = True
     has_backtrack = True
-    allow_backtrack = True
+    allow_loss_refinement = True
     _ParticlesClass = xt.Particles
 
 

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/particles/particles.py

@@ -1318,7 +1318,8 @@ class Particles(xo.HybridClass):
 
     @property
     def kin_ps(self):
-        out = np.sqrt((1 + self.delta) ** 2 - self.kin_px ** 2 - self.kin_py ** 2)
+        sqrt = self._context.nplike_lib.sqrt
+        out = sqrt((1 + self.delta) ** 2 - self.kin_px ** 2 - self.kin_py ** 2)
         return self._buffer.context.linked_array_type.from_array(
             out,
             mode='readonly',

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/tracker.py

@@ -143,9 +143,8 @@ class Tracker:
         if not _prebuilding_kernels:
             self._get_twiss_mask_markers() # to cache it
 
-        self._init_io_buffer(io_buffer)
-
         self.line = line
+        self._init_io_buffer(io_buffer)
         self.line.tracker = self
 
         if compile:
@@ -153,7 +152,13 @@ class Tracker:
 
     def _init_io_buffer(self, io_buffer=None):
         if io_buffer is None:
-            io_buffer = new_io_buffer(_context=self._context)
+            io_bufs = [ee.io_buffer for ee in self.line.elements if hasattr(ee, 'io_buffer')]
+            if len(io_bufs) == 0:
+                io_buffer = new_io_buffer(_context=self._context)
+            elif len(np.unique([id(buf) for buf in io_bufs])) > 1:
+                raise ValueError("Different io buffers found in elements!")
+            else:
+                io_buffer = io_bufs[0]
         self.io_buffer = io_buffer
 
     def _split_parts_for_collective_mode(self, line, _buffer):

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack/twiss.py

@@ -2824,6 +2824,76 @@ class TwissTable(Table):
 
         return Table(res)
 
+    def get_ibs_growth_rates(
+        self,
+        formalism: str,
+        total_beam_intensity: int = None,
+        gemitt_x: float = None,
+        nemitt_x: float = None,
+        gemitt_y: float = None,
+        nemitt_y: float = None,
+        sigma_delta: float = None,
+        bunch_length: float = None,
+        bunched: bool = True,
+        particles: xt.Particles = None,
+        **kwargs,
+    ):
+        """
+        Computes IntraBeam Scattering growth rates from the provided `xtrack.Line`.
+
+        Parameters
+        ----------
+        line : xtrack.Line
+            Line in which the IBS kick element will be installed.
+        formalism : str
+            Which formalism to use for the computation. Can be ``Nagaitsev``
+            or ``Bjorken-Mtingwa`` (also accepts ``B&M``), case-insensitively.
+        total_beam_intensity : int, optional
+            The beam intensity. Required if `particles` is not provided.
+        gemitt_x : float, optional
+            Horizontal geometric emittance in [m]. If `particles` is not
+            provided, either this parameter or `nemitt_x` is required.
+        nemitt_x : float, optional
+            Horizontal normalized emittance in [m]. If `particles` is not
+            provided, either this parameter or `gemitt_x` is required.
+        gemitt_y : float, optional
+            Vertical geometric emittance in [m]. If `particles` is not
+            provided, either this parameter or `nemitt_y` is required.
+        nemitt_y : float, optional
+            Vertical normalized emittance in [m]. If `particles` is not
+            provided, either this parameter or `gemitt_y` is required.
+        sigma_delta : float, optional
+            The momentum spread. Required if `particles` is not provided.
+        bunch_length : float, optional
+            The bunch length in [m]. Required if `particles` is not provided.
+        bunched : bool, optional
+            Whether the beam is bunched or not (coasting). Defaults to `True`.
+            Required if `particles` is not provided.
+        particles : xtrack.Particles
+            The particles to circulate in the line. If provided the emittances,
+            momentum spread and bunch length will be computed from the particles.
+            Otherwise explicit values must be provided (see above parameters).
+        **kwargs : dict
+            Keyword arguments are passed to the growth rates computation method of
+            the chosen IBS formalism implementation. See the IBS details from the
+            `xfields` package directly.
+
+        Returns
+        -------
+        IBSGrowthRates
+            An ``IBSGrowthRates`` object with the computed growth rates.
+        """
+        try:
+            from xfields.ibs import get_intrabeam_scattering_growth_rates
+        except ImportError:
+            raise ImportError("Please install xfields to use this feature.")
+        return get_intrabeam_scattering_growth_rates(
+            self, formalism, total_beam_intensity,
+            gemitt_x, nemitt_x, gemitt_y, nemitt_y,
+            sigma_delta, bunch_length, bunched,
+            particles, **kwargs,
+        )
+
     def get_R_matrix(self, start, end):
 
         assert self.values_at == 'entry', 'Not yet implemented for exit'

{xtrack-0.57.1 → xtrack-0.57.2}/xtrack.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.57.1
+Version: 0.57.2
 Summary: Tracking library for particle accelerators
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xtrack

xtrack-0.57.1/xtrack/_version.py

@@ -1 +0,0 @@
-__version__ = '0.57.1'