setu.NRE#

class setu.NRE(classifier, theta_embedding=None, x_embedding=None, x_transform=None, theta_transform=None, condition_transform=None, x_dim=0, condition_dim=0, trained=False)[source]#

Bases: Module

Neural Ratio Estimator using NRE-C algorithm.

Estimates log r(x, theta) = log p(x|theta) - log p(x). Use create_nre() to create an instance and fit_nre() to train it.

classifier#: ResNet classifier network.

theta_embedding#: MLP embedding for parameters.

x_embedding#: MLP embedding for observations.

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

trained#: Whether the NRE has been trained.

Parameters:

classifier (ResNetClassifier)
theta_embedding (MLP | None)
x_embedding (MLP | None)
x_transform (Standardizer | None)
theta_transform (Standardizer | None)
condition_transform (Standardizer | None)
x_dim (int)
condition_dim (int)
trained (bool)

log_ratio(x, theta, conditions=None)[source]#

Compute log r(x, theta) = log p(x|theta) - log p(x).

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 NRE was trained with conditions.

Return type:

Array

Returns:

Log ratio as scalar.

Raises:

ValueError – If NRE has not been trained yet.
ValueError – If conditions required but not provided, or vice versa.

to_pymc(theta, x_observed, *, conditions=None, name='nre_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-ratio in chunks to reduce memory.

Returns:

PyMC Potential adding log r(x_observed, theta) to model.

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

Convert to PyMC hierarchical likelihood.

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).
theta_shape (tuple[int, ...] | None) – Explicit theta shape when theta.type.shape contains None.
chunk_size (int | None) – If set, evaluate log-ratio in chunks to reduce memory.
mask (Array | None) – Boolean array matching x_observed leading dims.
trial_dim (str | None) – Name of trial dimension (innermost level). When omitted for a 2D x_observed, a single trial per subject is assumed and the trial axis is inserted automatically.
subject_dim (str | None) – Name of subject dimension.
group_dim (str | None) – Name of group dimension (optional).
name (str) – Name for the PyMC potential.
dims (tuple[str, ...] | None) – PyMC dimension names (for coords).

Returns:

PyMC Potential adding sum of log r(x_observed, theta) to model.

to_numpyro(theta, x_observed, *, conditions=None, name='nre_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='nre_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.