Make diagnostic plots work on population parameters as an option

[sschildhauer]

It will be important to visualize traceplots, rhat values, etc. for the population parameters moving forward.