ensemblesweep is a small library based on libEnsemble
for parallel parameter sweeps of objective functions or executables.
In the project directory:
pip install .-
Specify any number of iterable inputs to a
Dataobject. -
Specify an objective function to the
Sweepobject. Provide the input data. -
Run the sweep. All instances are run in parallel, up to the number of available cores.
import numpy as np
from ensemblesweep import Sweep, Data
# Define the function to be evaluated
def my_function(core, edge, x):
return float(edge) * np.sum(x) + core
data = Data(
core=random.sample(["rocm", "cuda", "cpu"]*6, 6),
edge=random.sample(range(10), 4),
x = np.random.uniform(lb, ub, (3, 3))
)
# Create a sweep object
sweep = Sweep(
objective_function=my_function,
input_data=data,
)
if __name__ == "__main__":
sweep.run()
print(sweep.results)- Specify an executable to the
Sweepobject.- The
Dataparameters will be passed as arguments to the executable in the order they are defined. - The output of the executable will be read:
- from the file specified by
objective_output(if provided). - or, if not provided, the last line of the executable's stdout.
- from the file specified by
- The
import random
from ensemblesweep import Sweep, Data
# an objective can also be an *executable*
objective_path = os.path.join(os.path.dirname(__file__), "forces.x")
data = Data(
num_particles=random.sample(range(100, 1000, 10), 10),
num_steps = [10, 100, 1000],
rand_seed = 100,
kill_rate = [0.01, 0.1, 0.5]
)
# each of the samples will be given to the executable as arguments, in the exact order. Parallel runs.
sweep = Sweep(
objective_executable=objective_path,
objective_output="forces.stat",
input_data=data,
)
sweep.run()Sweep a Python function:
ensemblesweep py --func mod.func --var "x=0:1:10" --var "y=1,2" --workers 4Sweep an executable:
ensemblesweep exe --app ./sim.x --var "m=1,2,3" --out-file out.statA concurrent.futures-style API. A libEnsemble run blocks, so submit_sweep returns one future for a whole batch. Individual parameter points are still evaluated concurrently.
from ensemblesweep import Data, SweepExecutor
data = Data(x=[1, 2, 3], y=[10, 20])
def my_function(x, y):
return x * y
with SweepExecutor(nworkers=4) as executor:
future = executor.submit_sweep(
objective_function=my_function,
input_data=data,
)
batch = future.result()
for result in batch:
print(result.params, result.outputs, result.eval_time)You can also submit an existing Sweep object, including a partial batch:
from ensemblesweep import Sweep, SweepExecutor
sweep = Sweep(objective_function=my_function, input_data=data)
with SweepExecutor(nworkers=4) as executor:
future = executor.submit(sweep, n=10)
batch = future.result()SweepResults is a collection of SweepResult objects.
Each SweepResult has:
indexparams(the input parameters)outputs(the objective function's return values)eval_timestatus
Both the simple Sweep API and the concurrent.futures API return SweepResults:
# Simple API
sweep = Sweep(objective_function=my_function, input_data=data)
sweep.run()
results = sweep.results # SweepResults (all completed)
# Executor API
with SweepExecutor(nworkers=4) as executor:
future = executor.submit(sweep, n=10)
results = future.result() # SweepResults (this batch only)Iterate over SweepResult objects:
for result in results:
print(result.index, result.params, result.outputs, result.eval_time)Index into results:
print(results[0]) # single SweepResult
print(results[:3]) # list of SweepResultExport:
results.to_numpy() # raw libEnsemble array
results.to_pandas() # DataFrame with index + params + outputsThe SweepResult dataclass shape:
SweepResult(
index=0,
params={"x": 1, "y": 10},
outputs={"output": 10.0},
eval_time=2.8e-06,
status="completed",
)- Run a subset of the total points. This also updates the estimated time (in seconds) to run the entire sweep.
if __name__ == "__main__":
sweep.run(10)
# estimated time is in seconds, for the entire sweep
if sweep.estimated_time() > 10:
print("Estimated time is too long, only doing a little bit")
sweep.run(10) # run 10 concurrently. Max concurrency is number of cores
else:
print("Estimated time is acceptable, continuing")
sweep.run()- Estimate the runtime for a subset of the total points. Parallelism is considered.
sweep.estimated_time(4)