SweepExpMPI

class sweepexp.sweepexp_mpi.SweepExpMPI(func: Callable, parameters: dict[str, list], save_path: Path | str | None = None, **kwargs: any)

Bases: SweepExp

Run a parameter sweep in parallel using MPI.

Parameters

funcCallable

The experiment function to run. The function should take the parameters as keyword arguments and return a dictionary with the return values.

parametersdict[str, list]

The parameters to sweep over. The keys are the parameter names and the values are lists of the parameter values.

save_pathPath | str | None

The path to save the results to. Supported file formats are: ‘.zarr’, ‘.nc’, ‘.cdf’, ‘.pkl’. The ‘.zarr’ and ‘.nc’ formats only support numeric and boolean data. Only the ‘.pkl’ format supports saving data of any type.

Description

The SweepExpMPI class can be used to run a custom experiment function with different parameter combinations. The results of the experiments are saved as an xarray dataset. The dataset can be saved to a file and loaded later to continue the experiments. All parameter combinations are run in parallel using mpi4py.

SweepExp supports the following additional features: - Custom arguments: Add custom arguments to the experiment function. - UUID: Pass a unique identifier to the experiment function. - Auto save: Automatically save the results after each experiment. - Timeit: Measure the duration of each experiment. - Priorities: Run experiments with higher priority first.

Examples

my_experiment.py

from sweepexp import SweepExpParallel

# Create a simple experiment function
def my_experiment(x: int, y: float) -> dict:
    return {"sum": x + y, "product": x * y}

# Initialize the SweepExp object
sweep = SweepExpParallel(
    func=my_experiment,
    parameters={"x": [1, 2, 3], "y": [4, 5, 6]},
)

# Run the sweep
sweep.run()

To run the experiment on 4 CPUs using MPI, use the following command:

mpiexec -n 4 python3 my_experiment.py

Or, alternatively any other MPI launcher, e.g., mpirun, srun, etc.

__init__(func: Callable, parameters: dict[str, list], save_path: Path | str | None = None, **kwargs: any) → None

Methods

__init__(func, parameters[, save_path])
add_custom_argument(name, default_value)	Add custom arguments to the experiment function.
load(save_path)	Load the xarray dataset from a file.
reset_status(states)	Reset the status of experiments to 'N' (not started).
run([status, max_workers])	Run all experiments with the status 'N' (not started).
save([mode])	Save the xarray dataset to the save path.

Attributes

auto_save	Whether to automatically save the results after each finished experiment.
custom_arguments	Custom arguments of the experiment function.
data	The data of the experiment.
duration	The duration of each experiment.
enable_priorities	Whether to enable priorities for the experiments.
func	The experiment function to run.
parameters	The parameters to sweep over.
pass_uuid	Whether to pass the uuid to the experiment function.
priority	The priority of each experiment.
save_path	Path to save the results to.
shape	The shape of the parameter grid.
status	The status of each parameter combination.
timeit	Whether to measure the duration of each experiment.
uuid	The uuid of each parameter combination.

save(mode: Literal['x', 'w'] = 'x') → None

Save the xarray dataset to the save path.

Parameters

modeLiteral[“x”, “w”] (default=”x”)

What to do if there is already data at the save path. If “x”, raise a FileExistsError. If “w”, overwrite the data.

run(status: str | list[str] | None = 'N', max_workers: int | None = None) → xr.Dataset

Run all experiments with the status ‘N’ (not started).

Parameters

statusstr | list[str] | None

The status of the experiments to run. If None, all experiments with status ‘N’ (not started) are run.

max_workersint | None

The maximum number of workers to run in parallel. If None, the number of workers is set to the number of CPUs available. Only relevant for parallel mode (‘mode=parallel’).

Returns

xr.Dataset

The xarray dataset with the results of the experiments.

Examples

from sweepexp import SweepExp
sweep = SweepExp(...)  # Initialize the sweep

# Run all experiments with status 'N'
sweep.run()

# Run all experiments with status 'C'
sweep.run("C")

# Run all experiments with status 'S' and 'F'
sweep.run(["S", "F"])