Reference

HPO

Optuna-backed hyperparameter optimization: the run_hpo entry point, five loss functions built for judging forecasts of (possibly chaotic) dynamics, and helpers for inspecting finished studies.

Optional dependency

run_hpo and the study utilities require optuna (pip install "resdag[hpo]"). The loss functions import without it.

hpo

Hyperparameter Optimization

This module provides Optuna-based hyperparameter optimization for ESN models. It supports multiple loss functions specialized for time series forecasting and chaotic systems.

Installation

This module requires the optional optuna dependency::

pip install resdag[hpo]

or::

pip install optuna

FUNCTION	DESCRIPTION
run_hpo	Run hyperparameter optimization study.
get_loss	Get a loss function by name.
get_study_summary	Generate summary of completed study.
resolve_pruner	Resolve a pruner key/instance to an :class:optuna.pruners.BasePruner.

Pruners

Pass pruner= to :func:run_hpo to enable Optuna early stopping, fed by the per-trial intermediate forecast-horizon reports. Registry keys (PRUNERS):

• "asha": ASHA / SuccessiveHalvingPruner
• "hyperband": Hyperband pruner
• "median": Median pruner (recommended starting point)
• "threshold": Threshold pruner
• "none": NopPruner (no pruning, the default)

Loss Functions

The following loss functions are available:

• "efh": Expected Forecast Horizon (default, recommended for chaotic systems)
• "forecast_horizon": Forecast Horizon Loss (contiguous valid steps)
• "lyapunov": Lyapunov-weighted Loss (exponential decay)
• "standard": Standard Loss (mean geometric mean error)
• "soft_horizon": Soft valid horizon (cumulative survival probability)

Examples:

Basic HPO workflow:

>>> from resdag.hpo import run_hpo, get_study_summary
>>> from resdag.models import ott_esn
>>>
>>> def model_creator(reservoir_size, spectral_radius):
...     return ott_esn(
...         reservoir_size=reservoir_size,
...         feedback_size=3,
...         output_size=3,
...         spectral_radius=spectral_radius,
...     )
>>>
>>> def search_space(trial):
...     return {
...         "reservoir_size": trial.suggest_int("reservoir_size", 100, 500, step=50),
...         "spectral_radius": trial.suggest_float("spectral_radius", 0.5, 1.5),
...     }
>>>
>>> def data_loader(trial):
...     return {
...         "warmup": warmup, "train": train, "target": target,
...         "f_warmup": f_warmup, "val": val,
...     }
>>>
>>> study = run_hpo(
...     model_creator=model_creator,
...     search_space=search_space,
...     data_loader=data_loader,
...     n_trials=50,
... )
>>> print(get_study_summary(study))

See Also

resdag.training.ESNTrainer : Training interface. resdag.utils.data : Data loading utilities.

run_hpo

run_hpo(model_creator: Callable[..., ESNModel], search_space: Callable[[Trial], dict[str, Any]], data_loader: Callable[[Trial], dict[str, Any]], n_trials: int, loss: str | LossProtocol = 'efh', loss_params: dict[str, Any] | None = None, targets_key: str = 'output', drivers_keys: list[str] | None = None, monitor_losses: list[str | LossProtocol] | None = None, monitor_params: dict[str, dict[str, Any]] | None = None, study_name: str | None = None, storage: str | None = None, warm_start: list[dict[str, Any]] | None = None, transfer_from: Study | str | None = None, sampler: BaseSampler | None = None, pruner: str | BasePruner | None = None, seed: int | None = None, device: str | device | None = None, n_workers: int = 1, verbosity: int = 1, catch_exceptions: bool = True, penalty_value: float = 10000000000.0, clip_value: float | None = None, prune_on_clip: bool = False, horizon_key: str | None = None, n_checkpoints: int = 5, torch_scoring: bool = True, trial_callbacks: list[TrialCallback] | None = None) -> Study

Run an Optuna hyperparameter optimization study for ESN models.

This function provides a complete HPO pipeline that handles model creation, training, forecasting, and evaluation with robust error handling.

When n_workers > 1, optimization is parallelized using real OS processes (not threads) that coordinate via shared file storage. By default, uses Optuna's JournalFileStorage which is designed for multi-process coordination. Pass a .db path or sqlite:/// URL to use SQLite instead (with WAL mode for concurrency).

PARAMETER	DESCRIPTION
model_creator	Function that creates a fresh model for each trial. Must accept all hyperparameters from search_space as keyword arguments. TYPE: Callable[..., ESNModel]
search_space	Function that defines the hyperparameter search space. Uses Optuna's trial.suggest_* methods and returns a dictionary of parameters. TYPE: Callable[[Trial], dict[str, Any]]
data_loader	Function that loads and returns data. Must return a dictionary with: "warmup": Warmup data (B, warmup_steps, D) "train": Training input (B, train_steps, D) "target": Training targets (B, train_steps, D) "f_warmup": Forecast warmup (B, warmup_steps, D) "val": Validation data (B, val_steps, D) TYPE: Callable[[Trial], dict[str, Any]]
n_trials	Total number of trials to run. TYPE: int
loss	Loss function to optimize. Can be: "efh": Expected Forecast Horizon (default, recommended) "forecast_horizon": Forecast Horizon Loss "lyapunov": Lyapunov-weighted Loss "standard": Standard Loss "soft_horizon": Soft valid horizon A custom callable following LossProtocol TYPE: str or LossProtocol DEFAULT: "efh"
loss_params	Additional keyword arguments for the loss function. TYPE: dict DEFAULT: None
targets_key	Name of the readout layer for training targets. TYPE: str DEFAULT: "output"
drivers_keys	List of driver input keys in data dict for input-driven models. TYPE: list[str] DEFAULT: None
monitor_losses	Additional loss functions to compute and log (but not optimize on). Can be loss names (e.g., "standard", "lyapunov") or callables. Results are stored as trial user attributes with prefix monitor_. TYPE: list[str | LossProtocol] DEFAULT: None
monitor_params	Keyword arguments for each monitor loss. Keys are loss function names (e.g., "lyapunov_weighted_loss"), values are kwargs dicts. TYPE: dict[str, dict[str, Any]] DEFAULT: None
study_name	Name for the study. If None, auto-generated from model_creator. TYPE: str DEFAULT: None
storage	Storage path or URL. Behaviour depends on the value: None: In-memory (single worker) or temp journal file (multi-worker). "study.log": Journal file storage (recommended for multi-worker). "study.db" or "sqlite:///study.db": SQLite with WAL mode. TYPE: str DEFAULT: None
warm_start	Known-good configurations to evaluate before any sampler-proposed trial. Each dict maps parameter name to a fixed value and is enqueued via :meth:optuna.study.Study.enqueue_trial. Enqueued in the parent process (pre-fork), so the warm-start trials are picked up by every worker when n_workers > 1. TYPE: list[dict] DEFAULT: None
transfer_from	A prior study (or a storage path/URL holding exactly one study) whose COMPLETE trials are copied into the new study via :meth:optuna.study.Study.add_trial, seeding the sampler without re-running the objective. Only parameters that overlap with the current search space are kept; mismatched parameters are dropped and logged at INFO level. Ingested in the parent process (pre-fork), so it works with n_workers > 1. TYPE: Study or str DEFAULT: None
sampler	Optuna sampler. Defaults to TPESampler with multivariate=True. TYPE: BaseSampler DEFAULT: None
pruner	Early-stopping policy fed by the per-trial intermediate forecast-horizon reports. Accepts a registry key — "asha" (ASHA / :class:~optuna.pruners.SuccessiveHalvingPruner), "hyperband", "median", "threshold", or "none" — a fully-configured :class:optuna.pruners.BasePruner instance, or None. None (the default) and "none" disable pruning (:class:optuna.pruners.NopPruner); "median" is a good first choice for forecasting studies because it stops trials whose intermediate loss is worse than the running median. A diverging configuration reports a large loss at an early forecast-horizon checkpoint and is pruned before the full horizon is evaluated, saving compute. Pruning is honored in both single- and multi-process modes and coexists with the clip_value / prune_on_clip path. TYPE: str or BasePruner DEFAULT: None
seed	Random seed for reproducibility. Seeds the Optuna sampler and, per trial, PyTorch / NumPy / Python random (seed + trial.number). The per-trial seed is also threaded into model_creator when it accepts a seed keyword argument, so the reservoir topology and input/feedback initializers — which build their own RNGs and do not read the legacy global NumPy state — reproduce as well. With a seed-aware model_creator, two run_hpo calls sharing the same seed therefore yield identical trials, parameters, and best_value. See Reproducible studies in the HPO reference for how to make a creator seed-aware. TYPE: int DEFAULT: None
device	Device to place models and data on (e.g., "cuda" or torch.device("cpu")). If None, uses default device. TYPE: str or device DEFAULT: None
n_workers	Number of parallel workers. When > 1, uses real OS processes (multiprocessing) that coordinate via shared file storage. The workers are started with fork on CPU-only studies, but with spawn when the study touches CUDA (device is a CUDA device, or CUDA is already initialised in the parent) because fork cannot re-initialise CUDA. spawn pickles the objective, so model_creator, search_space, data_loader (and any loss callable) must be top-level, importable functions — lambdas or closures defined inside another function are not picklable and will raise PicklingError. Single-worker (n_workers=1) studies run in-process and have no such restriction. TYPE: int DEFAULT: 1
verbosity	Logging verbosity: 0 = silent, 1 = normal, 2 = verbose. TYPE: int DEFAULT: 1
catch_exceptions	If True, catch exceptions and return penalty_value. TYPE: bool DEFAULT: True
penalty_value	Value returned for failed trials. TYPE: float DEFAULT: 1e10
clip_value	Upper bound for the objective value. When set and the raw loss exceeds this threshold, the value returned to Optuna is clamped to clip_value (or the trial is pruned, see prune_on_clip). The unclipped value is always stored as the "raw_loss" trial user attribute. When None (default), no clipping is applied and raw_loss == loss. TYPE: float DEFAULT: None
prune_on_clip	If True and clip_value is set, trials whose raw loss exceeds clip_value are pruned instead of returning the clipped value. TYPE: bool DEFAULT: False
horizon_key	Key in the data_loader output dict giving the forecast horizon (an int) for the autoregressive phase. When None (default), the horizon defaults to val.shape[1] (the number of validation steps). Set this to forecast a horizon that differs from the validation window length. TYPE: str DEFAULT: None
n_checkpoints	Number of growing forecast-horizon checkpoints at which each trial reports an intermediate loss and the configured pruner is consulted. The single forecast is sliced to each checkpoint (no repeated forecasts), and reports are keyed by 1-based checkpoint index. Values <= 1 disable intermediate reporting (one report at the full horizon), which also disables checkpoint-based pruning — set this >= 2 for pruner to have any early-stopping signal to act on. TYPE: int DEFAULT: 5
torch_scoring	If True, evaluate the loss on the forecast's device with torch.Tensor inputs, only moving the final scalar to the host, and fall back to NumPy when the loss rejects tensors. The five built-in losses are NumPy-based and always take the fallback; a custom torch-native loss benefits from staying on-device (no GPU→CPU→NumPy round-trip per checkpoint). Set False to always score in NumPy. TYPE: bool DEFAULT: True
trial_callbacks	Per-trial callbacks invoked after each successful evaluation. Each is called as callback(trial, context) where context is a dict with "params", "raw_loss" and "loss". A callback may call trial.report / raise optuna.TrialPruned to act as an extra pruning signal; any other exception it raises is logged and swallowed, so a callback can never fail a trial. Under n_workers > 1 with the spawn start method the callbacks must be picklable (top-level functions, not lambdas or closures). TYPE: list of callable DEFAULT: None

RETURNS	DESCRIPTION
Study	The completed study with all trial results.

RAISES	DESCRIPTION
ValueError	If n_trials is not positive.
TypeError	If model_creator, search_space, or data_loader are not callable.

Examples:

Basic usage:

>>> from resdag.hpo import run_hpo
>>> from resdag.models import ott_esn
>>>
>>> def model_creator(reservoir_size, spectral_radius):
...     return ott_esn(
...         reservoir_size=reservoir_size,
...         feedback_size=3,
...         output_size=3,
...         spectral_radius=spectral_radius,
...     )
>>>
>>> def search_space(trial):
...     return {
...         "reservoir_size": trial.suggest_int("reservoir_size", 100, 500, step=50),
...         "spectral_radius": trial.suggest_float("spectral_radius", 0.5, 1.5),
...     }
>>>
>>> def data_loader(trial):
...     # Load your data here
...     return {
...         "warmup": warmup, "train": train, "target": target,
...         "f_warmup": f_warmup, "val": val,
...     }
>>>
>>> study = run_hpo(
...     model_creator=model_creator,
...     search_space=search_space,
...     data_loader=data_loader,
...     n_trials=50,
...     loss="efh",
... )
>>> print(f"Best params: {study.best_params}")
>>> print(f"Best value: {study.best_value}")

Parallel with 4 workers:

>>> study = run_hpo(
...     model_creator=model_creator,
...     search_space=search_space,
...     data_loader=data_loader,
...     n_trials=100,
...     n_workers=4,
...     storage="study.log",
... )

See Also

LOSSES : Available loss functions. get_study_summary : Generate study summary. ESNTrainer : Training interface.

Source code in src/resdag/hpo/run.py

def run_hpo(
    model_creator: Callable[..., ESNModel],
    search_space: Callable[[optuna.Trial], dict[str, Any]],
    data_loader: Callable[[optuna.Trial], dict[str, Any]],
    n_trials: int,
    loss: str | LossProtocol = "efh",
    loss_params: dict[str, Any] | None = None,
    targets_key: str = "output",
    drivers_keys: list[str] | None = None,
    monitor_losses: list[str | LossProtocol] | None = None,
    monitor_params: dict[str, dict[str, Any]] | None = None,
    study_name: str | None = None,
    storage: str | None = None,
    warm_start: list[dict[str, Any]] | None = None,
    transfer_from: "optuna.Study | str | None" = None,
    sampler: optuna.samplers.BaseSampler | None = None,
    pruner: "str | optuna.pruners.BasePruner | None" = None,
    seed: int | None = None,
    device: str | torch.device | None = None,
    n_workers: int = 1,
    verbosity: int = 1,
    catch_exceptions: bool = True,
    penalty_value: float = 1e10,
    clip_value: float | None = None,
    prune_on_clip: bool = False,
    horizon_key: str | None = None,
    n_checkpoints: int = 5,
    torch_scoring: bool = True,
    trial_callbacks: list[TrialCallback] | None = None,
) -> optuna.Study:
    """Run an Optuna hyperparameter optimization study for ESN models.

    This function provides a complete HPO pipeline that handles model creation,
    training, forecasting, and evaluation with robust error handling.

    When ``n_workers > 1``, optimization is parallelized using real OS processes
    (not threads) that coordinate via shared file storage. By default, uses
    Optuna's ``JournalFileStorage`` which is designed for multi-process
    coordination. Pass a ``.db`` path or ``sqlite:///`` URL to use SQLite
    instead (with WAL mode for concurrency).

    Parameters
    ----------
    model_creator : Callable[..., ESNModel]
        Function that creates a fresh model for each trial. Must accept all
        hyperparameters from ``search_space`` as keyword arguments.
    search_space : Callable[[Trial], dict[str, Any]]
        Function that defines the hyperparameter search space. Uses Optuna's
        ``trial.suggest_*`` methods and returns a dictionary of parameters.
    data_loader : Callable[[Trial], dict[str, Any]]
        Function that loads and returns data. Must return a dictionary with:

        - ``"warmup"``: Warmup data (B, warmup_steps, D)
        - ``"train"``: Training input (B, train_steps, D)
        - ``"target"``: Training targets (B, train_steps, D)
        - ``"f_warmup"``: Forecast warmup (B, warmup_steps, D)
        - ``"val"``: Validation data (B, val_steps, D)

    n_trials : int
        Total number of trials to run.
    loss : str or LossProtocol, default="efh"
        Loss function to optimize. Can be:

        - ``"efh"``: Expected Forecast Horizon (default, recommended)
        - ``"forecast_horizon"``: Forecast Horizon Loss
        - ``"lyapunov"``: Lyapunov-weighted Loss
        - ``"standard"``: Standard Loss
        - ``"soft_horizon"``: Soft valid horizon
        - A custom callable following LossProtocol

    loss_params : dict, optional
        Additional keyword arguments for the loss function.
    targets_key : str, default="output"
        Name of the readout layer for training targets.
    drivers_keys : list[str], optional
        List of driver input keys in data dict for input-driven models.
    monitor_losses : list[str | LossProtocol], optional
        Additional loss functions to compute and log (but not optimize on).
        Can be loss names (e.g., ``"standard"``, ``"lyapunov"``) or callables.
        Results are stored as trial user attributes with prefix ``monitor_``.
    monitor_params : dict[str, dict[str, Any]], optional
        Keyword arguments for each monitor loss. Keys are loss function names
        (e.g., ``"lyapunov_weighted_loss"``), values are kwargs dicts.
    study_name : str, optional
        Name for the study. If ``None``, auto-generated from *model_creator*.
    storage : str, optional
        Storage path or URL. Behaviour depends on the value:

        - ``None``: In-memory (single worker) or temp journal file (multi-worker).
        - ``"study.log"``: Journal file storage (recommended for multi-worker).
        - ``"study.db"`` or ``"sqlite:///study.db"``: SQLite with WAL mode.

    warm_start : list[dict], optional
        Known-good configurations to evaluate **before** any sampler-proposed
        trial.  Each dict maps parameter name to a fixed value and is enqueued
        via :meth:`optuna.study.Study.enqueue_trial`.  Enqueued in the parent
        process (pre-fork), so the warm-start trials are picked up by every
        worker when ``n_workers > 1``.
    transfer_from : optuna.Study or str, optional
        A prior study (or a storage path/URL holding exactly one study) whose
        ``COMPLETE`` trials are copied into the new study via
        :meth:`optuna.study.Study.add_trial`, seeding the sampler without
        re-running the objective.  Only parameters that overlap with the current
        search space are kept; mismatched parameters are dropped and logged at
        ``INFO`` level.  Ingested in the parent process (pre-fork), so it works
        with ``n_workers > 1``.
    sampler : BaseSampler, optional
        Optuna sampler. Defaults to ``TPESampler`` with ``multivariate=True``.
    pruner : str or optuna.pruners.BasePruner, optional
        Early-stopping policy fed by the per-trial intermediate
        forecast-horizon reports.  Accepts a registry key — ``"asha"`` (ASHA /
        :class:`~optuna.pruners.SuccessiveHalvingPruner`), ``"hyperband"``,
        ``"median"``, ``"threshold"``, or ``"none"`` — a fully-configured
        :class:`optuna.pruners.BasePruner` instance, or ``None``.  ``None`` (the
        default) and ``"none"`` disable pruning (:class:`optuna.pruners.NopPruner`);
        ``"median"`` is a good first choice for forecasting studies because it
        stops trials whose intermediate loss is worse than the running median.
        A diverging configuration reports a large loss at an early
        forecast-horizon checkpoint and is pruned before the full horizon is
        evaluated, saving compute.  Pruning is honored in both single- and
        multi-process modes and coexists with the ``clip_value`` /
        ``prune_on_clip`` path.
    seed : int, optional
        Random seed for reproducibility. Seeds the Optuna sampler and, per
        trial, PyTorch / NumPy / Python ``random`` (``seed + trial.number``).
        The per-trial seed is also threaded into ``model_creator`` when it
        accepts a ``seed`` keyword argument, so the reservoir topology and
        input/feedback initializers — which build their own RNGs and do **not**
        read the legacy global NumPy state — reproduce as well. With a
        ``seed``-aware ``model_creator``, two ``run_hpo`` calls sharing the same
        ``seed`` therefore yield identical trials, parameters, and
        ``best_value``. See *Reproducible studies* in the HPO reference for how
        to make a creator seed-aware.
    device : str or torch.device, optional
        Device to place models and data on (e.g., ``"cuda"`` or
        ``torch.device("cpu")``). If ``None``, uses default device.
    n_workers : int, default=1
        Number of parallel workers. When ``> 1``, uses real OS processes
        (multiprocessing) that coordinate via shared file storage. The workers
        are started with ``fork`` on CPU-only studies, but with ``spawn`` when
        the study touches CUDA (``device`` is a CUDA device, or CUDA is already
        initialised in the parent) because ``fork`` cannot re-initialise CUDA.
        ``spawn`` pickles the objective, so ``model_creator``, ``search_space``,
        ``data_loader`` (and any ``loss`` callable) **must be top-level,
        importable functions** — lambdas or closures defined inside another
        function are not picklable and will raise ``PicklingError``. Single-worker
        (``n_workers=1``) studies run in-process and have no such restriction.
    verbosity : int, default=1
        Logging verbosity: ``0`` = silent, ``1`` = normal, ``2`` = verbose.
    catch_exceptions : bool, default=True
        If ``True``, catch exceptions and return *penalty_value*.
    penalty_value : float, default=1e10
        Value returned for failed trials.
    clip_value : float, optional
        Upper bound for the objective value.  When set and the raw loss
        exceeds this threshold, the value returned to Optuna is clamped to
        *clip_value* (or the trial is pruned, see *prune_on_clip*).  The
        unclipped value is always stored as the ``"raw_loss"`` trial user
        attribute.  When ``None`` (default), no clipping is applied and
        ``raw_loss == loss``.
    prune_on_clip : bool, default=False
        If ``True`` **and** *clip_value* is set, trials whose raw loss exceeds
        *clip_value* are pruned instead of returning the clipped value.
    horizon_key : str, optional
        Key in the ``data_loader`` output dict giving the forecast horizon (an
        ``int``) for the autoregressive phase.  When ``None`` (default), the
        horizon defaults to ``val.shape[1]`` (the number of validation steps).
        Set this to forecast a horizon that differs from the validation window
        length.
    n_checkpoints : int, default=5
        Number of growing forecast-horizon checkpoints at which each trial
        reports an intermediate loss and the configured *pruner* is consulted.
        The single forecast is sliced to each checkpoint (no repeated
        forecasts), and reports are keyed by 1-based checkpoint index.  Values
        ``<= 1`` disable intermediate reporting (one report at the full
        horizon), which also disables checkpoint-based pruning — set this ``>= 2``
        for *pruner* to have any early-stopping signal to act on.
    torch_scoring : bool, default=True
        If ``True``, evaluate the loss on the forecast's device with
        ``torch.Tensor`` inputs, only moving the final scalar to the host, and
        fall back to NumPy when the loss rejects tensors.  The five built-in
        losses are NumPy-based and always take the fallback; a **custom
        torch-native** ``loss`` benefits from staying on-device (no GPU→CPU→NumPy
        round-trip per checkpoint).  Set ``False`` to always score in NumPy.
    trial_callbacks : list of callable, optional
        Per-trial callbacks invoked after each successful evaluation.  Each is
        called as ``callback(trial, context)`` where *context* is a dict with
        ``"params"``, ``"raw_loss"`` and ``"loss"``.  A callback may call
        ``trial.report`` / raise ``optuna.TrialPruned`` to act as an extra
        pruning signal; any *other* exception it raises is logged and swallowed,
        so a callback can never fail a trial.  Under ``n_workers > 1`` with the
        ``spawn`` start method the callbacks must be picklable (top-level
        functions, not lambdas or closures).

    Returns
    -------
    optuna.Study
        The completed study with all trial results.

    Raises
    ------
    ValueError
        If *n_trials* is not positive.
    TypeError
        If *model_creator*, *search_space*, or *data_loader* are not callable.

    Examples
    --------
    Basic usage:

    >>> from resdag.hpo import run_hpo
    >>> from resdag.models import ott_esn
    >>>
    >>> def model_creator(reservoir_size, spectral_radius):
    ...     return ott_esn(
    ...         reservoir_size=reservoir_size,
    ...         feedback_size=3,
    ...         output_size=3,
    ...         spectral_radius=spectral_radius,
    ...     )
    >>>
    >>> def search_space(trial):
    ...     return {
    ...         "reservoir_size": trial.suggest_int("reservoir_size", 100, 500, step=50),
    ...         "spectral_radius": trial.suggest_float("spectral_radius", 0.5, 1.5),
    ...     }
    >>>
    >>> def data_loader(trial):
    ...     # Load your data here
    ...     return {
    ...         "warmup": warmup, "train": train, "target": target,
    ...         "f_warmup": f_warmup, "val": val,
    ...     }
    >>>
    >>> study = run_hpo(
    ...     model_creator=model_creator,
    ...     search_space=search_space,
    ...     data_loader=data_loader,
    ...     n_trials=50,
    ...     loss="efh",
    ... )
    >>> print(f"Best params: {study.best_params}")
    >>> print(f"Best value: {study.best_value}")

    Parallel with 4 workers:

    >>> study = run_hpo(
    ...     model_creator=model_creator,
    ...     search_space=search_space,
    ...     data_loader=data_loader,
    ...     n_trials=100,
    ...     n_workers=4,
    ...     storage="study.log",
    ... )

    See Also
    --------
    LOSSES : Available loss functions.
    get_study_summary : Generate study summary.
    ESNTrainer : Training interface.
    """
    # ── Configure logging ────────────────────────────────────────────
    # Set levels on the ``resdag`` and Optuna loggers only; never call
    # ``logging.basicConfig`` here, which would mutate the host application's
    # root logger and attach handlers as a global side effect.
    _configure_logging(verbosity)

    # ── Validate inputs ──────────────────────────────────────────────
    if n_trials <= 0:
        raise ValueError(f"n_trials must be positive, got {n_trials}")
    if not callable(model_creator):
        raise TypeError("model_creator must be callable")
    if not callable(search_space):
        raise TypeError("search_space must be callable")
    if not callable(data_loader):
        raise TypeError("data_loader must be callable")

    # ── Resolve device ───────────────────────────────────────────────
    if device is not None and not isinstance(device, torch.device):
        device = torch.device(device)

    # ── Resolve storage ──────────────────────────────────────────────
    # ``return_path=True`` hands back the concrete journal/SQLite path so that
    # multi-process dispatch can reconnect workers without reaching into
    # Optuna's private ``_backend._file_path`` attribute.
    resolved_storage, storage_path = resolve_storage(storage, n_workers, return_path=True)

    # ── Resolve loss function ────────────────────────────────────────
    loss_params = loss_params or {}
    base_loss = get_loss(loss)
    resolved_loss = partial(base_loss, **loss_params) if loss_params else base_loss

    loss_name = loss if isinstance(loss, str) else getattr(loss, "__name__", "custom")
    logger.info(f"Using loss function: {loss_name}")
    if loss_params:
        logger.info(f"Loss parameters: {loss_params}")

    # ── Resolve monitor losses ───────────────────────────────────────
    resolved_monitor_losses = None
    resolved_monitor_params = monitor_params
    if monitor_losses:
        resolved_monitor_losses = [get_loss(m) if isinstance(m, str) else m for m in monitor_losses]
        monitor_names = [
            m if isinstance(m, str) else getattr(m, "__name__", "custom") for m in monitor_losses
        ]
        logger.info(f"Monitoring additional losses: {monitor_names}")

        # Remap monitor_params keys from registry keys (e.g. "efh") to
        # function names (e.g. "expected_forecast_horizon") so the objective
        # can look them up by monitor_fn.__name__.
        if monitor_params:
            resolved_monitor_params = dict(monitor_params)
            for user_key, fn in zip(monitor_losses, resolved_monitor_losses):
                if isinstance(user_key, str):
                    func_name = getattr(fn, "__name__", user_key)
                    if user_key != func_name and user_key in resolved_monitor_params:
                        resolved_monitor_params[func_name] = resolved_monitor_params.pop(user_key)

    # ── Configure sampler ────────────────────────────────────────────
    if sampler is None:
        sampler = TPESampler(
            multivariate=True,
            warn_independent_sampling=False,
            seed=seed,
        )
        logger.info("Using TPESampler with multivariate optimization")

    # ── Configure pruner ─────────────────────────────────────────────
    # Resolved once here and threaded into the workers so single- and
    # multi-process runs share the same early-stopping policy.  Defaults to
    # ``NopPruner`` (no pruning) when ``pruner is None``.
    resolved_pruner = resolve_pruner(pruner)
    logger.info(f"Using pruner: {type(resolved_pruner).__name__}")

    # ── Create or load study ─────────────────────────────────────────
    if study_name is None:
        study_name = make_study_name(model_creator)
        logger.info(f"Auto-generated study name: {study_name}")

    study = optuna.create_study(
        direction="minimize",
        study_name=study_name,
        storage=resolved_storage,
        load_if_exists=True,
        sampler=sampler,
        pruner=resolved_pruner,
    )

    completed_trials = len([t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE])
    if completed_trials > 0:
        logger.info(f"Loaded existing study with {completed_trials} completed trials")
        logger.info(f"Best value so far: {study.best_value:.6f}")

    # ── Warm-start and transfer (pre-fork) ───────────────────────────
    # Both mutate the shared storage in the parent process, before any worker
    # is forked, so the enqueued / transferred trials are visible to every
    # worker when ``n_workers > 1``.
    #
    # Transferred trials are *seeds* (already evaluated), not budget: they land
    # in storage as COMPLETE, so the stop target and progress baseline are
    # offset by their count to keep ``remaining`` = the number of *new* trials
    # to run.  Warm-start trials, by contrast, are enqueued for evaluation and
    # consume budget like any sampler-proposed trial.
    n_transferred = 0
    if transfer_from is not None:
        n_transferred = transfer_trials(
            study, transfer_from, param_names=_discover_param_names(search_space)
        )
    if warm_start:
        apply_warm_start(study, warm_start)

    # Stop target and progress baseline shifted past the transferred seeds.
    stop_target = n_trials + n_transferred
    baseline_completed = completed_trials + n_transferred

    # ── Build objective function ─────────────────────────────────────
    objective = build_objective(
        model_creator=model_creator,
        search_space=search_space,
        data_loader=data_loader,
        loss_fn=resolved_loss,
        targets_key=targets_key,
        drivers_keys=drivers_keys,
        horizon_key=horizon_key,
        catch_exceptions=catch_exceptions,
        penalty_value=penalty_value,
        monitor_losses=resolved_monitor_losses,
        monitor_params=resolved_monitor_params,
        device=device,
        seed=seed,
        clip_value=clip_value,
        prune_on_clip=prune_on_clip,
        n_checkpoints=n_checkpoints,
        torch_scoring=torch_scoring,
        trial_callbacks=trial_callbacks,
    )

    # ── Run optimization ─────────────────────────────────────────────
    remaining = max(0, n_trials - completed_trials)
    if remaining > 0:
        logger.info(f"Starting optimization: {remaining} trials remaining")

        if n_workers > 1:
            study, resolved_storage = _dispatch_multiprocess(
                study=study,
                resolved_storage=resolved_storage,
                storage_path=storage_path,
                study_name=study_name,
                objective=objective,
                n_trials=stop_target,
                remaining=remaining,
                completed_trials=baseline_completed,
                n_workers=n_workers,
                seed=seed,
                pruner=resolved_pruner,
                verbosity=verbosity,
                device=device,
            )
        else:
            run_single(
                study=study,
                objective=objective,
                remaining=remaining,
                n_workers=1,
                verbosity=verbosity,
            )

        # Final summary
        done = len([t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE])
        logger.info(f"Optimization completed: {len(study.trials)} total trials")
        if done > 0:
            logger.info(f"Best value: {study.best_value:.6f}")
            logger.info(f"Best parameters: {study.best_params}")
    else:
        logger.info(f"All {n_trials} trials already completed")

    return study

---

Reproducible studies

Pass seed= to run_hpo to make a study reproducible end to end. The seed is applied in two places:

1. The Optuna sampler — so the sequence of sampled hyperparameters is identical across runs.
2. Each trial, where the per-trial seed seed + trial.number seeds PyTorch, NumPy, and Python's random, and is threaded into your model_creator.

That last step is the one that matters for the reservoir. The recurrent topology and the input/feedback initializers build their own RNGs (np.random.default_rng(...) / torch.Generator) and do not read NumPy's legacy global state — so seeding np.random alone leaves the reservoir weights varying run-to-run. To close the gap, make your creator accept a seed keyword and forward it to the reservoir (the premade factories forward **kwargs straight to ESNLayer, which fixes the whole reservoir):

from resdag.models import ott_esn

def model_creator(reservoir_size, spectral_radius, seed=None):
    # `seed` reaches ESNLayer/ESNCell, fixing the topology, the
    # feedback/input initializers, and the random bias.
    return ott_esn(
        reservoir_size=reservoir_size,
        feedback_size=3,
        output_size=3,
        spectral_radius=spectral_radius,
        seed=seed,
    )

study_a = run_hpo(model_creator, search_space, data_loader, n_trials=20, seed=42)
study_b = run_hpo(model_creator, search_space, data_loader, n_trials=20, seed=42)
assert study_a.best_value == study_b.best_value  # identical

run_hpo only injects seed when model_creator declares a seed parameter (or **kwargs); creators without one are called unchanged, and an explicit seed returned by search_space always wins. Reproducibility is exact for single-worker (n_workers=1) studies; with n_workers > 1, trial scheduling across processes is non-deterministic, so the set of evaluated configurations is reproducible but their order — and therefore best_value — may differ.

---

Pruning (early stopping)

Each trial forecasts once and is scored at a handful of growing forecast-horizon checkpoints, reporting the prefix loss at every one. A diverging configuration already shows a large loss at an early checkpoint — so an Optuna pruner can terminate it before paying for the full horizon. Pass pruner= to turn this on:

from resdag.hpo import run_hpo

study = run_hpo(
    model_creator, search_space, data_loader,
    n_trials=100,
    pruner="median",   # stop trials worse than the running median at each checkpoint
)

# Pruned trials are recorded but excluded from `best_value`:
print(get_study_summary(study))   # the summary lists the Pruned count

pruner accepts a registry key, a fully-configured optuna.pruners.BasePruner instance, or None:

Key	Pruner	When to reach for it
"none" (default)	NopPruner	No pruning.
"median"	MedianPruner	Good default — prunes trials worse than the running median.
"asha"	SuccessiveHalvingPruner (ASHA)	Aggressive, scales well with many parallel workers.
"hyperband"	HyperbandPruner	ASHA across several budgets; strong when the right budget is unknown.
"threshold"	ThresholdPruner	Prune on an absolute loss bound (default upper=1.0; pass a configured instance for a different bound).

Need to tune a pruner's knobs? Build the instance and pass it in:

import optuna

study = run_hpo(
    model_creator, search_space, data_loader, n_trials=100,
    pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=2),
)

Pruning is honored in both single- and multi-process (n_workers > 1) runs — the resolved pruner is attached to every worker's study. It also composes with the clip_value / prune_on_clip bound: the clip path prunes on the raw loss exceeding a ceiling, while the pruner reacts to the intermediate horizon reports; both can be active at once.

---

losses

Loss functions for hyperparameter optimization.

This module provides specialized loss functions for evaluating multi-step forecasts in reservoir computing applications. All functions operate on batched predictions with shape (B, T, D) where B is batch size, T is time steps, and D is the number of dimensions.

Available Losses

• "efh" : Expected Forecast Horizon (default, recommended for chaotic systems)
• "forecast_horizon" : Forecast Horizon Loss (contiguous valid steps)
• "lyapunov" : Lyapunov-weighted Loss (exponential decay for chaotic systems)
• "standard" : Standard Loss (mean geometric mean error)
• "soft_horizon" : Soft valid horizon (cumulative survival probability)

Scale conventions

The threshold-based horizon losses (:func:expected_forecast_horizon, :func:forecast_horizon, :func:soft_valid_horizon) compare a per-timestep error against a fixed threshold. For that comparison to be meaningful the error metric must be scale-free, otherwise the same threshold collapses the horizon to 0 or T depending on the absolute scale of the data. All of these losses therefore default to metric="nrmse" (RMSE normalised by the per-dimension standard deviation of y_true). Selecting a raw-scale metric ("rmse", "mse", "mae") for a threshold-based loss emits a :class:UserWarning, because the threshold is then only valid on unit-scale data.

Batch aggregation

Per-timestep errors of shape (B, T) are collapsed across the batch axis with the median in every loss (see :func:_aggregate_batch). The median is robust to outlier trajectories and — unlike the geometric mean — does not collapse to 0 when a single batch element happens to be perfect at a step. When a geometric mean is requested explicitly it is computed on clipped errors so it is zero-safe.

Example

from resdag.hpo import LOSSES, get_loss loss_fn = get_loss("efh") loss = loss_fn(y_true, y_pred, threshold=0.2)

See Also

resdag.hpo.objective : Objective builder that wraps these losses for Optuna. resdag.hpo.run : High-level HPO orchestrator.

LOSSES module-attribute

LOSSES: dict[str, Callable[..., float]] = {'efh': expected_forecast_horizon, 'forecast_horizon': forecast_horizon, 'lyapunov': lyapunov_weighted, 'standard': standard_loss, 'soft_horizon': soft_valid_horizon}

LossProtocol

Bases: Protocol

Protocol for HPO loss functions.

All loss functions must accept y_true and y_pred arrays of shape (B, T, D) as positional-only arguments and return a single float value to minimize.

Examples:

Implementing a custom loss:

>>> def my_loss(y_true, y_pred, /, **kwargs):
...     return float(np.mean((y_true - y_pred) ** 2))

get_loss

get_loss(key_or_callable: str | LossProtocol) -> LossProtocol

Get a loss function by key or return the callable directly.

PARAMETER	DESCRIPTION
key_or_callable	Either a string key from LOSSES (e.g., "efh") or a custom callable following the LossProtocol interface. TYPE: str or callable

RETURNS	DESCRIPTION
LossProtocol	The loss function callable.

RAISES	DESCRIPTION
KeyError	If the string key is not found in LOSSES.
TypeError	If the provided callable doesn't match LossProtocol.

Example

loss_fn = get_loss("efh") loss_fn = get_loss(my_custom_loss)

Source code in src/resdag/hpo/losses.py

def get_loss(key_or_callable: str | LossProtocol) -> LossProtocol:
    """Get a loss function by key or return the callable directly.

    Parameters
    ----------
    key_or_callable : str or callable
        Either a string key from LOSSES (e.g., "efh") or a custom callable
        following the LossProtocol interface.

    Returns
    -------
    LossProtocol
        The loss function callable.

    Raises
    ------
    KeyError
        If the string key is not found in LOSSES.
    TypeError
        If the provided callable doesn't match LossProtocol.

    Example
    -------
    >>> loss_fn = get_loss("efh")
    >>> loss_fn = get_loss(my_custom_loss)
    """
    if isinstance(key_or_callable, str):
        if key_or_callable not in LOSSES:
            available = ", ".join(LOSSES.keys())
            raise KeyError(f"Unknown loss '{key_or_callable}'. Available: {available}")
        return cast(LossProtocol, LOSSES[key_or_callable])

    if not callable(key_or_callable):
        raise TypeError("Loss must be a string key or a callable.")

    return key_or_callable

expected_forecast_horizon

expected_forecast_horizon(y_true: NDArray[floating], y_pred: NDArray[floating], /, metric: MetricType = 'nrmse', threshold: float = 0.2, softness: float = 0.04) -> float

Differentiable proxy for forecast horizon length.

This is the recommended loss for chaotic systems. It provides a smooth, differentiable approximation of the forecast horizon by using a soft threshold. The loss rewards models that keep errors below the threshold for as many consecutive steps as possible.

PARAMETER	DESCRIPTION
y_true	True values of shape (B, T, D). TYPE: ndarray
y_pred	Predicted values of shape (B, T, D). TYPE: ndarray
metric	Error metric to compute. TYPE: (rmse, mse, mae, nrmse) DEFAULT: "rmse"
threshold	Error threshold below which predictions are considered "good". TYPE: float DEFAULT: 0.2
softness	Controls the width of the soft threshold boundary. Smaller values create a harder threshold. Good default is ~10% of threshold. TYPE: float DEFAULT: 0.02

RETURNS	DESCRIPTION
float	Negative expected forecast horizon. Lower (more negative) is better.

WARNS	DESCRIPTION
UserWarning	If metric is not scale-free (i.e. not "nrmse"), since the fixed threshold is then only meaningful on unit-scale data.

Source code in src/resdag/hpo/losses.py

def expected_forecast_horizon(
    y_true: NDArray[np.floating],
    y_pred: NDArray[np.floating],
    /,
    metric: MetricType = "nrmse",
    threshold: float = 0.2,
    softness: float = 0.04,
) -> float:
    """Differentiable proxy for forecast horizon length.

    This is the recommended loss for chaotic systems. It provides a smooth,
    differentiable approximation of the forecast horizon by using a soft
    threshold. The loss rewards models that keep errors below the threshold
    for as many consecutive steps as possible.

    Parameters
    ----------
    y_true : ndarray
        True values of shape (B, T, D).
    y_pred : ndarray
        Predicted values of shape (B, T, D).
    metric : {"rmse", "mse", "mae", "nrmse"}, default="nrmse"
        Error metric to compute.
    threshold : float, default=0.2
        Error threshold below which predictions are considered "good".
    softness : float, default=0.02
        Controls the width of the soft threshold boundary. Smaller values
        create a harder threshold. Good default is ~10% of threshold.

    Returns
    -------
    float
        Negative expected forecast horizon. Lower (more negative) is better.

    Warns
    -----
    UserWarning
        If *metric* is not scale-free (i.e. not ``"nrmse"``), since the fixed
        *threshold* is then only meaningful on unit-scale data.
    """
    _warn_if_scale_unsafe(metric, "expected_forecast_horizon", threshold)
    errors = _compute_errors(y_true, y_pred, metric)  # (B, T)
    e_t = _aggregate_batch(errors, how="median")  # Robust across batch

    # Soft indicator of "good prediction" at each step
    good_t = expit((threshold - e_t) / softness)  # ∈ (0, 1)

    # Survival probability: product of all good indicators up to t
    # log_g = np.log(np.clip(good_t, 1e-12, 1.0))
    surv_t = np.cumprod(good_t)

    # Expected horizon length
    expected_horizon = np.sum(surv_t)

    return -float(expected_horizon)  # Minimize → maximize horizon

forecast_horizon

forecast_horizon(y_true: NDArray[floating], y_pred: NDArray[floating], /, metric: MetricType = 'nrmse', threshold: float = 0.2) -> float

Contiguous valid forecast horizon length.

Counts the number of consecutive time steps where the error stays below the threshold, starting from the beginning of the forecast.

PARAMETER	DESCRIPTION
y_true	True values of shape (B, T, D). TYPE: ndarray
y_pred	Predicted values of shape (B, T, D). TYPE: ndarray
metric	Error metric to compute. Defaults to the scale-free "nrmse" so the fixed threshold is comparable across datasets of any magnitude. TYPE: (rmse, mse, mae, nrmse) DEFAULT: "rmse"
threshold	Error threshold below which predictions are considered valid. TYPE: float DEFAULT: 0.2

RETURNS	DESCRIPTION
float	Negative of the valid horizon length. Lower is better.

WARNS	DESCRIPTION
UserWarning	If metric is not scale-free (i.e. not "nrmse"), since the fixed threshold is then only meaningful on unit-scale data.

Source code in src/resdag/hpo/losses.py

def forecast_horizon(
    y_true: NDArray[np.floating],
    y_pred: NDArray[np.floating],
    /,
    metric: MetricType = "nrmse",
    threshold: float = 0.2,
) -> float:
    """Contiguous valid forecast horizon length.

    Counts the number of consecutive time steps where the error stays below
    the threshold, starting from the beginning of the forecast.

    Parameters
    ----------
    y_true : ndarray
        True values of shape (B, T, D).
    y_pred : ndarray
        Predicted values of shape (B, T, D).
    metric : {"rmse", "mse", "mae", "nrmse"}, default="nrmse"
        Error metric to compute.  Defaults to the scale-free ``"nrmse"`` so the
        fixed *threshold* is comparable across datasets of any magnitude.
    threshold : float, default=0.2
        Error threshold below which predictions are considered valid.

    Returns
    -------
    float
        Negative of the valid horizon length. Lower is better.

    Warns
    -----
    UserWarning
        If *metric* is not scale-free (i.e. not ``"nrmse"``), since the fixed
        *threshold* is then only meaningful on unit-scale data.
    """
    _warn_if_scale_unsafe(metric, "forecast_horizon", threshold)
    errors = _compute_errors(y_true, y_pred, metric)  # (B, T)
    e_t = _aggregate_batch(errors, how="median")  # Robust across batch

    below = e_t < threshold
    if not below[0]:
        valid_len = 0
    else:
        valid_len = int(np.argmax(~below)) if (~below).any() else int(below.size)

    return -float(valid_len)

lyapunov_weighted

lyapunov_weighted(y_true: NDArray[floating], y_pred: NDArray[floating], /, metric: MetricType = 'nrmse', lyapunov_t: int = 64) -> float

Lyapunov-compensated time-weighted error.

Computes a time-weighted error where weights decay exponentially with characteristic time equal to the Lyapunov time. This compensates for expected exponential error growth in chaotic systems by applying a weight of exp(-t / lyapunov_t) at time step t.

At t = lyapunov_t, the weight equals exp(-1).

PARAMETER	DESCRIPTION
y_true	True values of shape (B, T, D). TYPE: ndarray
y_pred	Predicted values of shape (B, T, D). TYPE: ndarray
metric	Error metric to compute per time step. Defaults to the scale-free "nrmse" so the magnitude is comparable across datasets. TYPE: (rmse, mse, mae, nrmse) DEFAULT: "rmse"
lyapunov_t	Characteristic Lyapunov time (in time steps) controlling the exponential decay rate of the weights. TYPE: int DEFAULT: 64

RETURNS	DESCRIPTION
float	Exponentially time-weighted average error. Lower is better.

Notes

Errors are reduced across the batch with a zero-safe geometric mean (clipped to a small positive floor; see :func:_aggregate_batch).

Source code in src/resdag/hpo/losses.py

def lyapunov_weighted(
    y_true: NDArray[np.floating],
    y_pred: NDArray[np.floating],
    /,
    metric: MetricType = "nrmse",
    lyapunov_t: int = 64,
) -> float:
    """Lyapunov-compensated time-weighted error.

    Computes a time-weighted error where weights decay exponentially with
    characteristic time equal to the Lyapunov time. This compensates for
    expected exponential error growth in chaotic systems by applying a
    weight of exp(-t / lyapunov_t) at time step t.

    At t = lyapunov_t, the weight equals exp(-1).

    Parameters
    ----------
    y_true : ndarray
        True values of shape (B, T, D).
    y_pred : ndarray
        Predicted values of shape (B, T, D).
    metric : {"rmse", "mse", "mae", "nrmse"}, default="nrmse"
        Error metric to compute per time step.  Defaults to the scale-free
        ``"nrmse"`` so the magnitude is comparable across datasets.
    lyapunov_t : int, default=64
        Characteristic Lyapunov time (in time steps) controlling the
        exponential decay rate of the weights.

    Returns
    -------
    float
        Exponentially time-weighted average error. Lower is better.

    Notes
    -----
    Errors are reduced across the batch with a zero-safe geometric mean
    (clipped to a small positive floor; see :func:`_aggregate_batch`).
    """
    errors = _compute_errors(y_true, y_pred, metric)  # (B, T)
    e_t = _aggregate_batch(errors, how="gmean")  # (T,)

    t = np.arange(e_t.shape[0])
    weights = np.exp(-t / lyapunov_t)

    return float(np.sum(weights * e_t) / np.sum(weights))

standard_loss

standard_loss(y_true: NDArray[floating], y_pred: NDArray[floating], /, metric: MetricType = 'nrmse') -> float

Mean geometric mean error across all timesteps.

Simple baseline loss suitable for both stable and unstable systems.

PARAMETER	DESCRIPTION
y_true	True values of shape (B, T, D). TYPE: ndarray
y_pred	Predicted values of shape (B, T, D). TYPE: ndarray
metric	Error metric to compute. TYPE: (rmse, mse, mae, nrmse) DEFAULT: "rmse"

RETURNS	DESCRIPTION
float	Mean of geometric mean errors. Lower is better.

Notes

The geometric mean across the batch is zero-safe: errors are clipped to a small positive floor before reduction (see :func:_aggregate_batch) so a single perfect batch element cannot collapse a step to 0.

Source code in src/resdag/hpo/losses.py

def standard_loss(
    y_true: NDArray[np.floating],
    y_pred: NDArray[np.floating],
    /,
    metric: MetricType = "nrmse",
) -> float:
    """Mean geometric mean error across all timesteps.

    Simple baseline loss suitable for both stable and unstable systems.

    Parameters
    ----------
    y_true : ndarray
        True values of shape (B, T, D).
    y_pred : ndarray
        Predicted values of shape (B, T, D).
    metric : {"rmse", "mse", "mae", "nrmse"}, default="nrmse"
        Error metric to compute.

    Returns
    -------
    float
        Mean of geometric mean errors. Lower is better.

    Notes
    -----
    The geometric mean across the batch is zero-safe: errors are clipped to a
    small positive floor before reduction (see :func:`_aggregate_batch`) so a
    single perfect batch element cannot collapse a step to ``0``.
    """
    errors = _compute_errors(y_true, y_pred, metric)
    geom_mean = _aggregate_batch(errors, how="gmean")
    return float(np.mean(geom_mean))

soft_valid_horizon

soft_valid_horizon(y_true: NDArray[floating], y_pred: NDArray[floating], /, metric: MetricType = 'nrmse', threshold: float = 0.3, n: int = 6) -> float

Soft forecast horizon via cumulative survival probability.

At each timestep a soft indicator measures how "good" the prediction is using a Hill-function gate:

.. math::

g_t = \frac{1}{1 + (e_t / \theta)^n}

where :math:e_t is the median error at step t, :math:\theta is the threshold, and n controls the sharpness. A cumulative product of these indicators gives a survival probability that drops to zero once any step fails, mimicking a contiguous forecast horizon:

.. math::

H = \sum_t \prod_{i \leq t} g_i

The function is numerically safe: error ratios are clipped before exponentiation to prevent overflow.

PARAMETER	DESCRIPTION
y_true	True values of shape (B, T, D). TYPE: ndarray
y_pred	Predicted values of shape (B, T, D). TYPE: ndarray
metric	Error metric to compute per timestep. Defaults to the scale-free "nrmse" so the fixed threshold is comparable across datasets. TYPE: (rmse, mse, mae, nrmse) DEFAULT: "rmse"
threshold	Error threshold below which predictions are considered "good". TYPE: float DEFAULT: 0.3
n	Hill exponent controlling gate sharpness. Lower values (4-8) give smoother landscapes suitable for TPE optimisation; higher values (>12) approach a hard step function with poor gradient signal. TYPE: int DEFAULT: 6

RETURNS	DESCRIPTION
float	Negative expected horizon. Lower (more negative) is better.

WARNS	DESCRIPTION
UserWarning	If metric is not scale-free (i.e. not "nrmse"), since the fixed threshold is then only meaningful on unit-scale data.

Notes

Compared to :func:expected_forecast_horizon (which uses a sigmoid gate and batch-median), this loss uses a Hill-function gate and batch-median. The Hill gate has a sharper but still tunable transition and is symmetric around the threshold on a log-error scale.

See Also

expected_forecast_horizon : Sigmoid-gated variant with cumprod survival.

Source code in src/resdag/hpo/losses.py

def soft_valid_horizon(
    y_true: NDArray[np.floating],
    y_pred: NDArray[np.floating],
    /,
    metric: MetricType = "nrmse",
    threshold: float = 0.3,
    n: int = 6,
) -> float:
    """Soft forecast horizon via cumulative survival probability.

    At each timestep a soft indicator measures how "good" the prediction is
    using a Hill-function gate:

    .. math::

        g_t = \\frac{1}{1 + (e_t / \\theta)^n}

    where :math:`e_t` is the median error at step *t*, :math:`\\theta` is the
    threshold, and *n* controls the sharpness.  A cumulative product of these
    indicators gives a survival probability that drops to zero once any step
    fails, mimicking a contiguous forecast horizon:

    .. math::

        H = \\sum_t \\prod_{i \\leq t} g_i

    The function is numerically safe: error ratios are clipped before
    exponentiation to prevent overflow.

    Parameters
    ----------
    y_true : ndarray
        True values of shape (B, T, D).
    y_pred : ndarray
        Predicted values of shape (B, T, D).
    metric : {"rmse", "mse", "mae", "nrmse"}, default="nrmse"
        Error metric to compute per timestep.  Defaults to the scale-free
        ``"nrmse"`` so the fixed *threshold* is comparable across datasets.
    threshold : float, default=0.3
        Error threshold below which predictions are considered "good".
    n : int, default=6
        Hill exponent controlling gate sharpness.  Lower values (4-8) give
        smoother landscapes suitable for TPE optimisation; higher values
        (>12) approach a hard step function with poor gradient signal.

    Returns
    -------
    float
        Negative expected horizon.  Lower (more negative) is better.

    Warns
    -----
    UserWarning
        If *metric* is not scale-free (i.e. not ``"nrmse"``), since the fixed
        *threshold* is then only meaningful on unit-scale data.

    Notes
    -----
    Compared to :func:`expected_forecast_horizon` (which uses a sigmoid gate
    and batch-median), this loss uses a Hill-function gate and batch-median.
    The Hill gate has a sharper but still tunable transition and is symmetric
    around the threshold on a log-error scale.

    See Also
    --------
    expected_forecast_horizon : Sigmoid-gated variant with ``cumprod`` survival.
    """
    _warn_if_scale_unsafe(metric, "soft_valid_horizon", threshold)
    errors = _compute_errors(y_true, y_pred, metric)  # (B, T)
    e_t = _aggregate_batch(errors, how="median")  # (T,) — robust across batch

    # Numerically safe Hill gate: clip the ratio to avoid overflow in x**n
    ratio = np.clip(e_t / threshold, 0.0, 1e4)
    good_t = 1.0 / (1.0 + ratio**n)

    # Survival (cumulative product) — once one step fails, all later are ~0
    surv_t = np.cumprod(good_t)

    horizon = np.sum(surv_t)
    return -float(horizon)

---

utils

Utility functions for hyperparameter optimization.

Provides helpers for study management, naming, and result summarization.

See Also

resdag.hpo.run : High-level HPO orchestrator that uses these utilities.

get_study_summary

get_study_summary(study: Study, top_n: int = 5) -> str

Generate a human-readable summary of an Optuna study.

Creates a formatted text summary including study statistics, best trial information, and top N trials.

PARAMETER	DESCRIPTION
study	The Optuna study to summarize. TYPE: Study
top_n	Number of top-performing trials to include. TYPE: int DEFAULT: 5

RETURNS	DESCRIPTION
str	Formatted multi-line summary string.

Example

study = optuna.create_study()

... run optimization ...

print(get_study_summary(study))

Source code in src/resdag/hpo/utils.py

def get_study_summary(study: optuna.Study, top_n: int = 5) -> str:
    """Generate a human-readable summary of an Optuna study.

    Creates a formatted text summary including study statistics, best trial
    information, and top N trials.

    Parameters
    ----------
    study : optuna.Study
        The Optuna study to summarize.
    top_n : int, default=5
        Number of top-performing trials to include.

    Returns
    -------
    str
        Formatted multi-line summary string.

    Example
    -------
    >>> study = optuna.create_study()
    >>> # ... run optimization ...
    >>> print(get_study_summary(study))
    """
    lines = []
    lines.append("=" * 60)
    lines.append("Study Summary")
    lines.append("=" * 60)
    lines.append(f"Study Name: {study.study_name}")

    total_trials = len(study.trials)
    completed = len([t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE])
    pruned = len([t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED])
    failed = len([t for t in study.trials if t.state == optuna.trial.TrialState.FAIL])

    lines.append(f"Total Trials: {total_trials}")
    lines.append(f"  Completed: {completed}")
    lines.append(f"  Pruned: {pruned}")
    lines.append(f"  Failed: {failed}")
    lines.append("")

    if completed > 0:
        lines.append("-" * 60)
        lines.append("Best Trial")
        lines.append("-" * 60)
        lines.append(f"Trial Number: {study.best_trial.number}")
        lines.append(f"Value: {study.best_value:.6f}")
        lines.append("")
        lines.append("Parameters:")
        for k, v in study.best_params.items():
            if isinstance(v, float):
                lines.append(f"  {k}: {v:.6f}")
            else:
                lines.append(f"  {k}: {v}")
        lines.append("")

        # Top N trials. ``completed_trials`` are all COMPLETE, so ``t.value`` is
        # always a float at runtime; the ``inf`` fallback only satisfies the
        # ``float | None`` stub and never affects ordering (it would sort last).
        completed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
        sorted_trials = sorted(
            completed_trials, key=lambda t: t.value if t.value is not None else float("inf")
        )

        lines.append("-" * 60)
        lines.append(f"Top {min(top_n, len(sorted_trials))} Trials")
        lines.append("-" * 60)
        for i, trial in enumerate(sorted_trials[:top_n], 1):
            lines.append(f"{i}. Trial {trial.number}: {trial.value:.6f}")

    lines.append("=" * 60)
    return "\n".join(lines)

make_study_name

make_study_name(model_creator: Callable) -> str

Generate a stable study name from the model creator callable.

Creates a study name based on the callable's source file and name. It is robust to callables that lack the usual introspection metadata:

• functools.partial objects are unwrapped to their underlying function so the wrapped factory's name and source file are used (a bare partial raises TypeError in :func:inspect.getsourcefile).
• Callables without a usable __name__ (e.g. lambdas or callable class instances) fall back to their type name plus a short stable hash of their repr, so two logically-distinct creators do not collide onto the same study name and silently share persisted storage.

PARAMETER	DESCRIPTION
model_creator	The model creator callable to generate a name from. May be a plain function, a :class:functools.partial, a lambda, or any callable object. TYPE: Callable

RETURNS	DESCRIPTION
str	Study name in the format "filename:identifier". filename is "<unknown>" when the source file cannot be determined.

Example

def my_model_creator(units): ... return model make_study_name(my_model_creator) 'script:my_model_creator'

Source code in src/resdag/hpo/utils.py

def make_study_name(model_creator: Callable) -> str:
    """Generate a stable study name from the model creator callable.

    Creates a study name based on the callable's source file and name. It is
    robust to callables that lack the usual introspection metadata:

    - ``functools.partial`` objects are unwrapped to their underlying function so
      the wrapped factory's name and source file are used (a bare ``partial``
      raises ``TypeError`` in :func:`inspect.getsourcefile`).
    - Callables without a usable ``__name__`` (e.g. lambdas or callable class
      instances) fall back to their type name plus a short stable hash of their
      ``repr``, so two logically-distinct creators do not collide onto the same
      study name and silently share persisted storage.

    Parameters
    ----------
    model_creator : Callable
        The model creator callable to generate a name from. May be a plain
        function, a :class:`functools.partial`, a lambda, or any callable object.

    Returns
    -------
    str
        Study name in the format ``"filename:identifier"``. ``filename`` is
        ``"<unknown>"`` when the source file cannot be determined.

    Example
    -------
    >>> def my_model_creator(units):
    ...     return model
    >>> make_study_name(my_model_creator)
    'script:my_model_creator'
    """
    fn = _unwrap_callable(model_creator)

    try:
        src = inspect.getsourcefile(fn) or "<interactive>"
    except (TypeError, OSError):
        src = "<unknown>"

    name = getattr(fn, "__name__", None)
    if not name or name == "<lambda>":
        digest = hashlib.sha1(repr(model_creator).encode("utf-8")).hexdigest()[:8]
        name = f"{type(fn).__name__}-{digest}"

    return f"{Path(src).stem}:{name}"

get_best_params

get_best_params(study: Study) -> dict

Get the best parameters from a study.

Convenience function that returns the best trial's parameters.

PARAMETER	DESCRIPTION
study	The Optuna study. TYPE: Study

RETURNS	DESCRIPTION
dict	Dictionary of best hyperparameters.

RAISES	DESCRIPTION
ValueError	If no completed trials exist.

Source code in src/resdag/hpo/utils.py

def get_best_params(study: optuna.Study) -> dict:
    """Get the best parameters from a study.

    Convenience function that returns the best trial's parameters.

    Parameters
    ----------
    study : optuna.Study
        The Optuna study.

    Returns
    -------
    dict
        Dictionary of best hyperparameters.

    Raises
    ------
    ValueError
        If no completed trials exist.
    """
    completed = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
    if not completed:
        raise ValueError("No completed trials in study.")
    return study.best_params