setu.NLE#

class setu.NLE(flow, x_transform=None, theta_transform=None, condition_transform=None, x_dim=0, condition_dim=0, trained=False)[source]#

Bases: Module

Neural Likelihood Estimator using normalizing flows.

Estimates p(x|theta) or p(x|theta, conditions) using a conditional normalizing flow. Use create_nle() to create an instance and fit_nle() to train it.

flow#

FlowJAX conditional normalizing flow.

x_transform#

Standardizer for observations (fitted during training).

theta_transform#

Standardizer for parameters (fitted during training).

condition_transform#

Standardizer for conditions (fitted during training).

x_dim#

Dimension of observations.

condition_dim#

Dimension of experimental conditions (0 if none).

Parameters:
log_prob(x, theta, conditions=None)[source]#

Evaluate log p(x|theta) or log p(x|theta, conditions).

Note the argument order: log_prob takes x first, theta second, whereas to_pymc(theta, x_observed) and sample(theta, key, ...) take theta first. If x_dim == theta_dim, swapping the two silently returns a wrong (but finite) value rather than erroring, so pass them by keyword when in doubt.

Parameters:

  • x (Array) – Observation with shape (x_dim,).

  • theta (Array) – Parameter with shape (theta_dim,).

  • conditions (Array | None) – Experimental conditions with shape (condition_dim,). Required if NLE was trained with conditions.

Return type:

Array

Returns:

Log probability as scalar.

Raises:

  • ValueError – If NLE has not been trained yet.

  • ValueError – If conditions required but not provided, or vice versa.

  • ValueError – If conditions has wrong dimension.

sample(theta, key, conditions=None, *, n_samples)[source]#

Sample from p(x|theta) or p(x|theta, conditions).

Draws n_samples observations for a single parameter vector. The signature matches MixedNLE.sample (theta first, key second, n_samples keyword-only); unlike MixedNLE this does not accept a batched theta.

Parameters:

  • theta (Array) – Parameter with shape (theta_dim,).

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

  • conditions (Array | None) – Experimental conditions with shape (condition_dim,). Required if NLE was trained with conditions.

  • n_samples (int) – Number of samples to draw (keyword-only).

Return type:

Array

Returns:

Samples with shape (n_samples, x_dim).

Raises:

  • ValueError – If NLE has not been trained yet.

  • ValueError – If theta is not a single 1D parameter vector.

  • ValueError – If conditions required but not provided, or vice versa.

  • ValueError – If conditions has wrong dimension.

  • TypeError – If key is not a JAX random key (e.g. an integer passed positionally where key is expected).

to_pymc(theta, x_observed, *, conditions=None, name='nle_likelihood', dims=None, chunk_size=None)[source]#

Convert to PyMC likelihood.

Parameters:

  • theta – PyMC parameter variable

  • x_observed – Observed data (n_obs, x_dim) or (x_dim,)

  • conditions (Array | None) – Optional experimental conditions (n_obs, condition_dim)

  • name (str) – Name for the PyMC potential

  • dims (tuple[str, ...] | None) – PyMC dimension names (for coords)

  • chunk_size (int | None) – If set, evaluate log-prob in chunks to reduce memory.

Returns:

PyMC Potential adding log p(x_observed | theta, conditions) to model

to_pymc_hierarchical(theta, x_observed, *, conditions=None, trial_dim=None, subject_dim=None, group_dim=None, name='nle_likelihood_hierarchical', theta_shape=None, dims=None, chunk_size=None, mask=None)[source]#

Convert to PyMC hierarchical likelihood.

Automatically handles 2-level or 3-level hierarchies via nested vmap.

Parameters:

  • theta – PyMC parameter variable with dims matching hierarchy

  • x_observed – Observed data ([groups], subjects, trials, x_dim)

  • conditions (Array | None) – Optional conditions ([groups], subjects, trials, condition_dim)

  • trial_dim (str | None) – Name of trial dimension (innermost level)

  • subject_dim (str | None) – Name of subject dimension (middle level)

  • group_dim (str | None) – Name of group dimension (outermost level, optional)

  • name (str) – Name for the PyMC potential

  • theta_shape (tuple[int, ...] | None) – Explicit theta shape when theta.type.shape contains None (e.g., when using PyMC dims). Inferred from x_observed if not given.

  • dims (tuple[str, ...] | None) – PyMC dimension names (for coords)

  • chunk_size (int | None) – If set, evaluate log-prob in chunks to reduce memory.

  • mask (Array | None) – Boolean array matching x_observed leading dims. Where False, trial logp is zeroed (for ragged/unbalanced data).

Note

When sampling with BlackJAX using chain_method="vectorized", set progressbar=False in pm.sample(). The BlackJAX progress bar uses IO callbacks that are incompatible with vectorized chains.

Returns:

PyMC Potential adding sum of log p(x_observed | theta, conditions) to model

Parameters:

  • conditions (Array | None)

  • trial_dim (str | None)

  • subject_dim (str | None)

  • group_dim (str | None)

  • name (str)

  • theta_shape (tuple[int, ...] | None)

  • dims (tuple[str, ...] | None)

  • chunk_size (int | None)

  • mask (Array | None)

See setu.integration.pymc.to_pymc_hierarchical for full documentation.

to_numpyro(theta, x_observed, *, conditions=None, name='nle_likelihood')[source]#

Convert to NumPyro likelihood factor.

Parameters:

  • theta – JAX array from numpyro.sample(), shape (theta_dim,).

  • x_observed – Observed data (n_obs, x_dim) or (x_dim,).

  • conditions (Array | None) – Optional conditions (n_obs, condition_dim).

  • name (str) – Name for the NumPyro factor site.

to_numpyro_hierarchical(theta, x_observed, *, conditions=None, mask=None, name='nle_likelihood_hierarchical')[source]#

Convert to NumPyro hierarchical likelihood factor.

Parameters:

  • theta – Parameters from numpyro.sample(). 2-level: shape (n_subjects, theta_dim). 3-level: shape (n_groups, n_subjects, theta_dim).

  • x_observed – Observed data. 2-level: shape (n_subjects, n_trials, x_dim). 3-level: shape (n_groups, n_subjects, n_trials, x_dim).

  • conditions (Array | None) – Optional conditions matching x_observed leading dims.

  • mask (Array | None) – Boolean mask matching x_observed leading dims.

  • name (str) – Name for the NumPyro factor site.