C06 · How many galaxy clusters? — Bayesian Modelling in Practice

The brief

An astronomy group has measured the recession velocities of 82 galaxies in six well-separated conic sections of sky towards Corona Borealis. By Hubble's law velocity is a proxy for distance, so if galaxies were spread evenly through space the velocities would form one smooth hump. Theories of large-scale structure say otherwise: galaxies should sit in clusters separated by voids, which would show up as several humps with gaps in between.

The group asks three questions:

Do the velocities show distinct clusters separated by voids, and how many?
Which galaxies belong together? For the ones near a boundary, how sure are you?
Is the apparent void between the main body and the fastest galaxies real, or the kind of gap that 82 draws from a smooth distribution produce by chance?

They want answers with uncertainty, in language they can put in a paper.

How this notebook works

Each task states what to deliver, not how. Write your code in the YOUR CODE HERE cells.
Stuck? h.hint("task2") reveals hints one level at a time: nudge → approach → code skeleton. Try to get by on nudges.
h.check("task2", mu_main=...) compares your numbers with the reference solution.
A full worked solution lives in notebooks/solutions/. Open it only when you are done (or truly stuck).

import time

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
from scipy import stats

from pymc_challenges import Hints, data

RANDOM_SEED = 1986
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")

h = Hints("C06")
h.tasks()

C06 - Galaxies: how many clusters, and who belongs where?

task0 - Look at the data (2 hints)
task1 - The textbook mixture (3 hints, 1 check(s))
task2 - Fix the model (3 hints, 3 check(s))
task3 - Membership probabilities, two ways (3 hints, 2 check(s))
task4 - How many clusters? (3 hints, 2 check(s))
task5 - Unequal spreads and heavy tails (3 hints, 2 check(s))
task6 - Co-clustering and the void (3 hints, 3 check(s))

The data

A single column, dat: velocity in km/s. We sort it and work in thousands of km/s, so galaxy 0 is the nearest and galaxy 81 the most distant, and sensible prior scales are of order 1-10 rather than 1,000-10,000.

A historical footnote: this is the version distributed with R's MASS package, which contains a famous typo - the 78th value is 26690 where Roeder's paper has 26960. It makes no difference to anything below, but it is the reason published analyses of "the galaxy data" do not always agree to the last digit.

data.describe("galaxies")
galaxies = data.load("galaxies")
v = np.sort(galaxies["dat"].to_numpy() / 1000)
N = len(v)
print(N, "galaxies, from", v.min(), "to", v.max(), "thousand km/s")

galaxies: Velocities (km/s) of 82 galaxies in the Corona Borealis region.
  source: Postman, Huchra & Geller (1986) via Roeder (1990) and R package `MASS`.
  url:    https://vincentarelbundock.github.io/Rdatasets/csv/MASS/galaxies.csv
82 galaxies, from 9.172 to 34.279 thousand km/s

Task 0 · Look at the data

Deliver a plot that shows every single galaxy (82 points deserve better than a 10-bin histogram alone) and a list of the largest gaps between neighbouring velocities. How many groups do you see, how many galaxies in each, and how wide is each group?

Your turn Work this task in the notebook.

Task 1 · The textbook mixture

Fit a mixture of three Gaussians: three means, one standard deviation shared by all components, and three mixture weights. Write it the way a textbook would, with the same prior for every component - nothing in the model should distinguish "component 0" from "component 2". Use the default four chains.

Deliver

The summary table and whatever plots you need to decide whether you can trust it.
A diagnosis. Be precise: what exactly went wrong, is any individual chain wrong, and is there a quantity in the summary table that is perfectly fine - and why that one?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Do not stop at "r_hat is bad". r_hat compares chains with each other. Look at each chain on its own: the trace of the means, and the posterior mean of every component per chain. What is component 0 in chain 0? What is it in chain 2?

2 Approach

pm.NormalMixture("velocity", w=w, mu=mu, sigma=sigma, observed=v) with w = pm.Dirichlet("w", np.ones(3)) and mu = pm.Normal("mu", 20, 10, shape=3) is the whole model. For the per-chain view: idata.posterior["mu"].mean("draw") is a (chain, component) table, and az.plot_trace(idata, var_names=["mu"]) shows it.

The likelihood is unchanged if you swap the names of two components, and so is the prior. So the posterior has 3! = 6 identical copies of every mode. Which quantities in the model do not care about the names?

3 Skeleton

K = 3
with pm.Model(coords={"component": range(K)}) as naive_model:
    mu = pm.Normal("mu", ___, ___, dims="component")
    sigma = pm.HalfNormal("sigma", ___)
    w = pm.Dirichlet("w", np.ones(K), dims="component")
    pm.NormalMixture("velocity", w=w, mu=mu, sigma=sigma, observed=v)
    naive_idata = pm.sample(random_seed=RANDOM_SEED)

az.summary(naive_idata)
naive_idata.posterior["mu"].mean("draw").to_pandas()     # rows: chains, columns: components
float(az.rhat(naive_idata.posterior["mu"]).max())

Self-checks in the notebook (h.check("task1", …)):

max_rhat — largest r_hat among the three component means of the NAIVE model - it is supposed to be terrible (well above 1.1)

Task 2 · Fix the model

Deliver

A three-component mixture with the same likelihood and clean diagnostics. Fix the model, not the sampler settings. If your fix refuses to start sampling, find out which term of the log-probability is to blame before changing anything else.
Evidence for your Task 1 diagnosis: get an acceptable r_hat for the means out of the Task 1 posterior draws, without refitting anything.
One sentence each: does your fix change the fitted density? Would it still work if two clusters had nearly the same mean?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

The six copies of the posterior differ only in which component carries which name. Can you add a constraint that exactly one of the six copies satisfies? And for part 2: if every chain sits in a different copy, what operation on a single posterior draw maps all the copies onto the same one?

2 Approach

Constrain the means to be increasing: pm.Normal("mu", 20, 10, dims="component", transform=pm.distributions.transforms.ordered). The transform works with the first mean and the log of the gaps between neighbours, and the default starting point has all means equal - gaps of zero. model.initial_point() and model.point_logps() show the -inf. Give strictly increasing starting values through pm.sample(initvals={"mu": ...}); spreading them over the range of the data is a good default. (Prefer this to initval= on the variable: pm.compute_log_likelihood, which you need in Task 4, rejects models with non-default initial values.)

Part 2: sort the means within every draw (np.sort(..., axis=-1)), wrap the result back into a DataArray with (chain, draw, component) dims and call az.rhat on it.

3 Skeleton

ordered = pm.distributions.transforms.ordered
with pm.Model(coords={"component": range(K)}) as ordered_model:
    mu = pm.Normal("mu", 20, 10, dims="component", transform=ordered)
    sigma = pm.HalfNormal("sigma", 3)
    w = pm.Dirichlet("w", np.ones(K), dims="component")
    pm.NormalMixture("velocity", w=w, mu=mu, sigma=sigma, observed=v)
    ordered_idata = pm.sample(initvals={"mu": np.linspace(___, ___, K)}, random_seed=RANDOM_SEED)

naive_mu = naive_idata.posterior["mu"]
relabelled = naive_mu.copy(data=np.sort(naive_mu.values, axis=___))
az.rhat(relabelled)

Self-checks in the notebook (h.check("task2", …)):

mu_main — posterior mean of the mean of the big middle cluster, about 21.4 (thousand km/s)
w_main — posterior mean of the weight of the middle cluster, about 0.86
sigma — posterior mean of the shared standard deviation, about 2.1 (thousand km/s)

Task 3 · Which galaxies belong together?

The mixture likelihood you have been using sums over the unknown cluster of each galaxy, so the cluster assignments appear nowhere in the posterior. The astronomers want them.

Deliver

From the Task 2 posterior, $P(\text{galaxy } i \text{ belongs to cluster } k)$ for every galaxy and cluster. Plot it against velocity. Which galaxies are genuinely ambiguous?
The same model written the other way: with an explicit discrete assignment for every galaxy. Compare what the sampler had to do, how long it took, the effective sample sizes and the membership probabilities.
Which formulation would you use by default, and why?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Suppose somebody hands you the true weights, means and standard deviation. What is the probability that a galaxy at 16,084 km/s came from the middle cluster? That is one line of Bayes' rule. Now you do not have the true parameters - you have 4,000 plausible sets.

2 Approach

For one posterior draw, the "responsibility" of component k for galaxy i is w_k * N(v_i | mu_k, sigma) divided by the sum of that over k. Compute it for every draw by broadcasting to an array of shape (draws, galaxies, components) - do it on the log scale and subtract the row maximum before exponentiating - and average over draws.

The explicit version: z = pm.Categorical("z", p=w, dims="galaxy") and pm.Normal("velocity", mu[z], sigma, observed=v). PyMC assigns NUTS to the continuous parameters and a Metropolis-within-Gibbs step to z automatically (nutpie cannot do this, so it falls back to the PyMC sampler). Membership probabilities are then just the frequencies of z == k per galaxy. Compare az.ess of the shared parameters and the wall-clock time.

3 Skeleton

post = az.extract(ordered_idata, var_names=["mu", "sigma", "w"])
mu = post["mu"].values.T[:, None, :]            # (draws, 1, K)
w = post["w"].values.T[:, None, :]
sigma = post["sigma"].values[:, None, None]
log_r = np.log(w) + stats.norm.logpdf(v[None, :, None], mu, sigma)    # (draws, N, K)
log_r -= log_r.max(axis=-1, keepdims=True)
r = np.exp(log_r); r /= r.sum(axis=-1, keepdims=True)
membership = r.mean(axis=___)                   # (N, K)

with pm.Model(coords={"component": range(K), "galaxy": range(N)}) as explicit_model:
    ...                                          # same priors, ordered means
    z = pm.Categorical("z", p=___, dims="galaxy")
    pm.Normal("velocity", mu[___], sigma, observed=v, dims="galaxy")

Self-checks in the notebook (h.check("task3", …)):

p_main_16084 — P(galaxy at 16.084 belongs to the middle cluster) under the shared-sd K=3 model, about 0.95
p_main_26995 — P(galaxy at 26.995 belongs to the middle cluster) under the shared-sd K=3 model, about 0.9

Task 4 · How many clusters?

Three was the eye's answer. Let the data have a say: fit the same shared-spread mixture for K = 2, 3, 4, 5, 6.

Deliver

A table of sampler diagnostics per K (worst r_hat, smallest bulk ESS, divergences). Do not expect all of them to be pretty - say what the ugly ones are telling you.
The az.compare table, and how far you trust it with N = 82.
A posterior predictive check that separates at least one K from the rest: the fitted density over the data, plus a test statistic aimed at something the astronomers care about.
For the largest K: what do the posterior weights of the extra components look like?
Your answer to "how many clusters?", with the honest amount of hedging.

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Three different tools answer three different questions. LOO: does an extra component improve prediction of a held-out galaxy? PPC: can the model reproduce the features the astronomers care about (a far group, empty stretches)? Posterior weights: is the extra component actually being used? Expect them to agree about K = 2 and to shrug about the rest.

2 Approach

Wrap the model in a function of K (ordered means, initvals spread over the data range, pm.compute_log_likelihood at the end) and loop. az.compare({"K=2": idata2, ...}) ranks them; read elpd_diff against dse, and the warning column. With N = 82 a difference of 1-2 standard errors is not a finding.

Overfitted mixtures are hard to sample: when a component is not needed, the posterior has to decide whether to give it no weight, merge it with a neighbour or park it in a gap, and those are separate modes. A raised target_accept helps a little; mediocre r_hat for the larger K is itself a sign of redundant components.

PPC: the mixture density for a draw is sum_k w_k N(x | mu_k, sigma) on a grid - plot its mean and a band over the histogram. For a test statistic use pm.sample_posterior_predictive and count, per replicated survey, the galaxies above 30 (observed: 3).

3 Skeleton

def fit_mixture(K):
    with pm.Model(coords={"component": range(K), "galaxy": range(N)}) as model:
        mu = pm.Normal("mu", 20, 10, dims="component", transform=ordered)
        sigma = pm.HalfNormal("sigma", 3)
        w = pm.Dirichlet("w", ___, dims="component")
        pm.NormalMixture("velocity", w=w, mu=mu, sigma=sigma, observed=v, dims="galaxy")
        idata = pm.sample(initvals={"mu": ___},      # K ordered, well-separated starting means
                          target_accept=0.95, random_seed=RANDOM_SEED)
        pm.compute_log_likelihood(idata)
    return model, idata

fits = {f"K={k}": fit_mixture(k) for k in range(2, 7)}
az.compare({name: idata for name, (_, idata) in fits.items()})
az.loo(___).elpd - az.loo(___).elpd

Self-checks in the notebook (h.check("task4", …)):

elpd_gain_3_vs_2 — elpd_loo of K=3 minus elpd_loo of K=2 (shared sd), clearly positive - roughly 15-25
sigma_k6 — posterior mean of the shared sd in the K=6 model, about 1 - half of what it is for K=3

Task 5 · Clusters are not all the same size

One standard deviation for a group of 7 tightly packed galaxies and for a main body that spans 10,000 km/s is an assumption of convenience. Relax it, in two steps, staying with K = 3:

each cluster gets its own standard deviation;
additionally, the components get heavy tails (Student-t with a small, fixed number of degrees of freedom), the usual recipe for making a mixture robust to stragglers.

Deliver

A prior for the cluster spreads that an astronomer would sign off on. (Rich galaxy clusters have velocity dispersions of the order of 1,000 km/s.)
Both fits with their diagnostics. More flexible mixtures have more ways to go wrong: watch the initial values, and if a fit misbehaves work out what it is undecided about.
A LOO comparison with all the Task 4 models. What was K = 6 really doing?
What changed for the boundary galaxies of Task 3?

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Go back to your Task 0 table: the near group has a standard deviation several times smaller than the main body. A shared sigma has to compromise. What could a K = 6 model with a small shared sigma be doing that a K = 3 model with three sigmas does directly? And for the heavy-tailed version: if a component is allowed stragglers far from its centre, are three galaxies at 32-34 still a "cluster"?

2 Approach

Give sigma a dims="component" and an informative prior on the scale astronomers expect, e.g. pm.LogNormal("sigma", 0, 0.5) - median 1 (thousand km/s), 95% within roughly 0.4-2.7. A very wide prior on the spreads invites a degenerate solution in which one component becomes a flat background; bad starting values (all means inside the main body) invite another in which the chains disagree. Start the means spread over the data range.

Heavy tails: pm.Mixture("velocity", w=w, comp_dists=pm.StudentT.dist(nu=4, mu=mu, sigma=sigma, shape=K), observed=v). If r_hat complains here, look at the trace of mu chain by chain: this is no longer label switching (the means are ordered) but two different explanations of the data.

3 Skeleton

with pm.Model(coords={"component": range(3), "galaxy": range(N)}) as unequal_model:
    mu = pm.Normal("mu", 20, 10, dims="component", transform=ordered)
    sigma = pm.LogNormal("sigma", ___, ___, dims="component")
    w = pm.Dirichlet("w", np.ones(3), dims="component")
    pm.NormalMixture("velocity", w=w, mu=mu, sigma=sigma, observed=v, dims="galaxy")
    # heavy tails instead:
    # pm.Mixture("velocity", w=w, observed=v, dims="galaxy",
    #            comp_dists=pm.StudentT.dist(nu=4, mu=mu, sigma=sigma, shape=3))
    idata = pm.sample(initvals={"mu": np.linspace(v.min(), v.max(), 3)},
                      target_accept=0.95, random_seed=RANDOM_SEED)
    pm.compute_log_likelihood(idata)

In the responsibilities, sigma now has a component axis: post["sigma"].values.T[:, None, :].

Self-checks in the notebook (h.check("task5", …)):

sigma_near — posterior mean of the sd of the NEAR cluster (around 9.7) in the unequal-sd Gaussian model, about 0.6
sigma_main — posterior mean of the sd of the MAIN cluster in the unequal-sd Gaussian model, about 2.2

Task 6 · What the astronomers can say

Two deliverables for the paper. Use the model you would defend, and show how much the answers move under a credible alternative.

1. Co-clustering. Cluster labels are a modelling convenience; "are galaxies i and j in the same cluster?" is not. Deliver the posterior co-clustering probability for all pairs (a picture) and for the pairs listed below (numbers).

2. The void. No galaxy was observed between 27,500 and 31,500 km/s. Deliver the posterior predictive mass of that interval, the number of galaxies a new survey of the same size should expect to find in it, and the probability that such a survey finds it empty again. Compare with what a single smooth (one-component) population would predict. Then write the two or three sentences that go in the paper.

VOID = (27.5, 31.5)
PAIRS = [(16.084, 16.170), (16.084, 10.406), (16.084, 18.419), (26.995, 25.633), (26.995, 32.065)]

Your turn Work this task in the notebook.

Hints for this task — open one at a time

1 Nudge

Both deliverables are functions of the parameters, so both have a posterior - no new model is needed. For co-clustering: given the parameters, are the assignments of two different galaxies dependent? For the void: what is the probability that one new galaxy lands in the interval, and what does that imply for 82 of them?

2 Approach

Co-clustering: given theta the assignments are independent, so P(z_i = z_j | theta) = sum_k r_ik * r_jk with the responsibilities of Task 3; average over draws. np.einsum("sik,sjk->ij", r, r) / n_draws does all pairs at once.

Void: the predictive mass of [lo, hi] for one draw is sum_k w_k * (Phi((hi - mu_k) / sigma_k) - Phi((lo - mu_k) / sigma_k)). The expected count in a survey of N galaxies is N * mass; the probability that the interval is empty is (1 - mass) ** N. Compute these per draw and then average - (1 - mean) ** N is not the same thing. For the smooth alternative, a single Normal with the sample mean and sd is enough.

3 Skeleton

r = responsibilities(final_idata)                       # (draws, N, K)
cc = np.einsum("sik,sjk->ij", r, r) / r.shape[0]        # (N, N)
i, j = np.argmin(np.abs(v - 16.084)), np.argmin(np.abs(v - 18.419))
cc[i, j]

post = az.extract(final_idata, var_names=["mu", "sigma", "w"])
mu, sigma, w = post["mu"].values.T, post["sigma"].values.T, post["w"].values.T   # (draws, K)
mass = (w * (stats.norm.cdf((VOID[1] - mu) / sigma) - stats.norm.cdf(___))).sum(axis=1)
N * mass.mean(), ((1 - mass) ** N).mean()

Self-checks in the notebook (h.check("task6", …)):

coclust_16084_18419 — P(galaxies at 16.084 and 18.419 are in the same cluster) under the K=3 unequal-sd model - essentially 1
void_expected — expected number of galaxies inside the void (27.5-31.5) in a new survey of 82 under the K=3 unequal-sd model - about 0.8 with an informative prior on the sds, up to about 1.6 with a vague one
p_void_empty — posterior predictive probability that a new survey of 82 finds the void empty, K=3 unequal-sd model - about 0.5 (0.3 with a vague prior on the sds)

Going further

The famous typo: set the 78th velocity to Roeder's 26960 and rerun. Which of your numbers move, and which conclusions?
Fit an overfitted mixture, K = 10 with a sparse Dirichlet(0.1) prior on the weights, and report the posterior of the number of components holding at least one galaxy's worth of weight. Expect the sampler to struggle, and decide how much of the output you believe.
Verify the invariance claim: compute the co-clustering matrix from the naive Task 1 fit, broken r_hat and all, and compare it with the ordered model's.
Repeat the void analysis for the other gap (10,400-16,100 km/s), and do the look-elsewhere correction properly: simulate surveys from the single-Normal model and record the emptiest 4,000 km/s window in each.

h.progress()

C06: 0/20 hints used, 13/13 checks passed.

Reference solution

A complete, executed solution with commentary on every task — including the ways the obvious model fails. Try the hint ladder first.

Open the worked solution →

#The brief

#How this notebook works

#The data

#Task 0 · Look at the data

#Task 1 · The textbook mixture

#Task 2 · Fix the model

#Task 3 · Which galaxies belong together?

#Task 4 · How many clusters?

#Task 5 · Clusters are not all the same size

#Task 6 · What the astronomers can say

#Going further

Reference solution

The brief

How this notebook works

The data

Task 0 · Look at the data

Task 1 · The textbook mixture

Task 2 · Fix the model

Task 3 · Which galaxies belong together?

Task 4 · How many clusters?

Task 5 · Clusters are not all the same size

Task 6 · What the astronomers can say

Going further