Allow launching simulations in a different environment

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++ <https://github.com/Hi-PACE/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>`)

doc/source/examples/index.rst

@@ -18,4 +18,5 @@ Examples
 
    bo_basic
    bo_with_fbpic
+   bo_with_hipace
    bo_multitask_fbpic_waket

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

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()

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

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(

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()

tests/resources/env_script.sh

@@ -0,0 +1,2 @@
+#!/usr/bin/env bash
+export LIBE_TEST_SUB_ENV_VAR="testvalue"