Split up io methods into per-tool submodules

__init__.py

@@ -2,7 +2,7 @@
 pyIOTA accelerator control and analysis framework
 """
 
-__version__ = '0.6.1'
+__version__ = '0.6.2'
 __author__  = "Nikita Kuklev"
 
 # General imports - without these, most things will not work

examples/read_devices.py

@@ -1,22 +1,6 @@
 import pyIOTA.acnet as acnet
 import pyIOTA.iota as iota
 
-# devs = [acnet.frontends.BPMDevice(b) for b in iota.BPMS.HA]
-# ds = acnet.frontends.BPMDeviceSet(name='bpms', members=devs, enforce_array_length=1000)
-# ds.adapter = acnet.frontends.ACL(fallback=False)
-# ds.readonce()
-
-# devices = [acnet.frontends.BPMDevice(b) for b in iota.BPMS.HA][0:1]
-# ds = acnet.frontends.BPMDeviceSet(name='bpms', members=devices, enforce_array_length=1000)
-# ds.adapter = acnet.frontends.ACNETRelay()
-# ds.readonce()
-# print(ds)
-
-# devices = [acnet.frontends.DoubleDevice(b) for b in iota.QUADS.ALL_CURRENTS][0:1]
-# ds = acnet.frontends.DeviceSet(name='bpms', members=devices)
-# ds.adapter = acnet.frontends.ACNETRelay()
-# ds.readonce()
-
 devices = [acnet.frontends.DoubleDevice(b) for b in iota.QUADS.ALL_CURRENTS]
 ds = acnet.frontends.DeviceSet(name='bpms', members=devices)
 ds.adapter = acnet.frontends.ACNETRelay()

lattice/lattice.py

@@ -1,123 +1,20 @@
 import functools
 import pathlib
 import sys
-import numpy as np
-import pandas as pd
+
 
 """
 Holds information about the lattice, mostly optics
 """
 
 
 class Lattice:
-    omega62 = np.array(
-        [[0, 1, 0, 0, 0, 0], [-1, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0], [0, 0, -1, 0, 0, 0], [0, 0, 0, 0, 0, 1],
-         [0, 0, 0, 0, -1, 0]])
-    omega6 = np.array(
-        [[0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 1], [-1, 0, 0, 0, 0, 0], [0, -1, 0, 0, 0, 0],
-         [0, 0, -1, 0, 0, 0]])
-    omega4 = np.array([[0, 0, 1, 0], [0, 0, 0, 1], [-1, 0, 0, 0], [0, -1, 0, 0]])
-    omega42 = np.array([[0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1], [0, 0, -1, 0]])
-    omega2 = np.array([[0, 1], [-1, 0]])
-
-    def __init__(self):
-        self.bpm_s = None
-        self.bpm_beta = None
-        self.bpm_phase = None
-        self.bpm_alpha = None
-        if pathlib.Path.home().as_posix() not in sys.path: sys.path.insert(1, pathlib.Path.home().as_posix())
-        import pymadx
-
-    def load_optics(self, lattice_file):
-        iota = pymadx.Data.Tfs(lattice_file)
-        bpms = iota.GetElementsOfType('MONITOR')
-        print(f'Found {len(bpms)} BPMS: {bpms}')
-
-        bpm_s_temp = {bpm['NAME'].replace('BPM', 'B'): bpm['S'] for bpm in bpms}
-        bpm_s = {'IBA1C': bpm_s_temp['IBA1']}
-        del bpm_s_temp['IBA1']
-        bpm_s.update(bpm_s_temp)
-
-        bpm_beta_temp = {bpm['NAME'].replace('BPM', 'B'): (bpm['BETX'], bpm['BETY']) for bpm in bpms}
-        bpm_beta = {'IBA1C': bpm_beta_temp['IBA1']}
-        del bpm_beta_temp['IBA1']
-        bpm_beta.update(bpm_beta_temp)
-
-        bpm_alpha_temp = {bpm['NAME'].replace('BPM', 'B'): (bpm['ALFX'], bpm['ALFY']) for bpm in bpms}
-        bpm_alpha = {'IBA1C': bpm_alpha_temp['IBA1']}
-        del bpm_alpha_temp['IBA1']
-        bpm_alpha.update(bpm_alpha_temp)
-
-        bpm_phase = {bpm['NAME'].replace('BPM', 'B'): (bpm['MUX'], bpm['MUY']) for bpm in bpms}
-        bpm_phase['IBA1C'] = bpm_phase['IBA1']
-        del bpm_phase['IBA1']
-
-        self.bpm_s = bpm_s
-        self.bpm_beta = bpm_beta
-        self.bpm_phase = bpm_phase
-        self.bpm_alpha = bpm_alpha
-        return bpm_s, bpm_phase, bpm_beta, bpm_alpha
-
-    def load_transfer_maps(self, file, filetype='madx'):
-        # Read in matrices
-        if filetype == 'madx':
-            f = pymadx.Data.Tfs(file)
-            df = pd.DataFrame.from_dict(f.data, orient='index', columns=f.columns)
-            df = df.loc[:, 'NAME':'R66']
-
-            matrix_entries = ['R{}{}'.format(i, j) for i in range(1, 7) for j in range(1, 7)]
-            data_matrix = df.loc[:, matrix_entries].values.reshape((-1, 6, 6))
-            split_matrices = [np.asfortranarray(data_matrix[i, ...]) for i in range(len(df))]
-
-            elements = df.index
-            assert len(np.unique(elements)) == len(elements)
-            assert len([el for el in elements if 'START' in el]) == 1
-
-            df.loc[:, 'M'] = split_matrices
-            self.maps_transfer = df
-            self.maps_source = 'madx'
-            return df
-        elif filetype == '6dsim':
-            with open(file, 'r') as f:
-                tmaps = {}
-                lines = f.readlines()
-                lines = lines[14:]
-                for i, l in enumerate(lines):
-                    if i % 7 == 0:
-                        src, to = l.split('->')[0].strip(), l.split('->')[1].strip()
-                    else:
-                        j = i % 7 - 1
-                        if j == 0: matr = []
-                        # print(list(filter(None, l.strip().split(' '))))
-                        matr.append(np.array(list(filter(None, l.strip().split(' ')))).astype(np.float))
-                        if j == 5: tmaps[to] = (src, np.stack(matr));  # print(dataarr)
-                tmaps['start'] = ('end', np.identity(6))
-            self.maps_transfer = tmaps
-            self.maps_source = '6dsim'
-            return tmaps
-        else:
-            raise
-
-    def load_oneturn_maps(self, file, filetype='madx'):
-        # Read in matrices
-        if filetype == 'madx':
-            f = pymadx.Data.Tfs(file)
-            df = pd.DataFrame.from_dict(f.data, orient='index', columns=f.columns)
-            df = df.loc[:, 'NAME':'RE66']
-
-            matrix_entries = ['RE{}{}'.format(i, j) for i in range(1, 7) for j in range(1, 7)]
-            data_matrix = df.loc[:, matrix_entries].values.reshape((-1, 6, 6))
-            split_matrices = [data_matrix[i, ...] for i in range(len(df))]
-
-            elements = df.index
-            assert len(np.unique(elements)) == len(elements)
-            assert len([el for el in elements if 'START' in el]) == 1
-
-            df.loc[:, 'M'] = split_matrices
-            self.maps_oneturn = df
-            return df
-        else:
-            raise
+
+
+
+
+
+
 
     def get_transport_matrix(self, start, end, df=None):
         if df is None:

madx/__init__.py

@@ -0,0 +1 @@
+from .io import *

madx/io.py

@@ -0,0 +1,106 @@
+__all__ = ['MADXOptics']
+
+from pathlib import Path
+import numpy as np
+import pandas as pd
+
+
+class MADXOptics:
+    """
+    This class contains several previous-gen optics io methods.
+    There are DEPRECATED - it is strongly suggested to use native OCELOT extensions in lattice module.
+    """
+
+    def __init__(self):
+        import pymadx
+        self.bpm_s = None
+        self.bpm_beta = None
+        self.bpm_phase = None
+        self.bpm_alpha = None
+        # if pathlib.Path.home().as_posix() not in sys.path:
+        #     sys.path.insert(1, pathlib.Path.home().as_posix())
+
+    def load_optics(self, lattice_file: Path):
+        """
+        DEPRECATED
+        Loads optics from MADX TFS file
+        :param self:
+        :param lattice_file:
+        :return:
+        """
+        iota = pymadx.Data.Tfs(lattice_file)
+        bpms = iota.GetElementsOfType('MONITOR')
+        print(f'Found {len(bpms)} BPMS: {bpms}')
+
+        bpm_s_temp = {bpm['NAME'].replace('BPM', 'B'): bpm['S'] for bpm in bpms}
+        bpm_s = {'IBA1C': bpm_s_temp['IBA1']}
+        del bpm_s_temp['IBA1']
+        bpm_s.update(bpm_s_temp)
+
+        bpm_beta_temp = {bpm['NAME'].replace('BPM', 'B'): (bpm['BETX'], bpm['BETY']) for bpm in bpms}
+        bpm_beta = {'IBA1C': bpm_beta_temp['IBA1']}
+        del bpm_beta_temp['IBA1']
+        bpm_beta.update(bpm_beta_temp)
+
+        bpm_alpha_temp = {bpm['NAME'].replace('BPM', 'B'): (bpm['ALFX'], bpm['ALFY']) for bpm in bpms}
+        bpm_alpha = {'IBA1C': bpm_alpha_temp['IBA1']}
+        del bpm_alpha_temp['IBA1']
+        bpm_alpha.update(bpm_alpha_temp)
+
+        bpm_phase = {bpm['NAME'].replace('BPM', 'B'): (bpm['MUX'], bpm['MUY']) for bpm in bpms}
+        bpm_phase['IBA1C'] = bpm_phase['IBA1']
+        del bpm_phase['IBA1']
+
+        self.bpm_s = bpm_s
+        self.bpm_beta = bpm_beta
+        self.bpm_phase = bpm_phase
+        self.bpm_alpha = bpm_alpha
+        return bpm_s, bpm_phase, bpm_beta, bpm_alpha
+
+    def load_transfer_maps(self, file: Path):
+        """
+        DEPRECATED
+        Reads in MADX transfer matrix file. You should really just use OCELOT...
+        :param file:
+        :param filetype:
+        :return:
+        """
+        f = pymadx.Data.Tfs(file)
+        df = pd.DataFrame.from_dict(f.data, orient='index', columns=f.columns)
+        df = df.loc[:, 'NAME':'R66']
+
+        matrix_entries = ['R{}{}'.format(i, j) for i in range(1, 7) for j in range(1, 7)]
+        data_matrix = df.loc[:, matrix_entries].values.reshape((-1, 6, 6))
+        split_matrices = [np.asfortranarray(data_matrix[i, ...]) for i in range(len(df))]
+
+        elements = df.index
+        assert len(np.unique(elements)) == len(elements)
+        assert len([el for el in elements if 'START' in el]) == 1
+
+        df.loc[:, 'M'] = split_matrices
+        self.maps_transfer = df
+        return df
+
+    def load_oneturn_maps(self, file: Path):
+        """
+        DEPRECATED
+        Reads in MADX transfer matrix file. You should really just use OCELOT...
+        :param file:
+        :param filetype:
+        :return:
+        """
+        f = pymadx.Data.Tfs(file)
+        df = pd.DataFrame.from_dict(f.data, orient='index', columns=f.columns)
+        df = df.loc[:, 'NAME':'RE66']
+
+        matrix_entries = ['RE{}{}'.format(i, j) for i in range(1, 7) for j in range(1, 7)]
+        data_matrix = df.loc[:, matrix_entries].values.reshape((-1, 6, 6))
+        split_matrices = [data_matrix[i, ...] for i in range(len(df))]
+
+        elements = df.index
+        assert len(np.unique(elements)) == len(elements)
+        assert len([el for el in elements if 'START' in el]) == 1
+
+        df.loc[:, 'M'] = split_matrices
+        self.maps_oneturn = df
+        return df

sixdsim/io.py

@@ -384,6 +384,31 @@ def parse_knobs(fpath: Path, verbose: bool = False) -> Dict:
     return {k.name: k for k in knobs}
 
 
+def parse_transfer_maps(file: Path):
+    """
+    Parses outpout of sixdsim transfer map export. Note special, old format.
+    :param file:
+    :return: Dictionary with key as target element name
+    """
+    with open(file, 'r') as f:
+        tmaps = {}
+        lines = f.readlines()
+        lines = lines[14:]
+        for i, l in enumerate(lines):
+            if i % 7 == 0:
+                src, to = l.split('->')[0].strip(), l.split('->')[1].strip()
+            else:
+                j = i % 7 - 1
+                if j == 0:
+                    matr = []
+                # print(list(filter(None, l.strip().split(' '))))
+                matr.append(np.array(list(filter(None, l.strip().split(' ')))).astype(np.float))
+                if j == 5:
+                    tmaps[to] = (src, np.stack(matr));  # print(dataarr)
+        tmaps['start'] = ('end', np.identity(6))
+    return tmaps
+
+
 class AbstractKnob:
     """
     Superclass of all knobs, which are collections of KnobVariables representing setpoints of devices

tbt/constants.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+
+"""
+Omega matrices are used to check for symplecticity
+"""
+omega62 = np.array(
+    [[0, 1, 0, 0, 0, 0],
+     [-1, 0, 0, 0, 0, 0],
+     [0, 0, 0, 1, 0, 0],
+     [0, 0, -1, 0, 0, 0],
+     [0, 0, 0, 0, 0, 1],
+     [0, 0, 0, 0, -1, 0]])
+
+omega6 = np.array(
+    [[0, 0, 0, 1, 0, 0],
+     [0, 0, 0, 0, 1, 0],
+     [0, 0, 0, 0, 0, 1],
+     [-1, 0, 0, 0, 0, 0],
+     [0, -1, 0, 0, 0, 0],
+     [0, 0, -1, 0, 0, 0]])
+
+omega4 = np.array([[0, 0, 1, 0], [0, 0, 0, 1], [-1, 0, 0, 0], [0, -1, 0, 0]])
+
+omega42 = np.array([[0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, 1], [0, 0, -1, 0]])
+
+omega2 = np.array([[0, 1], [-1, 0]])