Plotting posteriors creatively from labelled arrays

The Palmer penguins are small, but they are multi-dimensional by nature: three species on three islands, two sexes, four measurements. A posterior for them is naturally an array with the dims (chain, draw, species, sex, measurement), and almost every useful display is a question about some slice, contrast, ranking or re-arrangement of that array.

This notebook is a gallery. We fit one modest multivariate model, then build a dozen displays, each from a few xarray operations and each answering a stated question. The point is not the pictures themselves but the habit: express the question by name (species, sex, measurement, sample), let xarray line up the dimensions, and the plotting code shrinks to a loop over labels. Siblings: E42 (dims, coords, broadcasting and time), E43 (transforming posteriors with coords); E16 is the general catalogue of uncertainty displays.

#1 · The data as a labelled Dataset

penguins_raw is the field record (D01 cleans it in detail). We need six columns: a short species name, island, sex, and the four measurements with readable names.

to_xarray() turns each column into a variable along penguin; set_coords then demotes the labels (species, island, sex) from data to coordinates, so they ride along with every selection and reduction and can be grouped on. to_dataarray("measurement") stacks the four measurements into one array with a new, labelled measurement dimension - the shape the model will want. A non-dimension coord records the units.

A multi-key groupby gives the design as a labelled table: one UniqueGrouper per coord, and the result is a species x island array. Adelie penguins live on all three islands; Chinstraps only on Dream, Gentoos only on Biscoe, so species and island cannot be separated for them and the model below uses species only.

Two penguins have no measurements at all and nine more have measurements but no recorded sex. .where(..., drop=True) splits them off with their coordinates attached; we will come back to the nine unsexed birds at the end, because the model can tell us their sex. We also drop the units coord from here on: later we rename measurement to measurement_b to build matrices, and a coord called units along both dims would clash.

Standardise by name. Each measurement is centred and scaled by its own mean and sd over all penguins. y.mean("penguin") has dim measurement only, so the subtraction broadcasts over penguins automatically - and the same two arrays convert every posterior quantity back to millimetres and grams later (z * scale + center). For a Bayesian this matters twice: one prior fits all four measurements, and "0.5" means half a population sd in every panel.

#2 · One multivariate model with named dimensions

Each penguin's four standardised measurements are one draw from a multivariate normal whose mean depends on species and sex and whose covariance depends on species:

$$ \begin{aligned} \mathbf z_i &\sim \text{MvNormal}\big(\boldsymbol\mu_{s[i],\,x[i]},\ \Sigma_{s[i]}\big), \qquad \Sigma_s = \text{diag}(\boldsymbol\sigma_s)\, R_s\, \text{diag}(\boldsymbol\sigma_s) \\ \mu_{s,x,m} &\sim \text{Normal}(0, 1), \qquad \sigma_{s,m} \sim \text{HalfNormal}(1), \qquad R_s \sim \text{LKJ}(2) \end{aligned} $$

In standardised units a group mean within about two population sds of the overall mean is a weak but sensible prior, and LKJ(2) gently prefers weaker correlations. The covariance matters for the gallery: it is what makes predictive ellipses tilt and what ties the four measurements of one penguin together. Each species gets its own LKJCholeskyCov; we stack the three Cholesky factors into one (species, 4, 4) tensor and index it per penguin, so the likelihood is a single MvNormal with dims (penguin, measurement). The correlation matrix gets two dims, measurement and measurement_b - a dim name cannot appear twice.

Clean and quick: r_hat at most 1.01, over 2,500 effective draws for every mean, sd and correlation, no divergences (the correlation matrices' constant diagonals have no r_hat and are skipped). From here on the model is done; the rest is xarray.

The one before/after. Which Gentoo body mass is this? The positional line needs you to remember that mu is (chain, draw, species, sex, measurement) and that Gentoo is 2, male is 1 and body mass is 3. The named line says what it means, and survives a new dimension.

#Everything the gallery needs, in five lines of arithmetic

az.extract stacks chains and draws into one sample dim. Then, by name:

• back to physical units: mu * scale + center (the measurement dims line up);
• covariance in units: corr * sd * sd_b, where sd_b is sd with its dim renamed to measurement_b - an outer product with no reshaping;
• Cholesky factors per draw and species: np.linalg.cholesky lifted by xr.apply_ufunc, with the two matrix dims declared as core dims (everything else is looped over);
• a new penguin of every species and sex: $\boldsymbol\mu + L\boldsymbol\varepsilon$, where xr.dot(L, eps, dim="measurement_b") is the matrix-vector product over the named dim.

The new penguins are the posterior predictive distribution: parameter uncertainty plus penguin-to-penguin variation. The means mu carry only the first.

#3 · xarray's own FacetGrid, straight from the posterior

Question: on which measurements do the species actually differ, and by how much? The posterior of mu in standardised units puts all four measurements on one scale (population sds). A kernel density per (species, sex, measurement) is one xr.apply_ufunc with vectorize=True: the function sees one 1-D array of draws at a time and returns a density on a grid, and xarray puts the labels back. Then .plot.line(hue=..., col=..., row=...) lays out the 2 x 4 grid with no plotting loop at all.

Read across a row: Gentoo is the odd one out on three of the four measurements - about 1.5 population sds above the other two species on flipper length and body mass and about 1.7 sds below them on bill depth - while Adelie and Chinstrap are within about half an sd of each other on those three (and overlap on bill depth) and far apart only on bill length, where Chinstraps are the longest of all. Read down a column: males sit above females on every measurement in every species. Every curve is narrow (about 0.2 sd wide): with 34-73 penguins per group the means are well determined. This is the posterior of the average penguin; the spread of individual penguins comes in section 4.

The second FacetGrid shows the other half of the model: the posterior mean correlation matrix of each species. xr.corr computes the same matrix from the raw data, ignoring species, by correlating z with a renamed copy of itself; xr.concat along the existing species dim with a new label adds it as a fourth panel. (xarray's 2-D plots refuse string coordinates, so we swap in positions 0-3 and put the names back as ticks.)

Within every species the four measurements are positively but modestly correlated (a bigger penguin is a little bigger all over): 0.05-0.31 for Adelie, up to 0.44 (flipper-mass) for Chinstrap, 0.27-0.39 for Gentoo. Pooled over species the picture changes completely: flipper and mass correlate at 0.87 and bill depth turns negative with the other three (down to -0.58): Gentoos are big with shallow bills. That is Simpson's paradox in one picture, and the reason the model has a covariance per species rather than one for everybody.

#4 · Ordered displays: ridgelines, rainclouds and a stacked group dim

Many displays want one categorical axis rather than two. stack(group=("species", "sex")) merges the two dims into one (a pandas MultiIndex). Plain string labels are easier to plot and to rename, so the helper below swaps the MultiIndex for "Adelie female"-style labels and keeps species and sex as non-dimension coords for colouring. sortby then orders the groups by any labelled quantity - here the posterior median body mass.

Ridgeline: how big is the average penguin of each group, in real units? One panel per measurement, one ridge per group in the same (body-mass) order, each ridge a posterior density scaled to the same height. The x axes are in millimetres and grams, which is what a field biologist would quote.

The ordering by body mass (bottom = lightest) makes the story of section 3 readable in physical units: Gentoo males (about 5.5 kg) and females (about 4.7 kg) are far ahead; the four Adelie and Chinstrap groups sit between 3.4 and 4.0 kg with Adelie and Chinstrap males overlapping. Bill length reorders the groups completely (Chinstraps longest), and bill depth puts the heavy Gentoos last. Ridges are narrow because these are means.

Raincloud: where does a single penguin fall? A raincloud stacks three layers per group: the "cloud" is the density of new penguins (posterior predictive), the thick bar is the 90% interval of the mean with its median as a white dot, and the "rain" is the observed penguins, jittered. The gap between the bar and the cloud is the difference between estimating a mean and predicting an individual - the most common confusion in reading posterior plots.

The 90% bars for the means are 150-250 g wide; the middle 90% of a cloud, a single penguin, spans about 1,000 g. An Adelie male and a Chinstrap male have almost the same predictive cloud, and a Gentoo female overlaps the heavier Adelie and Chinstrap males. In bill length the species separate much better (Adelie clouds mostly below 45 mm, Chinstrap above 42 mm), though an Adelie male and a Chinstrap female can still be confused. The rain sits inside the clouds, which is a quick visual posterior predictive check per group.

#5 · Contrasts and rankings: renaming a dim to compare every pair

Question: for every pair of groups, how sure are we that one is bigger? The trick is to make two copies of the same array that differ only in the name of the group dim. mug.rename(group="g1") - mug.rename(group="g2") then broadcasts into all 6 x 6 pairs for every measurement and draw - 36 contrasts x 4 measurements x 4000 draws in one line. (diff > 0).mean("sample") is the posterior probability that the row group is larger.

Almost every cell is saturated (0 or 1): with this much data, the ordering of group means is settled for most pairs. The informative cells are the few in between, and the matrix finds them without us having to guess which pairs to test: Adelie vs Chinstrap of the same sex in bill depth (0.58 for females, 0.83 for males: a coin flip and a lean), Adelie males vs Chinstrap females in flipper length (0.70), Adelie vs Chinstrap males in body mass (0.94) and Chinstrap vs Gentoo females in bill length (0.97). Because each matrix is ordered by body mass, the mass panel is a clean triangle, and the other panels show at a glance how differently the other measurements order the same groups.

Ranks: who is biggest, and how sure are we of the whole ordering? Per draw and per measurement we rank the six groups (1 = largest). xarray's .rank needs the optional bottleneck package, which is not installed here, so we lift scipy.stats.rankdata with apply_ufunc (the group dim is both the input and output core dim). Comparing the ranks with a rank array broadcasts into a (measurement, group, rank) table of probabilities.

A bump chart puts the four measurements on the x axis and rank on the y axis. Each line is a group through its posterior mean rank; each circle's area is the probability of that rank. Ties in the posterior show up as two half-size circles in one column.

The chart shows why "biggest penguin" is not one question. Gentoo males lead on flipper length and body mass, are 2nd in bill length and drop to 5th on bill depth, behind all four Adelie and Chinstrap groups; Chinstrap males are the longest-billed and (probably) the deepest-billed. The real doubts are where circles split: in bill depth Chinstrap and Adelie males share ranks 1-2 (83:17) and their females ranks 3-4 (58:42), and in flipper length Adelie males and Chinstrap females share ranks 4-5 (70:30). Everywhere else the ordering of the means is settled. A ranking is a derived quantity like any other: compute it per draw, then summarise.

#6 · Parallel coordinates: the posterior is a joint distribution

Question: which parameters drive the uncertainty in a derived quantity? Every display so far showed one quantity at a time. A draw, though, is a vector: its four means belong together. In parallel coordinates each draw is one polyline across the four measurements (each axis standardised over the draws so they share a scale). Colouring the lines by a derived quantity threads it through every axis: where the colours sort on an axis, that parameter drives the derived quantity; where they mix, it does not. Our derived quantity is the bill ratio, mean length / mean depth, a shape index (long, narrow bills score high), computed per draw by plain division of two .sels. We use Chinstrap females, the smallest group (34 penguins) and so the widest posterior.

The colours sort on the first two axes, in opposite directions: red (high ratio) lines are high on bill length and low on bill depth, blue ones the reverse, so they cross in an X between the two axes. On flipper length and body mass the colours are mixed: the ratio's uncertainty comes from the two bill means alone. The lines also cross a lot everywhere, i.e. the means are only weakly correlated a posteriori. A small calculation says what to expect: with $n$ penguins a mean's sampling covariance is $\Sigma / n$, so the posterior correlation of the four means should look like the within-species correlation of individuals. xr.corr over sample (between means, all 4000 draws) next to the model's correlation matrix (between penguins) checks it:

The two matrices agree to within 0.02, as the $\Sigma / n$ argument predicts. Why this matters: a derived quantity that combines several means (a bill ratio, a "size index", a difference of two measurements) must be computed per draw, as we did; combining the marginal intervals of its ingredients would ignore exactly the dependence this plot shows.

#7 · Posterior predictive ellipses, contours and an animation

Question: what do new penguins of each group look like in two dimensions, and how sure is the model about that shape? For a bivariate normal the $p$ contour is the unit circle stretched by the Cholesky factor: $\boldsymbol\mu + r_p L \,(\cos\theta, \sin\theta)$ with $r_p^2$ the $\chi^2_2$ quantile. In xarray that is the unit circle as an array with dims (measurement_b, angle), the radii as an array with dim level, and one xr.dot - an ellipse for every level, group and posterior draw at once, with no loop and no eigenvalues.

Left: 40 posterior draws of the 90% predictive ellipse per group over the observed penguins (circles female, triangles male), and the pooled least-squares line that ignores species. Right: the tidy route for a 2-D density - the predictive draws of new penguins go to seaborn's kdeplot for contours enclosing 50% and 90% of new penguins, per species.

On the left, the pooled line slopes down (longer bills are shallower) while every group's ellipse tilts up - Simpson's paradox again, now in data space. The spaghetti shows how well each ellipse's size and tilt is known: tight bundles for the Adelie and Gentoo groups, a little looser for the smaller Chinstrap groups (34 birds each). On the right, Gentoos form their own island of long flippers and high mass, while the Adelie and Chinstrap contours sit almost on top of each other: in these two measurements the two species are the same kind of penguin.

#Hypothetical outcome plot: one posterior draw per frame

Instead of drawing 40 ellipses at once, show them one at a time (Hullman et al. 2015; see E16). Each frame is one posterior draw: its six 90% ellipses and one replicated colony, i.e. a posterior predictive copy of all 333 penguins. The replicate comes from vectorised selection: mu.sel(species=y.species, sex=y.sex) looks up each observed penguin's group and returns an array along penguin - the named version of mu[sp_idx, sx_idx]; the matching Cholesky factor comes the same way, and xr.dot makes the correlated noise.

Played at 400 ms per frame, the ellipses shift only a little (most for the Chinstrap groups) while the replicated colony reshuffles completely each time: most of the variation you see is penguins, not parameters. A static band makes that distinction only to readers who already know it.

#8 · Quantile dotplots: twenty penguins you can count

Question: I will weigh the next penguin I catch. How likely is it to be over 4 kg? .quantile(levels, dim="sample") with 20 equally spaced levels turns each group's predictive distribution into 20 equally likely penguins, returned with a quantile dim. Stacked into a dotplot (Kay et al. 2016), "the probability" becomes "count the dots right of the line": each dot is one penguin in twenty, 5%.

Count, do not estimate: roughly half of the next Adelie males and a bit fewer Chinstrap males will weigh over 4 kg, almost every Gentoo will, and very few Adelie or Chinstrap females. The dot counts match the exact probabilities in the titles to within one dot (5%), which is all the precision a reader can use.

#9 · The tidy hand-off: to_dataframe() and seaborn

Some displays are simply what seaborn is built for - here a posterior predictive check as ECDFs of 30 replicated datasets against the data, faceted by species and measurement. To get there we make the replicates in xarray (vectorised .sel again, for 30 draws) and xr.concat them with the observed data along a new named dim source; to_dataframe() then gives exactly the long table seaborn wants, with species and sex as ordinary columns because they were coordinates.

In every one of the 12 panels the black observed ECDF runs inside (or along the edge of) the blue band of 30 replicates: at the level of each species' marginal distributions, the multivariate normal with sex-specific means reproduces the data. (A mixture of two sexes with one mean would show up here as a too-smooth replicate band around a two-humped observed curve.)

#10 · Interactive: every group x measurement cell with the numbers on hover

The 24 group means fit in a 6 x 4 heatmap. Colour shows the posterior mean in population sds (so all columns share a scale); hovering shows the numbers a reader would want to quote - the mean and 90% interval in physical units and the probability that males are larger - all computed by name beforehand and handed to plotly as arrays with matching (group, measurement) order.

Section 3 on one screen: the two Gentoo rows are red or pink (large) everywhere except bill depth, the Adelie female row is blue except bill depth, and Chinstrap males are the reddest cell in bill length. One trap on the way from xarray to any plotting library: .values drops the names, so the dimension order must be fixed by name (.transpose("group", "measurement")) right before it. Without that line this cell ran silently and drew a wrong heatmap with the axes swapped and two rows empty.

#11 · A posterior over a missing label: sexing the unsexed penguins

Nine penguins were measured but their sex was not recorded. The model says what males and females of each species look like, so for each of them and each posterior draw we can compute the probability of being male from Bayes' rule (with even prior odds):

$$ \log\frac{P(\text{male}\mid \mathbf z)}{P(\text{female}\mid \mathbf z)} = -\tfrac12 \left(d^2_{\text{male}} - d^2_{\text{female}}\right), \qquad d^2_x = (\mathbf z - \boldsymbol\mu_{s,x})^\top \Sigma_s^{-1} (\mathbf z - \boldsymbol\mu_{s,x}) $$

(the determinants cancel because both sexes share the species covariance). The xarray version reads like the formula: vectorised .sel(species=...) picks each bird's species parameters, the residual keeps a sex dim, and np.linalg.solve lifted by apply_ufunc with matrix and vector core dims gives $\Sigma^{-1}\mathbf r$ for every bird, sex and draw.

Seven of the nine are called with near certainty: six females (P(male) below 0.07 across the draws) and one male, #10. Penguin #12 is probably female (median 0.18, 90% 0.08-0.37), and #269 is a genuine toss-up whose bar runs from 0.17 to 0.74: the model cannot even say how uncertain it is, because that bird sits between the Gentoo ellipses and small changes in the parameters move it from one side to the other. The width of each bar is the parameter uncertainty; its position is the bird. This is the honest way to fill in a missing label: impute a probability (or better, keep the bird in the model with sex as a latent variable) rather than a hard guess.

#12 · The idioms on one page

#Try it yourself

1. Island within species. Adelie penguins live on all three islands. Give Adelie means an island dim (the others have one island each), and redo the P(row > column) matrix for the five species-island groups. Do Adelies differ between islands?
2. Rank the individuals, not the means. Redo the bump chart with the new-penguin draws new instead of mu. How much less certain is "a random Gentoo female is heavier than a random Adelie male" than the same statement about means?
3. A size-free shape. Compute the bill ratio (length / depth) per draw, for means and for new penguins, and make a raincloud of it by group. Which species is easiest to identify from bill shape alone, and does the answer change between means and individuals?