setu.fit_nre#

setu.fit_nre(nre, dataset, *, learning_rate=0.0005, max_epochs=1000, max_patience=50, batch_size=200, val_fraction=0.1, num_classes=20, gamma=1.0, gradient_clip=5.0, show_progress=True, epoch_callback=None, key)[source]#

Train NRE on simulation data using NRE-C loss.

Uses early stopping based on validation loss.

Parameters:

  • nre (NRE) – Untrained NRE instance from create_nre().

  • dataset (SimulationDataset) – Training data with theta-x pairs, optionally with conditions.

  • learning_rate (float) – Adam optimizer learning rate.

  • max_epochs (int) – Maximum training epochs.

  • max_patience (int) – Epochs without improvement before early stopping.

  • batch_size (int) – Mini-batch size for training.

  • val_fraction (float) – Fraction of data for validation.

  • num_classes (int) – Number of contrastive classes K for NRE-C.

  • gamma (float) – Weight ratio of joint to marginal in NRE-C loss.

  • gradient_clip (float) – Maximum gradient norm for clipping.

  • show_progress (bool) – Show tqdm progress bar.

  • epoch_callback (Callable[[int, float, float], None] | None) – Called after each epoch with (epoch, train_loss, val_loss). Can raise exceptions to stop training early.

  • key (Array) – JAX random key for training.

Return type:

NRETrainingResult

Returns:

NRETrainingResult containing trained NRE, losses, and best epoch.

Raises:

ValueError – If NRE condition_dim doesn’t match dataset condition_dim.