diff --git a/__init__.py b/__init__.py index 46a347a..1ad9757 100644 --- a/__init__.py +++ b/__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 diff --git a/examples/read_devices.py b/examples/read_devices.py index 4af3a42..edc3309 100644 --- a/examples/read_devices.py +++ b/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() diff --git a/lattice/lattice.py b/lattice/lattice.py index 3c3ae10..22dc882 100644 --- a/lattice/lattice.py +++ b/lattice/lattice.py @@ -1,8 +1,7 @@ import functools import pathlib import sys -import numpy as np -import pandas as pd + """ Holds information about the lattice, mostly optics @@ -10,114 +9,12 @@ 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: diff --git a/madx/__init__.py b/madx/__init__.py new file mode 100644 index 0000000..caad957 --- /dev/null +++ b/madx/__init__.py @@ -0,0 +1 @@ +from .io import * diff --git a/madx/io.py b/madx/io.py new file mode 100644 index 0000000..6bf1744 --- /dev/null +++ b/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 diff --git a/sixdsim/io.py b/sixdsim/io.py index 673c471..349c16a 100644 --- a/sixdsim/io.py +++ b/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 diff --git a/tbt/constants.py b/tbt/constants.py new file mode 100644 index 0000000..ad7f3b6 --- /dev/null +++ b/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]]) \ No newline at end of file