nexuml.tuning.optuna_tuner

nexuml.tuning.optuna_tuner

Optuna-based hyperparameter tuning for NexuML.

build_objective

build_objective(
    base_scenario: ScenarioSpec,
    search_space: dict[str, dict[str, Any]],
    metric_key: str = "val/loss",
    enable_progress_bar: bool = False,
    build_factory: Callable[..., ScenarioSpec]
    | None = None,
) -> Callable

Build an Optuna objective function from a base scenario and search space.

Parameters:

Name	Type	Description	Default
base_scenario	ScenarioSpec	The base ScenarioSpec to tune.	required
search_space	dict[str, dict[str, Any]]	Dict mapping dotted-path param names to Optuna suggest kwargs. Each entry must have a "type" key (float, int, categorical) and the remaining keys are forwarded to the Optuna suggest method. Example:: { "training.lr": {"type": "float", "low": 1e-5, "high": 1e-2, "log": True}, "training.max_epochs": {"type": "int", "low": 5, "high": 50}, "training.batch_size": {"type": "categorical", "choices": [32, 64, 128]}, }	required
metric_key	str	Lightning logged metric to optimise (must be logged with self.log()).	'val/loss'
enable_progress_bar	bool	Show Lightning progress bar during each trial.	False
build_factory	Callable[..., ScenarioSpec] | None	Optional factory to rebuild the scenario from sampled params.	None

Returns:

Name	Type	Description
Callable	Callable	objective(trial) -> float | list[float]

Raises:

Type	Description
ImportError	If the optuna package is not installed.

tune

tune(
    scenario: ScenarioSpec,
    search_space: dict[str, dict[str, Any]],
    tuning_spec: TuningSpec | None = None,
    metric_key: str = "val/loss",
    enable_progress_bar: bool = False,
    build_factory: Callable[..., ScenarioSpec]
    | None = None,
) -> Any

Run Optuna hyperparameter search.

Parameters:

Name	Type	Description	Default
scenario	ScenarioSpec	Base ScenarioSpec to tune.	required
search_space	dict[str, dict[str, Any]]	Dict mapping dotted-path param names to Optuna suggest kwargs.	required
tuning_spec	TuningSpec | None	TuningSpec controlling n_trials, storage, pruning, etc.	None
metric_key	str	Logged metric to optimise.	'val/loss'
enable_progress_bar	bool	Show progress bar during each trial.	False
build_factory	Callable[..., ScenarioSpec] | None	Optional factory to rebuild the scenario from sampled params.	None

Returns:

Type	Description
Any	Completed optuna.Study.

Raises:

Type	Description
ImportError	If the optuna package is not installed.