from functools import partial import glob import h5py from multiprocessing import Pool import numpy as np import os import shutil import subprocess from scipy.optimize import minimize_scalar as min_sc from A3PI_ImpactTemplate import ImpactTemplate def run_impactt(workflow, arg): #This workflow task performs the phase scan algorithm (to optimize the #lattice phase) and runs ImpactT with the given parameters defined in the #workflow configuration file sections for ImpactT. #Load ImpactT run parameters and path A3PI_mode = workflow.config.get('RUN_PARAMETERS', 'A3PI_mode') cores = workflow.config.get('RUN_PARAMETERS', 'impactt_cores') impactt_exe = workflow.config.get('PATHS', 'impactt_path') #Create an ImpactT template and write it to the file ImpactT.in or use #existing (static) ImpactT input file if workflow.config.has_option('RUN_PARAMETERS', 'impactt_input_file'): in_file = workflow.config.get('RUN_PARAMETERS', 'impactt_input_file') work_file = workflow.folder + '/' + in_file else: work_file = workflow.folder + '/ImpactT.in' impactt_temp = make_impactt_template(workflow) impactt_temp.write_file(work_file) if arg == '1': #Run ImpactT phase scan (single-particle) run_impact_phase_scan(workflow, work_file, thds=8) elif arg == '2': #Run ImpactT phase scan (single-particle) w/o normalization run_impact_phase_scan(workflow, work_file, thds=8, normalize=False) elif arg == '1a': #Run ImpactT phase scan (single-particle) run_impact_phase_scan(workflow, work_file, thds=8) return #Run phase scan only elif arg == '2a': #Run ImpactT phase scan (single-particle) w/o normalization run_impact_phase_scan(workflow, work_file, thds=8, normalize=False) return #Run phase scan only if A3PI_mode == 'single': shutil.copy(impactt_exe, workflow.folder) subprocess.run('srun -n ' + cores + ' --cpu-bind=cores ./' + os.path.split(impactt_exe)[1], shell=True, cwd=workflow.folder) elif A3PI_mode == 'worker': from libensemble.executors.executor import Executor exctr = Executor.executor os.chdir(workflow.folder) task = exctr.submit(app_name='impactt', num_procs=eval(cores), app_args=in_file, wait_on_run=True, extra_args='--cpu-bind=cores') task.wait() if task.success is False: task.kill() os.chdir('..') def run_impactz(workflow, arg): #This section simply is used to call ImpactZ, if a previous ImpactT run was #performed, the arg '1' can be used to pass particle information in the #file 'fort.50' to ImpactZ. A template parser is not implemented yet, so #A3PI cannot be used to generate ImpactZ input files for optimization A3PI_mode = workflow.config.get('RUN_PARAMETERS', 'A3PI_mode') cores = workflow.config.get('RUN_PARAMETERS', 'impactz_cores') impactz_exe = workflow.config.get('PATHS', 'impactz_path') if workflow.config.has_option('RUN_PARAMETERS', 'impactz_input_file'): in_file = workflow.config.get('RUN_PARAMETERS', 'impactz_input_file') else: print('ImpactZ parser not yet implemented.') raise if arg == '1': parse_particles(workflow.folder) if A3PI_mode == 'single': shutil.copy(impactz_exe, workflow.folder) subprocess.run('srun -n ' + cores + ' --cpu-bind=cores ./' + os.path.split(impactz_exe)[1], shell=True, cwd=workflow.folder) elif A3PI_mode == 'worker': from libensemble.executors.executor import Executor exctr = Executor.executor os.chdir(workflow.folder) task = exctr.submit(app_name='impactz', num_procs=eval(cores), app_args=in_file, wait_on_run=True, extra_args='--cpu-bind=cores') task.wait() if task.success is False: task.kill() os.chdir('..') def read_impact_output(workflow, arg): #Define function to read outputs from respective ImpactT fort.xx files. #Addtional parameters can be added by using the ImpactT user guide to locate #the appropriate file number and data column and using the same structure as #existing parameters as an example. param = arg.lower() if param == 'beam_energy': with open(workflow.folder + '/fort.18', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = tailline.split()[3] elif param == 'x_rms_size': with open(workflow.folder + '/fort.24', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = tailline.split()[3] elif param == 'x_rms_emittance': with open(workflow.folder + '/fort.24', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = tailline.split()[7] elif param == 'y_rms_size': with open(workflow.folder + '/fort.25', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = tailline.split()[3] elif param == 'y_rms_emittance': with open(workflow.folder + '/fort.25', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = tailline.split()[7] elif param == 'xy_rms_emittance': with open(workflow.folder + '/fort.24', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = float(tailline.split()[7]) with open(workflow.folder + '/fort.25', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = np.sqrt(float(data)**2 + float(tailline.split()[7])**2) elif param == 'z_rms_size': with open(workflow.folder + '/fort.26', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = tailline.split()[2] elif param == 'z_rms_emittance': with open(workflow.folder + '/fort.26', 'r') as file: tailline = file.readlines()[-1].replace('\n','') data = tailline.split()[6] else: print('Unknown output parameter listed for optimization.') raise return float(data) def make_impactt_template(workflow): #This ad-hoc routine creates an ImpactT template from quantities defined in #the workflow configuration file. impactt_temp = ImpactTemplate() #Create blank ImpactT template #Read ImpactT parameters from workflow config file procs = workflow.config.get('IMPACTT_PARAMETERS','procs') steps = workflow.config.get('IMPACTT_PARAMETERS','steps') parts = workflow.config.get('IMPACTT_PARAMETERS','parts') mesh = workflow.config.get('IMPACTT_PARAMETERS','mesh') dist = workflow.config.get('IMPACTT_PARAMETERS','dist') xdist = workflow.config.get('IMPACTT_PARAMETERS','xdist') ydist = workflow.config.get('IMPACTT_PARAMETERS','ydist') zdist = workflow.config.get('IMPACTT_PARAMETERS','zdist') beam = workflow.config.get('IMPACTT_PARAMETERS','beam') #Write Impact beam parameters to template impactt_temp.set_procs(procs) impactt_temp.set_steps(steps) impactt_temp.set_parts(parts) impactt_temp.set_mesh(mesh) impactt_temp.set_dist_type(dist) impactt_temp.set_dist('x',xdist) impactt_temp.set_dist('y',ydist) impactt_temp.set_dist('z',zdist) impactt_temp.set_beam(beam) #Read lattice from workflow config file and add to template lattice = workflow.config.items('IMPACTT_LATTICE') lattice.sort(key=lambda x: x[0]) for elem in range(len(lattice)): impactt_temp.add_element(lattice[elem][1]) return impactt_temp def impact_parse_input(file): #This routine reads an ImpactT input file and parses the file, ignoring #comment lines and empty rows, and exports the lines of data with open(file, 'r') as f: lines = f.readlines() numrows = len(lines) #Number of rows in input file ncflag = [0] * numrows #Flag to store non-commented rows ncrows = 0 #Number of non-commented rows in input file for i in range(numrows): if lines[i].strip() != '': #Ignore empty rows if not lines[i].startswith('!'): #Ignore commented rows ncrows += 1 ncflag[i] = ncrows return lines, ncflag, ncrows def run_impact_phase_scan(workflow, file, thds, normalize=True, cleanup=False): #This routine runs the phase scanning algorithm which optimizes lattice #element phases one by one, for maximum particle energy, from upstream to #downstream elements. The phase scan is done in a temporary folder and the #optimized phases are copied to the original input file at the end. impactt_exe = workflow.config.get('PATHS', 'impactt_path') if workflow.config.has_option('RUN_PARAMETERS', 'impactt_input_file'): in_file = workflow.config.get('RUN_PARAMETERS', 'impactt_input_file') else: in_file = 'ImpactT.in' #ImpactT run command used in single-core mode based on input file impact_cmd = './' + os.path.split(impactt_exe)[1] + ' ' + in_file + ' > a' phase_fs = [workflow.folder + '/phase_scan_' + str(i) for i in range(thds)] for phase_f in phase_fs: if not glob.glob(phase_f): os.mkdir(phase_f) shutil.copy(file, phase_f + '/' + in_file) shutil.copy(impactt_exe, phase_f) [lines, ncflag, ncrows] = impact_parse_input(phase_fs[0] + '/' + in_file) fids = [] #Used for storing file IDs used by ImpactT amps = dict() #Used for storing maximum field amplitudes in h5 files phases = dict() #Used for storing field phases in h5 files for elem in range(10, ncrows+1): #Scan lattice for needed data files lattice_line = lines[ncflag.index(elem)].split() if int(lattice_line[3]) in [101,104,105]: #Copy rfdata files if int(eval(lattice_line[8])) not in fids: fid = int(eval(lattice_line[8])) for phase_f in phase_fs: shutil.copy(workflow.folder + '/rfdata' + str(fid), phase_f) fids.append(fid) elif int(lattice_line[3]) == 111: #If hdf5 data, make rfdata file fid = int(eval(lattice_line[8])) if fid not in fids: fids.append(fid) print('Element file ID: ' + str(fid) + ' flagged as hdf5 file') make_t7_file(fid, workflow.folder) [Ez_max, Ez_phase] = construct_1d_field(fid, workflow.folder) amps[fid] = Ez_max phases[fid] = Ez_phase if not normalize: #Scale 1d fields by max of Ez line = lattice_line line[5] = '{:.9e}'.format(float(line[5])*Ez_max) lines[ncflag.index(elem)] = ' '.join(line) + '\n' print('1D field reconstruction of 3D hdf5 data saved to rfdata' + str(fid) + ' file (for phase scan)') for phase_f in phase_fs: shutil.copy(workflow.folder + '/rfdata' + str(fid), phase_f) else: Ez_max = amps[fid] if not normalize: #Scale 1d fields by max of Ez line = lattice_line line[5] = '{:.9e}'.format(float(line[5])*Ez_max) lines[ncflag.index(elem)] = ' '.join(line) + '\n' lattice_line[3] = '105' #Change element type for rfdata file lines[ncflag.index(elem)] = ' '.join(lattice_line) + '\n' #Adjust certain ImpactT parameters for phase scan routine line = lines[ncflag.index(1)].split() #Set number of processors to 1 x 1 line[0:2] = ['1', '1'] lines[ncflag.index(1)] = ' '.join(line) + '\n' line = lines[ncflag.index(3)].split() #Set number of particles to 1 line[1] = '1' lines[ncflag.index(3)] = ' '.join(line) + '\n' line = lines[ncflag.index(5)].split() #Set restart flag to 0 line[1] = '0' lines[ncflag.index(5)] = ' '.join(line) + '\n' line = lines[ncflag.index(9)].split() #Set beam current to 0 line[0] = '0.0' lines[ncflag.index(9)] = ' '.join(line) + '\n' for phase_f in phase_fs: with open(phase_f + '/' + in_file, 'w') as f: f.writelines(lines) print('Beginning optimization of lattice elements') print('==========================================') restart = False #Restart flag used for elements after first engs = [0] * thds if thds > 1: pool = Pool(thds) for elem in range(10,ncrows+1): #Loop of all lattice lines lattice_line = lines[ncflag.index(elem)].split() #Only optimize lattice elements with nonzero amplitude/frequency if int(lattice_line[3]) > 100 and \ float(lattice_line[5])!=0 and \ float(lattice_line[6])!=0: d_phase = lattice_line[7].replace('d','E') #Design phase row = ncflag.index(elem) #Current row of lattice element print('Optimizing phase for element in line: ' + str(row+1)) for phase_f in phase_fs: impact_add_restart_rows(row, phase_f, in_file) #Run scalar optimization to find phase with max particle energy min_sc_thread = partial(impact_minsc_thread, thds, row, restart, phase_fs, in_file, impact_cmd) if thds > 1: sol = pool.map(min_sc_thread, range(thds)) for i in range(thds): engs[i] = sol[i].fun thd_opt = engs.index(min(engs)) #Select best solution's thd. nfevs = [sol[i].nfev for i in range(thds)] else: sol = impact_minsc_thread(1, row, restart, phase_fs, in_file, impact_cmd, 0) thd_opt = 0 nfevs = sol.nfev #Shift optimized phase by design phase and rerun ImpactT phase = impact_offset_design_phase(row, phase_fs[thd_opt], d_phase, in_file) energy = -1*impact_phase_function(phase, row, restart, phase_fs[thd_opt], in_file, impact_cmd) impact_remove_restart_rows(row, phase_fs[thd_opt], in_file) shutil.copy(phase_fs[thd_opt] + '/fort.80', phase_fs[thd_opt] + '/fort_restart.80') #Pass best result of all threads to other threads for i in range(thds): if i is not thd_opt: shutil.copy(phase_fs[thd_opt] + '/' + in_file, phase_fs[i] + '/' + in_file) shutil.copy(phase_fs[thd_opt] + '/fort_restart.80', phase_fs[i] + '/fort_restart.80') print('Phase = ' + str(phase) + ' deg') print('Energy = ' + str(energy) + ' MeV') print('Iterations per thread = ' + str(nfevs)) print('==========================================') #Save the output phase space for next element and set restart flag if not restart: for phase_f in phase_fs: impact_set_restart_flag(phase_f, in_file) restart = True if thds > 1: pool.close() pool.join() #Finalize phase_scan process by copying the opt. phases to original file impact_finalize_phase_scan(workflow.folder, file, in_file, phase_fs[thd_opt], amps, phases) if cleanup: for phase_f in phase_fs: shutil.rmtree(phase_f) def make_t7_file(fid, folder): #This routine creates the necessary T7 file for use with openPMD formatted #hdf5 data, while the necessary data is contained within the hdf5 file #itself, a separate T7 file is needed for ImpactT #Read necessary mesh data from h5 files Ereal_f = h5py.File(folder +'/data' +str(fid) +'.h5','r') Ez_real = Ereal_f['data/200/fields/E_Real/z'] [xoff, yoff, zoff] = Ereal_f['data/200/fields/E_Real']\ .attrs['gridGlobalOffset'] [dx, dy, dz] = Ereal_f['data/200/fields/E_Real']\ .attrs['gridSpacing'] [xres, yres, zres] = Ez_real.shape xgrid = xoff + [dx*ind for ind in range(xres)] ygrid = yoff + [dy*ind for ind in range(yres)] zgrid = zoff + [dz*ind for ind in range(zres)] #Write the T7 file containing the Cartesian mesh information with open(folder + '/1T' + str(fid) + '.T7', 'w') as file: file.writelines( ["{:.9e}".format(xoff) + ' ' + "{:.9e}".format(xgrid[-1]) + ' ' + str(xres-1) + '\n' + "{:.9e}".format(yoff) + ' ' + "{:.9e}".format(ygrid[-1]) + ' ' + str(yres-1) + '\n' + "{:.9e}".format(0) + ' ' + "{:.9e}".format(zgrid[-1]-zgrid[0]) + ' ' + str(zres-1) + '\n']) def construct_1d_field(fid, folder): #This routine reads the Ez component from the data.h5 files and creates an #equivalent rfdata file used for the phase scan. The routine also returns #the maximum field amplitude and phase used to shift between the rfdata file #and original data.h5 file Ereal_f = h5py.File(folder +'/data' + str(fid) + '.h5', 'r') Eimag_f = h5py.File(folder +'/data' + str(fid+1) + '.h5', 'r') Ez_real = np.array(Ereal_f['data/200/fields/E_Real/z']) Ez_imag = np.array(Eimag_f['data/200/fields/E_Imag/z']) [xoff, yoff, zoff] = Ereal_f['data/200/fields/E_Real']\ .attrs['gridGlobalOffset'] [dx, dy, dz] = Ereal_f['data/200/fields/E_Real']\ .attrs['gridSpacing'] [xres, yres, zres] = Ez_real.shape xgrid = xoff + [dx*ind for ind in range(xres)] ygrid = yoff + [dy*ind for ind in range(yres)] zgrid = zoff + [dz*ind for ind in range(zres)] zlen = zgrid[-1] - zgrid[0] #Locate (x,y) index for z-axis (x=0, y=0) and extract on-axis Ez field x0 = xgrid.tolist().index(min(abs(xgrid))) y0 = ygrid.tolist().index(min(abs(ygrid))) Ez = Ez_real[x0, y0, :] + 1j*Ez_imag[x0, y0, :] Ez_max = np.max(abs(Ez)) #Ez maximum amplitude z0 = np.argmax(abs(Ez)) #Ez maximum amplitude position Ez_phase = np.angle(Ez[z0]) #Ez maximum amplitude phase Ez_norm = Ez * np.exp(-1j*Ez_phase) / Ez_max #Shift phase and normalize Ez0 = np.real(Ez_norm) #Remove negligible imag part of Ez_norm Ez0p = np.gradient(Ez0, dz, edge_order=2) #First derivative of Ez Ez0pp = np.gradient(Ez0p, dz, edge_order=2) #Second derivative of Ez Ez0ppp = np.gradient(Ez0pp, dz, edge_order=2) #Third derivative of Ez #Write rfdata file using 1D on-axis Ez field and its derivatives lines = [str(zres) + ' ' + "{:.6e}".format(0) + ' ' + "{:.6e}".format(zlen) + ' ' + "{:.6e}".format(zlen) + '\n'] for zind in range(len(zgrid)): lines.append(' ' + "{: .9e}".format(Ez0[zind]) + ' ' + "{: .9e}".format(Ez0p[zind]) + ' ' + "{: .9e}".format(Ez0pp[zind]) + ' ' + "{: .9e}".format(Ez0ppp[zind]) + '\n') lines.append('1 0.0 0.0 0.0 \n') #Ignore on-axis magnetic field lines.append('0.0 0.0 0.0 0.0 \n') with open(folder + '/rfdata' + str(fid), 'w') as file: file.writelines(lines) return Ez_max, Ez_phase*180/np.pi #Return amplitude and phase (deg) def impact_phase_function(phase, row, restart, folder, in_file, impact_cmd): #This function encapsulates the scalar function(phase) = energy impact_write_phase(phase, row, folder, in_file) #Write phase to file impact_run_phase(restart, folder, impact_cmd) #Run ImpactT eng = impact_read_energy(folder) #Read the energy output return -1*eng #The optimizer minimizes the negation of particle energy def impact_write_phase(phase, row, folder, in_file): #This routine writes the phase to the selected lattice element row [lines, ncflag, ncrows] = impact_parse_input(folder + '/' + in_file) line = lines[row].split() line[7] = str(phase) line = ' '.join(line) + '\n' lines[row] = line with open(folder + '/' + in_file, 'w') as file: file.writelines(lines) def impact_run_phase(restart, folder, impact_cmd): #This routine calls the ImpactT exectuable and resets the particles if restart: shutil.copy(folder + '/fort_restart.80', folder + '/fort.80') if os.name == 'nt': subprocess.run(impact_cmd, shell=False, cwd=folder) else: subprocess.run(impact_cmd, shell=True, cwd=folder) def impact_read_energy(folder): #This routine reads the output energy from the 'fort.18' file with open(folder + '/fort.18', 'r') as file: tailline = file.readlines()[-1].replace('\n','') eng = tailline.split()[3] return float(eng) def impact_minsc_thread(thds, row, restart, phase_fs, in_file, impact_cmd, thd): return min_sc(impact_phase_function, bounds=[thd*360/thds, (thd+1)*360/thds], args=(row, restart, phase_fs[thd], in_file, impact_cmd), method='bounded', options={'maxiter':20, 'xatol':1e-5}) def impact_offset_design_phase(row, folder, d_phase, in_file): #This routine shifts the phase in the input file by the design phase with open(folder + '/' + in_file, 'r') as file: lines = file.readlines() line = lines[row].split() phase = (float(line[7])+float(d_phase))%360 line[7] = str(phase) line = ' '.join(line) + '\n' lines[row] = line with open(folder + '/' + in_file, 'w') as file: file.writelines(lines) return phase def impact_add_restart_rows(row, folder, in_file): #This routine adds the 'output phase space' and 'stop simulation' flags [lines, ncflag, ncrows] = impact_parse_input(folder + '/' + in_file) line = lines[row].split() length = float(line[0]) zedge = float(line[4]) cav_end = length + zedge lines.insert(row+1,'0 0 80 -3 ' + str(cav_end) + ' ' + str(cav_end) + ' ' + str(cav_end) + ' / \n') lines.insert(row+2,'0 0 0 -99 ' + str(cav_end) + ' ' + str(cav_end) + ' ' + str(cav_end) + ' / \n') with open(folder + '/' + in_file, 'w') as file: file.writelines(lines) def impact_remove_restart_rows(row, folder, in_file): #This routine removes the 'output phase space' and 'stop simulation' flags [lines, ncflag, ncrows] = impact_parse_input(folder + '/' + in_file) lines.pop(row+1) lines.pop(row+1) with open(folder + '/' + in_file, 'w') as file: file.writelines(lines) def impact_set_restart_flag(folder, in_file): #This routine sets the restart flag in the input file to '1' [lines, ncflag, ncrows] = impact_parse_input(folder + '/' + in_file) line = lines[ncflag.index(5)].split() line[1] = '1' line = ' '.join(line) + '\n' lines[ncflag.index(5)] = line with open(folder + '/' + in_file, 'w') as file: file.writelines(lines) def impact_finalize_phase_scan(folder, file, in_file, phase_f, amps, phases): #This routine copies the optimized phases to the original ImpactT file and #removes the phase_scan temporary folder with open(file, 'r') as file: beam_lines = file.readlines() [lines, ncflag, ncrows] = impact_parse_input(phase_f + '/' + in_file) row_start = ncflag.index(10) #Index for lattice start for i in range(0,row_start): #Reset ImpactT non-lattice parameters lines[i] = beam_lines[i] for fid in phases: amp = amps[fid] phase = phases[fid] for elem in range(10, ncrows+1): line = lines[ncflag.index(elem)].split() if int(eval(line[3])) == 105: if int(eval(line[8])) == fid: line[3] = '111' line[5] = '{:.9e}'.format(float(line[5])/amp) line[7] = str((float(line[7])-float(phase))%360) line = ' '.join(line) + '\n' lines[ncflag.index(elem)] = line with open(phase_f + '/' + in_file, 'w') as file: file.writelines(lines) shutil.copyfile(phase_f + '/' + in_file, folder + '/' + in_file) def parse_particles(folder): #This script reads the particles from 'fort.50' and writes them to the new #file 'particle.in' for use in ImpactZ with a header row of particle total with open(folder + '/fort.50', 'r') as file: parts = file.readlines() num_parts = len(parts) with open(folder + '/particle.in', 'w') as file: file.writelines(str(num_parts) + '\n') file.writelines(parts)