Plotting analytical posterior alongside sampled posterior

General
1

Hi!

I have defined a simple Bayesian linear model from which I have created a trace using pm.sample (NUTS). I have also calculated the analytical posterior. The trace can be plotted using Arviz’ plot_dist, and I have successfully added the analytical posterior to the plot by calling plot on the return value of pc.get_viz. This works fine, but the analytical posterior is not added to the legend and there is also a color conflict. Sure, this could be adjusted manually, but I also want to do other things, like calculate the Wasserstein distance between them, and then they need to be in the same format.

The issue seems to be that the trace created by PyMC is an xarray.DataTree, whose posterior group is used by plot_dist. The group contains the chain of samples, whereas the analytical posterior contains y-values that correspond to an array of x-values.

I tried creating a DataTree out of the analytical posterior as follows:

reference = xr.DataTree.from_dict({
    "posterior": xr.Dataset(
        data_vars={
            "intercept": (("chain","draw"), intercept_analytical),
            "slope": (("chain", "draw"), slope_analytical)
        }
    )
})

but it doesn’t work. I honestly wouldn’t really expect it to, but I also don’t know how to fix it.

Do I need to somehow sample from the analytical solution? Is there another way of doing this?

Thank you to whomever takes the time reading this!