`research_helpers.sweep`

This module provides tools for planning and running parameter sweeps as scheduler job arrays and collecting the results. Planning, running, and status checks need only the standard library. collect and aggregate_runs require the sweep extra. See parameter sweeps and cluster jobs.

research_helpers.sweep.aggregate_runs(run_dirs, *, group_by=None, metrics=None, confidence=0.95)[source]

Combine the results of several runs of one grid, with a confidence interval per metric.

Parameters:

run_dirs (Iterable[Path | str]) – the run directories, one per independent draw.
group_by (Sequence[str] | None, default: None) – the columns identifying a parameter setting across runs. Defaults to the swept parameter names on the first run’s manifest.
metrics (Sequence[str] | None, default: None) – the measurements to aggregate. Defaults to every numeric column that is not a grouping column or engine bookkeeping.
confidence (float, default: 0.95) – the two-sided confidence level.

Return type:

DataFrame

Returns:

One row per parameter setting, carrying ‘n_runs’ and, for each metric, its mean, the half-width of the interval, and the standard error.

Raises:

ValueError – if no run directories were given, if they share no parameter setting, or if ‘group_by’ contains ‘combination_id’.

research_helpers.sweep.array_spec(indices)[source]

Render task indices the way ‘–array’ takes them, collapsing runs into ranges.

Parameters:: indices (Iterable[int]) – task indices, ascending.
Return type:: str
Returns:: For example ‘9-12,17’. Empty if there are no indices.

research_helpers.sweep.collect_results(run_dir, *, write=True, leading=None)[source]

Merge the part files into a table de-duplicating by combination id.

Duplicates are legitimate: a re-queued task can recompute a combination whose result was already flushed. The last record written wins.

Parameters:

run_dir (Path | str) – the run directory.
write (bool, default: True) – also write ‘results.csv’, and ‘results.parquet’ where an engine is installed.
leading (Iterable[str] | None, default: None) – columns to move to the front, so the table reads configuration first and measurements second. Defaults to the run’s own parameters, taken from the manifest. Pass an explicit sequence to override it, or an empty one to leave the order alone.

Return type:

DataFrame

Returns:

One row per combination.

research_helpers.sweep.leading_columns(run_dir)[source]

Return the configuration columns of a run, in the order they should be read.

The id first, then the constants, then the swept parameters. Everything else is a measurement.

Parameters:: run_dir (Path | str) – the run directory.
Return type:: list[str]
Returns:: The column names, or an empty list if the run has no manifest to describe it.

research_helpers.sweep.missing_task_indices(run_dir)[source]

Return the 1-based array indices whose slices still hold unfinished combinations.

Slices come from the manifest, not from the width of the array that was submitted, so an array narrower than the manifest plans leaves the tail simply unrun—nothing fails anywhere, the sweep stops short.

Parameters:: run_dir (Path | str) – the run directory.
Return type:: list[int]
Returns:: Sorted task indices with at least one combination still to run.

research_helpers.sweep.read_parts(run_dir)[source]

Read every result recorded by this run’s tasks.

Parameters:: run_dir (Path | str) – the run directory.
Return type:: list[dict[str, Any]]
Returns:: One dict per recorded result, in task order. A partial final line, left by a task killed mid-write, is skipped.
Raises:: FileNotFoundError – if the array has not run yet.

research_helpers.sweep.run_status(run_dir)[source]

Report how much of a planned sweep has finished.

Parameters:: run_dir (Path | str) – the run directory.
Return type:: dict[str, float | int | str]
Returns:: Counts, percentage complete, and the runtime seen so far. An unfinished run also carries the array width it was planned at and the indices to resubmit.

class research_helpers.sweep.Manifest(grid, combinations, n_tasks=1, metadata=<factory>, notes='', constants=<factory>, tuple_params=<factory>, artefact_key=None)[source]

Bases: object

The full description of a planned sweep, written once and read by every task.

Parameters:

grid (dict[str, list[Any]])
combinations (list[dict[str, Any]])
n_tasks (int)
metadata (dict[str, Any])
notes (str)
constants (dict[str, Any])
tuple_params (list[str])
artefact_key (str | None)

property n_combinations: int: Return how many parameter combinations the sweep covers.

save(run_dir)[source]

Write the manifest into ‘run_dir’ and return the path.

Return type:: Path
Parameters:: run_dir (Path | str)

classmethod load(run_dir)[source]

Read the manifest from ‘run_dir’.

Raises:: FileNotFoundError – if the sweep has not been planned.
Return type:: Manifest
Parameters:: run_dir (Path | str)

research_helpers.sweep.combination_id(params)[source]

Return a short, stable id for a parameter combination.

Derived from the values alone, so the same combination keeps the same id however the grid was expanded and on whichever machine.

Parameters:: params (dict[str, Any]) – the parameter values, excluding any id already assigned.
Return type:: str
Returns:: Twelve hex characters.

research_helpers.sweep.expand_grid(grid, constants=None, *, tuple_params=())[source]

Expand a parameter grid into the full Cartesian product of combinations.

Parameters:

grid (dict[str, list[Any]]) – parameter name to the list of values to sweep.
constants (dict[str, Any] | None, default: None) – fixed values added to every combination before it is hashed, e.g. the label of the dataset being swept. Part of the identity, so two datasets do not collide.
tuple_params (Iterable[str], default: ()) – parameter names whose values are tuples rather than lists.

Return type:

list[dict[str, Any]]

Returns:

One dict per combination, each carrying the constants and a ‘combination_id’. Ordering is deterministic: parameter names sorted, values in the order given.

Raises:

ValueError – if the grid is empty.

research_helpers.sweep.restore_tuples(params, tuple_params)[source]

Convert the named list-valued parameters back into tuples, in place.

JSON has no tuple, so a parameter whose value must be a tuple arrives back from the manifest as a list. Identity is unaffected, so this is safe to apply at any point: ‘json.dumps’ renders a tuple and a list to the same array, so a combination hashes the same either way.

Parameters:

params (dict[str, Any]) – one combination, modified in place.
tuple_params (Iterable[str]) – the parameter names whose values are tuples.

Return type:

dict[str, Any]

Returns:

The same dict, for chaining.

research_helpers.sweep.slice_bounds(n_items, n_tasks, task_index)[source]

Return the [start, end) slice of items belonging to a 1-based array task.

Parameters:

n_items (int) – total number of combinations.
n_tasks (int) – number of array tasks the sweep was planned for.
task_index (int) – 1-based task index.

Return type:

tuple[int, int]

Returns:

Start and end indices into the combination list.

Raises:

ValueError – if ‘task_index’ is outside 1..n_tasks.

class research_helpers.sweep.Unit(run_dir, task)[source]

Bases: object

A slice of a sweep.

Parameters:

run_dir (str)
task (int)

line()[source]

Render the unit as a line of a queue file.

Return type:: str

research_helpers.sweep.build_units(run_dirs, *, pending=False)[source]

Return every slice of every sweep, in run directory then task order.

Parameters:

run_dirs (Iterable[Path]) – planned run directories.
pending (bool, default: False) – only the slices still holding an unfinished combination, as ‘status’ finds them.

Return type:

list[Unit]

Returns:

The units.

research_helpers.sweep.chunk(units, per_task, max_array)[source]

Split units into groups small enough for one array of at most ‘max_array’ elements each.

Parameters:

units (list[Unit]) – the units.
per_task (int) – units each element runs.
max_array (int) – elements per array.

Return type:

list[list[Unit]]

Returns:

The groups, in order, each holding at most ‘per_task * max_array’ units.

Raises:

ValueError – if ‘per_task’ or ‘max_array’ is below 1.

research_helpers.sweep.expand_run_dirs(patterns)[source]

Expand run directory arguments.

Parameters:: patterns (Iterable[str | Path]) – run directories, or patterns matching them.
Return type:: list[Path]
Returns:: The planned run directories (those holding a manifest), sorted and without duplicates.
Raises:: FileNotFoundError – if a plain path is not a planned run directory, or a pattern matches none.

research_helpers.sweep.quarantine(units, *, apply=False, artefacts_dir=ARTEFACTS_DIR)[source]

Set aside every result of these units, so that a pending queue runs them again.

This is for units that two processes ran at once. When a scheduler re-queues an element whose node has stopped answering, the original process can keep running, and the same block then runs twice, from two nodes, appending to the same part files. On network storage that leaves runs of NUL bytes and lines spliced together.

Parameters:

units (Iterable[Unit]) – the units the duplicated elements ran (‘read_units’ of each element).
apply (bool, default: False) – rename the files. Without it, only count them.
artefacts_dir (str, default: ARTEFACTS_DIR) – the run directory’s subdirectory holding per-combination artefacts.

Return type:

Counter[str]

Returns:

Counts of ‘units’, ‘part files’, ‘damaged part files’ and ‘artefacts’ found, and of files ‘set aside’.

research_helpers.sweep.read_plan_list(path)[source]

Read a plan list: one ‘config<TAB>run_dir’ per line, blank lines and ‘#’ comments skipped.

Paths are taken as written, so relative ones resolve against the working directory.

Parameters:: path (Path | str) – the plan list.
Return type:: list[tuple[Path, Path]]
Returns:: The (config, run directory) pairs, in order.
Raises:: ValueError – if a line does not have exactly two tab-separated fields.

research_helpers.sweep.read_units(path, element, per_task)[source]

Return the units that an array element runs.

Parameters:

path (Path | str) – the queue file.
element (int) – the 1-based array element.
per_task (int) – units per element.

Return type:

list[Unit]

Returns:

Lines ‘(element - 1) * per_task + 1’ to ‘element * per_task’ of the file, as units. Empty past the end of the file.

Raises:

ValueError – if ‘element’ or ‘per_task’ is below 1, or a line is malformed.

research_helpers.sweep.select_run_dirs(run_dirs, *, where=None, where_not=None)[source]

Keep the sweeps whose manifests match, to queue a campaign in parts.

Parameters:

run_dirs (Iterable[Path]) – planned run directories.
where (Mapping[str, Collection[str]] | None, default: None) – key to accepted values. A sweep is kept only if, for every key, its value is one of these. Keys are looked up in the manifest’s metadata, then its constants, and values are compared as text.
where_not (Mapping[str, Collection[str]] | None, default: None) – key to rejected values. A sweep is dropped if any of its values is one of these.

Return type:

list[Path]

Returns:

The run directories kept, in the order given.

research_helpers.sweep.shuffled(units, seed)[source]

Return the units in a seeded random order, so that each element draws a mix of sweeps.

Parameters:

units (Iterable[Unit]) – the units.
seed (int) – the seed for the order.

Return type:

list[Unit]

Returns:

The same units, reordered.

research_helpers.sweep.write_queues(units, out, per_task, max_array=None, *, force=False)[source]

Write the queue files for a list of units, named ‘<out>.01.units’, ‘<out>.02.units’, …

An existing queue of the same stem is refused unless ‘force’ is set. A resumed queue takes a new stem. With ‘force’, every file of the old queue is removed first, so that a shorter queue never leaves an old file behind.

Parameters:

units (list[Unit]) – the units, in the order the elements should run them.
out (Path | str) – the queue’s stem.
per_task (int) – units each element runs. Part of the queue: submit it with the same number.
max_array (int | None, default: None) – elements per array, and so per file. Defaults to the project’s ‘max-array’.
force (bool, default: False) – replace an existing queue of this stem.

Return type:

list[tuple[Path, int]]

Returns:

For each file, its path and the array width to submit it with.

Raises:

FileExistsError – if a queue of this stem exists and ‘force’ is not set.

research_helpers.sweep.artefact_path(run_dir, combination_id)[source]

Return the location where a combination’s artefact is written.

Parameters:

run_dir (Path | str) – the run directory.
combination_id (str) – the combination’s id.

Return type:

Path

Returns:

The path, whether or not anything was written there.

research_helpers.sweep.completed_ids(run_dir)[source]

Return the combination ids already recorded in this run’s parts.

A task killed mid-write leaves a partial final line, which is skipped rather than treated as corruption.

Parameters:: run_dir (Path | str) – the run directory.
Return type:: set[str]
Returns:: The ids found.

research_helpers.sweep.jsonable(value)[source]

Convert result values into JSON-serialisable equivalents.

Parameters:: value (Any) – any result value, including nested containers.
Return type:: Any
Returns:: The same value with numpy scalars unwrapped and tuples and sets rendered as lists.

research_helpers.sweep.read_artefact(run_dir, combination_id)[source]

Read back what a combination put aside.

Parameters:

run_dir (Path | str) – the run directory.
combination_id (str) – the combination’s id.

Return type:

Any

Returns:

Whatever the evaluation returned under the manifest’s ‘artefact_key’.

Raises:

FileNotFoundError – if this combination recorded no artefact.

research_helpers.sweep.run_slice(run_dir, evaluate, *, task_index, n_tasks=None, context=None, resume=True, quiet=False, artefact_key=None)[source]

Run this task’s slice of the sweep, appending each result to its own part file.

Parameters:

run_dir (Path | str) – the run directory, holding the manifest.
evaluate (Callable[[dict[str, Any], Any], dict[str, Any]]) – called as ‘evaluate(params, context)’ for each combination, returning the measurements to record. The parameters are recorded alongside them automatically.
task_index (int) – 1-based array task index.
n_tasks (int | None, default: None) – total array tasks. Defaults to the manifest’s value.
context (Any, default: None) – whatever the evaluation needs, built once for the task.
resume (bool, default: True) – skip combinations already recorded in this run.
quiet (bool, default: False) – suppress per-combination progress.
artefact_key (str | None, default: None) – the key under which the evaluation returns output too bulky for a results row.

Return type:

Path

Returns:

The part file written.

research_helpers.sweep.task_index_from_env(explicit=None)[source]

Return the array task index, from the argument, the scheduler, or 1 for a local run.

Parameters:: explicit (int | None, default: None) – an index given on the command line, which wins.
Return type:: int
Returns:: The 1-based task index.

class research_helpers.sweep.Interval(mean, half_width, standard_error)[source]

Bases: NamedTuple

A mean and the interval around it.

Parameters:

mean (float)
half_width (float)
standard_error (float)

mean: float: Alias for field number 0

half_width: float: Alias for field number 1

standard_error: float: Alias for field number 2

property low: float: Return the lower bound.

property high: float: Return the upper bound.

research_helpers.sweep.confidence_interval(values, confidence=0.95)[source]

Return the mean of ‘values’ and the t-interval around it.

Parameters:

values (Iterable[float]) – the observations.
confidence (float, default: 0.95) – the two-sided confidence level, strictly between 0 and 1.

Return type:

Interval

Returns:

The mean, the half-width, and the standard error. All three are NaN for no observations, and the last two are NaN for exactly one.

Raises:

ValueError – if ‘confidence’ is not in (0, 1).

research_helpers.sweep.student_t_quantile(p, df)[source]

Return t such that P(T <= t) = ‘p’, for Student’s t with ‘df’ degrees of freedom.

Parameters:

p (float) – the cumulative probability, strictly between 0 and 1.
df (float) – degrees of freedom, greater than 0.

Return type:

float

Returns:

The quantile.

Raises:

ValueError – if ‘p’ is not in (0, 1), or ‘df’ is not positive.

class research_helpers.sweep.Sweep(evaluate=None, context=None)[source]

Bases: object

A parameter sweep: an evaluation, an optional per-task context, and a command line.

Parameters:

evaluate (Callable[[dict[str, Any], Any], dict[str, Any]] | None)
context (Callable[[Manifest, Path], Any] | None)

evaluate(function)[source]

Register the evaluation, called as ‘function(params, context)’ for each combination.

Return type:: Callable[[dict[str, Any], Any], dict[str, Any]]
Parameters:: function (Callable[[dict[str, Any], Any], dict[str, Any]])

context(function)[source]

Register what to build once per task, called as ‘function(manifest, run_dir)’.

Return type:: Callable[[Manifest, Path], Any]
Parameters:: function (Callable[[Manifest, Path], Any])

property evaluation: Callable[[dict[str, Any], Any], dict[str, Any]]

Return the registered evaluation.

Raises:: RuntimeError – if none was registered.

build_context(manifest, run_dir)[source]

Return the per-task context, or None if the sweep registered no context factory.

Return type:

Any

Parameters:

manifest (Manifest)
run_dir (Path)

runs_dir()[source]

Return where sweeps live.

Return type:: Path

plan(run_dir, grid, *, constants=None, metadata=None, notes='', n_tasks=None, max_tasks=None, tuple_params=(), artefact_key=None)[source]

Expand a grid and write the manifest that every task will read.

Parameters:

run_dir (Path | str) – directory to hold the manifest, parts and results.
grid (dict[str, list[Any]]) – parameter name to the values to sweep.
constants (dict[str, Any] | None, default: None) – fixed values on every combination, included in its identity.
metadata (dict[str, Any] | None, default: None) – anything a task needs, e.g. a dataset name or a seed count.
notes (str, default: '') – free text recorded with the sweep.
n_tasks (int | None, default: None) – array width. Defaults to the smaller of ‘max_tasks’ and the grid size.
max_tasks (int | None, default: None) – the cap applied when ‘n_tasks’ is not given. Defaults to the project’s ‘max-array’.
tuple_params (Sequence[str], default: ()) – parameters whose values are tuples rather than lists.
artefact_key (str | None, default: None) – the key under which the evaluation returns output too bulky for a results row.

Return type:

Manifest

Returns:

The manifest, already written.

plan_config(config, run_dir, n_tasks=None)[source]

Plan a sweep from a config file, as the ‘plan’ command does.

Parameters:

config (Path | str) – the sweep config, TOML, JSON or YAML.
run_dir (Path | str) – directory to hold the manifest, parts and results.
n_tasks (int | None, default: None) – array width, overriding the config’s ‘max_tasks’ cap.

Return type:

Manifest

Returns:

The manifest.

run_unit(unit, *, resume=True)[source]

Run a slice of a sweep, as an element of a queue does.

The context is rebuilt for every unit and never shared between slices, even between two of the same sweep.

Parameters:

unit (Unit) – the run directory and the slice’s task index.
resume (bool, default: True) – skip combinations already recorded.

Return type:

Path

Returns:

The part file.

main(argv=None)[source]

Run the sweep command line.

Parameters:: argv (list[str] | None, default: None) – command line arguments, or None to read ‘sys.argv’.
Return type:: int
Returns:: A process exit status.

research_helpers.sweep.estimate_runtime(manifest, evaluate, context=None, *, samples=DEFAULT_SAMPLES, seed=0)[source]

Time a random sample of combinations.

The sample is random, not strided. ‘research_helpers.sweep.grid.expand_grid’ orders combinations by the Cartesian product over sorted parameter names, so a stride lands on a systematically correlated subset.

Parameters:

manifest (Manifest) – the planned sweep.
evaluate (Callable[[dict[str, Any], Any], dict[str, Any]]) – the evaluation, called as for a real task.
context (Any, default: None) – whatever the evaluation needs.
samples (int, default: DEFAULT_SAMPLES) – how many combinations to time.
seed (int, default: 0) – seed for choosing the sample.

Return type:

list[float]

Returns:

The timings, sorted ascending.

research_helpers.sweep.read_config(path)[source]

Read a sweep config from TOML, JSON, or YAML.

TOML and JSON are standard library, YAML needs ‘pyyaml’.

Parameters:: path (Path | str) – the config file, read by extension.
Return type:: dict[str, Any]
Returns:: The parsed config – ‘grid’, and optionally ‘constants’, ‘metadata’, ‘notes’, ‘max_tasks’, ‘tuple_params’, and ‘artefact_key’.
Raises:: ImportError – for a YAML config when ‘pyyaml’ is not installed.