ptyrax.experiment

ptyrax.experiment#

Functions

get_experiment_index(hdf5_path)

Load the experiment index as a dictionary.

get_sweep_axes(hdf5_path)

Get sweep axis information from pooled HDF5.

iter_experiments(hdf5_path)

Iterate over experiments in a pooled HDF5 file.

load_experiment_by_name(hdf5_path, exp_name)

Load a specific experiment by name.

range_constructor(loader, node)

run_experiment(experiment_name[, dry_run])

Run or queue a DVC experiment sweep from a config with !range tags.

Classes

RangeTag(axis, values)

Custom YAML tag representing a parameter sweep range.
class ptyrax.experiment.RangeTag(axis, values)[source]#

Bases: object

Custom YAML tag representing a parameter sweep range.

Used in experiment config files with the !range YAML tag to specify values that should be swept over in a parameter study.

Variables:

  • axis – Name of the sweep axis (used for grouping), or None for auto.

  • values – List of values to sweep over.

Parameters:

  • axis (str | None)

  • values (List[Any])

ptyrax.experiment.get_experiment_index(hdf5_path)[source]#

Load the experiment index as a dictionary.

Parameters:

hdf5_path (str | Path) – Path to pooled experiment HDF5 file

Returns:

‘exp_names’, ‘indices’, ‘param_{name}’ for each parameter

Return type:

Dict with keys

Example

>>> index = get_experiment_index("experiment_sweep_xyz.hdf5")
>>> print(index["exp_names"])
>>> print(index["param_angle"])
ptyrax.experiment.get_sweep_axes(hdf5_path)[source]#

Get sweep axis information from pooled HDF5.

Parameters:

hdf5_path (str | Path) – Path to pooled experiment HDF5 file

Returns:

Dict with sweep metadata (sweep_id, n_experiments, parameter names)

Return type:

dict

Example

>>> axes = get_sweep_axes("experiment_sweep_xyz.hdf5")
>>> print(axes["sweep_id"])
>>> print(axes["param_names"])
ptyrax.experiment.iter_experiments(hdf5_path)[source]#

Iterate over experiments in a pooled HDF5 file.

Parameters:

hdf5_path (str | Path) – Path to pooled experiment HDF5 file

Yields:

Tuple of (experiment_name, experiment_group)

Return type:

Generator[tuple[str, Any], None, None]

Example

>>> for exp_name, exp_group in iter_experiments("experiment_sweep_xyz.hdf5"):
...     print(f"Processing {exp_name}")
...     losses = exp_group["training/scalars/0_loss/0_loss_total/value"][()]
ptyrax.experiment.load_experiment_by_name(hdf5_path, exp_name)[source]#

Load a specific experiment by name.

Parameters:

  • hdf5_path (str | Path) – Path to pooled experiment HDF5 file.

  • exp_name (str) – Experiment name to load.

Returns:

Nested dict with experiment data (scalars, images, etc.).

Return type:

dict

Example

>>> exp = load_experiment_by_name("experiment_sweep_xyz.hdf5", "leaky-lulu")
>>> losses = exp["training"]["scalars"]["0_loss"]["0_loss_total"]["value"]
ptyrax.experiment.range_constructor(loader, node)[source]#
Parameters:

  • loader (Any)

  • node (Any)

Return type:

RangeTag

ptyrax.experiment.run_experiment(experiment_name, dry_run=False)[source]#

Run or queue a DVC experiment sweep from a config with !range tags.

Parses the DVC stage’s config files, expands !range tags into a Cartesian product of parameter combinations, and queues one DVC experiment per combination.

Parameters:

  • experiment_name (str) – Name of the DVC stage to run.

  • dry_run (bool) – If True, log what would be queued without actually running.

Raises:

ValueError – If the experiment name is invalid or multiple dynamic config files are found.

Return type:

None

Contents