Sequence probability (data-driven Pgen)

Data-driven CDR3 generation/occurrence probability — the precursor- frequency / publicness axis — a model fit once on an external reference repertoire (it replaced an earlier OLGA/SONIA runtime path, since removed), so log_pgen(seq) is a fast, calibrated score for "how generatable / common is this CDR3" — lower = rarer precursor / more private.

On the B1-2 pilot the default k-mer model recovers the same alpha-chain publicness signal as OLGA (TRAV12-2 vs other TRAV AUROC ≈ 0.64, matching OLGA's 0.65–0.67) — with no GPL and no runtime dependency beyond numpy.

Two backends, one interface

Both implement SequenceProbabilityModel (fit / log_prob / save / load):

Backend	Deps	Notes
KmerProbabilityModel (default)	numpy only	Order-k Markov model over CDR3 AAs (length captured via an EOS symbol; add-alpha smoothing). Ships a default for each chain.
TCRpegProbabilityModel	pip install tcrsift[tcrpeg]	Wraps TCRpeg (autoregressive, PyTorch). Heavier, better-calibrated.

Shipped defaults & the GPL boundary

The default k-mer models (tcrsift/refseqs/kmer_background_{alpha,beta}.npz, ~300 KB each) are fit offline at build time on OLGA-generated synthetic repertoires (scripts/generate_kmer_background.py). OLGA (GPL-3.0) is used only to produce training sequences; tcrsift never imports it at runtime, so the package stays Apache-2.0.

To retrain on a different reference, fit and save your own:

from tcrsift.seqprob import KmerProbabilityModel
model = KmerProbabilityModel(order=3, chain="beta").fit(my_reference_cdr3s)
model.save("my_beta_background.npz")

Usage

from tcrsift.seqprob import score_log_pgen, load_background_model

# default shipped k-mer background:
clones["log_pgen_beta"] = score_log_pgen(clones, chain="beta")

# explicit model / TCRpeg backend:
clones["log_pgen_alpha"] = score_log_pgen(clones, chain="alpha", backend="tcrpeg")

CLI:

tcrsift log-pgen clones.csv -o clones_pgen.csv --chain both          # k-mer
tcrsift log-pgen clones.csv -o out.csv --backend tcrpeg --chain beta # TCRpeg

The resulting log_pgen_<chain> column plugs directly into the in-silico filter (insilico_filter) as a Ppost^low-style predicate.

seqprob

Data-driven CDR3 sequence-probability models — the publicness axis.

The data-driven, dependency-light publicness axis (it replaced an earlier OLGA/SONIA runtime path, since removed) for the precursor-frequency / publicness measure. Instead of a fixed, allele-masked GPL prior, a background generation/occurrence model is fit once on an external reference repertoire and reused; log_pgen(seq) is then a fast, dependency-light, calibrated score for "how generatable / common is this CDR3" — lower = more private / rarer precursor.

Two interchangeable backends behind one :class:SequenceProbabilityModel interface:

• :class:KmerProbabilityModel — an order-k Markov model over CDR3 amino acids (numpy-only, no GPL, the default). The shipped default models in :mod:tcrsift.refseqs are fit offline on OLGA-generated synthetic repertoires (OLGA used once at build time to produce training sequences — never at runtime, so tcrsift stays Apache-2.0).
• :class:TCRpegProbabilityModel — wraps TCRpeg (Jiang & Li 2023), an autoregressive deep model. Optional extra: pip install tcrsift[tcrpeg].

Both are trained on an external reference (not the experiment's own clones) so the probability is a genuine background, not circular with the selection target.

SequenceProbabilityModel

Bases: ABC

A fittable per-sequence log-probability model over CDR3 strings.

Source code in tcrsift/seqprob.py

class SequenceProbabilityModel(abc.ABC):
    """A fittable per-sequence log-probability model over CDR3 strings."""

    @abc.abstractmethod
    def fit(self, sequences: Iterable[str]) -> SequenceProbabilityModel:
        """Train on an iterable of CDR3 amino-acid strings. Returns self."""

    @abc.abstractmethod
    def log_prob(self, sequences: Iterable[str]) -> np.ndarray:
        """Natural-log probability per sequence (NaN for unscorable input)."""

    @abc.abstractmethod
    def save(self, path) -> None:
        """Persist the fitted model to ``path``."""

    @classmethod
    @abc.abstractmethod
    def load(cls, path) -> SequenceProbabilityModel:
        """Load a model previously written by :meth:`save`."""

fit abstractmethod

fit(sequences: Iterable[str]) -> SequenceProbabilityModel

Train on an iterable of CDR3 amino-acid strings. Returns self.

Source code in tcrsift/seqprob.py

@abc.abstractmethod
def fit(self, sequences: Iterable[str]) -> SequenceProbabilityModel:
    """Train on an iterable of CDR3 amino-acid strings. Returns self."""

log_prob abstractmethod

log_prob(sequences: Iterable[str]) -> np.ndarray

Natural-log probability per sequence (NaN for unscorable input).

Source code in tcrsift/seqprob.py

@abc.abstractmethod
def log_prob(self, sequences: Iterable[str]) -> np.ndarray:
    """Natural-log probability per sequence (NaN for unscorable input)."""

save abstractmethod

save(path) -> None

Persist the fitted model to path.

Source code in tcrsift/seqprob.py

@abc.abstractmethod
def save(self, path) -> None:
    """Persist the fitted model to ``path``."""

load abstractmethod classmethod

load(path) -> SequenceProbabilityModel

Load a model previously written by :meth:save.

Source code in tcrsift/seqprob.py

@classmethod
@abc.abstractmethod
def load(cls, path) -> SequenceProbabilityModel:
    """Load a model previously written by :meth:`save`."""

KmerProbabilityModel

Bases: SequenceProbabilityModel

Order-k Markov model over CDR3 amino acids (numpy-only).

log P(CDR3) = Σ_i log P(a_i | a_{i-k} … a_{i-1}) with the sequence padded by order BOS sentinels and terminated by EOS, so both the composition and the length are captured. Add-alpha (Laplace) smoothing keeps unseen contexts from giving -inf.

Parameters are a dense (N_SYM**order, N_SYM) log-probability table, compact enough to ship: the shipped defaults are order 2 (~20-32 KB float32 per chain); order 3 would be ~1 MB.

Source code in tcrsift/seqprob.py

class KmerProbabilityModel(SequenceProbabilityModel):
    """Order-``k`` Markov model over CDR3 amino acids (numpy-only).

    ``log P(CDR3) = Σ_i log P(a_i | a_{i-k} … a_{i-1})`` with the sequence
    padded by ``order`` BOS sentinels and terminated by EOS, so both the
    composition *and* the length are captured. Add-``alpha`` (Laplace)
    smoothing keeps unseen contexts from giving ``-inf``.

    Parameters are a dense ``(N_SYM**order, N_SYM)`` log-probability table,
    compact enough to ship: the shipped defaults are order 2 (~20-32 KB
    float32 per chain); order 3 would be ~1 MB.
    """

    def __init__(self, *, order: int = 2, alpha: float = 1.0, chain: str = ""):
        if order < 1:
            raise ValueError(f"order must be >= 1, got {order}")
        self.order = int(order)
        self.alpha = float(alpha)
        self.chain = chain
        self.n_train = 0
        self.n_contexts = N_SYM**self.order
        self._logp: np.ndarray | None = None  # (n_contexts, N_SYM)

    # -- context id helpers ------------------------------------------------
    def _context_id(self, ctx: list[int]) -> int:
        cid = 0
        for s in ctx:
            cid = cid * N_SYM + s
        return cid

    def fit(self, sequences: Iterable[str]) -> KmerProbabilityModel:
        counts = np.zeros((self.n_contexts, N_SYM), dtype=np.float64)
        n = 0
        skipped = 0
        for seq in sequences:
            ids = _encode(seq)
            if ids is None:
                skipped += 1
                continue
            padded = [BOS] * self.order + ids + [EOS]
            for i in range(self.order, len(padded)):
                cid = self._context_id(padded[i - self.order:i])
                counts[cid, padded[i]] += 1.0
            n += 1
        if n == 0:
            raise ValueError("KmerProbabilityModel.fit: no scorable sequences")
        counts += self.alpha
        totals = counts.sum(axis=1, keepdims=True)
        self._logp = np.log(counts / totals).astype(np.float32)
        self.n_train = n
        if skipped:
            logger.info(
                "KmerProbabilityModel.fit: skipped %d non-AA sequences", skipped
            )
        return self

    def _log_prob_one(self, seq: str) -> float:
        ids = _encode(seq)
        if ids is None:
            return float("nan")
        padded = [BOS] * self.order + ids + [EOS]
        lp = 0.0
        for i in range(self.order, len(padded)):
            cid = self._context_id(padded[i - self.order:i])
            lp += float(self._logp[cid, padded[i]])
        return lp

    def log_prob(self, sequences: Iterable[str]) -> np.ndarray:
        if self._logp is None:
            raise RuntimeError("KmerProbabilityModel is not fitted")
        return np.array([self._log_prob_one(s) for s in sequences], dtype=float)

    def save(self, path) -> None:
        if self._logp is None:
            raise RuntimeError("KmerProbabilityModel is not fitted")
        np.savez_compressed(
            path,
            logp=self._logp,
            order=np.int64(self.order),
            alpha=np.float64(self.alpha),
            n_train=np.int64(self.n_train),
            chain=np.array(self.chain),
        )

    @classmethod
    def load(cls, path) -> KmerProbabilityModel:
        with np.load(path, allow_pickle=False) as data:
            model = cls(
                order=int(data["order"]),
                alpha=float(data["alpha"]),
                chain=str(data["chain"]),
            )
            model._logp = data["logp"].astype(np.float32)
            model.n_train = int(data["n_train"])
        if model._logp.shape != (model.n_contexts, N_SYM):
            raise ValueError(
                f"loaded k-mer table shape {model._logp.shape} != expected "
                f"{(model.n_contexts, N_SYM)} for order {model.order}"
            )
        return model

TCRpegProbabilityModel

Bases: SequenceProbabilityModel

TCRpeg-backed CDR3 probability (optional [tcrpeg] extra).

Wraps the autoregressive TCRpeg model (Jiang & Li 2023). Heavier (PyTorch) but better-calibrated than the k-mer Markov model. Trained on the same external reference. Lazy import; raises :class:ImportError with an install hint when the extra is missing.

Source code in tcrsift/seqprob.py

class TCRpegProbabilityModel(SequenceProbabilityModel):
    """TCRpeg-backed CDR3 probability (optional ``[tcrpeg]`` extra).

    Wraps the autoregressive TCRpeg model (Jiang & Li 2023). Heavier
    (PyTorch) but better-calibrated than the k-mer Markov model. Trained on
    the same external reference. Lazy import; raises :class:`ImportError`
    with an install hint when the extra is missing.
    """

    _INSTALL_HINT = (
        "TCRpeg (+ torch) is required for the TCRpeg sequence-probability "
        "backend but is not installed. Install with:\n\n"
        "    pip install tcrsift[tcrpeg]\n\n"
        "Or use the numpy-only KmerProbabilityModel (the default backend)."
    )

    def __init__(
        self,
        *,
        max_length: int = 30,
        embedding_size: int = 32,
        hidden_size: int = 64,
        num_layers: int = 1,
        device: str = "cpu",
        epochs: int = 20,
        batch_size: int = 1000,
        lr: float = 1e-3,
        chain: str = "",
        embedding_path: str | None = None,
    ):
        self.max_length = max_length
        self.embedding_size = embedding_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.device = device
        self.epochs = epochs
        self.batch_size = batch_size
        self.lr = lr
        self.chain = chain
        self.embedding_path = embedding_path
        self.n_train = 0
        self._model = None

    @staticmethod
    def available() -> bool:
        import importlib.util

        return (
            importlib.util.find_spec("tcrpeg") is not None
            and importlib.util.find_spec("torch") is not None
        )

    def _require(self) -> None:
        if not self.available():
            raise ImportError(self._INSTALL_HINT)

    def _resolve_embedding_path(self) -> str:
        """Absolute path to the AA embedding TCRpeg needs.

        TCRpeg's default ``embedding_path`` is relative to the CWD; the file
        actually ships inside the installed ``tcrpeg/data/`` dir, so resolve
        to that bundled copy unless the caller gave an explicit path.
        """
        if self.embedding_path is not None:
            return self.embedding_path
        import os as _os

        import tcrpeg  # pylint: disable=import-error

        return _os.path.join(
            _os.path.dirname(tcrpeg.__file__),
            "data", f"embedding_{self.embedding_size}.txt",
        )

    def _new_model(self, sequences: list[str] | None = None):
        from tcrpeg.TCRpeg import TCRpeg  # pylint: disable=import-error

        model = TCRpeg(
            max_length=self.max_length,
            embedding_size=self.embedding_size,
            hidden_size=self.hidden_size,
            num_layers=self.num_layers,
            device=self.device,
            load_data=sequences is not None,
            path_train=sequences,
            embedding_path=self._resolve_embedding_path(),
        )
        model.create_model()
        return model

    def fit(self, sequences: Iterable[str]) -> TCRpegProbabilityModel:
        self._require()
        seqs = [s for s in sequences if isinstance(s, str) and s]
        if not seqs:
            raise ValueError("TCRpegProbabilityModel.fit: no scorable sequences")
        self._model = self._new_model(seqs)
        self._model.train_tcrpeg(
            epochs=self.epochs, batch_size=self.batch_size, lr=self.lr,
        )
        self.n_train = len(seqs)
        return self

    def log_prob(self, sequences: Iterable[str]) -> np.ndarray:
        self._require()
        if self._model is None:
            raise RuntimeError("TCRpegProbabilityModel is not fitted")
        seqs = list(sequences)
        scorable = [isinstance(s, str) and bool(s) and _encode(s) is not None
                    for s in seqs]
        out = np.full(len(seqs), np.nan)
        good = [s for s, ok in zip(seqs, scorable) if ok]
        if good:
            # sampling_tcrpeg already returns natural-log probabilities.
            logs = np.asarray(self._model.sampling_tcrpeg(good), dtype=float)
            j = 0
            for i, ok in enumerate(scorable):
                if ok:
                    out[i] = logs[j]
                    j += 1
        return out

    def save(self, path) -> None:
        self._require()
        if self._model is None:
            raise RuntimeError("TCRpegProbabilityModel is not fitted")
        self._model.save(str(path))

    @classmethod
    def load(cls, path, **kwargs) -> TCRpegProbabilityModel:
        obj = cls(**kwargs)
        obj._require()
        from tcrpeg.TCRpeg import TCRpeg  # pylint: disable=import-error

        model = TCRpeg(
            max_length=obj.max_length,
            embedding_size=obj.embedding_size,
            hidden_size=obj.hidden_size,
            num_layers=obj.num_layers,
            device=obj.device,
            embedding_path=obj._resolve_embedding_path(),
        )
        model.create_model(load=True, path=str(path))
        obj._model = model
        return obj

load_background_model

load_background_model(chain: str = 'beta', backend: str = 'kmer', role: str = 'ppost') -> SequenceProbabilityModel

Load (and cache) a shipped default background model.

role is "ppost" (default — fit on an observed repertoire, the post-selection publicness measure) or "pgen" (fit on an OLGA-generated reference, pre-selection generation probability). Only the "kmer" backend ships defaults. Role-pure: raises :class:FileNotFoundError when the requested role isn't shipped for that chain (e.g. no observed-α ppost) — it never silently returns the Pgen model in place of Ppost. Callers decide how to degrade.

Source code in tcrsift/seqprob.py

def load_background_model(
    chain: str = "beta", backend: str = "kmer", role: str = "ppost",
) -> SequenceProbabilityModel:
    """Load (and cache) a shipped default background model.

    ``role`` is ``"ppost"`` (default — fit on an *observed* repertoire, the
    post-selection publicness measure) or ``"pgen"`` (fit on an
    OLGA-generated reference, pre-selection generation probability). Only the
    ``"kmer"`` backend ships defaults. **Role-pure**: raises
    :class:`FileNotFoundError` when the requested role isn't shipped for that
    chain (e.g. no observed-α ppost) — it never silently returns the Pgen
    model in place of Ppost. Callers decide how to degrade.
    """
    chain = chain.lower()
    key = (chain, backend, role)
    if key in _MODEL_CACHE:
        return _MODEL_CACHE[key]
    path = _default_model_path(chain, backend, role)
    if not path.is_file():
        raise FileNotFoundError(
            f"no shipped {backend} {role} model for chain {chain!r} at {path}"
        )
    model = BACKENDS[backend].load(str(path))
    _MODEL_CACHE[key] = model
    return model

score_log_pgen

score_log_pgen(df: DataFrame, *, chain: str = 'beta', cdr3_col: str | None = None, backend: str = 'kmer', model: SequenceProbabilityModel | None = None, out_col: str = 'log_pgen') -> pd.Series

Per-clone log Pgen (generated-repertoire background). See :func:score_log_prob.

Source code in tcrsift/seqprob.py

def score_log_pgen(
    df: pd.DataFrame,
    *,
    chain: str = "beta",
    cdr3_col: str | None = None,
    backend: str = "kmer",
    model: SequenceProbabilityModel | None = None,
    out_col: str = "log_pgen",
) -> pd.Series:
    """Per-clone log Pgen (generated-repertoire background). See
    :func:`score_log_prob`."""
    return score_log_prob(
        df, chain=chain, cdr3_col=cdr3_col, backend=backend, role="pgen",
        model=model, out_col=out_col,
    )