structural-topic-model 0.2.0__py3-none-any.whl

pystm/__init__.py

@@ -0,0 +1,31 @@
+"""pystm: Python implementation of the Structural Topic Model.
+
+A port of the R ``stm`` package (Roberts, Stewart & Tingley) with an API
+modeled on scikit-learn's ``LatentDirichletAllocation``.
+"""
+
+from .diagnostics import (
+    TopicCorrelations,
+    check_residuals,
+    exclusivity,
+    semantic_coherence,
+    topic_corr,
+)
+from .effects import EstimatedEffects, estimate_effect
+from .model_selection import eval_heldout, make_heldout, search_k
+from .stm import StructuralTopicModel
+
+__all__ = [
+    "StructuralTopicModel",
+    "estimate_effect",
+    "EstimatedEffects",
+    "search_k",
+    "make_heldout",
+    "eval_heldout",
+    "topic_corr",
+    "TopicCorrelations",
+    "semantic_coherence",
+    "exclusivity",
+    "check_residuals",
+]
+__version__ = "0.2.0"

pystm/_estep.py

@@ -0,0 +1,159 @@
+"""Variational E-step (port of STMestep.R, STMlncpp.R and STMCfuns.cpp).
+
+For each document the variational posterior over eta (the K-1 dimensional
+logistic-normal document-topic parameter) is approximated by a Laplace
+approximation: the mode ``lambda`` is found with BFGS and the covariance
+``nu`` is the inverse Hessian at the mode.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.linalg import cho_solve, cholesky, solve
+from scipy.optimize import minimize
+
+# Numerical guards absent from the R/C++ original: BFGS line searches can
+# probe points where exp(eta) over/underflows or a document word has
+# (numerically) zero probability under every topic, poisoning the search
+# with inf/nan.  Clipping eta and flooring the logs/denominators keeps the
+# objective finite without affecting values in the normal range.
+_ETA_CLIP = 200.0
+_TINY = 1e-300
+
+
+def _expeta(eta):
+    return np.append(np.exp(np.clip(eta, -_ETA_CLIP, _ETA_CLIP)), 1.0)
+
+
+def neg_lhood(eta, beta_d, doc_ct, mu_d, siginv):
+    """Negative collapsed objective for one document (lhoodcpp)."""
+    expeta = _expeta(eta)
+    ndoc = doc_ct.sum()
+    word_probs = np.maximum(expeta @ beta_d, _TINY)
+    part1 = np.log(word_probs) @ doc_ct - ndoc * np.log(expeta.sum())
+    diff = eta - mu_d
+    part2 = 0.5 * diff @ siginv @ diff
+    return part2 - part1
+
+
+def neg_grad(eta, beta_d, doc_ct, mu_d, siginv):
+    """Gradient of :func:`neg_lhood` (gradcpp)."""
+    expeta = _expeta(eta)
+    EB = beta_d * expeta[:, None]
+    denom = np.maximum(EB.sum(axis=0), _TINY)
+    part1 = EB @ (doc_ct / denom) - (doc_ct.sum() / expeta.sum()) * expeta
+    part2 = siginv @ (eta - mu_d)
+    return part2 - part1[:-1]
+
+
+def hessian_phi_bound(eta, beta_d, doc_ct, mu_d, siginv, sigmaentropy):
+    """Compute the Laplace covariance, token assignments and bound (hpbcpp).
+
+    Returns ``(phis, nu, bound)`` where ``phis`` is K x V_d expected token
+    counts per topic, ``nu`` is the (K-1) x (K-1) posterior covariance of
+    eta, and ``bound`` is the document's contribution to the global ELBO.
+    """
+    expeta = _expeta(eta)
+    theta = expeta / expeta.sum()
+    ndoc = doc_ct.sum()
+    sqrtct = np.sqrt(doc_ct)
+
+    EB = beta_d * expeta[:, None]
+    EB *= (sqrtct / np.maximum(EB.sum(axis=0), _TINY))[None, :]
+
+    hess = EB @ EB.T - ndoc * np.outer(theta, theta)
+    # turn EB into phi (expected token counts per topic and word)
+    EB *= sqrtct[None, :]
+    np.fill_diagonal(hess, np.diag(hess) - (EB.sum(axis=1) - ndoc * theta))
+    hess = hess[:-1, :-1] + siginv
+
+    try:
+        L = cholesky(hess, lower=True)
+    except np.linalg.LinAlgError:
+        # not positive definite: enforce diagonal dominance as in hpbcpp
+        dvec = np.diag(hess).copy()
+        magnitudes = np.abs(hess).sum(axis=1) - np.abs(dvec)
+        dvec = np.maximum(dvec, magnitudes)
+        np.fill_diagonal(hess, dvec)
+        L = cholesky(hess, lower=True)
+
+    det_term = -np.log(np.diag(L)).sum()
+    nu = cho_solve((L, True), np.eye(hess.shape[0]))
+
+    diff = eta - mu_d
+    bound = (
+        np.log(np.maximum(theta @ beta_d, _TINY)) @ doc_ct
+        + det_term
+        - 0.5 * diff @ siginv @ diff
+        - sigmaentropy
+    )
+    return EB, nu, bound
+
+
+def optimize_document(eta, beta_d, doc_ct, mu_d, siginv, sigmaentropy,
+                      max_optim_iter=500):
+    """Infer one document's variational parameters (logisticnormalcpp)."""
+    res = minimize(
+        neg_lhood,
+        eta,
+        args=(beta_d, doc_ct, mu_d, siginv),
+        jac=neg_grad,
+        method="BFGS",
+        options={"maxiter": max_optim_iter},
+    )
+    eta_hat = res.x if np.isfinite(res.x).all() else eta
+    return (hessian_phi_bound(eta_hat, beta_d, doc_ct, mu_d, siginv,
+                              sigmaentropy), eta_hat)
+
+
+def decompose_sigma(sigma):
+    """Precompute the inverse and entropy term shared by all documents."""
+    try:
+        chol_u = cholesky(sigma, lower=False)
+        sigmaentropy = np.log(np.diag(chol_u)).sum()
+        siginv = cho_solve((chol_u, False), np.eye(sigma.shape[0]))
+    except np.linalg.LinAlgError:
+        sigmaentropy = 0.5 * np.linalg.slogdet(sigma)[1]
+        siginv = solve(sigma, np.eye(sigma.shape[0]), assume_a="sym")
+    return siginv, sigmaentropy
+
+
+def estep(docs, beta_index, update_mu, beta, lambda_old, mu, sigma,
+          max_optim_iter=500):
+    """Run the E-step over all documents and accumulate sufficient stats.
+
+    Parameters mirror estep() in STMestep.R.  ``docs`` is a list of
+    ``(word_indices, counts)`` pairs, ``beta`` a list with one K x V matrix
+    per content level (always length 1 here), ``mu`` is (K-1,) when shared
+    or (N, K-1) when document specific.
+
+    Returns ``(sigma_ss, beta_ss, bound, lambda_)``.
+    """
+    K, V = beta[0].shape
+    N = len(docs)
+    A = len(beta)
+
+    sigma_ss = np.zeros((K - 1, K - 1))
+    beta_ss = [np.zeros((K, V)) for _ in range(A)]
+    bound = np.empty(N)
+    lambda_ = np.empty((N, K - 1))
+
+    siginv, sigmaentropy = decompose_sigma(sigma)
+
+    for i, (words, counts) in enumerate(docs):
+        aspect = beta_index[i]
+        init = lambda_old[i]
+        mu_d = mu[i] if update_mu else mu
+        beta_d = np.ascontiguousarray(beta[aspect][:, words])
+
+        (phis, nu, bnd), eta_hat = optimize_document(
+            init, beta_d, counts, mu_d, siginv, sigmaentropy,
+            max_optim_iter=max_optim_iter,
+        )
+
+        sigma_ss += nu
+        beta_ss[aspect][:, words] += phis
+        bound[i] = bnd
+        lambda_[i] = eta_hat
+
+    return sigma_ss, beta_ss, bound, lambda_

pystm/_mnreg.py

@@ -0,0 +1,221 @@
+"""Content covariate M-step (port of STMmnreg.R).
+
+The SAGE-style topic-word update with content covariates is estimated via
+Distributed Multinomial Regression (Taddy 2013): the multinomial is
+factorized into independent Poisson regressions, one per vocabulary word,
+each with an L1 penalty.  The R package solves these with glmnet; here we
+implement an equivalent lasso-penalized Poisson solver (IRLS + coordinate
+descent over a regularization path with information-criterion selection).
+
+The solver exploits the structure of the problem heavily.  The design
+matrix consists of three groups of indicator columns (topic main effects,
+aspect main effects, topic-by-aspect interactions).  Columns within a
+group touch disjoint rows, so a coordinate-descent pass over a whole
+group can be performed as one vectorized update; and since every word
+shares the same design, all V regressions are advanced simultaneously.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.special import xlogy
+
+
+def _poisson_deviance(Y, mu):
+    """Per-word Poisson deviance, with the y=0 terms handled."""
+    return 2.0 * (xlogy(Y, Y) - xlogy(Y, mu) - (Y - mu)).sum(axis=(0, 1))
+
+
+def _soft_threshold(rho, lam):
+    return np.sign(rho) * np.maximum(np.abs(rho) - lam, 0.0)
+
+
+class _StructuredPoissonLasso:
+    """Distributed Poisson lasso with the STM content design.
+
+    Data is held as (A, K, V) arrays.  ``lam`` arguments are per-word
+    penalty levels of shape (V,).
+    """
+
+    def __init__(self, Y3, offsets3, use_aspect, use_inter):
+        self.Y3 = Y3
+        self.offsets3 = offsets3
+        self.A, self.K, self.V = Y3.shape
+        self.n = self.A * self.K
+        self.use_aspect = use_aspect
+        self.use_inter = use_inter
+        self.b_topic = np.zeros((self.K, self.V))
+        self.b_aspect = np.zeros((self.A, self.V)) if use_aspect else None
+        self.b_inter = (
+            np.zeros((self.A, self.K, self.V)) if use_inter else None
+        )
+
+    def linpred(self):
+        eta = self.offsets3 + self.b_topic[None, :, :]
+        if self.use_aspect:
+            eta = eta + self.b_aspect[:, None, :]
+        if self.use_inter:
+            eta = eta + self.b_inter
+        return np.clip(eta, -50.0, 30.0)
+
+    def _block_update(self, R, W, Wsum, b, lam, axis):
+        """One exact CD pass over a group of disjoint indicator columns.
+
+        ``axis`` is the (A, K, V) axis summed over to aggregate a column's
+        rows (None for the interaction block where each row is its own
+        column).  Updates ``R`` in place and returns (new_b, max_delta).
+        """
+        if axis is None:
+            num = R + W * b
+            denom = Wsum
+        else:
+            num = R.sum(axis=axis) + Wsum * b
+            denom = Wsum
+        rho = num / self.n
+        b_new = _soft_threshold(rho, lam)
+        with np.errstate(invalid="ignore", divide="ignore"):
+            b_new = np.where(denom / self.n > 1e-12,
+                             b_new / (denom / self.n), 0.0)
+        delta = b_new - b
+        max_delta = np.abs(delta).max() if delta.size else 0.0
+        if max_delta > 0:
+            if axis is None:
+                R -= W * delta
+            else:
+                R -= W * np.expand_dims(delta, axis)
+        return b_new, max_delta
+
+    def fit_one_lambda(self, lam, tol, max_irls, max_sweeps):
+        """Solve at one penalty level, warm-starting from current state."""
+        for _ in range(max_irls):
+            eta = self.linpred()
+            mu = np.exp(eta)
+            W = mu
+            R = self.Y3 - mu
+            Wsum_topic = W.sum(axis=0)
+            Wsum_aspect = W.sum(axis=1) if self.use_aspect else None
+            outer_delta = 0.0
+            for _ in range(max_sweeps):
+                d = 0.0
+                self.b_topic, d1 = self._block_update(
+                    R, W, Wsum_topic, self.b_topic, lam, axis=0)
+                d = max(d, d1)
+                if self.use_aspect:
+                    self.b_aspect, d1 = self._block_update(
+                        R, W, Wsum_aspect, self.b_aspect, lam, axis=1)
+                    d = max(d, d1)
+                if self.use_inter:
+                    self.b_inter, d1 = self._block_update(
+                        R, W, W, self.b_inter, lam, axis=None)
+                    d = max(d, d1)
+                outer_delta = max(outer_delta, d)
+                if d < tol:
+                    break
+            if outer_delta < tol:
+                break
+
+    def df(self):
+        out = (self.b_topic != 0).sum(axis=0)
+        if self.use_aspect:
+            out = out + (self.b_aspect != 0).sum(axis=0)
+        if self.use_inter:
+            out = out + (self.b_inter != 0).sum(axis=(0, 1))
+        return out
+
+    def coef_rows(self):
+        """Stack coefficients as the R package's kappa params (p, V)."""
+        rows = [self.b_topic]
+        if self.use_aspect:
+            rows.append(self.b_aspect)
+        if self.use_inter:
+            rows.append(self.b_inter.reshape(self.n, self.V))
+        return np.vstack(rows)
+
+    def state(self):
+        return (self.b_topic.copy(),
+                None if self.b_aspect is None else self.b_aspect.copy(),
+                None if self.b_inter is None else self.b_inter.copy())
+
+    def set_state(self, state):
+        self.b_topic, self.b_aspect, self.b_inter = (
+            state[0].copy(),
+            None if state[1] is None else state[1].copy(),
+            None if state[2] is None else state[2].copy(),
+        )
+
+
+def mnreg(beta_ss, wcounts, *, interactions=True, nlambda=250,
+          lambda_min_ratio=0.001, ic_k=2.0, tol=1e-4,
+          max_irls=4, max_sweeps=8):
+    """Update beta and kappa from the E-step expected counts (mnreg in R).
+
+    Only the (default) fixed-intercept variant is implemented: the
+    intercept of each word's Poisson regression is fixed to the background
+    log-probability ``m``.
+
+    Returns ``(beta, kappa)`` where ``beta`` is a list of A matrices of
+    shape (K, V) and ``kappa`` is a dict with the baseline ``m`` and the
+    selected deviation coefficients ``params`` of shape (p, V).
+    """
+    A = len(beta_ss)
+    K, V = beta_ss[0].shape
+    use_aspect = A > 1
+    use_inter = interactions and A > 1
+
+    Y3 = np.stack(beta_ss)  # (A, K, V)
+    m = np.log(wcounts) - np.log(wcounts.sum())
+    row_totals = np.maximum(Y3.sum(axis=2), 1e-10)  # (A, K)
+    offsets3 = m[None, None, :] + np.log(row_totals)[:, :, None]
+
+    solver = _StructuredPoissonLasso(Y3, offsets3, use_aspect, use_inter)
+    n = A * K
+
+    mu0 = np.exp(np.clip(offsets3, -50.0, 30.0))
+    nulldev = _poisson_deviance(Y3, mu0)
+    # per-word lambda_max: max over columns of |score| at the null model
+    R0 = Y3 - mu0
+    scores = [np.abs(R0.sum(axis=0))]            # topic columns (K, V)
+    if use_aspect:
+        scores.append(np.abs(R0.sum(axis=1)))    # aspect columns (A, V)
+    if use_inter:
+        scores.append(np.abs(R0).reshape(n, V))  # interaction columns
+    lambda_max = np.vstack(scores).max(axis=0) / n
+    lambda_max = np.maximum(lambda_max, 1e-10)
+
+    rel_path = np.exp(np.linspace(0.0, np.log(lambda_min_ratio), nlambda))
+    best_ic = nulldev.copy()  # path point 0: all coefficients zero
+    best_state = solver.state()
+    any_improved = False
+
+    for step in rel_path[1:]:
+        lam = lambda_max * step
+        solver.fit_one_lambda(lam, tol=tol, max_irls=max_irls,
+                              max_sweeps=max_sweeps)
+        mu = np.exp(solver.linpred())
+        dev = 2.0 * (xlogy(Y3, Y3) - xlogy(Y3, mu) - (Y3 - mu)).sum(axis=(0, 1))
+        ic = dev + ic_k * solver.df()
+        improved = ic < best_ic
+        if improved.any():
+            any_improved = True
+            cur = solver.state()
+            best_state[0][:, improved] = cur[0][:, improved]
+            if cur[1] is not None:
+                best_state[1][:, improved] = cur[1][:, improved]
+            if cur[2] is not None:
+                best_state[2][:, :, improved] = cur[2][:, :, improved]
+            best_ic[improved] = ic[improved]
+
+    final = _StructuredPoissonLasso(Y3, offsets3, use_aspect, use_inter)
+    if any_improved:
+        final.set_state(best_state)
+
+    linpred = final.linpred().reshape(n, V) - (
+        np.log(row_totals).reshape(n)[:, None]
+    )
+    linpred -= linpred.max(axis=1, keepdims=True)
+    explinpred = np.exp(linpred)
+    beta_full = explinpred / explinpred.sum(axis=1, keepdims=True)
+
+    beta = [beta_full[a * K:(a + 1) * K] for a in range(A)]
+    kappa = {"m": m, "params": final.coef_rows()}
+    return beta, kappa

pystm/_mstep.py

@@ -0,0 +1,96 @@
+"""M-step updates (port of STMmu.R, STMsigma.R, STMoptbeta.R)."""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.linalg import cho_solve, cholesky
+
+
+class PrevalenceRegressionError(RuntimeError):
+    pass
+
+
+def vb_variational_reg(Y, X, Xcorr=None, b0=1.0, d0=1.0, max_iter=1000):
+    """Variational linear regression with a half-Cauchy hyperprior.
+
+    Port of vb.variational.reg(); the first column of ``X`` is assumed to
+    be the (unpenalized) intercept.
+    """
+    if Xcorr is None:
+        Xcorr = X.T @ X
+    XYcorr = X.T @ Y
+
+    N, D = X.shape
+    an = (1 + N) / 2
+    w = np.zeros(D)
+    error_prec = 1.0
+    cn = D
+    dn = 1.0
+    Ea = cn / dn
+    ba = 1.0
+
+    for _ in range(max_iter):
+        w_old = w
+
+        prior_diag = np.full(D, Ea)
+        prior_diag[0] = 0.0
+        invV = error_prec * Xcorr + np.diag(prior_diag)
+        L = cholesky(invV, lower=True)
+        V = cho_solve((L, True), np.eye(D))
+        w = error_prec * (V @ XYcorr)
+
+        sse = np.sum((X @ w - Y) ** 2)
+        bn = 0.5 * (sse + np.trace(Xcorr @ V)) + ba
+        error_prec = an / bn
+        ba = 1.0 / (error_prec + b0)
+
+        da = 2.0 / (Ea + d0)
+        dn = 2.0 * da + (w[1:] @ w[1:] + np.diag(V)[1:].sum())
+        Ea = cn / dn
+
+        if np.abs(w - w_old).sum() < 1e-4:
+            return w
+
+    raise PrevalenceRegressionError(
+        "Prevalence regression failed to converge within the iteration "
+        "limit. You can raise it with gamma_max_iter."
+    )
+
+
+def opt_mu(lambda_, covar=None, max_iter=1000):
+    """Update the prevalence model (opt.mu, modes CTM and Pooled).
+
+    Returns ``(mu, gamma)``.  Without covariates (CTM mode) ``mu`` is the
+    shared (K-1,) mean and ``gamma`` is None.  With covariates ``mu`` is
+    (N, K-1), the per-document prior means, and ``gamma`` is (P, K-1).
+    """
+    if covar is None:
+        return lambda_.mean(axis=0), None
+
+    Xcorr = covar.T @ covar
+    gamma = np.column_stack([
+        vb_variational_reg(lambda_[:, k], covar, Xcorr=Xcorr, max_iter=max_iter)
+        for k in range(lambda_.shape[1])
+    ])
+    mu = covar @ gamma
+    return mu, gamma
+
+
+def opt_sigma(nu, lambda_, mu, sigprior):
+    """Update the global covariance matrix (opt.sigma).
+
+    ``mu`` is (K-1,) for the shared mean or (N, K-1) for covariate models.
+    """
+    if mu.ndim == 1:
+        diff = lambda_ - mu[None, :]
+    else:
+        diff = lambda_ - mu
+    sigma = (diff.T @ diff + nu) / lambda_.shape[0]
+    return sigprior * np.diag(np.diag(sigma)) + (1 - sigprior) * sigma
+
+
+def opt_beta(beta_ss):
+    """Update the topic-word distributions (opt.beta, LDA-beta mode)."""
+    row_sums = beta_ss[0].sum(axis=1, keepdims=True)
+    row_sums[row_sums == 0] = 1.0
+    return [beta_ss[0] / row_sums]

pystm/_spectral.py

@@ -0,0 +1,151 @@
+"""Spectral initialization via anchor words (port of spectral.R).
+
+Implements the method of Arora et al. (2013): build the word co-occurrence
+gram matrix, greedily select anchor words, then recover the topic-word
+matrix with RecoverL2.  The R package solves the simplex-constrained
+regression exactly with quadprog by default; here we use a penalized NNLS
+formulation which matches it closely.  The exponentiated gradient
+algorithm (R's ``recoverEG=TRUE`` option) is available as an alternative.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.optimize import nnls
+from scipy.sparse import csr_matrix
+
+
+def gram(mat: csr_matrix) -> np.ndarray:
+    """Word co-occurrence gram matrix from a sparse doc-term matrix."""
+    nd = np.asarray(mat.sum(axis=1)).ravel()
+    keep = nd >= 2  # undefined for docs with fewer than 2 tokens
+    mat = mat[keep]
+    nd = nd[keep]
+    divisor = nd * (nd - 1)
+
+    htilde = mat.multiply(1.0 / np.sqrt(divisor)[:, None]).tocsr()
+    hhat = np.asarray(mat.multiply(1.0 / divisor[:, None]).sum(axis=0)).ravel()
+    Q = (htilde.T @ htilde).toarray()
+    Q[np.diag_indices_from(Q)] -= hhat
+    return Q
+
+
+def fast_anchor(Qbar: np.ndarray, K: int) -> np.ndarray:
+    """Greedy anchor word selection by stabilized Gram-Schmidt."""
+    Qbar = Qbar.copy()
+    basis = np.zeros(K, dtype=np.int64)
+    row_squared_sums = (Qbar**2).sum(axis=1)
+
+    for i in range(K):
+        basis[i] = int(np.argmax(row_squared_sums))
+        max_val = row_squared_sums[basis[i]]
+        Qbar[basis[i]] *= 1.0 / np.sqrt(max_val)
+
+        inner_products = Qbar @ Qbar[basis[i]]
+        project = np.outer(inner_products, Qbar[basis[i]])
+        project[basis[: i + 1]] = 0.0
+        Qbar -= project
+        row_squared_sums = (Qbar**2).sum(axis=1)
+        row_squared_sums[basis[: i + 1]] = 0.0
+    return basis
+
+
+def expgrad(X, y, XtX=None, alpha=None, tol=1e-7, max_iter=500):
+    """Exponentiated gradient for simplex-constrained least squares."""
+    if alpha is None:
+        alpha = np.full(X.shape[0], 1.0 / X.shape[0])
+    if XtX is None:
+        XtX = X @ X.T
+    ytX = y @ X.T
+
+    eta = 50.0
+    sse_old = np.inf
+    for _ in range(max_iter):
+        grad = ytX - alpha @ XtX
+        sse = grad @ grad
+        grad = 2.0 * eta * grad
+        alpha = alpha * np.exp(grad - grad.max())
+        alpha = alpha / alpha.sum()
+        if abs(np.sqrt(sse_old) - np.sqrt(sse)) < tol:
+            break
+        sse_old = sse
+    return alpha
+
+
+def recover_l2(Qbar, anchors, wprob, solver="nnls"):
+    """Recover the K x V topic-word matrix from the anchor rows.
+
+    Each word's row of ``Qbar`` is expressed as a convex combination of
+    the anchor rows.  ``solver="nnls"`` enforces the sum-to-one constraint
+    through a heavily weighted penalty row in a non-negative least-squares
+    problem (the analogue of the exact quadprog solve used by the R
+    package); ``solver="expgrad"`` uses exponentiated gradient descent
+    (R's ``recoverEG=TRUE``).
+    """
+    X = Qbar[anchors]
+    XtX = X @ X.T
+    K = len(anchors)
+    anchor_pos = {a: idx for idx, a in enumerate(anchors)}
+
+    if solver == "nnls":
+        penalty = 1000.0
+        X_aug = np.vstack([X.T, np.full(K, penalty)])
+
+    weights = np.empty((Qbar.shape[0], K))
+    for i in range(Qbar.shape[0]):
+        if i in anchor_pos:
+            vec = np.zeros(K)
+            vec[anchor_pos[i]] = 1.0
+            weights[i] = vec
+        elif solver == "nnls":
+            solution, _ = nnls(X_aug, np.append(Qbar[i], penalty))
+            solution = np.maximum(solution, np.finfo(np.float64).eps)
+            weights[i] = solution / solution.sum()
+        else:
+            solution = expgrad(X, Qbar[i], XtX)
+            solution = np.maximum(solution, np.finfo(np.float64).eps)
+            weights[i] = solution
+
+    A = weights * wprob[:, None]
+    A = A.T / A.sum(axis=0)[None, :].T
+    return A
+
+
+def spectral_init(X: csr_matrix, K: int, max_vocab: int | None = 10000,
+                  solver: str = "nnls") -> np.ndarray:
+    """Spectral initialization of beta (the Spectral branch of stm.init)."""
+    V = X.shape[1]
+    if K >= V:
+        raise ValueError(
+            "Spectral initialization cannot be used when K >= vocabulary size."
+        )
+
+    wprob = np.asarray(X.sum(axis=0)).ravel().astype(np.float64)
+    wprob /= wprob.sum()
+
+    keep = None
+    if max_vocab is not None and V > max_vocab:
+        keep = np.argsort(wprob)[::-1][:max_vocab]
+        X = X[:, keep]
+        wprob = wprob[keep]
+
+    Q = gram(X)
+    Qsums = Q.sum(axis=1)
+    if np.any(Qsums == 0):
+        nonzero = Qsums != 0
+        keep = np.where(nonzero)[0] if keep is None else keep[nonzero]
+        Q = Q[np.ix_(nonzero, nonzero)]
+        Qsums = Qsums[nonzero]
+        wprob = wprob[nonzero]
+    Qbar = Q / Qsums[:, None]
+
+    anchors = fast_anchor(Qbar, K)
+    beta = recover_l2(Qbar, anchors, wprob, solver=solver)
+
+    if keep is not None:
+        # reintroduce dropped words with a small amount of mass
+        beta_new = np.zeros((K, V))
+        beta_new[:, keep] = beta
+        beta_new += 0.001 / V
+        beta = beta_new / beta_new.sum(axis=1, keepdims=True)
+    return beta

pystm/_utils.py

@@ -0,0 +1,41 @@
+"""Shared numerical utilities (port of STMfunctions.R)."""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.sparse import csr_matrix, issparse
+from scipy.special import logsumexp
+
+
+def safelog(x: np.ndarray, min_value: float = -1000.0) -> np.ndarray:
+    """log(x) with -inf (and anything below ``min_value``) clamped."""
+    with np.errstate(divide="ignore"):
+        out = np.log(x)
+    return np.maximum(out, min_value)
+
+
+def row_softmax(mat: np.ndarray) -> np.ndarray:
+    """Row-wise softmax of a 2-d array."""
+    return np.exp(mat - logsumexp(mat, axis=1, keepdims=True))
+
+
+def to_doc_list(X) -> list[tuple[np.ndarray, np.ndarray]]:
+    """Convert a document-term count matrix to the stm document format.
+
+    Returns one ``(word_indices, word_counts)`` pair per document, holding
+    the column indices of the document's distinct terms and their counts
+    (the two rows of the R package's document matrices, zero-indexed).
+    """
+    X = csr_matrix(X) if not issparse(X) else X.tocsr()
+    if (X.data < 0).any():
+        raise ValueError("X must contain non-negative counts.")
+    if np.any(X.data != np.round(X.data)):
+        raise ValueError("X must contain integer counts.")
+    docs = []
+    for i in range(X.shape[0]):
+        start, end = X.indptr[i], X.indptr[i + 1]
+        words = X.indices[start:end].astype(np.int64)
+        counts = X.data[start:end].astype(np.float64)
+        keep = counts > 0
+        docs.append((words[keep], counts[keep]))
+    return docs