Cannot get log-likelihood from Logistic Regression in Pyro

[FrederikWarburg]

I'm trying to get log-likelihood estimates via arviz from my logistic regression model

My model looks like this:

def logistic_regression(X, obs=None):
   alpha = pyro.sample("alpha", dist.Normal(0, 5.)) # Prior for the bias/intercept
   beta  = pyro.sample("beta", dist.Normal(torch.zeros(X.shape[1]), 
                                            5.*torch.ones(X.shape[1]))) # Priors for the regression coeffcients
    
   with pyro.plate("data"):
       logits = alpha + X.matmul(beta)
       y = pyro.sample("y", dist.Bernoulli(logits=logits), obs=obs)
        
   return y

I train it via standard MCMC, like this:

nuts_kernel = NUTS(logistic_regression,jit_compile=True, ignore_jit_warnings=True)
mcmc_logistic = MCMC(nuts_kernel, num_samples=1000, warmup_steps=500, num_chains=2)
mcmc_logistic.run(X_train, y_train)
mcmc_logistic.summary()

And then I would like to convert it to an arviz module.

pyro_data_logit = az.from_pyro(
    mcmc_logistic,
    prior=prior_logit,
    posterior_predictive=posterior_predictive_logit
)

The problem is that this object does not contain a log-likelihood, which seems weird. The missing log-likehood is problematic as I can use other functionalities such as .loo

I'm able to calculate the log-likelihood by hand, like

alpha_hat = samples_logistic["alpha"].mean(axis=0).detach().numpy()
beta_hat = samples_logistic["beta"].mean(axis=0).detach().numpy()
y_hat = sigmoid(alpha_hat + np.dot(X_test, beta_hat))
log_likelihood_logit = stats.binom.logpmf(1, n = 1, p =y_hat)

By arviz.frompyro does not accept manually specified log-likelihood as argument (arviz.frompystan accept this argument).

I'm using arviz version '0.7.0' and pyro version '1.2.1'

[OriolAbril]

Hi! Thanks for reporting.

az.from_pyro uses pyro.Predictive.get_vectorized_trace() to automatically retrieve and store pointwise log likelihood samples. To get the log_likelihood group populated, ArviZ must be able to execute this function, which as stated in Pyro docs:

... requires that the model has all batch dims correctly annotated via plate.

In your case, I think you need to annotate the dimension of beta X.shape[1] with plate to allow ArviZ to automatically populate log_likelihood group.

If this were not to work (I am not a Pyro expert), you can always combine from_pyro with from_dict and do something like:

pyro_data_logit = az.concat(
    az.from_pyro(...), 
    az.from_dict(log_likelihood={"y": log_likelihood_logit}
)

Note that this requires log_likelihood_logit to follow ArviZ dimension order: chain, draw, *shape, otherwise they will not be annotated correctly.