setu.tune_mixed_nle#

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

Tune MixedNLE hyperparameters with Optuna.

Searches over the continuous-flow and CategoricalMADE architecture parameters, returning the best MixedNLE by validation loss.

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

Parameters:

  • dataset (SimulationDataset) – Training data (parameters and mixed 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.

  • num_discrete (int) – Number of discrete observation columns (the trailing columns, per the continuous-first layout).

  • num_categories (Array) – Number of categories per discrete column, shape (num_discrete,).

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

  • search_space (MixedNLESearchSpace | None) – Hyperparameter ranges to search. A default MixedNLESearchSpace 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.

  • 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 MixedNLE, its hyperparameters, and the completed Optuna study.