In [ ]:

    

import numpy as np
import scipy.optimize as scyopt
import matplotlib.pyplot as plt
import cProfile

import pyaccel, sirius
import mathphys as mp
%load_ext autoreload
%autoreload 1
%aimport lattice_errors


    

In [ ]:

    

#import sirius.SI_V10 as v10
#acc      = v10.create_accelerator()
#fam_data = v10.get_family_data(acc)

acc      = sirius.SI_V12.create_accelerator()
fam_data = sirius.SI_V12.get_family_data(acc)


    

In [ ]:

    

um, mrad, percent = 1e-6, 1e-3, 1e-2
config = dict({'mags':dict(),'girder':dict()})
mags = []
dics = dict()
dics['labels'] = ['qfa','qda','qdb1','qdb2','qfb','qf1','qf2','qf3','qf4']
dics['x']     = 40 * um * 1
dics['y']     = 40 * um * 1
dics['roll']  = 0.30 * mrad * 1
dics['excit'] = 0.05 * percent * 1
mags.append(dics)

dics = dict()
dics['labels']   = ['sfa','sda','sf1j','sd1j','sd2j','sd3j','sf2j',
                    'sdb','sfb','sf1k','sd1k','sd2k','sd3k','sf2k']
dics['x']     = 40 * um * 1
dics['y']     = 40 * um * 1
dics['roll']  = 0.30 * mrad * 1
dics['excit'] = 0.05 * percent * 1
mags.append(dics)

dics = dict()
dics['labels']= ['bc_lf','bc_hf']
dics['x']     = 40 * um * 1
dics['y']     = 40 * um * 1
dics['roll']  = 0.30 * mrad * 1
dics['excit'] = 0.05 * percent * 1
dics['k_dip'] = 0.10 * percent * 1
mags.append(dics)

dics = dict()
dics['labels'] = ['b1','b2']
dics['nrsegs'] = [1,1]
dics['x']      = 40 * um * 1
dics['y']      = 40 * um * 1
dics['roll']   = 0.30 * mrad * 1
dics['excit']  = 0.05 * percent * 1
dics['k_dip']  = 0.10 * percent * 1
mags.append(dics)

girder = dict()
girder['x']    = 80 * um * 1
girder['y']    = 80 * um * 1
girder['roll'] = 0.3 * mrad * 1

errors = lattice_errors.generate_errors(acc,mags,girder,fam_data=fam_data,nr_mach=20,cutoff=1,rndtype='gauss',seed=19880419)


    

In [ ]:

    

machine = lattice_errors.apply_erros(acc,errors,increment=1.0)


    

In [ ]:

    

r = lattice_errors.calc_respm_cod(acc,
                                  fam_data['bpm']['index'],
                                  fam_data['ch']['index'],
                                  fam_data['cv']['index'],
                                  symmetry=10)

bba_ind = lattice_errors.get_bba_ind(acc,fam_data['bpm']['index'])

cor_conf = dict({'bpms':fam_data['bpm']['index'],
                 'hcms':fam_data['ch']['index'],
                 'vcms':fam_data['cv']['index'],
                 'bpm_err':dict({'sigma':(20e-6,20e-6),'cutoff':1})})
machine2 = [machine[i][:] for i in range(len(machine))]
gcodx, gcody = lattice_errors.correct_cod(machine2,ind_bba=bba_ind,sext_ramp=[0,1],respm=r,**cor_conf)


    

In [ ]:

    

for acc in machine2: 
    print(pyaccel.optics.calc_emittance_coupling(acc)*100)


    

In [ ]:

    

def fitEllipse(x,y):
    x = x[:,np.newaxis]
    y = y[:,np.newaxis]
    D = np.hstack((x*x, x*y, y*y, x, y, np.ones_like(x)))
    S = np.dot(D.T,D)
    C = np.zeros([6,6])
    C[0,2] = C[2,0] = 2; C[1,1] = -1
    E, V =  np.linalg.eig(np.linalg.solve(S, C))
    n = np.argmax(np.abs(E))
    a = V[:,n]
    return a
def ellipse_center(a):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    num = b*b-a*c
    x0=(c*d-b*f)/num
    y0=(a*f-b*d)/num
    return np.array([x0,y0])
def ellipse_angle_of_rotation( a ):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    return 0.5*np.arctan(2*b/(a-c))
def ellipse_axis_length( a ):
    b,c,d,f,g,a = a[1]/2, a[2], a[3]/2, a[4]/2, a[5], a[0]
    up = 2*(a*f*f+c*d*d+g*b*b-2*b*d*f-a*c*g)
    down1=(b*b-a*c)*( (c-a)*np.sqrt(1+4*b*b/((a-c)*(a-c)))-(c+a))
    down2=(b*b-a*c)*( (a-c)*np.sqrt(1+4*b*b/((a-c)*(a-c)))-(c+a))
    res1=np.sqrt(up/down1)
    res2=np.sqrt(up/down2)
    return np.array([res1, res2])


    

In [ ]:

    

acc = machine2[17]
orb = pyaccel.tracking.findorbit4(acc)
rin = np.array([2e-5+orb[0],0+orb[1],1e-8+orb[2],0+orb[3],0,0])
rout, *_ = pyaccel.tracking.ringpass(acc,rin, nr_turns = 100,
                                    turn_by_turn = 'closed',
                                    element_offset = 0)
r = np.dstack([rin[None,:,None],rout])


    

In [ ]:

    

pars = fitEllipse(r[0][2],r[0][3])
c    = ellipse_center(pars)
ang  = ellipse_angle_of_rotation(pars)
a,b  = ellipse_axis_length(pars)


    

In [ ]:

    

R = np.arange(0,2*np.pi, 0.01)
xx = c[0] + a*np.cos(R)*np.cos(ang) - b*np.sin(R)*np.sin(ang)
yy = c[1] + a*np.cos(R)*np.sin(ang) + b*np.sin(R)*np.cos(ang)


    

In [ ]:

    

plt.plot(r[0][2],r[0][3],'.')
plt.plot(xx,yy, color = 'red')
plt.show()


    

In [ ]:

    

pos = pyaccel.lattice.find_spos(acc)
twiss,*_ = pyaccel.optics.calc_twiss(acc)


    

In [ ]:

    

plt.plot(pos,twiss.etax)
plt.plot(pos,twiss.etay)
plt.show()


    

In [ ]: