xtrack 0.69.1__tar.gz → 0.69.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.69.1
+Version: 0.69.2
 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.2}/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.2}/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')

xtrack-0.69.2/xtrack/_version.py

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

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

@@ -3458,6 +3458,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 +3492,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
@@ -4653,6 +4672,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 +4694,32 @@ 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, 'expr': expr})
 
-        return xd.Table({'name': name, 'value': value})
+    def expr(self, var):
+        raise NotImplementedError # Untested
+        if isinstance(var,str):
+            ref=self.line._xdeps_vref[var]
+        elif is_expr(var):
+            ref=var
+        else:
+            raise ValueError(f"`{var}` not valid, must be str or expr")
+        expr=ref._expr
+        if expr is None:
+            raise NameError(f"`{var}` does not have any expression")
+        return expr
+
+    def eval(self, expr):
+        raise NotImplementedError # Untested
+        return self.line._xdeps_eval.eval(expr)
+
+    def value(self, expr):
+        raise NotImplementedError # Untested
+        return eval(self)._get_value()
 
     def __contains__(self, key):
         if self.line._xdeps_vref is None:

{xtrack-0.69.1 → xtrack-0.69.2}/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.2/xtrack.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.69.1
+Version: 0.69.2
 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.2}/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'