setu.tune_nre#

setu.tune_nre(dataset, *, x_dim=None, theta_dim=None, condition_dim=0, search_space=None, n_trials=50, val_fraction=0.1, batch_size=200, max_epochs=1000, max_patience=30, pruning=True, seed=None, study_name=None, storage=None, show_progress=True, baseline_params=None, key)[source]#

Tune NRE hyperparameters with Optuna.

Searches over classifier architecture, embedding size, number of contrastive classes, and training parameters, returning the best NRE by validation loss.

Requires optuna: pip install setu-sbi[tuning].

Parameters:

  • dataset (SimulationDataset) – Training data (parameters and observations) to tune on.

  • x_dim (int | None) – Observation dimension. Inferred from dataset if None.

  • theta_dim (int | None) – Parameter dimension. Inferred from dataset if None.

  • condition_dim (int) – Per-trial condition dimension; 0 for no conditions.

  • search_space (NRESearchSpace | None) – Hyperparameter ranges to search. A default NRESearchSpace is used if None.

  • n_trials (int) – Number of Optuna trials (architectures) to evaluate.

  • val_fraction (float) – Fraction of dataset held out to compute validation loss.

  • batch_size (int) – Minibatch size used when training each trial. Must be at least twice the upper bound of the num_classes search range.

  • max_epochs (int) – Maximum training epochs per trial.

  • max_patience (int) – Early-stopping patience, in epochs without validation improvement.

  • pruning (bool) – If True, prune unpromising trials early with Optuna’s median pruner.

  • seed (int | None) – Seed for the Optuna sampler, for reproducible searches.

  • study_name (str | None) – Name for the Optuna study.

  • storage (str | None) – Optuna storage URL (e.g. an SQLite path) for persisting the study; in-memory if None.

  • show_progress (bool) – If True, display Optuna’s progress bar.

  • baseline_params (dict[str, Any] | None) – Hyperparameters for a baseline trial enqueued before the search, for comparison against the tuned result.

  • key (Array) – JAX PRNG key for model initialization and training.

Return type:

TuningResult

Returns:

A TuningResult holding the best NRE, its hyperparameters, and the completed Optuna study.

Raises:

ValueError – If batch_size is smaller than twice the maximum num_classes in the search space.