setu.MixedNLE#

class setu.MixedNLE(discrete_net, continuous_flow, combiner, x_transform=None, theta_transform=None, condition_transform=None, x_dim=0, num_continuous=0, num_discrete=0, condition_dim=0, log_transform_continuous=False, trained=False)[source]#

Bases: Module

Mixed Neural Likelihood Estimator for discrete/continuous data.

Combines CategoricalMADE for discrete variables and normalizing flow for continuous variables. Models the joint distribution:

p(x|theta) = p(x_disc|theta) * p(x_cont|x_disc, theta)

Data layout convention: continuous variables FIRST, discrete LAST.

x = [x_cont_0, …, x_cont_k, x_disc_0, …, x_disc_m]

Use create_mixed_nle() to create an instance and fit_mixed_nle() to train it.

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

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

Accepts single inputs or batches. Batched inputs are evaluated with an internal vmap; leading dimensions broadcast against each other (a single theta can be paired with a batch of x, and vice versa).

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) – Mixed observation with shape (x_dim,) or (batch, x_dim). Continuous variables must be in first columns, discrete in last columns.

  • theta (Array) – Parameters with shape (theta_dim,) or (batch, theta_dim).

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

Return type:

Array

Returns:

Log probability as scalar for single inputs, or shape (batch,) if any input is batched.

Raises:

  • ValueError – If MixedNLE has not been trained yet.

  • ValueError – If trailing dimensions or batch sizes don’t match.

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

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

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

Two modes:

  • Batched theta (batch, theta_dim): one sample per row, returns (batch, x_dim). Do not pass n_samples.

  • Single theta (theta_dim,) with n_samples: draws n_samples for the same theta, returns (n_samples, x_dim). This mirrors NLE.sample(theta, key, n_samples=…).

Parameters:

  • theta (Array) – Parameters (batch, theta_dim) or (theta_dim,).

  • key (Array) – JAX random key.

  • conditions (Array | None) – Experimental conditions (batch, condition_dim) or (condition_dim,). Required if MixedNLE was trained with conditions.

  • n_samples (int | None) – Number of samples for a single theta. Required if theta is 1D; must be omitted if theta is batched.

Return type:

Array

Returns:

Samples (batch, x_dim) or (n_samples, x_dim).

Raises:

  • ValueError – If MixedNLE has not been trained yet.

  • ValueError – If theta dimensionality and n_samples are inconsistent.

  • 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, chunk_size=None, name='mixed_nle_likelihood', dims=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). Required if MixedNLE was trained with conditions.

  • chunk_size (int | None) – If provided, process observations in chunks of this size to reduce peak memory usage.

  • name (str) – Name for the PyMC potential

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

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='mixed_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. Discrete columns of x_observed are validated to contain integers.

Parameters:

  • theta – PyMC parameter variable with dims matching hierarchy

  • x_observed – Observed data ([groups], subjects, trials, x_dim) with continuous variables in first columns, discrete in last

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

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

Returns:

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

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

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