xtrack 0.69.1__tar.gz → 0.69.3__tar.gz

{xtrack-0.69.1/xtrack.egg-info → xtrack-0.69.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.69.1
+Version: 0.69.3
 Summary: Tracking library for particle accelerators
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xtrack
@@ -22,5 +22,9 @@ Requires-Dist: nafflib; extra == "tests"
 Requires-Dist: PyHEADTAIL; extra == "tests"
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-mock; extra == "tests"
+Provides-Extra: notebooks
+Requires-Dist: jupyter; extra == "notebooks"
+Requires-Dist: ipympl; extra == "notebooks"
+Requires-Dist: xplt; extra == "notebooks"
 
 Tracking library for particle accelerators

{xtrack-0.69.1 → xtrack-0.69.3}/setup.py

@@ -49,5 +49,6 @@ setup(
         ],
     extras_require={
         'tests': ['cpymad', 'nafflib', 'PyHEADTAIL', 'pytest', 'pytest-mock'],
-        },
+        'notebooks': ['jupyter', 'ipympl', 'xplt'],
+    },
     )

{xtrack-0.69.1 → xtrack-0.69.3}/tests/test_environment.py

@@ -1377,6 +1377,26 @@ def test_assemble_ring_repeated_elements():
     tw_one_cell_stripped = tw_one_cell.rows[:-1] # remove _end_point
     xo.assert_allclose(tw_one_cell_stripped.betx, tw_one_cell_ref.betx, atol=0, rtol=5e-4)
 
+    cell_selected = arc.select(start='mid::2', end='mid::3')
+    tw_cell_selected = cell_selected.twiss4d()
+    tw_cell_selected_stripped = tw_cell_selected.rows[:-1] # remove _end_point
+    xo.assert_allclose(tw_cell_selected_stripped.betx, tw_one_cell_ref.betx, atol=0, rtol=5e-4)
+
+    tt_ring2 = ring2.get_table(attr=True)
+    assert tt_ring2['name', 39] == 'mq.f::1'
+    assert tt_ring2['name', 48] == 'mq.f::2'
+    assert ring2.element_names[39] == 'mq.f'
+    assert ring2.element_names[48] == 'mq.f'
+
+    ring2.replace_all_repeated_elements()
+    tt_ring2_after = ring2.get_table(attr=True)
+    assert tt_ring2_after['name', 39] == 'mq.f.1'
+    assert tt_ring2_after['name', 48] == 'mq.f.2'
+    assert ring2.element_names[39] == 'mq.f.1'
+    assert ring2.element_names[48] == 'mq.f.2'
+    assert str(ring2.ref['mq.f.1'].k1._expr) == "vars['kqf']"
+    assert str(ring2.ref['mq.f.2'].k1._expr) == "vars['kqf']"
+    assert ring2.get('mq.f.1') is not ring2.get('mq.f.2')
 
     # import matplotlib.pyplot as plt
     # plt.close('all')
@@ -1762,3 +1782,74 @@ def test_select_in_multiline():
             == "((-vars['acbch7.r1b1']) + vars['aaa'])")
     assert line_sel.get('mcbch.7r1.b1').knl[0] == 1e-6
     assert line.get('mcbch.7r1.b1').knl[0] == 1e-6
+
+@pytest.mark.parametrize('container_type', ['env', 'line'])
+def test_inpection_methods(container_type):
+
+    env = xt.Environment()
+
+    env.vars({
+        'k.1': 1.,
+        'a': 2.,
+        'b': '2 * a + k.1',
+    })
+
+    line = env.new_line([
+        env.new('bb', xt.Bend, k0='2 * b', length=3+env.vars['a'] + env.vars['b'],
+            h=5., ksl=[0, '3*b']),
+    ])
+
+    ee = {'env': env, 'line': line}[container_type]
+
+    # Line/Env methods (get, set, eval, get_expr, new_expr, info)
+    assert ee.get('b') == 2 * 2 + 1
+    assert ee.get('bb') is env.element_dict['bb']
+
+    assert str(ee.get_expr('b')) == "((2.0 * vars['a']) + vars['k.1'])"
+
+    assert ee.eval('3*a - sqrt(k.1)') == 5
+
+    ne = ee.new_expr('sqrt(3*a + 3)')
+    assert xd.refs.is_ref(ne)
+    assert str(ne) == "f.sqrt(((3.0 * vars['a']) + 3.0))"
+
+    ee.info('bb') # Check that it works
+    ee.info('b')
+    ee.info('a')
+
+    ee.set('c', '6*a')
+    assert ee.get('c') == 6 * 2
+
+    # Line/Env containers (env[...], env.ref[...]
+    assert ee['b'] == 2 * 2 + 1
+    assert type(ee['bb']).__name__ == 'View'
+    assert ee['bb'].__class__.__name__ == 'Bend'
+
+    # Vars methods (get, set, eval, get_expr, new_expr, info, get_table)
+    assert ee.vars.get('b') == 2 * 2 + 1
+
+    assert str(ee.vars.get_expr('b')) == "((2.0 * vars['a']) + vars['k.1'])"
+
+    assert ee.vars.eval('3*a - sqrt(k.1)') == 5
+
+    ne = ee.vars.new_expr('sqrt(3*a + 3)')
+    assert xd.refs.is_ref(ne)
+    assert str(ne) == "f.sqrt(((3.0 * vars['a']) + 3.0))"
+
+    ee.vars.info('b')
+    ee.vars.info('a')
+
+    ee.vars.set('d', '7*a')
+    assert ee.vars.get('d') == 7 * 2
+
+    assert xd.refs.is_ref(ee.vars['b'])
+
+    # View methods get_expr, get_value, get_info, get_table (for now)
+    assert xd.refs.is_ref(ee['bb'].get_expr('k0'))
+    assert str(ee['bb'].get_expr('k0')) == "(2.0 * vars['b'])"
+    assert ee['bb'].get_expr('k0')._value == 2 * (2 * 2 + 1)
+    assert ee['bb'].get_value('k0') == 2 * (2 * 2 + 1)
+
+    tt = ee['bb'].get_table()
+    assert tt['value', 'k0'] == 2 * (2 * 2 + 1)
+    assert tt['expr', 'k0'] == "(2.0 * vars['b'])"

{xtrack-0.69.1 → xtrack-0.69.3}/tests/test_survey.py

@@ -12,6 +12,11 @@ tilted = True
 orientation = 'acw'
 transform_to_actual_elements = True
 
+if hasattr(np, 'trapezoid'): # numpy >= 2.0
+    trapz = np.trapezoid
+else:
+    trapz = np.trapz
+
 @for_all_test_contexts
 @pytest.mark.parametrize(
     'slice_mode',
@@ -103,17 +108,17 @@ def test_survey_slicing(test_context, slice_mode, tilted, orientation,
 
     if not tilted and orientation == 'acw':
         assert_allclose(np.abs(sv.Y), 0, rtol=0, atol=1e-14)
-        assert_allclose(np.trapezoid(sv.X, sv.Z), -4.818 , rtol=0, # anti-clockwise
+        assert_allclose(trapz(sv.X, sv.Z), -4.818 , rtol=0, # anti-clockwise
                         atol=(2e-3 if slice_mode is not None else 0.5))
     elif not tilted and orientation == 'cw':
         assert_allclose(np.abs(sv.Y), 0, rtol=0, atol=1e-14)
-        assert_allclose(np.trapezoid(sv.X, sv.Z), 4.818 , rtol=0, # clockwise
+        assert_allclose(trapz(sv.X, sv.Z), 4.818 , rtol=0, # clockwise
                         atol=(2e-3 if slice_mode is not None else 0.5))
     elif tilted and orientation == 'acw':
         assert_allclose(np.abs(sv.X), 0, rtol=0, atol=1e-14)
-        assert_allclose(np.trapezoid(sv.Y, sv.Z), -4.818 , rtol=0, # anti-clockwise
+        assert_allclose(trapz(sv.Y, sv.Z), -4.818 , rtol=0, # anti-clockwise
                         atol=(2e-3 if slice_mode is not None else 0.5))
     elif tilted and orientation == 'cw':
         assert_allclose(np.abs(sv.X), 0, rtol=0, atol=1e-14)
-        assert_allclose(np.trapezoid(sv.Y, sv.Z), 4.818 , rtol=0, # clockwise
+        assert_allclose(trapz(sv.Y, sv.Z), 4.818 , rtol=0, # clockwise
                         atol=(2e-3 if slice_mode is not None else 0.5))

xtrack-0.69.3/xtrack/_version.py

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

{xtrack-0.69.1 → xtrack-0.69.3}/xtrack/environment.py

@@ -40,6 +40,12 @@ class Environment:
         self._drift_counter = 0
         self.ref = EnvRef(self)
 
+    @property
+    def manager(self):
+        if not hasattr(self, '_var_management') or self._var_management is None:
+            self._init_var_management()
+        return self._var_management['manager']
+
     def new_line(self, components=None, name=None):
         out = xt.Line()
         out.particle_ref = self.particle_ref
@@ -198,6 +204,11 @@ Environment.replace_replica = xt.Line.replace_replica
 Environment.__getitem__ = xt.Line.__getitem__
 Environment.set = xt.Line.set
 Environment.get = xt.Line.get
+Environment.eval = xt.Line.eval
+Environment.info = xt.Line.info
+Environment.get_expr = xt.Line.get_expr
+Environment.new_expr = xt.Line.new_expr
+
 
 class Place:
 
@@ -356,7 +367,7 @@ def _resolve_s_positions(seq_all_places, env):
             aux_s_center.append(ss)
     aux_tt['s_center'] = np.concatenate([aux_s_center, [0]])
 
-    i_sorted = np.argsort(aux_s_center, stable=True)
+    i_sorted = np.argsort(aux_s_center, kind='stable')
 
     name_sorted = [str(aux_tt.name[ii]) for ii in i_sorted]
 

{xtrack-0.69.1 → xtrack-0.69.3}/xtrack/line.py

@@ -3405,7 +3405,8 @@ class Line:
 
     def __add__(self, other):
         self._env_if_needed
-        assert isinstance(other, Line), 'Only Line can be added to Line'
+        #assert isinstance(other, Line), 'Only Line can be added to Line'
+        assert other.__class__.__name__=="Line", 'Only Line can be added to Line'
         assert other.env is self.env, 'Lines must be in the same environment'
         out = self.env.new_line(
             components=list(self.element_names) + list(other.element_names))
@@ -3458,6 +3459,25 @@ class Line:
             if isinstance(self.element_dict[nn], xt.Replica):
                 self.replace_replica(nn)
 
+    def replace_all_repeated_elements(self, separator='.', mode='clone'):
+
+        self._env_if_needed()
+        env = self.env
+
+        self.discard_tracker()
+        unique_names = list(set(self.element_names))
+        aux_dict = {nn: [] for nn in unique_names}
+        for ii, nn in enumerate(self.element_names):
+            aux_dict[nn].append(ii)
+
+        for nn in unique_names:
+            if len(aux_dict[nn]) > 1:
+                i_rep = 0
+                for ii in aux_dict[nn]:
+                    while (new_name := nn + separator + str(i_rep)) in self.element_dict:
+                        i_rep += 1
+                    env.new(new_name, nn, mode=mode)
+                    self.element_names[ii] = new_name
 
     def select(self, start=None, end=None, name=None):
 
@@ -3473,7 +3493,7 @@ class Line:
 
         self._env_if_needed()
 
-        out = self.env.new_line(components=list(tt.name), name=name)
+        out = self.env.new_line(components=list(tt.env_name), name=name)
 
         if hasattr(self, '_in_multiline') and self._in_multiline is not None:
             out.env._var_management = None
@@ -3536,6 +3556,35 @@ class Line:
         else:
             raise KeyError(f'Element or variable {key} not found')
 
+    def info(self, key, limit=12):
+        if key in self.element_dict:
+            return self[key].get_info()
+        elif key in self.vars:
+            return self.vars.info(key, limit=limit)
+        else:
+            raise KeyError(f'Element or variable {key} not found')
+
+#    def get_value(self, key):
+#        if key in self.element_dict:
+#            return self.element_dict[key].get_value()
+#        elif key in self.vars:
+#            return self.vars.get_value(key)
+#        else:
+#            raise KeyError(f'Element or variable {key} not found')
+
+    @property
+    def manager(self):
+        return self._xdeps_manager
+
+    def eval(self, expr):
+        return self.vars.eval(expr)
+
+    def new_expr(self, expr):
+        return self.vars.new_expr(expr)
+
+    def get_expr(self, vars):
+        return self.vars.get_expr(vars)
+
     def _env_if_needed(self):
         if not hasattr(self, 'env') or self.env is None:
             self.env = xt.Environment(element_dict=self.element_dict,
@@ -4653,6 +4702,10 @@ class LineVars:
         out = list(self.line._xdeps_vref._owner.keys()).copy()
         return out
 
+    def __iter__(self):
+        raise NotImplementedError('Use keys() method') # Untested
+        return self.line._xdeps_vref._owner.__iter__()
+
     def update(self, other):
         if self.line._xdeps_vref is None:
             raise RuntimeError(
@@ -4671,10 +4724,23 @@ class LineVars:
         if self.line._xdeps_vref is None:
             raise RuntimeError(
                 f'Cannot access variables as the line has no xdeps manager')
-        name = np.array([kk for kk in list(self.keys()) if kk != '__vary_default'])
+        name = np.array([kk for kk in list(self.keys()) if kk != '__vary_default'], dtype=object)
         value = np.array([self.line._xdeps_vref[kk]._value for kk in name])
+        expr  = np.array([str(self.line._xdeps_vref[str(kk)]._expr) for kk in name])
 
-        return xd.Table({'name': name, 'value': value})
+        return xd.Table({'name': name, 'value': value, 'expr': expr})
+
+    def new_expr(self, expr):
+        return self.line._xdeps_eval.eval(expr)
+
+    def eval(self, expr):
+        return self.new_expr(expr)._get_value()
+
+    def info(self, var, limit=10):
+        return self[var]._info(limit=limit)
+
+    def get_expr(self, var):
+        return self[var]._expr
 
     def __contains__(self, key):
         if self.line._xdeps_vref is None:
@@ -4799,6 +4865,9 @@ class LineVars:
         else:
             self[name] = value
 
+    def get(self, name):
+        return self[name]._value
+
 class ActionVars(Action):
 
     def __init__(self, line):

{xtrack-0.69.1 → xtrack-0.69.3}/xtrack/survey.py

@@ -168,13 +168,21 @@ class SurveyTable(Table):
         return out
 
     def plot(self, element_width=None, legend=True, **kwargs):
+        # Shallow copy of self
+        out_sv_table = SurveyTable.__new__(SurveyTable)
+        out_sv_table.__dict__.update(self.__dict__)
+        out_sv_table._data = self._data.copy()
+
+        # Removing the count for repeated elements
+        out_sv_table.name = np.array([nn.split('::')[0] for nn in out_sv_table.name])
+
         if element_width is None:
             x_range = max(self.X) - min(self.X)
             y_range = max(self.Y) - min(self.Y)
             z_range = max(self.Z) - min(self.Z)
             element_width = max([x_range, y_range, z_range]) * 0.03
         import xplt
-        xplt.FloorPlot(self, self.line, element_width=element_width, **kwargs)
+        xplt.FloorPlot(out_sv_table, self.line, element_width=element_width, **kwargs)
         if legend:
             import matplotlib.pyplot as plt
             plt.legend()

{xtrack-0.69.1 → xtrack-0.69.3/xtrack.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.69.1
+Version: 0.69.3
 Summary: Tracking library for particle accelerators
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xtrack
@@ -22,5 +22,9 @@ Requires-Dist: nafflib; extra == "tests"
 Requires-Dist: PyHEADTAIL; extra == "tests"
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-mock; extra == "tests"
+Provides-Extra: notebooks
+Requires-Dist: jupyter; extra == "notebooks"
+Requires-Dist: ipympl; extra == "notebooks"
+Requires-Dist: xplt; extra == "notebooks"
 
 Tracking library for particle accelerators

{xtrack-0.69.1 → xtrack-0.69.3}/xtrack.egg-info/requires.txt

@@ -5,6 +5,11 @@ tqdm
 xobjects
 xdeps
 
+[notebooks]
+jupyter
+ipympl
+xplt
+
 [tests]
 cpymad
 nafflib

xtrack-0.69.1/xtrack/_version.py

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