After LDA: newer topic models, problem by problem

data.describe("upworthy_exploratory")
uw = data.load("upworthy_exploratory")
STOP = set("""a about above after again against all am an and any are aren't as at be because been before
being below between both but by can can't cannot could couldn't did didn't do does doesn't doing don't down
during each few for from further had hadn't has hasn't have haven't having he he'd he'll he's her here here's
hers herself him himself his how how's i i'd i'll i'm i've if in into is isn't it it's its itself let's me
more most mustn't my myself no nor not of off on once only or other ought our ours ourselves out over own same
shan't she she'd she'll she's should shouldn't so some such than that that's the their theirs them themselves
then there there's these they they'd they'll they're they've this those through to too under until up very
was wasn't we we'd we'll we're we've were weren't what what's when when's where where's which while who who's
whom why why's with won't would wouldn't you you'd you'll you're you've your yours yourself yourselves just
like get gets got one two really thing things way ever even also make makes made know see go going goes say
says said will new something every much many want wants yes""".split())


def tokens(headline):
    words = re.findall(r"[a-z][a-z']+", headline.lower().replace("’", "'"))
    return [w for w in words if w not in STOP and len(w) > 2]


heads = uw.drop_duplicates(["test_id", "headline"])
docs = heads.groupby("test_id")["headline"].apply(lambda h: sum((tokens(x) for x in h), []))
doc_freq = Counter(w for d in docs for w in set(d))
fold = {w: w[:-1] for w in doc_freq if w.endswith("s") and doc_freq.get(w[:-1], 0) >= 20}
docs = docs.map(lambda d: [fold.get(w, w) for w in d])
doc_freq = Counter(w for d in docs for w in set(d))
vocab = np.array(sorted(w for w, c in doc_freq.items() if 20 <= c <= 0.1 * len(docs)))
w_index = {w: i for i, w in enumerate(vocab)}
X_all = np.zeros((len(docs), len(vocab)), dtype=np.int64)
for d, ws in enumerate(docs):
    for w in ws:
        if w in w_index:
            X_all[d, w_index[w]] += 1
keep = X_all.sum(1) >= 3
X_txt, test_ids = X_all[keep], docs.index.to_numpy()[keep]
@numba.njit
def _gibbs_sweeps(doc, word, z, ndk, nkw, nk, alpha, beta, n_sweeps, seed):
    np.random.seed(seed)
    K, V = nkw.shape
    cum = np.empty(K)
    for _ in range(n_sweeps):
        for i in range(doc.shape[0]):
            d, w, k = doc[i], word[i], z[i]
            ndk[d, k] -= 1
            nkw[k, w] -= 1
            nk[k] -= 1
            total = 0.0
            for j in range(K):
                total += (ndk[d, j] + alpha) * (nkw[j, w] + beta) / (nk[j] + V * beta)
                cum[j] = total
            u = np.random.random() * total
            k = 0
            while cum[k] < u:
                k += 1
            z[i] = k
            ndk[d, k] += 1
            nkw[k, w] += 1
            nk[k] += 1


def lda_gibbs(X, K, alpha=0.1, beta=0.05, burn=2000, draws=40, thin=10, seed=0):
    """Collapsed Gibbs LDA on a document x word count matrix.

    Returns posterior draws of phi (draw, topic, word) and theta (draw, doc, topic), and the
    log-likelihood of the tokens at each saved sweep.
    """
    g = np.random.default_rng(seed)
    d_i, w_i = np.nonzero(X)
    n = X[d_i, w_i]
    doc, word = np.repeat(d_i, n), np.repeat(w_i, n)
    z = g.integers(K, size=doc.size)
    D, V = X.shape
    ndk, nkw = np.zeros((D, K), np.int64), np.zeros((K, V), np.int64)
    np.add.at(ndk, (doc, z), 1)
    np.add.at(nkw, (z, word), 1)
    nk = nkw.sum(1)
    _gibbs_sweeps(doc, word, z, ndk, nkw, nk, alpha, beta, burn, seed)
    phis, thetas, lls = [], [], []
    for s in range(draws):
        _gibbs_sweeps(doc, word, z, ndk, nkw, nk, alpha, beta, thin, seed + 1 + s)
        phi_s = g.gamma(nkw + beta)
        theta_s = g.gamma(ndk + alpha)
        phis.append(phi_s / phi_s.sum(1, keepdims=True))
        thetas.append(theta_s / theta_s.sum(1, keepdims=True))
        lls.append(np.sum(X * np.log(thetas[-1] @ phis[-1])))
    return np.array(phis), np.array(thetas), np.array(lls)


@numba.njit
def _foldin_sweeps(doc, word, z, ndk, phi, alpha, n_sweeps, seed):
    np.random.seed(seed)
    K = phi.shape[0]
    cum = np.empty(K)
    for _ in range(n_sweeps):
        for i in range(doc.shape[0]):
            d, w, k = doc[i], word[i], z[i]
            ndk[d, k] -= 1
            total = 0.0
            for j in range(K):
                total += (ndk[d, j] + alpha) * phi[j, w]
                cum[j] = total
            u = np.random.random() * total
            k = 0
            while cum[k] < u:
                k += 1
            z[i] = k
            ndk[d, k] += 1


tests = uw.groupby("test_id").agg(clicks=("clicks", "sum"), impressions=("impressions", "sum"),
                                  week=("test_week", "first")).loc[test_ids]
logit_ctr = np.log((tests["clicks"] + 0.5) / (tests["impressions"] - tests["clicks"] + 0.5)).to_numpy()
year, wk = tests["week"].to_numpy() // 100, tests["week"].to_numpy() % 100
month = (year - year.min()) * 12 + np.minimum((wk - 1) * 12 // 52, 11)
month_idx, month_codes = pd.factorize(month, sort=True)
y_rel = logit_ctr - pd.Series(logit_ctr).groupby(month_idx).transform("mean").to_numpy()
cv_fold = np.random.default_rng(2).integers(5, size=len(y_rel))

def cv_r2(F, lam=1.0):
    pred = np.zeros_like(y_rel)
    for f in range(5):
        tr = cv_fold != f
        A = np.c_[np.ones(tr.sum()), F[tr]]
        b = np.linalg.solve(A.T @ A + lam * np.eye(A.shape[1]), A.T @ y_rel[tr])
        pred[~tr] = np.c_[np.ones((~tr).sum()), F[~tr]] @ b
    return 1 - np.mean((y_rel - pred) ** 2) / np.var(y_rel)

def cosine(a, b):
    return np.sum(a * b, -1) / np.linalg.norm(a, axis=-1) / np.linalg.norm(b, axis=-1)

def stability(phis):
    """phis: list of (K, V) topic matrices from restarts. Per topic of run 0: min cosine to its match."""
    ref = phis[0]; sims = []
    for other in phis[1:]:
        perm = linear_sum_assignment(-cosine(ref[:, None], other[None]))[1]
        sims.append(cosine(ref, other[perm]))
    return np.min(sims, 0)

Xb = (X_txt > 0).astype(float)
n_docs = len(Xb)
doc_freq_w = Xb.sum(0)
co_doc = Xb.T @ Xb
def npmi(phi, top=10):
    """Mean NPMI of top-word pairs (document co-occurrence in the corpus), per topic."""
    out = []
    for k in range(len(phi)):
        t = np.argsort(-phi[k])[:top]; vals = []
        for i in range(top):
            for j in range(i + 1, top):
                pij = co_doc[t[i], t[j]] / n_docs
                pi, pj = doc_freq_w[t[i]] / n_docs, doc_freq_w[t[j]] / n_docs
                vals.append(-1.0 if pij == 0 else np.log(pij / (pi * pj)) / -np.log(pij))
        out.append(np.mean(vals))
    return np.array(out)

def diversity(phi, top=25):
    tops = np.concatenate([np.argsort(-p)[:top] for p in phi])
    return len(np.unique(tops)) / len(tops)

def top_words(phi, n=6):
    return [" / ".join(vocab[np.argsort(-p)[:n]]) for p in phi]

def top_sets(phi, n=10):
    return [set(np.argsort(-p)[:n]) for p in phi]

def stability_jaccard(phis, n=10):
    """Per topic of run 0: the smallest (over other runs) Jaccard overlap of top-n words with its best match."""
    ref = top_sets(phis[0], n); out = []
    for other in phis[1:]:
        o = top_sets(other, n)
        J = np.array([[len(a & b) / len(a | b) for b in o] for a in ref])
        perm = linear_sum_assignment(-J)[1]
        out.append(J[np.arange(len(ref)), perm])
    return np.min(out, 0)
print(f"{len(X_txt):,} tests x {len(vocab)} words, {X_txt.sum():,} tokens (median {np.median(X_txt.sum(1)):.0f} per test)")

@numba.njit
def _dmm_sweeps(doc_ptr, word, z, m, nkw, nk, alpha, beta, n_sweeps, seed):
    np.random.seed(seed)
    K, V = nkw.shape
    D = len(z)
    logp = np.empty(K)
    for _ in range(n_sweeps):
        for d in range(D):
            a, b = doc_ptr[d], doc_ptr[d + 1]
            k = z[d]
            m[k] -= 1
            for i in range(a, b):
                nkw[k, word[i]] -= 1
            nk[k] -= b - a
            for j in range(K):
                lp = np.log(m[j] + alpha)
                # tokens of the document are sorted by word, so repeated words are adjacent
                rep = 0
                for i in range(a, b):
                    if i > a and word[i] == word[i - 1]:
                        rep += 1
                    else:
                        rep = 0
                    lp += np.log(nkw[j, word[i]] + beta + rep)
                for i in range(b - a):
                    lp -= np.log(nk[j] + V * beta + i)
                logp[j] = lp
            mx = logp.max()
            tot = 0.0
            for j in range(K):
                logp[j] = np.exp(logp[j] - mx)
                tot += logp[j]
            u = np.random.random() * tot
            k = 0
            acc = logp[0]
            while acc < u:
                k += 1
                acc += logp[k]
            z[d] = k
            m[k] += 1
            for i in range(a, b):
                nkw[k, word[i]] += 1
            nk[k] += b - a

def dmm_gibbs(X, K, alpha=0.1, beta=0.05, burn=300, draws=20, thin=10, seed=0):
    g = np.random.default_rng(seed)
    D, V = X.shape
    d_i, w_i = np.nonzero(X); n = X[d_i, w_i]
    word = np.repeat(w_i, n); doc_of = np.repeat(d_i, n)
    doc_ptr = np.r_[0, np.cumsum(X.sum(1))]
    z = g.integers(K, size=D)
    m = np.bincount(z, minlength=K).astype(np.int64)
    nkw = np.zeros((K, V), np.int64); np.add.at(nkw, (z[doc_of], word), 1); nk = nkw.sum(1)
    _dmm_sweeps(doc_ptr, word, z, m, nkw, nk, alpha, beta, burn, seed)
    phis, resp = [], np.zeros((D, K))
    for s in range(draws):
        _dmm_sweeps(doc_ptr, word, z, m, nkw, nk, alpha, beta, thin, seed + 1 + s)
        ph = g.gamma(nkw + beta); phis.append(ph / ph.sum(1, keepdims=True))
        # soft cluster membership given the topics: p(z_d = k | x_d)
        lr = np.log(m + alpha)[None] + X @ np.log(phis[-1]).T
        lr = np.exp(lr - lr.max(1, keepdims=True)); resp += lr / lr.sum(1, keepdims=True)
    return np.array(phis), resp / draws, m

@numba.njit
def _btm_sweeps(w1, w2, z, nk, nkw, alpha, beta, n_sweeps, seed):
    np.random.seed(seed)
    K, V = nkw.shape
    cum = np.empty(K)
    for _ in range(n_sweeps):
        for b in range(len(z)):
            k = z[b]; i, j = w1[b], w2[b]
            nk[k] -= 1; nkw[k, i] -= 1; nkw[k, j] -= 1
            tot = 0.0
            for t in range(K):
                nw = 2.0 * nk[t]
                tot += (nk[t] + alpha) * (nkw[t, i] + beta) * (nkw[t, j] + beta) / ((nw + V * beta) * (nw + 1 + V * beta))
                cum[t] = tot
            u = np.random.random() * tot
            k = 0
            while cum[k] < u:
                k += 1
            z[b] = k; nk[k] += 1; nkw[k, i] += 1; nkw[k, j] += 1

def biterms(X):
    d_list, a_list, b_list = [], [], []
    for d in range(len(X)):
        toks = np.repeat(np.arange(X.shape[1]), X[d])
        ii, jj = np.triu_indices(len(toks), 1)
        d_list.append(np.full(len(ii), d)); a_list.append(toks[ii]); b_list.append(toks[jj])
    return np.concatenate(d_list), np.concatenate(a_list), np.concatenate(b_list)

def btm_gibbs(X, K, alpha=None, beta=0.01, burn=300, draws=20, thin=10, seed=0):
    g = np.random.default_rng(seed)
    bd, w1, w2 = biterms(X)
    alpha = 50 / K if alpha is None else alpha
    V = X.shape[1]
    z = g.integers(K, size=len(w1)); nk = np.bincount(z, minlength=K).astype(np.int64)
    nkw = np.zeros((K, V), np.int64); np.add.at(nkw, (z, w1), 1); np.add.at(nkw, (z, w2), 1)
    _btm_sweeps(w1, w2, z, nk, nkw, alpha, beta, burn, seed)
    phis, thetas = [], np.zeros((len(X), K))
    for s in range(draws):
        _btm_sweeps(w1, w2, z, nk, nkw, alpha, beta, thin, seed + 1 + s)
        ph = (nkw + beta) / (2 * nk[:, None] + V * beta); phis.append(ph)
        pz = (nk + alpha) / (nk.sum() + K * alpha)
        pb = pz[None] * ph[:, w1].T * ph[:, w2].T; pb /= pb.sum(1, keepdims=True)   # p(z | biterm)
        th = np.zeros((len(X), K)); np.add.at(th, bd, pb)
        nb = np.bincount(bd, minlength=len(X))[:, None]
        thetas += np.where(nb > 0, th / np.maximum(nb, 1), pz[None])
    return np.array(phis), thetas / draws, len(w1)

t0 = time.perf_counter()
dmm_runs = [dmm_gibbs(X_txt, 6, seed=s) for s in range(4)]
secs = (time.perf_counter() - t0) / 4
dmm_phis = [r[0].mean(0) for r in dmm_runs]
print(score("DMM (GSDMM), 6 clusters", dmm_phis, [r[1] for r in dmm_runs], secs).round(3))
print("documents per cluster, first run:", dmm_runs[0][2])
show_topics(dmm_phis[0], title="DMM clusters:")

t0 = time.perf_counter()
btm_runs = [btm_gibbs(X_txt, 6, seed=s) for s in range(4)]
secs = (time.perf_counter() - t0) / 4
btm_phis = [r[0].mean(0) for r in btm_runs]
print(f"\n{btm_runs[0][2]:,} biterms")
print(score("Biterm topic model, 6 topics", btm_phis, [r[1] for r in btm_runs], secs).round(3))
show_topics(btm_phis[0], title="BTM topics:")

def cooccurrence(X):
    """Arora et al. (2013) unbiased word co-occurrence estimate from documents with >= 2 tokens."""
    X = X[X.sum(1) >= 2].astype(float)
    n = X.sum(1)
    w = 1 / (n * (n - 1))
    Q = (X * w[:, None]).T @ X - np.diag((X * w[:, None]).sum(0))
    return Q / Q.sum()

def find_anchors(Qbar, K, candidates):
    """Greedy farthest-point anchors in the row space of Qbar (Gram-Schmidt), among candidate words."""
    rows = Qbar[candidates].copy()
    anchors = [int(np.argmax(np.linalg.norm(rows, axis=1)))]
    basis = []
    for _ in range(K - 1):
        v = rows[anchors[-1]] - (rows[anchors[0]] if len(anchors) > 1 else 0)
        if len(anchors) == 1:
            # second anchor: farthest from the first
            d = np.linalg.norm(rows - rows[anchors[0]], axis=1)
            anchors.append(int(np.argmax(d)))
            b = rows[anchors[1]] - rows[anchors[0]]
            basis.append(b / np.linalg.norm(b))
            continue
        resid = rows - rows[anchors[0]]
        for b in basis:
            resid = resid - np.outer(resid @ b, b)
        anchors.append(int(np.argmax(np.linalg.norm(resid, axis=1))))
        b = resid[anchors[-1]] / np.linalg.norm(resid[anchors[-1]])
        basis.append(b)
    return np.array(candidates)[anchors]

def recover(Q, anchors, lam=1e3):
    """RecoverL2: each word's conditional row as a convex combination of the anchor rows."""
    p_w = Q.sum(1)
    Qbar = Q / p_w[:, None]
    A_rows = Qbar[anchors]                          # K x V
    M = np.vstack([A_rows.T, np.sqrt(lam) * np.ones(len(anchors))])
    C = np.array([nnls(M, np.r_[Qbar[i], np.sqrt(lam)])[0] for i in range(len(Q))])
    C = C / C.sum(1, keepdims=True)
    A = C * p_w[:, None]                            # p(word, topic) up to normalisation
    return (A / A.sum(0)).T                         # K x V

t0 = time.perf_counter()
Q = cooccurrence(X_txt)
Qbar = Q / Q.sum(1, keepdims=True)
anchor_words = find_anchors(Qbar, 6, np.flatnonzero(doc_freq_w >= 50))
phi_anchor = recover(Q, anchor_words)
secs = time.perf_counter() - t0
print(f"anchor words: {', '.join(vocab[anchor_words])}  ({secs:.2f} s)")
d_i, w_i = np.nonzero(X_txt)
tok_doc, tok_word = np.repeat(d_i, X_txt[d_i, w_i]), np.repeat(w_i, X_txt[d_i, w_i])
theta_anchor = np.zeros((len(X_txt), 6))
for s in range(5):
    z0 = np.random.default_rng(s).integers(6, size=tok_doc.size)
    ndk0 = np.zeros((len(X_txt), 6), np.int64)
    np.add.at(ndk0, (tok_doc, z0), 1)
    _foldin_sweeps(tok_doc, tok_word, z0, ndk0, np.clip(phi_anchor, 1e-9, None), 0.1, 50, s)
    theta_anchor += (ndk0 + 0.1) / (ndk0.sum(1, keepdims=True) + 0.6) / 5
print(score("Anchor words (spectral)", [phi_anchor, phi_anchor], [theta_anchor], secs).round(3))
show_topics(phi_anchor, title="Anchor-word topics:")

def gibbs_from(X, phi0, K, alpha=0.1, beta=0.05, burn=2000, draws=20, thin=10, seed=0):
    """LDA collapsed Gibbs whose token topics start from a fold-in under phi0 (not at random)."""
    g = np.random.default_rng(seed)
    d_i, w_i = np.nonzero(X); n = X[d_i, w_i]; doc, word = np.repeat(d_i, n), np.repeat(w_i, n)
    D, V = X.shape
    z = g.integers(K, size=doc.size); ndk = np.zeros((D, K), np.int64); np.add.at(ndk, (doc, z), 1)
    _foldin_sweeps(doc, word, z, ndk, np.clip(phi0, 1e-9, None), alpha, 30, seed)
    nkw = np.zeros((K, V), np.int64); np.add.at(nkw, (z, word), 1); nk = nkw.sum(1)
    _gibbs_sweeps(doc, word, z, ndk, nkw, nk, alpha, beta, burn, seed + 7)
    phis, thetas = [], []
    for s in range(draws):
        _gibbs_sweeps(doc, word, z, ndk, nkw, nk, alpha, beta, thin, seed + 100 + s)
        phis.append((nkw + beta) / (nk[:, None] + V * beta)); thetas.append((ndk + alpha) / (ndk.sum(1, keepdims=True) + K * alpha))
    return np.array(phis), np.array(thetas)

t0 = time.perf_counter()
anchored = [gibbs_from(X_txt, phi_anchor, 6, seed=s) for s in range(4)]
secs = (time.perf_counter() - t0) / 4
anch_phis = [r[0].mean(0) for r in anchored]
print(score("LDA, Gibbs started from anchors", anch_phis, [r[1].mean(0) for r in anchored], secs).round(3))
print("per-topic stability (Jaccard):", np.round(stability_jaccard(anch_phis), 2))


def token_loglik(phi, theta):
    return np.sum(X_txt * np.log(theta @ phi))


ll_random = [token_loglik(r[0].mean(0), r[1].mean(0)) for r in lda_runs]
ll_anchor = [token_loglik(r[0].mean(0), r[1].mean(0)) for r in anchored]
print("token log-likelihood, random starts: ", np.round(ll_random))
print("token log-likelihood, anchor starts: ", np.round(ll_anchor))
show_topics(anch_phis[0], title="Anchor-started LDA, first run:")

@numba.njit
def _gibbs_sweeps_asym(doc, word, z, ndk, nkw, nk, alpha, beta, n_sweeps, seed):
    np.random.seed(seed)
    K, V = nkw.shape
    cum = np.empty(K)
    for _ in range(n_sweeps):
        for i in range(doc.shape[0]):
            d, w, k = doc[i], word[i], z[i]
            ndk[d, k] -= 1; nkw[k, w] -= 1; nk[k] -= 1
            total = 0.0
            for j in range(K):
                total += (ndk[d, j] + alpha[j]) * (nkw[j, w] + beta) / (nk[j] + V * beta)
                cum[j] = total
            u = np.random.random() * total
            k = 0
            while cum[k] < u:
                k += 1
            z[i] = k; ndk[d, k] += 1; nkw[k, w] += 1; nk[k] += 1

def lda_asym(X, K, alpha0=0.1, beta=0.05, burn=2000, every=20, draws=20, thin=10, seed=0):
    g = np.random.default_rng(seed)
    d_i, w_i = np.nonzero(X); n = X[d_i, w_i]; doc, word = np.repeat(d_i, n), np.repeat(w_i, n)
    D, V = X.shape
    z = g.integers(K, size=doc.size)
    ndk = np.zeros((D, K), np.int64); nkw = np.zeros((K, V), np.int64)
    np.add.at(ndk, (doc, z), 1); np.add.at(nkw, (z, word), 1); nk = nkw.sum(1)
    alpha = np.full(K, alpha0); nd = ndk.sum(1); path = []
    for it in range(burn // every):
        _gibbs_sweeps_asym(doc, word, z, ndk, nkw, nk, alpha, beta, every, seed + it)
        if it >= 5:   # Minka's fixed-point update of an asymmetric Dirichlet (Wallach, Mimno & McCallum 2009)
            for _ in range(5):
                num = (digamma(ndk + alpha) - digamma(alpha)).sum(0)
                den = (digamma(nd + alpha.sum()) - digamma(alpha.sum())).sum()
                alpha = np.maximum(alpha * num / den, 1e-4)
        path.append(alpha.copy())
    phis, thetas = [], []
    for s in range(draws):
        _gibbs_sweeps_asym(doc, word, z, ndk, nkw, nk, alpha, beta, thin, seed + 10_000 + s)
        phis.append((nkw + beta) / (nk[:, None] + V * beta)); thetas.append((ndk + alpha) / (nd[:, None] + alpha.sum()))
    return np.array(phis), np.array(thetas), np.array(path), nk

@numba.njit
def _table_counts(ndk, ab, seed):
    """Antoniak draw: tables m_dk for n_dk customers in a CRP with concentration alpha * beta_k."""
    np.random.seed(seed)
    D, K = ndk.shape
    m = np.zeros(K)
    for d in range(D):
        for k in range(K):
            for j in range(ndk[d, k]):
                if np.random.random() < ab[k] / (ab[k] + j):
                    m[k] += 1
    return m

def hdp_weak_limit(X, K_max=40, alpha=1.0, gamma=1.0, eta=0.05, burn=2000, every=10, draws=20, thin=10, seed=0):
    """Truncated (weak-limit) HDP-LDA: global weights beta ~ Dir(gamma/K_max), documents ~ Dir(alpha * beta)."""
    g = np.random.default_rng(seed)
    d_i, w_i = np.nonzero(X); n = X[d_i, w_i]; doc, word = np.repeat(d_i, n), np.repeat(w_i, n)
    D, V = X.shape
    z = g.integers(K_max, size=doc.size)
    ndk = np.zeros((D, K_max), np.int64); nkw = np.zeros((K_max, V), np.int64)
    np.add.at(ndk, (doc, z), 1); np.add.at(nkw, (z, word), 1); nk = nkw.sum(1)
    beta = np.full(K_max, 1 / K_max); used = []
    for it in range(burn // every):
        _gibbs_sweeps_asym(doc, word, z, ndk, nkw, nk, alpha * beta, eta, every, seed + it)
        m = _table_counts(ndk, alpha * beta, seed + 50_000 + it)
        beta = g.dirichlet(gamma / K_max + m)
        beta = np.maximum(beta, 1e-12); beta /= beta.sum()
        used.append(np.sum(nk > 0.01 * nk.sum()))
    phis, thetas, betas = [], [], []
    for s in range(draws):
        _gibbs_sweeps_asym(doc, word, z, ndk, nkw, nk, alpha * beta, eta, thin, seed + 90_000 + s)
        phis.append((nkw + eta) / (nk[:, None] + V * eta))
        thetas.append((ndk + alpha * beta) / (ndk.sum(1, keepdims=True) + alpha))
        betas.append(beta.copy())
    return np.array(phis), np.array(thetas), np.array(betas), np.array(used), nk

t0 = time.perf_counter()
phi_sym30, theta_sym30, _ = lda_gibbs(X_txt, 30, seed=5, draws=10)
tok_sym30 = (theta_sym30.mean(0) * X_txt.sum(1)[:, None]).sum(0)
phi_asym, theta_asym, alpha_path, nk_asym = lda_asym(X_txt, 30, seed=0)
hdp40 = hdp_weak_limit(X_txt, K_max=40, seed=0)
hdp100 = hdp_weak_limit(X_txt, K_max=100, seed=1)
print(f"four fits: {time.perf_counter() - t0:.0f} s")
spectra = {"LDA, 30 topics, symmetric prior": tok_sym30, "LDA, 30 topics, learned asymmetric prior": nk_asym,
           "HDP, truncated at 40": hdp40[4], "HDP, truncated at 100": hdp100[4]}
for name, counts in spectra.items():
    share = np.sort(counts / counts.sum())[::-1]
    print(f"{name:42s} topics with > 1% of tokens: {np.sum(share > 0.01):3d};  largest topic: {share[0]:.0%}")
print("click-through CV R2: HDP-40", round(cv_r2(hdp40[1].mean(0)), 3), "| HDP-100", round(cv_r2(hdp100[1].mean(0)), 3),
      "| asymmetric-30", round(cv_r2(theta_asym.mean(0)), 3), "| symmetric-30", round(cv_r2(theta_sym30.mean(0)), 3))

fig, ax = plt.subplots(figsize=(9, 4))
for j, (name, counts) in enumerate(spectra.items()):
    share = np.sort(counts / counts.sum())[::-1]
    ax.plot(np.arange(1, len(share) + 1), share, "o-", ms=3, lw=1.6, color=SERIES[j], label=name)
ax.axhline(0.01, color=MUTED, lw=0.8, ls=":")
ax.text(95, 0.0105, "1% of tokens", fontsize=8, color=MUTED, ha="right", va="bottom")
ax.set(xscale="log", yscale="log", xlabel="topic rank", ylabel="share of tokens",
       title="Topic sizes, largest first: nothing switches off")
ax.legend(fontsize=8);

import jax
import jax.numpy as jnp
import optax

def init_params(key, V, K, H=100):
    ks = jax.random.split(key, 5)
    glorot = lambda k, a, b: jax.random.normal(k, (a, b)) * np.sqrt(2 / (a + b))
    return {"W1": glorot(ks[0], V, H), "b1": jnp.zeros(H), "W2": glorot(ks[1], H, H), "b2": jnp.zeros(H),
            "Wmu": glorot(ks[2], H, K), "bmu": jnp.zeros(K), "Wlv": glorot(ks[3], H, K), "blv": jnp.zeros(K),
            "B": glorot(ks[4], K, V)}

def prior(K, alpha=1.0):
    # Laplace approximation of a symmetric Dirichlet(alpha) in the softmax basis (Srivastava & Sutton 2017)
    var = (1 / alpha) * (1 - 2 / K) + (1 / K) * (1 / alpha)
    return 0.0, var

def encode(p, x, key, drop, train):
    h = jax.nn.softplus(x @ p["W1"] + p["b1"])
    h = jax.nn.softplus(h @ p["W2"] + p["b2"])
    if train:
        h = h * jax.random.bernoulli(key, 1 - drop, h.shape) / (1 - drop)
    return h @ p["Wmu"] + p["bmu"], h @ p["Wlv"] + p["blv"]

def batchnorm(a):
    return (a - a.mean(0)) / jnp.sqrt(a.var(0) + 1e-5)

def loss(p, x, key, K, drop=0.2):
    k1, k2 = jax.random.split(key)
    mu, lv = encode(p, x, k1, drop, True)
    eps = jax.random.normal(k2, mu.shape)
    theta = jax.nn.softmax(mu + jnp.exp(0.5 * lv) * eps)
    logits = batchnorm(theta @ p["B"])                    # product of experts: mix in logit space
    rec = -(x * jax.nn.log_softmax(logits)).sum(1)
    m0, v0 = prior(K)
    kl = 0.5 * (jnp.exp(lv) / v0 + (mu - m0) ** 2 / v0 - 1 - lv + jnp.log(v0)).sum(1)
    return (rec + kl).mean()

def fit_prodlda(X, K, seed=0, epochs=200, batch=200, lr=2e-3):
    key = jax.random.PRNGKey(seed)
    key, k0 = jax.random.split(key)
    p = init_params(k0, X.shape[1], K)
    opt = optax.adam(lr, b1=0.99)
    state = opt.init(p)
    Xj = jnp.asarray(X, jnp.float32)
    @jax.jit
    def step(p, state, xb, key):
        l, grads = jax.value_and_grad(loss)(p, xb, key, K)
        upd, state = opt.update(grads, state, p)
        return optax.apply_updates(p, upd), state, l
    g = np.random.default_rng(seed); losses = []
    for ep in range(epochs):
        perm = g.permutation(len(X))
        for i in range(0, len(X) - batch + 1, batch):
            key, kb = jax.random.split(key)
            p, state, l = step(p, state, Xj[perm[i:i + batch]], kb)
        losses.append(float(l))
    mu, _ = encode(p, Xj, key, 0.0, False)
    return np.asarray(p["B"]), np.asarray(jax.nn.softmax(mu)), np.array(losses)

t0 = time.perf_counter()
prod_runs = [fit_prodlda(X_txt, 6, seed=s, epochs=300, lr=5e-3) for s in range(4)]
secs = (time.perf_counter() - t0) / 4
prod_phis = []
for B, _, _ in prod_runs:
    e = np.exp(B - B.max(1, keepdims=True))
    prod_phis.append(e / e.sum(1, keepdims=True))
print("final training loss per run:", [round(float(np.mean(r[2][-10:])), 1) for r in prod_runs])
print(score("ProdLDA (neural, JAX), 6 topics", prod_phis, [r[1] for r in prod_runs], secs).round(3))
show_topics(prod_phis[0], title="ProdLDA, first run:")
print("typical largest topic share of a document:", round(float(np.median(prod_runs[0][1].max(1))), 2))