Source code for vdjtools.preprocess.filter

"""Clonotype filtering (pure polars).

Reimplements the legacy vdjtools clonotype-filter family:

- :func:`filter_productive` — AIRR-productive rearrangements (supersedes ``filter_functional``,
  the legacy ``FunctionalClonotypeFilter`` / ``isCoding``).
- :func:`filter_frequency` — ``FilterByFrequency`` (``FrequencyFilter`` + ``QuantileFilter``).
- :func:`filter_segment` — ``FilterBySegment`` (``VFilter`` / ``DFilter`` / ``JFilter``).
- :func:`filter_by_sample` — ``ApplySampleAsFilter`` (``IntersectionClonotypeFilter``).

Every filter recomputes ``frequency`` within the surviving subset by default, and
:func:`filter_productive` exposes that as ``recompute_frequencies=False`` for a caller who wants
the file's own frequencies left alone.

**On the word "functional".** It is IMGT's, and IMGT applies it to a *germline gene* (F / ORF / P),
not to a rearrangement. What this module filters is AIRR **productivity** — a property of the
rearranged sequence: no stop codon, junction in frame. The two are orthogonal; a perfectly
productive rearrangement can use a pseudogene V. :func:`filter_productive` is therefore the name,
and :func:`filter_functional` is kept as a deprecated alias.
"""
from __future__ import annotations

import polars as pl

import warnings

from ..io.schema import (
    JUNCTION_AA,
    COUNT,
    D_CALL,
    FREQ,
    J_CALL,
    PRODUCTIVE,
    STOP_CODON,
    V_CALL,
    VJ_IN_FRAME,
    column_names,
    strip_allele,
    recompute_frequency,
)

#: Characters that mark a non-coding CDR3 amino-acid string. ``*`` is the stop
#: codon; the lowercase nucleotides and ``# ~ _ ?`` are the legacy out-of-frame
#: markers emitted when a junction cannot be cleanly translated (legacy
#: ``CommonUtil.OOF_SYMBOLS_POSSIBLE`` + ``STOP_CHAR``).
_NONCODING_CHARS = r"[*atgc#~_?]"


def _productive_from_annotation(df: pl.DataFrame) -> "tuple[pl.Expr, str] | None":
    """Read productivity off the file's own AIRR columns, if it carries any.

    Preference order is the composite first, then its components: ``productive`` states the whole
    of it (open reading frame, no defect in start codon / splicing / regulatory elements, no
    internal stop, junction in frame), while ``stop_codon`` and ``vj_in_frame`` each state a part.
    Re-deriving any of this from ``junction_aa`` when the file already says it is how a caller ends
    up disagreeing with their own annotation without being told.
    """
    cols = set(column_names(df))
    if PRODUCTIVE in cols:
        return pl.col(PRODUCTIVE).cast(pl.Boolean, strict=False), PRODUCTIVE
    parts, used = [], []
    if STOP_CODON in cols:
        parts.append(~pl.col(STOP_CODON).cast(pl.Boolean, strict=False))
        used.append(STOP_CODON)
    if VJ_IN_FRAME in cols:
        parts.append(pl.col(VJ_IN_FRAME).cast(pl.Boolean, strict=False))
        used.append(VJ_IN_FRAME)
    if not parts:
        return None
    expr = parts[0]
    for e in parts[1:]:
        expr = expr & e
    return expr, "+".join(used)


[docs] def productive_mask(df: pl.DataFrame) -> "tuple[pl.Expr, str]": """The productivity predicate for this frame, and the evidence it rests on. Returns ``(expr, source)`` where ``source`` is the AIRR column(s) used, or ``"junction_aa"`` when none were present and productivity had to be derived from the amino-acid string. The fallback reads a stop codon as ``*`` and an out-of-frame junction as one of the legacy markers. It is a proxy: it cannot see a defect in a splicing site or a regulatory element, which AIRR's ``productive`` can. """ ann = _productive_from_annotation(df) if ann is not None: expr, src = ann return expr.fill_null(False), src return (pl.col(JUNCTION_AA).is_not_null() & ~pl.col(JUNCTION_AA).str.contains(_NONCODING_CHARS)), JUNCTION_AA
[docs] def filter_productive(df: pl.DataFrame, keep: str = "productive", *, recompute_frequencies: bool = True) -> pl.DataFrame: """Keep only AIRR-productive rearrangements (or only the complement). A rearrangement is *productive* when it can encode a receptor chain. Where the frame carries the AIRR annotation columns (``productive``, or ``stop_codon`` / ``vj_in_frame``) those are authoritative; otherwise productivity is derived from ``junction_aa``, where a stop codon is ``*`` and an out-of-frame junction carries one of the legacy markers ``[atgc#~_?]``. A null ``junction_aa`` is treated as non-productive. This is **not** IMGT functionality. That is a property of the germline gene (F / ORF / P) and is orthogonal to this one — see :func:`filter_functional_genes`. Args: df: A clonotype frame. keep: ``"productive"`` (default) or ``"nonproductive"`` for the complement. recompute_frequencies: Renormalise ``frequency`` over the survivors. **Default ``True``**, which is the legacy behaviour and what almost every caller wants. Pass ``False`` to leave the file's own frequencies untouched — useful when the frequencies are the quantity of interest and must stay comparable to the unfiltered file. Returns: The filtered frame. Raises: ValueError: If ``keep`` is not ``"productive"`` or ``"nonproductive"``. """ if keep not in ("productive", "nonproductive"): raise ValueError(f"keep must be 'productive' or 'nonproductive'; got {keep!r}") mask, _ = productive_mask(df) out = df.filter(mask if keep == "productive" else ~mask) return recompute_frequency(out) if recompute_frequencies else out
[docs] def filter_functional(df: pl.DataFrame, keep: str = "coding") -> pl.DataFrame: """Deprecated alias for :func:`filter_productive`. Use that instead. "Functional" is IMGT's word for a germline gene; this function filters rearrangements, which AIRR calls *productive*. Kept for one release so existing callers do not break silently. """ warnings.warn( "filter_functional() is deprecated; use filter_productive(). 'functional' is IMGT's term " "for a germline gene (F/ORF/P), while this filters rearrangements, which AIRR calls " "'productive'. Map keep='coding' -> keep='productive', " "keep='noncoding' -> keep='nonproductive'.", DeprecationWarning, stacklevel=2) if keep not in ("coding", "noncoding"): raise ValueError(f"keep must be 'coding' or 'noncoding'; got {keep!r}") return filter_productive(df, "productive" if keep == "coding" else "nonproductive")
[docs] def filter_functional_genes(df: pl.DataFrame, *, segments: tuple[str, ...] = ("V", "J"), keep: tuple[str, ...] = ("F",), organism: str = "human", locus: str | None = None, recompute_frequencies: bool = True) -> pl.DataFrame: """Keep rearrangements whose germline gene calls are IMGT-functional. This is the **other** axis, and the one that actually deserves the word *functional*. IMGT classifies a germline gene as **F** (functional), **ORF** (an open reading frame, but a defect in splicing, regulatory elements or conserved-residue hydropathy — *not* functional), or **P** (pseudogene: a defect in the ORF itself). See https://www.imgt.org/IMGTindex/functionality.php It is orthogonal to :func:`filter_productive`. A rearrangement can be perfectly in frame with no stop codon — AIRR-productive — while using a pseudogene V; and a functional V gene can rearrange out of frame. Filtering one says nothing about the other. A call this function cannot resolve against the germline reference is **kept**, not dropped: an unrecognised gene name means our reference is incomplete or the caller uses a different nomenclature, and silently discarding those rows would be a vocabulary bug reported as biology. Args: df: A clonotype frame. segments: Which calls to check — any of ``"V"``, ``"D"``, ``"J"``. keep: IMGT functionality codes to keep. ``("F",)`` is strict; ``("F", "ORF")`` is the common looser choice. organism: Passed to the germline reference. locus: Locus to load the reference for. Inferred from the V calls when omitted. recompute_frequencies: Renormalise ``frequency`` over the survivors. Default ``True``. Returns: The filtered frame. """ from ..model.reference import load_germline if locus is None: v = df[V_CALL].drop_nulls() if not v.len(): return df locus = str(v[0])[:3].upper() germ = load_germline(locus, organism) ok = {seg: set(germ.filter((pl.col("segment") == seg) & pl.col("functionality").is_in(list(keep)))["gene"].to_list()) for seg in segments} known = {seg: set(germ.filter(pl.col("segment") == seg)["gene"].to_list()) for seg in segments} mask = pl.lit(True) for seg, col in (("V", V_CALL), ("D", D_CALL), ("J", J_CALL)): if seg not in segments or col not in column_names(df): continue gene = strip_allele(pl.col(col)) # unknown -> kept: an unrecognised name is a vocabulary gap, not a pseudogene mask = mask & (gene.is_in(list(ok[seg])) | ~gene.is_in(list(known[seg])) | gene.is_null()) out = df.filter(mask) return recompute_frequency(out) if recompute_frequencies else out
#: Default junction_aa length bounds, INCLUSIVE, in amino acids. #: #: A CDR3 shorter than 5 aa cannot span the Cys104..Phe118 anchors with any diversity between them, #: and a junction longer than 60 aa is beyond anything the germline can produce -- both are almost #: always a misparse or a chimeric read rather than a receptor. Deliberately wide: this is a #: sanity bound, not a biological filter, and real junctions sit far inside it. MIN_JUNCTION_AA, MAX_JUNCTION_AA = 5, 60
[docs] def filter_length(df: pl.DataFrame, *, min_len: int = MIN_JUNCTION_AA, max_len: int = MAX_JUNCTION_AA, keep: str = "within", recompute_frequencies: bool = True) -> pl.DataFrame: """Keep clonotypes whose ``junction_aa`` length is within bounds, **inclusive**. Both bounds are inclusive: ``min_len=5`` keeps a 5-mer, ``max_len=60`` keeps a 60-mer. This is a data-sanity filter, not a receptor-biology one. It catches misparses, truncated reads and chimeras; it does not encode a claim about what lengths are immunologically interesting. Note that neither this nor :func:`filter_productive` filters on length by default -- nothing upstream in this package has ever imposed a length bound, so switching this on will change counts on any corpus that carries junk. **CDR3 vs junction.** These bounds are on ``junction_aa``, which *includes* the Cys104 and Phe118 anchors and is therefore two residues longer than the IMGT CDR3. Subtract 2 if you are reasoning in CDR3 lengths. A null ``junction_aa`` is dropped by ``keep="within"`` -- an absent junction has no length. Args: df: A clonotype frame. min_len: Shortest ``junction_aa`` to keep, inclusive. max_len: Longest ``junction_aa`` to keep, inclusive. keep: ``"within"`` (default) or ``"outside"`` for the complement -- useful for inspecting what a bound would discard before committing to it. recompute_frequencies: Renormalise ``frequency`` over the survivors. Default ``True``. Returns: The filtered frame. Raises: ValueError: If ``keep`` is unknown, or ``min_len`` exceeds ``max_len``. """ if keep not in ("within", "outside"): raise ValueError(f"keep must be 'within' or 'outside'; got {keep!r}") if min_len > max_len: raise ValueError(f"min_len ({min_len}) exceeds max_len ({max_len})") n = pl.col(JUNCTION_AA).str.len_chars() within = pl.col(JUNCTION_AA).is_not_null() & (n >= min_len) & (n <= max_len) out = df.filter(within if keep == "within" else ~within) return recompute_frequency(out) if recompute_frequencies else out
[docs] def filter_frequency(df: pl.DataFrame, min_freq: float | None = None, top_quantile: float | None = None) -> pl.DataFrame: """Keep abundant clonotypes by frequency threshold and/or top quantile. Reimplements ``FilterByFrequency`` (a composite of ``FrequencyFilter`` and ``QuantileFilter``). Both criteria, when given, are combined with AND: - ``min_freq``: keep clonotypes with ``frequency >= min_freq``. - ``top_quantile``: keep the top clonotypes (by ``duplicate_count``) whose *cumulative* original frequency, including the clonotype itself, is at most ``top_quantile`` of the full-sample total frequency. This matches the legacy ``QuantileFilter``: it walks the count-sorted sample accumulating frequency and drops the first clonotype that would push the running fraction above the threshold (so ``top_quantile=0.25`` keeps roughly the top 25% of the read mass). The denominator is the full-sample frequency total (~1.0), and only clonotypes that already passed ``min_freq`` contribute to the cumulative (legacy filters short-circuit in the order count/freq/quantile). Args: df: A clonotype frame with ``duplicate_count`` and ``frequency`` columns. min_freq: Minimum per-clonotype frequency (e.g. legacy default ``0.01``). ``None`` disables it. top_quantile: Top read-mass quantile to retain (e.g. legacy default ``0.25``). ``None`` disables it. Returns: The filtered frame, sorted by descending ``duplicate_count``, with ``frequency`` recomputed. """ out = df.sort(COUNT, descending=True, maintain_order=True) total_freq = out[FREQ].sum() # legacy parent.getFreqAsInInput(): full-sample total if min_freq is not None: out = out.filter(pl.col(FREQ) >= min_freq) if top_quantile is not None and total_freq: cutoff = top_quantile * total_freq out = out.filter(pl.col(FREQ).cum_sum() <= cutoff) return recompute_frequency(out)
def _segment_matches(col: str, names: list[str]) -> pl.Expr: """Boolean expression: does this segment call match any query name (prefix)? Incomplete query names act as wildcards (legacy ``getAtFuzzy``): a match is a prefix match, so ``TRBV12`` matches ``TRBV12-3*01`` (allele-insensitive) while ``TRBV12-3*01`` matches only itself. The match is per **comma-token**, not on the raw string: an AIRR call may be an ambiguity tie like ``IGHV3-23*01,IGHV3-23D*01``, and a ``starts_with`` on the whole string only ever tests the FIRST gene -- so ``filter_segment(v=["IGHV3-23D"])`` would match none of the tens of thousands of rows that name IGHV3-23D only in a tie. """ # Regex over the whole call: `name` at the start of the string OR right after a comma # (+ optional whitespace). re.escape guards names containing regex metacharacters. import re pat = "(^|,)\\s*(" + "|".join(re.escape(n) for n in names) + ")" return pl.col(col).str.contains(pat)
[docs] def filter_segment(df: pl.DataFrame, v: list[str] | None = None, d: list[str] | None = None, j: list[str] | None = None, keep: bool = True) -> pl.DataFrame: """Keep or remove clonotypes by V/D/J segment membership. Reimplements ``FilterBySegment``. A clonotype *matches* when its V segment is in ``v`` **and** its D segment in ``d`` **and** its J segment in ``j`` (only the lists that are supplied constrain; unsupplied loci always pass). Matching is a prefix match, so incomplete names act as wildcards and are allele-insensitive (``TRBV12`` matches ``TRBV12-3*01``). Args: df: A clonotype frame with ``v_call`` / ``d_call`` / ``j_call`` columns. v: V-segment query names (prefixes). ``None`` leaves V unconstrained. d: D-segment query names. ``None`` leaves D unconstrained. j: J-segment query names. ``None`` leaves J unconstrained. keep: If ``True`` (default) keep matching clonotypes; if ``False`` remove them (legacy ``--negative``). Returns: The filtered frame with ``frequency`` recomputed. """ match = pl.lit(True) for col, names in ((V_CALL, v), (D_CALL, d), (J_CALL, j)): if names: match = match & _segment_matches(col, names) out = df.filter(match if keep else ~match) return recompute_frequency(out)
[docs] def filter_by_sample(df: pl.DataFrame, other: pl.DataFrame, keep: bool = True, key: "tuple[str, ...]" = (JUNCTION_AA, V_CALL, J_CALL)) -> pl.DataFrame: """Keep or remove clonotypes by exact-key presence in another sample. Reimplements ``ApplySampleAsFilter`` / ``IntersectionClonotypeFilter``: build the key set of ``other`` and keep (or, with ``keep=False``, remove) the clonotypes of ``df`` whose key is present in it. Matching is an exact match on the ``key`` columns. Args: df: The clonotype frame to filter. other: The filter sample; only its ``key`` columns are used. keep: If ``True`` (default) keep clonotypes present in ``other``; if ``False`` remove them (legacy ``--negative``). key: Columns forming the match key (default ``("junction_aa", "v_call", "j_call")`` — legacy "strict"-style at the aa level). Returns: The filtered frame with ``frequency`` recomputed. """ key = list(key) keyset = other.select(key).unique() out = df.join(keyset, on=key, how="semi" if keep else "anti") return recompute_frequency(out)