def __init__(self,
args,
tempdir,
scaling=4,
overwrite=True,
verbose=False,
summary=False):
self.cons = constraintsClass()
self.beam = rbf.beam()
self.twiss = rtf.twiss()
self.tmpdir = tempdir
self.verbose = verbose
self.summary = summary
self.parameters = list(args)
self.dirname = os.path.basename(self.tmpdir)
self.framework = Framework(self.dirname, clean=False, verbose=verbose)
self.framework.loadSettings('Lattices/clara400_v12_v3.def')
self.framework.change_Lattice_Code('All', 'elegant')
if not os.name == 'nt':
self.framework.defineASTRACommand([
'mpiexec', '-np',
str(3 * scaling), '/opt/ASTRA/astra_MPICH2.sh'
])
self.framework.defineGeneratorCommand(['/opt/ASTRA/generator.sh'])
self.framework.defineCSRTrackCommand([
'/opt/OpenMPI-1.4.3/bin/mpiexec', '-n',
str(3 * scaling), '/opt/CSRTrack/csrtrack_openmpi.sh'
])
self.framework.generator.number_of_particles = 2**(3 * scaling)
else:
self.framework.generator.number_of_particles = 2**(3 * 3)
self.framework.defineElegantCommand(['elegant'])
self.framework.setElementType('quadrupole', 'k1l', self.parameters)
def create_base_files(scaling):
framework = Framework('basefiles_' + str(scaling), overwrite=True)
framework.loadSettings('Lattices/clara400_v12_v3.def')
if not os.name == 'nt':
framework.defineASTRACommand(scaling=scaling)
framework.defineCSRTrackCommand(scaling=scaling)
framework.generator.number_of_particles = 2**(3 * scaling)
# framework.change_Lattice_Code('VBC', 'ASTRA')
framework.track(endfile="S07") #startfile='VBC')
def create_base_files(scaling, changes=None):
framework = Framework('basefiles_' + str(scaling), overwrite=True)
framework.loadSettings('FEBE_Single_L01.def')
if not os.name == 'nt':
framework.defineASTRACommand(scaling=scaling)
framework.generator.number_of_particles = 2**(3 * scaling)
if not changes is None:
framework.load_changes_file(changes)
framework.track()
def example_1():
# Define a new framework instance, in directory 'C2V'.
# "clean" will empty (delete everything!) the directory if true
# "verbose" will print a progressbar is true
lattice = fw.Framework('GPT_Test', clean=False, verbose=False)
# Load a lattice definition. By default the frameowrk looks in "OnlineModel/MasterLattice/" to find things
# This example loads a lattice with a CLARA 400Hz gun, and tracking to VELA BA1
lattice.loadSettings(
'./CLARA_FE/CLARA_BA1_Gaussian/CLA10-BA1_TOMP_ASTRA.def')
# This is a scaling parameter
scaling = 5
# This defines the location of tracking codes - On windows this uses versions in "OnlineModel/MasterLattice/Codes",
# but you will need to define these yourself on linux
lattice.defineASTRACommand(location=['/opt/ASTRA/astra_MPICH2.sh'])
lattice.defineGeneratorCommand(location=['/opt/ASTRA/generator.sh'])
lattice.defineCSRTrackCommand(location=['/opt/CSRTrack/csrtrack.sh'])
# This defines the number of particles to create at the gun (this is "ASTRA generator" which creates distributions)
lattice.generator.number_of_particles = 2**(3 * scaling)
# This tracks the beam based on the definitions in the lattice file loaded using "lattice.loadSettings"
lattice.change_Lattice_Code('All', 'gpt')
lattice['S02'].prefix = '../CLARA_FE/CLARA_FE_Gaussian/TOMP_SETUP_-10/'
lattice['S02'].sample_interval = 2**(3 * 1)
lattice.track(files=['S02', 'C2V'],
preprocess=True,
write=True,
track=True,
postprocess=True)
def __init__(self, lattice='Lattices/clara400_v12_v3.def', scaling=5):
super(Optimise_transverse, self).__init__()
self.lattice_file = lattice
self.framework = fw.Framework(None)
self.framework.loadSettings(self.lattice_file)
self.parameter_names = [q for q in self.framework.getElementType('quadrupole','objectname')]
self.parameters = [[q, 'k1l'] for q in self.framework.getElementType('quadrupole','objectname')]
self.cons = constraintsClass()
self.changes = None
self.resultsDict = {}
# ******************************************************************************
CLARA_dir = os.path.relpath(__file__+'/../../')
self.POST_INJECTOR = True
CREATE_BASE_FILES = False
if self.POST_INJECTOR and CREATE_BASE_FILES:
self.optfunc = partial(self.OptimisingFunction, scaling=scaling, post_injector=self.POST_INJECTOR, basefiles='../CLARA/basefiles_'+str(scaling)+'/', CLARA_dir=CLARA_dir, lattice=self.lattice_file)
else:
self.optfunc = partial(self.OptimisingFunction, scaling=scaling, post_injector=self.POST_INJECTOR, CLARA_dir=CLARA_dir, lattice=self.lattice_file)
# ******************************************************************************
if not self.POST_INJECTOR:
best = injector_startingvalues + best
elif CREATE_BASE_FILES:
for i in [scaling]:
pass
# optfunc(injector_startingvalues + best, scaling=scaling, post_injector=False, verbose=False, runGenesis=False, dir='nelder_mead/basefiles_'+str(i))
self.set_CLARA_directory(CLARA_dir)
self.save_parameters = []
self.changes = None
# self.base_files = '../../../../basefiles_' + str(int(scaling)) + '/'
self.verbose = False
self.best_changes = './transverse_best_changes.yaml'
def __init__(self):
object.__init__(self)
self.my_name = "data"
self.parameterDict = collections.OrderedDict()
self.lattices = ['INJ', 'S02', 'C2V', 'EBT', 'BA1']
[
self.parameterDict.update({l: collections.OrderedDict()})
for l in self.lattices
]
# self.latticeDict = self.parameterDict['lattice']
self.parameterDict['scan'] = collections.OrderedDict()
self.scanDict = self.parameterDict['scan']
self.scanDict['parameter_scan'] = False
self.parameterDict['simulation'] = collections.OrderedDict()
self.simulationDict = self.parameterDict['simulation']
self.parameterDict['generator'] = collections.OrderedDict()
self.generatorDict = self.parameterDict['generator']
self.scannableParametersDict = collections.OrderedDict()
self.Framework = Fw.Framework(directory=directory_summary,
clean=False,
verbose=True)
print('self.Framework.master_lattice_location = ',
self.Framework.master_lattice_location)
self.Framework.defineElegantCommand(location=[
self.Framework.master_lattice_location + 'Codes/elegant'
])
os.environ[
'RPN_DEFNS'] = self.Framework.master_lattice_location + 'Codes/defns.rpn'
self.Framework.loadSettings('Lattices/CLA10-BA1_OM.def')
self.my_name = "data"
self.get_data()
self.initialise_data()
def create_base_files(scaling, changes=None):
framework = Framework('basefiles_' + str(scaling), overwrite=True)
framework.loadSettings('afterglow.def')
framework.generator.number_of_particles = 2**(3 * scaling)
if not changes is None:
framework.load_changes_file(changes)
framework['S02'].prefix = '../../CLARA/basefiles_' + str(
int(scaling)) + '/'
framework.track(startfile='S02')
def create_base_files_20pC(scaling):
framework = Framework('basefiles_' + str(scaling) + '_20pC',
overwrite=True)
framework.loadSettings('Lattices/clara400_v12_v3.def')
if not os.name == 'nt':
framework.defineASTRACommand(scaling=scaling)
framework.defineCSRTrackCommand(scaling=scaling)
framework.generator.number_of_particles = 2**(3 * scaling)
framework.generator.charge = 20e-12
framework['CLA-HRG1-GUN-SOL'].field_amplitude = 0.322
framework['CLA-L01-CAV-SOL-01'].field_amplitude = 0.08
framework['CLA-L01-CAV-SOL-02'].field_amplitude = -0.08
framework['CLA-L02-CAV'].phase = 20
framework['CLA-L03-CAV'].phase = 30
# framework['bunch_compressor'].angle = 0.1
# framework['L02'].sample_interval = 2**(3*2)
framework.track(endfile="S07") #, startfile='VBC')
def __init__(self,
args,
tempdir,
scaling=4,
overwrite=True,
verbose=False,
summary=False,
post_injector=True,
parameterfile=[]):
self.cons = constraintsClass()
self.beam = rbf.beam()
self.twiss = rtf.twiss()
self.scaling = scaling
self.tmpdir = tempdir
self.verbose = verbose
self.summary = summary
self.overwrite = overwrite
self.post_injector = post_injector
self.parameters = list(args)
self.dirname = os.path.basename(self.tmpdir)
self.framework = Framework(self.dirname, clean=False)
self.framework.loadSettings('Lattices/clara400_v12_v3.def')
if isinstance(parameterfile, (list, tuple)):
for p in parameterfile:
self.framework.loadParametersFile(p)
elif parameterfile is not None:
self.framework.loadParametersFile(parameterfile)
if not os.name == 'nt':
self.framework.defineASTRACommand([
'mpiexec', '-np',
str(3 * scaling), '/opt/ASTRA/astra_MPICH2.sh'
])
self.framework.defineGeneratorCommand(['/opt/ASTRA/generator.sh'])
self.framework.defineCSRTrackCommand([
'/opt/OpenMPI-1.4.3/bin/mpiexec', '-n',
str(3 * scaling), '/opt/CSRTrack/csrtrack_openmpi.sh'
])
self.framework.generator.number_of_particles = 2**(3 * scaling)
else:
self.framework.generator.number_of_particles = 2**(3 * 3)
self.framework.setElementType('quadrupole', 'k1', self.parameters)
def get_Cavity_Phases(settings, dir='.'):
rcplattice = Framework('.', clean=False, verbose=False, overwrite=False)
rcplattice.loadSettings(settings)
files = filter(os.path.isfile, glob.glob(dir+"/*.log"))
files.sort(key=lambda x: os.path.getmtime(x))
cavity_phases = []
cavity_energy_gain = []
for f in files:
section = os.path.basename(f).replace('.log','')
# print 'section = ', section
cavs = rcplattice[section].getElementType('cavity', setting='objectName')
# phases = rcplattice['injector400'].getElementType('cavity', setting='phase')
with open(f) as file:
n = 0
for line in file:
ans = cavity_regex.findall(line)
if len(ans) > 0:
cavity_phases.append([cavs[n], float(ans[-1][-1])])
cavity_energy_gain.append([cavs[n], float(ans[-1][1])])
n+=1
return cavity_phases, cavity_energy_gain
def examples_2():
lattice = fw.Framework('example', clean=False, verbose=True)
lattice.loadSettings('FEBE/FEBE.def')
if not os.name == 'nt': # i.e. we are on linux
scaling = 5
lattice.defineASTRACommand(scaling=scaling)
# lattice.defineASTRACommand(['/opt/ASTRA/astra.sh'])
lattice.defineCSRTrackCommand(scaling=scaling)
lattice.generator.number_of_particles = 2**(3 * scaling)
else:
lattice.generator.number_of_particles = 2**(3 * 5)
lattice['S02'].prefix = '../CLARA/basefiles_5/'
lattice.change_Lattice_Code('All', 'elegant')
# lattice.track(startfile='S02')
lattice.read_Lattice(
'test_lattice', {
'code': 'elegant',
'charge': {
'cathode': False,
'space_charge_mode': '3D'
},
'input': {},
'output': {
'start_element': 'CLA-S07-MARK-01',
'end_element': 'CLA-FEB-W-START-01'
}
})
lattice['test_lattice'].preProcess()
lattice['test_lattice'].optimisation = fw.elegantOptimisation(
lattice['test_lattice'],
variables={'test': {
'item': 'betax',
'lower': 1,
'upper': 10
}})
print(lattice['test_lattice'].optimisation.commandObjects['test'].write())
lattice['test_lattice'].write()
lattice['test_lattice'].run()
lattice['test_lattice'].postProcess()
def __init__(self, args, tempdir, scaling=4, overwrite=True, verbose=False, summary=False, clean=False):
global preparameters, variables
self.cons = constraintsClass()
self.beam = rbf.beam()
self.twiss = rtf.twiss()
self.tmpdir = tempdir
self.scaling = scaling
self.verbose = verbose
self.summary = summary
self.parameters = preparameters + list(args)
self.dirname = os.path.basename(self.tmpdir)
self.framework = Framework(self.dirname, clean=clean, verbose=False)
self.framework.loadSettings('Lattices/clara400_v12_FEBE.def')
if not os.name == 'nt':
self.framework.defineASTRACommand(['mpiexec','-np',str(3*scaling),'/opt/ASTRA/astra_MPICH2.sh'])
self.framework.defineGeneratorCommand(['/opt/ASTRA/generator.sh'])
self.framework.defineCSRTrackCommand(['/opt/OpenMPI-1.4.3/bin/mpiexec','-n',str(3*scaling),'/opt/CSRTrack/csrtrack_openmpi.sh'])
# self.framework.defineElegantCommand(['mpiexec','-np',str(3*scaling),'Pelegant'])
# else:
self.framework.defineElegantCommand(['elegant'])
# self.framework.setElementType('quadrupole','k1', self.parameters)
[setattr(a, 'k1', 2*a['k1']) for a in variables]
def setup_lattice(self, inputargs, tempdir, *args, **kwargs):
self.dirname = tempdir
self.input_parameters = list(inputargs)
self.npart = 2**(3 * self.scaling)
ncpu = self.scaling * 3
self.framework = fw.Framework(self.dirname,
overwrite=self.overwrite,
verbose=self.verbose)
if not os.name == 'nt':
self.framework.defineGeneratorCommand(['/opt/ASTRA/generator'])
self.framework.defineASTRACommand(
['mpiexec', '-np',
str(ncpu), '/opt/ASTRA/astra_MPICH2.sh'])
self.framework.defineCSRTrackCommand([
'/opt/OpenMPI-1.4.3/bin/mpiexec', '-n',
str(ncpu), '/opt/CSRTrack/csrtrack_openmpi.sh'
])
self.framework.defineElegantCommand(['elegant'])
# if os.name == 'nt':
# self.framework.defineElegantCommand(['mpiexec','-np','10','pelegant'])
self.framework.loadSettings(self.lattice_file)
self.framework.change_Lattice_Code('All', 'elegant')
### Define starting lattice
if self.start_lattice is None:
if 'POSTINJ' in self.framework.latticeObjects:
self.start_lattice = 'POSTINJ'
else:
self.start_lattice = self.framework.latticeObjects[
0].objectname
### Apply any pre-tracking changes to elements
if self.changes is not None:
if isinstance(self.changes, (tuple, list)):
for c in self.changes:
self.framework.load_changes_file(c)
else:
self.framework.load_changes_file(self.changes)
### Apply input arguments to element definitions
''' Apply arguments: [[element, parameter, value], [...]] '''
for e, p, v in self.input_parameters:
self.framework.modifyElement(e, p, v)
### Save the changes to the run directory
self.framework.save_changes_file(filename=self.framework.subdirectory +
'/changes.yaml',
elements=self.input_parameters)
def __init__(self, V_MAG_Ctrl=None, C_S01_MAG_Ctrl=None,
C_S02_MAG_Ctrl=None, C2V_MAG_Ctrl=None, V_RF_Ctrl=None,
C_RF_Ctrl=None, L01_RF_Ctrl=None, messages=False):
QThread.__init__(self)
self.showMessages = messages
self.VELA_MAG_Controller = V_MAG_Ctrl
self.CLA_MAG_S01_Controller = C_S01_MAG_Ctrl
self.CLA_MAG_S02_Controller = C_S02_MAG_Ctrl
self.C2V_MAG_Controller = C2V_MAG_Ctrl
self.VELA_LLRF_Controller = V_RF_Ctrl
self.CLA_LLRF_Controller = C_RF_Ctrl
self.CLA_L01_LLRF_Controller = L01_RF_Ctrl
self.startElement = 'Null'
self.stopElement = 'Null'
self.initDistribFile = 'Null'
self.initCharge = 0.0
self.pathway = Framework.Framework(subdir='.', overwrite='overwrite')
def load_best(self, filename):
with open(filename, 'r') as infile:
data = dict(yaml.load(infile, Loader=yaml.UnsafeLoader))
# best = [data[n]['k1l'] for n in parameter_names]
best = []
for n, p in self.parameter_names:
if n in data:
best.append(data[n][p])
elif n == 'bunch_compressor' and p == 'set_angle':
best.append(data['CLA-VBC-MAG-DIP-01']['angle'])
else:
print(n, p)
if not hasattr(self, 'framework'):
self.framework = fw.Framework(None)
self.framework.loadSettings(self.lattice)
best.append(self.framework[n][p])
self.best = best
return best
def __init__(self, V_MAG_Ctrl=None, C_S01_MAG_Ctrl=None,
C_S02_MAG_Ctrl=None, C2V_MAG_Ctrl=None, LRRG_RF_Ctrl=None,
HRRG_RF_Ctrl=None, L01_RF_Ctrl=None, messages=False):
QThread.__init__(self)
self.showMessages = messages
self.V_MAG_Ctrl = V_MAG_Ctrl
self.C_S01_MAG_Ctrl = C_S01_MAG_Ctrl
self.C_S02_MAG_Ctrl = C_S02_MAG_Ctrl
self.C2V_MAG_Ctrl = C2V_MAG_Ctrl
self.HRRG_RF_Ctrl = HRRG_RF_Ctrl
self.LRRG_RF_Ctrl = LRRG_RF_Ctrl
self.L01_RF_Ctrl = L01_RF_Ctrl
self.startElement = None
self.stopElement = None
self.initDistribFile = 'Null'
self.initDistrib = None
self.initCharge = 0.0
self.pathway = Framework.Framework(subdir='.', overwrite='overwrite')
def setup_lattice(self, inputargs, tempdir, *args, **kwargs):
self.dirname = tempdir
self.input_parameters = list(inputargs)
self.framework = fw.Framework(self.dirname, clean=self.clean, overwrite=self.overwrite, verbose=self.verbose)
# print('self.ncpu = ', self.ncpu)
if not os.name == 'nt':
self.framework.defineASTRACommand(ncpu=self.ncpu)
self.framework.defineCSRTrackCommand(ncpu=self.ncpu)
self.framework.defineElegantCommand(ncpu=self.elegant_ncpu)
# if os.name == 'nt':
# self.framework.defineElegantCommand(['mpiexec','-np','10','pelegant'])
self.framework.loadSettings(self.lattice_file)
if hasattr(self, 'change_to_elegant') and self.change_to_elegant:
self.framework.change_Lattice_Code('All','elegant', exclude=['injector400'])
elif hasattr(self, 'change_to_astra') and self.change_to_astra:
self.framework.change_Lattice_Code('All','ASTRA', exclude=['injector400'])
elif hasattr(self, 'change_to_gpt') and self.change_to_gpt:
self.framework.change_Lattice_Code('All','GPT', exclude=['injector400'])
### Define starting lattice
if self.start_lattice is None:
self.start_lattice = self.framework[list(self.framework.latticeObjects)[0]].objectname
### Apply any pre-tracking changes to elements
if self.changes is not None:
if isinstance(self.changes, (tuple, list)):
for c in self.changes:
self.framework.load_changes_file(c)
else:
self.framework.load_changes_file(self.changes)
### Apply input arguments to element definitions
''' Apply arguments: [[element, parameter, value], [...]] '''
for e, p, v in self.input_parameters:
if e == 'startcharge':
print('Setting start charge to ', v)
self.startcharge = v
else:
self.framework.modifyElement(e, p, v)
### Save the changes to the run directory
self.framework.save_changes_file(filename=self.framework.subdirectory+'/changes.yaml', elements=self.input_parameters)
def create_base_files(scaling):
framework = fw.Framework('basefiles_' + str(scaling), clean=False)
framework.loadSettings('CLA10-BA1_TOMP.def')
# framework.loadSettings('Lattices/clara10_v12.def')
if not os.name == 'nt':
framework.defineASTRACommand(
['mpiexec', '-np',
str(4 * scaling), '/opt/ASTRA/astra_MPICH2.sh'])
framework.defineGeneratorCommand(['/opt/ASTRA/generator.sh'])
framework.defineCSRTrackCommand([
'/opt/OpenMPI-1.4.3/bin/mpiexec', '-n',
str(4 * scaling), '/opt/CSRTrack/csrtrack_openmpi.sh'
])
framework.defineElegantCommand(['elegant'])
framework.generator.number_of_particles = 2**(3 * scaling)
framework.modifyElement('CLA-LRG1-GUN-CAV', 'phase', -5)
# framework.modifyElement('CLA-HRG1-GUN-SOL', 'field_amplitude', gunsol)
# framework.modifyElement('CLA-L01-CAV', 'field_amplitude', abs(linac1field))
framework.modifyElement('CLA-L01-CAV', 'phase', +10)
# framework.track(files=['generator','injector10'])#startfile='VBC')
framework.track(files=['L01', 'S02', 'C2V', 'VELA'])
def __init__(self, lattice='Lattices/clara400_v12_v3_elegant_jkj.def'):
self.CLARA_dir = os.path.relpath(os.path.dirname(os.path.dirname( os.path.abspath(__file__))))# if CLARA_dir is None else CLARA_dir
self.cons = constraintsClass()
self.beam = rbf.beam()
self.twiss = rtf.twiss()
self.lattice_file = lattice
self.base_files = None
self.verbose = False
self.start_lattice = None
self.scaling = 6
self.overwrite = True
self.clean = False
self.post_injector = True
self.startcharge = None
self.changes = None
self.sample_interval = 1
self.npart=2**(3*self.scaling)
self.ncpu = self.scaling*3
self.elegant_ncpu = 1
self.genesis_ncpu = 2
self.doTracking = True
self.change_to_elegant = True
self.framework = fw.Framework('.', clean=False, overwrite=False, verbose=False)
################################ CSRTrack #####################################
# framework = Framework('VBC_CSRTrack')
# if not os.name == 'nt':
# framework.defineGeneratorCommand(['/opt/ASTRA/generator'])
# framework.defineASTRACommand(['mpiexec','-np',str(ncpu),'/opt/ASTRA/astra_MPICH2.sh'])
# framework.defineCSRTrackCommand(['/opt/OpenMPI-1.4.3/bin/mpiexec','-n',str(ncpu),'/opt/CSRTrack/csrtrack_openmpi.sh'])
# framework.defineElegantCommand(['elegant'])
#
# framework.loadSettings('Lattices/clara400_v12_v3.def')
# framework['VBC'].file_block['input']['prefix'] = '../basefiles_5/'
# framework.track(startfile='VBC', endfile='S07')
################################ ELEGANT ######################################
framework = Framework('VBC_Elegant')
if not os.name == 'nt':
# framework.defineGeneratorCommand(['/opt/ASTRA/generator'])
framework.defineASTRACommand(ncpu=ncpu)
framework.defineCSRTrackCommand(ncpu=ncpu)
framework.loadSettings('Lattices/clara400_v12_v3.def')
framework.change_Lattice_Code('VBC', 'elegant')
framework.change_Lattice_Code('S06', 'elegant')
framework.change_Lattice_Code('L04', 'elegant')
framework.change_Lattice_Code('S07', 'elegant')
framework['VBC'].file_block['input']['prefix'] = '../basefiles_5/'
framework.track(startfile='VBC', endfile='S07')
################################ ASTRA ######################################
def __init__(self,
args,
tempdir,
scaling=4,
overwrite=True,
verbose=False,
summary=False,
post_injector=True):
self.cons = constraintsClass()
self.beam = rbf.beam()
self.twiss = rtf.twiss()
self.scaling = scaling
self.tmpdir = tempdir
self.verbose = verbose
self.summary = summary
self.overwrite = overwrite
self.post_injector = post_injector
# print 'self.post_injector = ', self.post_injector
''' if only post-injector optimisation'''
if self.post_injector:
linac2field, linac2phase, linac3field, linac3phase, fhcfield, fhcphase, linac4field, linac4phase, bcangle = args
self.parameters = dict(
zip([
'linac2field', 'linac2phase', 'linac3field', 'linac3phase',
'fhcfield', 'fhcphase', 'linac4field', 'linac4phase',
'bcangle'
], args))
else:
''' including injector parameters '''
gunphase, gunsol, linac1field, linac1phase, linac1sol1, linac1sol2, linac2field, linac2phase, linac3field, linac3phase, fhcfield, fhcphase, linac4field, linac4phase, bcangle = args
self.parameters = dict(
zip([
'gunphase', 'gunsol', 'linac1field', 'linac1phase',
'linac1sol1', 'linac1sol2', 'linac2field', 'linac2phase',
'linac3field', 'linac3phase', 'fhcfield', 'fhcphase',
'linac4field', 'linac4phase', 'bcangle'
], args))
self.npart = 2**(3 * scaling)
ncpu = scaling * 3
if self.post_injector:
self.sbandlinacfields = np.array(
[linac2field, linac3field, linac4field])
else:
self.sbandlinacfields = np.array(
[linac1field, linac2field, linac3field, linac4field])
self.dirname = os.path.basename(self.tmpdir)
self.framework = Framework(self.dirname,
overwrite=overwrite,
verbose=verbose)
if not os.name == 'nt':
self.framework.defineGeneratorCommand(['/opt/ASTRA/generator'])
self.framework.defineASTRACommand(
['mpiexec', '-np',
str(ncpu), '/opt/ASTRA/astra_MPICH2.sh'])
self.framework.defineCSRTrackCommand([
'/opt/OpenMPI-1.4.3/bin/mpiexec', '-n',
str(ncpu), '/opt/CSRTrack/csrtrack_openmpi.sh'
])
self.framework.defineElegantCommand(['elegant'])
self.framework.loadSettings('Lattices/clara400_v12_v3.def')
if not self.post_injector:
self.framework.generator.particles = self.npart
self.framework.modifyElement('CLA-HRG1-GUN-CAV', 'phase', gunphase)
self.framework.modifyElement('CLA-HRG1-GUN-SOL', 'field_amplitude',
gunsol)
self.framework.modifyElement('CLA-L01-CAV', 'field_amplitude',
abs(linac1field))
self.framework.modifyElement('CLA-L01-CAV', 'phase', linac1phase)
self.framework.modifyElement('CLA-L01-CAV-SOL-01',
'field_amplitude', linac1sol1)
self.framework.modifyElement('CLA-L01-CAV-SOL-02',
'field_amplitude', linac1sol2)
self.framework.modifyElement('CLA-L02-CAV', 'field_amplitude',
abs(linac2field))
self.framework.modifyElement('CLA-L02-CAV', 'phase', linac2phase)
self.framework.modifyElement('CLA-L03-CAV', 'field_amplitude',
abs(linac3field))
self.framework.modifyElement('CLA-L03-CAV', 'phase', linac3phase)
self.framework.modifyElement('CLA-L4H-CAV', 'field_amplitude',
abs(fhcfield))
self.framework.modifyElement('CLA-L4H-CAV', 'phase', fhcphase)
self.framework.modifyElement('CLA-L04-CAV', 'field_amplitude',
abs(linac4field))
self.framework.modifyElement('CLA-L04-CAV', 'phase', linac4phase)
self.framework['bunch_compressor'].set_angle(abs(bcangle))
def create_base_files(scaling):
framework = Framework('basefiles_' + str(scaling), overwrite=True)
framework.loadSettings('Lattices/claraX400_v12_80MVm.def')
if not os.name == 'nt':
framework.defineASTRACommand(
['mpiexec', '-np',
str(3 * scaling), '/opt/ASTRA/astra_MPICH2.sh'])
framework.defineGeneratorCommand(['/opt/ASTRA/generator.sh'])
framework.defineCSRTrackCommand([
'/opt/OpenMPI-1.4.3/bin/mpiexec', '-n',
str(3 * scaling), '/opt/CSRTrack/csrtrack_openmpi.sh'
])
framework.generator.number_of_particles = 2**(3 * scaling)
framework.track(files=['generator', 'injector400', 'S02'])
class fitnessFunc():
def __init__(self,
args,
tempdir,
scaling=4,
overwrite=True,
verbose=False,
summary=False,
post_injector=True):
self.cons = constraintsClass()
self.beam = rbf.beam()
self.twiss = rtf.twiss()
self.scaling = scaling
self.tmpdir = tempdir
self.verbose = verbose
self.summary = summary
self.overwrite = overwrite
self.post_injector = post_injector
# print 'self.post_injector = ', self.post_injector
''' if only post-injector optimisation'''
if self.post_injector:
linac2field, linac2phase, linac3field, linac3phase, fhcfield, fhcphase, linac4field, linac4phase, bcangle = args
self.parameters = dict(
zip([
'linac2field', 'linac2phase', 'linac3field', 'linac3phase',
'fhcfield', 'fhcphase', 'linac4field', 'linac4phase',
'bcangle'
], args))
else:
''' including injector parameters '''
gunphase, gunsol, linac1field, linac1phase, linac1sol1, linac1sol2, linac2field, linac2phase, linac3field, linac3phase, fhcfield, fhcphase, linac4field, linac4phase, bcangle = args
self.parameters = dict(
zip([
'gunphase', 'gunsol', 'linac1field', 'linac1phase',
'linac1sol1', 'linac1sol2', 'linac2field', 'linac2phase',
'linac3field', 'linac3phase', 'fhcfield', 'fhcphase',
'linac4field', 'linac4phase', 'bcangle'
], args))
self.npart = 2**(3 * scaling)
ncpu = scaling * 3
if self.post_injector:
self.sbandlinacfields = np.array(
[linac2field, linac3field, linac4field])
else:
self.sbandlinacfields = np.array(
[linac1field, linac2field, linac3field, linac4field])
self.dirname = os.path.basename(self.tmpdir)
self.framework = Framework(self.dirname,
overwrite=overwrite,
verbose=verbose)
if not os.name == 'nt':
self.framework.defineGeneratorCommand(['/opt/ASTRA/generator'])
self.framework.defineASTRACommand(
['mpiexec', '-np',
str(ncpu), '/opt/ASTRA/astra_MPICH2.sh'])
self.framework.defineCSRTrackCommand([
'/opt/OpenMPI-1.4.3/bin/mpiexec', '-n',
str(ncpu), '/opt/CSRTrack/csrtrack_openmpi.sh'
])
self.framework.defineElegantCommand(['elegant'])
self.framework.loadSettings('Lattices/clara400_v12_v3.def')
if not self.post_injector:
self.framework.generator.particles = self.npart
self.framework.modifyElement('CLA-HRG1-GUN-CAV', 'phase', gunphase)
self.framework.modifyElement('CLA-HRG1-GUN-SOL', 'field_amplitude',
gunsol)
self.framework.modifyElement('CLA-L01-CAV', 'field_amplitude',
abs(linac1field))
self.framework.modifyElement('CLA-L01-CAV', 'phase', linac1phase)
self.framework.modifyElement('CLA-L01-CAV-SOL-01',
'field_amplitude', linac1sol1)
self.framework.modifyElement('CLA-L01-CAV-SOL-02',
'field_amplitude', linac1sol2)
self.framework.modifyElement('CLA-L02-CAV', 'field_amplitude',
abs(linac2field))
self.framework.modifyElement('CLA-L02-CAV', 'phase', linac2phase)
self.framework.modifyElement('CLA-L03-CAV', 'field_amplitude',
abs(linac3field))
self.framework.modifyElement('CLA-L03-CAV', 'phase', linac3phase)
self.framework.modifyElement('CLA-L4H-CAV', 'field_amplitude',
abs(fhcfield))
self.framework.modifyElement('CLA-L4H-CAV', 'phase', fhcphase)
self.framework.modifyElement('CLA-L04-CAV', 'field_amplitude',
abs(linac4field))
self.framework.modifyElement('CLA-L04-CAV', 'phase', linac4phase)
self.framework['bunch_compressor'].set_angle(abs(bcangle))
def between(self, value, minvalue, maxvalue, absolute=True):
if absolute:
result = max([minvalue, min([maxvalue, abs(value)])])
else:
result = np.sign(value) * max(
[minvalue, min([maxvalue, abs(value)])])
return result
def calculateBeamParameters(self):
bcangle = self.framework['bunch_compressor'].angle
# print 'bcangle = ', bcangle
try:
# if abs(bcangle) < 0.01 or abs(bcangle) > 0.175:
# raise ValueError
if self.overwrite:
startS = self.framework['S02'].startObject['position_start'][2]
if self.post_injector:
self.framework['S02'].file_block['input'][
'prefix'] = '../basefiles_' + str(self.scaling) + '/'
self.framework.track(startfile='S02',
endfile='S07') #startfile='FMS')
else:
self.framework.track(endfile='S07')
# self.beam.read_astra_beam_file(self.dirname+'/S07.4928.001')
self.beam.read_HDF5_beam_file(self.dirname +
'/CLA-S07-MARK-03.hdf5')
self.beam.slices = 10
self.beam.bin_time()
sigmat = 1e12 * np.std(self.beam.t)
sigmap = np.std(self.beam.p)
meanp = np.mean(self.beam.p)
emitx = 1e6 * self.beam.normalized_mve_horizontal_emittance
emity = 1e6 * self.beam.normalized_mve_horizontal_emittance
density = self.beam.density
fitp = 100 * sigmap / meanp
fhcfield = self.parameters['fhcfield']
peakI, peakIMomentumSpread, peakIEmittanceX, peakIEmittanceY, peakIMomentum, peakIDensity = self.beam.mvesliceAnalysis(
)
chirp = self.beam.chirp
constraintsList = {
'peakI_min': {
'type': 'greaterthan',
'value': abs(peakI),
'limit': 600,
'weight': 60
},
'peakI_max': {
'type': 'lessthan',
'value': abs(peakI),
'limit': 750,
'weight': 10
},
# 'peakIMomentumSpread': {'type': 'lessthan', 'value': peakIMomentumSpread, 'limit': 0.1, 'weight': 2},
# 'peakIEmittanceX': {'type': 'lessthan', 'value': 1e6*peakIEmittanceX, 'limit': 0.75, 'weight': 5},
# 'peakIEmittanceY': {'type': 'lessthan', 'value': 1e6*peakIEmittanceY, 'limit': 0.75, 'weight': 5},
'peakIMomentum': {
'type': 'equalto',
'value': 1e-6 * peakIMomentum,
'limit': 220,
'weight': 20
},
'sband_linac fields': {
'type': 'lessthan',
'value': 1e-6 * self.sbandlinacfields,
'limit': 32,
'weight': 200
},
# 'xband_linac fields': {'type': 'lessthan', 'value': 1e-6*self.xbandlinacfields, 'limit': 100, 'weight': 100},
'4hc field': {
'type': 'lessthan',
'value': 1e-6 * fhcfield,
'limit': 35,
'weight': 100
},
# 'horizontal emittance': {'type': 'lessthan', 'value': emitx, 'limit': 2, 'weight': 0},
# 'vertical emittance': {'type': 'lessthan', 'value': emity, 'limit': 2, 'weight': 0},
# 'momentum_spread': {'type': 'lessthan', 'value': fitp, 'limit': 0.1, 'weight': 2},
'chirp': {
'type': 'equalto',
'value': abs(chirp),
'limit': 0.75,
'weight': 5
},
'correct_chirp': {
'type': 'lessthan',
'value': chirp,
'limit': 0,
'weight': 100
},
'peakI_volume': {
'type': 'greaterthan',
'value': peakIDensity,
'limit': 1e32,
'weight': 5
},
'volume': {
'type': 'greaterthan',
'value': density,
'limit': 1e30,
'weight': 5
},
}
# self.twiss.read_astra_emit_files(self.dirname+'/S07.Zemit.001')
# constraintsList5 = {
# 'last_exn_5': {'type': 'lessthan', 'value': 1e6*self.twiss['enx'], 'limit': 0.75, 'weight': 1},
# 'last_eyn_5': {'type': 'lessthan', 'value': 1e6*self.twiss['eny'], 'limit': 0.75, 'weight': 1},
# }
# constraintsList = merge_two_dicts(constraintsList, constraintsList5)
fitness = self.cons.constraints(constraintsList)
if self.verbose:
print self.cons.constraintsList(constraintsList)
if self.summary:
np.save('summary_constraints.txt',
self.cons.constraintsList(constraintsList))
# self.astra.createHDF5Summary(reference='Longitudinal_GA')
print fitness, 1e-6 * peakIMomentum, abs(
peakI
), 1e6 * peakIEmittanceX, 1e6 * peakIEmittanceY, chirp, peakIDensity, density
return fitness
except Exception as e:
print(e)
return 1e6
print(e)
return 1e6
def optfunc(args, dir=None, **kwargs):
if dir == None:
with TemporaryDirectory(dir=os.getcwd()) as tmpdir:
fit = fitnessFunc(args, tmpdir, **kwargs)
fitvalue = fit.calculateBeamParameters()
else:
fit = fitnessFunc(args, dir, **kwargs)
fitvalue = fit.calculateBeamParameters()
return (fitvalue, )
framework = Framework('longitudinal_best', overwrite=False)
framework.loadSettings('Lattices/claraX400_v12_80MVm.def')
parameters = []
''' if including injector'''
# parameters.append(framework.getElement('CLA-HRG1-GUN-CAV', 'phase'))
# parameters.append(framework.getElement('CLA-HRG1-GUN-SOL', 'field_amplitude'))
# parameters.append(framework.getElement('CLA-L01-CAV', 'field_amplitude'))
# parameters.append(framework.getElement('CLA-L01-CAV', 'phase'))
# parameters.append(framework.getElement('CLA-L01-CAV-SOL-01', 'field_amplitude'))
# parameters.append(framework.getElement('CLA-L01-CAV-SOL-02', 'field_amplitude'))
''' always '''
parameters.append(framework.getElement('CLA-L02-CAV', 'field_amplitude'))
parameters.append(framework.getElement('CLA-L02-CAV', 'phase'))
parameters.append(framework.getElement('CLA-L03-CAV', 'field_amplitude'))
parameters.append(framework.getElement('CLA-L03-CAV', 'phase'))
parameters.append(framework.getElement('CLA-L4H-CAV', 'field_amplitude'))
def load_lattice():
global lattice
lattice = Framework('example', clean=False, verbose=False)
lattice.loadSettings('split.def')
if not os.name == 'nt':
scaling = 4
lattice.defineASTRACommand(['mpiexec','-np',str(4*scaling),'/opt/ASTRA/astra_MPICH2.sh'])
# lattice.defineASTRACommand(['/opt/ASTRA/astra.sh'])
lattice.defineGeneratorCommand(['/opt/ASTRA/generator.sh'])
lattice.defineCSRTrackCommand(['/opt/OpenMPI-1.4.3/bin/mpiexec','-n',str(3*scaling),'/opt/CSRTrack/csrtrack_openmpi.sh'])
lattice.generator.number_of_particles = 2**(3*scaling)
else:
lattice.generator.number_of_particles = 2**(3*3)
lattice.defineElegantCommand(['elegant'])
def create_base_file(settings='./clara400_v12_v3.def'):
cavity_absolute_phases, cavity_absolute_momenta = rcp.get_Cavity_Phases(dir='../basefiles_4', settings='./clara400_v12_v3.def')
lattice = Framework('test', clean=False, verbose=False)
lattice.loadSettings(settings)
if not os.name == 'nt':
scaling = 4
lattice.defineASTRACommand(['mpiexec','-np',str(3*scaling),'/opt/ASTRA/astra_MPICH2.sh'])
# lattice.defineASTRACommand(['/opt/ASTRA/astra.sh'])
lattice.defineGeneratorCommand(['/opt/ASTRA/generator.sh'])
lattice.defineCSRTrackCommand(['/opt/OpenMPI-1.4.3/bin/mpiexec','-n',str(3*scaling),'/opt/CSRTrack/csrtrack_openmpi.sh'])
lattice.generator.number_of_particles = 2**(3*scaling)
else:
lattice.generator.number_of_particles = 2**(3*3)
lattice.defineElegantCommand(['elegant'])
for lat in lattice.latticeObjects:
print 'lattice = ', lat
if hasattr(lattice[lat], 'headers') and not lat == 'generator':
lattice.setSubDirectory('basefiles')
lattice[lat].headers['newrun']['auto_phase'] = True
lattice[lat].headers['newrun']['Track_All'] = False
lattice[lat].file_block['input']['prefix'] = '../test/'
lattice.track(files=[lat], postprocess=False)
lattice.setSubDirectory('test')
cavities = lattice[lat].getElementType('cavity')
if len(cavities) > 0:
absolute_phases, absolute_momenta = rcp.get_Cavity_Phase('basefiles/'+lat+'.log')
print zip(cavities, absolute_phases)
for cav, absphase in zip(cavities, absolute_phases):
cav['phase'] = absphase + cav['phase']
lattice[lat].headers['newrun']['auto_phase'] = False
lattice[lat].headers['newrun']['Track_All'] = True
lattice.track(files=[lat], postprocess=True)
def optimise_Lattice(phase=4, q=70, do_optimisation=False):
global dir, lattice, scaling, bestdelta
bestdelta = 1e10
dir = './SETUP/TOMP_SETUP'
lattice = Framework(dir, clean=False, verbose=False)
lattice.loadSettings('CLA10-BA1_TOMP_ASTRA.def')
lattice.change_Lattice_Code('S02', 'elegant')
lattice.change_Lattice_Code('C2V', 'elegant')
scaling = 6
if not os.name == 'nt':
lattice.defineASTRACommand(scaling=(scaling))
lattice.defineCSRTrackCommand(scaling=(scaling))
lattice.define_gpt_command(scaling=(scaling))
lattice['L01'].file_block['input']['prefix'] = '../basefiles_' + str(
scaling) + '_' + str(q) + '/'
quads = 0.107 * np.array([
21.11058462, -11.36377551, 24.69336696, -22.63264054, 56.07039682,
-51.58739658
])
quads = [
# lattice.getElement('CLA-S02-MAG-QUAD-01', 'k1l'),
# lattice.getElement('CLA-S02-MAG-QUAD-02', 'k1l'),
# lattice.getElement('CLA-S02-MAG-QUAD-03', 'k1l'),
# lattice.getElement('CLA-S02-MAG-QUAD-04', 'k1l'),
# lattice.getElement('CLA-C2V-MAG-QUAD-01', 'k1l'),
# lattice.getElement('CLA-C2V-MAG-QUAD-02', 'k1l'),
lattice.getElement('EBT-INJ-MAG-QUAD-07', 'k1l'),
lattice.getElement('EBT-INJ-MAG-QUAD-08', 'k1l'),
lattice.getElement('EBT-INJ-MAG-QUAD-09', 'k1l'),
lattice.getElement('EBT-INJ-MAG-QUAD-10', 'k1l'),
lattice.getElement('EBT-INJ-MAG-QUAD-11', 'k1l'),
lattice.getElement('EBT-INJ-MAG-QUAD-15', 'k1l'),
lattice.getElement('EBT-BA1-MAG-QUAD-01', 'k1l'),
lattice.getElement('EBT-BA1-MAG-QUAD-02', 'k1l'),
lattice.getElement('EBT-BA1-MAG-QUAD-03', 'k1l'),
lattice.getElement('EBT-BA1-MAG-QUAD-04', 'k1l'),
lattice.getElement('EBT-BA1-MAG-QUAD-05', 'k1l'),
lattice.getElement('EBT-BA1-MAG-QUAD-06', 'k1l'),
lattice.getElement('EBT-BA1-MAG-QUAD-07', 'k1l'),
]
quads = np.array([
1.813883341270411, -1.657984997152041, 2.436078305927638,
-1.3805803719324001, 5.757521138743945, -4.774512264530579,
1.5759640821849459, -1.5658635139945611, 0.4665090704522008,
-0.3701448642096947, 0.14139459095219498, 0.1294045271940972,
-0.26170990073523187, -0.972243028906357, 1.2999713197810956,
1.3880410074706666, 0.043801044316032774, -1.3992446074305926,
1.3604179976621753
])
quadsCHICANE = quads[:6]
# quadsCHICANE = np.array([ 1.81966806, -1.60694003, 2.52387281, -1.34318408, 5.75696896,
# -4.86231975,])
lattice['S02'].sample_interval = 2**(3 * 3)
if do_optimisation:
if phase is not None:
lattice['S02'].file_block['input'][
'prefix'] = '../TOMP_SETUP_' + str(phase) + '_' + str(q) + '/'
lattice['S02'].file_block['output']['end_element'] = 'CLA-C2V-MARK-02'
val = optFuncChicane(quadsCHICANE)
quadsCHICANE = setChicane(quadsCHICANE)
val = optFuncChicane(quadsCHICANE)
print val
print('Chicane = ', quadsCHICANE)
quads[:6] = quadsCHICANE
# exit()
lattice['S02'].file_block['output'][
'end_element'] = 'EBT-BA1-DIA-FCUP-01'
val = optFuncVELA(quads)
while val > 10:
print('################ CURRENT VALUE OF OPTIMISATION IS ', val,
' ##################')
quads = setVELA(quads)
val = optFuncVELA(quads)
else:
val = optFuncVELA(quads)
lattice['S02'].file_block['output']['end_element'] = 'EBT-BA1-DIA-FCUP-01'
lattice['S02'].sample_interval = 1
return quads
def test_ASTRA_Elegant_C2V(quads, phase=4, charge=70):
global dir, lattice, scaling, bestdelta
dir = './SETUP/TOMP_SETUP'
lattice = Framework(dir, clean=True, verbose=False)
lattice.loadSettings('CLA10-BA1_TOMP_ASTRA.def')
lattice.change_Lattice_Code('C2V', 'elegant')
scaling = 6
if not os.name == 'nt':
lattice.defineASTRACommand(scaling=(scaling))
lattice.defineCSRTrackCommand(scaling=(scaling))
lattice.define_gpt_command(scaling=(scaling))
lattice['S02'].file_block['input'][
'prefix'] = '../SETUP/TOMP_SETUP_' + str(phase) + '_' + str(
charge) + '/'
lattice['S02'].sample_interval = 1 #2**(3*2)
dir = './test_EleC2V_' + str(phase) + '_' + str(charge)
lattice.setSubDirectory(dir)
lattices = lattice.latticeObjects.values()
# for l in lattices:
# if hasattr(l, 'headers'):
# if isinstance(l.headers, (dict)):
# if 'charge' in l.headers.keys():
# l.headers['charge'].space_charge_mode = False
optFuncVELA(quads, track=False)
lattice.track(startfile='S02', track=True)
################################ ELEGANT ######################################
# framework = Framework('CLA10BA1', clean=True)
# if not os.name == 'nt':
# framework.defineGeneratorCommand(['/opt/ASTRA/generator'])
# framework.defineASTRACommand(['mpiexec','-np',str(ncpu),'/opt/ASTRA/astra_MPICH2.sh'])
# framework.defineCSRTrackCommand(['/opt/OpenMPI-1.4.3/bin/mpiexec','-n',str(ncpu),'/opt/CSRTrack/csrtrack_openmpi.sh'])
# framework.defineElegantCommand(['elegant'])
#
# framework.loadSettings('Lattices/CLA10-BA1.def')
# framework.change_Lattice_Code('All','elegant')
# framework['S02'].prefix = '../../basefiles_4/'
# framework.track(startfile='S02', endfile='BA1_dipole')
basedir = 'dipole_scan'
framework = Framework(basedir, clean=True, verbose=False)
if not os.name == 'nt':
framework.defineGeneratorCommand(['/opt/ASTRA/generator'])
framework.defineASTRACommand(
['mpiexec', '-np',
str(ncpu), '/opt/ASTRA/astra_MPICH2.sh'])
framework.defineCSRTrackCommand([
'/opt/OpenMPI-1.4.3/bin/mpiexec', '-n',
str(ncpu), '/opt/CSRTrack/csrtrack_openmpi.sh'
])
framework.defineElegantCommand(['elegant'])
framework.loadSettings('Lattices/CLA10-BA1.def')
framework.change_Lattice_Code('All', 'elegant')
framework['S02'].prefix = '../../basefiles_4/'
etax = []
lattice['test_lattice'].preProcess()
lattice['test_lattice'].optimisation = fw.elegantOptimisation(
lattice['test_lattice'],
variables={'test': {
'item': 'betax',
'lower': 1,
'upper': 10
}})
print(lattice['test_lattice'].optimisation.commandObjects['test'].write())
lattice['test_lattice'].write()
lattice['test_lattice'].run()
lattice['test_lattice'].postProcess()
# eample_2()
# exit()
lattice = fw.Framework('example', clean=False, verbose=True)
lattice.loadSettings('Lattices/clara400_v12_v3.def')
lattice['S02'].prefix = '../CLARA/basefiles_4/'
# lattice.setSubDirectory('example_astra')
# lattice.track(startfile='FMS', endfile='FMS')
# lattice.setSubDirectory('example_elegant')
lattice.change_Lattice_Code('S02', 'ASTRA')
# lattice.load_changes_file('./CLARA/Elegant_Genesis/Short-240/transverse_best_changes.yaml')
# lattice.modifyElement('CLA-S07-DCP-01', 'scale_kick', 1e9)
lattice.track(startfile='S02', endfile='S02')
# # lattice.track(startfile='S02')#,preprocess=True, track=False, postprocess=False)
# lattice['bunch_compressor'].set_angle(-0.0)
# lattice.track(startfile='S02', endfile='S07')#,preprocess=True, track=False, postprocess=False)