Merge pull request #50 from willrogers/element-index

Add lattice index to Element class.
DiamondLightSource · Nov 16, 2017 · 77464d6 · 77464d6
2 parents ed00e52 + 31a68fd
commit 77464d6
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/pytac/element.py b/pytac/element.py
@@ -15,24 +15,28 @@ class Element(object):
         type_ (str): type of the element
         length (number): length of the element in metres
         s (float): the element's start position within the lattice in metres
+        index (int): the element's index within the ring, starting at 1
         cell (int): the element's cell within the lattice
         families (set): the families this element is a member of
 
     """
-    def __init__(self, name, length, element_type, s=None, cell=None):
+    def __init__(self, name, length, element_type,
+                 s=None, index=None, cell=None):
         """
         Args:
             name (int): Unique identifier for the element in the ring.
             length (float): The length of the element.
             element_type (str): Type of the element.
             s (float): Position of the start of the element in the ring
+            index (float): Index of the element in the ring, starting at 1
             cell (int): lattice cell this element is wihin
 
         """
         self.name = name
         self.type_ = element_type
         self.length = length
         self.s = s
+        self.index = index
         self.cell = cell
         self.families = set()
         self._uc = dict()

diff --git a/pytac/load_csv.py b/pytac/load_csv.py
@@ -90,15 +90,18 @@ def load(mode, control_system=None, directory=None):
                                  'data')
     lat = lattice.Lattice(mode, control_system, 3000)
     s = 0
+    index = 1
     with open(os.path.join(directory, mode, 'elements.csv')) as elements:
         csv_reader = csv.DictReader(elements)
         for item in csv_reader:
             length = float(item['length'])
             cell = int(item['cell']) if item['cell'] else None
-            e = element.Element(item['name'], length, item['type'], s, cell)
+            e = element.Element(item['name'], length,
+                                item['type'], s, index, cell)
             e.add_to_family(item['type'])
             lat.add_element(e)
             s += length
+            index += 1
 
     with open(os.path.join(directory, mode, 'devices.csv')) as devices:
         csv_reader = csv.DictReader(devices)

diff --git a/test/test_load.py b/test/test_load.py
@@ -24,6 +24,7 @@ def test_element_details_loaded(lattice):
     quad = lattice.get_elements('quad')[0]
     assert quad.cell == 1
     assert quad.s == 1.0
+    assert quad.index == 2
 
 
 def test_devices_loaded(lattice):