Trace individual dimensions using PYMC3's traceplot?

Question

I have used PYMC3 to perform inference on a Bayesian logistic regression model. I want to find the posterior over the weights $\beta \in \mathbb{R}^K$ given a Gaussian prior $\mathcal{N} \sim (0,I)$ .

This is using real, not simulated data represented by the design matrix $X \in \mathbb{R}^{N \times K}$ and binary outcome vector $y\in \mathbb{R}^N$.

My Python code is as follows:

with pm.Model() as logistic_model:
    # 14 dimensional Gaussian prior
    beta = pm.Normal('beta', mu = 0, sd = 10, shape = X.shape[1])

    # Expected value of outcome with sigmoid link
    logit_p = t.dot(X, beta)
    pm.Bernoulli('likelihood', logit_p=logit_p, observed=y)

    with logistic_model:
    # Inference
    trace = pm.sample(2000)

I then proceed to use

pm.traceplot(trace) 

to get a visualisation of the posteriors. This outputs the following:

I'm curious; how I would go about plotting each dimension of the posterior on separate sub-plots? It looks rather messy having them all on one!

Apologies for the dollar signs - I'm new to this and unsure how to display math correctly.

Thanks

Answer 1

You can use arviz, which can be installed with

pip install arviz

Then you can use

import arviz as az

az.traceplot(trace)

Which works as you describe.

Note arviz is still somewhat unstable, but should replace pymc3's plotting functions within the next few month.