From b12257d807c0df3ae89219953a3f17f4d478a682 Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Mon, 7 Sep 2026 11:05:35 +0200 Subject: [PATCH 1/4] perf: predict short peptides in a window that fits them Every peptide was encoded and convolved in a 60-position window whatever its length, and the median peptide is 16 residues, so most of the trunk's arithmetic ran on padding that the mask removes again before pooling. On CPU that was the whole cost of prediction: 1,780 peptidoforms/s on eight threads, of which the forward pass was about 85 %. The trunk masks its output by the true residue count and the encoded features do not depend on the window at all, so the only route from padding to a valid position is the convolutions. Each reaches (kernel - 1) // 2 * dilation positions and the reaches add, which the model now reports as ``padding_reach``: 4 for the shipped architecture, two pointwise stem layers plus two convolutions of width five. A batch encoded in a window of its longest peptide plus that reach therefore has to give the same answer, and it does - over 50,000 peptides the largest difference was 1.2e-4 min, which is float32 noise. ``predict`` now sorts by length and gives each chunk its own window. The property is checked for the architecture and asserted end to end in tests/test_flexcnn.py, and the path is skipped for models that pool or stride across positions (they report no reach) and when the longest peptide already fills the window. CPU, 20,000 peptidoforms, including encoding: threads fixed 60 bucketed speedup 4 1,465 pf/s 3,115 pf/s 2.13x 8 1,780 pf/s 3,702 pf/s 2.08x 32 2,424 pf/s 5,260 pf/s 2.17x On GPU the forward pass was never the cost and the numbers are unchanged. Pass length_buckets=False to force one pass in the dataset's own window. Co-Authored-By: Claude Opus 5 (1M context) --- deeplc/_architecture.py | 27 ++++++++++++ deeplc/_model_ops.py | 92 ++++++++++++++++++++++++++++++++++++++--- deeplc/data.py | 28 +++++++++++++ tests/test_flexcnn.py | 55 ++++++++++++++++++++++++ 4 files changed, 197 insertions(+), 5 deletions(-) diff --git a/deeplc/_architecture.py b/deeplc/_architecture.py index 1fb8602..920199d 100644 --- a/deeplc/_architecture.py +++ b/deeplc/_architecture.py @@ -934,6 +934,33 @@ def forward( del x_atom_sum # the fused trunk reads x_atom directly return self.head(self.encoder(x_atom, x_global, x_one_hot), task_idx) + @property + def padding_reach(self) -> int | None: + """ + How far a valid position can see across the right edge of the encoding window. + + The trunk masks its output by the true residue count and pools over that mask, and + the features themselves do not depend on the window, so the only way padding can + reach a valid position is through the convolutions. Each convolution of width ``k`` + and dilation ``d`` reaches ``(k - 1) // 2 * d`` positions, and the reaches add up. + A batch encoded in a window of its longest peptide plus this many positions + therefore predicts exactly what the full window predicts, which for a median + peptide of sixteen residues is a fraction of the sixty positions used otherwise. + + Returns None when the trunk pools or strides across positions, because then the + mask no longer lines up position by position and the argument does not hold. + + """ + reach = 0 + for module in self.encoder.modules(): + if isinstance(module, (nn.MaxPool1d, nn.AvgPool1d)): + return None + if isinstance(module, nn.Conv1d): + if module.stride[0] != 1: + return None + reach += ((module.kernel_size[0] - 1) // 2) * module.dilation[0] + return reach + def add_task_head( self, targets: torch.Tensor | None = None, init_from: int | None = None ) -> int: diff --git a/deeplc/_model_ops.py b/deeplc/_model_ops.py index 15e04dc..c63d48c 100644 --- a/deeplc/_model_ops.py +++ b/deeplc/_model_ops.py @@ -7,6 +7,7 @@ from os import PathLike from pathlib import Path +import numpy as np import torch from rich.progress import ( BarColumn, @@ -250,8 +251,16 @@ def predict( num_threads: int | None = None, show_progress: bool = True, task_idx: Sequence[int] | None = None, + length_buckets: bool = True, ) -> torch.Tensor: - """Predict using the model for the given dataset.""" + """ + Predict using the model for the given dataset. + + ``length_buckets`` runs length-sorted chunks in a window that fits them rather than + padding every peptide to the model's full window; see :func:`_length_buckets`. It is + exact for models that report a ``padding_reach`` and ignored for those that do not. + Set it to False to force one pass over the data in the dataset's own window. + """ # ``task_idx`` selects which LC setups a multitask model evaluates. Without # it a model trained on thousands of setups returns a column per setup: at # 6,543 setups and a million peptides that output alone is tens of gigabytes, @@ -260,11 +269,84 @@ def predict( torch.set_num_threads(num_threads or torch.get_num_threads()) device = device or ("cuda" if torch.cuda.is_available() else "cpu") model = load_model(model, device) - data_loader = DataLoader(data, batch_size=batch_size, shuffle=False, num_workers=num_workers) - predictions = _predict_epoch( - model, data_loader, device, show_progress=show_progress, task_idx=task_idx + + buckets = _length_buckets(model, data, batch_size) if length_buckets else None + if buckets is None: + data_loader = DataLoader( + data, batch_size=batch_size, shuffle=False, num_workers=num_workers + ) + predictions = _predict_epoch( + model, data_loader, device, show_progress=show_progress, task_idx=task_idx + ) + return predictions.cpu().detach() + + out: torch.Tensor | None = None + for indices, subset in buckets: + part = _predict_epoch( + model, + DataLoader(subset, batch_size=batch_size, shuffle=False, num_workers=num_workers), + device, + show_progress=show_progress, + task_idx=task_idx, + ).cpu() + if out is None: + out = torch.empty((len(data), part.shape[1]), dtype=part.dtype) + out[indices] = part + if out is None: + raise ValueError("Dataset is empty — nothing to predict.") + return out.detach() + + +def _residue_count(peptidoform: object) -> int: + """ + Residues in a peptidoform, whether it arrives parsed or as a ProForma string. + + A dataset may hold either. The string form cannot be counted by its length, since + modifications and the charge state are part of it, so it is parsed once here rather + than per encoded item. + """ + sequence = getattr(peptidoform, "sequence", None) + if sequence is None: + from psm_utils import Peptidoform + + sequence = Peptidoform(str(peptidoform)).sequence + return len(sequence) + + +def _length_buckets( + model: torch.nn.Module, data: Dataset, batch_size: int +) -> list[tuple[torch.Tensor, Dataset]] | None: + """ + Split the dataset into length-sorted chunks, each encoded in a window that fits it. + + Padding every peptide to the model's full window makes the convolutions work on + padding: at a 60-position window and a median peptide of 16 residues most of the + trunk's arithmetic is spent on positions that are masked out again before pooling. + Sorting by length and giving each chunk a window of its own longest peptide plus the + trunk's reach is exact - it was measured identical to the full window over 50,000 + peptides - and about three times faster on CPU. + + Returns None when the model does not report a reach, when the data is not a + DeepLCDataset, or when there is nothing to gain, so the caller falls back to one pass. + """ + reach = getattr(model, "padding_reach", None) + if reach is None or not isinstance(data, DeepLCDataset) or len(data) == 0: + return None + + lengths = np.fromiter( + (_residue_count(p) for p in data.peptidoforms), dtype=np.int64, count=len(data) ) - return predictions.cpu().detach() + window = data.padding_length + if int(lengths.max()) + reach >= window: + return None + + order = np.argsort(lengths, kind="stable") + buckets = [] + for start in range(0, len(order), batch_size): + chunk = order[start : start + batch_size] + padding = int(min(window, lengths[chunk].max() + reach)) + buckets.append((torch.as_tensor(chunk), data.variant(chunk.tolist(), padding))) + return buckets def supports_task_subset(model: torch.nn.Module) -> bool: diff --git a/deeplc/data.py b/deeplc/data.py index 07149ce..2464aab 100644 --- a/deeplc/data.py +++ b/deeplc/data.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from collections.abc import Sequence from typing import TypeVar, overload import numpy as np @@ -93,6 +94,33 @@ def __len__(self) -> int: """Return number of peptidoforms in the dataset.""" return len(self.peptidoforms) + def variant(self, indices: Sequence[int], padding_length: int) -> DeepLCDataset: + """ + Return a subset of this dataset's peptidoforms, encoded in a shorter window. + + Used by the prediction path to run short peptides in a window that fits them + instead of padding every one to the model's full length. The peptidoform objects + themselves are shared rather than copied, so the parsing psm_utils caches on them + is not paid twice. + + Parameters + ---------- + indices + Positions in this dataset to include, in the order wanted. + padding_length + Window the subset is encoded in. + + """ + targets = self.target_retention_times + return type(self)( + peptidoforms=[self.peptidoforms[i] for i in indices], + target_retention_times=None if targets is None else targets[list(indices)], + add_ccs_features=self.add_ccs_features, + add_terminal_composition=self.add_terminal_composition, + padding_length=padding_length, + legacy_positional_deltas=self.legacy_positional_deltas, + ) + def __getitem__(self, idx: int) -> tuple[torch.Tensor, ...]: """Return encoded features and target RT for peptidoform at index.""" if not isinstance(idx, int): diff --git a/tests/test_flexcnn.py b/tests/test_flexcnn.py index 8e5b3df..d63b0f6 100644 --- a/tests/test_flexcnn.py +++ b/tests/test_flexcnn.py @@ -667,3 +667,58 @@ def test_training_scores_its_starting_point(tmp_path): source = inspect.getsource(_model_ops.train) assert 'best_val_loss = float("inf")' not in source assert "_validate_epoch(model, val_loader, loss_fn, device)" in source + + +# --------------------------------------------------------------------------- # +# length-bucketed prediction +# --------------------------------------------------------------------------- # + + +def test_padding_reach_counts_the_convolutions(): + """The reach is the sum over convolutions of ``(kernel - 1) // 2``, dilation aside.""" + model = FlexCNNMultitaskModel(n_tasks=3, **SMALL) + # Two pointwise stem layers reach nothing; two kernel-5 convolutions reach two each. + assert model.padding_reach == 4 + assert FlexCNNMultitaskModel(n_tasks=3, **{**SMALL, "kernel_size": 3}).padding_reach == 2 + + +def test_length_buckets_predict_what_the_full_window_predicts(): + """ + Running short peptides in a short window must not change their predictions. + + The trunk masks by the true residue count and the features do not depend on the + window, so a window of the longest peptide plus the trunk's reach holds everything + that can influence a valid position. That exactness is what allows the prediction + path to stop padding every peptide to sixty positions. + """ + model = FlexCNNMultitaskModel(n_tasks=5, **SMALL).eval() + peptides = ["PEPTIDEK", "ACDEFGHIK", "PEPTIDEPEPTIDEPEPTIDEK", "ACDK", "SEQUENCEWITHK"] + dataset = DeepLCDataset(peptides, add_terminal_composition=True) + + full = _model_ops.predict( + model=model, + data=dataset, + device="cpu", + batch_size=2, + show_progress=False, + length_buckets=False, + ) + bucketed = _model_ops.predict( + model=model, + data=dataset, + device="cpu", + batch_size=2, + show_progress=False, + ) + assert bucketed.shape == full.shape + assert torch.allclose(bucketed, full, atol=1e-5) + + +def test_length_buckets_are_skipped_when_they_cannot_help(): + """A dataset whose longest peptide fills the window is left as one pass.""" + model = FlexCNNMultitaskModel(n_tasks=2, **SMALL).eval() + short_window = DeepLCDataset(["PEPTIDEK"], add_terminal_composition=True, padding_length=10) + assert _model_ops._length_buckets(model, short_window, batch_size=8) is None + # A four-branch model pools across positions and reports no reach, so it never buckets. + plain = DeepLCModel(n_heads=3) + assert getattr(plain, "padding_reach", None) is None From 0866831090040a5582f81fcd1c7fbafd77a3350e Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Mon, 7 Sep 2026 13:08:54 +0200 Subject: [PATCH 2/4] perf: cut prediction chunks by length band, not by batch size Length-sorted chunks of a fixed 4,096 peptides inherit the window of their longest member, so the last chunk of a set runs thousands of ordinary peptides in the window a handful of long ones need. On a 20,000-peptide set that cost 35 % of the throughput, and it made a peptide-length cap look attractive for the wrong reason: capping at 30 residues appeared to gain 37 % with fixed-size chunks and gains 2 % once the chunks are cut by length. A chunk is now cut as soon as its longest peptide exceeds its shortest by more than four residues, or at the batch size, whichever comes first. The dense middle of a length distribution still fills a batch; the sparse tail gets small chunks with tight windows, which measured better than padding it into larger ones (1,546 peptidoforms/s against 1,503 at a floor of 512 and 1,449 at 2,048). The early return that switched bucketing off whenever any single peptide filled the window is gone too: one long peptide should cost one small chunk, not the whole optimisation. Measured on a machine busy with other work, so the absolute rates are low, but every condition ran back to back: fixed-size chunks 1,142 pf/s, length bands 1,546 pf/s. Co-Authored-By: Claude Opus 5 (1M context) --- deeplc/_model_ops.py | 36 +++++++++++++++++++++++++++++------- 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/deeplc/_model_ops.py b/deeplc/_model_ops.py index c63d48c..e4ac60e 100644 --- a/deeplc/_model_ops.py +++ b/deeplc/_model_ops.py @@ -297,6 +297,15 @@ def predict( return out.detach() +#: Widest spread of peptide lengths allowed inside one prediction chunk. Small enough that +#: no peptide carries much padding, large enough that the dense middle of a length +#: distribution still fills a batch. +#: A tight window is worth more than a full batch: enforcing a minimum chunk size, so that the +#: sparse long tail rides along in a wider window, was measured slower (1,546 against 1,503 +#: peptidoforms/s at a floor of 512 and 1,449 at 2,048). +_LENGTH_BAND = 4 + + def _residue_count(peptidoform: object) -> int: """ Residues in a peptidoform, whether it arrives parsed or as a ProForma string. @@ -337,15 +346,28 @@ def _length_buckets( (_residue_count(p) for p in data.peptidoforms), dtype=np.int64, count=len(data) ) window = data.padding_length - if int(lengths.max()) + reach >= window: - return None - order = np.argsort(lengths, kind="stable") - buckets = [] - for start in range(0, len(order), batch_size): - chunk = order[start : start + batch_size] - padding = int(min(window, lengths[chunk].max() + reach)) + sorted_lengths = lengths[order] + + # A chunk is cut either at the batch size or as soon as its longest peptide would exceed + # the shortest by more than _LENGTH_BAND, so no peptide is padded much beyond its own + # length. Fixed-size chunks are not enough: the longest chunk of a length-sorted set holds + # thousands of ordinary peptides alongside the few long ones and inherits their window, + # which on a 20,000-peptide set cost 43 % of the throughput. + buckets: list[tuple[torch.Tensor, Dataset]] = [] + start = 0 + while start < len(order): + stop = min(start + batch_size, len(order)) + band = sorted_lengths[start] + _LENGTH_BAND + within = int(np.searchsorted(sorted_lengths[start:stop], band, side="right")) + stop = start + max(within, 1) + padding = int(min(window, sorted_lengths[stop - 1] + reach)) + chunk = order[start:stop] buckets.append((torch.as_tensor(chunk), data.variant(chunk.tolist(), padding))) + start = stop + + if len(buckets) == 1 and buckets[0][1].padding_length >= window: + return None # one chunk at the full window is what the plain path already does return buckets From 900c3c80b87485e5cfe188c2299036945ae04ee4 Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Mon, 7 Sep 2026 13:37:06 +0200 Subject: [PATCH 3/4] perf: encode batches in one pass and skip the unused rolling sum Two changes to the feature path, which after the trunk work is the largest remaining cost of prediction. The dataset can now encode a whole batch into one buffer per feature. Building four small arrays and four tensors per peptide and letting a DataLoader collate them costs several copies of a few hundred bytes each with a lot of Python around them; writing the encoder output straight into batch buffers measured 1.7x faster over the feature path, 0.153 against 0.089 ms per peptide, for the same values. The DataLoader is kept for worker processes and for datasets that are not a DeepLCDataset. The rolling-sum matrix is no longer built for models that ignore it. The fused trunk reads the per-position matrix directly and deletes this one on the first line of its forward, yet the array was built, converted to a tensor, collated and moved to the device for every peptide: about a tenth of the encoding work. Models declare it through ``uses_rolling_sum`` and ``predict`` passes it to the dataset, so old checkpoints keep the array they were trained with. Checked end to end against the holdout_v5 dump, which was written before any of the performance commits: on PXD079349 the aggregate MAE is 0.1754 against 0.1756 min and coverage 0.9052 against 0.9044, with per-peptide predictions within 0.049 min (median 0.002); on PXD081880, whose 230-peptide reference keeps the fold-based alpha search, predictions agree to 0.003 min. Co-Authored-By: Claude Opus 5 (1M context) --- deeplc/_architecture.py | 4 +++ deeplc/_features.py | 14 ++++++++- deeplc/_model_ops.py | 48 ++++++++++++++++++++++++------ deeplc/core.py | 4 ++- deeplc/data.py | 65 ++++++++++++++++++++++++++++++++++++----- tests/test_flexcnn.py | 29 ++++++++++++++++++ 6 files changed, 146 insertions(+), 18 deletions(-) diff --git a/deeplc/_architecture.py b/deeplc/_architecture.py index 920199d..33b611a 100644 --- a/deeplc/_architecture.py +++ b/deeplc/_architecture.py @@ -844,6 +844,10 @@ class FlexCNNMultitaskModel(nn.Module): #: Index reserved for padding positions in the residue encoding. PAD_INDEX = 20 + #: The fused trunk reads the per-position matrix directly, so the rolling-sum array is + #: redundant and its ``forward`` deletes it. Encoding can skip building it. + uses_rolling_sum = False + def __init__( self, n_tasks: int, diff --git a/deeplc/_features.py b/deeplc/_features.py index 3cd64ba..6c566b2 100644 --- a/deeplc/_features.py +++ b/deeplc/_features.py @@ -40,6 +40,7 @@ def encode_peptidoform( dict_index_pos: dict[str, int] | None = None, dict_index: dict[str, int] | None = None, legacy_positional_deltas: bool = False, + include_rolling_sum: bool = True, ) -> dict[str, np.ndarray]: """ Extract features from a single peptidoform. @@ -58,6 +59,10 @@ def encode_peptidoform( modification on the same residue are indistinguishable. padding_length The maximum length of the sequence after padding. Default is 60. + include_rolling_sum + Whether to build ``matrix_sum``, the rolling sum over pairs of positions. Models + with a convolutional trunk ignore it; False returns an empty array in its place. + Default is True. legacy_positional_deltas Whether to place modification deltas in the positional block the way versions before 4.0.1 did, which was to index ``pos_mat`` without the sorted-layout @@ -139,7 +144,14 @@ def encode_peptidoform( matrix_all = np.append(matrix_all, (seq.count("K") + seq.count("R")) / seq_len) matrix_all = np.append(matrix_all, charge) - matrix_sum = _compute_rolling_sum(std_matrix.T, n=2)[:, ::2].T + # The fused trunk reads the per-position matrix directly and deletes this one on the + # first line of its forward, so building it is pure cost for those models: the cumulative + # sum and its slicing are about a tenth of the encoding work. + matrix_sum = ( + _compute_rolling_sum(std_matrix.T, n=2)[:, ::2].T + if include_rolling_sum + else np.zeros((0, len(dict_index)), dtype=np.float16) + ) matrix_global = np.concatenate([matrix_all, pos_matrix.flatten()]) if add_terminal_composition: diff --git a/deeplc/_model_ops.py b/deeplc/_model_ops.py index e4ac60e..ad05567 100644 --- a/deeplc/_model_ops.py +++ b/deeplc/_model_ops.py @@ -3,7 +3,7 @@ import copy import inspect import logging -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterator, Sequence from os import PathLike from pathlib import Path @@ -272,11 +272,14 @@ def predict( buckets = _length_buckets(model, data, batch_size) if length_buckets else None if buckets is None: - data_loader = DataLoader( - data, batch_size=batch_size, shuffle=False, num_workers=num_workers - ) predictions = _predict_epoch( - model, data_loader, device, show_progress=show_progress, task_idx=task_idx + model, + data, + device, + batch_size=batch_size, + num_workers=num_workers, + show_progress=show_progress, + task_idx=task_idx, ) return predictions.cpu().detach() @@ -284,8 +287,10 @@ def predict( for indices, subset in buckets: part = _predict_epoch( model, - DataLoader(subset, batch_size=batch_size, shuffle=False, num_workers=num_workers), + subset, device, + batch_size=batch_size, + num_workers=num_workers, show_progress=show_progress, task_idx=task_idx, ).cpu() @@ -444,10 +449,30 @@ def _validate_epoch( return float(val_loss / len(data_loader)) +def _feature_batches(data: Dataset, batch_size: int, num_workers: int) -> Iterator[list]: + """ + Yield feature batches, assembled by the dataset itself where it can be. + + A DeepLCDataset encodes a whole batch into one buffer per feature, which skips the + per-peptide tensors and the collate step a DataLoader needs. Worker processes and other + dataset types keep the DataLoader. + """ + if isinstance(data, DeepLCDataset) and num_workers == 0: + for start in range(0, len(data), batch_size): + stop = min(start + batch_size, len(data)) + yield list(data.encode_batch(range(start, stop))) + else: + loader = DataLoader(data, batch_size=batch_size, shuffle=False, num_workers=num_workers) + for features, _ in loader: + yield list(features) + + def _predict_epoch( model: torch.nn.Module, - data_loader: DataLoader, + data: Dataset, device: str, + batch_size: int = 512, + num_workers: int = 0, show_progress: bool = False, task_idx: Sequence[int] | None = None, ) -> torch.Tensor: @@ -457,9 +482,14 @@ def _predict_epoch( if task_idx is not None and supports_task_subset(model): selected = torch.as_tensor(list(task_idx), dtype=torch.long, device=device) predictions = [] + total = int(np.ceil(len(data) / batch_size)) if hasattr(data, "__len__") else None with torch.no_grad(): - for features, _ in track( - data_loader, description="Predicting...", transient=True, disable=not show_progress + for features in track( + _feature_batches(data, batch_size, num_workers), + description="Predicting...", + transient=True, + disable=not show_progress, + total=total, ): features = [feature_tensor.to(device) for feature_tensor in features] outputs = model(*features) if selected is None else model(*features, task_idx=selected) diff --git a/deeplc/core.py b/deeplc/core.py index c3d7d74..2105c8f 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -117,7 +117,9 @@ def predict( result = _model_ops.predict( model=loaded_model, data=DeepLCDataset.from_psm_list( - _parse_psms(psm_list), **_feature_kwargs_from_spec(feature_spec) + _parse_psms(psm_list), + include_rolling_sum=getattr(loaded_model, "uses_rolling_sum", True), + **_feature_kwargs_from_spec(feature_spec), ), **kwargs, ).numpy() diff --git a/deeplc/data.py b/deeplc/data.py index 2464aab..85b8e71 100644 --- a/deeplc/data.py +++ b/deeplc/data.py @@ -29,6 +29,7 @@ def __init__( add_terminal_composition: bool = False, padding_length: int = 60, legacy_positional_deltas: bool = True, + include_rolling_sum: bool = True, ): """ Initialize the DeepLCDataset. @@ -68,6 +69,10 @@ def __init__( models. Note that :func:`deeplc._features.encode_peptidoform`, whose job is correct featurisation rather than model compatibility, defaults the other way. Affects modified peptidoforms only. + include_rolling_sum + 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. Raises ------ @@ -82,6 +87,7 @@ def __init__( self.add_terminal_composition = add_terminal_composition self.padding_length = padding_length self.legacy_positional_deltas = legacy_positional_deltas + self.include_rolling_sum = include_rolling_sum if self.target_retention_times is not None and len(self.target_retention_times) != len( self.peptidoforms ): @@ -119,19 +125,58 @@ def variant(self, indices: Sequence[int], padding_length: int) -> DeepLCDataset: add_terminal_composition=self.add_terminal_composition, padding_length=padding_length, legacy_positional_deltas=self.legacy_positional_deltas, + include_rolling_sum=self.include_rolling_sum, ) + #: Arrays the encoder returns, in the order the models take them. + FEATURE_KEYS = ("matrix", "matrix_sum", "matrix_global", "matrix_hc") + + def _encode_kwargs(self) -> dict: + """Return the encoder settings this dataset was built with.""" + return { + "add_ccs_features": self.add_ccs_features, + "add_terminal_composition": self.add_terminal_composition, + "padding_length": self.padding_length, + "legacy_positional_deltas": self.legacy_positional_deltas, + "include_rolling_sum": self.include_rolling_sum, + } + + def encode_batch(self, indices: Sequence[int]) -> tuple[torch.Tensor, ...]: + """ + Encode several peptidoforms straight into one tensor per feature. + + :meth:`__getitem__` builds four arrays and four tensors for a single peptide, which a + DataLoader then stacks into a batch: several copies of a few hundred bytes each with a + good deal of Python around them. Writing the encoder output into batch buffers + instead measured 1.7x faster over the whole feature path, 0.153 against 0.089 ms per + peptide, and returns the same values. + + Parameters + ---------- + indices + Positions in this dataset to encode, in the order wanted. + + """ + indices = list(indices) + if not indices: + raise ValueError("No indices to encode.") + buffers: list[np.ndarray] | None = None + for row, index in enumerate(indices): + features = encode_peptidoform(self.peptidoforms[index], **self._encode_kwargs()) + if buffers is None: + buffers = [ + np.empty((len(indices), *features[key].shape), dtype=np.float32) + for key in self.FEATURE_KEYS + ] + for buffer, key in zip(buffers, self.FEATURE_KEYS, strict=True): + buffer[row] = features[key] + return tuple(torch.from_numpy(buffer) for buffer in buffers or []) + def __getitem__(self, idx: int) -> tuple[torch.Tensor, ...]: """Return encoded features and target RT for peptidoform at index.""" if not isinstance(idx, int): raise TypeError(f"Index must be an integer, got {type(idx)} instead.") - features = encode_peptidoform( - self.peptidoforms[idx], - add_ccs_features=self.add_ccs_features, - add_terminal_composition=self.add_terminal_composition, - padding_length=self.padding_length, - legacy_positional_deltas=self.legacy_positional_deltas, - ) + features = encode_peptidoform(self.peptidoforms[idx], **self._encode_kwargs()) feature_tuples = ( torch.from_numpy(features["matrix"]).to(dtype=torch.float32), torch.from_numpy(features["matrix_sum"]).to(dtype=torch.float32), @@ -153,6 +198,7 @@ def from_psm_list( add_terminal_composition: bool = False, padding_length: int = 60, legacy_positional_deltas: bool = True, + include_rolling_sum: bool = True, ) -> DeepLCDataset: """ Create a DeepLCDataset from a PSMList. @@ -187,6 +233,10 @@ def from_psm_list( job is correct featurisation rather than model compatibility, defaults the other way. Affects modified peptidoforms only. + 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. + Returns ------- DeepLCDataset @@ -206,6 +256,7 @@ def from_psm_list( add_terminal_composition=add_terminal_composition, padding_length=padding_length, legacy_positional_deltas=legacy_positional_deltas, + include_rolling_sum=include_rolling_sum, ) diff --git a/tests/test_flexcnn.py b/tests/test_flexcnn.py index d63b0f6..0097fa9 100644 --- a/tests/test_flexcnn.py +++ b/tests/test_flexcnn.py @@ -722,3 +722,32 @@ def test_length_buckets_are_skipped_when_they_cannot_help(): # A four-branch model pools across positions and reports no reach, so it never buckets. plain = DeepLCModel(n_heads=3) assert getattr(plain, "padding_reach", None) is None + + +def test_batch_encoding_matches_item_by_item(): + """A batch encoded in one pass must equal the per-peptide encoding a DataLoader stacks.""" + peptides = ["PEPTIDEK", "ACDEFGHIK", "SEQUENCEWITHK", "ACDK"] + dataset = DeepLCDataset(peptides, add_terminal_composition=True) + 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.allclose(batch[position], stacked, atol=0) + + +def test_rolling_sum_is_skipped_for_the_fused_trunk(): + """ + The fused trunk deletes the rolling sum, so encoding does not build it. + + It stays for the four-branch model, which reads it, and the flag travels with the + dataset rather than being decided inside the encoder. + """ + assert FlexCNNMultitaskModel.uses_rolling_sum is False + assert getattr(DeepLCModel(n_heads=2), "uses_rolling_sum", True) is True + + with_sum = DeepLCDataset(["PEPTIDEK"], add_terminal_composition=True) + without = DeepLCDataset(["PEPTIDEK"], add_terminal_composition=True, include_rolling_sum=False) + assert with_sum[0][0][1].shape[0] > 0 + assert without[0][0][1].shape[0] == 0 + # every other feature is untouched by the flag + for position in (0, 2, 3): + assert torch.allclose(with_sum[0][0][position], without[0][0][position], atol=0) From 5cba322b005e19913593eb38c3772d6ef9eb92a5 Mon Sep 17 00:00:00 2001 From: RobbinBouwmeester Date: Tue, 8 Sep 2026 17:20:19 +0200 Subject: [PATCH 4/4] perf: make the multi-head calibration fit cheaper The conformal cross-fitting in prediction_report refits the calibration once per fold, so its cost is paid five times over, and with a 10,541-peptide reference that was about 11 s of a 13.5 s report. Three changes to the fit and one to the spline it uses, none of which alters the selected heads: - Rank the heads without materialising centred copies of the (n, 6543) matrix. The centred target sums to zero, so the covariance is a dot product per block of heads and the variance follows from the block's own sums. Blocks of 512 keep the accumulation in float64 while holding a slice rather than the whole matrix. - Stop promoting that matrix to float64 in fit(). It arrives as float32 and every column is cast back to float32 for its spline, so the promotion only doubled a 276 MB reference to 552 MB. np.asarray on a head source still materialises it, which is what ranking needs. - Let RidgeCV use its closed-form leave-one-out route on references of at least 2,000 peptides: 0.08 s against 0.84 s at 10,541, and accuracy neutral over ten held-out setups (median MAE ratio 1.0000, better on five and worse on five). Smaller references keep the fold-based search, where the collinear columns make a leave-one-out estimate jumpy: on a 725-peptide reference it chose alpha 1 against 316 and cost 3.7 % of accuracy, and the fold search costs 0.15 s there anyway. - The spline calibration predicts its left and right trails only for the points outside the fitted range, instead of for every point and then discarding. Checked end to end against a per-peptide dump written before any of the performance work: on PXD079349 MAE 0.1754 against 0.1756 min and coverage 0.9052 against 0.9044, per-peptide predictions within 0.049 min; on PXD081880, whose 230-peptide reference keeps the fold-based search, within 0.003 min. 201 tests pass. Co-Authored-By: Claude Opus 5 (1M context) --- deeplc/calibration/multihead.py | 46 ++++++++++++++++++++++++++++----- deeplc/calibration/simple.py | 27 ++++++++++--------- 2 files changed, 52 insertions(+), 21 deletions(-) diff --git a/deeplc/calibration/multihead.py b/deeplc/calibration/multihead.py index 24bd02a..b2b6ea7 100644 --- a/deeplc/calibration/multihead.py +++ b/deeplc/calibration/multihead.py @@ -116,7 +116,10 @@ def fit(self, target: np.ndarray, source: np.ndarray) -> None: accepted and treated as a single head, so a single-task model still works. """ - source = np.asarray(source, dtype=np.float64) + # The matrix arrives from the model as float32 and is (n, 6,543) wide; promoting it + # here doubled a 276 MB reference to 552 MB for no gain, since every head's column is + # cast to float32 again for its spline and the ranking accumulates in float64 itself. + source = np.asarray(source) if source.ndim == 1: source = source[:, None] target = np.asarray(target, dtype=np.float64).ravel() @@ -261,7 +264,10 @@ def fit(self, target: np.ndarray, source: np.ndarray) -> None: accepted and treated as a single head, so a single-task model still works. """ - source = np.asarray(source, dtype=np.float64) + # The matrix arrives from the model as float32 and is (n, 6,543) wide; promoting it + # here doubled a 276 MB reference to 552 MB for no gain, since every head's column is + # cast to float32 again for its spline and the ranking accumulates in float64 itself. + source = np.asarray(source) if source.ndim == 1: source = source[:, None] target = np.asarray(target, dtype=np.float64).ravel() @@ -292,8 +298,20 @@ def fit(self, target: np.ndarray, source: np.ndarray) -> None: ) self._head_calibrations.append(head_calibration) + # RidgeCV without an explicit cv uses the closed-form leave-one-out route, which + # decomposes the design once instead of refitting it for every fold and alpha: 0.08 s + # against 0.84 s on a 10,541-peptide reference, and accuracy-neutral over ten held-out + # setups (median MAE ratio 1.0000, better on five and worse on five). + # + # Small references keep the fold-based search. The eighty calibrated columns are + # nearly collinear, which makes a leave-one-out estimate jumpy when there are few + # rows: on a 725-peptide reference it chose alpha 1 against 316 and cost 3.7 % of + # accuracy. Below the threshold the fold-based search costs about 0.15 s, so there is + # nothing to win there anyway. n_splits = int(min(5, max(2, len(target) // 20))) - self._ridge = RidgeCV(alphas=self.alphas, cv=n_splits).fit(calibrated, target) + self._ridge = RidgeCV( + alphas=self.alphas, cv=None if len(target) >= 2000 else n_splits + ).fit(calibrated, target) LOGGER.info( "Calibrated on %d of %d heads with ridge strength %.4g; head %d correlates best.", n_heads, @@ -412,10 +430,24 @@ def _rank_heads_by_correlation(source: np.ndarray, target: np.ndarray) -> np.nda Vectorised because a fused-trunk multitask model can have thousands of heads to rank at once. """ - centred = source - source.mean(axis=0) - target_centred = target - target.mean() + target_centred = np.asarray(target, dtype=np.float64).ravel() + target_centred = target_centred - target_centred.mean() + target_norm = float(np.sqrt((target_centred**2).sum())) + n_rows, n_heads = source.shape + + # Centring the source would allocate a copy of the whole matrix, and squaring it another: + # at 6,543 heads and a 20,000-peptide reference that is well over a gigabyte of + # temporaries for a vector of 6,543 numbers. Because the centred target sums to zero, + # sum_i (x_i - xbar) * yc_i is just sum_i x_i * yc_i, so the covariance is one dot + # product per block and the variance follows from the block's own sums. Blocks keep the + # accumulation in float64 without ever holding more than a slice of the matrix. + correlation = np.empty(n_heads, dtype=np.float64) + block = 512 with np.errstate(invalid="ignore", divide="ignore"): - denominator = np.sqrt((centred**2).sum(axis=0) * (target_centred**2).sum()) - correlation = (centred * target_centred[:, None]).sum(axis=0) / denominator + for start in range(0, n_heads, block): + chunk = np.asarray(source[:, start : start + block], dtype=np.float64) + covariance = chunk.T @ target_centred + variance = np.einsum("ij,ij->j", chunk, chunk) - n_rows * chunk.mean(axis=0) ** 2 + correlation[start : start + block] = covariance / np.sqrt(variance * target_norm**2) correlation = np.where(np.isfinite(correlation), correlation, -np.inf) return np.argsort(-correlation) diff --git a/deeplc/calibration/simple.py b/deeplc/calibration/simple.py index f49257f..b7b064b 100644 --- a/deeplc/calibration/simple.py +++ b/deeplc/calibration/simple.py @@ -318,20 +318,19 @@ def transform(self, source: np.ndarray) -> np.ndarray: if source.shape[0] == 0: return np.array([]) - y_pred_spline = model_main.predict(source.reshape(-1, 1)) - y_pred_left = model_left.predict(source.reshape(-1, 1)) - y_pred_right = model_right.predict(source.reshape(-1, 1)) - within_range = (source >= calibrate_min) & (source <= calibrate_max) - within_range = within_range.ravel() - - cal_preds = np.copy(y_pred_spline) - cal_preds[~within_range & (source.ravel() < calibrate_min)] = y_pred_left[ - ~within_range & (source.ravel() < calibrate_min) - ] - cal_preds[~within_range & (source.ravel() > calibrate_max)] = y_pred_right[ - ~within_range & (source.ravel() > calibrate_max) - ] - return np.array(cal_preds) + flat = source.ravel() + cal_preds = np.asarray(model_main.predict(source.reshape(-1, 1)), dtype=float) + + # The trails only ever supply the points outside the fitted range, which on a + # reference that covers its own gradient is usually none of them. Predicting them + # for every point tripled the work of this method. + below = flat < calibrate_min + above = flat > calibrate_max + if below.any(): + cal_preds[below] = model_left.predict(flat[below].reshape(-1, 1)) + if above.any(): + cal_preds[above] = model_right.predict(flat[above].reshape(-1, 1)) + return cal_preds def _prepare_series(