Source code for mir.cli

"""``mir`` command-line interface — turn receptor tables into embeddings.

Two commands cover the two scales mirpy embeds at:

* ``mir embed clonotypes SAMPLE``   — one repertoire's clonotype table → a per-clonotype
  TCREMP embedding table (``e0…``), the input to clustering / ML.
* ``mir embed repertoires SAMPLE…`` — a *dataset* of clonotype tables → one repertoire
  vector ``Φ(S)`` (``phi0…``) per sample, per chain, on one shared basis (so the rows are
  mutually comparable / MMD-able).

Inputs are any format ``vdjtools.io.read`` sniffs (AIRR TSV, vdjtools, MiXCR, immunoSEQ,
parquet, …). Output is TSV (default / ``.tsv``) or Parquet (``.parquet`` — recommended for
the wide raw embedding); ``-o -`` (or no ``-o``) writes TSV to stdout.

Run ``mir embed clonotypes -h`` / ``mir embed repertoires -h`` for the full flag list.
"""
from __future__ import annotations

import argparse
import sys

import polars as pl

import mir


# --- IO helpers ------------------------------------------------------------
def _read(path: str, *, productive_only: bool = True,
          recompute_frequencies: bool = True) -> pl.DataFrame:
    """Read a clonotype file into a normalized AIRR frame, **productive rearrangements only**.

    ``v_identity`` is kept by name: it is the one field the signature needs that the canonical
    schema does not carry, so without it ``vsig:shm:IGH:mean_v_identity`` is not merely absent
    but uncomputable, and ships as a nan column on files that do have it.

    **Non-productive rearrangements are removed on every read, and that cannot be turned off.**
    This is where mirpy deliberately differs from vdjtools, where the same filter is optional. The
    reason is specific: a stop codon is *in* seqtree's amino-acid alphabet, so an unfiltered frame
    does not crash the distance code -- it returns a finite, meaningless number and contaminates
    the geometry silently. mirpy cannot embed a non-productive sequence *meaningfully*, which is a
    stronger and more useful statement than cannot embed it at all.

    Use ``vdjtools`` if you want the non-productive fraction: ``vdjtools filter --nonproductive``,
    or :func:`vdjtools.preprocess.filter_productive` in Python.

    Args:
        path: Clonotype file, any format vdjtools sniffs.
        productive_only: Must stay ``True``. Present so that passing ``False`` raises with an
            explanation rather than silently doing the wrong thing.
        recompute_frequencies: Renormalise ``frequency`` over the surviving rearrangements.

    Raises:
        ValueError: If ``productive_only`` is ``False``.
    """
    from vdjtools import io
    from vdjtools.preprocess import filter_productive, productive_mask

    if not productive_only:
        raise ValueError(
            "mirpy cannot embed non-productive rearrangements meaningfully: a stop codon is in "
            "seqtree's alphabet, so it does not raise in the distance code -- it returns a finite, "
            "meaningless value and contaminates the geometry silently. productive_only=False is "
            "therefore refused here, unlike in vdjtools where the same filter is optional. Use "
            "`vdjtools filter --nonproductive` or vdjtools.preprocess.filter_productive() if the "
            "non-productive fraction is what you want.")
    df = io.read(path, keep=("v_identity",))
    n0 = df.height
    _, src = productive_mask(df)
    df = filter_productive(df, recompute_frequencies=recompute_frequencies)
    if df.height < n0:
        print(f"[mir] dropped {n0 - df.height} non-productive rearrangement(s) of {n0} "
              f"(productivity read from {src}); mirpy always does this", file=sys.stderr)
    return df


def _with_locus(df: pl.DataFrame) -> pl.DataFrame:
    """Ensure a ``locus`` column (derive from ``v_call`` when absent)."""
    if "locus" in df.columns and df["locus"].null_count() < df.height:
        return df
    from vdjtools import io

    try:
        return io.add_locus(df)
    except Exception:
        # Fallback: the IMGT locus is the v_call's first 3 characters (TRB/TRA/IGH/…).
        return df.with_columns(pl.col("v_call").str.slice(0, 3).alias("locus"))


def _emb_frame(X, prefix: str) -> pl.DataFrame:
    """(N, d) float matrix → a polars frame with columns ``{prefix}0…{prefix}{d-1}``."""
    return pl.from_numpy(X, schema=[f"{prefix}{i}" for i in range(X.shape[1])])


def _per_locus_path(path: str, locus: str) -> str:
    """Insert ``locus`` before the extension: ``mmd.tsv`` → ``mmd.TRB.tsv``.

    A plain ``str.replace(".", …, 1)`` would split on the first dot *anywhere* — mangling
    ``./mmd.tsv`` into ``.TRB./mmd.tsv`` — and would silently no-op on an extensionless path,
    letting one locus overwrite another's matrix.
    """
    from pathlib import Path

    p = Path(path)
    return str(p.with_name(f"{p.stem}.{locus}{p.suffix}"))


def _write(df: pl.DataFrame, path: str | None) -> None:
    if path is None or path == "-":
        sys.stdout.write(df.write_csv(separator="\t"))
    elif path.endswith(".parquet"):
        df.write_parquet(path)
    else:
        df.write_csv(path, separator="\t")


def _sample_id(path: str) -> str:
    """Sample id = filename up to the first dot (``P1.TRB.tsv.gz`` → ``P1``)."""
    import os

    return os.path.basename(path).split(".")[0]


def _pick_locus(df: pl.DataFrame, requested: str | None) -> str:
    """Resolve ``--locus`` to a canonical IMGT locus, or infer it when the file has only one."""
    loci = [x for x in df["locus"].unique().to_list() if x]
    if requested:
        # Through the alias table, so `--locus beta` matches the data's `TRB` rows rather
        # than silently selecting nothing.
        from mir.aliases import normalize_locus_alias

        try:
            return normalize_locus_alias(requested)
        except ValueError as exc:
            raise SystemExit(str(exc)) from None
    if len(loci) == 1:
        return loci[0]
    raise SystemExit(
        f"multiple loci present ({', '.join(sorted(loci))}); pass --locus to pick one"
    )


def _require_productive(enabled: bool) -> None:
    """mirpy has no --no-filter-functional. Refuse it, and say where to go instead."""
    if not enabled:
        raise SystemExit(
            "[mir] --no-filter-functional is not available in mirpy. A stop codon is in seqtree's "
            "alphabet, so an unfiltered frame does not crash -- it embeds to a finite, meaningless "
            "distance and contaminates the geometry silently. Use `vdjtools filter "
            "--nonproductive` if the non-productive fraction is what you want.")


def _apply_functional_filter(sub: pl.DataFrame, enabled: bool) -> pl.DataFrame:
    """Drop unusable clonotypes unless disabled.

    Uses ``vdjtools.signature.blocks.sanitise``, NOT ``preprocess.filter_functional``. The two are
    not the same predicate and the difference matters here: ``filter_functional`` is a denylist
    (``[*atgc#~_?]``) and keeps the ambiguity codes ``X``/``B``/``Z``, while ``sanitise`` is an
    anchored allowlist of the 20 standard amino acids and drops them. ``TCREmp.embed`` refuses
    exactly what ``sanitise`` drops, so filtering with the weaker predicate here would leave rows
    the embedder then rejects -- the default path has to clear its own guard.

    Neither predicate filters on LENGTH. A two-residue junction and a sixty-residue one both
    survive; the question asked is only whether the string is a plain amino-acid string.
    """
    _require_productive(enabled)
    from vdjtools.signature.blocks import sanitise

    n0 = sub.height
    sub, dropped = sanitise(sub)
    if sub.height < n0:
        print(f"[mir] filtered {n0 - sub.height} unusable clonotype(s) "
              f"({100 * dropped:.2f}% of reads: stop codon, out-of-frame marker, ambiguity code "
              "or non-positive count)", file=sys.stderr)
    return sub


# --- commands --------------------------------------------------------------
def cmd_clonotypes(a: argparse.Namespace) -> None:
    from mir.embedding.pca import pca_denoise
    from mir.embedding.tcremp import TCREmp

    df = _with_locus(_read(a.input))
    locus = _pick_locus(df, a.locus)
    sub = df.filter(pl.col("locus") == locus)
    if sub.is_empty():
        raise SystemExit(f"no clonotypes for locus {locus!r} in {a.input}")
    sub = _apply_functional_filter(sub, a.filter_functional)
    if sub.is_empty():
        raise SystemExit(f"no coding clonotypes remain for locus {locus!r} after functional filtering")

    model = TCREmp.from_defaults(a.species, locus, n_prototypes=a.n_prototypes,
                                 mode=a.mode, replicate=a.replicate, threads=a.threads)
    X = model.embed(sub)
    if a.pca:
        X = pca_denoise(X, n_components=a.pca)

    if (a.output is None or a.output == "-" or a.output.endswith(".tsv")) and X.shape[1] > 500:
        print(f"[mir] {X.shape[1]} embedding columns — consider --pca K or a .parquet output.",
              file=sys.stderr)

    id_cols = [c for c in ("junction_aa", "v_call", "j_call", "duplicate_count") if c in sub.columns]
    out = sub.select(id_cols).hstack(_emb_frame(X, "e"))
    _write(out, a.output)
    print(f"[mir] embedded {X.shape[0]} {locus} clonotypes → {X.shape[1]}-d "
          f"({'PCA ' if a.pca else ''}table)", file=sys.stderr)


def cmd_repertoires(a: argparse.Namespace) -> None:
    from collections import defaultdict

    from mir.embedding.tcremp import TCREmp
    from mir.repertoire import fit_repertoire_space, mmd_matrix, sample_embedding

    blocks = tuple(b.strip() for b in a.blocks.split(",") if b.strip())
    n_rff_second = a.n_rff_second if "second" in blocks else 0

    # Load every sample, split its clonotypes by locus.
    by_locus: dict[str, list] = defaultdict(list)
    for path in a.input:
        df = _with_locus(_read(path))
        sid = _sample_id(path)
        for locus in [x for x in df["locus"].unique().to_list() if x]:
            if a.locus and locus != a.locus:
                continue
            sub = _apply_functional_filter(df.filter(pl.col("locus") == locus), a.filter_functional)
            if sub.is_empty():
                print(f"[mir] {sid}/{locus}: no coding clonotypes after functional filtering, "
                      "skipping this sample/locus", file=sys.stderr)
                continue
            by_locus[locus].append((sid, sub))

    if not by_locus:
        raise SystemExit("no samples/loci to embed (check inputs / --locus)")

    rows: list[dict] = []
    vectors: list = []
    for locus in sorted(by_locus):
        items = by_locus[locus]
        model = TCREmp.from_defaults(a.species, locus, n_prototypes=a.n_prototypes,
                                    replicate=a.replicate, threads=a.threads)
        pooled = pl.concat([sub for _, sub in items])
        space = fit_repertoire_space(model, pooled, n_rff=a.n_rff, n_rff_second=n_rff_second,
                                     n_components=a.n_components, seed=a.seed)
        embs = [sample_embedding(space, sub, weight=a.weight, blocks=blocks) for _, sub in items]
        for (sid, sub), se in zip(items, embs):
            rows.append({"sample_id": sid, "locus": locus, "n_clonotypes": sub.height})
            vectors.append(se.vector)
        if a.mmd:
            D = mmd_matrix(embs, unbiased=True)
            ids = [sid for sid, _ in items]
            mmd_df = pl.DataFrame({"sample_id": ids}).with_columns(
                [pl.Series(ids[j], D[:, j]) for j in range(len(ids))])
            out = a.mmd if len(by_locus) == 1 else _per_locus_path(a.mmd, locus)
            _write(mmd_df, out)
        print(f"[mir] {locus}: {len(items)} samples → Φ dim {len(embs[0].vector)}", file=sys.stderr)

    import numpy as np

    meta = pl.DataFrame(rows)
    out = meta.hstack(_emb_frame(np.vstack(vectors), "phi"))
    _write(out, a.output)


def cmd_signature(a: argparse.Namespace) -> None:
    from vdjtools.signature import presets as P

    from mir.signature import assemble, columns, describe

    # --preset picks BOTH the tier and the column subset. mirpy is where the two halves meet, so
    # unlike `vdjtools signature` every preset resolves here in full.
    keep = None
    if getattr(a, "preset", None):
        try:
            spec = P.get(a.preset)
        except KeyError as e:
            raise SystemExit(str(e)) from None
        a.tier, keep = spec.tier, spec.columns()
        print(f"[mir] preset {spec.name!r} [{spec.rank}]: {len(keep)} columns, "
              f"suggested scaling {spec.scaling}", file=sys.stderr)

    if a.describe:
        d = describe(a.tier)
        _write(d.filter(pl.col("column").is_in(keep)) if keep else d, a.output)
        return

    # A sample is one file, or several files sharing a sample id — a donor sequenced on TRA and
    # TRB is one signature with both loci filled, not two half-empty ones.
    from collections import defaultdict

    samples: dict[str, dict[str, pl.DataFrame]] = defaultdict(dict)
    for path in a.input:
        df = _with_locus(_read(path))
        sid = _sample_id(path)
        for locus in [x for x in df["locus"].unique().to_list() if x]:
            sub = df.filter(pl.col("locus") == locus)
            if sub.height:
                samples[sid][locus] = sub

    if not samples:
        raise SystemExit("no samples to sign (check inputs)")

    scale = None
    if a.standardize == "reference" and a.scale:
        from mir.signature.scale import load_scale
        scale = load_scale(a.scale)

    out = assemble.signature_cohort(samples, tier=a.tier, species=a.species, weight=a.weight,
                                    standardize=a.standardize, scale=scale,
                                    n_jobs=a.threads, columns=keep)
    n_cols = len(keep) if keep else len(columns(a.tier))
    print(f"[mir] {out.height} samples x {n_cols} columns "
          f"({a.preset or a.tier}, standardize={a.standardize})", file=sys.stderr)
    _write(out, a.output)


def cmd_presets(a: argparse.Namespace) -> None:
    """List the feature presets, or explain one.

    `recommended` — use unless you have a reason not to. `specific` — correct for a stated purpose
    and wrong outside it. `avoid` — a control or a measured dead end, named so that picking it is
    deliberate rather than accidental.
    """
    from vdjtools.signature import presets as P

    if not a.name:
        _write(P.table().select("preset", "rank", "columns", "halves", "scaling", "summary"),
               a.output)
        return
    try:
        spec = P.get(a.name)
    except KeyError as e:
        raise SystemExit(str(e)) from None
    print(f"{spec.name}  [{spec.rank}]  {spec.n_columns} columns  tier={spec.tier}  "
          f"halves={'+'.join(spec.sig)}  scaling={spec.scaling}\n")
    for label, text in (("summary", spec.summary), ("features", spec.features),
                        ("how it is computed", spec.how), ("use cases", spec.use_cases),
                        ("notes", spec.notes)):
        if text:
            print(f"{label}:\n  {text}\n")


# --- parser ----------------------------------------------------------------
[docs] def build_parser() -> argparse.ArgumentParser: """Return the ``mir`` argument parser (``embed clonotypes`` / ``embed repertoires``).""" p = argparse.ArgumentParser(prog="mir", description=__doc__.splitlines()[0]) p.add_argument("--version", action="version", version=f"mir {mir.__version__}") sub = p.add_subparsers(dest="cmd", required=True) embed = sub.add_parser("embed", help="compute embeddings").add_subparsers(dest="what", required=True) c = embed.add_parser("clonotypes", help="repertoire → per-clonotype embedding table") c.add_argument("input", help="a clonotype table (AIRR/vdjtools/MiXCR/parquet/…)") c.add_argument("-o", "--output", help="output .tsv/.parquet (default: stdout TSV)") c.add_argument("--species", default="human") c.add_argument("--locus", help="chain to embed (inferred if the file has one locus)") c.add_argument("--n-prototypes", type=int, default=None, help="prototype count (default: per-chain preset)") c.add_argument("--mode", default="vjcdr3", choices=("vjcdr3", "cdr123")) c.add_argument("--replicate", type=int, default=0, metavar="R", help="prototype draw: 0 = the default set; r>0 = an independent disjoint draw of the same size, for prototype-sensitivity runs (embeddings across draws are NOT comparable)") c.add_argument("--pca", type=int, default=None, metavar="K", help="PCA-denoise the embedding to K dims (compact table)") c.add_argument("--filter-functional", action=argparse.BooleanOptionalAction, default=True, help="drop non-coding clonotypes (stop codon / out-of-frame junction_aa) " "before embedding (default: on)") c.add_argument("--threads", type=int, default=0, help="0 = all cores") c.set_defaults(func=cmd_clonotypes) r = embed.add_parser("repertoires", help="dataset of clonotype tables → per-sample Φ(S), by chain") r.add_argument("input", nargs="+", help="one clonotype file per repertoire (sample id = filename stem)") r.add_argument("-o", "--output", help="output .tsv/.parquet (default: stdout TSV)") r.add_argument("--species", default="human") r.add_argument("--locus", help="restrict to one chain (default: all loci present, one basis each)") r.add_argument("--n-prototypes", type=int, default=None) r.add_argument("--replicate", type=int, default=0, metavar="R", help="prototype draw: 0 = the default set; r>0 = an independent disjoint draw of the same size, for prototype-sensitivity runs (embeddings across draws are NOT comparable)") r.add_argument("--weight", default="log2p1", choices=("log2p1", "duplicate_count", "distinct", "log1p", "anscombe"), help="clone-size weight g (frequencies w = g(a)/Σg): log2p1 g=log2(1+a) " "(default), duplicate_count g=a (linear), distinct g=1 (presence), " "log1p g=ln(1+a), anscombe g=√(a+3/8)") r.add_argument("--blocks", default="mean,diversity", help="Φ blocks: mean,diversity[,second] (second = heavy HLA-interaction block)") r.add_argument("--n-rff", type=int, default=1024, help="mean-block RFF dimension") r.add_argument("--n-rff-second", type=int, default=128, help="second-moment RFF dimension (if used)") r.add_argument("--n-components", type=int, default=None, help="clonotype-PCA dims for the shared basis (default: preset)") r.add_argument("--mmd", metavar="OUT", help="also write the per-chain pairwise unbiased-MMD matrix") r.add_argument("--filter-functional", action=argparse.BooleanOptionalAction, default=True, help="drop non-coding clonotypes (stop codon / out-of-frame junction_aa) " "before embedding (default: on)") r.add_argument("--threads", type=int, default=0, help="0 = all cores") r.add_argument("--seed", type=int, default=0) r.set_defaults(func=cmd_repertoires) s = sub.add_parser( "signature", help="clonotype tables → the portable repertoire signature (one row/sample)", formatter_class=argparse.RawDescriptionHelpFormatter, description=( "One repertoire in, one row of named features out - ready for a classifier.\n" "\n" "Emits BOTH halves of the portable repertoire signature as one vector: the\n" "`vsig` statistics (from vdjtools) and the `rsig` embedding geometry (here).\n" "Fixed, named and positional, so column i means the same thing in every\n" "matrix anyone computes - this is the object you hand a collaborator.\n" "Reads AIRR Rearrangement, native vdjtools, Parquet and the usual\n" "third-party exports, auto-detected. Writes TSV, or Parquet if -o ends\n" "in .parquet.\n" ), epilog=( "START HERE\n" " # a cohort; files sharing a sample id join into one multi-locus sample\n" " mir signature --preset classify cohort/*.tsv.gz -o sig.parquet\n" "\n" " # which columns am I about to get? reads no input at all\n" " mir signature --preset classify --describe\n" "\n" " # all cores, one process per sample\n" " mir signature --preset classify --threads 0 cohort/*.tsv.gz -o sig.parquet\n" "\n" "PICK A PRESET rather than columns by hand (`mir presets` lists all):\n" " compact smallest vector that still describes a repertoire (n >= 50)\n" " classify general-purpose; the usual random-forest / boosting input\n" " transfer for a model that must work on ANOTHER LAB's samples\n" "\n" "--preset overrides --tier; with neither you get all of --tier (standard).\n" "Unlike `vdjtools signature`, every preset resolves here in full, because\n" "mirpy is where the two halves meet.\n" "\n" "GOTCHAS\n" " * CDR3 vs junction. The reader prefers AIRR `junction_aa` (anchors\n" " INCLUDED) and falls back to IMGT `cdr3_aa` (anchors excluded), so a\n" " file carrying only `cdr3_aa` is two residues short everywhere --\n" " shifting the length, k-mer and Pgen features. Check your headers.\n" " * Do not PCA-project the result. Plain scaling beat projection at\n" " every rank tested.\n" " * --standardize reference (the default) is what makes your vector\n" " comparable with someone else's. Use 'none' only to inspect raw\n" " block values.\n" ), ) s.add_argument("input", nargs="*", help="one or more clonotype files; files sharing a sample id (the name up to " "the first dot) are joined into one multi-locus sample") s.add_argument("-o", "--output", help="output .tsv/.parquet (default: stdout TSV)") s.add_argument("--tier", default="standard", choices=("core", "standard", "full"), help="column set; the narrower tiers are exact index subsets of the wider ones") s.add_argument("--species", default="human") s.add_argument("--weight", default="log2p1", choices=("log2p1", "duplicate_count", "distinct", "log1p", "anscombe"), help="clone-size weight g (default log2p1)") s.add_argument("--standardize", default="reference", choices=("reference", "none"), help="'reference' rescales every column against the bundled reference so the " "vector is comparable with anyone else's (default); 'none' emits raw " "block values") s.add_argument("--scale", default=None, help="path to an alternative scale artifact (default: the bundled one)") s.add_argument("--preset", default=None, help="named feature set; overrides --tier (see `mir presets`)") s.add_argument("--threads", type=int, default=1, help="worker processes over samples; 0 = every core (default 1)") s.add_argument("--describe", action="store_true", help="print the column dictionary for --tier/--preset and exit; reads no input") s.set_defaults(func=cmd_signature) q = sub.add_parser( "presets", help="list the named feature sets with their rankings", formatter_class=argparse.RawDescriptionHelpFormatter, description=( "List the named feature sets for `mir signature`, with their rankings.\n" ), epilog=( " mir presets # the table: name, rank, width, halves, scaling\n" " mir presets classify # one preset in full: what, how, when, caveats\n" "\n" "Ranks tell you how much to trust a choice:\n" " recommended use one of these unless you have a reason not to\n" " specific correct for a stated purpose and wrong outside it\n" " avoid a control or a measured dead end, named so that picking\n" " it is deliberate\n" "\n" "The `halves` column says whether a preset needs `vsig` (statistics),\n" "`rsig` (embedding geometry), or both. `mir signature` serves all three;\n" "`vdjtools signature` emits only the `vsig:` columns.\n" ), ) q.add_argument("name", nargs="?", help="show one preset in full") q.add_argument("-o", "--output", help="output .tsv/.parquet (default: stdout TSV)") q.set_defaults(func=cmd_presets) return p
[docs] def main(argv: list[str] | None = None) -> None: """Parse ``argv`` (default ``sys.argv[1:]``) and run the requested ``mir`` subcommand.""" args = build_parser().parse_args(argv) args.func(args)
if __name__ == "__main__": main()