From e2400d68adbc8d1053abf4726f7ca4b439935a5c Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Wed, 9 Sep 2026 10:54:01 +0200 Subject: [PATCH 1/2] perf: encode the residue-only part of a batch in one pass Encoding is a quarter of the prediction pipeline and every peptide walked through it alone, filling four small arrays a value at a time. The part that depends only on residue identity does not need that: the composition matrix and the one-hot residues are a table gather and a scatter over the batch, the positional block is the same gather at different rows, and the global vector is a sum along the batch's position axis. Bucketed encoding of 20,000 peptides: 0.1837 to 0.0252 ms per peptide, 7.3x, and core.predict on CPU 2,979 to 3,734 peptidoforms/s. encode_peptidoform stays the definition of a feature. Modifications are not reimplemented: they sit on a quarter of a typical peptide list and their placement carries a legacy quirk a shipped model was trained against, so _apply_modifications and its terminal counterpart are called on views into the batch arrays. Exactness comes from reusing that code rather than matching it. Two cases fall back to the per-peptide encoder, both divergences found by testing for them rather than hypotheticals: - Peptides under four residues. The reference's negative positional indices wrap around the sequence - with one residue, seq[seq_len - 2] is seq[-1], that same residue - and the batched form had skipped them. - Peptides carrying a residue the one-hot block has no slot for, such as selenocysteine, whose atoms the reference still reads from pyteomics. Feature layouts the route does not cover at all are declined by supports(): the rolling-sum matrix the four-branch model reads, and the collision cross section extras. Those keep the per-peptide path, as does vectorised_encoding=False on the dataset, which forces it for everything. tests/test_batch_features.py compares the two routes value for value and dtype for dtype over random batches at four windows, truncation, both termini, two modifications, either feature switch, the short and unusual residues above, and through the dataset in both modes. It needs no data files, so CI runs it. Also fixes a crash it surfaced. In _apply_composition_to_matrices the outer try guards KeyError and IndexError, but the inner branch - the one that strips isotope brackets, so C[13] and N[15] as TMT and SILAC labels carry - guarded only KeyError. A labelled modification beyond the padding window therefore raised IndexError where the unlabelled path warned and carried on. 225 tests pass, ruff clean. Co-Authored-By: Claude Opus 5 (1M context) --- deeplc/_batch_features.py | 218 +++++++++++++++++++++++++++++++++++ deeplc/_features.py | 6 + deeplc/data.py | 32 +++++ tests/test_batch_features.py | 209 +++++++++++++++++++++++++++++++++ 4 files changed, 465 insertions(+) create mode 100644 deeplc/_batch_features.py create mode 100644 tests/test_batch_features.py 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..bab0d26 --- /dev/null +++ b/tests/test_batch_features.py @@ -0,0 +1,209 @@ +""" +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) From 07cd9b5663ad66fbeb4240b6b0239dd3fae9d3a4 Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Wed, 9 Sep 2026 11:43:06 +0200 Subject: [PATCH 2/2] test: assert a CCS dataset never enters the batched encoder IM2Deep holds a CCS model trained against the pre-4.0.1 encoding and reaches DeepLC through exactly one call, `DeepLCDataset.from_psm_list(psm_list, add_ccs_features=True)`. Two things decline the batched route for that call, the CCS extras and the rolling sum, which `from_psm_list` leaves on. The test asserts the outcome rather than either reason, so removing one of them still fails here. Verified separately end to end: IM2Deep 2.0.2, unmodified, returns bit-identical CCS values on this branch and on the released 4.1.0, over twelve peptidoforms covering side-chain modifications, both termini, an isotope label, a short peptide and selenocysteine. Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_batch_features.py | 38 ++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/test_batch_features.py b/tests/test_batch_features.py index bab0d26..6114bde 100644 --- a/tests/test_batch_features.py +++ b/tests/test_batch_features.py @@ -207,3 +207,41 @@ def test_a_label_beyond_the_window_warns_instead_of_raising(): ) 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