setu.fit_mixed_nle#

setu.fit_mixed_nle(mixed_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 MixedNLE on simulation data.

Uses early stopping based on validation loss.

Parameters:

  • mixed_nle (MixedNLE) – Untrained MixedNLE instance from create_mixed_nle().

  • dataset (SimulationDataset) – Training data with theta-x pairs. x must have continuous columns first, discrete columns last.

  • learning_rate (float) – AdamW optimizer learning rate.

  • weight_decay (float) – L2 regularization weight. Prevents memorization of discrete variable patterns.

  • max_epochs (int) – Maximum training epochs (safety limit).

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

MixedNLETrainingResult

Returns:

MixedNLETrainingResult containing trained MixedNLE, losses, and best epoch.