pycopykat 0.2.0__py3-none-any.whl

pycopykat/__init__.py

@@ -0,0 +1,8 @@
+"""pycopykat: a Python implementation of CopyKAT (Gao et al. 2021) that reproduces R copykat."""
+from pycopykat.config import CopykatConfig
+from pycopykat.io.h5ad import copykat_by_batch
+from pycopykat.pipeline import copykat
+from pycopykat.result import CopykatResult
+
+__all__ = ["CopykatConfig", "CopykatResult", "copykat", "copykat_by_batch"]
+__version__ = "0.2.0"

pycopykat/_rcompat/__init__.py

@@ -0,0 +1 @@
+"""Pieces of R reproduced with R's arithmetic and random stream."""

pycopykat/_rcompat/hclust.py

@@ -0,0 +1,158 @@
+"""R's ``hclust`` (``ward.D`` and ``ward.D2``) and ``cutree``.
+
+``hclust`` follows R's Fortran routine (F. Murtagh's nearest-neighbour list, ``hclust.f``): at
+every step the pair with the smallest nearest-neighbour dissimilarity is merged (the first one on
+ties), the dissimilarities to the new cluster are updated with the Lance-Williams formula in the
+same order of operations, and the nearest neighbours are refreshed as R does, including the fix
+for neighbours to the left of the new cluster. ``ward.D2`` squares the dissimilarities first and
+reports square roots as heights. ``cutree(k)`` undoes the last ``k - 1`` merges and numbers the
+groups in the order in which the observations first appear, as R does.
+
+The update of each step is vectorised over the remaining clusters.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+_INF = 1e300
+
+
+@dataclass
+class HClust:
+    """An R ``hclust`` result: ``merge`` (n-1 x 2, R's signed convention), ``height`` and
+    ``order`` (1-based, as in R), plus the raw merge pairs (0-based)."""
+
+    merge: np.ndarray
+    height: np.ndarray
+    order: np.ndarray
+    pairs: np.ndarray
+    labels: np.ndarray | None = None
+
+    @property
+    def n(self) -> int:
+        return self.pairs.shape[0] + 1
+
+    def cutree(self, k: int) -> np.ndarray:
+        return cutree(self.pairs, k)
+
+
+def hclust(d: np.ndarray, method: str = "ward.D", labels=None) -> HClust:
+    """``hclust(as.dist(d), method)`` for a square symmetric dissimilarity matrix ``d``."""
+    if method not in ("ward.D", "ward.D2"):
+        raise ValueError("method must be 'ward.D' or 'ward.D2'")
+    D = np.array(d, dtype=np.float64, copy=True)
+    n = D.shape[0]
+    if n < 2:
+        raise ValueError("must have n >= 2 objects to cluster")
+    if method == "ward.D2":
+        D = D * D
+    np.fill_diagonal(D, _INF)
+    idx = np.arange(n)
+    membr = np.ones(n)
+    active = np.ones(n, dtype=bool)
+    # nearest neighbour to the right of each i (first minimum)
+    upper = np.where(idx[None, :] > idx[:, None], D, np.inf)
+    nn = np.argmin(upper, axis=1)
+    disnn = upper[idx, nn]
+    del upper
+    nn[n - 1] = -1
+    disnn[n - 1] = np.inf
+    pairs = np.empty((n - 1, 2), dtype=np.int64)
+    crit = np.empty(n - 1)
+    for step in range(n - 1):
+        cand = np.where(active[: n - 1], disnn[: n - 1], np.inf)
+        im = int(np.argmin(cand))
+        dmin = cand[im]
+        jm = int(nn[im])
+        i2, j2 = min(im, jm), max(im, jm)
+        pairs[step] = (i2, j2)
+        crit[step] = np.sqrt(dmin) if method == "ward.D2" else dmin
+        active[j2] = False
+        ks = np.flatnonzero(active)
+        ks = ks[ks != i2]
+        if ks.size:
+            d12 = D[i2, j2]
+            mk = membr[ks]
+            new = ((membr[i2] + mk) * D[i2, ks] + (membr[j2] + mk) * D[j2, ks]) - mk * d12
+            new = new / ((membr[i2] + membr[j2]) + mk)
+            D[i2, ks] = new
+            D[ks, i2] = new
+            right = ks > i2
+            if right.any():
+                r = int(np.argmin(new[right]))
+                dnew, jj = new[right][r], int(ks[right][r])
+            else:
+                dnew, jj = np.inf, -1
+            left = ~right
+            if left.any():
+                kl, vl = ks[left], new[left]
+                closer = vl < disnn[kl]
+                disnn[kl[closer]] = vl[closer]
+                nn[kl[closer]] = i2
+        else:
+            dnew, jj = np.inf, -1
+        membr[i2] += membr[j2]
+        disnn[i2] = dnew
+        nn[i2] = jj
+        rows = np.flatnonzero(active[: n - 1] & ((nn[: n - 1] == i2) | (nn[: n - 1] == j2)))
+        if rows.size:
+            sub = D[rows]
+            ok = active[None, :] & (idx[None, :] > rows[:, None])
+            sub = np.where(ok, sub, np.inf)
+            jmin = np.argmin(sub, axis=1)
+            dm = sub[np.arange(rows.size), jmin]
+            none = ~np.isfinite(dm)
+            nn[rows] = np.where(none, -1, jmin)
+            disnn[rows] = np.where(none, np.inf, dm)
+    merge, order = _hcass2(pairs + 1, n)
+    return HClust(merge=merge, height=crit, order=order, pairs=pairs,
+                  labels=None if labels is None else np.asarray(labels))
+
+
+def _hcass2(pairs1: np.ndarray, n: int) -> tuple[np.ndarray, np.ndarray]:
+    """R's ``hcass2``: the signed merge matrix and the leaf order (both 1-based)."""
+    ia = pairs1[:, 0].tolist()
+    ib = pairs1[:, 1].tolist()
+    iia, iib = list(ia), list(ib)
+    last = {}                      # cluster representative -> step that last merged it
+    for i in range(n - 1):
+        a, b = ia[i], ib[i]
+        iia[i] = last.get(a, -a)
+        iib[i] = last.get(b, -b)
+        last[min(a, b)] = i + 1
+    for i in range(n - 1):
+        if iia[i] > 0 and iib[i] < 0:
+            iia[i], iib[i] = iib[i], iia[i]
+        if iia[i] > 0 and iib[i] > 0:
+            iia[i], iib[i] = min(iia[i], iib[i]), max(iia[i], iib[i])
+    order = [iia[n - 2], iib[n - 2]]
+    for i in range(n - 2, 0, -1):
+        j = order.index(i)
+        order[j:j + 1] = [iia[i - 1], iib[i - 1]]
+    return np.column_stack([iia, iib]).astype(np.int64), -np.asarray(order, dtype=np.int64)
+
+
+def cutree(pairs: np.ndarray, k: int) -> np.ndarray:
+    """``cutree(tree, k)`` from the 0-based merge pairs: 1-based group numbers in the order of
+    first appearance."""
+    n = pairs.shape[0] + 1
+    if not 1 <= k <= n:
+        raise ValueError(f"k must be between 1 and {n}")
+    parent = np.arange(n)
+
+    def find(x):
+        while parent[x] != x:
+            parent[x] = parent[parent[x]]
+            x = parent[x]
+        return x
+
+    for a, b in pairs[: n - k]:
+        ra, rb = find(int(a)), find(int(b))
+        parent[rb] = ra
+    roots = np.array([find(i) for i in range(n)])
+    _, first, inv = np.unique(roots, return_index=True, return_inverse=True)
+    rank = np.empty(first.size, dtype=np.int64)
+    rank[np.argsort(first)] = np.arange(1, first.size + 1)
+    return rank[inv]

pycopykat/_rcompat/mixtools.py

@@ -0,0 +1,95 @@
+"""mixtools ``normalmixEM`` as CopyKAT calls it: three components with a common standard deviation
+(``arbvar = FALSE``), starting from given weights, means and sigma, run on several vectors at once.
+
+Each EM iteration follows mixtools: ``lambda = colMeans(post)``; each mean is
+``sum(post[, j] * x) / (n * lambda[j])``; the common sigma is
+``sqrt(sum(post * res) / (n * sum(lambda)))``; the posteriors, squared residuals and
+log-likelihood come from mixtools' ``normpost`` C routine (scaled by the component with the
+smallest standardised residual, the log-likelihood added in order). The loop stops when the
+log-likelihood gain is at most ``epsilon`` or after ``maxit`` iterations.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+_LD = np.longdouble
+_HALF_LOG_2PI = 0.91893853320467274178
+
+
+@dataclass
+class MixFit:
+    lam: np.ndarray       # (B, 3)
+    mu: np.ndarray        # (B, 3)
+    sigma: np.ndarray     # (B,)
+    loglik: np.ndarray    # (B,)
+    iterations: np.ndarray
+
+
+def _normpost(x, mu, sigma, lam):
+    """mixtools ``normpost`` for B vectors at once: posteriors and squared residuals (B, 3, n) and
+    the log-likelihood (B,)."""
+    B, n = x.shape
+    lsr = lam / sigma[:, None]                                   # LamSigRatio (B, 3)
+    llsr = np.log(lsr)
+    res2 = x[:, None, :] - mu[:, :, None]
+    res2 *= res2
+    den = (2.0 * sigma) * sigma
+    work = res2 / den[:, None, None]
+    w0, w1, w2 = work[:, 0], work[:, 1], work[:, 2]
+    # first component with the smallest standardised residual (j == 0 || r < min)
+    lt1 = w1 < w0
+    mn = np.where(lt1, w1, w0)
+    lt2 = w2 < mn
+    mn = np.where(lt2, w2, mn)
+    minj = np.where(lt2, 2, np.where(lt1, 1, 0))
+    rb = np.arange(B)[:, None]
+    lsr_min = lsr[rb, minj]
+    post = np.empty_like(work)
+    rowsum = np.ones((B, n))
+    for j in range(3):
+        wj = (lsr[:, j][:, None] / lsr_min) * np.exp(mn - work[:, j])
+        is_min = minj == j
+        wj[is_min] = 1.0
+        post[:, j] = wj
+        rowsum += np.where(is_min, 0.0, wj)
+    post /= rowsum[:, None, :]
+    terms = (np.log(rowsum) - mn) + llsr[rb, minj]
+    terms[:, 0] += -float(n) * _HALF_LOG_2PI
+    ll = np.cumsum(terms, axis=1)[:, -1]
+    return post, res2, ll
+
+
+def normalmix3(x: np.ndarray, sigma0: np.ndarray, mu0=(-0.2, 0.0, 0.2), lam0=None,
+               epsilon: float = 1e-8, maxit: int = 1000) -> MixFit:
+    """``normalmixEM(x[b], lambda = rep(1, 3)/3, mu = mu0, sigma = sigma0[b], arbvar = FALSE,
+    ECM = FALSE, maxit = maxit)`` for each row ``x[b]``."""
+    x = np.atleast_2d(np.asarray(x, dtype=np.float64))
+    B, n = x.shape
+    lam_init = np.full(3, 1.0) / 3.0 if lam0 is None else np.asarray(lam0, dtype=np.float64)
+    lam_init = lam_init / lam_init.sum()
+    lam = np.tile(lam_init, (B, 1))
+    mu = np.tile(np.asarray(mu0, dtype=np.float64), (B, 1))
+    sigma = np.asarray(sigma0, dtype=np.float64).reshape(B).copy()
+    post, res, ll = _normpost(x, mu, sigma, lam)
+    iters = np.zeros(B, dtype=np.int64)
+    act = np.ones(B, dtype=bool)
+    while act.any():
+        a = np.flatnonzero(act)
+        xa, pa, ra = x[a], post[a], res[a]
+        lam_a = (pa.sum(axis=2, dtype=_LD) / n).astype(np.float64)
+        mu_a = (pa * xa[:, None, :]).sum(axis=2, dtype=_LD).astype(np.float64) / (float(n) * lam_a)
+        num = (pa * ra).reshape(a.size, 3 * n).sum(axis=1, dtype=_LD)
+        sl = lam_a.astype(_LD)
+        sl = (sl[:, 0] + sl[:, 1]) + sl[:, 2]
+        sig_a = np.sqrt(num.astype(np.float64) / (float(n) * sl.astype(np.float64)))
+        if np.any(sig_a < 1e-8):
+            raise FloatingPointError("normalmixEM: a variance went to zero (R restarts at random)")
+        p2, r2, ll2 = _normpost(xa, mu_a, sig_a, lam_a)
+        diff = ll2 - ll[a]
+        lam[a], mu[a], sigma[a] = lam_a, mu_a, sig_a
+        post[a], res[a], ll[a] = p2, r2, ll2
+        iters[a] += 1
+        act[a] = (diff > epsilon) & (iters[a] < maxit)
+    return MixFit(lam=lam, mu=mu, sigma=sigma, loglik=ll, iterations=iters)

pycopykat/_rcompat/rbase.py

@@ -0,0 +1,112 @@
+"""Pieces of base R with R's arithmetic.
+
+``mean`` adds in long double and corrects with a second pass (summary.c); ``var``/``sd`` use that
+mean and a long-double sum of squares (cov.c); ``median`` averages the two middle values for even
+lengths; ``qnorm`` is R's ``qnorm5`` (Wichura's AS 241).
+"""
+from __future__ import annotations
+
+import numpy as np
+
+_LD = np.longdouble
+
+
+def _seq_sum(m: np.ndarray, axis: int) -> np.ndarray:
+    """Long-double sum along ``axis``, adding the slices in order (R's LDOUBLE loops)."""
+    m = np.moveaxis(np.asarray(m), axis, 0)
+    s = np.zeros(m.shape[1:], dtype=_LD)
+    for i in range(m.shape[0]):
+        s += m[i]
+    return s
+
+
+def r_means(m: np.ndarray, axis: int = 0) -> np.ndarray:
+    """``apply(m, axis + 1 == 2 ? 2 : 1, mean)``: R ``mean`` of each slice along ``axis``."""
+    m = np.asarray(m, dtype=np.float64)
+    n = m.shape[axis]
+    s = _seq_sum(m, axis) / n
+    ok = np.isfinite(s.astype(np.float64))
+    if ok.any():
+        d = np.moveaxis(m, axis, 0).astype(_LD) - s
+        t = np.zeros_like(s)
+        for i in range(d.shape[0]):
+            t += d[i]
+        s = np.where(ok, s + t / n, s)
+    return s.astype(np.float64)
+
+
+def r_mean(x) -> float:
+    """``mean(x)`` for a double vector."""
+    return float(r_means(np.asarray(x, dtype=np.float64).reshape(-1, 1), 0)[0])
+
+
+def r_sds(m: np.ndarray, axis: int = 0) -> np.ndarray:
+    """``sd`` of each slice along ``axis`` (cov.c: long-double mean, then the squares added in
+    long double and divided by n - 1)."""
+    m = np.asarray(m, dtype=np.float64)
+    n = m.shape[axis]
+    mu = r_means(m, axis)
+    d = np.moveaxis(m, axis, 0) - mu
+    ss = _seq_sum(d * d, 0)
+    return np.sqrt((ss / (n - 1)).astype(np.float64))
+
+
+def r_sd(x) -> float:
+    """``sd(x)``."""
+    return float(r_sds(np.asarray(x, dtype=np.float64).reshape(-1, 1), 0)[0])
+
+
+def r_medians(m: np.ndarray, axis: int = 0) -> np.ndarray:
+    """``median`` of each slice along ``axis`` (the mean of the two middle values for even n)."""
+    return np.median(np.asarray(m, dtype=np.float64), axis=axis)
+
+
+def qnorm(p, log=np.log) -> np.ndarray:
+    """``qnorm(p)`` (lower tail): R's ``qnorm5``. ``p`` must lie in (0, 1).
+
+    ``log`` evaluates the tails (|p - 0.5| > 0.425); NumPy's ``log`` can differ from the C library's
+    in the last bit."""
+    p = np.asarray(p, dtype=np.float64)
+    q = p - 0.5
+    val = np.empty_like(p)
+    c = np.abs(q) <= 0.425
+    if c.any():
+        qq = q[c]
+        r = 0.180625 - qq * qq
+        val[c] = (qq * (((((((r * 2509.0809287301226727 + 33430.575583588128105) * r
+                              + 67265.770927008700853) * r + 45921.953931549871457) * r
+                            + 13731.693765509461125) * r + 1971.5909503065514427) * r
+                          + 133.14166789178437745) * r + 3.387132872796366608)
+                  / (((((((r * 5226.495278852854561 + 28729.085735721942674) * r
+                          + 39307.89580009271061) * r + 21213.794301586595867) * r
+                        + 5394.1960214247511077) * r + 687.1870074920579083) * r
+                      + 42.313330701600911252) * r + 1.0))
+    t = ~c
+    if t.any():
+        qt, pt = q[t], p[t]
+        r = np.sqrt(-log(np.where(qt > 0, 1.0 - pt, pt)))
+        out = np.empty_like(r)
+        a = r <= 5.0
+        if a.any():
+            ra = r[a] - 1.6
+            out[a] = ((((((((ra * 7.7454501427834140764e-4 + 0.0227238449892691845833) * ra
+                            + 0.24178072517745061177) * ra + 1.27045825245236838258) * ra
+                          + 3.64784832476320460504) * ra + 5.7694972214606914055) * ra
+                        + 4.6303378461565452959) * ra + 1.42343711074968357734)
+                      / (((((((ra * 1.05075007164441684324e-9 + 5.475938084995344946e-4) * ra
+                              + 0.0151986665636164571966) * ra + 0.14810397642748007459) * ra
+                            + 0.68976733498510000455) * ra + 1.6763848301838038494) * ra
+                          + 2.05319162663775882187) * ra + 1.0))
+        b = ~a
+        if b.any():
+            rb = r[b] - 5.0
+            out[b] = ((((((((rb * 2.01033439929228813265e-7 + 2.71155556874348757815e-5) * rb
+                            + 0.0012426609473880784386) * rb + 0.026532189526576123093) * rb
+                          + 0.29656057182850489123) * rb + 1.7848265399172913358) * rb
+                        + 5.4637849111641143699) * rb + 6.6579046435011037772)
+                      / (((((((rb * 2.04426310338993978564e-15 + 1.4215117583164458887e-7) * rb
+                              + 1.8463183175100546818e-5) * rb + 7.868691311456132591e-4) * rb
+                            + 0.0148753612908506148525) * rb + 0.13692988092273580531) * rb
+                          + 0.59983220655588793769) * rb + 1.0))
+        val[t] = np.where(qt < 0.0, -out, out)
+    return val