Doc: optimas Usage

Document how to use `optimas` with ImpactX.

examples/optimize_triplet/README.rst

@@ -56,6 +56,19 @@ Run
          :language: python3
          :caption: You can copy this file from ``tests/python/test_xopt.py``.
 
+   .. tab-item:: optimas
+
+      This example uses `optimas <https://optimas.readthedocs.io>`__ (method: `Bayesian Optimization <https://optimas.readthedocs.io/en/latest/api/_autosummary/optimas.generators.AxSingleFidelityGenerator.html>`__) to find the quadrupole strengths by minimizing the objective.
+
+      Machine-learning based, surrogate optimization like Bayesian Optimization (BO) works well for highly dimensional inputs and/or to find global minima in an objective that has potentially many local minima, where conventional optimizers can get stuck.
+      At the same time, the BO is prone to over-explore an objective (at the cost of finding a point closer to the global minima).
+
+      This example can be run as a **Python** script: ``python3 tests/python/test_optimas.py``.
+
+      .. literalinclude:: test_optimas.py
+         :language: python3
+         :caption: You can copy this file from ``tests/python/test_optimas.py``.
+
 
 Analyze
 -------

examples/optimize_triplet/test_optimas.py

@@ -0,0 +1 @@
+../../tests/python/test_optimas.py

tests/python/test_optimas.py

@@ -0,0 +1,261 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 The ImpactX Community
+#
+# Authors: Axel Huebl
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import importlib
+
+import amrex.space3d as amr
+import impactx
+import numpy as np
+import pytest
+from impactx import ImpactX, distribution, elements
+
+# configure the test
+verbose = True
+max_steps = 60
+
+
+def build_lattice(parameters: dict, write_particles: bool) -> list:
+    """
+    Create the quadrupole triplet.
+
+    Parameters
+    ----------
+    parameters: dict
+      quadrupole strengths k of quad 1/3 and quad 2.
+
+    write_particles: bool
+      write the particles in a beam monitor at the beginning and
+      end of the simulation
+
+    Returns
+    -------
+    A lattice for ImpactX: a list of impactx.elements.
+    """
+    q1_k, q2_k = parameters["q1_k"], parameters["q2_k"]
+
+    ns = 10  # number of slices per ds in the element
+
+    # enforce a mirror symmetry of the triplet
+    line = [
+        elements.Drift(ds=2.7, nslice=ns),
+        elements.Quad(ds=0.1, k=q1_k, nslice=ns),
+        elements.Drift(ds=1.4, nslice=ns),
+        elements.Quad(ds=0.2, k=q2_k, nslice=ns),
+        elements.Drift(ds=1.4, nslice=ns),
+        elements.Quad(ds=0.1, k=q1_k, nslice=ns),
+        elements.Drift(ds=2.7, nslice=ns),
+    ]
+
+    if write_particles:
+        monitor = elements.BeamMonitor("monitor", backend="h5")
+        line = [monitor] + line + [monitor]
+
+    return line
+
+
+def run(parameters: dict, write_particles=False, write_reduced=False) -> dict:
+    """
+    Run an ImpactX simulation with a new set of lattice parameters.
+
+    Parameters
+    ----------
+    parameters: dict
+      quadrupole strengths k of quad 1/3 and quad 2.
+
+    write_particles: bool
+      write the particles in a beam monitor at the beginning and
+      end of the simulation
+
+    write_reduced: bool
+      write the reduced diagnositcs of ImpactX to a file.
+
+    Returns
+    -------
+    A dictionary with reduced diagnositcs of ImpactX, characterizing
+    the beam at the end of the simulation.
+    """
+    pp_amrex = amr.ParmParse("amrex")
+    pp_amrex.add("verbose", 0)
+
+    sim = ImpactX()
+
+    sim.verbose = 0
+
+    # set numerical parameters and IO control
+    sim.particle_shape = 2  # B-spline order
+    sim.space_charge = False
+    sim.diagnostics = write_reduced
+    sim.slice_step_diagnostics = write_reduced
+
+    # domain decomposition & space charge mesh
+    sim.init_grids()
+
+    # load a 2 GeV electron beam with an initial
+    # unnormalized rms emittance of 5 nm
+    kin_energy_MeV = 2.0e3  # reference energy
+    bunch_charge_C = 100.0e-12  # used with space charge
+    npart = 10000  # number of macro particles
+
+    #   reference particle
+    ref = sim.particle_container().ref_particle()
+    ref.set_charge_qe(1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+    #   particle bunch
+    distr = distribution.Waterbag(
+        lambdaX=2.0e-4,
+        lambdaY=2.0e-4,
+        lambdaT=3.1622776602e-5,
+        lambdaPx=1.1180339887e-5,
+        lambdaPy=1.1180339887e-5,
+        lambdaPt=3.1622776602e-5,
+        muxpx=0.894427190999916,
+        muypy=-0.894427190999916,
+        mutpt=0.0,
+    )
+    sim.add_particles(bunch_charge_C, distr, npart)
+
+    # design the accelerator lattice
+    sim.lattice.extend(build_lattice(parameters, write_particles=write_particles))
+
+    # run simulation
+    sim.evolve()
+
+    # in situ calculate the reduced beam characteristics
+    beam = sim.particle_container()
+    rbc = beam.reduced_beam_characteristics()
+
+    # clean shutdown
+    sim.finalize()
+
+    return rbc
+
+
+def analyze_simulation(simulation_directory, output_params) -> dict:
+    """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.
+        This is the L2 norm of alpha and beta of the beam at the end of the
+    simulation.
+    """
+    if verbose:
+        print(f"Analyze simulation in {simulation_directory}")
+
+    output_params = ...
+    alpha_x, alpha_y, beta_x, beta_y = (
+        output_params["alpha_x"],
+        output_params["alpha_y"],
+        output_params["beta_x"],
+        output_params["beta_y"],
+    )
+    if verbose:
+        print(
+            f"  -> alpha_x={alpha_x}, alpha_y={alpha_y}, beta_x={beta_x}, beta_y={beta_y}"
+        )
+    alpha_beta_is = np.array([alpha_x, alpha_y, beta_x, beta_y])
+
+    beta_x_goal = 0.55
+    beta_y_goal = beta_x_goal
+    alpha_beta_goal = np.array([0, 0, beta_x_goal, beta_y_goal])
+
+    error = np.sum((alpha_beta_is - alpha_beta_goal) ** 2)
+
+    output_params["error"] = error
+
+    return output_params
+
+
+@pytest.mark.skipif(
+    importlib.util.find_spec("optimas") is None, reason="optimas is not available"
+)
+def test_optimas():
+    from optimas.core import Objective, Parameter, VaryingParameter
+    from optimas.evaluators import TemplateEvaluator
+    from optimas.explorations import Exploration
+    from optimas.generators import AxSingleFidelityGenerator
+
+    # Create varying parameters and objectives
+    var_q1_k = VaryingParameter("q1_k", -10, 0)
+    var_q2_k = VaryingParameter("q2_k", 0, 10)
+    obj = Objective("f", minimize=False)
+
+    # Define additional parameters to analyze
+    alpha_x = Parameter("alpha_x")
+    alpha_y = Parameter("alpha_y")
+    beta_x = Parameter("beta_x")
+    beta_y = Parameter("beta_y")
+
+    # Create generator
+    gen = AxSingleFidelityGenerator(
+        varying_parameters=[var_q1_k, var_q2_k],
+        objectives=[obj],
+        analyzed_parameters=[alpha_x, alpha_y, beta_x, beta_y],
+        n_init=4,
+    )
+
+    # Create evaluator
+    ev = TemplateEvaluator(
+        sim_template="template_simulation_script.py",  # TODO
+        analysis_func=analyze_simulation,
+    )
+
+    # Create exploration
+    exp = Exploration(generator=gen, evaluator=ev, max_evals=max_steps, sim_workers=1)
+
+    # run exploration
+    exp.run()
+
+    # Print all trials
+    if verbose:
+        print(exp)
+
+    # Select the best result
+    # ... TODO ...
+
+    # Print the optimization result # ... TODO ...
+    # print("Optimal parameters for k:", best_ks)
+    # print("L2 norm of alpha & beta at the optimum:", best_run["error"].values[0])
+
+    # analytical result:
+    #   k: -3.5, 2.75
+    #   alpha & beta: 0, 0, 0.55, 0.55
+
+    # final run w/ detailed I/O on # ... TODO ...
+    # rbc = run(best_ks, write_particles=True, write_reduced=True)
+    # alpha_x, alpha_y, beta_x, beta_y = (
+    #     rbc["alpha_x"],
+    #     rbc["alpha_y"],
+    #     rbc["beta_x"],
+    #     rbc["beta_y"],
+    # )
+    # print(f"alpha_x={alpha_x} alpha_y={alpha_y}\n beta_x={beta_x}     beta_y={beta_y}")
+
+
+if __name__ == "__main__":
+    # Call MPI_Init and MPI_Finalize only once:
+    if impactx.Config.have_mpi:
+        from mpi4py import MPI  # noqa
+
+    test_optimas()