<commit_before>"""Element family definitions"""

from . import lattice as _lattice
import pyaccel as _pyaccel


def families_dipoles():
    return ['bend',]

def families_quadrupoles():
    return ['qfa', 'qda', 'qfb', 'qdb1', 'qdb2', 'qf1', 'qf2', 'qf3', 'qf4',]

def families_sextupoles():
    return ['sda', 'sfa', 'sdb', 'sfb',
            'sd1j', 'sf1j', 'sd2j', 'sd3j', 'sf2j',
            'sd1k', 'sf1k', 'sd2k', 'sd3k', 'sf2k', ]

def families_horizontal_correctors():
    return ['chf', 'chs',]

def families_vertical_correctors():
    return ['cvf', 'cvs',]

def families_skew_correctors():
    return ['qs',]

def families_rf():
    return ['cav',]

def get_family_data(lattice):
    """Get pyaccel lattice model index and segmentation for each family name

    Keyword argument:
    lattice -- lattice model

    Returns dict.
    """
    latt_dict=_pyaccel.lattice.find_dict(lattice,'fam_name')
    data={}
    for key in latt_dict.keys():
        if key in _family_segmentation.keys():
            data[key]={'index' : latt_dict[key], 'nr_segs' : _family_segmentation[key] , 'families' : key}

    for key in data.keys():
        if data[key]['nr_segs'] != 1:
            new_index=[]
            j=0
            for i in range(len(data[key]['index'])//data[key]['nr_segs']):
                new_index.append(data[key]['index'][j:j+data[key]['nr_segs']])
                j += data[key]['nr_segs']
            data[key]['index']=new_index

    # chs - slow horizontal correctors
    data['chs']={}
    data['chs']['index']=[]
    data['chs']['families'] = ['sfa','sfb','sd1j','sd2j','sd1k','sd2k']
    for family in data['chs']['families']:
        data['chs']['index'] = data['chs']['index'] + data[family]['index']
    data['chs']['index']=sorted(data['chs']['index'])

    # cvs - slow vertical correctors
    data['cvs']={}
    data['cvs']['index']=[]
    data['cvs']['families'] = ['sfa','sfb','sd1j','sd3j','sd1k','sd3k']
    for family in data['cvs']['families']:
        data['cvs']['index'] = data['cvs']['index'] + data[family]['index']<commit_msg>Fix in SI.V10<commit_after>"""Element family definitions"""

from . import lattice as _lattice
import pyaccel as _pyaccel


def families_dipoles():
    return ['bend',]

def families_quadrupoles():
    return ['qfa', 'qda', 'qfb', 'qdb1', 'qdb2', 'qf1', 'qf2', 'qf3', 'qf4',]

def families_sextupoles():
    return ['sda', 'sfa', 'sdb', 'sfb',
            'sd1j', 'sf1j', 'sd2j', 'sd3j', 'sf2j',
            'sd1k', 'sf1k', 'sd2k', 'sd3k', 'sf2k', ]

def families_horizontal_correctors():
    return ['chf', 'chs',]

def families_vertical_correctors():
    return ['cvf', 'cvs',]

def families_skew_correctors():
    return ['qs',]

def families_rf():
    return ['cav',]

def get_family_data(lattice):
    """Get pyaccel lattice model index and segmentation for each family name

    Keyword argument:
    lattice -- lattice model

    Returns dict.
    """
    latt_dict=_pyaccel.lattice.find_dict(lattice,'fam_name')
    data={}
    for key in latt_dict.keys():
        if key in _family_segmentation.keys():
            data[key]={'index' : latt_dict[key], 'nr_segs' : _family_segmentation[key] , 'families' : key}

    for key in data.keys():
        if data[key]['nr_segs'] != 1:
            new_index=[]
            j=0
            for i in range(len(data[key]['index'])//data[key]['nr_segs']):
                new_index.append(data[key]['index'][j:j+data[key]['nr_segs']])
                j += data[key]['nr_segs']
            data[key]['index']=new_index

    # chs - slow horizontal correctors
    data['chs']={}
    data['chs']['index']=[]
    data['chs']['families'] = ['sfa','sfb','sd1j','sf2j','sd1k','sf2k']
    for family in data['chs']['families']:
        data['chs']['index'] = data['chs']['index'] + data[family]['index']
    data['chs']['index']=sorted(data['chs']['index'])

    # cvs - slow vertical correctors
    data['cvs']={}
    data['cvs']['index']=[]
    data['cvs']['families'] = ['sfa','sfb','sd1j','sd3j','sd1k','sd3k']
    for family in data['cvs']['families']:
        data['cvs']['index'] = data['cvs']['index'] + data[family]['index']
{'fraction_non_alphanumeric': 0.21173346209560598, 'fraction_numerical': 0.013037180106228875, 'mean_word_length': 3.5728476821192054, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6312370', 'n_tokens_mistral': 1519, 'n_tokens_neox': 1485, 'n_words': 295}