xtrack 0.73.4__tar.gz → 0.73.5__tar.gz

{xtrack-0.73.4/xtrack.egg-info → xtrack-0.73.5}/PKG-INFO

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

{xtrack-0.73.4 → xtrack-0.73.5}/tests/test_line.py

@@ -5,6 +5,7 @@
 
 import pathlib
 import pickle
+import math
 
 import numpy as np
 import pytest
@@ -1138,3 +1139,159 @@ def test_line_table_unique_names():
         if name == '_end_point': continue
         assert line[name] == line[env_name]
 
+
+def test_extend_knl_ksl():
+
+    classes_to_check = ['Bend', 'Quadrupole', 'Sextupole', 'Octupole', 'Solenoid',
+                        'Multipole']
+
+    for cc in classes_to_check:
+
+        nn1 = 'test1_'+cc.lower()
+        nn2 = 'test2_'+cc.lower()
+        env = xt.Environment()
+        env.new(nn1, cc, length=10, knl=[
+                1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], ksl=[3, 2, 1])
+        env.new(nn2, cc, length=10, ksl=[
+                1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], knl=[3, 2, 1], order=11)
+
+        assert env[nn1].__class__.__name__ == cc
+        assert env[nn1].order == 11
+        assert len(env[nn1].knl) == 12
+        assert len(env[nn1].ksl) == 12
+        xo.assert_allclose(env[nn1].knl, [1, 2, 3, 4, 5,
+                           6, 7, 8, 9, 10, 11, 12], rtol=0, atol=1e-15)
+        xo.assert_allclose(env[nn1].ksl, [3, 2, 1, 0, 0,
+                           0, 0, 0, 0, 0, 0, 0], rtol=0, atol=1e-15)
+        xo.assert_allclose(env[nn1].inv_factorial_order,
+                           1/math.factorial(11), rtol=0, atol=1e-15)
+
+        assert env[nn2].__class__.__name__ == cc
+        assert env[nn2].order == 11
+        assert len(env[nn2].ksl) == 12
+        assert len(env[nn2].knl) == 12
+        xo.assert_allclose(env[nn2].ksl, [1, 2, 3, 4, 5,
+                           6, 7, 8, 9, 10, 11, 12], rtol=0, atol=1e-15)
+        xo.assert_allclose(env[nn2].knl, [3, 2, 1, 0, 0,
+                           0, 0, 0, 0, 0, 0, 0], rtol=0, atol=1e-15)
+        xo.assert_allclose(env[nn2].inv_factorial_order,
+                           1/math.factorial(11), rtol=0, atol=1e-15)
+
+    env.vars.default_to_zero = True
+    line = env.new_line(components=[
+        env.new('b1', xt.Bend, length=1, knl=[
+                'a', 'b', 'c'], ksl=['d', 'e', 'f']),
+        env.new('q1', xt.Quadrupole, length=1, knl=[
+                'a', 'b', 'c'], ksl=['d', 'e', 'f']),
+        env.new('s1', xt.Sextupole, length=1, knl=[
+                'a', 'b', 'c'], ksl=['d', 'e', 'f']),
+        env.new('o1', xt.Octupole, length=1, knl=[
+                'a', 'b', 'c'], ksl=['d', 'e', 'f']),
+        env.new('s2', xt.Solenoid, length=1, knl=[
+                'a', 'b', 'c'], ksl=['d', 'e', 'f']),
+        env.new('m1', xt.Multipole, length=1, knl=[
+                'a', 'b', 'c'], ksl=['d', 'e', 'f']),
+    ])
+
+    env['a'] = 3.
+    env['b'] = 2.
+    env['c'] = 1.
+    env['d'] = 4.
+    env['e'] = 5.
+    env['f'] = 6.
+
+    element_names = ['b1', 'q1']
+    order = 10
+
+    line.extend_knl_ksl(order=order, element_names=element_names)
+
+    assert line['b1'].order == order
+    assert line['q1'].order == order
+    assert line['s1'].order == 5
+    assert line['o1'].order == 5
+    assert line['s2'].order == 5
+    assert line['m1'].order == 2
+
+    xo.assert_allclose(line['b1'].inv_factorial_order,
+                       1/math.factorial(order), rtol=0, atol=1e-15)
+    xo.assert_allclose(line['q1'].inv_factorial_order,
+                       1/math.factorial(order), rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s1'].inv_factorial_order,
+                       1/math.factorial(5), rtol=0, atol=1e-15)
+    xo.assert_allclose(line['o1'].inv_factorial_order,
+                       1/math.factorial(5), rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s2'].inv_factorial_order,
+                       1/math.factorial(5), rtol=0, atol=1e-15)
+    xo.assert_allclose(line['m1'].inv_factorial_order,
+                       1/math.factorial(2), rtol=0, atol=1e-15)
+
+    xo.assert_allclose(line['b1'].knl, [3., 2., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['b1'].ksl, [4., 5., 6., 0.,
+                       0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['q1'].knl, [3., 2., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['q1'].ksl, [4., 5., 6., 0.,
+                       0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s1'].knl, [3., 2., 1.,
+                       0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s1'].ksl, [4., 5., 6.,
+                       0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['o1'].knl, [3., 2., 1.,
+                       0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['o1'].ksl, [4., 5., 6.,
+                       0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s2'].knl, [3., 2., 1.,
+                       0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s2'].ksl, [4., 5., 6.,
+                       0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['m1'].knl, [3., 2., 1.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['m1'].ksl, [4., 5., 6.], rtol=0, atol=1e-15)
+
+    line.extend_knl_ksl(order=11)
+
+    assert line['b1'].order == 11
+    assert line['q1'].order == 11
+    assert line['s1'].order == 11
+    assert line['o1'].order == 11
+    assert line['s2'].order == 11
+    assert line['m1'].order == 11
+    assert line['b1'].inv_factorial_order == 1/math.factorial(11)
+    assert line['q1'].inv_factorial_order == 1/math.factorial(11)
+    assert line['s1'].inv_factorial_order == 1/math.factorial(11)
+    assert line['o1'].inv_factorial_order == 1/math.factorial(11)
+    assert line['s2'].inv_factorial_order == 1/math.factorial(11)
+    assert line['m1'].inv_factorial_order == 1/math.factorial(11)
+    xo.assert_allclose(line['b1'].knl, [3., 2., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['b1'].ksl, [4., 5., 6., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['q1'].knl, [3., 2., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['q1'].ksl, [4., 5., 6., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s1'].knl, [3., 2., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s1'].ksl, [4., 5., 6., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['o1'].knl, [3., 2., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['o1'].ksl, [4., 5., 6., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s2'].knl, [3., 2., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['s2'].ksl, [4., 5., 6., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['m1'].knl, [3., 2., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['m1'].ksl, [4., 5., 6., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+
+    # test an expression
+    line['b'] = 100
+    line['f'] = 200
+
+    xo.assert_allclose(line['o1'].knl, [3., 100., 1., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)
+    xo.assert_allclose(line['o1'].ksl, [4., 5., 200., 0.,
+                       0., 0., 0., 0., 0., 0., 0., 0.], rtol=0, atol=1e-15)

{xtrack-0.73.4 → xtrack-0.73.5}/tests/test_new_madloader.py

@@ -19,6 +19,7 @@ def test_simple_parser():
     if (version>=50401){option,-rbarc;};  ! to be ignored
     
     third = 1 / 3;
+    power = 3^4;
     hello := third * twopi;
     mb.l = 1;
     qd.l = 0.5;
@@ -45,6 +46,7 @@ def test_simple_parser():
     expected: MadxOutputType = {
         'vars': {
             'third': {'deferred': False, 'expr': '(1.0 / 3.0)'},
+            'power': {'deferred': False, 'expr': '(3.0 ** 4.0)'},
             'hello': {'deferred': True, 'expr': '(third * twopi)'},
             'mb.l': {'deferred': False, 'expr': 1.0},
             'qd.l': {'deferred': False, 'expr': 0.5},

xtrack-0.73.5/xtrack/_version.py

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

{xtrack-0.73.4 → xtrack-0.73.5}/xtrack/beam_elements/elements.py

@@ -819,7 +819,6 @@ class Bend(BeamElement):
 
         if model is not None:
             self.model = model
-        self.order = order
 
     def to_dict(self, copy_to_cpu=True):
         out = super().to_dict(copy_to_cpu=copy_to_cpu)
@@ -995,8 +994,6 @@ class Sextupole(BeamElement):
 
         self.xoinitialize(**kwargs)
 
-        self.order = order
-
     def to_dict(self, copy_to_cpu=True):
         out = super().to_dict(copy_to_cpu=copy_to_cpu)
 
@@ -1089,8 +1086,6 @@ class Octupole(BeamElement):
 
         self.xoinitialize(**kwargs)
 
-        self.order = order
-
     def to_dict(self, copy_to_cpu=True):
         out = super().to_dict(copy_to_cpu=copy_to_cpu)
 
@@ -1180,8 +1175,6 @@ class Quadrupole(BeamElement):
 
         self.xoinitialize(**kwargs)
 
-        self.order = order
-
     def to_dict(self, copy_to_cpu=True):
         out = super().to_dict(copy_to_cpu=copy_to_cpu)
 
@@ -1316,8 +1309,6 @@ class Solenoid(BeamElement):
 
         self.xoinitialize(**kwargs)
 
-        self.order = order
-
     @property
     def order(self):
         return self._order

{xtrack-0.73.4 → xtrack-0.73.5}/xtrack/line.py

@@ -3203,6 +3203,50 @@ class Line:
         self._check_valid_tracker()
         stop_internal_logging_for_elements_of_type(self.tracker, element_type)
 
+    def extend_knl_ksl(self, order, element_names=None):
+
+        """
+        Extend the order of the knl and ksl attributes of the elements.
+
+        Parameters
+        ----------
+        order: int
+            New order of the knl and ksl attributes.
+        element_names: list of str
+            Names of the elements to extend. If None, all elements having `knl`
+            and `ksl` attributes are extended.
+
+        """
+
+        self.discard_tracker()
+
+        if element_names is None:
+            element_names = []
+            for nn in self.element_names:
+                if hasattr(self.get(nn), 'knl'):
+                    element_names.append(nn)
+
+        for nn in element_names:
+            if self.get(nn).order > order:
+                raise ValueError(f'Order of element {nn} is smaller than {order}')
+
+        for nn in element_names:
+            ee = self.get(nn)
+
+            if ee.order == order:
+                continue
+
+            new_knl = [vv for vv in ee.knl] + [0] * (order - len(ee.knl) + 1)
+            new_ksl = [vv for vv in ee.ksl] + [0] * (order - len(ee.ksl) + 1)
+
+            dct = ee.to_dict()
+            dct.pop('order', None)
+            dct['knl'] = new_knl
+            dct['ksl'] = new_ksl
+
+            new_ee = ee.__class__.from_dict(dct, _buffer=ee._buffer)
+            self.env.elements[nn] = new_ee
+
     def remove_markers(self, inplace=True, keep=None):
         """
         Remove markers from the line

{xtrack-0.73.4 → xtrack-0.73.5}/xtrack/mad_parser/loader.py

@@ -298,14 +298,14 @@ class MadxLoader:
 
 
             if (isinstance(self.env[name], BeamElement) and not self.env[name].isthick
-                 and length and not isinstance(self.env[name], xt.Marker)):
+                    and length and not isinstance(self.env[name], xt.Marker)):
                 drift_name = f'drift_{name}'
                 self.env.new(drift_name, 'Drift', length=f'({length}) / 2')
                 name = builder.new_line([drift_name, name, drift_name])
             builder.place(name, **el_params)
         else:
             if (isinstance(self.env[parent], BeamElement) and not self.env[parent].isthick
-                and length and not isinstance(self.env[parent], xt.Marker)):
+                    and length and not isinstance(self.env[parent], xt.Marker)):
                 drift_name = f'{name}_drift'
                 self.env.new(drift_name, 'Drift', force=True, length=f'({length}) / 2')  # Not sure why `force` is needed
                 at, from_ = el_params.pop('at', None), el_params.pop('from_', None)

{xtrack-0.73.4 → xtrack-0.73.5}/xtrack/mad_parser/madx.lark

@@ -96,7 +96,7 @@ _ENDSEQUENCE: "endsequence"i
 _SEQEDIT: "seqedit"i
 _ENDEDIT: "endedit"i
 _LINE: "line"i
-NAME: /[A-z_\\.][A-z0-9_\\.%]*/
+NAME: /[A-Za-z_\\.][A-Za-z0-9_\\.%]*/
 STRING: /"[^"]*"/
 _COMMENT: "!" /[^\n]*/
        | "//" /[^\n]*/

{xtrack-0.73.4 → xtrack-0.73.5}/xtrack/mad_writer.py

@@ -284,7 +284,7 @@ def to_madx_sequence(line, name='seq', mode='sequence'):
         elements_str = ""
         for nn in line.element_names:
             el = line[nn]
-            el_str = xsuite_to_mad_conveters[type(el)](nn, line)
+            el_str = xsuite_to_mad_conveters[el.__class__](nn, line)
             elements_str += f"{nn}: {el_str};\n"
         line_str = f'{name}: line=(' + ', '.join(line.element_names) + ');'
         machine_str = elements_str + line_str
@@ -307,7 +307,7 @@ def to_madx_sequence(line, name='seq', mode='sequence'):
 
         for nn in line.element_names:
             el = line.element_dict[nn]
-            el_str = xsuite_to_mad_conveters[type(el)](nn, line)
+            el_str = xsuite_to_mad_conveters[el.__class__](nn, line)
             if nn + '_tilt_entry' in line.element_dict:
                 el_str += ", " + mad_assignment('tilt',
                             _ge(line.element_refs[nn + '_tilt_entry'].angle) / 180. * np.pi)

{xtrack-0.73.4 → xtrack-0.73.5/xtrack.egg-info}/PKG-INFO

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

xtrack-0.73.4/xtrack/_version.py

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