lib_tw.py

"""
    Classes and functions for TraceWin

    - Note for clean-up

      * Old LATTICE and MATCH classes can be deleted once trajectory
        correction is established with new classes.
"""

from __future__ import print_function

# ---- Lib

import numpy

from struct import pack, unpack
from itertools import chain

from subprocess import check_output
import os
import sys

from . import lib_tw_elem


# -------- Classes

# ---- Lattice and project


class LATTICE:
    """
    """

    def __init__(self, file_name_lat, file_name_fmap=[], freq=352.21, gamma=1.0):
        """
        :param file_name_lat: name of lattice file
        :type file_name_lat: str
        :param file_name_fmap: list of field map file(-s)
        :type file_name_fmap: list or str
        :param freq: RF frequency
        :type freq: float
        :param gamma: relativistic gamma
        :type gamma: float
        """

        # In case file_name_fmap is str
        if isinstance(file_name_fmap, str):
            file_name_fmap = [file_name_fmap]

        # Elem/comm class dict
        dic_cls = {
            "DRIFT": lib_tw_elem.DRIFT,
            "QUAD": lib_tw_elem.QUAD,
            "THIN_STEERING": lib_tw_elem.THIN_STEERING,
            "GAP": lib_tw_elem.GAP,
            "DTL_CEL": lib_tw_elem.DTL_CEL,
            #
            "BEND": lib_tw_elem.BEND,
            "EDGE": lib_tw_elem.EDGE,
            "APERTURE": lib_tw_elem.APERTURE,
            "DIAG_POSITION": lib_tw_elem.DIAG_POSITION,
            #
            "STEERER": lib_tw_elem.STEERER,
            "CHOPPER": lib_tw_elem.CHOPPER,
            #
            "ADJUST": lib_tw_elem.ADJUST,
            "FREQ": lib_tw_elem.FREQ,
            "MARKER": lib_tw_elem.MARKER,
            #
            "ERROR_BEAM_STAT": lib_tw_elem.ERROR_BEAM_STAT,
            "ERROR_BEAM_DYN": lib_tw_elem.ERROR_BEAM_DYN,
            "ERROR_QUAD_NCPL_STAT": lib_tw_elem.ERROR_QUAD_NCPL_STAT,
            "ERROR_QUAD_CPL_STAT": lib_tw_elem.ERROR_QUAD_CPL_STAT,
            "ERROR_CAV_NCPL_STAT": lib_tw_elem.ERROR_CAV_NCPL_STAT,
            "ERROR_CAV_NCPL_DYN": lib_tw_elem.ERROR_CAV_NCPL_DYN,
            "ERROR_CAV_CPL_STAT": lib_tw_elem.ERROR_CAV_CPL_STAT,
            "ERROR_CAV_CPL_DYN": lib_tw_elem.ERROR_CAV_CPL_DYN,
            "ERROR_STAT_FILE": lib_tw_elem.ERROR_STAT_FILE,
            "ERROR_DYN_FILE": lib_tw_elem.ERROR_DYN_FILE,
        }

        # Field map dict
        dic_fmap = {}
        for file_name_fmap_i in file_name_fmap:
            name_fmap = ".".join(
                file_name_fmap_i.split("/")[-1].split(".")[:-1]
            )  # Remove / and extension
            dic_fmap[name_fmap] = lib_tw_elem.FIELD_MAP_DATA(file_name_fmap_i)
        # In case the lattice file is in a different folder:
        basefolder = os.path.dirname(file_name_lat)
        if basefolder:
            _update_field_map_dict(dic_fmap, basefolder)

        # Go through the lat file
        with open(file_name_lat) as file:
            lst = []
            for lin in file:
                lin = lin.partition(";")[0]  # Remove comment
                if lin.split() != []:  # Remove empty line
                    # Split a line
                    if ":" in lin:
                        name = lin.partition(":")[0].split()[0]
                        typ = lin.partition(":")[2].split()[0].upper()
                        para = lin.partition(":")[2].split()[1:]
                    else:
                        name = ""
                        typ = lin.split()[0].upper()
                        para = lin.split()[1:]
                    # Map to a class
                    if typ == "FIELD_MAP":
                        lst.append(lib_tw_elem.FIELD_MAP(name, typ, para, dic_fmap))
                    elif typ in dic_cls.keys():
                        lst.append(dic_cls[typ](name, typ, para))
                    elif "DIAG" in typ:
                        lst.append(lib_tw_elem.DIAG(name, typ, para))
                    else:
                        lst.append(lib_tw_elem.COMM(name, typ, para))

                    # in case of field map path, update dictionary with new path
                    if typ == "FIELD_MAP_PATH":
                        _update_field_map_dict(dic_fmap, para[0])

                    # Break the loop
                    if typ == "END":
                        break

        # Instances
        self.gamma = gamma
        self.freq = freq
        self.lst = lst
        self.fmap = dic_fmap  # Maybe not needed ...

        # Assign idx, idx_elem, s, freq, apt
        self.update_idx()

    def get_correctors_idx(self, i):
        """
        Get correctors associated to corrector index i
        """
        import logging

        found = False
        correctors = []
        for element in self.lst:
            if found:
                # WARNING I worry that there could be an inactive comment/element between the ADJUST and actual corrector
                logging.debug(
                    "Found element {} for corrector family {}".format(element.typ, i)
                )
                correctors.append(element)
                found = False
            if isinstance(element, lib_tw_elem.COMM):
                if element.typ in ["ADJUST", "ADJUST_STEERER"]:
                    if int(element.para[0]) == i:  # next element is the corrector
                        found = True
        return correctors

    def get_elem_idx(self, i):
        """
        Get a TraceWin index number

        Note: We start counting from 0, TW starts from 1
        """
        for element in self.lst:
            if element.idx_elem == i - 1:
                return element

    def get_steerer_for(self, idx_elem):
        """
        Returns the steerer object for an element (e.g. quad)
        """
        previous = None
        for element in self.lst:
            if element.idx_elem + 1 == idx_elem:
                if previous.typ == "STEERER":
                    return previous
                return None
            previous = element

    def update_idx(self):

        # Assign/update idx, idx_elem, s, freq
        for i in range(len(self.lst)):
            if i == 0:
                self.lst[0].idx = -1
                self.lst[0].idx_elem = -1
                self.lst[0].s = 0.0
                self.lst[0].freq = self.freq
            if i != 0:
                self.lst[i].idx = self.lst[i - 1].idx
                self.lst[i].idx_elem = self.lst[i - 1].idx_elem
                self.lst[i].s = self.lst[i - 1].s
                self.lst[i].freq = self.lst[i - 1].freq
            self.lst[i].update_idx()

        # Assign apt (using apt of the previous elem for diag elem)
        for i in range(len(self.lst)):
            try:
                if self.lst[i].apt is None: