Nutpie doesn't compute element-wise log-likelihood

[AlexAndorra]

The elemwise log-likelihood is not stored in the InferenceData that nutpie returns, even when asking for it. The following for instance doesn't error out, but doesn't add a log_likelihood group to the trace (whereas it does when using the default PyMC sampler):

y = np.array([28, 8, -3, 7, -1, 1, 18, 12])
sigma = np.array([15, 10, 16, 11, 9, 11, 10, 18])
J = len(y)

with pm.Model() as pooled:
    mu = pm.Normal("mu", 0, sigma=10)
    obs = pm.Normal("obs", mu, sigma=sigma, observed=y)

    trace_p = pm.sample(nuts_sampler="nutpie", idata_kwargs={"log_likelihood": True}) # doesn't store
    trace_p_ = pm.sample(idata_kwargs={"log_likelihood": True}) # does store

In PyMC it's not that big of a deal (although it adds friction to the user workflow), as one can just do:

with pooled:
    pm.compute_log_likelihood(trace_p)

But that may be a small issue for Bambi users, which are usually less advanced (cc @tomicapretto). They'd have to do pooled.compute_log_likelihood(trace_p), which takes much more time to compute

[djlacombe]

I was searching for this exact issue when using Bambi alongside nutpie. If I use the following code, the log-likelihood is not calcualted:

model.fit(nuts_sampler='nutpie',draws=20000,cores=10,idata_kwargs={"log_likelihood": True})

I also tried the following:

model.fit(nuts_sampler='nutpie',draws=20000,cores=10,idata_kwargs={"log_likelihood": True, "extend_inference_data":True)

and this did not work either.

Thanks for opening up this issue.

[AlexAndorra]

Glad this was useful @djlacombe . I'm actually wondering it this was solved by a recent PR on Bambi (still shouldn't work on when using PyMC directly) 🤔

Could you update to this latest version of Bambi, try it out and report back please?

[aseyboldt]

Thanks for reminding me.
Sounds like we should add the elemwise logp values as a deterministics (based on an argument to the compilation function).
Then we can sort the results into the correct arviz section after sampling.

[aseyboldt]

#74 is related.

[djlacombe]

@AlexAndorra

I updated Bambi to the 0.15 version and ran the two lines of code in my original post separately, i.e.:

model.fit(nuts_sampler='nutpie',draws=20000,cores=10,idata_kwargs={"log_likelihood": True})

and

model.fit(nuts_sampler='nutpie',draws=20000,cores=10,idata_kwargs={"log_likelihood": True, "extend_inference_data":True)

which sampled just fine but did not produce the log-posterior in the results structure. I think it's actually being calculated because the sampling was finished pretty quickly, it's just that it's not being saved.

[AlexAndorra]

Ok so that means it's on nutpie's side, as @aseyboldt was saying

[AlexAndorra]

Sounds like we should add the elemwise logp values as a deterministics (based on an argument to the compilation function).
Then we can sort the results into the correct arviz section after sampling.

That doesn't sound to hard. @djlacombe , do you feel like trying a PR out?

[djlacombe]

@AlexAndorra I appreciate the confidence, but I'm not sure if I have the skills to do this.

[AlexAndorra]

That's why we're here for @djlacombe -- answer your questions and guide you along the way.
Unless @aseyboldt thinks it's too complex for a beginner-level issue