SweepExp

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

Bases: object

Run a parameter sweep sequentially.

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

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 SweepExp 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 sequentially. For parallel execution, consider using the ‘SweepExpParallel’ or ‘SweepExpMPI’ classes.

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

from sweepexp import SweepExp

# 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 = SweepExp(
    func=my_experiment,
    parameters={"x": [1, 2, 3], "y": [4, 5, 6]},
)

# Run the sweep
sweep.run()

__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.

add_custom_argument(name: str, default_value: any) → None

Add custom arguments to the experiment function.

Parameters

namestr

The name of the argument.

default_valueany

The default value of the argument.

Examples

from sweepexp import SweepExp

# Create a function that takes a custom argument
def my_experiment(param1: int, custom: float) -> dict:
    return {"product": param1 * custom}

sweep = SweepExp(
    func=my_experiment,
    parameters={"param1": [1, 2, 3]},
)

# Add a custom argument
sweep.add_custom_argument("custom", 2.0)
# Set the custom argument to 3.0 for the second experiment
sweep.data["custom"].data[1] = 3.0
# Run the sweep
sweep.run()

run(status: str | list[str] | None = 'N', max_workers: int | None = None) → xarray.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"])

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.

static load(save_path: Path | str) → xarray.DataArray

Load the xarray dataset from a file.

Parameters

save_pathPath | str

The path to the file to load the data from. Must have one of the following extensions: ‘.zarr’, ‘.nc’, ‘.cdf’, ‘.pkl’.

Returns

dataxr.DataArray

The xarray dataset with the data.

Note

Not the SweepExp object is returned, but only the data.

Examples

# Create a SweepExp object
sweep = SweepExp(
    func=lambda x: {"y": x},
    parameters={"x": [1, 2, 3]},
    save_path="my_data.zarr"
)
# Run the sweep and save it
sweep.run()
sweep.save()
# Load the data from the file
df = SweepExp.load("my_data.zarr")

reset_status(states: str | list[str] | None) → None

Reset the status of experiments to ‘N’ (not started).

Parameters

stateslist[str] | None

The states to reset. If None, all states with status ‘C’ (completed) and ‘F’ (failed) are reset

Examples

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

# Reset all experiments with status 'C' and 'F' to 'N'
sweep.reset_status()
# Reset all experiments with status 'C' to 'N'
sweep.reset_status("C")
# Reset all experiments with status 'S' and 'F' to 'N'
sweep.reset_status(["S", "F"])

property func: Callable

The experiment function to run.

property parameters: dict[str, list]

The parameters to sweep over.

property custom_arguments: set[str]

Custom arguments of the experiment function.

property save_path: Path | None

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.

property pass_uuid: bool

Whether to pass the uuid to the experiment function.

Description

When the pass_uuid property is set to True, a uuid (unique identifier) will be assigned to each experiment. This uuid is passed to the experiment function as an argument. Make sure to add a uuid argument to the function.

Examples

from sweepexp import SweepExp

# Create a function that takes the uuis an an argument and write
# something to a file with the uuid in the name
def my_experiment(x: int, uuid: str) -> dict:
    with open(f"output_{uuid}.txt", "w") as file:
        file.write(f"Experiment with x={x} and uuid={uuid}.")
    return {}

sweep = SweepExp(
    func=my_experiment,
    parameters={"x": [1, 2, 3]},
)

# Enable the uuid
sweep.pass_uuid = True
# Run the sweep
sweep.run()

property auto_save: bool

Whether to automatically save the results after each finished experiment.

property timeit: bool

Whether to measure the duration of each experiment.

Description

When the timeit property is set to True, a new variable ‘duration’ is added to the data. This variable stores the duration of each experiment in seconds.

property enable_priorities: bool

Whether to enable priorities for the experiments.

property shape: tuple[int]

The shape of the parameter grid.

property data: xarray.Dataset

The data of the experiment.

property uuid: xarray.DataArray

The uuid of each parameter combination.

Description

The uuid is a string that uniquely identifies each experiment. By default, uuids are disbabled to save memory. They can be enabled by setting the ‘pass_uuid’ property to True:

sweep = SweepExp(...)
sweep.pass_uuid = True

property status: xarray.DataArray

The status of each parameter combination.

Possible values are: - ‘N’: not started - ‘C’: completed - ‘F’: failed - ‘S’: skip

property duration: xarray.DataArray

The duration of each experiment.

Description

The duration is the time it took to run the experiment with the given parameter combination. The duration is only available if the ‘timeit’ property is set to True. Make sure to enable it with:

sweep = SweepExp(...)
sweep.timeit = True

property priority: xarray.DataArray

The priority of each experiment.

Description

The priority is an integer that determines the order in which the experiments are run. Experiments with higher priority are run first. Make sure to enable the priorities before accessing them with:

sweep = SweepExp(...)
sweep.enable_priorities = True