pydecontx 0.1.0__py3-none-any.whl

pydecontx/__init__.py

@@ -0,0 +1,45 @@
+"""pydecontx -- pure-Python port of the Bioconductor ``decontX`` package.
+
+DecontX (Yang et al., *Genome Biology* 2020) identifies and removes
+ambient-RNA / cross-contamination from droplet single-cell RNA-seq.
+Each cell's observed UMI counts are modelled as a Bayesian
+two-component multinomial mixture of a *native* gene distribution
+``phi`` and a *contamination* distribution ``eta`` (a weighted blend of
+every other population). Inference is by variational EM and yields a
+per-cell contamination fraction and a decontaminated count matrix.
+
+Main entry points
+-----------------
+- :func:`decontx`              -- run DecontX on a count matrix / AnnData.
+- :class:`DecontXResult`       -- the result object.
+- :func:`simulate_contamination` -- simulate a contaminated dataset.
+
+Lower-level routines (variational-EM building blocks) are available in
+:mod:`pydecontx._core` and the Dirichlet MLE in
+:mod:`pydecontx._dirichlet`.
+"""
+from __future__ import annotations
+
+from ._core import (
+    calculate_native_matrix,
+    decontx_em,
+    decontx_initialize,
+    decontx_loglik,
+)
+from ._dirichlet import fit_dirichlet
+from .decontx import DecontXResult, decontx
+from .simulate import simulate_contamination
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "decontx",
+    "DecontXResult",
+    "simulate_contamination",
+    "decontx_initialize",
+    "decontx_em",
+    "decontx_loglik",
+    "calculate_native_matrix",
+    "fit_dirichlet",
+    "__version__",
+]

pydecontx/_core.py

@@ -0,0 +1,223 @@
+"""Core variational-EM routines for DecontX.
+
+Pure-Python / numpy / scipy port of the C++ inner loops of the
+Bioconductor ``decontX`` package (``src/DecontX.cpp``):
+
+* :func:`decontx_initialize` -- initial native (``phi``) and contamination
+  (``eta``) gene distributions from a random ``theta``.
+* :func:`decontx_em` -- one variational-EM step updating ``phi``, ``eta``,
+  ``theta`` and (optionally) the Dirichlet hyper-parameter ``delta``.
+* :func:`decontx_loglik` -- the two-component multinomial log-likelihood.
+* :func:`calculate_native_matrix` -- the decontaminated count matrix.
+
+All matrices follow the R convention: ``counts`` is genes-by-cells
+(rows = genes, columns = cells); ``phi`` / ``eta`` are genes-by-clusters.
+``z`` holds 1-based integer cluster labels (one per cell).
+"""
+from __future__ import annotations
+
+import numpy as np
+import scipy.sparse as sp
+
+from ._dirichlet import fit_dirichlet
+
+__all__ = [
+    "decontx_initialize",
+    "decontx_em",
+    "decontx_loglik",
+    "calculate_native_matrix",
+]
+
+
+def _as_csc(counts) -> sp.csc_matrix:
+    """Return ``counts`` as a float CSC sparse matrix (genes x cells)."""
+    if sp.issparse(counts):
+        return counts.tocsc().astype(float)
+    return sp.csc_matrix(np.asarray(counts, dtype=float))
+
+
+def decontx_initialize(counts, theta, z, pseudocount: float = 1e-20):
+    """Initialise the native (``phi``) and contamination (``eta``) matrices.
+
+    Port of the C++ ``decontXInitialize``. ``phi[:, k]`` accumulates
+    ``theta_j * counts`` over all cells ``j`` assigned to cluster ``k``;
+    ``eta`` is the row-sum complement (every *other* cluster's signal),
+    and both are column-normalised to proportions.
+
+    Parameters
+    ----------
+    counts : (G, C) array or sparse matrix
+        Gene-by-cell UMI counts.
+    theta : (C,) array
+        Initial native proportion for each cell.
+    z : (C,) int array
+        1-based cluster label for each cell.
+    pseudocount : float
+        Added to every cell of ``phi``/``eta`` before normalising.
+
+    Returns
+    -------
+    dict with keys ``phi`` and ``eta`` -- (G, K) numpy arrays.
+    """
+    counts = _as_csc(counts)
+    theta = np.asarray(theta, dtype=float)
+    z = np.asarray(z, dtype=int)
+    G, C = counts.shape
+    K = int(z.max())
+
+    phi = np.full((G, K), pseudocount, dtype=float)
+    indptr, indices, data = counts.indptr, counts.indices, counts.data
+    for j in range(C):
+        k = z[j] - 1
+        start, end = indptr[j], indptr[j + 1]
+        rows = indices[start:end]
+        vals = data[start:end] * theta[j]
+        np.add.at(phi[:, k], rows, vals)
+
+    phi_rowsum = phi.sum(axis=1)
+    eta = phi_rowsum[:, None] - phi
+
+    phi = phi / phi.sum(axis=0, keepdims=True)
+    eta = eta / eta.sum(axis=0, keepdims=True)
+    return {"phi": phi, "eta": eta}
+
+
+def decontx_em(counts, counts_colsums, theta, eta, phi, z,
+               estimate_eta: bool = True, estimate_delta: bool = True,
+               delta=(10.0, 10.0), pseudocount: float = 1e-20):
+    """One variational-EM update of the DecontX model.
+
+    Port of the C++ ``decontXEM``. For every observed transcript the
+    variational native/contaminant responsibility is
+
+    ``p_native  = (phi[i,k] + pc) * (theta[j] + pc)``
+    ``p_contam  = (eta[i,k] + pc) * (1 - theta[j] + pc)``
+    ``normp     = p_native / (p_native + p_contam)``
+
+    (the non-log form -- exact for a two-component mixture and what the
+    C++ code uses). The native mass ``normp * x`` is accumulated into the
+    new ``phi`` by cluster; ``eta`` is the row-sum complement. ``theta``
+    is then the posterior mean of a Beta/Dirichlet with concentration
+    ``delta``, which is itself re-estimated by :func:`fit_dirichlet`.
+
+    Returns a dict with the updated ``phi``, ``eta``, ``theta``,
+    ``delta`` and the per-cell ``contamination`` fraction.
+    """
+    counts = _as_csc(counts)
+    theta = np.asarray(theta, dtype=float)
+    counts_colsums = np.asarray(counts_colsums, dtype=float)
+    phi = np.asarray(phi, dtype=float)
+    eta = np.asarray(eta, dtype=float)
+    z = np.asarray(z, dtype=int)
+    delta = np.asarray(delta, dtype=float)
+
+    G, C = counts.shape
+    K = phi.shape[1]
+
+    new_phi = np.zeros((G, K), dtype=float)
+    native_total = np.zeros(C, dtype=float)
+
+    indptr, indices, data = counts.indptr, counts.indices, counts.data
+    for j in range(C):
+        k = z[j] - 1
+        start, end = indptr[j], indptr[j + 1]
+        if start == end:
+            continue
+        rows = indices[start:end]
+        x = data[start:end]
+        p_native = (phi[rows, k] + pseudocount) * (theta[j] + pseudocount)
+        p_contam = (eta[rows, k] + pseudocount) * (1.0 - theta[j] + pseudocount)
+        normp = p_native / (p_native + p_contam)
+        px = normp * x
+        np.add.at(new_phi[:, k], rows, px)
+        native_total[j] = px.sum()
+
+    if estimate_eta:
+        phi_rowsum = new_phi.sum(axis=1)
+        new_eta = phi_rowsum[:, None] - new_phi
+    else:
+        new_eta = eta
+
+    new_phi = new_phi / new_phi.sum(axis=0, keepdims=True)
+    if estimate_eta:
+        new_eta = new_eta / new_eta.sum(axis=0, keepdims=True)
+
+    # Update theta (and optionally its Dirichlet hyper-parameter delta).
+    contamination_prop = (counts_colsums - native_total) / counts_colsums
+    native_prop = 1.0 - contamination_prop
+    new_delta = delta
+    if estimate_delta:
+        theta_raw = np.column_stack([native_prop, contamination_prop])
+        new_delta = fit_dirichlet(theta_raw)["alpha"]
+
+    new_theta = (native_total + new_delta[0]) / (counts_colsums
+                                                 + np.sum(new_delta))
+
+    return {
+        "phi": new_phi,
+        "eta": new_eta,
+        "theta": new_theta,
+        "delta": new_delta,
+        "contamination": contamination_prop,
+    }
+
+
+def decontx_loglik(counts, theta, eta, phi, z, pseudocount: float = 1e-20):
+    """Two-component multinomial log-likelihood of the DecontX model.
+
+    Port of the C++ ``decontXLogLik``:
+    ``ll = sum_{i,j} x_{ij} * log(phi*theta + eta*(1-theta) + pc)``.
+    """
+    counts = _as_csc(counts)
+    theta = np.asarray(theta, dtype=float)
+    phi = np.asarray(phi, dtype=float)
+    eta = np.asarray(eta, dtype=float)
+    z = np.asarray(z, dtype=int)
+
+    loglik = 0.0
+    indptr, indices, data = counts.indptr, counts.indices, counts.data
+    for j in range(counts.shape[1]):
+        k = z[j] - 1
+        start, end = indptr[j], indptr[j + 1]
+        if start == end:
+            continue
+        rows = indices[start:end]
+        x = data[start:end]
+        mix = (phi[rows, k] * theta[j]
+               + eta[rows, k] * (1.0 - theta[j]) + pseudocount)
+        loglik += float(np.sum(x * np.log(mix)))
+    return loglik
+
+
+def calculate_native_matrix(counts, theta, eta, phi, z,
+                            pseudocount: float = 1e-20) -> sp.csc_matrix:
+    """Return the decontaminated (native) count matrix.
+
+    Port of the C++ ``calculateNativeMatrix``: each observed entry is
+    scaled by its variational native responsibility ``normp``. Values
+    may be non-integer; round for integer counts.
+    """
+    counts = _as_csc(counts).copy()
+    theta = np.asarray(theta, dtype=float)
+    phi = np.asarray(phi, dtype=float)
+    eta = np.asarray(eta, dtype=float)
+    z = np.asarray(z, dtype=int)
+
+    indptr, indices, data = counts.indptr, counts.indices, counts.data
+    out = data.copy()
+    for j in range(counts.shape[1]):
+        k = z[j] - 1
+        start, end = indptr[j], indptr[j + 1]
+        if start == end:
+            continue
+        rows = indices[start:end]
+        p_native = np.log(phi[rows, k] + pseudocount) + np.log(theta[j]
+                                                               + pseudocount)
+        p_contam = np.log(eta[rows, k] + pseudocount) + np.log(
+            1.0 - theta[j] + pseudocount)
+        normp = np.exp(p_native) / (np.exp(p_contam) + np.exp(p_native))
+        out[start:end] = data[start:end] * normp
+
+    native = sp.csc_matrix((out, indices.copy(), indptr.copy()),
+                           shape=counts.shape)
+    return native

pydecontx/_dirichlet.py

@@ -0,0 +1,96 @@
+"""Dirichlet maximum-likelihood estimation.
+
+A pure-Python port of ``MCMCprecision::fit_dirichlet`` -- the fixed-point
+iteration of Minka (2000) used by R ``decontX`` to update the per-cell
+``delta`` (theta) concentration parameters at every EM iteration.
+"""
+from __future__ import annotations
+
+import numpy as np
+from scipy.special import digamma, polygamma
+
+# Euler-Mascheroni constant ( -psi(1) ).
+_EULER_GAMMA = -digamma(1.0)
+
+
+def inv_digamma(y: np.ndarray, iter: int = 5) -> np.ndarray:
+    """Inverse of the digamma function (Minka, 2000).
+
+    Mirrors ``MCMCprecision::inv_digamma``: starts from ``exp(y) + 0.5``
+    (or ``-1/(y + gamma)`` when ``y < -2.22``) and applies a few Newton
+    steps using the trigamma function.
+    """
+    y = np.asarray(y, dtype=float)
+    x = np.exp(y) + 0.5
+    mask = y < -2.22
+    x[mask] = -1.0 / (y[mask] + _EULER_GAMMA)
+    for _ in range(iter):
+        x = x - (digamma(x) - y) / polygamma(1, x)
+    return x
+
+
+def dirichlet_fp(alpha: np.ndarray, logx_mean: np.ndarray,
+                 maxit: int = 100000, abstol: float = 1e-5) -> np.ndarray:
+    """Fixed-point iteration for Dirichlet MLE (Minka, 2000).
+
+    Port of ``MCMCprecision::dirichlet_fp``.
+    """
+    alpha = np.asarray(alpha, dtype=float).copy()
+    logx_mean = np.asarray(logx_mean, dtype=float)
+    cnt = 0
+    diff = 1.0
+    while (diff > abstol) and (cnt < maxit):
+        alpha0 = alpha
+        alpha = inv_digamma(digamma(np.sum(alpha0)) + logx_mean)
+        cnt += 1
+        diff = np.max(np.abs(alpha0 - alpha))
+    return alpha
+
+
+def fit_dirichlet(x: np.ndarray, const: float | None = None,
+                  maxit: int = 100000, abstol: float = 0.1) -> dict:
+    """Estimate the parameters of a Dirichlet distribution.
+
+    Pure-Python port of ``MCMCprecision::fit_dirichlet`` (Minka 2000
+    fixed-point algorithm). ``x`` is a 2-D array of Dirichlet samples,
+    one observation per row.
+
+    Returns a dict with keys ``alpha`` (the estimated concentration
+    vector) and ``sum`` (its sum). The default ``abstol`` of 0.1 matches
+    the R default used inside ``decontX``.
+    """
+    x = np.asarray(x, dtype=float)
+    if x.ndim != 2:
+        raise ValueError("'x' must be a 2-D array of Dirichlet samples.")
+
+    if x.min() == 0:
+        if const is None:
+            const = x[x > 0].min() * 0.01
+        x = (x + const) / (1.0 + 2.0 * const)
+
+    x = x / x.sum(axis=1, keepdims=True)
+    logx_mean = np.mean(np.log(x), axis=0)
+
+    # Heuristic starting values (method of moments).
+    x_mean = np.mean(x, axis=0)
+    x_squares = np.mean(x ** 2, axis=0)
+    denom = x_squares - x_mean ** 2
+    with np.errstate(divide="ignore", invalid="ignore"):
+        xi = (x_mean - x_squares) / denom
+    alpha0 = xi * x_mean
+
+    alpha = None
+    try:
+        cand = dirichlet_fp(np.maximum(0.01, alpha0), logx_mean,
+                            maxit=maxit, abstol=abstol)
+        if not np.any(np.isnan(cand)):
+            alpha = cand
+    except Exception:  # pragma: no cover - fall through to random init
+        alpha = None
+
+    if alpha is None or np.any(np.isnan(alpha)):
+        rng = np.random.default_rng()
+        alpha = dirichlet_fp(rng.uniform(0.5, 1.0, size=len(alpha0)),
+                             logx_mean, maxit=maxit, abstol=abstol)
+
+    return {"alpha": alpha, "sum": float(np.sum(alpha))}