def bump_4cors(lat: MagneticLattice, cor_list: List[Element], marker: Union[Marker, Monitor],
x: float = 0.0001, xp: float = 0., energy: float = 14.) -> np.array:
"""
Bump with 4 correctors.
Function calculates correctors angles (strength) to make closed bump with (x, x') or (y, y') at the marker position.
Calculated angles are applied to corresponding correctors (MagneticLattice is mutated).
:param lat: MagneticLattice
:param cor_list: list of 4 correctors, All for correctors must be in the same plane
:param marker: Marker or Monitor. It must be between 2d and 3d corrector
:param x: x or y beam coordinate at Marker position [m]
:param xp: x' or y' beam coordinate at Marker position [rad]
:param energy: Beam energy [GeV]
:return: numpy array of corrector angles in [rad]
"""
# lattice analysis, calculate R matrices between different devices
if len(cor_list) != 4:
raise TypeError("It must be 4 correctors")
if len([cor for cor in cor_list if isinstance(cor, Hcor)]) == 4:
x_idx, xp_idx = 0, 1
elif len([cor for cor in cor_list if isinstance(cor, Vcor)]) == 4:
x_idx, xp_idx = 2, 3
else:
raise TypeError("Not all correctors belong to Y- or X-plane")
Rs = []
sorted_list = [cor_list[0], cor_list[1], marker, cor_list[2], cor_list[3]]
for i in range(4):
sequence = copy.deepcopy(lat.get_sequence_part(start=sorted_list[i], stop=sorted_list[i + 1]))
sequence[0].l /= 2.
sequence[-1].l /= 2.
lat1 = MagneticLattice(sequence)
R1 = lattice_transfer_map(lat1, energy)
Rs.append(R1)
R21 = np.dot(Rs[1], Rs[0])
M = np.array([[R21[x_idx, xp_idx], Rs[1][x_idx, xp_idx]],
[R21[xp_idx, xp_idx], Rs[1][xp_idx, xp_idx]]])
Minv = np.linalg.inv(M)
X = np.array([x, xp])
cor_strengths = np.dot(Minv, X)
angle_1, angle_2 = cor_strengths
R43 = np.dot(Rs[3], Rs[2])
Xf = np.array([0., 0.])
M43 = np.array([[R43[x_idx, x_idx], R43[x_idx, xp_idx]],
[R43[xp_idx, x_idx], R43[xp_idx, xp_idx]]])
X4 = np.dot(M43, X)
angle_3 = (Xf[0] - X4[0]) / Rs[3][x_idx, xp_idx]
angle_4 = -(X4[1] + Rs[3][xp_idx, xp_idx] * angle_3 - Xf[1])
a = np.array([angle_1, angle_2, angle_3, angle_4])
for i, cor in enumerate(cor_list):
cor.angle = a[i]
lat.update_transfer_maps()
return a
def closed_orbit(lattice, eps_xy=1.e-7, eps_angle=1.e-7, energy=0):
__author__ = 'Sergey Tomin'
"""
Searching of initial coordinates (p0) by iteration method.
For initial conditions p uses exact solution of equation p = M*p + B
:param lattice: class MagneticLattice
:param eps_xy: tolerance on coordinates of beam in the start and end of lattice
:param eps_angle: tolerance on the angles of beam in the start and end of lattice
:return: class Particle
"""
R = lattice_transfer_map(lattice, energy)
smult = SecondOrderMult()
ME = np.eye(4) - R[:4, :4]
P = np.dot(np.linalg.inv(ME), lattice.B[:4])
def errf(x):
X = np.array([[x[0]], [x[1]], [x[2]], [x[3]], [0], [0]])
smult.numpy_apply(X, R, lattice.T)
X += lattice.B
err = np.sum([1000 * (X[i, 0] - x[i])**2 for i in range(4)])
return err
res = fmin(errf, P, xtol=1e-8, maxiter=2e3, maxfun=2.e3)
return Particle(x=res[0], px=res[1], y=res[2], py=res[3])
def merger(lat, remaining_types=[], remaining_elems=[], init_energy=0.):
"""
Function to compress the lattice excluding elements by type or by individual elements
:param lat: MagneticLattice
:param remaining_types: list, the type of the elements which needed to be untouched
others will be "compress" to Matrix element
e.g. [Monitor, Quadrupole, Bend, Hcor]
:param remaining_elems: list of elements (ids (str) or object)
:param init_energy: initial energy
:return: New MagneticLattice
"""
_logger.debug("element numbers before: " + str(len(lat.sequence)))
lattice_analysis = []
merged_elems = []
for elem in lat.sequence:
if elem.__class__ in remaining_types or elem.id in remaining_elems or elem in remaining_elems:
lattice_analysis.append(merged_elems)
merged_elems = []
lattice_analysis.append([elem])
elif elem.__class__ == Edge and ((Bend in remaining_types) or
(SBend in remaining_types) or
(RBend in remaining_types)):
continue
else:
merged_elems.append(elem)
if len(merged_elems) != 0:
lattice_analysis.append(merged_elems)
seq = []
E = init_energy
for elem_list in lattice_analysis:
if len(elem_list) == 1:
E += elem_list[0].transfer_map.delta_e
seq.append(elem_list[0])
elif len(elem_list) == 0:
continue
else:
delta_e = np.sum([elem.transfer_map.delta_e for elem in elem_list])
lattice = MagneticLattice(elem_list, method=lat.method)
R = lattice_transfer_map(lattice, energy=E)
m = Matrix()
m.r = lattice.R
m.t = lattice.T
m.b = lattice.B
m.l = lattice.totalLen
m.delta_e = delta_e
E += delta_e
seq.append(m)
new_lat = MagneticLattice(seq, method=lat.method)
_logger.debug("element numbers after: " + str(len(new_lat.sequence)))
return new_lat
def error_func(x):
for i in range(len(varz)):
if varz[i].__class__ == Quadrupole:
varz[i].k1 = x[i]
varz[i].transfer_map = lat.method.create_tm(
varz[i]) # create_transfer_map(varz[i])
R = lattice_transfer_map(lat, beam.E)[0:2, 0:2]
# print
# R
err = np.linalg.norm(np.abs(R - target_matrix)**2)
# print
# 'iteration error: ', err
return err
def errf(x):
tw_loc = deepcopy(tw)
tw0 = deepcopy(tw)
'''
parameter to be varied is determined by variable class
'''
for i in range(len(vars)):
if isinstance(vars[i], Drift):
if x[i] < 0:
# print('negative length in match')
return weights('negative_length')
vars[i].l = x[i]
vars[i].transfer_map = lat.method.create_tm(vars[i])
if isinstance(vars[i], Quadrupole):
vars[i].k1 = x[i]
vars[i].transfer_map = lat.method.create_tm(vars[i])
if isinstance(vars[i], Solenoid):
vars[i].k = x[i]
vars[i].transfer_map = lat.method.create_tm(vars[i])
if isinstance(vars[i], (RBend, SBend, Bend)):
if vary_bend_angle:
vars[i].angle = x[i]
else:
vars[i].k1 = x[i]
vars[i].transfer_map = lat.method.create_tm(vars[i])
if isinstance(vars[i], list):
if isinstance(vars[i][0], Twiss) and isinstance(
vars[i][1], str):
k = vars[i][1]
tw_loc.__dict__[k] = x[i]
if isinstance(
vars[i], tuple
): # all quads strength in tuple varied simultaneously
for v in vars[i]:
v.k1 = x[i]
v.transfer_map = lat.method.create_tm(v)
err = 0.0
if "periodic" in constr.keys():
if constr["periodic"]:
tw_loc = periodic_twiss(tw_loc,
lattice_transfer_map(lat, tw.E))
tw0 = deepcopy(tw_loc)
if tw_loc is None:
print("########")
return weights('periodic')
# save reference points where equality is asked
ref_hsh = {} # penalties on two-point inequalities
for e in constr.keys():
if e == 'periodic': continue
if e == 'total_len': continue
for k in constr[e].keys():
if isinstance(constr[e][k], list):
if constr[e][k][0] == '->':
# print 'creating reference to', constr[e][k][1].id
ref_hsh[constr[e][k][1]] = {k: 0.0}
# evaluating global and point penalties
tw_loc.s = 0
for e in lat.sequence:
tw_loc = e.transfer_map * tw_loc
if 'global' in constr.keys():
for c in constr['global'].keys():
if isinstance(constr['global'][c], list):
v1 = constr['global'][c][1]
if constr['global'][c][0] == '<':
if tw_loc.__dict__[c] > v1:
err = err + weights(k) * (tw_loc.__dict__[c] -
v1)**2
if constr['global'][c][0] == '>':
if tw_loc.__dict__[c] < v1:
err = err + weights(k) * (tw_loc.__dict__[c] -
v1)**2
if 'delta' in constr.keys():
if e in constr['delta'].keys():
tw_k = constr['delta'][e][0]
constr['delta'][e][1] = tw_loc.__dict__[tw_k]
if e in ref_hsh.keys():
ref_hsh[e] = deepcopy(tw_loc)
if e in constr.keys():
for k in constr[e].keys():
if isinstance(constr[e][k], list):
v1 = constr[e][k][1]
if constr[e][k][0] == '<':
if tw_loc.__dict__[k] > v1:
err = err + weights(k) * (tw_loc.__dict__[k] -
v1)**2
if constr[e][k][0] == '>':
if tw_loc.__dict__[k] < v1:
err = err + weights(k) * (tw_loc.__dict__[k] -
v1)**2
if constr[e][k][0] == 'a<':
if np.abs(tw_loc.__dict__[k]) > v1:
err = err + weights(k) * (tw_loc.__dict__[k] -
v1)**2
if constr[e][k][0] == 'a>':
if np.abs(tw_loc.__dict__[k]) < v1:
err = err + weights(k) * (tw_loc.__dict__[k] -
v1)**2
if constr[e][k][0] == '->':
try:
if len(constr[e][k]) > 2:
dv1 = float(constr[e][k][2])
else:
dv1 = 0.0
err += (tw_loc.__dict__[k] -
(ref_hsh[v1].__dict__[k] + dv1))**2
if tw_loc.__dict__[k] < v1:
err = err + (tw_loc.__dict__[k] - v1)**2
except:
print(
'constraint error: rval should precede lval in lattice'
)
if tw_loc.__dict__[k] < 0:
err += (tw_loc.__dict__[k] - v1)**2
elif isinstance(constr[e][k], str):
pass
else:
err = err + weights(k) * (constr[e][k] -
tw_loc.__dict__[k])**2
if "total_len" in constr.keys():
total_len = constr["periodic"]
err = err + weights('total_len') * (tw_loc.s - total_len)**2
if 'delta' in constr.keys():
delta_dict = constr['delta']
elems = []
for e in delta_dict.keys():
if isinstance(e, Element):
elems.append(e)
delta_err = delta_dict["weight"] * (delta_dict[elems[0]][1] -
delta_dict[elems[1]][1] -
delta_dict["val"])**2
err = err + delta_err
if min_i5:
''' evaluating integral parameters
'''
I1, I2, I3, I4, I5 = radiation_integrals(lat, tw0, nsuperperiod=1)
err += I5 * weights('i5')
Je = 2 + I4 / I2
Jx = 1 - I4 / I2
Jy = 1
if Je < 0 or Jx < 0 or Jy < 0: err = 100000.0
# c1, c2 = natural_chromaticity(lat, tw0)
# err += ( c1**2 + c2**2) * 1.e-6
if verbose:
print('iteration error:', err)
return err
def errf(x):
p_array0 = deepcopy(p_array)
tws = get_envelope(p_array0, bounds=bounds)
tw_loc = deepcopy(tws)
tw0 = deepcopy(tws)
'''
parameter to be varied is determined by variable class
'''
for i in range(len(vars)):
if vars[i].__class__ == Drift:
if x[i] < 0:
# print('negative length in match')
return weights('negative_length')
pass
vars[i].l = x[i]
vars[i].transfer_map = lat.method.create_tm(vars[i])
if vars[i].__class__ == Quadrupole:
vars[i].k1 = x[i]
vars[i].transfer_map = lat.method.create_tm(vars[i])
if vars[i].__class__ == Solenoid:
vars[i].k = x[i]
vars[i].transfer_map = lat.method.create_tm(vars[i])
if vars[i].__class__ in [RBend, SBend, Bend]:
if vary_bend_angle:
vars[i].angle = x[i]
else:
vars[i].k1 = x[i]
vars[i].transfer_map = lat.method.create_tm(vars[i])
if vars[i].__class__ == list:
if vars[i][0].__class__ == Twiss and vars[i][
1].__class__ == str:
k = vars[i][1]
tw_loc.__dict__[k] = x[i]
if vars[i].__class__ == tuple: # all quads strength in tuple varied simultaneously
for v in vars[i]:
v.k1 = x[i]
v.transfer_map = lat.method.create_tm(v)
err = 0.0
if "periodic" in constr.keys():
if constr["periodic"] is True:
tw_loc = periodic_twiss(tw_loc,
lattice_transfer_map(lat, tw_loc.E))
tw0 = deepcopy(tw_loc)
if tw_loc is None:
print("########")
return weights('periodic')
# save reference points where equality is asked
ref_hsh = {} # penalties on two-point inequalities
for e in constr.keys():
if e == 'periodic': continue
if e == 'total_len': continue
for k in constr[e].keys():
if constr[e][k].__class__ == list:
if constr[e][k][0] == '->':
# print 'creating reference to', constr[e][k][1].id
ref_hsh[constr[e][k][1]] = {k: 0.0}
# print 'references:', ref_hsh.keys()
# evaluating global and point penalties
# tw_loc.s = 0
# print("start = ", get_envelope(p_array0))
navi.go_to_start()
tws_list, p_array0 = track(lat,
p_array0,
navi,
print_progress=False,
bounds=bounds)
s = np.array([tw.s for tw in tws_list])
# print("stop = ", tws_list[-1])
L = 0.
for e in lat.sequence:
indx = (np.abs(s - L)).argmin()
L += e.l
tw_loc = tws_list[indx]
if 'global' in constr.keys():
# print 'there is a global constraint', constr['global'].keys()
for c in constr['global'].keys():
if constr['global'][c].__class__ == list:
# print 'list'
v1 = constr['global'][c][1]
if constr['global'][c][0] == '<':
if tw_loc.__dict__[c] > v1:
err = err + weights(k) * (tw_loc.__dict__[c] -
v1)**2
if constr['global'][c][0] == '>':
# print '> constr'
if tw_loc.__dict__[c] < v1:
err = err + weights(k) * (tw_loc.__dict__[c] -
v1)**2
if e in ref_hsh.keys():
# print 'saving twiss for', e.id
ref_hsh[e] = deepcopy(tw_loc)
if e in constr.keys():
for k in constr[e].keys():
# print(k)
if constr[e][k].__class__ == list:
v1 = constr[e][k][1]
if constr[e][k][0] == '<':
if tw_loc.__dict__[k] > v1:
err = err + weights(k) * (tw_loc.__dict__[k] -
v1)**2
if constr[e][k][0] == '>':
if tw_loc.__dict__[k] < v1:
err = err + weights(k) * (tw_loc.__dict__[k] -
v1)**2
if constr[e][k][0] == 'a<':
if np.abs(tw_loc.__dict__[k]) > v1:
err = err + weights(k) * (tw_loc.__dict__[k] -
v1)**2
if constr[e][k][0] == 'a>':
if np.abs(tw_loc.__dict__[k]) < v1:
err = err + weights(k) * (tw_loc.__dict__[k] -
v1)**2
if constr[e][k][0] == '->':
try:
# print 'huh', k, e.id, float(constr[e][k][2])
if len(constr[e][k]) > 2:
dv1 = float(constr[e][k][2])
else:
dv1 = 0.0
# print 'weiter'
err += (tw_loc.__dict__[k] -
(ref_hsh[v1].__dict__[k] + dv1))**2
if tw_loc.__dict__[k] < v1:
err = err + (tw_loc.__dict__[k] - v1)**2
except:
print(
'constraint error: rval should precede lval in lattice'
)
if tw_loc.__dict__[k] < 0:
# print 'negative constr (???)'
err += (tw_loc.__dict__[k] - v1)**2
else:
# print "safaf", constr[e][k] , tw_loc.__dict__[k], k, e.id, x
err = err + weights(k) * (constr[e][k] -
tw_loc.__dict__[k])**2
# print err
for v in vars:
print(v.id, v.k1)
if "total_len" in constr.keys():
total_len = constr["periodic"]
err = err + weights('total_len') * (tw_loc.s - total_len)**2
if min_i5:
''' evaluating integral parameters
'''
I1, I2, I3, I4, I5 = radiation_integrals(lat, tw0, nsuperperiod=1)
err += I5 * weights('i5')
Je = 2 + I4 / I2
Jx = 1 - I4 / I2
Jy = 1
if Je < 0 or Jx < 0 or Jy < 0: err = 100000.0
# c1, c2 = natural_chromaticity(lat, tw0)
# err += ( c1**2 + c2**2) * 1.e-6
if verbose:
print('iteration error:', err)
return err