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

pystm/diagnostics.py

@@ -0,0 +1,168 @@
+"""Model diagnostics (ports of semanticCoherence.R, exclusivity.R,
+residuals.R and topicCorr.R)."""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy import stats
+from scipy.sparse import csc_matrix, csr_matrix, issparse
+from scipy.stats import rankdata
+
+from ._utils import safelog
+
+
+def _as_csr(X):
+    return csr_matrix(X) if not issparse(X) else X.tocsr()
+
+
+def _semcoh_one_beta(X, logbeta, M):
+    """Semantic coherence per topic for one beta (semCoh1beta in R)."""
+    K = logbeta.shape[0]
+    top_words = np.argsort(-logbeta, axis=1)[:, :M]
+    wordlist, positions = np.unique(top_words, return_inverse=True)
+    labels = positions.reshape(K, M)
+
+    sub = csc_matrix(X[:, wordlist])
+    sub.data = np.minimum(sub.data, 1.0)  # binarize
+    cross = (sub.T @ sub).toarray()  # document co-occurrence counts
+
+    result = np.zeros(K)
+    for k in range(K):
+        idx = labels[k]
+        for a in range(M):
+            for b in range(M):
+                m_, l_ = idx[a], idx[b]
+                if m_ > l_:
+                    result[k] += (np.log(0.01 + cross[m_, l_])
+                                  - np.log(cross[l_, l_] + 0.01))
+    return result
+
+
+def semantic_coherence(model, X, content=None, M=10):
+    """Semantic coherence (Mimno et al. 2011) per topic.
+
+    Higher is better; the metric checks that a topic's top ``M`` words
+    co-occur within documents.  For content covariate models pass the
+    content labels; the score is then the document-weighted average over
+    the aspect-specific betas, as in the R package.
+    """
+    X = _as_csr(X)
+    if model.aspect_components_ is None:
+        return _semcoh_one_beta(X, safelog(model.components_), M)
+    if content is None:
+        raise ValueError(
+            "The model was fitted with a content covariate; pass the "
+            "matching content labels."
+        )
+    levels = model.content_levels_
+    index = np.searchsorted(levels, np.asarray(content).ravel())
+    result = np.zeros(model.n_components)
+    for a in range(len(levels)):
+        subset = index == a
+        if not subset.any():
+            continue
+        logbeta = safelog(model.aspect_components_[a])
+        result += _semcoh_one_beta(X[subset], logbeta, M) * subset.sum()
+    return result / X.shape[0]
+
+
+def exclusivity(model, M=10, frexw=0.7):
+    """FREX-based exclusivity per topic (exclusivity in R).
+
+    Not defined for content covariate models (matching the R package).
+    """
+    if model.aspect_components_ is not None:
+        raise ValueError(
+            "Exclusivity calculation is only designed for models without "
+            "content covariates."
+        )
+    w = frexw
+    tbeta = model.components_.T  # (V, K)
+    mat = tbeta / tbeta.sum(axis=1, keepdims=True)
+
+    ex = np.apply_along_axis(rankdata, 0, mat) / mat.shape[0]
+    fr = np.apply_along_axis(rankdata, 0, tbeta) / mat.shape[0]
+    frex = 1.0 / (w / ex + (1 - w) / fr)
+    index = np.argsort(-tbeta, axis=0)[:M]
+    return np.array([
+        frex[index[:, k], k].sum() for k in range(tbeta.shape[1])
+    ])
+
+
+def check_residuals(model, X, content=None, tol=0.01):
+    """Multinomial dispersion of the residuals (Taddy 2012).
+
+    Under a correctly specified model the dispersion is 1; values above 1
+    suggest the number of topics is too small.  Returns a dict with
+    ``dispersion``, ``pvalue`` and ``df``.
+    """
+    X = _as_csr(X)
+    n, V = X.shape
+    K = model.n_components
+    theta = model.theta_
+    if theta.shape[0] != n:
+        raise ValueError("X must be the corpus the model was fitted on.")
+
+    beta = model._beta_list()
+    if model.content_levels_ is not None:
+        if content is None:
+            raise ValueError(
+                "The model was fitted with a content covariate; pass the "
+                "matching content labels."
+            )
+        index = np.searchsorted(model.content_levels_,
+                                np.asarray(content).ravel())
+    else:
+        index = np.zeros(n, dtype=np.int64)
+
+    d = n * (K - 1) + K * (V - 1)
+    D = 0.0
+    Nhat = 0
+    for i in range(n):
+        row = X.getrow(i)
+        q = theta[i] @ beta[index[i]]  # (V,)
+        m = row.sum()
+        Nhat += int((q * m > tol).sum())
+        x = np.zeros(V)
+        x[row.indices] = row.data
+        denom = m * q * (1 - q)
+        D += ((x**2 - 2 * x * q * m) / denom).sum() + (m * q / (1 - q)).sum()
+
+    df = Nhat - V - d
+    with np.errstate(invalid="ignore"):
+        dispersion = D / df
+        pvalue = stats.chi2.sf(D, df) if df > 0 else np.nan
+    return {"dispersion": dispersion, "pvalue": pvalue, "df": df}
+
+
+class TopicCorrelations:
+    """Result of :func:`topic_corr` (class topicCorr in R).
+
+    Attributes
+    ----------
+    cor : ndarray (K, K)
+        Correlation matrix with entries below ``cutoff`` (in absolute
+        value) set to zero.
+    posadj : ndarray (K, K)
+        Adjacency matrix of positive correlations above the cutoff.
+    poscor : ndarray (K, K)
+        Correlations masked to the positive adjacency structure.
+    """
+
+    def __init__(self, cor, posadj, poscor):
+        self.cor = cor
+        self.posadj = posadj
+        self.poscor = poscor
+
+
+def topic_corr(model, cutoff=0.01):
+    """Topic correlation graph from theta (topicCorr method="simple").
+
+    The R package's "huge" method (semiparametric graphical model
+    selection) depends on the huge package and is not implemented.
+    """
+    cormat = np.corrcoef(model.theta_, rowvar=False)
+    posadj = (cormat > cutoff).astype(float)
+    poscor = cormat * posadj
+    cor = np.where(np.abs(cormat) > cutoff, cormat, 0.0)
+    return TopicCorrelations(cor=cor, posadj=posadj, poscor=poscor)

pystm/effects.py

@@ -0,0 +1,203 @@
+"""Covariate effect estimation (port of estimateEffect.R / thetaPosterior.R).
+
+Regressions where topic proportions are the outcome, propagating the
+measurement uncertainty of theta via the method of composition: draw
+theta from the variational posterior, run the OLS, repeat, then pool.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+from scipy import stats
+from sklearn.utils import check_random_state
+
+from ._utils import row_softmax
+
+
+def _global_sigma(model):
+    """Global approximation to the per-document posterior covariance.
+
+    Subtracts the contribution of deviations from the prior mean out of
+    the topic covariance, leaving the (average) variational covariance
+    (thetapost.global in R).
+    """
+    lambda_ = model.eta_
+    mu = model.mu_
+    diff = lambda_ - (mu[None, :] if mu.ndim == 1 else mu)
+    covariance = (diff.T @ diff) / lambda_.shape[0]
+    sigma = model.sigma_ - covariance
+    # guard against indefiniteness from the subtraction
+    evals, evecs = np.linalg.eigh(sigma)
+    if evals[0] <= 0:
+        evals = np.maximum(evals, 1e-10)
+        sigma = (evecs * evals) @ evecs.T
+    return sigma
+
+
+def _draw_theta(model, rng):
+    """One draw of theta for every document (Global approximation)."""
+    sigma = _global_sigma(model)
+    chol = np.linalg.cholesky(sigma)
+    z = rng.standard_normal(model.eta_.shape)
+    eta = model.eta_ + z @ chol.T
+    return row_softmax(
+        np.column_stack([eta, np.zeros(eta.shape[0])])
+    )
+
+
+class _QRRegression:
+    """OLS with a cached QR decomposition (qr.lm / summary.qr.lm in R)."""
+
+    def __init__(self, xmat, prior=None):
+        self.n_obs = xmat.shape[0]
+        p = xmat.shape[1]
+        if prior is not None:
+            if np.isscalar(prior):
+                prior = np.diag(np.full(p, float(prior)))
+            xmat = np.vstack([xmat, np.linalg.cholesky(prior).T])
+        if np.linalg.matrix_rank(xmat) < p:
+            warnings.warn(
+                "Covariate matrix is singular; adding a small ridge prior "
+                "(1e-5) for numerical stability.", stacklevel=3,
+            )
+            xmat = np.vstack([xmat, np.sqrt(1e-5) * np.eye(p)])
+        self.xmat = xmat
+        self.q, self.r = np.linalg.qr(xmat)
+        self.rinv = np.linalg.inv(self.r)
+        self.df_residual = xmat.shape[0] - p
+
+    def fit(self, y):
+        if y.shape[0] != self.xmat.shape[0]:
+            y = np.concatenate(
+                [y, np.zeros(self.xmat.shape[0] - y.shape[0])]
+            )
+        coef = self.rinv @ (self.q.T @ y)
+        resid = y - self.xmat @ coef
+        resvar = (resid @ resid) / self.df_residual
+        vcov = resvar * (self.rinv @ self.rinv.T)
+        return coef, vcov
+
+
+class EstimatedEffects:
+    """Result of :func:`estimate_effect`.
+
+    Attributes
+    ----------
+    parameters : dict
+        Maps topic index to a list of ``(coef, vcov)`` pairs, one per
+        composition draw.
+    topics : list of int
+        Topics for which effects were estimated (0-based).
+    n_obs : int
+        Number of documents.
+    n_params : int
+        Number of regression coefficients (including the intercept).
+    """
+
+    def __init__(self, parameters, topics, n_obs, n_params):
+        self.parameters = parameters
+        self.topics = topics
+        self.n_obs = n_obs
+        self.n_params = n_params
+
+    def summary(self, topics=None, nsim=500, random_state=None):
+        """Pooled coefficient tables (summary.estimateEffect in R).
+
+        Returns a dict mapping topic index to a record array with fields
+        ``estimate``, ``std_error``, ``t_value`` and ``p_value``, one row
+        per regression coefficient.
+        """
+        rng = check_random_state(random_state)
+        topics = self.topics if topics is None else list(topics)
+        dtype = [("estimate", float), ("std_error", float),
+                 ("t_value", float), ("p_value", float)]
+        tables = {}
+        for k in topics:
+            if k not in self.parameters:
+                raise ValueError(f"Topic {k} was not estimated.")
+            sims = np.vstack([
+                rng.multivariate_normal(est, vcov, size=nsim)
+                for est, vcov in self.parameters[k]
+            ])
+            est = sims.mean(axis=0)
+            se = sims.std(axis=0, ddof=1)
+            tval = est / se
+            rdf = self.n_obs - self.n_params
+            p = 2 * stats.t.sf(np.abs(tval), rdf)
+            table = np.zeros(len(est), dtype=dtype)
+            table["estimate"] = est
+            table["std_error"] = se
+            table["t_value"] = tval
+            table["p_value"] = p
+            tables[k] = table
+        return tables
+
+
+def estimate_effect(model, prevalence, topics=None,
+                    uncertainty="Global", nsims=25, prior=None,
+                    random_state=None):
+    """Regress topic proportions on covariates (estimateEffect in R).
+
+    Parameters
+    ----------
+    model : fitted StructuralTopicModel
+    prevalence : array-like of shape (n_samples, n_covariates)
+        Covariate design matrix for the regression; an intercept column
+        is added automatically.  Should normally contain (at least) the
+        covariates used when fitting the model.  Categorical variables
+        must be encoded numerically beforehand.
+    topics : iterable of int, optional
+        0-based topic indices to estimate effects for (default: all).
+    uncertainty : {"Global", "None"}, default="Global"
+        "Global" draws theta from the variational posterior using a
+        globally shared covariance approximation; "None" uses the MAP
+        theta without measurement uncertainty.  (The R package's "Local"
+        method is not implemented.)
+    nsims : int, default=25
+        Number of method-of-composition draws ("Global" only).
+    prior : float or ndarray, optional
+        Ridge penalty (scalar or full precision matrix) added to the
+        regression for numerical stability.
+    """
+    if not hasattr(model, "theta_"):
+        raise ValueError("model must be a fitted StructuralTopicModel.")
+    if uncertainty not in ("Global", "None"):
+        raise ValueError(
+            "uncertainty must be 'Global' or 'None' ('Local' is not "
+            "implemented; 'Global' is the recommended method)."
+        )
+    rng = check_random_state(random_state)
+    K = model.theta_.shape[1]
+    topics = list(range(K)) if topics is None else list(topics)
+    if any(k < 0 or k >= K for k in topics):
+        raise ValueError("topics must be 0-based indices below n_components.")
+
+    xmat = np.asarray(prevalence, dtype=np.float64)
+    if xmat.ndim == 1:
+        xmat = xmat[:, None]
+    if xmat.shape[0] != model.theta_.shape[0]:
+        raise ValueError(
+            "prevalence has a different number of rows than the fitted "
+            "documents."
+        )
+    if not np.allclose(xmat[:, 0], 1.0):
+        xmat = np.column_stack([np.ones(xmat.shape[0]), xmat])
+
+    reg = _QRRegression(xmat, prior=prior)
+    if uncertainty == "None":
+        nsims = 1
+
+    parameters = {k: [] for k in topics}
+    for _ in range(nsims):
+        if uncertainty == "None":
+            theta = model.theta_
+        else:
+            theta = _draw_theta(model, rng)
+        for k in topics:
+            parameters[k].append(reg.fit(theta[:, k]))
+
+    return EstimatedEffects(parameters, topics,
+                            n_obs=model.theta_.shape[0],
+                            n_params=xmat.shape[1])

pystm/model_selection.py

@@ -0,0 +1,166 @@
+"""Choosing the number of topics (ports of searchK.R and heldout.R)."""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.sparse import csr_matrix, issparse, lil_matrix
+from scipy.special import gammaln
+from sklearn.utils import check_random_state
+
+from ._utils import safelog
+from .diagnostics import check_residuals, exclusivity, semantic_coherence
+from .stm import StructuralTopicModel
+
+
+def make_heldout(X, N=None, proportion=0.5, random_state=None):
+    """Hold out a fraction of tokens for document completion (make.heldout).
+
+    Selects ``N`` documents (default 10%) and removes ``proportion`` of
+    each one's tokens.  Held-out tokens whose word disappears entirely
+    from the training corpus are dropped (mirroring the R package's vocab
+    remapping).
+
+    Returns a dict with ``X_train`` (csr matrix), ``index`` (held-out
+    document ids) and ``docs`` (list of ``(word_indices, counts)``).
+    """
+    X = csr_matrix(X) if not issparse(X) else X.tocsr()
+    rng = check_random_state(random_state)
+    n_docs, V = X.shape
+    if N is None:
+        N = int(np.floor(0.1 * n_docs))
+    if not 0 < N <= n_docs:
+        raise ValueError("N must be between 1 and the number of documents.")
+    if not 0 < proportion < 1:
+        raise ValueError("proportion must be in (0, 1).")
+
+    index = np.sort(rng.choice(n_docs, size=N, replace=False))
+    X_train = lil_matrix(X.copy())
+    missing_docs = []
+    kept_index = []
+    for i in index:
+        row = X.getrow(i)
+        if row.indices.shape[0] < 2:
+            continue  # too few distinct words to split (as in R)
+        tokens = np.repeat(row.indices, row.data.astype(np.int64))
+        nsamp = max(1, int(np.floor(proportion * tokens.shape[0])))
+        nsamp = min(nsamp, tokens.shape[0] - 1)  # keep the doc non-empty
+        held = rng.choice(tokens.shape[0], size=nsamp, replace=False)
+        held_counts = np.bincount(tokens[held], minlength=V)
+        words = np.flatnonzero(held_counts)
+        for w in words:
+            X_train[i, w] -= held_counts[w]
+        missing_docs.append((words, held_counts[words].astype(np.float64)))
+        kept_index.append(i)
+    X_train = csr_matrix(X_train)
+    X_train.eliminate_zeros()
+
+    # drop held-out tokens of words that vanished from the training corpus
+    train_wcounts = np.asarray(X_train.sum(axis=0)).ravel()
+    final_docs, final_index = [], []
+    for i, (words, counts) in zip(kept_index, missing_docs):
+        keep = train_wcounts[words] > 0
+        if keep.any():
+            final_docs.append((words[keep], counts[keep]))
+            final_index.append(i)
+    return {
+        "X_train": X_train,
+        "index": np.asarray(final_index, dtype=np.int64),
+        "docs": final_docs,
+    }
+
+
+def eval_heldout(model, heldout, content=None):
+    """Heldout log-likelihood by document completion (eval.heldout).
+
+    Uses the fitted theta of each partially observed document to score
+    its held-out tokens.  Returns a dict with ``expected_heldout`` (the
+    mean over documents of the mean per-token log-probability) and the
+    per-document values ``doc_heldout``.
+    """
+    beta = model._beta_list()
+    if model.content_levels_ is not None:
+        if content is None:
+            raise ValueError(
+                "The model was fitted with a content covariate; pass the "
+                "matching content labels."
+            )
+        index_all = np.searchsorted(model.content_levels_,
+                                    np.asarray(content).ravel())
+    else:
+        index_all = np.zeros(model.theta_.shape[0], dtype=np.int64)
+
+    doc_scores = np.empty(len(heldout["docs"]))
+    for j, (doc_id, (words, counts)) in enumerate(
+            zip(heldout["index"], heldout["docs"])):
+        logprobs = safelog(
+            model.theta_[doc_id] @ beta[index_all[doc_id]][:, words]
+        )
+        doc_scores[j] = np.repeat(logprobs, counts.astype(np.int64)).mean()
+    return {
+        "expected_heldout": float(np.nanmean(doc_scores)),
+        "doc_heldout": doc_scores,
+    }
+
+
+def search_k(X, K_values, *, prevalence=None, content=None, N=None,
+             proportion=0.5, M=10, heldout_random_state=None,
+             model_params=None, verbose=False):
+    """Fit models over a grid of K and compute diagnostics (searchK).
+
+    Parameters
+    ----------
+    X : array-like or sparse matrix of shape (n_samples, n_features)
+    K_values : iterable of int
+        Topic numbers to evaluate.
+    prevalence, content : optional
+        Covariates forwarded to :meth:`StructuralTopicModel.fit`.
+    N, proportion : heldout construction parameters (see make_heldout).
+    M : int, default=10
+        Number of top words for exclusivity / semantic coherence.
+    model_params : dict, optional
+        Extra keyword arguments for the StructuralTopicModel constructor.
+
+    Returns
+    -------
+    dict of arrays keyed by "K", "heldout", "residual", "bound",
+    "lbound", "exclus", "semcoh", "em_its" (exclusivity and semantic
+    coherence are omitted for content covariate models, as in R).
+    """
+    K_values = list(K_values)
+    model_params = dict(model_params or {})
+    heldout = make_heldout(X, N=N, proportion=proportion,
+                           random_state=heldout_random_state)
+
+    results = {key: [] for key in
+               ("K", "heldout", "residual", "bound", "lbound",
+                "exclus", "semcoh", "em_its")}
+    for K in K_values:
+        if verbose:
+            print(f"searchK: fitting K={K} ...")
+        model = StructuralTopicModel(n_components=K, **model_params)
+        model.fit(heldout["X_train"], prevalence=prevalence, content=content)
+
+        results["K"].append(K)
+        results["heldout"].append(
+            eval_heldout(model, heldout, content=content)["expected_heldout"]
+        )
+        results["residual"].append(
+            check_residuals(model, heldout["X_train"],
+                            content=content)["dispersion"]
+        )
+        bound = max(model.bound_)
+        results["bound"].append(bound)
+        results["lbound"].append(bound + gammaln(K + 1))
+        if content is None:
+            results["exclus"].append(
+                float(np.mean(exclusivity(model, M=M, frexw=0.7)))
+            )
+            results["semcoh"].append(
+                float(np.mean(semantic_coherence(model, heldout["X_train"],
+                                                 M=M)))
+            )
+        results["em_its"].append(model.n_iter_)
+
+    if content is not None:
+        del results["exclus"], results["semcoh"]
+    return {key: np.asarray(val) for key, val in results.items()}