API reference#

Every public symbol lives in the flat setu namespace: import it as from setu import NLE regardless of which submodule defines it. The tables below group the surface by purpose for readability; each symbol is documented once, and the linked page carries the full signature, parameters, and return type.

Neural Likelihood Estimation (NLE)#

Learn the conditional density \(p(x \mid \theta)\) with a normalizing flow.

`NLE`	Neural Likelihood Estimator using normalizing flows.
`TrainingResult`	Result of NLE training.
`create_nle`	Create NLE with sensible defaults.
`fit_nle`	Train NLE on simulation data.

Mixed Neural Likelihood Estimation (MixedNLE)#

NLE for mixed discrete + continuous data (a flow for the continuous part, a CategoricalMADE for the discrete part).

`MixedNLE`	Mixed Neural Likelihood Estimator for discrete/continuous data.
`MixedNLETrainingResult`	Result of MixedNLE training.
`create_mixed_nle`	Create MixedNLE with sensible defaults.
`fit_mixed_nle`	Train MixedNLE on simulation data.

Discrete Neural Likelihood Estimation (DiscreteNLE)#

NLE for pure-discrete observations: a categorical head for bounded counts, or a CountMADE (ZINB / NegBinom / Poisson) for count data.

`DiscreteNLE`	Discrete Neural Likelihood Estimator for categorical or count data.
`DiscreteNLETrainingResult`	Result of DiscreteNLE training.
`create_discrete_nle`	Create a DiscreteNLE for categorical or count data.
`fit_discrete_nle`	Train DiscreteNLE on simulation data.

Neural Ratio Estimation (NRE)#

Learn a likelihood-to-evidence ratio with a classifier instead of a normalized density.

`NRE`	Neural Ratio Estimator using NRE-C algorithm.
`NRETrainingResult`	Result of NRE training.
`create_nre`	Create NRE with sensible defaults.
`fit_nre`	Train NRE on simulation data using NRE-C loss.

Data#

Simulation datasets and shape helpers. See Shape conventions for the (G, S, T, *E) conventions these follow.

`SimulationDataset`	Dataset of (theta, x) simulation pairs with optional conditions.
`ValidationDataset`	Dataset for estimator validation with optional conditions.
`simulate_dataset`	Generate a SimulationDataset from a simulator and prior.
`one_hot`	One-hot encode integer IDs.
`broadcast_conditions`	Broadcast subject-level conditions to trial-level for hierarchical models.

Serialization#

Save and load trained estimators (Equinox leaves plus a JSON metadata sidecar).

`save_estimator`	Save a trained NLE, MixedNLE, or DiscreteNLE to disk.
`load_estimator`	Load a trained NLE, MixedNLE, or DiscreteNLE from disk.

Transforms#

Z-scoring transforms applied to parameters and observations.

Standardizer

Z-scoring transform for data standardization.

Metrics#

Sample-based distances for comparing estimator draws against ground truth.

`c2st`	Classifier Two-Sample Test using MLP.
`mmd`	Maximum Mean Discrepancy with Gaussian (RBF) kernel.
`sliced_wasserstein`	Sliced Wasserstein distance.
`total_variation`	Total variation distance between two discrete sample distributions.

Simulators#

Reference simulators and their analytical likelihoods, used in tests and tutorials.

`linear_gaussian`	Sample x from linear Gaussian likelihood: x ~ N(theta + shift, cov).
`sample_prior`	Sample theta from multivariate normal prior.
`analytical_log_likelihood`	Compute analytical log p(x\|theta) for linear Gaussian.
`analytical_log_ratio`	Compute analytical log ratio log r(x, theta) = log p(x\|theta) - log p(x).
`analytical_marginal_log_prob`	Compute analytical log p(x) for linear Gaussian.

PyMC integration#

Export a trained estimator as a PyMC likelihood term. See Shape conventions for the (theta, x_observed) argument order.

`to_pymc`	Convert trained likelihood estimator (NLE or MixedNLE) to PyMC likelihood.
`to_pymc_hierarchical`	Convert trained likelihood estimator (NLE or MixedNLE) to PyMC hierarchical likelihood.
`mixed_nle_to_pymc`	Convert trained MixedNLE to PyMC likelihood.
`discrete_nle_to_pymc`	Convert trained DiscreteNLE to PyMC likelihood.
`nre_to_pymc`	Convert trained NRE to PyMC likelihood using GenericLogpOp.
`nre_to_pymc_hierarchical`	Convert trained NRE to PyMC hierarchical likelihood.

NumPyro integration#

Export a trained estimator as a NumPyro factor.

`to_numpyro`	Add NLE/MixedNLE log-likelihood as a factor in the active NumPyro model.
`to_numpyro_hierarchical`	Add hierarchical NLE/MixedNLE log-likelihood as a factor in NumPyro.
`discrete_nle_to_numpyro`	Add DiscreteNLE log-likelihood as a factor in the active NumPyro model.
`nre_to_numpyro`	Add NRE log-ratio as a factor in the active NumPyro model.
`nre_to_numpyro_hierarchical`	Add hierarchical NRE log-ratio as a factor in NumPyro.

Leave-one-out cross-validation#

compute_log_likelihood

Compute per-observation log-likelihoods and add to arviz trace.

Tuning#

Optuna-based hyperparameter search. Requires the tuning extra (pip install setu-sbi[tuning]).

`tune_nle`	Tune NLE hyperparameters with Optuna.
`tune_nre`	Tune NRE hyperparameters with Optuna.
`tune_mixed_nle`	Tune MixedNLE hyperparameters with Optuna.
`TuningResult`	Result of hyperparameter tuning.
`NLESearchSpace`	Search space for NLE hyperparameter tuning.
`NRESearchSpace`	Search space for NRE hyperparameter tuning.
`MixedNLESearchSpace`	Search space for MixedNLE hyperparameter tuning.

Predictive checks#

Prior and posterior predictive sampling.

`prior_predictive`	Run prior predictive simulations organized by condition group.
`posterior_predictive`	Run posterior predictive simulations.
`PriorPredictiveResult`	Result of prior predictive simulations, one dataset per condition.

Validation#

Diagnostics that confirm a learned likelihood is trustworthy before export.

`validate_nle`
`validate_estimator`
`ValidationConfig`	Configuration for validation check thresholds and parameters.
`ValidationResult`	Complete estimator validation diagnostics.
`ValidationError`	Raised when validation fails in strict mode.