setu.fit_nle#

setu.fit_nle(nle, dataset, *, learning_rate=0.001, max_epochs=1000, max_patience=50, batch_size=256, val_fraction=0.1, show_progress=True, epoch_callback=None, key)[source]#

Train NLE on simulation data.

Uses early stopping based on validation loss. Training stops when validation loss doesn’t improve for max_patience epochs. Returns the best model (lowest validation loss).

Parameters:

  • nle (NLE) – Untrained NLE instance from create_nle().

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

  • learning_rate (float) – Adam optimizer learning rate.

  • max_epochs (int) – Maximum training epochs. Early stopping typically ends sooner.

  • 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 (used for early stopping).

  • 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:

TrainingResult

Returns:

TrainingResult containing trained NLE, losses, and best epoch.

Raises:

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