How languages carve up colour: the World Color Survey

import io
import logging
import time
import warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
from IPython.display import HTML, Image, display
from matplotlib import animation
from matplotlib.collections import LineCollection, PolyCollection

from pymc_challenges import data

RANDOM_SEED = 77
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):
    """Figures full of coloured chips are far smaller as JPEG than PNG."""
    buf = io.BytesIO()
    fig.savefig(buf, format="jpeg", dpi=dpi, pil_kwargs={"quality": quality})
    plt.close(fig)
    display(Image(buf.getvalue()))

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


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)


ROWS = "ABCDEFGHIJ"
LAB = chips[["L", "a", "b"]].to_numpy()
RGB = lab_to_srgb(LAB)
ROW = chips.row.map(ROWS.index).to_numpy()
COL = chips.col.to_numpy()
XPOS = np.where(COL == 0, 0.0, COL + 0.7)          # achromatic column on the left, then a gap
YPOS = ROW.astype(float)
CHIP_AT = {(r, c): i for i, (r, c) in enumerate(zip(ROW, COL))}
DARK = LAB[:, 0] < 55                              # glyphs in white on dark chips


def edges():
    """Neighbouring chip pairs and the line segment(s) of their shared edge (the hue wrap gets two)."""
    E, segs = [], []
    for r in range(1, 9):
        for c in range(1, 41):
            i = CHIP_AT[(r, c)]
            j = CHIP_AT[(r, c % 40 + 1)]
            x = XPOS[i] + 0.5
            s = [[(x, r - 0.5), (x, r + 0.5)]]
            if c == 40:                             # 40 | 1: draw at both ends of the row
                s.append([(XPOS[j] - 0.5, r - 0.5), (XPOS[j] - 0.5, r + 0.5)])
            E.append((i, j))
            segs.append(s)
            if r < 8:
                E.append((i, CHIP_AT[(r + 1, c)]))
                segs.append([[(XPOS[i] - 0.5, r + 0.5), (XPOS[i] + 0.5, r + 0.5)]])
    for r in range(9):
        E.append((CHIP_AT[(r, 0)], CHIP_AT[(r + 1, 0)]))
        segs.append([[(-0.5, r + 0.5), (0.5, r + 0.5)]])
    return np.array(E), segs


EDGES, EDGE_SEGS = edges()
SQUARES = np.array([[(x - 0.47, y - 0.47), (x + 0.47, y - 0.47), (x + 0.47, y + 0.47), (x - 0.47, y + 0.47)]
                    for x, y in zip(XPOS, YPOS)])


def chip_grid(ax, face=RGB, edge_w=None, edge_color="k", labels=True, bg="0.55"):
    """The WCS palette: one square per chip in `face` colours; optional edge lines of width edge_w."""
    ax.add_collection(PolyCollection(SQUARES, facecolors=face, edgecolors="none"))
    if edge_w is not None:
        segs, lw = [], []
        for s, w in zip(EDGE_SEGS, edge_w):
            segs += s
            lw += [w] * len(s)
        keep = np.array(lw) > 0.05
        kept = [s for s, k in zip(segs, keep) if k]
        ax.add_collection(LineCollection(kept, linewidths=np.array(lw)[keep] + 1.6, colors="w", capstyle="round"))
        ax.add_collection(LineCollection(kept, linewidths=np.array(lw)[keep], colors=edge_color, capstyle="round"))
    ax.set_xlim(-0.6, 41.3)
    ax.set_ylim(9.6, -0.6)
    ax.set_aspect("equal")
    ax.set_facecolor(bg)
    if labels:
        ax.set_yticks(range(10), list(ROWS), fontsize=7)
        ax.set_xticks([0] + [c + 0.7 for c in (1, 10, 20, 30, 40)], ["0", "1", "10", "20", "30", "40"], fontsize=7)
    else:
        ax.set_xticks([])
        ax.set_yticks([])
    for s in ax.spines.values():
        s.set_visible(False)


fig = plt.figure(figsize=(16, 5.6))
gs = fig.add_gridspec(1, 2, width_ratios=[3.2, 1])
ax = fig.add_subplot(gs[0])
chip_grid(ax)
ax.set_title("The 330 WCS chips: 40 hue columns x 8 lightness rows (B-I), and 10 greys (column 0)", fontsize=11)
ax = fig.add_subplot(gs[1])
ax.scatter(LAB[:, 1], LAB[:, 2], c=RGB, s=10 + 0.3 * LAB[:, 0], edgecolor="0.3", lw=0.2)
for c, m in ((1, "s"), (40, "D")):
    k = [CHIP_AT[(r, c)] for r in range(1, 9)]
    ax.scatter(LAB[k, 1], LAB[k, 2], s=70, marker=m, facecolor="none", edgecolor="k", lw=1.1, label=f"column {c}")
ax.legend(fontsize=8, loc="lower left")
ax.set_xlabel("a*  (green - red)")
ax.set_ylabel("b*  (blue - yellow)")
ax.set_aspect("equal")
ax.set_title("The same chips in CIELAB", fontsize=11)
show_jpeg(fig)

def lang_data(L, share=1 / 3):
    """Naming matrix (speakers x chips, -1 = missing) recoded to major terms + 'other', and the foci."""
    rows = np.flatnonzero(SPK_LANG == L)
    Y = NAMING[rows].astype(int)
    tl = terms[terms.lang == L].sort_values("code")
    major = (tl.n_speakers >= share * len(rows)).to_numpy()
    remap = np.full(tl.code.max() + 1, major.sum())
    remap[tl.code.to_numpy()[major]] = np.arange(major.sum())
    Yr = np.where(Y < 0, -1, remap[np.maximum(Y, 0)])
    names = tl.abbrev[major].tolist() + (["other"] if (~major).any() else [])
    fr = np.flatnonzero(wcs["foci_lang"] == L)
    spk_row = {s: i for i, s in enumerate(SPK_NR[rows])}
    ft = remap[wcs["foci_code"][fr]]
    ok = ft < major.sum()                                   # foci of major terms only
    foci = (np.array([spk_row.get(s, -1) for s in wcs["foci_speaker"][fr]])[ok], ft[ok], wcs["foci_chip"][fr][ok] - 1)
    gloss = dict(zip(tl.abbrev, tl.term))
    return dict(L=L, name=langs.name[L], Y=Yr, terms=names, T=len(names), S=len(rows), foci=foci,
                gloss={t: gloss.get(t, "") for t in names})


def counts(Y, T):
    """chip x term counts over speakers."""
    return np.stack([(Y == t).sum(0) for t in range(T)], 1)


NAF = lang_data(77)
cnt = counts(NAF["Y"], NAF["T"])
print(f"{NAF['name']}: {NAF['S']} speakers, terms {NAF['terms']}, transcriptions: {NAF['gloss']}")
print("responses per term:", dict(zip(NAF["terms"], cnt.sum(0).tolist())))


def glyphs(ax, term_idx, alpha, names, size=7.5, scale_size=False):
    for i in range(330):
        a = float(np.clip(alpha[i], 0, 1))
        ax.text(XPOS[i], YPOS[i] + 0.04, names[term_idx[i]], ha="center", va="center",
                fontsize=size * (0.55 + 0.45 * a) if scale_size else size, color="w" if DARK[i] else "k",
                alpha=0.15 + 0.85 * a, fontweight="bold")


fig, ax = plt.subplots(figsize=(15, 4.4))
chip_grid(ax)
share = cnt / cnt.sum(1, keepdims=True)
glyphs(ax, share.argmax(1), share.max(1), NAF["terms"])
ax.set_title(f"{NAF['name']}: the majority term for each chip (letter), opacity = share of the {NAF['S']} speakers who gave it")
show_jpeg(fig)

Lc, ac, bc = (LAB[:, 0] - 50) / 50, LAB[:, 1] / 50, LAB[:, 2] / 50


def standardise(F):
    return np.c_[np.ones(len(F)), (F - F.mean(0)) / F.std(0)]


FEATURES = {
    "CIELAB quadratic": standardise(np.c_[Lc, ac, bc, Lc**2, ac**2, bc**2, ac * bc, Lc * ac, Lc * bc]),
}
F_LAB = FEATURES["CIELAB quadratic"]
FEAT_NAMES = ["1", "L", "a", "b", "L2", "a2", "b2", "ab", "La", "Lb"]


def softmax(eta):
    e = np.exp(eta - eta.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)


def exemplar_colours(P):
    """A display colour per term: the CIELAB mean of the chips, weighted by P(term | chip)^8 (its core)."""
    w = P**8 + 1e-300
    return lab_to_srgb((w.T @ LAB) / w.sum(0)[:, None])


def posterised(ax, P, names=None, edge=None, title=None, glyph=True, labels=False):
    """Each chip painted in the exemplar colour of its most probable term."""
    face = exemplar_colours(P)[P.argmax(1)]
    chip_grid(ax, face=face, edge_w=edge, labels=labels)
    if glyph and names is not None:
        for i in range(330):
            ax.text(XPOS[i], YPOS[i] + 0.04, names[P[i].argmax()], ha="center", va="center", fontsize=5.5,
                    color="w" if lab_to_srgb_L(face[i]) < 0.5 else "k")
    if title:
        ax.set_title(title, fontsize=9)


def lab_to_srgb_L(rgb):
    return 0.2126 * rgb[0] + 0.7152 * rgb[1] + 0.0722 * rgb[2]


rng_p = np.random.default_rng(RANDOM_SEED + 1)
fig, axs = plt.subplots(2, 2, figsize=(15, 5.4))
prior_stats = []
for k in range(200):
    W = rng_p.normal(0, 3, (10, 4))
    W -= W.mean(1, keepdims=True)
    P = softmax(F_LAB @ W)
    prior_stats.append((P.max(1).mean(), len(np.unique(P.argmax(1)))))
    if k < 4:
        posterised(axs.flat[k], P, names=["t1", "t2", "t3", "t4"], title=f"prior draw {k + 1}")
show_jpeg(fig)
ps = np.array(prior_stats)
print(f"prior (4 terms): mean probability of the most likely term {ps[:, 0].mean():.2f}; "
      f"terms that win at least one chip: {np.bincount(ps[:, 1].astype(int), minlength=5)[1:]} (1, 2, 3, 4 terms, of 200 draws)")

theta = 2 * np.pi * COL / 40
chrom = (COL > 0).astype(float)
FEATURES["grid, hue as a line"] = standardise(np.c_[ROW, COL, ROW**2, COL**2, ROW * COL, 1 - chrom])
FEATURES["grid, hue as a circle"] = standardise(np.c_[ROW, chrom * np.cos(theta), chrom * np.sin(theta),
                                                      chrom * np.cos(2 * theta), chrom * np.sin(2 * theta),
                                                      ROW * chrom * np.cos(theta), ROW * chrom * np.sin(theta),
                                                      ROW**2, 1 - chrom])


def naming_model(Y, names, F=F_LAB, hier=True, n_rand=4, mask=None, w_sd=3.0):
    """Softmax naming model. hier=False pools all speakers; hier=True gives every speaker their own
    intercept and linear (L, a, b) weights, drawn around the language's."""
    S, C = Y.shape
    T, K = len(names), F.shape[1]
    ok = (Y >= 0) if mask is None else (Y >= 0) & mask
    onehot = np.zeros((S, C, T))
    s_i, c_i = np.nonzero(ok)
    onehot[s_i, c_i, Y[s_i, c_i]] = 1
    feat = (FEAT_NAMES if K == 10 else [f"f{k}" for k in range(K)])
    coords = {"term": names, "speaker": np.arange(S), "feat": feat, "rfeat": feat[:n_rand]}
    with pm.Model(coords=coords) as m:
        W = pm.ZeroSumNormal("W", w_sd, dims=("feat", "term"))
        if hier:
            sd = pm.HalfNormal("sd", 1.0, dims=("rfeat", "term"))
            z = pm.ZeroSumNormal("z", 1.0, dims=("speaker", "rfeat", "term"))
            Ws = W[None] + pt.concatenate([sd[None] * z, pt.zeros((S, K - n_rand, T))], axis=1)
            eta = pt.einsum("ck,skt->sct", pt.as_tensor(F), Ws)
            pm.Potential("loglik", ((eta - pt.logsumexp(eta, axis=-1, keepdims=True)) * onehot).sum())
        else:
            eta = pt.dot(F, W)
            pm.Potential("loglik", ((eta - pt.logsumexp(eta, axis=-1, keepdims=True)) * onehot.sum(0)).sum())
    return m


def fit(model, label, draws=1000):
    t0 = time.time()
    idata = pm.sample(draws=draws, model=model, random_seed=RANDOM_SEED, progressbar=False)
    st = idata.sample_stats
    msg = (f"{label:34s} {time.time() - t0:4.0f} s | divergences {int(st['diverging'].sum())}, mean tree depth "
           f"{float(st['depth'].mean()):.1f}, max r_hat(W) {float(az.rhat(idata.posterior['W']).max()):.3f}, "
           f"min ESS(W) {float(az.ess(idata.posterior['W']).min()):.0f}")
    if "sd" in idata.posterior:
        msg += (f", max r_hat(sd) {float(az.rhat(idata.posterior['sd']).max()):.3f}, "
                f"min ESS(sd) {float(az.ess(idata.posterior['sd']).min()):.0f}")
    print(msg)
    return idata


def draws_of(idata, name, n=400):
    """(n, ...) evenly thinned posterior draws of one variable, sample axis first."""
    x = idata.posterior[name].stack(s=("chain", "draw")).transpose("s", ...).values
    return x[np.linspace(0, len(x) - 1, n).astype(int)]


rng_split = np.random.default_rng(RANDOM_SEED + 2)
perm_spk = rng_split.permutation(NAF["S"])
TRAIN, TEST = np.sort(perm_spk[:20]), np.sort(perm_spk[20:])
Ytr, Yte = NAF["Y"][TRAIN], NAF["Y"][TEST]


def heldout_score(P_draws, Y):
    """Mean log posterior-predictive probability per response (P_draws: draws x chips x terms),
    overall and per held-out speaker."""
    Pm = P_draws.mean(0)
    s, c = np.nonzero(Y >= 0)
    lp = np.log(Pm[c, Y[s, c]])
    return lp.mean(), np.array([lp[s == k].mean() for k in range(len(Y))]), lp


pooled_P, grid_rows = {}, []
for fname, F in FEATURES.items():
    idata = fit(naming_model(Ytr, NAF["terms"], F=F, hier=False), f"pooled, {fname}")
    Wd = draws_of(idata, "W")
    pooled_P[fname] = softmax(np.einsum("ck,dkt->dct", F, Wd))
    grid_rows.append((fname, heldout_score(pooled_P[fname], Yte)))
base, base_all = grid_rows[0][1][1], grid_rows[0][1][2]
print("\nlog score per response on the 9 held-out speakers (higher is better); difference to CIELAB, "
      "mean ± se over the 9 speakers:")
for fname, (sc, per_spk, lp) in grid_rows:
    dlt = per_spk - base
    print(f"  {fname:24s} {sc:.3f}   diff {dlt.mean():+.3f} ± {dlt.std(ddof=1) / np.sqrt(len(dlt)):.3f}  "
          f"(summed over {len(lp)} responses: {(lp - base_all).sum():+.0f} nats)")

def entropy(p, axis=-1):
    return -(p * np.log(np.clip(p, 1e-300, 1))).sum(axis)


def new_speaker_probs(idata, F=F_LAB, n_draws=400, n_new=20, seed=0, n_rand=4):
    """For each of n_draws posterior draws: choice probabilities (n_new x chips x terms) of new speakers."""
    Wd, sdd = draws_of(idata, "W", n_draws), draws_of(idata, "sd", n_draws)
    g = np.random.default_rng(seed)
    T = Wd.shape[-1]
    for d in range(n_draws):
        z = g.standard_normal((n_new, n_rand, T))
        z -= z.mean(-1, keepdims=True)                       # zero-sum over terms, as in the model
        Ws = np.repeat(Wd[d][None], n_new, 0)
        Ws[:, :n_rand] += sdd[d][None] * z
        yield softmax(np.einsum("ck,nkt->nct", F, Ws))


def sample_names(P, g):
    """One categorical draw per row of P (..., terms)."""
    u = g.random(P.shape[:-1])[..., None]
    return np.minimum((P.cumsum(-1) < u).sum(-1), P.shape[-1] - 1)


def per_speaker_counts(Y, T):
    return np.stack([(Y == t).sum(1) for t in range(T)], 1)


T_N, S_N = NAF["T"], NAF["S"]
MASK = NAF["Y"] >= 0
obs_cnt = per_speaker_counts(NAF["Y"], T_N)
g = np.random.default_rng(RANDOM_SEED + 3)
P_pool = softmax(np.einsum("ck,dkt->dct", F_LAB, draws_of(idata_pool, "W", 200)))
rep = {"pooled": [], "hierarchical": []}
for d, p_new in enumerate(new_speaker_probs(idata_hier, n_draws=200, n_new=S_N, seed=RANDOM_SEED + 4)):
    for name, P in (("pooled", np.broadcast_to(P_pool[d], (S_N, 330, T_N))), ("hierarchical", p_new)):
        y = np.where(MASK, sample_names(P, g), -1)
        rep[name].append(per_speaker_counts(y, T_N))
rep = {k: np.array(v) for k, v in rep.items()}                          # draws x speakers x terms
rows_ = []
for t, tn in enumerate(NAF["terms"]):
    o = obs_cnt[:, t].std()
    row = {"term": tn, "observed sd": o}
    for k, r in rep.items():
        sds = r[:, :, t].std(1)
        row[f"{k}: 90% interval"] = f"{np.quantile(sds, 0.05):.1f} - {np.quantile(sds, 0.95):.1f}"
        row[f"{k}: P(rep >= obs)"] = (sds >= o).mean()
    rows_.append(row)
print(f"B: {(obs_cnt[:, -1] == 0).sum()} of {S_N} speakers never use it; the others use it for "
      f"{obs_cnt[obs_cnt[:, -1] > 0, -1].min()}-{obs_cnt[:, -1].max()} chips")
print("standard deviation across speakers of the number of chips given each term:")
print(pd.DataFrame(rows_).set_index("term").round(2))

fig, axs = plt.subplots(1, T_N, figsize=(15, 3.6))
for t, ax in enumerate(axs):
    hi = max(obs_cnt[:, t].max(), np.quantile(rep["hierarchical"][:, :, t], 0.995)) + 5
    bins = np.linspace(0, hi, 26)
    for k, col in (("pooled", "C0"), ("hierarchical", "C1")):
        ax.hist(rep[k][:, :, t].ravel(), bins=bins, weights=np.full(rep[k][:, :, t].size, 1 / rep[k][:, :, t].size),
                histtype="step", lw=2, color=col, label=f"{k} (replicated speakers)")
    ax.hist(obs_cnt[:, t], bins=bins, weights=np.full(S_N, 1 / S_N), color="0.3", alpha=0.45, label="observed speakers")
    ax.set_title(f"{NAF['terms'][t]} ({NAF['gloss'][NAF['terms'][t]]})", fontsize=10)
axs[0].set_ylabel("share of speakers")
fig.supxlabel("number of chips (of 330) one speaker gives the term", fontsize=10)
axs[0].legend(fontsize=7)
fig.suptitle("Nafaanra: how many chips each speaker gives each term - data vs the two models");

lo, hi = np.quantile(PN, [0.05, 0.95], axis=0)
f_spk, f_term, f_chip = NAF["foci"]
fig, axs = plt.subplots(T_N, 1, figsize=(15, 3.35 * T_N))
safest = P_mean.argmax(0)
for t, ax in enumerate(axs):
    chip_grid(ax, face=np.full((330, 3), 0.62), labels=t == T_N - 1)
    r_hi, r_lo = 0.49 * np.sqrt(hi[:, t]), 0.49 * np.sqrt(lo[:, t])
    for i in range(330):
        if r_hi[i] > 0.04:
            ax.add_patch(plt.Circle((XPOS[i], YPOS[i]), r_hi[i], fill=False, lw=0.7, color="0.15"))
        if r_lo[i] > 0.02:
            ax.add_patch(plt.Circle((XPOS[i], YPOS[i]), r_lo[i], color=RGB[i], lw=0))
    fc = np.bincount(f_chip[f_term == t], minlength=330)
    k = fc > 0
    ax.scatter(XPOS[k], YPOS[k], s=10 + 9 * fc[k], color="k", edgecolor="w", lw=0.8, zorder=5)
    ax.scatter(XPOS[safest[t]], YPOS[safest[t]], marker="*", s=260, color="w", edgecolor="k", lw=1.2, zorder=6)
    ax.set_title(f"{NAF['terms'][t]} ({NAF['gloss'][NAF['terms'][t]]}): share of speakers using it (disc: 5% bound, "
                 f"ring: 95% bound), foci (black dots), safest exemplar (star)", fontsize=10)
show_jpeg(fig)

rows_ = []
for t, tn in enumerate(NAF["terms"]):
    fc = f_chip[f_term == t]
    c0 = safest[t]
    top = np.bincount(fc, minlength=330).argmax() if len(fc) else -1
    fcen = LAB[fc].mean(0) if len(fc) else np.full(3, np.nan)
    rows_.append({"term": tn, "safest exemplar": chips.grid[c0],
                  "P(new speaker uses it)": f"{P_mean[c0, t]:.2f} [{lo[c0, t]:.2f}, {hi[c0, t]:.2f}]",
                  "foci (speakers)": f"{len(fc)} ({len(np.unique(f_spk[f_term == t]))})",
                  "most chosen focus": chips.grid[top] if top >= 0 else "-",
                  "speakers choosing it": len(np.unique(f_spk[(f_term == t) & (f_chip == top)])) if top >= 0 else 0,
                  "P at that focus": f"{P_mean[top, t]:.2f}" if top >= 0 else "-",
                  "CIELAB distance exemplar - foci mean": f"{np.linalg.norm(LAB[c0] - fcen):.0f}" if len(fc) else "-"})
print(pd.DataFrame(rows_).set_index("term").to_string())

TERM_RGB = exemplar_colours(P_mean)
g = np.random.default_rng(RANDOM_SEED + 6)
frames = []
for d, p in enumerate(new_speaker_probs(idata_hier, n_draws=36, n_new=1, seed=RANDOM_SEED + 7)):
    y = sample_names(p[0], g)
    frames.append(y)
frames = np.array(frames)
print("chips per term across the 36 simulated speakers (min - max):",
      {tn: f"{(frames == t).sum(1).min()} - {(frames == t).sum(1).max()}" for t, tn in enumerate(NAF["terms"])})

fig, ax = plt.subplots(figsize=(9, 3.0), dpi=72, layout="none")
fig.subplots_adjust(left=0.02, right=0.98, top=0.86, bottom=0.03)
polys = PolyCollection(SQUARES, facecolors=TERM_RGB[frames[0]], edgecolors="none")
ax.add_collection(polys)
lines = LineCollection([], colors="k", linewidths=1.6)
ax.add_collection(lines)
ax.set_xlim(-0.6, 41.3)
ax.set_ylim(9.6, -0.6)
ax.set_aspect("equal")
ax.axis("off")
ttl = ax.set_title("", fontsize=10)


def draw_frame(f):
    y = frames[f]
    polys.set_facecolors(TERM_RGB[y])
    diff = np.flatnonzero(y[i_e] != y[j_e])
    lines.set_segments([s for e in diff for s in EDGE_SEGS[e]])
    ttl.set_text(f"new Nafaanra speaker {f + 1}: " + ", ".join(f"{tn} {(y == t).sum()}" for t, tn in enumerate(NAF['terms'])))
    return polys, lines, ttl


anim = animation.FuncAnimation(fig, draw_frame, frames=len(frames), interval=600)
plt.close(fig)
plt.rcParams["animation.frame_format"] = "jpeg"
display(HTML(anim.to_jshtml(default_mode="loop")))
plt.rcParams["animation.frame_format"] = "png"

NB = [[] for _ in range(330)]
for (a_, b_) in EDGES:
    NB[a_].append(b_)
    NB[b_].append(a_)


def isolated(y):
    """Chips named, with a non-missing name, differently from every named neighbour."""
    return sum(1 for i in range(330) if y[i] >= 0 and all(y[j] != y[i] for j in NB[i] if y[j] >= 0)
               and any(y[j] >= 0 for j in NB[i]))


g = np.random.default_rng(RANDOM_SEED + 12)
sim = []
for p in new_speaker_probs(idata_hier, n_draws=200, n_new=1, seed=RANDOM_SEED + 13):
    sim.append(np.where(MASK[g.integers(S_N)], sample_names(p[0], g), -1))
iso_sim = np.array([isolated(y) for y in sim])
iso_obs = np.array([isolated(y) for y in NAF["Y"]])
print(f"isolated chips per speaker: real median {np.median(iso_obs):.0f} (range {iso_obs.min()}-{iso_obs.max()}), "
      f"simulated median {np.median(iso_sim):.0f} (90% range {np.quantile(iso_sim, 0.05):.0f}-{np.quantile(iso_sim, 0.95):.0f}); "
      f"share of simulated speakers with at least the real median: {(iso_sim >= np.median(iso_obs)).mean():.2f}")

real_pos = int(g.integers(9))
real_spk = int(g.integers(S_N))
fig, axs = plt.subplots(3, 3, figsize=(15, 4.9), layout="none")
fig.subplots_adjust(left=0.01, right=0.99, top=0.93, bottom=0.01, wspace=0.03, hspace=0.12)
k_sim = 0
for pos, ax in enumerate(axs.flat):
    y = NAF["Y"][real_spk] if pos == real_pos else sim[k_sim]
    k_sim += pos != real_pos
    face = np.where((y >= 0)[:, None], TERM_RGB[np.maximum(y, 0)], 0.55)
    diff = [e for e in range(len(EDGES)) if y[i_e[e]] != y[j_e[e]] and min(y[i_e[e]], y[j_e[e]]) >= 0]
    chip_grid(ax, face=face, edge_w=np.isin(np.arange(len(EDGES)), diff) * 0.9, labels=False)
    ax.set_title(f"speaker {pos + 1}", fontsize=9)
fig.suptitle("Eight simulated Nafaanra speakers and one real one (grey chips: no response)", fontsize=11)
show_jpeg(fig)

def shifted(k):
    """Index of the chip k hue columns to the left of each chip (greys stay put)."""
    return np.array([i if COL[i] == 0 else CHIP_AT[(ROW[i], (COL[i] - 1 - k) % 40 + 1)] for i in range(330)])


SHIFT = np.array([shifted(k) for k in range(40)])


def nmi(a, b):
    J = np.zeros((a.max() + 1, b.max() + 1))
    np.add.at(J, (a, b), 1)
    J /= J.sum()
    pa, pb = J.sum(1), J.sum(0)
    nz = J > 0
    mi = (J[nz] * np.log(J[nz] / np.outer(pa, pb)[nz])).sum()
    return mi / np.sqrt(entropy(pa) * entropy(pb))


g = np.random.default_rng(RANDOM_SEED + 8)
names_ = [LD[L]["name"] for L in E_LANGS]
M = np.zeros((len(E_LANGS), len(E_LANGS)))
text = np.empty_like(M, dtype=object)
excess = []
for a_, La in enumerate(E_LANGS):
    for b_, Lb in enumerate(E_LANGS):
        if b_ <= a_:
            continue
        ma, mb = SUM[La]["P"].argmax(2), SUM[Lb]["P"].argmax(2)
        n = min(len(ma), len(mb))
        vals = np.array([nmi(ma[d], mb[d]) for d in range(n)])
        rot = np.array([nmi(ma[d][SHIFT[g.integers(1, 40)]], mb[d]) for d in range(n)])
        M[a_, b_] = M[b_, a_] = vals.mean()
        text[a_, b_] = f"{vals.mean():.2f}\n[{np.quantile(vals, 0.05):.2f}, {np.quantile(vals, 0.95):.2f}]"
        text[b_, a_] = f"rotated:\n{rot.mean():.2f}"
        excess.append((names_[a_], names_[b_], vals.mean(), rot.mean(), (vals > np.median(rot)).mean()))
print(pd.DataFrame(excess, columns=["language 1", "language 2", "NMI", "NMI, hue-rotated", "P(NMI > rotated median)"])
      .round(2).to_string(index=False))

fig, ax = plt.subplots(figsize=(15, 4.6))
chip_grid(ax, edge_w=7 * np.mean([SUM[L]["edge"] for L in E_LANGS], 0))
ax.set_title("Where the five languages put boundaries: line width = P(a new speaker names the two chips differently), "
             "averaged over languages", fontsize=10)
show_jpeg(fig)
AVG_EDGE = np.mean([SUM[L]["edge"] for L in E_LANGS], 0)
vert = (COL[i_e] == COL[j_e]) & (COL[i_e] > 0)                     # neighbours in lightness (same hue column)
hue = (ROW[i_e] == ROW[j_e])                                       # neighbours in hue (same row)
lb = pd.Series(AVG_EDGE[vert]).groupby([f"{ROWS[r]}|{ROWS[r + 1]}" for r in ROW[i_e][vert]]).mean()
hb = pd.Series(AVG_EDGE[hue]).groupby([f"{c}|{c % 40 + 1}" for c in COL[i_e][hue]]).mean()
print("lightness boundaries (mean over the 40 hue columns):", lb.round(2).to_dict())
print("strongest hue boundaries (mean over rows B-I):", hb.sort_values(ascending=False).head(6).round(2).to_dict(),
      f"| median hue edge {hb.median():.2f}")

fig, ax = plt.subplots(figsize=(7.5, 6))
ax.imshow(np.where(np.eye(len(E_LANGS)) > 0, np.nan, M), cmap="Greys", vmin=0, vmax=1)
for a_ in range(len(E_LANGS)):
    for b_ in range(len(E_LANGS)):
        if a_ != b_:
            ax.text(b_, a_, text[a_, b_], ha="center", va="center", fontsize=7,
                    color="w" if (b_ > a_ and M[a_, b_] > 0.55) else "k")
ax.set_xticks(range(len(E_LANGS)), names_, rotation=40, fontsize=8)
ax.set_yticks(range(len(E_LANGS)), names_, fontsize=8)
ax.set_title("NMI between majority maps\nabove: posterior mean [90%]; below: after a random hue rotation", fontsize=10)
ax.grid(False);