diff --git a/deeplc/_batch_features.py b/deeplc/_batch_features.py new file mode 100644 index 0000000..437533b --- /dev/null +++ b/deeplc/_batch_features.py @@ -0,0 +1,218 @@ +""" +Batched feature encoding: the residue-only base for a whole batch in one pass. + +:func:`deeplc._features.encode_peptidoform` stays the definition of what a feature is. What +this module adds is a faster route to the same numbers for a batch of peptidoforms, and it +only accelerates the part that depends on residue identity alone: + +* the per-position atom composition matrix and the one-hot residue matrix are a table gather + and a scatter over the batch; +* the positional atom block reads the first few and last few residues of each peptide, which + is the same gather with different rows; +* the global vector is a sum along the batch's position axis. + +Modifications are not reimplemented here. They occur on a quarter of a typical peptide list, +and their placement has a legacy quirk that a shipped model was trained against, so they are +applied by calling :func:`deeplc._features._apply_modifications` and its terminal counterpart +on **views into the batch arrays**. The exactness therefore comes from reusing that code +rather than from matching it. + +Anything the batched route is not verified for falls back to +:func:`~deeplc._features.encode_peptidoform` per peptide, which keeps its warnings and its +exceptions rather than approximating them: peptides shorter than four residues, whose +negative positional indices wrap around the sequence in the reference encoder, and peptides +carrying a residue the one-hot block has no slot for, such as selenocysteine. Feature layouts +the route does not cover at all are rejected by :func:`supports`, and the caller keeps the +per-peptide path for those. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np +from psm_utils import Peptidoform +from pyteomics import mass + +from deeplc._features import ( + DEFAULT_DICT_AA, + DEFAULT_DICT_INDEX, + DEFAULT_DICT_INDEX_POS, + DEFAULT_POSITIONS, + DEFAULT_POSITIONS_NEG, + DEFAULT_POSITIONS_POS, + _apply_modifications, + _apply_terminal_modifications, + _terminal_composition, + encode_peptidoform, +) + +#: Residues the one-hot block has a slot for, in its own order. +RESIDUES: tuple[str, ...] = tuple(sorted(DEFAULT_DICT_AA, key=lambda key: DEFAULT_DICT_AA[key])) +CODE: dict[str, int] = {residue: index for index, residue in enumerate(RESIDUES)} +#: Row reserved for padding in the gather tables, holding zeros. +PAD = len(RESIDUES) +N_ATOMS = len(DEFAULT_DICT_INDEX) +_POSITIONS = DEFAULT_POSITIONS_POS | DEFAULT_POSITIONS_NEG +POS_ROWS = max(_POSITIONS) - min(_POSITIONS) + 1 +POS_OFFSET = min(_POSITIONS) + +#: Shortest peptide the batched route takes. Below four residues the reference encoder's +#: negative positional indices wrap around the sequence - with a single residue, +#: ``seq[seq_len - 2]`` is ``seq[-1]``, that same residue - and such peptides are rare enough +#: that reproducing the wrap in vectorised form is not worth the risk of getting it wrong. +MIN_LENGTH = 4 + + +def _gather_table(atom_index: dict[str, int]) -> np.ndarray: + """Atom counts per residue under one atom ordering, with a zero row for padding.""" + table = np.zeros((PAD + 1, N_ATOMS), dtype=np.float16) + for residue, row in CODE.items(): + for atom, count in mass.std_aa_comp[residue].items(): + column = atom_index.get(atom) + if column is not None: + table[row, column] = count + return table + + +COMPOSITION = _gather_table(DEFAULT_DICT_INDEX) +POSITIONAL = _gather_table(DEFAULT_DICT_INDEX_POS) +#: An unmodified peptidoform contributes nothing to the terminal block: the reference reads +#: the ``n_term`` and ``c_term`` properties, which are empty, and returns zeros. The backbone +#: termini are already inside the residue compositions. +TERMINI = np.zeros((2, N_ATOMS), dtype=np.float16) + + +def supports(add_ccs_features: bool, include_rolling_sum: bool) -> bool: + """ + Whether the batched route covers a feature layout. + + It does not build the rolling-sum matrix, which the four-branch model reads, and it does + not build the collision cross section extras, which need the precursor charge and a + handful of residue fractions. A caller asking for either keeps the per-peptide path. + """ + return not add_ccs_features and not include_rolling_sum + + +def _as_peptidoform(peptidoform: Peptidoform | str) -> Peptidoform: + """Parse a ProForma string if that is what a dataset holds, else pass it through.""" + return peptidoform if hasattr(peptidoform, "sequence") else Peptidoform(str(peptidoform)) + + +def encode_batch_features( + peptidoforms: Sequence[Peptidoform | str], + padding_length: int = 60, + add_terminal_composition: bool = True, + legacy_positional_deltas: bool = True, +) -> dict[str, np.ndarray]: + """ + Encode a batch of peptidoforms, gathering the base and patching the modified ones. + + Parameters + ---------- + peptidoforms + The peptidoforms to encode, parsed or as ProForma strings. + padding_length + Window the per-position matrices are padded or truncated to. + add_terminal_composition + Whether to append the N- and C-terminal group compositions to the global vector. + legacy_positional_deltas + Whether modification deltas go into the positional block the way versions before + 4.0.1 placed them, which every released model was trained against. + + Returns + ------- + dict of str to numpy.ndarray + ``matrix``, ``matrix_global`` and ``matrix_hc``, batched along the first axis and + equal value for value, dtype included, to stacking what + :func:`~deeplc._features.encode_peptidoform` returns for each peptidoform. + + """ + parsed = [_as_peptidoform(p) for p in peptidoforms] + n = len(parsed) + sequences = [p.sequence for p in parsed] + lengths = np.fromiter( + (min(len(sequence), padding_length) for sequence in sequences), dtype=np.int64, count=n + ) + reference_rows = { + row + for row, sequence in enumerate(sequences) + if len(sequence) < MIN_LENGTH or any(residue not in CODE for residue in sequence) + } + + residues = np.full((n, padding_length), PAD, dtype=np.int64) + for row, (sequence, length) in enumerate(zip(sequences, lengths, strict=True)): + if row not in reference_rows: + residues[row, :length] = [CODE[residue] for residue in sequence[:length]] + + matrix = COMPOSITION[residues] + one_hot = np.zeros((n, padding_length, len(DEFAULT_DICT_AA)), dtype=np.float16) + rows, columns = np.nonzero(residues != PAD) + one_hot[rows, columns, residues[rows, columns]] = 1.0 + + positional = np.zeros((n, POS_ROWS, N_ATOMS), dtype=np.float16) + for position in sorted(DEFAULT_POSITIONS_POS): + taken = lengths > position + positional[taken, position - POS_OFFSET] = POSITIONAL[residues[taken, position]] + for position in sorted(DEFAULT_POSITIONS_NEG): + taken = lengths + position >= 0 + positional[taken, position - POS_OFFSET] = POSITIONAL[ + residues[np.nonzero(taken)[0], (lengths + position)[taken]] + ] + + terminal = np.repeat(TERMINI[None, :, :], n, axis=0) if add_terminal_composition else None + + for row, peptidoform in enumerate(parsed): + if row in reference_rows: + continue + tokens = peptidoform.parsed_sequence + if any(token[1] is not None for token in tokens): + _apply_modifications( + matrix[row], + positional[row], + tokens, + int(lengths[row]), + DEFAULT_DICT_INDEX, + DEFAULT_DICT_INDEX_POS, + DEFAULT_POSITIONS, + legacy_positional_deltas, + ) + properties = peptidoform.properties + if properties.get("n_term") or properties.get("c_term"): + _apply_terminal_modifications( + matrix[row], + positional[row], + peptidoform, + int(lengths[row]), + DEFAULT_DICT_INDEX, + DEFAULT_DICT_INDEX_POS, + DEFAULT_POSITIONS, + legacy_positional_deltas, + ) + if terminal is not None: + terminal[row] = _terminal_composition(peptidoform, DEFAULT_DICT_INDEX) + + # float64, because the reference promotes when it appends the integer length to a float16 + # sum. The dataset casts to float32 afterwards either way, but the dtype is observable. + blocks = [ + matrix.sum(axis=1).astype(np.float64), + lengths[:, None].astype(np.float64), + positional.reshape(n, -1).astype(np.float64), + ] + if terminal is not None: + blocks.append(terminal.reshape(n, -1).astype(np.float64)) + matrix_global = np.concatenate(blocks, axis=1) + + for row in sorted(reference_rows): + features = encode_peptidoform( + parsed[row], + add_terminal_composition=add_terminal_composition, + padding_length=padding_length, + legacy_positional_deltas=legacy_positional_deltas, + include_rolling_sum=False, + ) + matrix[row] = features["matrix"] + one_hot[row] = features["matrix_hc"] + matrix_global[row] = features["matrix_global"] + + return {"matrix": matrix, "matrix_global": matrix_global, "matrix_hc": one_hot} diff --git a/deeplc/_features.py b/deeplc/_features.py index 6c566b2..0b69bda 100644 --- a/deeplc/_features.py +++ b/deeplc/_features.py @@ -349,6 +349,12 @@ def _apply_composition_to_matrices( except KeyError: warnings.warn(f"Ignoring atom {atom_comp} at pos {i}", stacklevel=2) continue + except IndexError: + # Same guard as the branch below, which this one was missing: a modification + # placed beyond the padding window has no row to add itself to. Reachable + # with an isotope-labelled modification, so a TMT or SILAC label, on a + # peptide longer than the window, where it raised instead of warning. + warnings.warn(f"Index error for atom {atom_comp} at pos {i}", stacklevel=2) except IndexError: warnings.warn(f"Index error for atom {atom_comp} at pos {i}", stacklevel=2) diff --git a/deeplc/data.py b/deeplc/data.py index 85b8e71..6774fd6 100644 --- a/deeplc/data.py +++ b/deeplc/data.py @@ -11,6 +11,7 @@ from psm_utils import Peptidoform, PSMList from torch.utils.data import Dataset, Subset +from deeplc._batch_features import encode_batch_features, supports from deeplc._features import encode_peptidoform _DatasetT = TypeVar("_DatasetT", bound=Dataset) @@ -30,6 +31,7 @@ def __init__( padding_length: int = 60, legacy_positional_deltas: bool = True, include_rolling_sum: bool = True, + vectorised_encoding: bool = True, ): """ Initialize the DeepLCDataset. @@ -73,6 +75,11 @@ def __init__( Whether to build the rolling-sum matrix. A convolutional trunk reads the per-position matrix directly and ignores this one, so building it is pure cost for such a model; False puts an empty array in its place. Default is True. + vectorised_encoding + Whether :meth:`encode_batch` may gather the residue-only part of a batch in one + pass instead of encoding peptide by peptide. It produces the same values, dtype + included, and falls back per peptide for anything it does not cover; False forces + the per-peptide encoder for everything. Default is True. Raises ------ @@ -88,6 +95,7 @@ def __init__( self.padding_length = padding_length self.legacy_positional_deltas = legacy_positional_deltas self.include_rolling_sum = include_rolling_sum + self.vectorised_encoding = vectorised_encoding if self.target_retention_times is not None and len(self.target_retention_times) != len( self.peptidoforms ): @@ -126,6 +134,7 @@ def variant(self, indices: Sequence[int], padding_length: int) -> DeepLCDataset: padding_length=padding_length, legacy_positional_deltas=self.legacy_positional_deltas, include_rolling_sum=self.include_rolling_sum, + vectorised_encoding=self.vectorised_encoding, ) #: Arrays the encoder returns, in the order the models take them. @@ -160,6 +169,25 @@ def encode_batch(self, indices: Sequence[int]) -> tuple[torch.Tensor, ...]: indices = list(indices) if not indices: raise ValueError("No indices to encode.") + + if self.vectorised_encoding and supports(self.add_ccs_features, self.include_rolling_sum): + # The residue-only part of the batch is a table gather; modifications and + # anything the gather does not cover go through the per-peptide encoder, so the + # values are the same either way. See deeplc._batch_features. + features = encode_batch_features( + [self.peptidoforms[index] for index in indices], + padding_length=self.padding_length, + add_terminal_composition=self.add_terminal_composition, + legacy_positional_deltas=self.legacy_positional_deltas, + ) + empty = np.zeros((len(indices), 0, features["matrix"].shape[-1]), dtype=np.float32) + return ( + torch.from_numpy(features["matrix"].astype(np.float32)), + torch.from_numpy(empty), + torch.from_numpy(features["matrix_global"].astype(np.float32)), + torch.from_numpy(features["matrix_hc"].astype(np.float32)), + ) + buffers: list[np.ndarray] | None = None for row, index in enumerate(indices): features = encode_peptidoform(self.peptidoforms[index], **self._encode_kwargs()) @@ -199,6 +227,7 @@ def from_psm_list( padding_length: int = 60, legacy_positional_deltas: bool = True, include_rolling_sum: bool = True, + vectorised_encoding: bool = True, ) -> DeepLCDataset: """ Create a DeepLCDataset from a PSMList. @@ -236,6 +265,8 @@ def from_psm_list( include_rolling_sum Whether to build the rolling-sum matrix. Models with a convolutional trunk ignore it, and :func:`deeplc.core.predict` sets this from the model. + vectorised_encoding + Whether a batch may be encoded in one pass; see the class docstring. Returns ------- @@ -257,6 +288,7 @@ def from_psm_list( padding_length=padding_length, legacy_positional_deltas=legacy_positional_deltas, include_rolling_sum=include_rolling_sum, + vectorised_encoding=vectorised_encoding, ) diff --git a/tests/test_batch_features.py b/tests/test_batch_features.py new file mode 100644 index 0000000..6114bde --- /dev/null +++ b/tests/test_batch_features.py @@ -0,0 +1,247 @@ +""" +The batched encoder must agree with the per-peptide one, value for value and dtype for dtype. + +:func:`deeplc._features.encode_peptidoform` is the definition of a feature; everything here +compares the batched route against it. Two of these cases were divergences found that way +during development rather than hypotheticals: selenocysteine, whose atoms the reference reads +from pyteomics even though the one-hot block has no slot for it, and peptides shorter than +four residues, where the reference's negative positional indices wrap around the sequence. +Both now take the per-peptide fallback, which is why they agree. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from psm_utils import Peptidoform + +from deeplc._batch_features import MIN_LENGTH, encode_batch_features, supports +from deeplc._features import encode_peptidoform +from deeplc.data import DeepLCDataset + +FEATURES = ("matrix", "matrix_global", "matrix_hc") +RESIDUES = "ACDEFGHIKLMNPQRSTVWY" +MODIFICATIONS = ("", "[UNIMOD:4]", "[UNIMOD:35]", "[UNIMOD:1]", "[UNIMOD:21]") + + +def assert_agrees(peptidoforms, padding_length=20, terminal=True, legacy=True): + """Compare the batched route against stacking the per-peptide encoder.""" + parsed = [Peptidoform(p) if isinstance(p, str) else p for p in peptidoforms] + one_by_one = [ + encode_peptidoform( + p, + add_terminal_composition=terminal, + padding_length=padding_length, + legacy_positional_deltas=legacy, + include_rolling_sum=False, + ) + for p in parsed + ] + batched = encode_batch_features( + parsed, + padding_length=padding_length, + add_terminal_composition=terminal, + legacy_positional_deltas=legacy, + ) + for key in FEATURES: + reference = np.stack([features[key] for features in one_by_one]) + assert batched[key].shape == reference.shape, key + assert batched[key].dtype == reference.dtype, key + assert np.array_equal(batched[key], reference), key + + +def random_peptidoforms(count, seed, min_length=4, max_length=18): + """A spread of peptides with modifications at a rate like a real peptide list.""" + rng = np.random.RandomState(seed) + out = [] + for _ in range(count): + length = int(rng.randint(min_length, max_length + 1)) + residues = list(rng.choice(list(RESIDUES), size=length)) + if rng.rand() < 0.3: + site = int(rng.randint(0, length)) + residues[site] = residues[site] + str(rng.choice(MODIFICATIONS[1:])) + sequence = "".join(residues) + if rng.rand() < 0.1: + sequence = "[UNIMOD:1]-" + sequence + if rng.rand() < 0.05: + sequence = sequence + "-[UNIMOD:2]" + out.append(Peptidoform(f"{sequence}/2")) + return out + + +def test_agrees_on_a_random_batch(): + """The everyday case: a few hundred peptides, a third of them modified.""" + assert_agrees(random_peptidoforms(400, seed=0)) + + +@pytest.mark.parametrize("padding_length", [8, 20, 30, 60]) +def test_agrees_at_every_window(padding_length): + """Windows vary per batch because prediction is length-bucketed.""" + assert_agrees(random_peptidoforms(120, seed=1), padding_length=padding_length) + + +def test_agrees_when_peptides_are_truncated(): + """A peptide longer than the window is truncated, modifications included.""" + assert_agrees( + ["A" * 30 + "/2", "A" * 24 + "C[UNIMOD:4]" + "AAAAA/2", "PEPTIDEK/2"], + padding_length=20, + ) + + +@pytest.mark.parametrize( + "sequence", + [ + "PEPTIDEK/2", + "C[UNIMOD:4]EPTIDEC[UNIMOD:4]/2", + "PEPTC[UNIMOD:4]IDEM[UNIMOD:35]K/2", + "[UNIMOD:1]-PEPTIDEK/2", + "PEPTIDEK-[UNIMOD:2]/2", + "[UNIMOD:1]-PEPTIDEK-[UNIMOD:2]/2", + "PEPUIDEK/2", + "PEPOIDEK/2", + "AAAA/2", + ], + ids=[ + "plain", + "modified first and last", + "two modifications", + "n-terminal", + "c-terminal", + "both termini", + "selenocysteine", + "pyrrolysine", + "shortest on the fast path", + ], +) +def test_agrees_on_one_peptide(sequence): + """Each of these is a case the batched route has to hand over or handle exactly.""" + assert_agrees([sequence]) + + +@pytest.mark.parametrize("length", [1, 2, 3]) +def test_agrees_below_the_fast_path_minimum(length): + """ + Short peptides go to the per-peptide encoder, so its wrap-around behaviour is preserved. + + With one residue the reference reads ``seq[seq_len - 2]``, which Python resolves to the + last residue rather than to nothing; the batched route would have skipped it. + """ + assert length < MIN_LENGTH + assert_agrees(["A" * length + "/2", "PEPTIDEK/2"]) + + +def test_agrees_without_terminal_composition_and_without_legacy_deltas(): + """Both feature-layout switches the route supports.""" + batch = random_peptidoforms(60, seed=2) + assert_agrees(batch, terminal=False) + assert_agrees(batch, legacy=False) + + +def test_unknown_residues_raise_the_same_way(): + """An ambiguous residue has no composition, and both routes say so identically.""" + for sequence in ("PEPBIDEK/2", "PEPZIDEK/2", "PEPXIDEK/2"): + with pytest.raises(KeyError): + encode_peptidoform( + Peptidoform(sequence), add_terminal_composition=True, padding_length=20 + ) + with pytest.raises(KeyError): + encode_batch_features([Peptidoform(sequence)], padding_length=20) + + +def test_layouts_the_route_does_not_cover_are_declined(): + """The rolling sum and the collision cross section extras keep the per-peptide path.""" + assert supports(add_ccs_features=False, include_rolling_sum=False) + assert not supports(add_ccs_features=True, include_rolling_sum=False) + assert not supports(add_ccs_features=False, include_rolling_sum=True) + + +def test_dataset_batch_matches_item_by_item_either_way(): + """ + Through the dataset, the switch must not change what comes out. + + ``encode_batch`` picks the route; ``__getitem__`` always uses the per-peptide encoder, so + comparing the two covers the wiring as well as the encoder. + """ + peptides = [str(p) for p in random_peptidoforms(50, seed=3)] + for vectorised in (True, False): + dataset = DeepLCDataset( + peptides, + add_terminal_composition=True, + padding_length=20, + include_rolling_sum=False, + vectorised_encoding=vectorised, + ) + batch = dataset.encode_batch(range(len(peptides))) + for position in range(4): + stacked = torch.stack([dataset[i][0][position] for i in range(len(peptides))]) + assert torch.equal(batch[position], stacked), (vectorised, position) + + +def test_dataset_keeps_the_per_peptide_path_for_the_rolling_sum(): + """A model that reads the rolling sum still gets it, batched the old way.""" + peptides = [str(p) for p in random_peptidoforms(20, seed=4)] + dataset = DeepLCDataset( + peptides, add_terminal_composition=True, padding_length=20, include_rolling_sum=True + ) + batch = dataset.encode_batch(range(len(peptides))) + assert batch[1].shape[1] > 0 + stacked = torch.stack([dataset[i][0][1] for i in range(len(peptides))]) + assert torch.equal(batch[1], stacked) + + +def test_a_label_beyond_the_window_warns_instead_of_raising(): + """ + An isotope-labelled modification past the padding window must not raise. + + The unlabelled path already warned and carried on; the branch that strips isotope + brackets, so ``C[13]`` and ``N[15]`` as TMT and SILAC labels carry, did not, and raised + IndexError from inside the encoder instead. Both routes go through the same helper, so + one test covers both. + """ + sequence = "A" * 24 + "K[UNIMOD:737]" + "AAAAA/2" # TMT6plex, which carries 13C + with pytest.warns(UserWarning): + reference = encode_peptidoform( + Peptidoform(sequence), add_terminal_composition=True, padding_length=20, + legacy_positional_deltas=True, include_rolling_sum=False, + ) + assert np.isfinite(reference["matrix"]).all() + assert_agrees([sequence], padding_length=20) + + +def test_a_ccs_dataset_keeps_the_per_peptide_path(): + """ + The collision cross section layout must never reach the batched route. + + IM2Deep holds a CCS model trained against the pre-4.0.1 encoding and reaches DeepLC + only through ``DeepLCDataset.from_psm_list(psm_list, add_ccs_features=True)``, so this + is the call that has to keep the per-peptide encoder. Two things decline it - the CCS + extras and the rolling sum, which ``from_psm_list`` leaves on - and this asserts the + outcome rather than either reason, so removing one of them still fails here. + """ + from psm_utils import PSM, PSMList + + import deeplc.data as data_module + + peptidoforms = ["AC[UNIMOD:4]DEK/2", "[UNIMOD:737]-PEPTIDEK/2", "PEPTM[UNIMOD:35]IDEKR/3"] + psm_list = PSMList( + psm_list=[ + PSM(peptidoform=Peptidoform(p), spectrum_id=str(i), retention_time=float(i)) + for i, p in enumerate(peptidoforms) + ] + ) + dataset = DeepLCDataset.from_psm_list(psm_list, add_ccs_features=True) + assert dataset.vectorised_encoding, "the switch is on; the layout is what declines it" + + calls = [] + original = data_module.encode_batch_features + data_module.encode_batch_features = lambda *args, **kwargs: calls.append(1) + try: + batch = dataset.encode_batch(range(len(peptidoforms))) + finally: + data_module.encode_batch_features = original + + assert not calls, "a CCS dataset must not enter the batched encoder" + for position in range(4): + stacked = torch.stack([dataset[i][0][position] for i in range(len(peptidoforms))]) + assert torch.equal(batch[position], stacked), position