mmm4microbiome 0.1.0__py3-none-any.whl

mmm4microbiome/__init__.py

@@ -0,0 +1,47 @@
+"""
+mmm4microbiome: read a mixed-membership sweep as microbiome.
+
+Downstream of ``stripesankey`` (>= 0.4.1): ``stripesankey.prepare`` and
+``stripesankey.sweep`` write a base folder, and everything here reads it. A
+component is called **MC** (microbial component), and its H row is mapped to
+taxa through ``feature_annotation.csv``. Fitting, consensus and the Sankey
+widget itself stay in stripesankey::
+
+    import stripesankey as ss
+    import mmm4microbiome as mmm
+
+    ss.prepare(base, counts, feature_annotation, metadata, level="Genus")
+    ss.sweep(base, k_range=range(2, 11))
+
+    mmm.stripesankey_mb(base, metadata_df=meta, hover_field="Country")
+    mmm.heatmap_mc(base, 5)
+    mmm.compare_groups(base, "K5MC2", "K5MC0")
+    mmm.onestop(base, meta)
+"""
+
+__version__ = "0.1.0"
+
+from ._base import (dominant_samples, fit_level, k_values, load_clusters, load_counts,
+                    load_mc, mc_feature_table, parse_mc, taxonomy_for_features)
+from .sankey import stripesankey_mb, metasankey_mb
+from .sample import SampleHeatmapInline, heatmap_sample, bar_sample
+from .features import heatmap_mc, mc_profile
+from .da_barplot import da_barplot
+from .compare import compare_groups, CompareReport, DAA_METHODS
+from .onestop import onestop
+
+__all__ = [
+    # the base folder
+    "k_values", "fit_level", "load_mc", "load_clusters", "load_counts",
+    "mc_feature_table", "taxonomy_for_features", "parse_mc", "dominant_samples",
+    # the sweep
+    "stripesankey_mb", "metasankey_mb",
+    # one k, sample by sample
+    "SampleHeatmapInline", "heatmap_sample", "bar_sample",
+    # what the MCs are made of
+    "heatmap_mc", "mc_profile",
+    # two groups compared
+    "compare_groups", "CompareReport", "DAA_METHODS", "da_barplot",
+    # everything in one page
+    "onestop",
+]

mmm4microbiome/_base.py

@@ -0,0 +1,352 @@
+"""
+Read the base folder stripesankey writes, the microbiome way.
+
+``stripesankey.prepare`` + ``stripesankey.sweep`` leave one folder per
+analysis::
+
+    base/intermediate/features.csv            sample x feature counts that were fitted
+    base/intermediate/feature_annotation.csv  feature id x taxonomic ranks
+    base/intermediate/manifest.json           the level the counts were collapsed to
+    base/results/W/W_{k}.csv                  sample x LC loadings, one table per k
+    base/results/H/H_{k}.csv                  LC x feature profiles, one table per k
+    base/results/consensus.json               the consensus clusters across k
+    base/results/run_info.json                k list, parameters, convergence
+
+Everything in mmm4microbiome reads through this module. A component (a column
+of W, a row of H) is called LC upstream and **MC** (microbial component) here.
+
+Two taxonomic levels are involved, and they are different things:
+
+* the **fit level** -- what ``prepare(level=)`` collapsed the counts to, i.e.
+  the features the model was fitted on (``level=None`` keeps one feature per
+  ASV);
+* the **display level** -- ``feature_level`` in this package: what a figure
+  names. By default the fit level, except that an ASV (or Species) fit is
+  shown at Genus. A coarser level sums its child probabilities; a level finer
+  than the fit cannot be recovered and falls back to the fit level.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import warnings
+
+import numpy as np
+import pandas as pd
+
+#: Ranks from fine to coarse; a display level can only be at or above the fit.
+_RANKS = ["ASV", "Species", "Genus", "Family", "Order", "Class", "Phylum", "Kingdom", "Domain"]
+#: Annotation columns that hold one value per ASV, i.e. an ASV-level fit.
+_ASV_COLUMNS = {"genus_id", "randomid", "asv", "feature id", "feature_id"}
+_RANK_PREFIX = re.compile(r"^[a-z]__")
+_LC_LABEL = re.compile(r"[LMlm][Cc]\s*(\d+)$")
+# 'K5MC2', 'K5_MC2', 'k5 mc2', 'K=5:MC2', ...
+_MC_SELECTOR = re.compile(r"^\s*[Kk]?\s*[=:]?\s*(\d+)\s*[_ :=\-]*[Mm][Cc]\s*[=:]?\s*(\d+)\s*$")
+
+
+# --------------------------------------------------------------------------- #
+#  JSON helpers
+# --------------------------------------------------------------------------- #
+def _np_default(o):
+    if isinstance(o, np.generic):
+        return o.item()
+    if isinstance(o, np.ndarray):
+        return o.tolist()
+    raise TypeError(f"not JSON serialisable: {type(o).__name__}")
+
+
+def _json_safe(obj):
+    """Round-trip through JSON so numpy scalars / arrays become plain Python."""
+    return json.loads(json.dumps(obj, default=_np_default))
+
+
+# --------------------------------------------------------------------------- #
+#  the base folder
+# --------------------------------------------------------------------------- #
+def _require_sweep(base) -> None:
+    if not os.path.isdir(os.path.join(str(base), "results", "W")):
+        raise FileNotFoundError(
+            f"{base} has no results/W, so it is not a stripesankey base folder yet. "
+            "Run stripesankey.prepare(base, counts, feature_annotation, metadata) and "
+            "stripesankey.sweep(base, k_range=...) first.")
+
+
+def k_values(base) -> list[int]:
+    """The k values with a ``results/W/W_{k}.csv`` in ``base``."""
+    _require_sweep(base)
+    wdir = os.path.join(str(base), "results", "W")
+    found = (re.fullmatch(r"W_(\d+)\.csv", name) for name in os.listdir(wdir))
+    return sorted(int(m.group(1)) for m in found if m)
+
+
+def _check_k(base, k) -> int:
+    k = int(k)
+    ks = k_values(base)
+    if k not in ks:
+        raise KeyError(f"k={k} is not in {base}; the sweep has k={ks}")
+    return k
+
+
+def _mc_index(label, fallback: int) -> int:
+    """'LC3' / 'MC3' / 'K5_MC3' -> 3; any other label -> its position."""
+    m = _LC_LABEL.search(str(label))
+    return int(m.group(1)) if m else fallback
+
+
+def load_mc(base, k) -> pd.DataFrame:
+    """W at one k as a sample x MC table (columns ``MC0 .. MC{k-1}``).
+
+    Values are as stored; ``stripesankey.sweep`` writes rows that sum to 1.
+    """
+    import stripesankey as ss
+
+    k = _check_k(base, k)
+    W = ss.load_W(str(base), [k])[k].copy()
+    W.index = W.index.map(str)
+    W.columns = [f"MC{_mc_index(c, i)}" for i, c in enumerate(W.columns)]
+    return W
+
+
+def load_clusters(base) -> dict:
+    """``{sample: cluster}`` from ``results/consensus.json`` (``-1`` = noise); ``{}`` if absent."""
+    import stripesankey as ss
+
+    try:
+        s2c = ss.load_consensus(str(base)).get("sample_to_cluster") or {}
+    except FileNotFoundError:
+        return {}
+    return {str(s): int(c) for s, c in s2c.items()}
+
+
+def _manifest(base) -> dict:
+    import stripesankey as ss
+
+    try:
+        return ss.load_manifest(str(base))
+    except FileNotFoundError:
+        return {}
+
+
+def _annotation(base) -> pd.DataFrame | None:
+    path = os.path.join(str(base), "intermediate", "feature_annotation.csv")
+    if not os.path.exists(path):
+        return None
+    ann = pd.read_csv(path, index_col=0)
+    ann.index = ann.index.map(str)
+    return ann
+
+
+def load_counts(base) -> pd.DataFrame:
+    """The sample x feature count table the sweep was fitted on."""
+    import stripesankey as ss
+
+    try:
+        X = ss.load_features(str(base))
+    except FileNotFoundError:
+        raise FileNotFoundError(
+            f"{base}/intermediate/features.csv not found -- it is written by "
+            "stripesankey.prepare(). Pass table= explicitly for a sweep fitted elsewhere.")
+    X.index = X.index.map(str)
+    X.columns = X.columns.map(str)
+    return X
+
+
+# --------------------------------------------------------------------------- #
+#  taxonomic levels and names
+# --------------------------------------------------------------------------- #
+def _rank(name) -> str:
+    """A level name or annotation column -> its rank ('genus_ID' -> 'ASV', 'genus' -> 'Genus')."""
+    if name is None:
+        return "ASV"
+    s = str(name)
+    low = s.lower()
+    if low in _ASV_COLUMNS:
+        return "ASV"
+    if low == "genus_based":
+        return "Genus"
+    for r in _RANKS:
+        if low == r.lower():
+            return r
+    return s
+
+
+def _finer(a: str, b: str) -> bool:
+    """Is rank ``a`` strictly finer than rank ``b`` (unknown ranks never are)?"""
+    return a in _RANKS and b in _RANKS and _RANKS.index(a) < _RANKS.index(b)
+
+
+def fit_level(base) -> str:
+    """The rank the model was fitted on, from ``manifest.json`` ('ASV' for ``level=None``)."""
+    import stripesankey as ss
+
+    m = _manifest(base)
+    if m:
+        return _rank(m.get("feature_column"))
+    try:
+        return _rank(ss.load_run_info(str(base)).get("level"))
+    except FileNotFoundError:
+        return "ASV"
+
+
+def _display_level(requested, fit_rank: str) -> str:
+    """The level a figure names: the request, or the default for this fit."""
+    if requested is None:
+        return "Genus" if fit_rank in ("ASV", "Species") else fit_rank
+    target = _rank(requested)
+    if _finer(target, fit_rank):
+        warnings.warn(f"feature_level={requested!r} is finer than the {fit_rank!r} level "
+                      f"the model was fitted on; showing {fit_rank!r}.", stacklevel=4)
+        return fit_rank
+    return target
+
+
+def _column_for(rank: str, columns) -> str | None:
+    """The annotation column holding ``rank`` (case-insensitive)."""
+    lower = {str(c).lower(): c for c in columns}
+    if rank == "ASV":
+        return lower.get("genus_id") or lower.get("randomid") or lower.get("asv")
+    return lower.get(rank.lower()) or (lower.get("genus_based") if rank == "Genus" else None)
+
+
+def _display_name(label) -> str:
+    """'g__Prevotella' -> 'Prevotella' (the level is named by the figure, not the label)."""
+    return _RANK_PREFIX.sub("", str(label)).strip()
+
+
+def _is_unclassified(label) -> bool:
+    """Missing, a bare rank prefix ('g__'), 'Unclassified...' (case-sensitive), or
+    the marker prepare() writes for a missing annotation ('unassigned_<level>')."""
+    if label is None:
+        return True
+    s = str(label)
+    if s.strip().lower() in ("", "nan", "na", "none") or s.startswith("unassigned_"):
+        return True
+    name = _display_name(s)
+    return name == "" or name.startswith("Unclassified")
+
+
+def _labels_at(base, level: str, fit_rank: str):
+    """Fitted token -> its label at ``level``, or None when the annotation lacks that level."""
+    ann = _annotation(base)
+    target = _column_for(level, ann.columns) if ann is not None else None
+    if target is None:
+        warnings.warn(f"feature_level={level!r} is not a column of feature_annotation.csv; "
+                      f"showing the {fit_rank!r} level the model was fitted on.", stacklevel=4)
+        return None
+    token_col = _manifest(base).get("feature_column")
+    if token_col and token_col in ann.columns:          # tokens = values of a rank column
+        sub = ann[[token_col, target]].dropna(subset=[token_col])
+        return sub.groupby(sub[token_col].astype(str))[target].first()
+    return ann[target]                                  # tokens = feature ids (ASV fit)
+
+
+def mc_feature_table(base, k, feature_level=None, *, show_unclassified: bool = True):
+    """What every MC at ``k`` is made of, at one taxonomic level.
+
+    Returns ``(table, level)``: ``table`` is MC x taxon (index ``MC0 ..``),
+    each row the MC's own probability distribution over the fitted features
+    (H, row-normalised), summed into ``level`` when that level is coarser than
+    the fit. Taxon names lose their rank prefix (``g__``). With
+    ``show_unclassified=False`` the unclassified taxa are dropped; the
+    remaining probabilities are not renormalised.
+    """
+    import stripesankey as ss
+
+    k = _check_k(base, k)
+    if not os.path.exists(os.path.join(str(base), "results", "H", f"H_{k}.csv")):
+        raise FileNotFoundError(
+            f"{base}/results/H/H_{k}.csv not found. MC features need H; a sweep "
+            "fitted elsewhere and saved as W only has none.")
+    H = ss.load_H(str(base), k)
+    P = H.div(H.sum(axis=1).replace(0, np.nan), axis=0).fillna(0.0)
+    P.index = [f"MC{_mc_index(n, i)}" for i, n in enumerate(P.index)]
+    P.columns = P.columns.map(str)
+
+    fit = fit_level(base)
+    level = _display_level(feature_level, fit)
+    if level != fit:
+        labels = _labels_at(base, level, fit)
+        if labels is None:
+            level = fit
+        else:
+            names = []
+            for t in P.columns:
+                v = labels.get(t, None)
+                names.append(str(v) if v is not None and pd.notna(v) else f"unassigned_{level}")
+            P = P.T.groupby(pd.Index(names), sort=False).sum().T
+
+    shown = [_display_name(c) or f"unassigned_{level}" for c in P.columns]
+    if len(set(shown)) < len(shown):
+        P = P.T.groupby(pd.Index(shown), sort=False).sum().T
+    else:
+        P.columns = shown
+    if not show_unclassified:
+        P = P.loc[:, [not _is_unclassified(c) for c in P.columns]]
+    P.index.name = "MC"
+    P.columns.name = level
+    return P, level
+
+
+def _top_rows(row: pd.Series, top_n: int, show_unclassified: bool) -> list:
+    """The best ``top_n`` taxa of one MC, as ``[[name, prob], ...]``."""
+    out = []
+    for name, p in row[row > 0].sort_values(ascending=False).items():
+        if not show_unclassified and _is_unclassified(name):
+            continue
+        out.append([str(name), round(float(p), 5)])
+        if len(out) >= top_n:
+            break
+    return out
+
+
+def taxonomy_for_features(base) -> pd.DataFrame | None:
+    """The feature annotation, keyed by the features that were fitted.
+
+    For an ASV fit that is ``feature_annotation.csv`` itself. For a fit
+    collapsed to a rank (say Genus) the fitted features are genus names, so the
+    table is re-keyed by them and keeps only the columns that hold one value per
+    genus -- the ranks at and above it -- plus the rank itself.
+    """
+    ann = _annotation(base)
+    if ann is None:
+        return None
+    token_col = _manifest(base).get("feature_column")
+    if not token_col or token_col not in ann.columns:
+        return ann
+    sub = ann.dropna(subset=[token_col]).copy()
+    key = sub[token_col].astype(str)
+    grouped = sub.groupby(key)
+    keep = [c for c in sub.columns
+            if c != token_col and int(grouped[c].nunique(dropna=True).max() or 0) <= 1]
+    tax = grouped[keep].first() if keep else pd.DataFrame(index=sorted(key.unique()))
+    tax[token_col] = tax.index
+    tax.index.name = token_col
+    return tax
+
+
+# --------------------------------------------------------------------------- #
+#  MC selectors
+# --------------------------------------------------------------------------- #
+def parse_mc(selector) -> tuple[int, int]:
+    """``'K5MC2'`` (also ``'K5_MC2'``, ``'k5 mc2'``) -> ``(5, 2)``."""
+    m = _MC_SELECTOR.match(selector) if isinstance(selector, str) else None
+    if not m:
+        raise ValueError(f"could not parse {selector!r}; expected something like 'K5MC2' "
+                         "(K = the model, MC = the 0-based component index)")
+    return int(m.group(1)), int(m.group(2))
+
+
+def dominant_samples(base, k, mc, high_threshold: float = 0.5) -> list[str]:
+    """Samples dominated by one MC: ``argmax(W) == mc`` and that probability ``>= high_threshold``.
+
+    The rule metasankey draws its nodes with, so a selector picks what the
+    figure shows.
+    """
+    W = load_mc(base, k)
+    if not 0 <= int(mc) < W.shape[1]:
+        raise KeyError(f"MC{mc} is out of range for k={k}; valid: MC0..MC{W.shape[1] - 1}")
+    arr = W.to_numpy(dtype=float)
+    best, top = arr.argmax(axis=1), arr.max(axis=1)
+    return [s for s, b, t in zip(W.index, best, top) if b == int(mc) and t >= high_threshold]