setu.fit_discrete_nle#

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

Train DiscreteNLE on simulation data.

Uses early stopping based on validation loss. No x-standardization is applied (discrete data should not be z-scored).

Parameters:

  • discrete_nle (DiscreteNLE) – Untrained DiscreteNLE from create_discrete_nle().

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

  • learning_rate (float) – AdamW optimizer learning rate.

  • weight_decay (float) – L2 regularization weight.

  • max_epochs (int) – Maximum training epochs.

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

  • batch_size (int) – Mini-batch size.

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

  • 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).

  • key (Array) – JAX random key.

Return type:

DiscreteNLETrainingResult

Returns:

DiscreteNLETrainingResult with trained estimator and losses.