diff --git a/doc/source/examples/bo_with_hipace.rst b/doc/source/examples/bo_with_hipace.rst new file mode 100644 index 00000000..490b3eae --- /dev/null +++ b/doc/source/examples/bo_with_hipace.rst @@ -0,0 +1,72 @@ +.. _bo-with-hipace: + +Optimization with HiPACE++ +========================== + + +Description +~~~~~~~~~~~ + +This examples shows how to perform a Bayesian optimization of a PWFA using +`HiPACE++ `_. + +The setup is a simple driver-witness configuration where the witness is +optimized to maximize the objetive + +.. math:: + + f = \frac{\sqrt{Q} E_{MED}}{100 E_{MAD}} + + +where :math:`Q` is the beam charge, :math:`E_{MED}` is the median energy, and +:math:`E_{MAD}` is the median absolute deviation energy spread. The only +optimization parameter is the charge: + +- ``'witness_charge'``: parameter in the range :math:`[0.05, 1.]` in units of + :math:`\mathrm{nC}`. + +The optimization is carried out using an +:class:`~optimas.generators.AxSingleFidelityGenerator` and a +:class:`~optimas.evaluators.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 example is set up to make use of a system of 4 GPUs, where each HiPACE++ +simulation uses a 2 GPUs and 2 simulations are carried out in parallel. + +If HiPACE++ is installed in a different environment than ``optimas``, make +sure to specify ``env_script`` and ``env_mpi`` in the ``TemplateEvaluator``. + +Scripts +~~~~~~~ + +The files needed to run the optimization should be located in a folder +(named e.g., ``optimization``) with the following structure: + +.. code-block:: bash + + optimization + ├── run_example.py + ├── template_simulation_script.py + └── analysis_script.py + +The optimization is started by executing: + +.. code-block:: bash + + python run_example.py + +The scripts needed to run this example can be seen below. + +.. literalinclude:: ../../../examples/hipace/run_example.py + :language: python + :caption: run_example.py (:download:`download <../../../examples/hipace/run_example.py>`) + +.. literalinclude:: ../../../examples/hipace/template_simulation_script + :language: python + :caption: template_simulation_script.py (:download:`download <../../../examples/hipace/template_simulation_script>`) + +.. literalinclude:: ../../../examples/hipace/analysis_script.py + :language: python + :caption: analysis_script.py (:download:`download <../../../examples/hipace/analysis_script.py>`) diff --git a/doc/source/examples/index.rst b/doc/source/examples/index.rst index 37f66cbe..f498fc76 100644 --- a/doc/source/examples/index.rst +++ b/doc/source/examples/index.rst @@ -18,4 +18,5 @@ Examples bo_basic bo_with_fbpic + bo_with_hipace bo_multitask_fbpic_waket diff --git a/examples/hipace/analysis_script.py b/examples/hipace/analysis_script.py new file mode 100644 index 00000000..ca9eece2 --- /dev/null +++ b/examples/hipace/analysis_script.py @@ -0,0 +1,105 @@ +""" +Contains the function that analyzes the simulation results, +after the simulation was run. +""" +import os +from openpmd_viewer.addons import LpaDiagnostics +import numpy as np +from scipy.constants import e + + +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. + """ + # Open simulation diagnostics. + d = LpaDiagnostics(os.path.join(simulation_directory, 'diags/hdf5')) + + # Get beam particles with `u_z >= 10` and transverse offset no larger than + # 15 µm in `x` and `y`. + uz, w = d.get_particle(['uz', 'w'], iteration=d.iterations[-1], + species='witness') + + # Convert charge to pC. + q = w.sum()*e*1e12 + + # Analyze distribution and fill in the output data. + if len(w) < 2: # Need at least 2 particles to calculate energy spread + output_params['f'] = 0 + else: + med, mad = weighted_mad(uz/2, w) + output_params['f'] = np.sqrt(q)*med/mad/100 + output_params['charge'] = q + output_params['energy_med'] = med + output_params['energy_mad'] = mad + + return output_params + + +def weighted_median(data, weights): + """ + Compute the weighted quantile of a 1D numpy array. + Parameters + ---------- + data : ndarray + Input array (one dimension). + weights : ndarray + Array with the weights of the same size of `data`. + quantile : float + Quantile to compute. It must have a value between 0 and 1. + Returns + ------- + quantile_1D : float + The output value. + """ + quantile = .5 + # Check the data + if not isinstance(data, np.matrix): + data = np.asarray(data) + if not isinstance(weights, np.matrix): + weights = np.asarray(weights) + nd = data.ndim + if nd != 1: + raise TypeError("data must be a one dimensional array") + ndw = weights.ndim + if ndw != 1: + raise TypeError("weights must be a one dimensional array") + if data.shape != weights.shape: + raise TypeError("the length of data and weights must be the same") + if ((quantile > 1.) or (quantile < 0.)): + raise ValueError("quantile must have a value between 0. and 1.") + # Sort the data + ind_sorted = np.argsort(data) + sorted_data = data[ind_sorted] + sorted_weights = weights[ind_sorted] + # Compute the auxiliary arrays + Sn = np.cumsum(sorted_weights) + # TODO: Check that the weights do not sum zero + # assert Sn != 0, "The sum of the weights must not be zero" + Pn = (Sn-0.5*sorted_weights)/Sn[-1] + # Get the value of the weighted median + return np.interp(quantile, Pn, sorted_data) + + +def weighted_mad(x, w): + med = weighted_median(x, w) + mad = weighted_median(np.abs(x-med), w) + return med, mad diff --git a/examples/hipace/run_example.py b/examples/hipace/run_example.py new file mode 100644 index 00000000..5ec594bf --- /dev/null +++ b/examples/hipace/run_example.py @@ -0,0 +1,66 @@ +""" +This example optimizes a PWFA stage using HiPACE++. + +The HiPACE++ simulations are performed using the template defined in the +`template_simulation_script` file. + +In addition to the objective `f`, three additional parameters +are analyzed for each simulation and including in the optimization +history. The calculation of `f` and the additional parameters is performed +in the `analyze_simulation` function, which for convenience is here defined in +the `analysis_script.py` file. +""" +from optimas.core import Parameter, VaryingParameter, Objective +from optimas.generators import AxSingleFidelityGenerator +from optimas.evaluators import TemplateEvaluator +from optimas.explorations import Exploration + +from analysis_script import analyze_simulation + + +# Create varying parameters and objectives. +var_1 = VaryingParameter('witness_charge', 0.05, 1.) +obj = Objective('f', minimize=False) + + +# Define additional parameters to analyze. +energy_med = Parameter('energy_med') +energy_mad = Parameter('energy_mad') +charge = Parameter('charge') + + +# Create generator. +gen = AxSingleFidelityGenerator( + varying_parameters=[var_1], + objectives=[obj], + analyzed_parameters=[energy_med, energy_mad, charge], + n_init=4 +) + + +# Create evaluator. +ev = TemplateEvaluator( + sim_template='template_simulation_script', + analysis_func=analyze_simulation, + executable='/path/to/build/bin/hipace', + n_gpus=2, # Use 2 GPUs per simulation. + # Uncomment if HiPACE is installed in a different environment than optimas. + # env_script='/path/to/profile.hipace', + # Uncomment if `env_script` loads a different MPI to that used by optimas. + # env_mpi='openmpi' +) + + +# Create exploration. +exp = Exploration( + generator=gen, + evaluator=ev, + max_evals=100, + sim_workers=2 +) + + +# 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() diff --git a/examples/hipace/template_simulation_script b/examples/hipace/template_simulation_script new file mode 100644 index 00000000..8e5718a6 --- /dev/null +++ b/examples/hipace/template_simulation_script @@ -0,0 +1,53 @@ +max_step = 300 +amr.n_cell = 256 256 256 + +amr.max_level = 0 + +hipace.max_time = 0.3/clight +diagnostic.output_period = 30 +hipace.verbose = 1 + +hipace.numprocs_x = 1 +hipace.numprocs_y = 1 + +hipace.depos_order_xy = 2 +hipace.dt = adaptive +hipace.nt_per_betatron = 30 + +geometry.is_periodic = true true false # Is periodic? +geometry.prob_lo = -250.e-6 -250.e-6 -250.e-6 # physical domain +geometry.prob_hi = 250.e-6 250.e-6 110.e-6 + +beams.names = driver witness + +driver.position_mean = 0. 0. 0. +driver.position_std = 2.e-6 2.e-6 30.e-6 +driver.injection_type = fixed_weight +driver.num_particles = 1000000 +driver.total_charge = .6e-9 +driver.u_mean = 0. 0. 1000. +driver.u_std = 2. 2. 10. +driver.do_symmetrize = 1 + +witness.position_mean = 0. 0. -160.e-6 +witness.position_std = 2.e-6 2.e-6 5.e-6 +witness.injection_type = fixed_weight +witness.num_particles = 1000000 +witness.total_charge = {{witness_charge}}e-9 +witness.u_mean = 0. 0. 1000. +witness.u_std = 2. 2. 10. +witness.do_symmetrize = 1 + +plasmas.names = electron ion + +electron.density(x,y,z) = 2.e22 +electron.ppc = 1 1 +electron.u_mean = 0.0 0.0 0. +electron.element = electron + +ion.density(x,y,z) = 2.e22 +ion.ppc = 1 1 +ion.u_mean = 0.0 0.0 0. +ion.element = H + +diagnostic.diag_type = xz diff --git a/optimas/evaluators/template_evaluator.py b/optimas/evaluators/template_evaluator.py index 9c362b54..f5103c25 100644 --- a/optimas/evaluators/template_evaluator.py +++ b/optimas/evaluators/template_evaluator.py @@ -35,6 +35,15 @@ class TemplateEvaluator(Evaluator): n_gpus : int, optional The number of GPUs that will be made available for each evaluation. By default, 0. + env_script : str, optional + The full path of a shell script to set up the environment for the + launched simulation. This is useful when the simulation needs to run + in a different environment than optimas. The script should start with a + shebang. + env_mpi : str, optional + If the `env_script` loads an MPI different than the one in the optimas + environment, indicate it here. Possible values are "mpich", "openmpi", + "aprun", "srun", "jsrun", "msmpi". """ def __init__( self, @@ -43,7 +52,9 @@ def __init__( executable: Optional[str] = None, sim_files: Optional[List[str]] = None, n_procs: Optional[int] = None, - n_gpus: Optional[int] = None + n_gpus: Optional[int] = None, + env_script: Optional[str] = None, + env_mpi: Optional[str] = None, ) -> None: super().__init__( sim_function=run_template_simulation, @@ -53,6 +64,8 @@ def __init__( self.sim_template = sim_template self.analysis_func = analysis_func self.executable = executable + self.env_script = env_script + self.env_mpi = env_mpi self.sim_files = [] if sim_files is None else sim_files self._app_name = 'sim' @@ -83,6 +96,8 @@ def get_sim_specs( sim_specs['user']['analysis_func'] = self.analysis_func sim_specs['user']['sim_template'] = os.path.basename(self.sim_template) sim_specs['user']['app_name'] = self._app_name + sim_specs['user']['env_script'] = self.env_script + sim_specs['user']['env_mpi'] = self.env_mpi return sim_specs def get_libe_specs(self) -> Dict: @@ -119,8 +134,7 @@ def _register_app(self) -> None: 'An executable must be provided for non-Python simulations') assert os.path.exists(self.executable), ( 'Executable {} does not exist.'.format(self.executable)) - executable_path = './' + self.executable - self.sim_files.append(self.executable) + executable_path = os.path.abspath(self.executable) # Register app. Executor.executor.register_app( diff --git a/optimas/sim_functions.py b/optimas/sim_functions.py index 4ec180e7..9023e7d5 100644 --- a/optimas/sim_functions.py +++ b/optimas/sim_functions.py @@ -1,4 +1,5 @@ import os + import jinja2 import numpy as np @@ -52,11 +53,15 @@ def run_template_simulation(H, persis_info, sim_specs, libE_info): # Passed to command line in addition to the executable. exctr = Executor.executor # Get Executor - task = exctr.submit(app_name=app_name, - app_args=sim_script, - stdout='out.txt', - stderr='err.txt', - ) + # Launch simulation. + task = exctr.submit( + app_name=app_name, + app_args=sim_script, + stdout='out.txt', + stderr='err.txt', + env_script=user_specs['env_script'], + mpi_runner_type=user_specs['env_mpi'] + ) # Wait for simulation to complete task.wait() diff --git a/tests/resources/env_script.sh b/tests/resources/env_script.sh new file mode 100644 index 00000000..29b00642 --- /dev/null +++ b/tests/resources/env_script.sh @@ -0,0 +1,2 @@ +#!/usr/bin/env bash +export LIBE_TEST_SUB_ENV_VAR="testvalue" diff --git a/tests/resources/template_simulation_script.py b/tests/resources/template_simulation_script.py index 02304572..5661ecf2 100644 --- a/tests/resources/template_simulation_script.py +++ b/tests/resources/template_simulation_script.py @@ -2,10 +2,16 @@ Dummy simulation template used for testing. It takes x0 and x1 as input parameters and stores the result in `result.txt`. """ +import os import numpy as np +test_env_var = os.getenv('LIBE_TEST_SUB_ENV_VAR') + # 2D function with multiple minima result = -( {{x0}} + 10*np.cos({{x0}}) )*( {{x1}} + 5*np.cos({{x1}}) ) with open('result.txt', 'w') as f: - f.write("%f" %result) + output = [str(result) + '\n'] + if test_env_var is not None: + output.append(test_env_var) + f.writelines(output) diff --git a/tests/test_env_script.py b/tests/test_env_script.py new file mode 100644 index 00000000..561ecde6 --- /dev/null +++ b/tests/test_env_script.py @@ -0,0 +1,67 @@ +import os + +import numpy as np + +from optimas.explorations import Exploration +from optimas.generators import RandomSamplingGenerator +from optimas.evaluators import TemplateEvaluator +from optimas.core import VaryingParameter, Objective, Parameter + + +def analysis_func(sim_dir, output_params): + """Analysis function used by the template evaluator.""" + # Read back result from file + with open('result.txt') as f: + result = f.readlines() + f = float(result[0]) + test_var = result[1] + output_params['f'] = f + output_params['test_var'] = test_var + + +def test_env_script(): + # Define variables and objectives. + var1 = VaryingParameter('x0', -50., 5.) + var2 = VaryingParameter('x1', -5., 15.) + obj = Objective('f', minimize=False) + test_var = Parameter('test_var', dtype='U10') + + # Define variables and objectives. + gen = RandomSamplingGenerator( + varying_parameters=[var1, var2], + objectives=[obj], + analyzed_parameters=[test_var] + ) + + # Create template evaluator. + ev = TemplateEvaluator( + sim_template=os.path.join( + os.path.abspath(os.path.dirname(__file__)), + 'resources', + 'template_simulation_script.py' + ), + analysis_func=analysis_func, + env_script=os.path.join( + os.path.abspath(os.path.dirname(__file__)), + 'resources', + 'env_script.sh' + ), + ) + + # Create exploration. + exploration = Exploration( + generator=gen, + evaluator=ev, + max_evals=10, + sim_workers=2, + exploration_dir_path='./tests_output/test_env_script' + ) + + # Run exploration. + exploration.run() + + assert np.all(exploration.history['test_var'] == 'testvalue') + + +if __name__ == '__main__': + test_env_script()