def place_quads(pos, ssz):
madx = Madx(stdout=False)
madx.option(echo=False, warn=False, info=False, debug=False, verbose=False)
madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
madx.input('BRHO := BEAM->PC * 3.3356;')
madx.call(file='ps_mu.seq')
madx.call(file='ps_ss_mod.seq')
madx.call(file='ps_50LeQ.str')
madx.call(file='ps_pro_bare_machine.str')
madx.call(file='remove_elements.seq')
madx.input('seqedit, sequence = PS;')
madx.input('select, flag=seqedit, class = MQNAAIAP;')
madx.input('select, flag=seqedit, class = MQNABIAP;')
madx.input('select, flag=seqedit, class = MQSAAIAP;')
madx.input('select, flag=seqedit, class = QNSD;')
madx.input('select, flag=seqedit, class = QNSF;')
madx.input('use, sequence = PS;')
madx.input('seqedit,sequence = PS;flatten;endedit;')
madx.input('seqedit,sequence = PS;remove, element=SELECTED;endedit;')
madx.input('endedit;')
madx.input('use, sequence = PS;')
madx.input('seqedit, sequence = PS;')
for s_idx in pos:
if s_idx == 99:
madx.input('PR.QDN00: MULTIPOLE, KNL:={0,kd};')
madx.input('install, element=PR.QDN00, at=' + str(623.834315) +
';')
elif (s_idx % 2) == 1:
madx.input('PR.QDN%02d: MULTIPOLE, KNL:={0,kd};' % (s_idx + 1))
madx.input('install, element=PR.QDN%02d, at=' % (s_idx + 1) +
str(ssz[s_idx]) + ';')
else:
madx.input('PR.QFN%02d: MULTIPOLE, KNL:={0,kf};' % (s_idx + 1))
madx.input('install, element=PR.QFN%02d, at=' % (s_idx + 1) +
str(ssz[s_idx]) + ';')
madx.input('endedit;')
madx.input('use, sequence=PS;')
madx.input('''
match, sequence=PS;
vary, name= kd, step= 0.00001;
vary, name= kf, step= 0.00001;
global,sequence=PS,Q1= 6.10;
global,sequence=PS,Q2= 6.10;
jacobian, calls = 50000, tolerance=1.0e-15;
endmatch;
''')
madx.twiss()
return madx
def place_quads_wmarkers_FODO(int_steps):
madx = Madx(stdout=False)
madx.option(echo=False, warn=False, info=False, debug=False, verbose=False)
#madx.input('beam, particle=proton, energy=7000;')
madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
madx.input('BRHO := BEAM->PC * 3.3356;')
madx.call(file='fodo_ring.seq')
madx.input('seqedit,sequence = FODO_ring;flatten;endedit;')
madx.input('use, sequence=FODO_ring;')
madx.input('select, flag=makethin, CLASS=SBEND, THICK= false, SLICE =' +
str(int_steps - 2) + ';')
madx.input('makethin, sequence=FODO_ring;')
madx.input('use, sequence=FODO_ring;')
madx.twiss()
posz = [
i for i, elem in enumerate(madx.table.twiss.name)
if elem.startswith('q')
]
return madx, posz
def place_quads_wmarkers(int_steps, pos, s_pos, ssz):
madx = Madx(stdout=False)
madx.option(echo=False, warn=False, info=False, debug=False, verbose=False)
madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
madx.input('BRHO := BEAM->PC * 3.3356;')
madx.call(file='ps_mu.seq')
madx.call(file='ps_ss_mod.seq')
madx.call(file='ps_50LeQ.str')
madx.call(file='ps_pro_bare_machine.str')
madx.call(file='remove_elements.seq')
madx.input('seqedit, sequence = PS;')
madx.input('select, flag=seqedit, class = MQNAAIAP;')
madx.input('select, flag=seqedit, class = MQNABIAP;')
madx.input('select, flag=seqedit, class = MQSAAIAP;')
madx.input('select, flag=seqedit, class = QNSD;')
madx.input('select, flag=seqedit, class = QNSF;')
madx.input('use, sequence = PS;')
madx.input('seqedit,sequence = PS;flatten;endedit;')
madx.input('seqedit,sequence = PS;remove, element=SELECTED;endedit;')
madx.input('endedit;')
madx.input('use, sequence = PS;')
madx.input('seqedit, sequence = PS;')
for i in range(100):
madx.input('MARK%02d: MARKER;' % (i + 1))
madx.input('install, element= MARK%02d, at=' % (i + 1) + str(ssz[i]) +
';')
for i in s_pos:
madx.input('MARK%02d_2: MARKER;' % (i + 1))
madx.input('install, element= MARK%02d_2, at=' % (i + 1) +
str(ssz[i] + 0.01) + ';')
madx.input('endedit;')
madx.input('use, sequence=PS;')
madx.input('select, flag=makethin, CLASS=SBEND, THICK= false, SLICE =' +
str(int_steps) + ';')
madx.input('makethin, sequence=PS;')
madx.input('use, sequence=PS;')
madx.twiss()
posz = [
i for i, elem in enumerate(madx.table.twiss.name)
if elem.startswith('mark') and not (elem.endswith('_2:1'))
]
madx.input('seqedit, sequence = PS;')
for s_idx in pos:
if s_idx == 99:
madx.input('PR.QDN00: MULTIPOLE, KNL:={0,kd};')
madx.input('replace, element=MARK100, by=PR.QDN00;')
elif (s_idx % 2) == 1:
madx.input('PR.QDN%02d: MULTIPOLE, KNL:={0,kd};' % (s_idx + 1))
madx.input('replace, element=MARK%02d, by=PR.QDN%02d;' %
(s_idx + 1, s_idx + 1))
else:
madx.input('PR.QFN%02d: MULTIPOLE, KNL:={0,kf};' % (s_idx + 1))
madx.input('replace, element=MARK%02d, by=PR.QFN%02d;' %
(s_idx + 1, s_idx + 1))
madx.input('endedit;')
madx.input('use, sequence=PS;')
madx.input('''
match, sequence=PS;
vary, name= kd, step= 0.00001;
vary, name= kf, step= 0.00001;
global,sequence=PS,Q1= 6.10;
global,sequence=PS,Q2= 6.10;
jacobian, calls = 50000, tolerance=1.0e-15;
endmatch;
''')
madx.twiss()
return madx, posz
class OptEnv(gym.Env):
def __init__(self, ss, sz, focusing_list, solver, _n_iter, best_approx):
self.dof = len(ss)
self.rew = 2000
self.x_prev = ss
self.sz = sz
self.focusing_list = focusing_list
self.l = len(self.focusing_list)
self.solver = solver
self._n_iter = _n_iter
self.t0 = time.time()
self.timer = [0]
self.values = [1000]
# Spawn MAD-X process
self.reset()
self.best_approx = best_approx
if self.solver == 'ZOOpt':
dim = self.dof
#dim_bounds = self.Bounds_maker()
dimobj = Dimension(dim, [[0, 628.3185]] * dim, [True] * dim)
self.parameter = Parameter(budget=self._n_iter,
init_samples=[ss],
exploration_rate=0.25)
self.step = Objective(self.step, dimobj)
elif self.solver == 'BOBYQA': # currently broken
self.upper = np.multiply(np.ones((self.dof), ), 628.3185)
self.lower = np.multiply(self.upper, 0)
elif self.solver == 'Bayesian': # currently unfinished
dim = self.dof
x = [
'x1', 'x2', 'x3', 'x4', 'x5', 'x6', 'x7', 'x8', 'x9', 'x10',
'x11', 'x12', 'x13', 'x14', 'x15', 'x16', 'x17', 'x18', 'x19',
'x20', 'x21', 'x22', 'x23', 'x24', 'x25', 'x26', 'x27', 'x28',
'x29', 'x30', 'x31', 'x32', 'x33', 'x34', 'x35', 'x36', 'x37',
'x38', 'x39', 'x40'
] #,
#'x41' ,'x42' ,'x43' ,'x44' ,'x45' ,'x46' ,'x47' ,'x48' ,'x49' , 'x50']
bounds = {}
for n in range(self.dof):
bounds[x[n]] = (0, 628.3185)
self.pbounds = bounds
self.optimizer = BayesianOptimization(
f=self.step2,
pbounds=self.pbounds,
random_state=2,
)
else:
self.bounds = [[0, 628.3185]] * self.dof
def step(self, x):
if self.solver == 'ZOOpt':
x = x.get_x()
else:
pass
self.madx.input('EXIT;')
self.reset()
c1 = get_match_and_beta.getBeating(x, self.sz, self.madx,
self.focusing_list,
self.best_approx)
a = (c1.match_and_beat())
output = a[0] * 100
if output < self.rew:
plt.figure(1)
plt.clf()
self.x_best = a[5]
self.rew = output
plt.plot(a[3], a[4])
print('new minimum found, saved to: img/' + self.solver + '.png')
plt.tight_layout()
plt.savefig('img/' + self.solver + '.png')
self.x_prev = a[2]
# If MAD-X has failed re-spawn process
if a[1]:
self.reset()
self.timer.append(time.time() - self.t0)
self.values.append(self.rew)
# Objective function with a = [beam_size_x, beam_size_y, beam_size_z, loss, percentage, error_flag]
#print('\n locations = ', self.x_prev,'\n')
print('output =', output)
# If objective function is best so far, update x_best with new best parameters
return output
def reset(self):
"""
If MAD-X fails, re-spawn process
"""
self.madx = Madx(stdout=False)
self.madx.option(echo=False,
warn=False,
info=False,
debug=False,
verbose=False)
self.madx.input('BEAM, PARTICLE=PROTON, PC = 2.14')
self.madx.input('BRHO := BEAM->PC * 3.3356;')
self.madx.call(file='ps_mu.seq')
self.madx.call(file='ps_ss_mod.seq')
self.madx.call(file='ps_50LeQ.str')
self.madx.call(file='ps_pro_bare_machine.str')
self.madx.call(file='remove_elements.seq')
self.madx.input('seqedit, sequence = PS;')
self.madx.input('select, flag=seqedit, class = MQNAAIAP;')
self.madx.input('select, flag=seqedit, class = MQNABIAP;')
self.madx.input('select, flag=seqedit, class = MQSAAIAP;')
self.madx.input('select, flag=seqedit, class = QNSD;')
self.madx.input('select, flag=seqedit, class = QNSF;')
self.madx.input('use, sequence = PS;')
self.madx.input('seqedit,sequence = PS;flatten;endedit;')
self.madx.input(
'seqedit,sequence = PS;remove, element=SELECTED;endedit;')
self.madx.input('endedit;')
self.madx.input('use, sequence = PS;')
def Bounds_maker(self):
size_section = 628.3185 / self.dof
bounds = [[(i - 3) * size_section, i * size_section]
for i in range(3, 41)]
bounds = np.concatenate(
([[0, size_section * 2]
], bounds, [[size_section * 38, size_section * 40]]), 0)
return bounds
def step2(self, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13,
x14, x15, x16, x17, x18, x19, x20, x21, x22, x23, x24, x25, x26,
x27, x28, x29, x30, x31, x32, x33, x34, x35, x36, x37, x38, x39,
x40): #, x41 ,x42 ,x43 ,x44 ,x45 ,x46 ,x47 ,x48 ,x49 , x50):
x = [
x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15,
x16, x17, x18, x19, x20, x21, x22, x23, x24, x25, x26, x27, x28,
x29, x30, x31, x32, x33, x34, x35, x36, x37, x38, x39, x40
] #, x41 ,x42 ,x43 ,x44 ,x45 ,x46 ,x47 ,x48 ,x49 , x50]
self.madx.input('EXIT;')
self.reset()
c1 = get_match_and_beta.getBeating(x, self.sz, self.madx,
self.focusing_list,
self.best_approx)
a = (c1.match_and_beat())
output = a[0] * 100
if output < self.rew:
plt.clf()
self.x_best = a[2]
self.rew = output
plt.plot(a[3], a[4])
print('new minimum found, saved to: img/' + self.solver + '.png')
plt.tight_layout()
plt.savefig('img/' + self.solver + '.png')
self.x_prev = a[2]
# If MAD-X has failed re-spawn process
if a[1]:
self.reset()
self.timer.append(time.time() - self.t0)
self.values.append(self.rew)
# print('Iteration' + str(len(self.loss_all)))
# print('Sigmax =' + str(a[0]) + ', Sigmay=' + str(a[1]))
# print('Total =' + str(a[0] ** 2 + a[1] ** 2))
# print('Percentage <5 um =' + str(a[4]))
# Objective function with a = [beam_size_x, beam_size_y, beam_size_z, loss, percentage, error_flag]
#print('\n locations = ', self.x_prev,'\n')
print('output =', output)
# If objective function is best so far, update x_best with new best parameters
output = -1 * output
return output
@author: wietse
"""
import numpy as np
import pickle
#import pandas
from zoopt import ExpOpt
from scipy.interpolate import interp1d
from cpymad.madx import Madx
from cl2pd import madx as cl2madx
#%%
madx = Madx()
madx.option(echo=False)
madx.input('BEAM, PARTICLE=PROTON, PC = 2.14;')
madx.input('BRHO := BEAM->PC * 3.3356;')
# call sequence of main units
with open('ps_mu.seq', 'r') as main_unit:
MU_sequence = main_unit.read()
madx.input(MU_sequence)
#call sequence of straight section elements
with open('ps_ss_mod.seq', 'r') as SS:
SS_sequence = SS.read()
madx.input(SS_sequence)
# call general strength file
with open('ps_mod.str', 'r') as main_str:
PS_strength = main_str.read()
madx.input(PS_strength)
# call configuration strength file
