analyze_e1.py

# -*- coding: utf-8 -*-
"""
Automatically generated by Colaboratory.


## Analyze_e1

Example for the MAXM05 course

(Sverker Werin)
Adding a test-line in gitHub to test how to pull it into the main.
"""

#!pip install accelerator-toolbox # only for Google Colab

import at
import numpy as np
import matplotlib.pyplot as plt

"""Define the lattice elements."""

# Straight Section definitions
D      =    at.elements.Drift('D', 1.5  );

# Quadrupoles 
QF     =    at.elements.Quadrupole('QF'  , 0.2, 0.30000);

"""Define the lattice."""

L = [ D ]
L1 = 5*L
lattice =  [D,D,D, QF]+ L1 
print(lattice)

"""Define a particle X [x, xp, y ,yp, dp, s]"""

X0=np.zeros((6,1))
X0[0]=0.001 # x
X0[1]=0.0 # xp

"""Track through the lattice"""

# Find reference points for output. One for each element.
length = np.size(lattice)
refpts = np.r_[0:length + 1]

# Track the particle in X0 through the lattice.
X_out = at.lattice_pass(lattice, X0, nturns=1,refpts=refpts)

"""Get the longitudinal position of each element exit."""

s = at.lattice.get_s_pos(lattice)

"""Plot the particle position along the lattice"""

plt.plot(s,X_out[0,0,:,0])
axes = plt.gca()
axes.set_xlabel('s (m)')
axes.set_ylabel('x (m)')
axes.set_title('Particle position in x')
plt.show()

"""Plot the angle of particle path along lattice."""

plt.plot(s,X_out[1,0,:,0])
axes = plt.gca()
axes.set_xlabel('s (m)')
axes.set_ylabel('x p')
axes.set_title('Angle of particle path')
plt.show()

"""Calculate transfer matrices of different elements"""

element=0
m66=at.find_elem_m66(lattice[element])

print('Transfer matrix of element',element,' = ')
print (m66)

element=3
m66=at.find_elem_m66(lattice[element])


print('Transfer matrix of element',element,' = ')
print (m66)

"""Calculate the transfer matrix of the complete lattice"""

a=at.find_m44(lattice,0)
print(a[0])