<commit_before>
import numpy as _np
import matplotlib.pyplot as plt
import scipy.stats as _scystat

import pyaccel
import mathphys as _mp

def generate_errors(name,accel,config,fam_data,nr_mach=20,cutoff=1,rndtype='gauss'):
    """Generates random errors to be applied in the model by the function apply_errors.

 INPUTS:
   name       : a string to be used as name to save the configuration data
   accel      : is the accelerator model.
   config     : dictionary with configuration of the errors with keys
                'mags' and/or 'girders'. All of them are optional:
     'mags'   : dictionary with arbitrary named keys whose values are
                dictionaries with keys:
       'labels': list of family names
       'sigma' : dictionary. possible keys:'x','y','roll','excit','k_dip','yaw','pitch'.
                 Values are errors definition (1-sigma for gauss, max for uniform)

     'girder' : dictionary with definition of girder errors. Possible keys:
                'x','y','roll','excit','k_dip','yaw','pitch'.
                Values are errors definition (1-sigma for gauss, max for uniform)
   fam_data : dictionary with families data created by function
              accelerator.get_family_data().
   nr_mach : generate errors for this number of machines.
   cutoff  : number of sigmas to cutoff the distribution (default is infinity)
   rndtype : type of distribution. Possible values: 'sin' and 'gauss' (default is 'gauss')

 OUTPUT:
    errors : dictionary with keys: 'x','y','roll','yaw','pitch','excit','k_dip'.
      Each key is a list with dimension nr_mach x len(acc)
      with errors generated for the elements defined by the inputs. If an
      element errors has contributions from 'mags' and 'girder', the value
      present in this output will be the sum of them.

 EXAMPLES:
 >>> acc = sirius.si.SI_V10.create_accelerator()
 >>> um, mrad, percent = 1e-6, 1e-3, 1e-2
 >>> config = dict({'mags':dict(),'girder':dict()})
 >>> config['mags']['quads'] = dict({'labels':list(),'sigma':dict()})
 >>> config['mags']['quads']['labels'] += ['qfa','qdb2','qfb']
 >>> config['mags']['quads']['sigma']['x']     = 40 * um * 1
 >>> config['mags']['quads']['sigma']['y']     = 40 * um * 1
 >>> config['mags']['quads']['sigma']['roll']  = 0.20 * mrad * 1
 >>> config['mags']['quads']['sigma']['excit'] = 0.05 * percent * 1
 >>> config['mags']['dips'] = dict({'labels':list(),'sigma':dict()})
 >>> config['mags']['dips']['labels'] += ['b1','b2']
 >>> config['mags']['dips']['sigma']['x']     = 40 * um * 1
 >>> config['mags']['dips']['sigma']['y']     = 40 * um * 1
 >>> config['girder']['x']     = 100 * um * 1
 >>> config['girder']['y']     = 100 * um * 1;
 >>> config['girder']['roll']  =0.20 * mrad * 1;
 >>> config['girder']['yaw']   =  20 * mrad * 0;
 >>> config['girder']['pitch'] =  20 * mrad * 0;
 >>> fam_data = sirius.si.SI_V10.get_family_data(acc)
 >>> errors = generate_errors('test',acc,config,fam_data,nr_mach=20,cutoff=2)
"""

    #define the random numbers generator function
    if rndtype.lower().startswith('gauss'):
        random_numbers = _scystat.truncnorm(-cutoff,cutoff).rvs
    else:
        random_numbers = _scystat.uniform(loc=-1,scale=2).rvs

    errors = dict()
    for errtype in ['x','y','roll','yaw','pitch','excit','k_dip']:
        errors[errtype] = _np.zeros((nr_mach, len(acc)))


    _mp.save_pickle(name+'_generate_errors_input',
                    config=config,nr_mach=nr_mach,cutoff=cutoff,rndtype=rndtype)

    if 'mags' in config:
        for mtype in config['mags'].keys():
            for errtype in config['mags'][mtype]['sigma']:
                for fam_name in config['mags'][mtype]['labels']:
                    idx = _np.array(fam_data[fam_name]['index']).T
                    rnd = random_numbers((nr_mach,idx.shape[1])).repeat(idx.shape[0],axis=1)
                    errors[errtype][:,idx.ravel()] += rnd * config['mags'][mtype]['sigma'][errtype]

    if 'girder' in config:
        for errtype in config['girder']:
            for gir_name in fam_data['girder']:
                idx = _np.array(fam_data['girder'][gir_name]).T
                rnd = random_numbers((nr_mach,idx.shape[1])).repeat(idx.shape[0],axis=1)
                errors[errtype][:,idx.ravel()] += rnd * config['girder'][errtype]

    return errors

def apply_erros(name, machine, errors, increment=1.0):
    """ Apply the errors generated by generate_errors to the ring model.

 INPUTS:
   name     : name of the file to save input parameters
   machine  : might be a model of the ring or a list of models of
              the ring
   errors   : structure of errors to be applied (for more details see
              generate_errors help
   increment: float defining the fraction of the errors which will be
              additively applied to the machines.

 OUTPUT:
   machine  : lis of ring models with errors.
"""
    nr_mach = errors['x'].shape[0]

    _mp.save_pickle([name,'_apply_errors_input'],
                    errors=errors, increment=increment)

    if not isinstance(machine,list): machine = nr_mach*[machine]

    if len(machine) != nr_mach:
        print('DifferentSizes: Incompatibility between errors and'+
                ' machine lengths.\n Using minimum of both.')
        nr_mach = min([len(machine),nr_mach])

    return NotImplemented
    
    dim = get_dim(machine{1});

    ids_idx = _pyaccel.findcells(machine{1}, 'PassMethod', 'LNLSThickEPUPass');
    sext_idx = [i for i in range(len(machine[1]))
                            if machine[1][i].polynom_b[2] != 0.0]

    print('    ------------------------------- \n')
    print('   |   codx [mm]   |   cody [mm]   |\n')
    print('   | (max)   (rms) | (max)   (rms) |\n')
    print('---|-------------------------------|\n')
    %fprintf('001| 13.41   14.32 | 13.41   14.32 | 13.41   14.32 | 13.41   14.32 |\n');
    for i=1:nr_machines
        machine{i}    = apply_errors_one_machine(machine{i}, errors, i, increment);
        the_ring = setcellstruct(machine{i}, 'PolynomB', sext_idx, 0, 1, 3);
        the_ring = turn_ids_off(the_ring, ids_idx);
        [codx, cody] = calc_cod(the_ring, dim);
        [x_max_all,x_rms_all] = get_max_rms(codx,1e3);
        [y_max_all,y_rms_all] = get_max_rms(cody,1e3);
        fprintf('%03i| %5.2f   %5.2f | %5.2f   %5.2f | %5.2f   %5.2f | %5.2f   %5.2f |\n',i,x_max_all,x_rms_all,x_max_bpm,x_rms_bpm,y_max_all,y_rms_all,y_max_bpm,y_rms_bpm);
    end
    fprintf('------------------------------------------------------------------- \n');
<commit_msg>continuing implementation of lattice_errors<commit_after>
import numpy as _np
import matplotlib.pyplot as plt
import scipy.stats as _scystat

import pyaccel as _pyaccel
import mathphys as _mp

def generate_errors(name,accel,config,fam_data,nr_mach=20,cutoff=1,rndtype='gauss'):
    """Generates random errors to be applied in the model by the function apply_errors.

 INPUTS:
   name       : a string to be used as name to save the configuration data
   accel      : is the accelerator model.
   config     : dictionary with configuration of the errors with keys
                'mags' and/or 'girders'. All of them are optional:
     'mags'   : dictionary with arbitrary named keys whose values are
                dictionaries with keys:
       'labels': list of family names
       'sigma' : dictionary. possible keys:'x','y','roll','excit','k_dip','yaw','pitch'.
                 Values are errors definition (1-sigma for gauss, max for uniform)

     'girder' : dictionary with definition of girder errors. Possible keys:
                'x','y','roll','excit','k_dip','yaw','pitch'.
                Values are errors definition (1-sigma for gauss, max for uniform)
   fam_data : dictionary with families data created by function
              accelerator.get_family_data().
   nr_mach : generate errors for this number of machines.
   cutoff  : number of sigmas to cutoff the distribution (default is infinity)
   rndtype : type of distribution. Possible values: 'sin' and 'gauss' (default is 'gauss')

 OUTPUT:
    errors : dictionary with keys: 'x','y','roll','yaw','pitch','excit','k_dip'.
      Each key is a list with dimension nr_mach x len(acc)
      with errors generated for the elements defined by the inputs. If an
      element errors has contributions from 'mags' and 'girder', the value
      present in this output will be the sum of them.

 EXAMPLES:
 >>> acc = sirius.si.SI_V10.create_accelerator()
 >>> um, mrad, percent = 1e-6, 1e-3, 1e-2
 >>> config = dict({'mags':dict(),'girder':dict()})
 >>> config['mags']['quads'] = dict({'labels':list(),'sigma':dict()})
 >>> config['mags']['quads']['labels'] += ['qfa','qdb2','qfb']
 >>> config['mags']['quads']['sigma']['x']     = 40 * um * 1
 >>> config['mags']['quads']['sigma']['y']     = 40 * um * 1
 >>> config['mags']['quads']['sigma']['roll']  = 0.20 * mrad * 1
 >>> config['mags']['quads']['sigma']['excit'] = 0.05 * percent * 1
 >>> config['mags']['dips'] = dict({'labels':list(),'sigma':dict()})
 >>> config['mags']['dips']['labels'] += ['b1','b2']
 >>> config['mags']['dips']['sigma']['x']     = 40 * um * 1
 >>> config['mags']['dips']['sigma']['y']     = 40 * um * 1
 >>> config['girder']['x']     = 100 * um * 1
 >>> config['girder']['y']     = 100 * um * 1;
 >>> config['girder']['roll']  =0.20 * mrad * 1;
 >>> config['girder']['yaw']   =  20 * mrad * 0;
 >>> config['girder']['pitch'] =  20 * mrad * 0;
 >>> fam_data = sirius.si.SI_V10.get_family_data(acc)
 >>> errors = generate_errors('test',acc,config,fam_data,nr_mach=20,cutoff=2)
"""

    #define the random numbers generator function
    if rndtype.lower().startswith('gauss'):
        random_numbers = _scystat.truncnorm(-cutoff,cutoff).rvs
    else:
        random_numbers = _scystat.uniform(loc=-1,scale=2).rvs

    errors = dict()
    for errtype in ['x','y','roll','yaw','pitch','excit','k_dip']:
        errors[errtype] = _np.zeros((nr_mach, len(acc)))


    _mp.save_pickle(name+'_generate_errors_input',
                    config=config,nr_mach=nr_mach,cutoff=cutoff,rndtype=rndtype)

    if 'mags' in config:
        for mtype in config['mags'].keys():
            for errtype in config['mags'][mtype]['sigma']:
                for fam_name in config['mags'][mtype]['labels']:
                    idx = _np.array(fam_data[fam_name]['index']).T
                    rnd = random_numbers((nr_mach,idx.shape[1])).repeat(idx.shape[0],axis=1)
                    errors[errtype][:,idx.ravel()] += rnd * config['mags'][mtype]['sigma'][errtype]

    if 'girder' in config:
        for errtype in config['girder']:
            for gir_name in fam_data['girder']:
                idx = _np.array(fam_data['girder'][gir_name]).T
                rnd = random_numbers((nr_mach,idx.shape[1])).repeat(idx.shape[0],axis=1)
                errors[errtype][:,idx.ravel()] += rnd * config['girder'][errtype]

    return errors

def apply_erros(name, machine, errors, increment=1.0):
    """ Apply the errors generated by generate_errors to the ring model.

 INPUTS:
   name     : name of the file to save input parameters
   machine  : might be a model of the ring or a list of models of
              the ring
   errors   : structure of errors to be applied (for more details see
              generate_errors help
   increment: float defining the fraction of the errors which will be
              additively applied to the machines.

 OUTPUT:
   machine  : lis of ring models with errors.
"""
    nr_mach = errors['x'].shape[0]

    _mp.save_pickle([name,'_apply_errors_input'],
                    errors=errors, increment=increment)

    machs = []
    if not isinstance(machine,list):
        for i in range(nr_mach):
            machs.append(_pyaccel.accelerator.Accelerator(machine))
    machine = machs

    if len(machine) != nr_mach:
        print('DifferentSizes: Incompatibility between errors and'+
                ' machine lengths.\n Using minimum of both.')
        nr_mach = min([len(machine),nr_mach])

    return None
    dim = get_dim(machine{1});

    ids_idx  = [i for in range(len(machine[0]))
                  if machine[0][i].pass_method.startswith('kicktable_pass')]
    sext_idx = [i for i in range(len(machine[0]))
                  if machine[0][i].polynom_b[2] != 0.0]

    print('    ------------------------------- ')
    print('   |   codx [mm]   |   cody [mm]   |')
    print('   | (max)   (rms) | (max)   (rms) |')
    print('---|-------------------------------|')
   #print('001| 13.41   14.32 | 13.41   14.32 |');
    return None
    for i in range(nr_mach):
        apply_errors_one_machine(machine[i], errors, i, increment)
        ring = machine[i][:]
        _pyccel.lattice.set_attribute(ring, 'polynom_b',sext_idx, 0.0, 1, 3)
        ring = turn_ids_off(ring, ids_idx)
        codx, cody = calc_cod(ring, dim)
        x_max_all, x_rms_all = get_max_rms(codx,1e3)
        y_max_all, y_rms_all = get_max_rms(cody,1e3)
        print('{0:03d}| {1:5.2f}   {2:5.2f} | {3:5.2f}   {4:5.2f} |'.format(
                                    i,x_max_all,x_rms_all,y_max_all,y_rms_all))
    print(36*'-')


    def apply_errors_one_machine(the_ring, errors, ii, fraction):
        funs={'x':_pyaccel.lattice.add_error_misalignment_x,
              'y':_pyaccel.lattice.add_error_misalignment_y,
              'roll':_pyaccel.lattice.add_error_rotation_roll,
              'yaw':_pyaccel.lattice.add_error_rotation_yaw,
              'pitch':_pyaccel.lattice.add_error_rotation_pitch,
              'excit':_pyaccel.lattice.add_error_excitation_main,
              'k_dip':_pyaccel.lattice.add_error_excitation_kdip}

        for errtype in ['x','y','roll','yaw','pitch','excit','k_dip']:
            errors[errtype] = _np.zeros((nr_mach, len(acc)))
            err = fraction * errors[errtype][ii,:]
            idx = err.nonzero()
            funs[errtype](the_ring, idx, err[idx])

{'fraction_non_alphanumeric': 0.17005625270445696, 'fraction_numerical': 0.01875090148564835, 'mean_word_length': 2.8745459625593743, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 123, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '950029', 'n_tokens_mistral': 4828, 'n_tokens_neox': 4584, 'n_words': 1089}