xtrack 0.45.0__tar.gz → 0.46.0__tar.gz

{xtrack-0.45.0/xtrack.egg-info → xtrack-0.46.0}/PKG-INFO

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

{xtrack-0.45.0 → xtrack-0.46.0}/xtrack/__init__.py

@@ -12,7 +12,8 @@ from .tracker_data import TrackerData
 from .line import Line, Node, freeze_longitudinal, _temp_knobs, EnergyProgram
 from .tracker import Tracker
 from .match import (Vary, Target, TargetList, VaryList, TargetInequality, Action,
-                    TargetRelPhaseAdvance, TargetSet, START, END)
+                    TargetRelPhaseAdvance, TargetSet, START, END,
+                    GreaterThan, LessThan)
 from .targets import (TargetLuminosity, TargetSeparationOrthogonalToCrossing,
                       TargetSeparation)
 from .twiss import TwissInit, TwissTable

xtrack-0.46.0/xtrack/_version.py

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

{xtrack-0.45.0 → xtrack-0.46.0}/xtrack/line.py

@@ -10,7 +10,7 @@ import json
 from contextlib import contextmanager
 from copy import deepcopy
 from pprint import pformat
-from typing import List, Literal, Optional
+from typing import List, Literal, Optional, Collection
 
 import numpy as np
 from scipy.constants import c as clight
@@ -22,9 +22,10 @@ import xpart as xp
 import xtrack as xt
 import xdeps as xd
 from .compounds import CompoundContainer, CompoundType, Compound, SlicedCompound
+from .progress_indicator import progress
 from .slicing import Slicer
 
-from .survey import survey_from_tracker
+from .survey import survey_from_line
 from xtrack.twiss import (compute_one_turn_matrix_finite_differences,
                           find_closed_orbit_line, twiss_line,
                           compute_T_matrix_line,
@@ -171,19 +172,13 @@ class Line:
 
         if isinstance(dct['elements'], dict):
             elements = {}
-            num_elements = len(dct['elements'].keys())
-            for ii, (kk, ee) in enumerate(dct['elements'].items()):
-                if ii % 100 == 0:
-                    _print('Loading line from dict: '
-                        f'{round(ii/num_elements*100):2d}%  ',end="\r", flush=True)
+            for ii, (kk, ee) in enumerate(
+                    progress(dct['elements'].items(), desc='Loading line from dict')):
                 elements[kk] = _deserialize_element(ee, class_dict, _buffer)
         elif isinstance(dct['elements'], list):
             elements = []
-            num_elements = len(dct['elements'])
-            for ii, ee in enumerate(dct['elements']):
-                if ii % 100 == 0:
-                    _print('Loading line from dict: '
-                        f'{round(ii/num_elements*100):2d}%  ',end="\r", flush=True)
+            for ii, ee in enumerate(
+                    progress(dct['elements'], desc='Loading line from dict')):
                 elements.append(_deserialize_element(ee, class_dict, _buffer))
         else:
             raise ValueError('Field `elements` must be a dict or a list')
@@ -733,6 +728,7 @@ class Line:
         turn_by_turn_monitor=None,
         freeze_longitudinal=False,
         time=False,
+        with_progress=False,
         **kwargs):
 
         """
@@ -770,7 +766,11 @@ class Line:
         time: bool, optional
             If True, the time taken for tracking is recorded and can be retrieved
             in `line.time_last_track`.
-
+        with_progress: bool or int, optional
+            If truthy, a progress bar is displayed during tracking. If an integer
+            is provided, it is used as the number of turns between two updates
+            of the progress bar. If True, 100 is taken by default. By default,
+            equals to False and no progress bar is displayed.
         """
 
         self._check_valid_tracker()
@@ -783,6 +783,7 @@ class Line:
             turn_by_turn_monitor=turn_by_turn_monitor,
             freeze_longitudinal=freeze_longitudinal,
             time=time,
+            with_progress=with_progress,
             **kwargs)
 
     def slice_thick_elements(self, slicing_strategies):
@@ -970,6 +971,12 @@ class Line:
         compute_R_element_by_element=None,
         compute_lattice_functions=None,
         compute_chromatic_properties=None,
+        ele_init=None,
+        x=None, px=None, y=None, py=None, zeta=None, delta=None,
+        betx=None, alfx=None, bety=None, alfy=None, bets=None,
+        dx=None, dpx=None, dy=None, dpy=None, dzeta=None,
+        mux=None, muy=None, muzeta=None,
+        ax_chrom=None, bx_chrom=None, ay_chrom=None, by_chrom=None,
         _continue_if_lost=None,
         _keep_tracking_data=None,
         _keep_initial_particles=None,
@@ -1142,7 +1149,7 @@ class Line:
             Survey table.
         """
 
-        return survey_from_tracker(self.tracker, X0=X0, Y0=Y0, Z0=Z0, theta0=theta0,
+        return survey_from_line(self, X0=X0, Y0=Y0, Z0=Z0, theta0=theta0,
                                    phi0=phi0, psi0=psi0, element0=element0,
                                    reverse=reverse)
 
@@ -2618,13 +2625,7 @@ class Line:
         i_prev_aperture = elements_df[elements_df['is_aperture']].index[0]
         i_next_aperture = 0
 
-        for iee in range(i_prev_aperture, num_elements):
-
-            if iee % 100 == 0:
-                _print(
-                    f'Checking aperture: {round(iee/num_elements*100):2d}%  ',
-                    end="\r", flush=True)
-
+        for iee in progress(range(i_prev_aperture, num_elements), desc='Checking aperture'):
             if dont_need_aperture[elements_df.loc[iee, 'name']]:
                 continue
 
@@ -3154,6 +3155,7 @@ class Line:
                          fields=['hxl', 'hyl', 'length', 'radiation_flag',
                                  'delta_taper', 'voltage', 'frequency',
                                  'lag', 'lag_taper',
+                                 'k1',
                                 ('knl', 0), ('ksl', 0), ('knl', 1), ('ksl', 1),
                                 ('knl', 2), ('ksl', 2),
                                 ])

{xtrack-0.45.0 → xtrack-0.46.0}/xtrack/mad_loader.py

@@ -27,20 +27,15 @@ Loader.add_<name>(mad_elem,line,buffer) to add a new element to line
 
 if the want to control how the xobject is created
 """
-import abc
-import functools
-import re
-from itertools import zip_longest
-from typing import List, Iterable, Iterator, Tuple, Union
+from typing import List, Union
 
 import numpy as np
-from math import tan
 
-import xtrack, xobjects
+import xobjects
+import xtrack
 from .compounds import Compound
-
 from .general import _print
-
+from .progress_indicator import progress
 
 # Generic functions
 
@@ -691,10 +686,11 @@ class MadLoader:
             madeval = None
             self.Builder = ElementBuilder
 
-        nelem = len(self.sequence.expanded_elements)
-
-        for ii, el in enumerate(self.iter_elements(madeval=madeval)):
-
+        for ii, el in enumerate(progress(
+                self.iter_elements(madeval=madeval),
+                desc=f'Converting sequence "{self.sequence.name}"',
+                total=len(self.sequence.expanded_elements)),
+        ):
             # for each mad element create xtract elements in a buffer and add to a line
             converter = getattr(self, "convert_" + el.type, None)
             adder = getattr(self, "add_" + el.type, None)
@@ -715,14 +711,6 @@ class MadLoader:
                     f'Element {el.type} not supported,\nimplement "add_{el.type}"'
                     f" or convert_{el.type} in function in MadLoader"
                 )
-            if ii % 100 == 0:
-                _print(
-                    f'Converting sequence "{self.sequence.name}":'
-                    f' {round(ii/nelem*100):2d}%     ',
-                    end="\r",
-                    flush=True,
-                )
-        _print()
         return line
 
     def add_elements(

{xtrack-0.45.0 → xtrack-0.46.0}/xtrack/match.py

@@ -10,6 +10,7 @@ import xtrack as xt
 import xdeps as xd
 
 XTRACK_DEFAULT_TOL = 1e-10
+XTRACK_DEFAULT_SIGMA_REL = 0.01
 
 XTRACK_DEFAULT_WEIGHTS = {
     # For quantities not specified here the default weight is 1
@@ -148,19 +149,149 @@ class ActionTwiss(xd.Action):
                     ' for Multiline')
             self.line.compensate_radiation_energy_loss(verbose=False)
         if not self.allow_twiss_failure or not allow_failure:
-            return self.line.twiss(**self.kwargs)
+            out = self.line.twiss(**self.kwargs)
         else:
             try:
-                return self.line.twiss(**self.kwargs)
+                out = self.line.twiss(**self.kwargs)
             except Exception as ee:
                 if allow_failure:
                     return 'failed'
                 else:
                     raise ee
+        out.line = self.line
+        return out
+
+# Alternative transitions functions
+# def _transition_sigmoid_integral(x):
+#     x_shift = x - 3
+#     if x_shift > 10:
+#         return x_shift
+#     else:
+#         return np.log(1 + np.exp(x_shift))
+
+# def _transition_sin(x):
+#     if x < 0:
+#         return 0
+#     if x < 1.:
+#         return 2 /np.pi - 2 /np.pi * np.cos(np.pi * x / 2)
+#     else:
+#         return x + 2 / np.pi - 1
+
+def _poly(x):
+     return 3 * x**3 - 2 * x**4
+
+def _transition_poly(x):
+        x_cut = 1/16 + np.sqrt(33)/16
+        if x < 0:
+            return 0
+        if x < x_cut:
+            return _poly(x)
+        else:
+            return x - x_cut + _poly(x_cut)
+
+class GreaterThan:
+
+    _transition = staticmethod(_transition_poly)
+
+    def __init__(self, lower, mode='step', sigma=None,
+                 sigma_rel=XTRACK_DEFAULT_SIGMA_REL):
+        assert mode in ['step', 'smooth']
+        self.lower = lower
+        self._value = 0.
+        self.mode=mode
+        if mode == 'smooth':
+            assert sigma is not None or sigma_rel is not None
+            if sigma is not None:
+                assert sigma_rel is None
+                self.sigma = sigma
+            else:
+                assert sigma_rel is not None
+                self.sigma = np.abs(self.lower) * sigma_rel
+
+    def auxtarget(self, res):
+        '''Transformation applied to target value to obtain the corresponding
+        cost function.
+        '''
+        if self.mode == 'step':
+            if res < self.lower:
+                return res - self.lower
+            else:
+                return 0
+        elif self.mode == 'smooth':
+            return self.sigma * self._transition((self.lower - res) / self.sigma)
+        elif self.mode == 'auxvar':
+            raise NotImplementedError # experimental
+            return res - self.lower - self.vary.container[self.vary.name]**2
+        else:
+            raise ValueError(f'Unknown mode {self.mode}')
+
+    def __repr__(self):
+        return f'GreaterThan({self.lower:4g})'
+
+    # Part of the `auxvar` experimental code
+    # def _set_value(self, val, target):
+    #     self.lower = val
+    #     aux_vary_container = self.vary.container
+    #     aux_vary_container[self.vary.name] = 0
+    #     val = target.runeval()
+    #     if val > 0:
+    #         aux_vary_container[self.vary.name] = np.sqrt(val)
+    # def gen_vary(self, container):
+    #     self.vary = _gen_vary(container)
+    #     return self.vary
+
+class LessThan:
+
+    _transition = staticmethod(_transition_poly)
+
+    def __init__(self, upper, mode='step', sigma=None,
+                 sigma_rel=XTRACK_DEFAULT_SIGMA_REL):
+        assert mode in ['step', 'smooth']
+        self.upper = upper
+        self._value = 0.
+        self.mode=mode
+        if mode == 'smooth':
+            assert sigma is not None or sigma_rel is not None
+            if sigma is not None:
+                assert sigma_rel is None
+                self.sigma = sigma
+            else:
+                assert sigma_rel is not None
+                self.sigma = np.abs(self.upper) * sigma_rel
+
+    def auxtarget(self, res):
+        if self.mode == 'step':
+            if res > self.upper:
+                return self.upper - res
+            else:
+                return 0
+        elif self.mode == 'smooth':
+            return self.sigma * self._transition((res - self.upper) / self.sigma)
+        elif self.mode == 'auxvar':
+            raise NotImplementedError # experimental
+            return self.upper - res - self.vary.container[self.vary.name]**2
+        else:
+            raise ValueError(f'Unknown mode {self.mode}')
+
+    def __repr__(self):
+        return f'LessThan({self.upper:4g})'
+
+# part of the `auxvar` experimental code
+# def _gen_vary(container):
+#     for ii in range(10000):
+#         if f'auxvar_{ii}' not in container:
+#             vv = f'auxvar_{ii}'
+#             break
+#     else:
+#         raise RuntimeError('Too many auxvary variables')
+#     container[vv] = 0
+#     return xt.Vary(name=vv, container=container, step=1e-3)
+
 
 class Target(xd.Target):
     def __init__(self, tar=None, value=None, at=None, tol=None, weight=None, scale=None,
-                 line=None, action=None, tag='', optimize_log=False, **kwargs):
+                 line=None, action=None, tag='', optimize_log=False,
+                 **kwargs):
 
         for kk in kwargs:
             assert kk in ALLOWED_TARGET_KWARGS, (
@@ -180,6 +311,9 @@ class Target(xd.Target):
             xdtar = (tar, at)
         else:
             xdtar = tar
+
+        self._freeze_value = None
+
         xd.Target.__init__(self, tar=xdtar, value=value, tol=tol,
                             weight=weight, scale=scale, action=action, tag=tag,
                             optimize_log=optimize_log)
@@ -196,9 +330,38 @@ class Target(xd.Target):
         if self.line is not None:
             res = res[self.line]
         if callable(self.tar):
-            return self.tar(res)
+            out = self.tar(res)
         else:
-            return res[self.tar]
+            out = res[self.tar]
+
+        if self._freeze_value is not None:
+            return out
+
+        return out
+
+    def transform(self, val):
+        if hasattr(self.value, 'auxtarget'):
+            return self.value.auxtarget(val)
+        else:
+            return val
+
+    @property
+    def value(self):
+        if self._freeze_value is not None:
+            return self._freeze_value
+        else:
+            return self._user_value
+
+    @value.setter
+    def value(self, val):
+        self._user_value = val
+
+    def freeze(self):
+        self._freeze_value = True # to bypass inequality logic
+        self._freeze_value = self.runeval()
+
+    def unfreeze(self):
+        self._freeze_value = None
 
 class Vary(xd.Vary):
     def __init__(self, name, container=None, limits=None, step=None, weight=None,
@@ -234,53 +397,43 @@ class TargetInequality(Target):
 
     def __init__(self, tar, ineq_sign, rhs, at=None, tol=None, scale=None,
                  line=None, weight=None, tag=''):
-        Target.__init__(self, tar, value=0, at=at, tol=tol, scale=scale, line=line,
-                         weight=weight, tag=tag)
-        assert ineq_sign in ['<', '>'], ('ineq_sign must be either "<" or ">"')
-        self.ineq_sign = ineq_sign
-        self.rhs = rhs
 
-    def __repr__(self):
-        return f'TargetInequality({self.tar} {self.ineq_sign} {self.rhs}, tol={self.tol}, weight={self.weight})'
-
-    def eval(self, tw):
-        val = super().eval(tw)
-        if self.ineq_sign == '<' and val < self.rhs:
-            return 0
-        elif self.ineq_sign == '>' and val > self.rhs:
-            return 0
-        else:
-            return val - self.rhs
+        raise NotImplementedError('TargetInequality is not anymore supported. '
+            'Please use Target with `GreaterThan` `LessThan` instead. '
+            'For example, instead of '
+            'TargetInequality("x", "<", 0.1, at="ip1") '
+            'use '
+            'Target("x", LessThan(0.1), at="ip1")')
 
 class TargetRelPhaseAdvance(Target):
 
-    def __init__(self, tar, value, at_1=None, at_0=None, tag='',  **kwargs):
+    def __init__(self, tar, value, ele_stop=None, ele_start=None, tag='',  **kwargs):
 
         Target.__init__(self, tar=self.compute, value=value, tag=tag, **kwargs)
 
         assert tar in ['mux', 'muy'], 'Only mux and muy are supported'
         self.var = tar
-        if at_1 is None:
-            at_1 = '__ele_stop__'
-        if at_0 is None:
-            at_0 = '__ele_start__'
-        self.at_1 = at_1
-        self.at_0 = at_0
+        if ele_stop is None:
+            ele_stop = '__ele_stop__'
+        if ele_start is None:
+            ele_start = '__ele_start__'
+        self.ele_stop = ele_stop
+        self.ele_start = ele_start
 
     def __repr__(self):
-        return f'TargetPhaseAdvance({self.var}({self.at_1}) - {self.var}({self.at_0}), value={self.value}, tol={self.tol}, weight={self.weight})'
+        return f'TargetPhaseAdv({self.var}({self.ele_stop} - {self.ele_start}), val={self.value}, tol={self.tol}, weight={self.weight})'
 
     def compute(self, tw):
 
-        if self.at_1 == '__ele_stop__':
+        if self.ele_stop == '__ele_stop__':
             mu_1 = tw[self.var, -1]
         else:
-            mu_1 = tw[self.var, self.at_1]
+            mu_1 = tw[self.var, self.ele_stop]
 
-        if self.at_0 == '__ele_start__':
+        if self.ele_start == '__ele_start__':
             mu_0 = tw[self.var, 0]
         else:
-            mu_0 = tw[self.var, self.at_0]
+            mu_0 = tw[self.var, self.ele_start]
 
         return mu_1 - mu_0
 
@@ -298,6 +451,9 @@ def match_line(line, vary, targets, restore_if_fail=True, solver=None,
         else:
             targets_flatten.append(tt.copy())
 
+    aux_vary_container = {}
+    aux_vary = []
+
     action_twiss = None
     for tt in targets_flatten:
 
@@ -342,9 +498,20 @@ def match_line(line, vary, targets, restore_if_fail=True, solver=None,
             else:
                 tt.tol = default_tol
 
+        # part of the `auxvar` experimental code
+        # if isinstance(tt.value, (GreaterThan, LessThan)):
+        #     if tt.value.mode == 'auxvar':
+        #         aux_vary.append(tt.value.gen_vary(aux_vary_container))
+        #         aux_vary_container[aux_vary[-1].name] = 0
+        #         val = tt.runeval()
+        #         if val > 0:
+        #             aux_vary_container[aux_vary[-1].name] = np.sqrt(val)
+
     if not isinstance(vary, (list, tuple)):
         vary = [vary]
 
+    vary = list(vary) + aux_vary
+
     vary_flatten = _flatten_vary(vary)
     _complete_vary_with_info_from_line(vary_flatten, line)
 

{xtrack-0.45.0 → xtrack-0.46.0}/xtrack/multisetter/multisetter.py

@@ -114,9 +114,7 @@ class MultiSetter(xo.HybridClass):
     }
 
     def __init__(self, line, elements, field, index=None):
-
-        '''
-        Create object to efficiently set and get values of a specific field of
+        """Create object to efficiently set and get values of a specific field of
         several elements of a line.
 
         Parameters
@@ -129,8 +127,7 @@ class MultiSetter(xo.HybridClass):
             Name of the field to be mutated.
         index: int or None
             If the field is an array, the index of the array to be mutated.
-
-        '''
+        """
 
         if isinstance(line, xt.Tracker):
             tracker = line
@@ -142,6 +139,13 @@ class MultiSetter(xo.HybridClass):
         tracker_buffer = tracker._buffer
         line = tracker.line
 
+        if len(elements) == 0:
+            self._empty = True
+            self.xoinitialize(_context=context, offsets=[])
+            return
+
+        self._empty = False
+
         # Get dtype from first element
         el = line[elements[0]]
         dd = getattr(el.copy(_context=xo.context_default), field)
@@ -178,10 +182,9 @@ class MultiSetter(xo.HybridClass):
         }[self.dtype]
 
     def get_values(self):
-
-        '''
-        Get the values of the multisetter fields.
-        '''
+        """Get the values of the multisetter fields."""
+        if self._empty:
+            return self._context.zeros(0, dtype=np.float64)
 
         out = self._context.zeros(len(self.offsets), dtype=self.dtype)
         self._get_kernel.set_n_threads(len(self.offsets))
@@ -189,19 +192,19 @@ class MultiSetter(xo.HybridClass):
         return out
 
     def set_values(self, values):
-
-        '''
-        Set the values of the multisetter fields.
+        """Set the values of the multisetter fields.
 
         Parameters
         ----------
         values: np.ndarray
             Array of values to be set.
-        '''
+        """
+        if self._empty:
+            return
 
         self._set_kernel.set_n_threads(len(self.offsets))
         self._set_kernel(data=self, buffer=self._tracker_buffer.buffer,
-               input=xt.BeamElement._arr2ctx(self,values))
+               input=xt.BeamElement._arr2ctx(self, values))
 
 
 def _extract_offset(obj, field_name, index, dtype, xodtype):

xtrack-0.46.0/xtrack/progress_indicator.py

@@ -0,0 +1,93 @@
+from math import ceil
+from typing import cast, Iterable, Collection
+
+from .general import _print
+
+
+class DefaultProgressIndicator:
+    """Display progress of a task.
+
+    This is a simple progress indicator, with an API that is a subset of the
+    one of `tqdm`. It will provide simple progress information if tqdm is
+    missing on the system.
+
+    Parameters
+    ----------
+    iterable: Iterable
+        The iterable to iterate over.
+    desc: str
+        Description of the task.
+    total: int, optional
+        Total number of iterations, if unspecified `len(iterable)` is used.
+    miniters: int, optional
+        Minimum number of iterations between updates, by default 1.
+    unit_scale: int, optional
+        Unused, kept for compatibility with tqdm. To be interpreted as the
+        scale by which the iteration number and `total` are multiplied: i.e.,
+        one iteration corresponds to `unit_scale` events.
+    """
+    def __init__(
+            self,
+            iterable: Iterable,
+            desc: str,
+            total: int = None,
+            miniters: int = None,
+            unit_scale: int = None,
+    ):
+        self.iterable = iterable
+        self.desc = desc
+        self._iterator = None
+        self._iteration = 0
+        self._total = total or len(cast(Collection, iterable))
+        self._update_interval = miniters or ceil(self._total / 100)
+        self._unit_scale = unit_scale or 1
+        print(f'Init: interval {self._update_interval}, total {self._total}')
+
+    def __iter__(self):
+        self._iterator = iter(self.iterable)
+        return self
+
+    def __next__(self):
+        if self._iteration % self._update_interval == 0:
+            self._print_progress()
+
+        try:
+            return next(self._iterator)
+        except StopIteration:
+            self._print_progress()
+            _print()  # finish with a newline
+            raise
+        finally:
+            self._iteration += 1
+
+    def _print_progress(self):
+        percent = round(self._iteration * 100 / self._total)
+        index = self._iteration * self._unit_scale
+        scaled_total = self._total * self._unit_scale
+        message = f'{self.desc}: {percent:3}% ({index}/{scaled_total} iterations)'
+        _print(message, end='\r', flush=True)
+
+
+class _ProgressIndicatorConfig:
+    default_indicator_cls = DefaultProgressIndicator
+    default_options = {}
+
+
+_config = _ProgressIndicatorConfig()
+
+
+def set_default_indicator(indicator_cls, **options):
+    _config.default_indicator_cls = indicator_cls
+    _config.default_options = options
+
+
+def progress(iterable: Iterable, **options):
+    indicator = _config.default_indicator_cls
+    return indicator(iterable, **_config.default_options, **options)
+
+
+try:
+    from tqdm.autonotebook import tqdm
+    set_default_indicator(tqdm)
+except ModuleNotFoundError:
+    pass