setu.compute_log_likelihood#

setu.compute_log_likelihood(estimator, trace, x_observed, theta_var_name, *, conditions=None)[source]#

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

Needed because to_pymc_hierarchical uses a single pm.Potential that sums all log-likelihoods, making pm.compute_log_likelihood() and az.loo() unavailable. This function computes per-observation log-likelihoods from posterior samples and injects them into the trace.

Note: this function mutates trace in-place (adds a log_likelihood group) and returns the same object.

Parameters:

  • estimator (LikelihoodEstimator) – Trained NLE, MixedNLE, or DiscreteNLE.

  • trace (InferenceData) – ArviZ InferenceData with posterior samples.

  • x_observed (Array) – Observed data, shape (n_obs, x_dim).

  • theta_var_name (str) – Name of the per-subject theta variable in the posterior (e.g., "alpha"). Must have a per-subject dimension at axis 2.

  • conditions (Array | None) – Per-observation conditions, shape (n_obs, condition_dim). Required if estimator was trained with conditions.

Return type:

InferenceData

Returns:

The same InferenceData with a log_likelihood group added, compatible with az.loo().