Multi-objective optimization with ASTRA#
Description#
This example shows how to perform a multi-objective Bayesian optimization of beam parameters using ASTRA.
The setup is based on a beamline example from the ASTRA manual which can be found here.
Two optimization parameters are used:
the RF phase of the cavity
'RF_phase', which is varied in the range \([-2.5, 2.5]\),and the solenoid strength
'B_sol', which is varied in the range \([0.12, 0.38]\).
Two beam parameters are minimized:
the bunch_length,
and the transverse emittances
'emittance'in \(\mathrm{µm}\), which are combined into one single parameter: \(\log \epsilon_{n,x} \epsilon_{n,y}\) and where the logarithm is used for better optimization as the emittance can span over several orders of magnitude.
In addition, the transverse normalized emittances in \(x\) and \(y\) are stored as additional analyzed parameters 'emittance_x' and 'emittance_y'.
The optimization is carried out using an
AxSingleFidelityGenerator and a
TemplateEvaluator. In this case, the function
analyze_simulation that analyzes the output of each simulation is defined
in a separate file analysis_script.py and imported into the main
optimas script.
The ASTRA simulation template ASTRA_example.in requires additional files. These can be downloaded from the ASTRA website and are the input particle distribution Example.ini, the RF field profile 3_cell_L-Band.dat, and the solenoid field profile Solenoid.dat.
These files need to be passed to the TemplateEvaluator using the sim_files argument.
The path to the ASTRA executable needs to be specified in the TemplateEvaluator using the executable argument.
Scripts#
The files needed to run the optimization should be located in a folder
(named e.g., optimization) with the following structure:
optimization
├── run_optimization_serial_ASTRA.py
├── ASTRA_example.in
├── analysis_script.py
├── Example.ini
├── 3_cell_L-Band.dat
└── Solenoid.dat
The optimization is started by executing:
python run_optimization_serial_ASTRA.py
The main scripts needed to run this example can be seen below.
"""Example of multiobjective Bayesian optimization using the serial ASTRA version.
This example uses an ASTRA example from the ASTRA webpage:
https://www.desy.de/~mpyflo/EXAMPLES/Manual_Example/
In order to run this example, please download the files "3_cell_L-Band.dat",
"Solenoid.dat", and create the particle distribution "Example.ini" by running
the generator.in file.
Further optimas documentation and examples:
https://optimas.readthedocs.io/en/latest/index.html
"""
from optimas.generators import AxSingleFidelityGenerator
from optimas.evaluators import TemplateEvaluator
from optimas.explorations import Exploration
from gest_api.vocs import VOCS
from analysis_script import analyze_simulation
# Create VOCS object.
# name of parameter, lower bound of values to be explored,
# upper bound of values to be explored
vocs = VOCS(
variables={
"RF_phase": [-2.5, 2.5],
"B_sol": [0.12, 0.38],
},
objectives={
"bunch_length": "MINIMIZE",
"emittance": "MINIMIZE",
},
observables=["emittance_x", "emittance_y"],
)
# Create generator.
# Pick the generator to be used, here Single-fidelity Bayesian optimization.
gen = AxSingleFidelityGenerator(
vocs=vocs,
n_init=8,
)
# Create evaluator.
# sim_template is the ASTRA input template, here the parameters that are going
# to be varied need to be changed to the name given in var_1 etc. The format
# in the ASTRA input is e.g. Phi(1)={{RF_phase}}.
# analysis_func is the function that will analyze the output.
# sim_files contains the path to the particle distribution and other files
# needed for the ASTRA simulation like field maps.
# executable is the path to the ASTRA executable
ev = TemplateEvaluator(
sim_template="ASTRA_example.in",
analysis_func=analyze_simulation,
sim_files=[
"Example.ini",
"3_cell_L-Band.dat",
"Solenoid.dat",
],
executable="/path_to_ASTRA/Astra",
n_procs=1,
)
# Create exploration.
# max_evals is the maximum number of evaluations.
# max_evalvs / sim_worker is the number of simulation batches that are sent.
# sim_workers is the number of simulations that are launched in parallel.
# sim_workers should be smaller than the number of available CPU cores.
# In case you already have some data from an optimization run but would like
# to add further datapoints, set resume=True and max_evals to a higher number.
exp = Exploration(
generator=gen, evaluator=ev, max_evals=500, sim_workers=8, resume=False
)
# To safely perform exploration, run it in the block below (this is needed
# for some flavours of multiprocessing, namely spawn and forkserver)
if __name__ == "__main__":
exp.run()
&NEWRUN
Head=' Example of ASTRA users manual'
RUN=1
Distribution = 'Example.ini', Xoff=0.0, Yoff=0.0,
TRACK_ALL=T, Auto_phase=F
H_max=0.001, H_min=0.00
/
&OUTPUT
ZSTART=0.0, ZSTOP=2.5
Zemit=500, Zphase=1
RefS=T
EmitS=T, PhaseS=T
/
&CHARGE
LSPCH=F
Nrad=10, Cell_var=2.0, Nlong_in=10
min_grid=0.0
Max_Scale=0.05
/
&CAVITY
LEField=T,
File_Efield(1)='3_cell_L-Band.dat', C_pos(1)=0.3
Nue(1)=1.3, MaxE(1)=40.0, Phi(1)={{RF_phase}},
/
&SOLENOID
LBField=T,
File_Bfield(1)='Solenoid.dat', S_pos(1)=1.2
MaxB(1)={{B_sol}}, S_smooth(1)=10
/
"""Defines the analysis function that runs after the simulation.
This is an example calculation of the bunch length in µm, a combined
normalized transverse emittance, and the emittances in both transverse planes.
"""
import numpy as np
# Function to analyze the simulation result
def analyze_simulation(simulation_directory, output_params):
"""Analyze the simulation output.
This method analyzes the output generated by the simulation to
obtain the value of the optimization objective and other analyzed
parameters, if specified. The value of these parameters has to be
given to the `output_params` dictionary.
Parameters
----------
simulation_directory : str
Path to the simulation folder where the output was generated.
output_params : dict
Dictionary where the value of the objectives and analyzed parameters
will be stored. There is one entry per parameter, where the key
is the name of the parameter given by the user.
Returns
-------
dict
The `output_params` dictionary with the results from the analysis.
"""
try:
# Read back results from files
s, t, x_av, x_rms, xp_rms, em_n_x, x_xp = np.loadtxt(
simulation_directory + "/ASTRA_example.Xemit.001", unpack=True
)
s, t, y_av, y_rms, yp_rms, em_n_y, y_yp = np.loadtxt(
simulation_directory + "/ASTRA_example.Yemit.001", unpack=True
)
x, y, z, px, py, pz, t, charge, idx, flag = np.loadtxt(
simulation_directory + "/ASTRA_example.0250.001", unpack=True
)
z[1:] = z[1:] + z[0] # adding the position of the reference particle
output_params["bunch_length"] = np.std(z) * 1e6
output_params["emittance"] = np.log10(
em_n_x[-1] * em_n_y[-1] * 1e12
) # normalized emittances in µm, logarithm for better optimization
output_params["emittance_x"] = em_n_x[-1]
output_params["emittance_y"] = em_n_y[-1]
except Exception as exc:
logf = open("exception.log", "w")
logf.write(
"Failed to open or evaluate {0}: {1}\n".format(
str(simulation_directory), str(exc)
)
)
return output_params
# Not needed, for debugging only
if __name__ == "__main__":
analyze_simulation("path_to_simulation_result", {})