Simulators

WarpX

warpx_simf.run_warpx(H, persis_info, sim_specs, libE_info)

This function runs a WarpX simulation and returns quantity ‘f’ as well as other physical quantities measured in the run for convenience. Status check is done periodically on the simulation, provided by LibEnsemble.

warpx_simf.py
 1import os
 2import time
 3import numpy as np
 4
 5from libensemble.executors.executor import Executor
 6from libensemble.message_numbers import WORKER_DONE, TASK_FAILED
 7from read_sim_output import read_sim_output
 8from write_sim_input import write_sim_input
 9
10"""
11This file is part of the suite of scripts to use LibEnsemble on top of WarpX
12simulations. It defines a sim_f function that takes LibEnsemble history and
13input parameters, run a WarpX simulation and returns 'f'.
14"""
15
16
17def run_warpx(H, persis_info, sim_specs, libE_info):
18    """
19    This function runs a WarpX simulation and returns quantity 'f' as well as
20    other physical quantities measured in the run for convenience. Status check
21    is done periodically on the simulation, provided by LibEnsemble.
22    """
23
24    # Setting up variables needed for input and output
25    # keys              = variable names
26    # x                 = variable values
27    # libE_output       = what will be returned to libE
28
29    input_file = sim_specs["user"]["input_filename"]
30    time_limit = sim_specs["user"]["sim_kill_minutes"] * 60.0
31    machine_specs = sim_specs["user"]["machine_specs"]
32
33    exctr = Executor.executor  # Get Executor
34
35    # Modify WarpX input file with input parameters calculated by gen_f
36    # and passed to this sim_f.
37    write_sim_input(input_file, H["x"])
38
39    # Passed to command line in addition to the executable.
40    # Here, only input file
41    app_args = input_file
42    os.environ["OMP_NUM_THREADS"] = machine_specs["OMP_NUM_THREADS"]
43
44    # Launch the executor to actually run the WarpX simulation
45
46    use_gpus = machine_specs["name"] == "polaris"
47
48    task = exctr.submit(
49        app_name="warpx",
50        num_procs=machine_specs["cores"],
51        auto_assign_gpus=use_gpus,
52        match_procs_to_gpus=use_gpus,
53        app_args=app_args,
54        stdout="out.txt",
55        stderr="err.txt",
56        wait_on_start=True,
57    )
58
59    # Periodically check the status of the simulation
60    calc_status = exctr.polling_loop(task)
61
62    # Safety
63    time.sleep(0.2)
64
65    # Get output from a run and delete output files
66    warpx_out = read_sim_output(task.workdir)
67
68    # Excluding results - NAN - from runs where beam was lost
69    if warpx_out[0] != warpx_out[0]:
70        print(task.workdir, " output led to NAN values")
71
72    # Pass the sim output values to LibEnsemble.
73    # When optimization is ON, 'f' is then passed to the generating function
74    # gen_f to generate new inputs for next runs.
75    # All other parameters are here just for convenience.
76    libE_output = np.zeros(1, dtype=sim_specs["out"])
77    libE_output["f"] = warpx_out[0]
78    libE_output["energy_std"] = warpx_out[1]
79    libE_output["energy_avg"] = warpx_out[2]
80    libE_output["charge"] = warpx_out[3]
81    libE_output["emittance"] = warpx_out[4]
82    libE_output["ramp_down_1"] = H["x"][0][0]
83    libE_output["ramp_down_2"] = H["x"][0][1]
84    libE_output["zlens_1"] = H["x"][0][2]
85    libE_output["adjust_factor"] = H["x"][0][3]
86
87    return libE_output, persis_info, calc_status
Example usage
  1#!/usr/bin/env python
  2
  3"""
  4This file is part of the suite of scripts to use LibEnsemble on top of WarpX
  5simulations. It is the entry point script that runs LibEnsemble. Libensemble
  6then launches WarpX simulations.
  7
  8Execute locally via the following command:
  9    python run_libensemble_on_warpx.py --comms local --nworkers 3
 10On summit, use the submission script:
 11    bsub summit_submit_mproc.sh
 12
 13The number of concurrent evaluations of the objective function will be
 14nworkers=1 as one worker is for the persistent gen_f.
 15"""
 16
 17# Either 'random' or 'aposmm'
 18generator_type = "aposmm"
 19# Either 'local' or 'swing'
 20machine = "local"
 21
 22import sys
 23
 24import numpy as np
 25from warpx_simf import run_warpx  # Sim function from current directory
 26
 27# Import libEnsemble modules
 28from libensemble.libE import libE
 29
 30if generator_type == "random":
 31    from libensemble.alloc_funcs.give_sim_work_first import (
 32        give_sim_work_first as alloc_f,
 33    )
 34    from libensemble.gen_funcs.sampling import uniform_random_sample as gen_f
 35elif generator_type == "aposmm":
 36    import libensemble.gen_funcs
 37
 38    libensemble.gen_funcs.rc.aposmm_optimizers = "nlopt"
 39    from libensemble.alloc_funcs.persistent_aposmm_alloc import (
 40        persistent_aposmm_alloc as alloc_f,
 41    )
 42    from libensemble.gen_funcs.persistent_aposmm import aposmm as gen_f
 43else:
 44    print("you shouldn' hit that")
 45    sys.exit()
 46
 47import all_machine_specs
 48
 49from libensemble import logger
 50from libensemble.executors.mpi_executor import MPIExecutor
 51from libensemble.tools import add_unique_random_streams, parse_args, save_libE_output
 52
 53# Import machine-specific run parameters
 54if machine == "local":
 55    machine_specs = all_machine_specs.local_specs
 56elif machine == "polaris":
 57    machine_specs = all_machine_specs.polaris_specs
 58else:
 59    print("you shouldn' hit that")
 60    sys.exit()
 61
 62logger.set_level("INFO")
 63
 64nworkers, is_manager, libE_specs, _ = parse_args()
 65
 66# Set to full path of warp executable
 67sim_app = machine_specs["sim_app"]
 68
 69# Problem dimension. This is the number of input parameters exposed,
 70# that LibEnsemble will vary in order to minimize a single output parameter.
 71n = 4
 72
 73exctr = MPIExecutor()
 74exctr.register_app(full_path=sim_app, app_name="warpx")
 75
 76# State the objective function, its arguments, output, and necessary parameters
 77# (and their sizes). Here, the 'user' field is for the user's (in this case,
 78# the simulation) convenience. Feel free to use it to pass number of nodes,
 79# number of ranks per note, time limit per simulation etc.
 80sim_specs = {
 81    # Function whose output is being minimized. The parallel WarpX run is
 82    # launched from run_WarpX.
 83    "sim_f": run_warpx,
 84    # Name of input for sim_f, that LibEnsemble is allowed to modify.
 85    # May be a 1D array.
 86    "in": ["x"],
 87    "out": [
 88        # f is the single float output that LibEnsemble minimizes.
 89        ("f", float),
 90        # All parameters below are not used for calculation,
 91        # just output for convenience.
 92        # Final relative energy spread.
 93        ("energy_std", float, (1,)),
 94        # Final average energy, in MeV.
 95        ("energy_avg", float, (1,)),
 96        # Final beam charge.
 97        ("charge", float, (1,)),
 98        # Final beam emittance.
 99        ("emittance", float, (1,)),
100        # input parameter: length of first downramp.
101        ("ramp_down_1", float, (1,)),
102        # input parameter: Length of second downramp.
103        ("ramp_down_2", float, (1,)),
104        # input parameter: position of the focusing lens.
105        ("zlens_1", float, (1,)),
106        # Relative strength of the lens (1. is from
107        # back-of-the-envelope calculation)
108        ("adjust_factor", float, (1,)),
109    ],
110    "user": {
111        # name of input file
112        "input_filename": "inputs",
113        # Run timeouts after 3 mins
114        "sim_kill_minutes": 3,
115        # machine-specific parameters
116        "machine_specs": machine_specs,
117    },
118}
119
120# State the generating function, its arguments, output,
121# and necessary parameters.
122if generator_type == "random":
123    # Here, the 'user' field is for the user's (in this case,
124    # the RNG) convenience.
125    gen_specs = {
126        # Generator function. Will randomly generate new sim inputs 'x'.
127        "gen_f": gen_f,
128        # Generator input. This is a RNG, no need for inputs.
129        "in": [],
130        "out": [
131            # parameters to input into the simulation.
132            ("x", float, (n,))
133        ],
134        "user": {
135            # Total max number of sims running concurrently.
136            "gen_batch_size": nworkers,
137            # Lower bound for the n parameters.
138            "lb": np.array([2.0e-3, 2.0e-3, 0.005, 0.1]),
139            # Upper bound for the n parameters.
140            "ub": np.array([2.0e-2, 2.0e-2, 0.028, 3.0]),
141        },
142    }
143
144    alloc_specs = {
145        # Allocator function, decides what a worker should do.
146        # We use a LibEnsemble allocator.
147        "alloc_f": alloc_f,
148        "user": {
149            # If true wait for all sims to process before generate more
150            "batch_mode": True,
151            # Only one active generator at a time
152            "num_active_gens": 1,
153        },
154    }
155
156elif generator_type == "aposmm":
157    # Here, the 'user' field is for the user's (in this case,
158    # the optimizer) convenience.
159    gen_specs = {
160        # Generator function. Will randomly generate new sim inputs 'x'.
161        "gen_f": gen_f,
162        "persis_in": ["f", "x", "x_on_cube", "sim_id", "local_min", "local_pt"],
163        "out": [
164            # parameters to input into the simulation.
165            ("x", float, (n,)),
166            # x scaled to a unique cube.
167            ("x_on_cube", float, (n,)),
168            # unique ID of simulation.
169            ("sim_id", int),
170            # Whether this point is a local minimum.
171            ("local_min", bool),
172            # whether the point is from a local optimization run
173            # or a random sample point.
174            ("local_pt", bool),
175        ],
176        "user": {
177            # Number of sims for initial random sampling.
178            # Optimizer starts afterwards.
179            "initial_sample_size": max(nworkers - 1, 1),
180            # APOSMM/NLOPT optimization method
181            "localopt_method": "LN_BOBYQA",
182            "num_pts_first_pass": nworkers,
183            # Relative tolerance of inputs
184            "xtol_rel": 1e-3,
185            # Absolute tolerance of output 'f'. Determines when
186            # local optimization stops.
187            "ftol_abs": 3e-8,
188            # Lower bound for the n input parameters.
189            "lb": np.array([2.0e-3, 2.0e-3, 0.005, 0.1]),
190            # Upper bound for the n input parameters.
191            "ub": np.array([2.0e-2, 2.0e-2, 0.028, 3.0]),
192        },
193    }
194
195    alloc_specs = {"alloc_f": alloc_f}
196
197else:
198    print("you shouldn' hit that")
199    sys.exit()
200
201# Save H to file every N simulation evaluations
202libE_specs["save_every_k_sims"] = 100
203# Sim directory to be copied for each worker
204libE_specs["sim_input_dir"] = "sim"
205libE_specs["sim_dirs_make"] = True
206libE_specs["dedicated_mode"] = True
207
208sim_max = machine_specs["sim_max"]  # Maximum number of simulations
209exit_criteria = {"sim_max": sim_max}  # Exit after running sim_max simulations
210
211# Create a different random number stream for each worker and the manager
212persis_info = add_unique_random_streams({}, nworkers + 1)
213
214if __name__ == "__main__":
215
216    # Run LibEnsemble, and store results in history array H
217    H, persis_info, flag = libE(
218        sim_specs, gen_specs, exit_criteria, persis_info, alloc_specs, libE_specs
219    )
220
221    # Save results to numpy file
222    if is_manager:
223        save_libE_output(H, persis_info, __file__, nworkers)