Image segmentation with an unknown number of regions, checked against people

import io
import logging
import math
import time
import warnings

import arviz as az
import matplotlib.pyplot as plt
import numba as nb
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
import xarray as xr
from IPython.display import HTML, Image, display
from matplotlib import animation
from scipy import ndimage, stats
from scipy.special import gammaln

from pymc_challenges import data

RANDOM_SEED = 75
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)
warnings.filterwarnings("ignore", category=RuntimeWarning, module="arviz")
pd.set_option("display.width", 170)
pd.set_option("display.max_columns", 20)


def show_jpeg(fig, quality=85, dpi=90):
    """Photographs compress badly as PNG: show figures of photographs as JPEG."""
    buf = io.BytesIO()
    fig.savefig(buf, format="jpeg", dpi=dpi, pil_kwargs={"quality": quality})
    plt.close(fig)
    display(Image(buf.getvalue()))

M_SRGB = np.array([[0.4124, 0.3576, 0.1805], [0.2126, 0.7152, 0.0722], [0.0193, 0.1192, 0.9505]])
D65_WHITE_2DEG = np.array([0.95047, 1.0, 1.08883])


def xyz_to_lab(XYZ):
    d = 6 / 29
    f = np.where(XYZ > d**3, np.cbrt(np.maximum(XYZ, 1e-12)), XYZ / (3 * d**2) + 4 / 29)
    return np.stack([116 * f[..., 1] - 16, 500 * (f[..., 0] - f[..., 1]), 200 * (f[..., 1] - f[..., 2])], -1)


def srgb_to_lab(rgb):
    lin = np.where(rgb <= 0.04045, rgb / 12.92, ((rgb + 0.055) / 1.055) ** 2.4)
    return xyz_to_lab(lin @ M_SRGB.T / D65_WHITE_2DEG)


def lab_to_srgb(lab):
    d = 6 / 29
    fy = (lab[..., 0] + 16) / 116
    f = np.stack([fy + lab[..., 1] / 500, fy, fy - lab[..., 2] / 200], -1)
    XYZ = np.where(f > d, f**3, 3 * d**2 * (f - 4 / 29)) * D65_WHITE_2DEG
    lin = np.clip(XYZ @ np.linalg.inv(M_SRGB).T, 0, 1)
    return np.where(lin <= 0.0031308, 12.92 * lin, 1.055 * lin ** (1 / 2.4) - 0.055)


def mean_colour_paint(img, lbl):
    """Paint every region of a label map in its mean colour (averaged in CIELAB)."""
    lab = srgb_to_lab(img)
    lbl = np.unique(lbl, return_inverse=True)[1].reshape(lbl.shape)
    cnt = np.bincount(lbl.ravel())
    mean = np.stack([np.bincount(lbl.ravel(), lab[..., k].ravel()) / cnt for k in range(3)], 1)
    return lab_to_srgb(mean[lbl])


def label_boundary(lbl):
    """Boundary pixels of a label map: the right or lower neighbour has another label."""
    b = np.zeros(lbl.shape, bool)
    b[:, :-1] |= lbl[:, :-1] != lbl[:, 1:]
    b[:-1, :] |= lbl[:-1, :] != lbl[1:, :]
    return b


def thick(m):
    """Thicken 1-pixel boundary maps for display only."""
    return ndimage.grey_dilation(m, size=(2, 2))


R_TOL = 2.0     # matching tolerance in pixels at half resolution (BSDS: 0.0075 x diagonal = 2.2)


def near(bmap, r=R_TOL):
    """Pixels within distance r of a boundary pixel."""
    if not bmap.any():
        return np.zeros(bmap.shape, bool)
    return ndimage.distance_transform_edt(~bmap) <= r


def human_maps(i):
    hb = [label_boundary(human[i, a]) for a in range(n_human[i])]
    freq = np.mean([near(b) for b in hb], 0)       # fraction of people with a boundary within R_TOL
    return hb, freq


SHOW = int(np.where(ids == "228076")[0][0])       # the photograph we study in detail
fig, axs = plt.subplots(2, 4, figsize=(15, 5.6), layout="none")
fig.subplots_adjust(wspace=0.03, hspace=0.12, left=0.01, right=0.99, top=0.93, bottom=0.01)
for ax, i in zip(axs.ravel(), TEST):
    ax.imshow(images[i])
    ax.set_title(f"{ids[i]} - {n_human[i]} people", fontsize=9)
    ax.axis("off")
fig.suptitle("The eight test photographs (BSDS500 test split, half resolution)")
show_jpeg(fig)

@nb.njit
def _slic_assign(lab, centres, S, m, labels, dist):
    H, W, _ = lab.shape
    dist[:] = np.inf
    for c in range(centres.shape[0]):
        cy, cx = centres[c, 0], centres[c, 1]
        y0, y1 = max(0, int(cy - 2 * S)), min(H, int(cy + 2 * S) + 1)
        x0, x1 = max(0, int(cx - 2 * S)), min(W, int(cx + 2 * S) + 1)
        for y in range(y0, y1):
            for x in range(x0, x1):
                dc = 0.0
                for k in range(3):
                    dc += (lab[y, x, k] - centres[c, 2 + k]) ** 2
                d = dc + ((y - cy) ** 2 + (x - cx) ** 2) / S**2 * m**2
                if d < dist[y, x]:
                    dist[y, x] = d
                    labels[y, x] = c


@nb.njit
def _enforce_connectivity(labels, min_size):
    """Relabel 4-connected pieces; pieces smaller than min_size join a neighbouring piece."""
    H, W = labels.shape
    out = -np.ones((H, W), np.int64)
    qy, qx = np.empty(H * W, np.int64), np.empty(H * W, np.int64)
    dy, dx = np.array([-1, 1, 0, 0]), np.array([0, 0, -1, 1])
    new = 0
    for y0 in range(H):
        for x0 in range(W):
            if out[y0, x0] >= 0:
                continue
            adj = -1
            for k in range(4):
                yy, xx = y0 + dy[k], x0 + dx[k]
                if 0 <= yy < H and 0 <= xx < W and out[yy, xx] >= 0:
                    adj = out[yy, xx]
            lab = labels[y0, x0]
            out[y0, x0] = new
            qy[0], qx[0] = y0, x0
            head, tail = 0, 1
            while head < tail:
                y, x = qy[head], qx[head]
                head += 1
                for k in range(4):
                    yy, xx = y + dy[k], x + dx[k]
                    if 0 <= yy < H and 0 <= xx < W and out[yy, xx] < 0 and labels[yy, xx] == lab:
                        out[yy, xx] = new
                        qy[tail], qx[tail] = yy, xx
                        tail += 1
            if tail < min_size and adj >= 0:
                for t in range(tail):
                    out[qy[t], qx[t]] = adj
            else:
                new += 1
    return out


def slic(lab, S=8, m=10.0, n_iter=10):
    """SLIC superpixels (Achanta et al. 2012): k-means in (y, x, L, a, b), searched locally."""
    H, W, _ = lab.shape
    cy, cx = np.meshgrid(np.arange(S / 2, H, S), np.arange(S / 2, W, S), indexing="ij")
    cy, cx = cy.ravel(), cx.ravel()
    centres = np.column_stack([cy, cx, lab[cy.astype(int), cx.astype(int)]])
    labels, dist = np.zeros((H, W), np.int64), np.zeros((H, W))
    Y, X = np.mgrid[:H, :W]
    for _ in range(n_iter):
        _slic_assign(lab, centres, S, m, labels, dist)
        cnt = np.bincount(labels.ravel(), minlength=len(centres)).astype(float)
        ok = cnt > 0
        for j, v in enumerate([Y, X, lab[..., 0], lab[..., 1], lab[..., 2]]):
            centres[ok, j] = np.bincount(labels.ravel(), v.ravel(), minlength=len(centres))[ok] / cnt[ok]
    return _enforce_connectivity(labels, S * S // 4)


FEATS = ["L*", "a*", "b*", "texture (sd of L*)"]


def superpixels(img, S=8):
    """Superpixel map, features (mean L*a*b* and the sd of L* inside), pixel counts, and the
    adjacency graph with shared boundary lengths."""
    lab = srgb_to_lab(img)
    sp = slic(lab, S=S)
    cnt = np.bincount(sp.ravel())
    means = [np.bincount(sp.ravel(), lab[..., k].ravel()) / cnt for k in range(3)]
    sd = np.sqrt(np.maximum(np.bincount(sp.ravel(), (lab[..., 0] ** 2).ravel()) / cnt - means[0] ** 2, 0))
    a = np.concatenate([sp[:, :-1].ravel(), sp[:-1, :].ravel()])
    b = np.concatenate([sp[:, 1:].ravel(), sp[1:, :].ravel()])
    keep = a != b
    n = sp.max() + 1
    key, blen = np.unique(np.minimum(a[keep], b[keep]) * n + np.maximum(a[keep], b[keep]), return_counts=True)
    return dict(sp=sp, X=np.column_stack(means + [sd]), npix=cnt, edges=np.column_stack([key // n, key % n]),
                blen=blen)


t0 = time.time()
G = {i: superpixels(images[i]) for i in range(len(ids))}
print(f"superpixels for {len(ids)} photographs in {time.time() - t0:.1f} s; "
      f"{min(len(g['X']) for g in G.values())}-{max(len(g['X']) for g in G.values())} superpixels each, "
      f"{np.median([g['npix'].mean() for g in G.values()]):.0f} pixels on average")
g = G[SHOW]
print(f"photograph {ids[SHOW]}: {len(g['X'])} superpixels, {len(g['edges'])} neighbouring pairs")

@nb.njit
def _logdet_chol(A):
    """log det of a small SPD matrix by an in-place Cholesky (A is overwritten)."""
    D = A.shape[0]
    ld = 0.0
    for j in range(D):
        s = A[j, j]
        for k in range(j):
            s -= A[j, k] ** 2
        A[j, j] = math.sqrt(s)
        ld += 2 * math.log(A[j, j])
        for i in range(j + 1, D):
            s = A[i, j]
            for k in range(j):
                s -= A[i, k] * A[j, k]
            A[i, j] = s / A[j, j]
    return ld


def niw(m0, k0, nu0, Psi0, nmax):
    """A Normal-inverse-Wishart prior as a tuple, with the terms of the log marginal likelihood
    that depend only on the cluster size n tabulated for n = 0..nmax."""
    D = len(m0)
    n = np.arange(nmax + 1)
    gtab = (-n * D / 2 * np.log(np.pi) + nu0 / 2 * np.linalg.slogdet(Psi0)[1] + D / 2 * (np.log(k0) - np.log(k0 + n))
            + sum(gammaln((nu0 + n) / 2 - j / 2) - gammaln(nu0 / 2 - j / 2) for j in range(D)))
    return (np.asarray(m0, float), float(k0), float(nu0), np.asarray(Psi0, float), gtab)


@nb.njit
def log_ml(n, sx, sxx, pr):
    """log marginal likelihood of n feature vectors (sum sx, sum of outer products sxx) under a
    Normal-inverse-Wishart prior pr: the cluster's mean and covariance integrated out."""
    m0, k0, nu0, Psi0, gtab = pr
    D = m0.shape[0]
    if n == 0:
        return 0.0
    kn, nun = k0 + n, nu0 + n
    A = np.empty((D, D))
    for i in range(D):
        mi = (k0 * m0[i] + sx[i]) / kn
        for j in range(i + 1):
            mj = (k0 * m0[j] + sx[j]) / kn
            A[i, j] = Psi0[i, j] + sxx[i, j] + k0 * m0[i] * m0[j] - kn * mi * mj
    return gtab[n] - nun / 2 * _logdet_chol(A)


@nb.njit
def log_ml_plus(n, sx, sxx, ax, axx, pr):
    """log_ml(n, sx + ax, sxx + axx) without building the sums (the sampler's inner loop)."""
    m0, k0, nu0, Psi0, gtab = pr
    D = m0.shape[0]
    kn, nun = k0 + n, nu0 + n
    A = np.empty((D, D))
    for i in range(D):
        mi = (k0 * m0[i] + sx[i] + ax[i]) / kn
        for j in range(i + 1):
            mj = (k0 * m0[j] + sx[j] + ax[j]) / kn
            A[i, j] = Psi0[i, j] + sxx[i, j] + axx[i, j] + k0 * m0[i] * m0[j] - kn * mi * mj
    return gtab[n] - nun / 2 * _logdet_chol(A)


@nb.njit
def _sample_log(lp, K):
    mx = -np.inf
    for k in range(K):
        mx = max(mx, lp[k])
    tot = 0.0
    for k in range(K):
        lp[k] = math.exp(lp[k] - mx)
        tot += lp[k]
    u = np.random.random() * tot
    c = 0.0
    for k in range(K):
        c += lp[k]
        if u <= c:
            return k
    return K - 1


@nb.njit
def _find(par, i):
    while par[i] != i:
        par[i] = par[par[i]]
        i = par[i]
    return i


@nb.njit
def _drop_if_empty(k_old, K, z, cnt, sx, sxx, lml):
    """If cluster k_old is empty, move the last cluster into its slot (labels stay 0..K-1)."""
    if cnt[k_old] > 0:
        return K
    last = K - 1
    if k_old != last:
        cnt[k_old], lml[k_old] = cnt[last], lml[last]
        sx[k_old], sxx[k_old] = sx[last], sxx[last]
        for j in range(z.shape[0]):
            if z[j] == last:
                z[j] = k_old
    cnt[last], lml[last] = 0, 0.0
    sx[last], sxx[last] = 0.0, 0.0
    return K - 1


@nb.njit
def _sweep(X, xx, lml1, ptr, nbr, w, z, cnt, sx, sxx, lml, K, alpha, beta, tau, lw,
           pr, lp, lmlnew, pot, par, comp_n, members, start, fill):
    """One single-site Gibbs sweep, then one Swendsen-Wang block step, for the posterior raised
    to the power tau (tempering) with the likelihood weighted by lw (lw = 0: the prior)."""
    n, D = X.shape
    for i in np.random.permutation(n):
        k_old = z[i]
        cnt[k_old] -= 1
        sx[k_old] -= X[i]
        sxx[k_old] -= xx[i]
        lml[k_old] = log_ml(cnt[k_old], sx[k_old], sxx[k_old], pr)
        K = _drop_if_empty(k_old, K, z, cnt, sx, sxx, lml)
        pot[:K] = 0.0
        for e in range(ptr[i], ptr[i + 1]):
            pot[z[nbr[e]]] += w[e]
        for k in range(K):
            lmlnew[k] = log_ml_plus(cnt[k] + 1, sx[k], sxx[k], X[i], xx[i], pr)
            lp[k] = tau * (math.log(cnt[k]) + lw * (lmlnew[k] - lml[k]) + beta * pot[k])
        lp[K] = tau * (math.log(alpha) + lw * lml1[i])
        k_new = _sample_log(lp, K + 1)
        if k_new == K:
            K += 1
            lml[k_new] = lml1[i]
        else:
            lml[k_new] = lmlnew[k_new]
        z[i] = k_new
        cnt[k_new] += 1
        sx[k_new] += X[i]
        sxx[k_new] += xx[i]
    if beta <= 0:
        return K
    # Swendsen-Wang: bond equal-label neighbours with probability 1 - exp(-tau beta w) ...
    for i in range(n):
        par[i] = i
    for i in range(n):
        for e in range(ptr[i], ptr[i + 1]):
            j = nbr[e]
            if j > i and z[j] == z[i] and np.random.random() < 1 - math.exp(-tau * beta * w[e]):
                ri, rj = _find(par, i), _find(par, j)
                if ri != rj:
                    par[ri] = rj
    comp_n[:] = 0
    for i in range(n):
        comp_n[_find(par, i)] += 1
    start[0] = 0
    for r in range(n):
        start[r + 1] = start[r] + comp_n[r]
        fill[r] = start[r]
    for i in range(n):
        r = _find(par, i)
        members[fill[r]] = i
        fill[r] += 1
    # ... then Gibbs-update the label of each bonded block: given the bonds, the Potts factor is
    # constant, so only the partition prior and the marginal likelihoods matter
    cx_, cxx_ = np.zeros(D), np.zeros((D, D))
    for r in np.random.permutation(n):
        m = comp_n[r]
        if m < 2:
            continue
        cx_[:] = 0.0
        cxx_[:] = 0.0
        for t in range(start[r], start[r + 1]):
            cx_ += X[members[t]]
            cxx_ += xx[members[t]]
        k_old = z[members[start[r]]]
        cnt[k_old] -= m
        sx[k_old] -= cx_
        sxx[k_old] -= cxx_
        lml[k_old] = log_ml(cnt[k_old], sx[k_old], sxx[k_old], pr)
        K = _drop_if_empty(k_old, K, z, cnt, sx, sxx, lml)
        for k in range(K):
            lmlnew[k] = log_ml_plus(cnt[k] + m, sx[k], sxx[k], cx_, cxx_, pr)
            lp[k] = tau * (math.lgamma(cnt[k] + m) - math.lgamma(cnt[k]) + lw * (lmlnew[k] - lml[k]))
        lml_c = log_ml(m, cx_, cxx_, pr)
        lp[K] = tau * (math.log(alpha) + math.lgamma(m) + lw * lml_c)
        k_new = _sample_log(lp, K + 1)
        if k_new == K:
            K += 1
            lml[k_new] = lml_c
        else:
            lml[k_new] = lmlnew[k_new]
        for t in range(start[r], start[r + 1]):
            z[members[t]] = k_new
        cnt[k_new] += m
        sx[k_new] += cx_
        sxx[k_new] += cxx_
    return K


@nb.njit
def _log_post(z, cnt, lml, K, n, ptr, nbr, w, alpha, beta, lw):
    """Unnormalised log posterior of a partition: CRP x Potts x marginal likelihoods."""
    lp = math.lgamma(alpha) - math.lgamma(alpha + n)
    for k in range(K):
        lp += math.log(alpha) + math.lgamma(cnt[k]) + lw * lml[k]
    s = 0.0
    for i in range(n):
        for e in range(ptr[i], ptr[i + 1]):
            if z[nbr[e]] == z[i]:
                s += w[e]
    return lp + beta * s / 2


@nb.njit
def _swap(arr, r):
    tmp = arr[r].copy()
    arr[r] = arr[r + 1]
    arr[r + 1] = tmp


@nb.njit
def potts_dp(X, ptr, nbr, w, z0, alpha, beta, pr, taus, n_iter, thin, seed,
             n_anneal=0, lw=1.0):
    """Collapsed Gibbs + Swendsen-Wang for a CRP(alpha) x Potts(beta) partition prior with
    NIW-Gaussian clusters, with parallel tempering over inverse temperatures `taus` (taus[0] = 1).
    beta ramps up from 0 over the first n_anneal iterations. Returns the tau = 1 partitions (every
    thin-th iteration), their log posterior, and the swap acceptance rate of each adjacent pair."""
    np.random.seed(seed)
    n, D = X.shape
    R, kmax = taus.shape[0], n + 2
    Z = np.empty((R, n), np.int64)
    cnt = np.zeros((R, kmax), np.int64)
    sx, sxx = np.zeros((R, kmax, D)), np.zeros((R, kmax, D, D))
    lml, Ks = np.zeros((R, kmax)), np.zeros(R, np.int64)
    xx, lml1 = np.empty((n, D, D)), np.empty(n)
    for i in range(n):
        xx[i] = np.outer(X[i], X[i])
        lml1[i] = log_ml(1, X[i], xx[i], pr)
    mp = -np.ones(z0.max() + 1, np.int64)
    K0 = 0
    for i in range(n):
        if mp[z0[i]] < 0:
            mp[z0[i]] = K0
            K0 += 1
    for r in range(R):
        for i in range(n):
            k = mp[z0[i]]
            Z[r, i] = k
            cnt[r, k] += 1
            sx[r, k] += X[i]
            sxx[r, k] += xx[i]
        for k in range(K0):
            lml[r, k] = log_ml(cnt[r, k], sx[r, k], sxx[r, k], pr)
        Ks[r] = K0
    lp, lmlnew, pot = np.empty(kmax + 1), np.empty(kmax), np.zeros(kmax)
    par, comp_n, members = np.empty(n, np.int64), np.zeros(n, np.int64), np.empty(n, np.int64)
    start, fill, E = np.empty(n + 1, np.int64), np.empty(n, np.int64), np.empty(R)
    n_save = n_iter // thin
    out, out_lp = np.empty((n_save, n), np.int16), np.empty(n_save)
    acc, tries = np.zeros(max(R - 1, 1)), np.zeros(max(R - 1, 1))
    save = 0
    for it in range(n_iter):
        b = beta * min(1.0, (it + 1) / n_anneal) if n_anneal > 0 else beta
        for r in range(R):
            Ks[r] = _sweep(X, xx, lml1, ptr, nbr, w, Z[r], cnt[r], sx[r], sxx[r], lml[r], Ks[r], alpha, b,
                           taus[r], lw, pr, lp, lmlnew, pot, par, comp_n, members, start, fill)
        if R > 1:
            for r in range(R):
                E[r] = _log_post(Z[r], cnt[r], lml[r], Ks[r], n, ptr, nbr, w, alpha, b, lw)
            for r in range(it % 2, R - 1, 2):          # propose swaps between neighbouring temperatures
                tries[r] += 1
                if math.log(np.random.random()) < (taus[r] - taus[r + 1]) * (E[r + 1] - E[r]):
                    acc[r] += 1
                    _swap(Z, r)
                    _swap(cnt, r)
                    _swap(sx, r)
                    _swap(sxx, r)
                    _swap(lml, r)
                    Ks[r], Ks[r + 1] = Ks[r + 1], Ks[r]
        if (it + 1) % thin == 0:
            out[save] = Z[0]
            out_lp[save] = _log_post(Z[0], cnt[0], lml[0], Ks[0], n, ptr, nbr, w, alpha, b, lw)
            save += 1
    return out, out_lp, acc / np.maximum(tries, 1)


def csr(n, edges, w):
    """Neighbour lists of an undirected weighted graph (compressed sparse rows)."""
    src = np.concatenate([edges[:, 0], edges[:, 1]])
    dst = np.concatenate([edges[:, 1], edges[:, 0]])
    ww = np.concatenate([w, w]).astype(float)
    order = np.argsort(src, kind="stable")
    ptr = np.concatenate([[0], np.cumsum(np.bincount(src, minlength=n))])
    return ptr, dst[order].astype(np.int64), ww[order]


def log_post(z, X, edges, w, alpha, beta, prior, lw=1.0):
    """NumPy version of the log posterior (for enumeration and checks)."""
    n, D = X.shape
    z = np.unique(z, return_inverse=True)[1]
    lp = math.lgamma(alpha) - math.lgamma(alpha + n)
    for k in range(z.max() + 1):
        Xk = X[z == k]
        lp += math.log(alpha) + math.lgamma(len(Xk)) + lw * log_ml(len(Xk), Xk.sum(0), Xk.T @ Xk, prior)
    return lp + beta * (w * (z[edges[:, 0]] == z[edges[:, 1]])).sum()

def partitions(n):
    def rec(i, z, K):
        if i == n:
            yield z.copy()
            return
        for k in range(K + 1):
            z[i] = k
            yield from rec(i + 1, z, max(K, k + 1))
    yield from rec(0, np.zeros(n, np.int64), 0)


def canon(z):
    seen = {}
    return tuple(seen.setdefault(v, len(seen)) for v in z)


rng_t = np.random.default_rng(1)
Xt = np.concatenate([rng_t.normal(0, 1, (3, 2)), rng_t.normal(2.5, 1, (4, 2))])
edges_t = np.array([[i, i + 1] for i in range(6)] + [[0, 3], [2, 5]])
w_t = rng_t.uniform(0.5, 1.5, len(edges_t))
prior_t = niw(np.zeros(2), 0.1, 4.0, np.eye(2), 7)
ptr_t, nbr_t, ww_t = csr(7, edges_t, w_t)
parts = list(partitions(7))
index = {canon(z): j for j, z in enumerate(parts)}
rows = []
for beta_t in [0.0, 1.5]:
    lp_exact = np.array([log_post(z, Xt, edges_t, w_t, 1.0, beta_t, prior_t) for z in parts])
    p_exact = np.exp(lp_exact - lp_exact.max())
    p_exact /= p_exact.sum()
    for label, taus in [("Gibbs + SW", np.array([1.0])), ("Gibbs + SW, tempered", np.array([1.0, 0.6, 0.3]))]:
        Zt, _, _ = potts_dp(Xt, ptr_t, nbr_t, ww_t, np.zeros(7, np.int64), 1.0, beta_t, prior_t, taus, 100_000, 1, 3)
        freq = np.bincount([index[canon(z)] for z in Zt], minlength=len(parts)) / len(Zt)
        iid = rng_t.multinomial(len(Zt), p_exact) / len(Zt)     # what exact independent draws would give
        rows.append((beta_t, label, 0.5 * np.abs(freq - p_exact).sum(), 0.5 * np.abs(iid - p_exact).sum()))
print(f"{len(parts)} partitions of 7 nodes")
print(pd.DataFrame(rows, columns=["beta", "sampler", "TV(sampler, exact)", "TV(100k exact draws, exact)"]).round(4)
      .to_string(index=False))

Yg, Xg = np.mgrid[:H, :W]
truth = np.zeros((H, W), int)
truth[Yg > 0.55 * H + 12 * np.sin(Xg / 25)] = 1
truth[(Yg - 60) ** 2 + (Xg - 70) ** 2 < 32**2] = 2
truth[(abs(Yg - 95) < 28) & (abs(Xg - 175) < 38)] = 3
lab_true = np.array([[66, 10, 8], [52, 26, 26], [48, -20, 25], [72, 0, -25]], float)
rng_s = np.random.default_rng(7)
blotches = np.stack([ndimage.gaussian_filter(rng_s.normal(0, 1, (H, W)), 4) for _ in range(3)], -1)
lab_syn = lab_true[truth] + 6 * blotches / blotches.std() + rng_s.normal(0, 6, (H, W, 3))
img_syn = np.clip(lab_to_srgb(lab_syn), 0, 1)
g_syn = superpixels(img_syn)
truth_sp = np.array([np.bincount(truth[g_syn["sp"] == s]).argmax() for s in range(len(g_syn["X"]))])


def niw_prior(X, s=8.0):
    D = X.shape[1]
    nu0 = D + 2.0
    return niw(X.mean(0), 0.05, nu0, (nu0 - D - 1) * s**2 * np.eye(D), len(X))


TAUS = np.array([1.0, 0.9, 0.8, 0.7])


def fit(g, beta, chains=2, n_iter=600, thin=3, burn=100, n_anneal=150, alpha=1.0, s=8.0, taus=TAUS, seed=0, z0=None):
    """Run chains of the tempered sampler; returns (chains, draws, n) partitions and log posteriors."""
    n = len(g["X"])
    ptr, nbr, w = csr(n, g["edges"], g["blen"] / g["blen"].mean())
    Zs, lps = [], []
    for c in range(chains):
        start = np.random.default_rng(seed + c).integers(0, 20, n) if z0 is None else z0
        Z, lp, acc = potts_dp(g["X"], ptr, nbr, w, start, alpha, beta, niw_prior(g["X"], s), taus, n_iter, thin,
                              seed + c, n_anneal)
        Zs.append(Z[burn:])
        lps.append(lp[burn:])
    return np.array(Zs), np.array(lps), acc


def edge_probs(Z, edges):
    Z = Z.reshape(-1, Z.shape[-1])
    return np.mean(Z[:, edges[:, 0]] != Z[:, edges[:, 1]], axis=0)


def n_clusters(Z):
    return np.array([len(np.unique(z)) for z in Z.reshape(-1, Z.shape[-1])])


@nb.njit
def vi_dist(a, b, wt):
    """Variation of information between two partitions, nodes weighted by wt (summing to 1)."""
    J = np.zeros((a.max() + 1, b.max() + 1))
    for i in range(a.shape[0]):
        J[a[i], b[i]] += wt[i]
    pa, pb = J.sum(1), J.sum(0)
    h = 0.0
    for k in range(pa.shape[0]):
        if pa[k] > 0:
            h += pa[k] * math.log(pa[k])
    for l in range(pb.shape[0]):
        if pb[l] > 0:
            h += pb[l] * math.log(pb[l])
    for k in range(J.shape[0]):
        for l in range(J.shape[1]):
            if J[k, l] > 0:
                h -= 2 * J[k, l] * math.log(J[k, l])
    return h


t0 = time.time()
syn = {}
for beta in [0.0, 1.0, 2.0]:
    Zc, lpc, _ = fit(g_syn, beta)
    Z = Zc.reshape(-1, Zc.shape[-1]).astype(np.int64)
    wt = g_syn["npix"] / g_syn["npix"].sum()
    syn[beta] = dict(Z=Z, K=n_clusters(Z), vi=np.array([vi_dist(z, truth_sp, wt) for z in Z]), pe=edge_probs(Z, g_syn["edges"]))
print(f"synthetic fits in {time.time() - t0:.0f} s")
print(pd.DataFrame({f"beta = {b:g}": {"K (median)": np.median(v["K"]), "K 90% interval": np.quantile(v["K"], [0.05, 0.95]),
                                      "VI to truth (median)": np.median(v["vi"]).round(3)} for b, v in syn.items()}))