diff --git a/CHANGELOG.md b/CHANGELOG.md index 3cdaebb..357a4a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,95 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- `legacy_positional_deltas`, on `encode_peptidoform` and `DeepLCDataset`, reproducing + the placement of modification deltas in the positional block exactly as versions + before 4.0.1 did. Verified bit-identical to a v4.0.0 checkout across 4,760 feature + arrays from 1,190 peptidoforms covering lengths 2 to 70, every modified position, and + terminal modifications. Unmodified peptidoforms are unaffected either way. +- The feature layout a model expects is now resolved from the specification it carries, + in one place. A checkpoint that records no specification was written before 4.1.0, so + it also predates the 4.0.1 correction and is fed the encoding it was trained on. One + that records a specification is read literally, and every checkpoint this version + writes records the encoding it used. `DeepLCDataset` therefore defaults to the + pre-4.0.1 placement, since a dataset exists to feed a model; `encode_peptidoform`, + whose job is correct featurisation, still defaults to the corrected placement. + Newly trained models get the corrected placement. + +### Fixed + +- Predictions from models trained before the 4.0.1 feature correction no longer change. + 4.0.1 corrected where modification deltas land in the positional block, which altered + the encoding of every modified peptidoform, while every model released up to that + point had been trained on the old encoding. Modified peptides were therefore predicted + from input those models had never seen. All five bundled checkpoints are bare state + dicts and are now recognised as predating the correction, so their predictions match + v4.0.0 exactly again, with no retraining and no change to any calling code. + + This also covers models held by downstream packages. IM2Deep 2.0.2, unmodified, + reproduces its v4.0.0 CCS predictions exactly on this release; under 4.0.1 its modified + peptides had shifted by up to 7.6 A^2. Note that this means predictions for such models + differ from 4.0.1, which is the point: 4.0.1's change to them was not intended. + +- Fine-tuning onto a new LC setup for models with a low-rank multitask head, fitting the + setup's own `rank + 2` parameters with the encoder and pretrained setups frozen: 66 + values at rank 64, against roughly 1.7 million for an adapter over a 6,543-wide head + vector. `finetune()` previously refused this architecture. +- `MIN_FINETUNE_REFERENCE`, with a warning when fine-tuning is attempted on fewer + reference PSMs. Measured on six unseen LC setups, fine-tuning was worse than + calibration below roughly 500 reference peptidoforms and better above 700; the + validation split is widened automatically below the threshold so early stopping has + signal to work with. +- The new setup's `scale` and `shift` are solved by least squares on the reference data + before training, rather than learned. Left to the optimiser on a small reference set + they collapse: on a 133-minute gradient with 230 reference peptides the output range + shrank to 17 minutes and the error reached 91 minutes, with the correlation still above + 0.9 because the ordering was never what broke. Anchoring them brought that setup to + 2.2 minutes. + +- Fused-trunk multitask architecture (`FlexCNNMultitaskModel`), which merges atomic + composition with a learned residue embedding in a single convolutional trunk and pools + over the valid length instead of flattening four separate branches. Available as + `deeplc.core.FLEXCNN_MULTITASK_MODEL`; the bundled model was trained across 6,543 LC + setups and reaches 0.82 min test MAE and 0.24 min median against 1.26 min and 0.51 min + for the four-branch backbone on the same data and the same head. +- `FactorHead`, a low-rank multitask head where a setup owns only `rank + 2` parameters, + so adding an LC setup means fitting 66 values with the encoder frozen rather than + training a head. +- Self-describing checkpoints. A model file may now record its architecture, constructor + arguments, feature specification and target units, so loading no longer infers the + architecture from tensor shapes. Bare state dicts continue to load unchanged. +- `add_terminal_composition` on `DeepLCDataset` and `DeepLCDataset.from_psm_list`, + passed through to `encode_peptidoform`. + +### Fixed + +- Fine-tuning could return a model far worse than the one it started from. Three + defects compounded: `train()` began with an infinite best validation loss, so the + first epoch always became the best however bad it was; the output layer's scale was + learned rather than solved, though it is linear in its input; and the adapter's ReLU + stack is largely dead at its default initialisation, which left the activations rank + deficient and made a CUDA least-squares solve return non-finite values and silently + decline. On a 133-minute gradient with 230 reference peptides the adapter path + returned predictions spanning 1.3 to 27.2 minutes at an error of 92 minutes, with + the correlation still above 0.9 because only the scale was lost. Training now scores + its starting point, both adaptation paths solve their output layer on the reference + data first, and that solve is rank tolerant. The same setup now gives 1.63 minutes + for the adapter path and 2.05 for the low-rank head, against 1.47 and 1.28 for + calibration. +- A fine-tuned model whose validation error exceeds a large fraction of the reference + retention-time span is now reported at error level, since a collapsed fit leaves the + loss curve and the correlation looking unremarkable. + +### Changed + +- `predict()` loads the model before encoding features, so a model that records a feature + specification gets the features it was trained on. Previously the dataset was always + built with defaults. + ## [4.0.0] - 2026-07-24 ### Changed diff --git a/deeplc/_architecture.py b/deeplc/_architecture.py index f6e1110..1fb8602 100644 --- a/deeplc/_architecture.py +++ b/deeplc/_architecture.py @@ -416,6 +416,92 @@ def add_adapter(self, hidden_size: int = 256) -> None: ) self.adapter.to(self.heads.b2.device) + @torch.no_grad() + def solve_adapter_output( + self, + x_atom: torch.Tensor, + x_atom_sum: torch.Tensor, + x_global: torch.Tensor, + x_one_hot: torch.Tensor, + targets: torch.Tensor, + ) -> bool: + """ + Set the adapter's output layer by least squares on the reference data. + + The adapter is otherwise trained from a default initialisation, so its output + begins unrelated to minutes and, on a small reference set, can settle on a + degenerate fit that predicts every peptide near the mean retention time. Its + final layer is linear in its input, so the best output weights and bias for + the current earlier layers follow in closed form. + + Returns True when the solve was applied. Called after :meth:`add_adapter` and + before training. + + Parameters + ---------- + x_atom, x_atom_sum, x_global, x_one_hot + Encoded reference peptidoforms. + targets + Their observed retention times. + + """ + adapter = getattr(self, "adapter", None) + if adapter is None or not isinstance(adapter[-1], nn.Linear): + return False + + was_training = self.training + self.eval() + try: + x_atom_t = x_atom.transpose(1, 2) + x_atom_sum_t = x_atom_sum.transpose(1, 2) + x_one_hot_t = x_one_hot.transpose(1, 2) + concatenated = torch.cat( + [ + self.branch_a(x_atom_t), + self.branch_b(x_atom_sum_t), + self.branch_c(x_global), + self.branch_d(x_one_hot_t), + ], + dim=1, + ) + head_vector = self.heads(self.shared_trunk(concatenated)) + # Everything up to, but not including, the output layer. + penultimate = head_vector + for layer in list(adapter)[:-1]: + penultimate = layer(penultimate) + finally: + if was_training: + self.train() + + features = penultimate.detach().double() + y = targets.detach().reshape(-1).double() + if features.shape[0] != y.shape[0] or features.shape[0] < 3: + return False + + design = torch.cat([features, torch.ones_like(features[:, :1])], dim=1) + # Solved on the CPU with a rank-revealing driver, and ridge-regularised. The + # adapter's ReLU stack is largely dead at its default initialisation, so the + # activations arriving here are often rank deficient; the default CUDA driver + # requires full rank and returns non-finite values on such a system, which is + # how this solve silently declined to apply. + design = design.cpu() + target = y.unsqueeze(1).cpu() + ridge = 1e-6 * torch.eye(design.shape[1], dtype=design.dtype) + try: + gram = design.T @ design + ridge + solution = torch.linalg.solve(gram, design.T @ target).reshape(-1) + except Exception: # noqa: BLE001 - a singular system leaves the init alone + return False + if not torch.isfinite(solution).all(): + return False + solution = solution.to(features.device) + + output = adapter[-1] + output.weight.copy_(solution[:-1].reshape(1, -1).to(output.weight.dtype)) + if output.bias is not None: + output.bias.copy_(solution[-1].reshape(1).to(output.bias.dtype)) + return True + def freeze_backbone(self) -> None: """Freeze all parameters except the adapter.""" for name, param in self.named_parameters(): @@ -466,3 +552,528 @@ def forward( if adapter is not None: return adapter(out) # [batch, 1] return out + + +# --------------------------------------------------------------------------- +# Fused-trunk multitask architecture +# --------------------------------------------------------------------------- +# +# ``MultitaskDeepLCModel`` above runs four independent branches over four feature +# arrays and concatenates the flattened results, which is what DeepLC has always +# done. The architecture below instead fuses atomic composition with a learned +# residue embedding into a single convolutional trunk and pools it to a fixed +# vector. On the same training data and the same head it reaches roughly a third +# of the error of the four-branch backbone, and it is length-agnostic because it +# pools rather than flattens. + + +class _ConvSiLU(nn.Module): + """A single ``Conv1d`` with ``same`` padding followed by SiLU.""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int): + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding="same") + self.act = nn.SiLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Convolve and activate ``(batch, channels, length)``.""" + return self.act(self.conv(x)) + + +def _pointwise_stem(channels: int, layers: int, in_channels: int = 6) -> nn.Sequential: + """ + Stack of width-1 convolutions decoding each position's atom counts alone. + + Atomic composition determines the residue uniquely except for the Leu/Ile + pair, so this is the layer that can recover identity from the atom matrix + before any neighbour mixing happens. Recovering a residue from six counts is + a lookup rather than a linear map, so one layer is generally not enough. + + Returned as a bare ``Sequential`` rather than wrapped in a module, so the + parameter names stay ``stem.0.conv.weight`` and match the trained checkpoint. + """ + return nn.Sequential( + *[ + _ConvSiLU(in_channels if i == 0 else channels, channels, kernel_size=1) + for i in range(layers) + ] + ) + + +class InputNorm(nn.Module): + """ + Per-feature standardisation with buffers fitted on the training rows. + + The dense feature vector mixes raw atom counts, sequence length and + positional compositions, whose scales differ by an order of magnitude. + + A feature that never varies during training is left in raw units rather than + having its standard deviation clamped to a floor. Clamping would multiply any + non-zero test value by one over that floor: with phosphorus absent from + training, a single phosphate standardised to a value near a thousand and + destroyed the forward pass. Constant features must not become high-gain + inputs. + """ + + def __init__(self, n_features: int): + super().__init__() + self.register_buffer("mean", torch.zeros(n_features)) + self.register_buffer("std", torch.ones(n_features)) + + @torch.no_grad() + def fit(self, values: torch.Tensor) -> None: + """Set the buffers from ``(n_rows, n_features)`` of training features.""" + self.mean.copy_(values.mean(dim=0)) + std = values.std(dim=0) + self.std.copy_(torch.where(std < 1e-3, torch.ones_like(std), std)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Standardise ``(batch, n_features)``.""" + return (x - self.mean) / self.std + + +class FactorHead(nn.Module): + """ + Low-rank multitask head: a per-setup embedding dotted with a shared trunk. + + ``pred[:, j] = (proj(trunk) . embedding[j]) * scale[j] + shift[j]`` + + Unlike :class:`BatchedHeads`, which gives every LC setup its own hidden + projection, the projection here is shared and only ``rank + 2`` parameters + belong to a setup. Adding a setup therefore means fitting 66 values at + rank 64, with the encoder frozen, rather than retraining a head. + + ``scale`` and ``shift`` are an affine map on that setup's output. Training + normalises each setup's retention times to a fixed range, so a raw training + checkpoint predicts in that normalised space; because both transforms are + affine they compose, and packaging folds the normalisation into ``scale`` and + ``shift`` so that a shipped model returns minutes directly. Whether that has + happened is recorded as ``target_units`` in the checkpoint rather than + assumed. + + Parameters + ---------- + trunk_dim + Width of the shared trunk output. + n_tasks + Number of LC setups the model was trained on. + rank + Size of the per-setup embedding. + + """ + + def __init__(self, trunk_dim: int, n_tasks: int, rank: int = 64): + super().__init__() + self.n_tasks = n_tasks + self.rank = rank + self.proj = nn.Linear(trunk_dim, rank) + self.embedding = nn.Parameter(torch.zeros(n_tasks, rank)) + self.scale = nn.Parameter(torch.ones(n_tasks)) + self.shift = nn.Parameter(torch.zeros(n_tasks)) + + def add_task(self, targets: torch.Tensor | None = None, init_from: int | None = None) -> None: + """ + Attach parameters for one new LC setup and switch to single-task output. + + The new setup gets its own ``rank + 2`` parameters as separate tensors rather + than extra rows in the existing ones, so freezing the pretrained setups is a + matter of ``requires_grad`` and does not need per-row gradient masking. + + The embedding starts at the mean of the trained setups, which is the least + committed starting point available. + + ``scale`` and ``shift`` start from the mean of the trained setups rather than + from the target retention times. Setting them directly from the targets looks + natural and is wrong by more than an order of magnitude: the dot product they + multiply is in the normalised space the model was trained in, roughly 0 to 100, + not in minutes, so seeding ``scale`` with a spread measured in minutes + overshoots by about thirty-fold and the fit starts hundreds of minutes away. + The trained setups' own values are already the right magnitude for mapping that + dot product onto a real gradient. + + Parameters + ---------- + targets + Observed retention times for the new setup. Used only to nudge the + initial offset toward the right part of the gradient. + init_from + Index of an existing setup to copy from, instead of the mean. Useful when + a similar gradient is known. + + """ + with torch.no_grad(): + if init_from is None: + start = self.embedding.mean(dim=0, keepdim=True).clone() + scale = self.scale.mean().reshape(1).clone() + shift = self.shift.mean().reshape(1).clone() + else: + start = self.embedding[init_from : init_from + 1].clone() + scale = self.scale[init_from].reshape(1).clone() + shift = self.shift[init_from].reshape(1).clone() + + if targets is not None and targets.numel() > 1: + # Centre the offset on the observed gradient while leaving the slope + # at a trained magnitude, so the fit starts on the right window. + shift = shift + (targets.mean() - shift) + + self.new_embedding = nn.Parameter(start) + self.new_scale = nn.Parameter(scale) + self.new_shift = nn.Parameter(shift) + + def project_new_task(self, trunk: torch.Tensor) -> torch.Tensor: + """Return the new setup's prediction before scale and shift are applied.""" + return (self.proj(trunk) @ self.new_embedding.t()).squeeze(-1) + + @torch.no_grad() + def solve_new_task_affine(self, projected: torch.Tensor, targets: torch.Tensor) -> None: + """ + Set ``scale`` and ``shift`` by least squares instead of learning them. + + The two are linear in the prediction, so the best values for a given + embedding follow in closed form and do not need gradient descent. Leaving + them to the optimiser is how fine-tuning fails on a small reference set: with + few points the scale decays toward zero and every prediction collapses onto + the mean retention time. On one 133-minute gradient with 230 reference + peptides that produced a 17-minute output range and a 91-minute error, while + the correlation stayed above 0.9 because the ordering was never the problem. + + Solving them first also gives the embedding a sane starting error to descend + from, rather than one dominated by a mis-scaled output. + """ + x = projected.detach().reshape(-1).double() + y = targets.detach().reshape(-1).double() + if x.numel() < 3 or torch.std(x) < 1e-9: + return + design = torch.stack([x, torch.ones_like(x)], dim=1) + solution = torch.linalg.lstsq(design, y.unsqueeze(1)).solution.reshape(-1) + slope, intercept = solution[0], solution[1] + if not torch.isfinite(slope) or not torch.isfinite(intercept): + return + self.new_scale.copy_(slope.reshape(1).to(self.new_scale.dtype)) + self.new_shift.copy_(intercept.reshape(1).to(self.new_shift.dtype)) + + def freeze_pretrained(self) -> None: + """Freeze everything except the newly added setup's parameters.""" + for parameter in self.parameters(): + parameter.requires_grad = False + for name in ("new_embedding", "new_scale", "new_shift"): + parameter = getattr(self, name, None) + if parameter is not None: + parameter.requires_grad = True + + @property + def has_new_task(self) -> bool: + """Whether :meth:`add_task` has been called.""" + return getattr(self, "new_embedding", None) is not None + + def forward(self, trunk: torch.Tensor, task_idx: torch.Tensor | None = None) -> torch.Tensor: + """ + Map trunk output to one prediction per task. + + Parameters + ---------- + trunk + Shape ``(batch, trunk_dim)``. + task_idx + Optional task indices to evaluate. Without it every task is + returned, which for a model trained on thousands of setups is a wide + matrix; pass the subset when only a few setups are of interest. + + Returns + ------- + torch.Tensor + Shape ``(batch, n_tasks)``, or ``(batch, len(task_idx))``. + + """ + projected = self.proj(trunk) + if self.has_new_task: + # Fine-tuned onto one setup: return that column alone, so the shape + # matches a single-output model and the training loop needs no change. + return (projected @ self.new_embedding.t()) * self.new_scale + self.new_shift + if task_idx is None: + return projected @ self.embedding.t() * self.scale + self.shift + return (projected @ self.embedding[task_idx].t()) * self.scale[task_idx] + self.shift[ + task_idx + ] + + +class FlexCNNMultitaskModel(nn.Module): + """ + Multitask RT model fusing atom composition and residue identity in one trunk. + + Atom counts pass through a pointwise stem, are concatenated with a learned + residue embedding, and the result runs through a convolutional stack that is + masked and pooled over the valid length. The pooled vector is concatenated + with the global feature vector and the residue counts, and a small MLP feeds + :class:`FactorHead`. + + The forward signature matches :class:`MultitaskDeepLCModel` so the two are + interchangeable in the prediction path, but ``x_atom_sum`` is unused: the + convolutional trunk sees the per-position matrix directly, which makes the + rolling-sum array redundant. + + The global feature vector must be the 67-dimensional form produced with + ``add_terminal_composition=True``; the 55-dimensional default will fail on + shape at the first dense layer. + + Parameters + ---------- + n_tasks + Number of LC setups. + global_dim + Length of the global feature vector. + embed_dim + Width of the residue embedding. + channels + Output channels of each convolution stage. + kernel_size + Convolution width. + stem_channels + Width of the pointwise stem, or 0 to feed raw atom counts. + stem_layers + Number of pointwise stem layers. + width + Width of the dense trunk. + depth + Number of dense layers. + rank + Per-setup embedding size in the head. + + """ + + #: Index reserved for padding positions in the residue encoding. + PAD_INDEX = 20 + + def __init__( + self, + n_tasks: int, + global_dim: int = 67, + embed_dim: int = 16, + channels: tuple[int, ...] = (512, 512), + kernel_size: int = 5, + stem_channels: int = 128, + stem_layers: int = 2, + width: int = 256, + depth: int = 3, + rank: int = 64, + ): + super().__init__() + # Kept so the model can describe itself when saved, rather than leaving + # the serialiser to know how each architecture is constructed. + self._encoder_kwargs = { + "global_dim": global_dim, + "embed_dim": embed_dim, + "channels": tuple(channels), + "kernel_size": kernel_size, + "stem_channels": stem_channels, + "stem_layers": stem_layers, + "width": width, + "depth": depth, + } + self._head_kwargs = {"rank": rank} + self.n_tasks = n_tasks + #: Feature layout this instance expects. Replaced by the value recorded in + #: a checkpoint when one is loaded. + self.feature_spec: dict | None = { + "name": "global67_terminal" if global_dim == 67 else f"global{global_dim}", + "global_dim": global_dim, + "add_terminal_composition": global_dim == 67, + "add_ccs_features": False, + "padding_length": 60, + # Trained after the 4.0.1 correction to positional modification deltas, + # so it wants the corrected placement rather than the compatibility + # default DeepLCDataset applies for checkpoints that declare nothing. + "legacy_positional_deltas": False, + } + self.target_units: str | None = None + self.task_names: list[str] | None = None + + self.encoder = _FlexCNNEncoder( + global_dim=global_dim, + embed_dim=embed_dim, + channels=channels, + kernel_size=kernel_size, + stem_channels=stem_channels, + stem_layers=stem_layers, + width=width, + depth=depth, + ) + self.head = FactorHead(trunk_dim=width, n_tasks=n_tasks, rank=rank) + + def forward( + self, + x_atom: torch.Tensor, + x_atom_sum: torch.Tensor, + x_global: torch.Tensor, + x_one_hot: torch.Tensor, + task_idx: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Predict retention time for every LC setup. + + Parameters + ---------- + x_atom + Shape ``(batch, length, 6)``, per-position atomic composition. + x_atom_sum + Unused; accepted so the signature matches the four-branch model. + x_global + Shape ``(batch, 67)``, global feature vector with terminal + composition. + x_one_hot + Shape ``(batch, length, 20)``, one-hot residue encoding. + task_idx + Optional subset of task indices to evaluate. + + Returns + ------- + torch.Tensor + Shape ``(batch, n_tasks)``, or ``(batch, len(task_idx))``. + + """ + 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) + + def add_task_head( + self, targets: torch.Tensor | None = None, init_from: int | None = None + ) -> int: + """ + Prepare the model to be fine-tuned onto one new LC setup. + + Adds ``rank + 2`` trainable parameters and freezes everything else, so + adapting to a setup costs 66 values at rank 64 rather than retraining a head + or fitting an adapter over the full head vector. Returns the number of + trainable parameters, which callers log to make the cost visible. + """ + self.head.add_task(targets=targets, init_from=init_from) + self.head.freeze_pretrained() + for parameter in self.encoder.parameters(): + parameter.requires_grad = False + return sum(p.numel() for p in self.parameters() if p.requires_grad) + + @torch.no_grad() + def solve_new_task_affine( + self, features: tuple[torch.Tensor, ...], targets: torch.Tensor + ) -> None: + """ + Anchor the new setup's affine parameters on the reference data. + + Call after :meth:`add_task_head` and before training, passing the reference + features and their observed retention times. + """ + trunk = self.encoder(features[0], features[2], features[3]) + self.head.solve_new_task_affine(self.head.project_new_task(trunk), targets) + + def describe(self) -> dict: + """ + Return a serialisable description of this model, including its weights. + + Saving this rather than a bare state dict means a checkpoint can be + reloaded without guessing the architecture from tensor names, and it + keeps the knowledge of how to rebuild a model with the model rather than + in the serialiser. + """ + return { + "architecture": type(self).__name__, + "encoder_kwargs": dict(self._encoder_kwargs), + "head_kwargs": dict(self._head_kwargs), + "n_tasks": self.n_tasks, + "feature_spec": self.feature_spec, + "target_units": self.target_units, + "task_names": self.task_names, + "state_dict": self.state_dict(), + } + + +class _FlexCNNEncoder(nn.Module): + """Convolutional trunk shared by every LC setup in :class:`FlexCNNMultitaskModel`.""" + + PAD_INDEX = 20 + + def __init__( + self, + global_dim: int, + embed_dim: int, + channels: tuple[int, ...], + kernel_size: int, + stem_channels: int, + stem_layers: int, + width: int, + depth: int, + ): + super().__init__() + self.embed = nn.Embedding(self.PAD_INDEX + 1, embed_dim, padding_idx=self.PAD_INDEX) + self.stem = _pointwise_stem(stem_channels, stem_layers) if stem_channels > 0 else None + + in_channels = (stem_channels if stem_channels > 0 else 6) + embed_dim + blocks = [] + for out_channels in channels: + blocks.append(_ConvSiLU(in_channels, out_channels, kernel_size)) + in_channels = out_channels + self.blocks = nn.ModuleList(blocks) + + # Sum and max pooling, masked to the valid length. Sum is extensive in + # peptide length and max is not, so the pair carries both. + pooled_dim = in_channels * 2 + self.pool_norm = nn.LayerNorm(pooled_dim) + + dense_dim = global_dim + self.PAD_INDEX + self.norm = InputNorm(dense_dim) + layers: list[nn.Module] = [] + sizes = [dense_dim + pooled_dim] + [width] * depth + for i in range(len(sizes) - 1): + layers.append(nn.Linear(sizes[i], sizes[i + 1])) + layers.append(nn.SiLU()) + self.net = nn.Sequential(*layers) + self.trunk_dim = width + + def residue_indices(self, x_one_hot: torch.Tensor) -> torch.Tensor: + """ + Convert a one-hot residue matrix to integer indices. + + Padding positions are all-zero rows, for which ``argmax`` returns 0 and + would collide with the first residue, so they are set to + :attr:`PAD_INDEX` explicitly. + """ + idx = x_one_hot.argmax(dim=2) + return idx.masked_fill(x_one_hot.sum(dim=2) == 0, self.PAD_INDEX) + + def residue_counts(self, idx: torch.Tensor) -> torch.Tensor: + """ + Residue counts per peptidoform, shape ``(batch, 20)``. + + Built with ``scatter_add`` rather than by materialising a + ``(batch, length, 21)`` one-hot tensor, which dominated inference time + for every model that uses counts. + """ + idx = idx.long().clamp(0, self.PAD_INDEX) + counts = torch.zeros( + idx.shape[0], self.PAD_INDEX + 1, device=idx.device, dtype=torch.float32 + ) + counts.scatter_add_(1, idx, torch.ones_like(idx, dtype=torch.float32)) + return counts[:, : self.PAD_INDEX] # drop the padding column + + def forward( + self, x_atom: torch.Tensor, x_global: torch.Tensor, x_one_hot: torch.Tensor + ) -> torch.Tensor: + """Encode a batch to ``(batch, trunk_dim)``.""" + idx = self.residue_indices(x_one_hot) + valid = (idx != self.PAD_INDEX).unsqueeze(1) # (batch, 1, length) + + atom = x_atom.float().transpose(1, 2) + if self.stem is not None: + atom = self.stem(atom) + hidden = torch.cat([atom, self.embed(idx).transpose(1, 2)], dim=1) + + for block in self.blocks: + hidden = block(hidden) + hidden = hidden * valid + + summed = hidden.sum(dim=2) + maxed = torch.nan_to_num( + hidden.masked_fill(~valid, float("-inf")).max(dim=2).values, neginf=0.0 + ) + pooled = self.pool_norm(torch.cat([summed, maxed], dim=1)) + + dense = torch.cat([x_global.float(), self.residue_counts(idx)], dim=1) + return self.net(torch.cat([self.norm(dense), pooled], dim=1)) diff --git a/deeplc/_features.py b/deeplc/_features.py index d546995..3cd64ba 100644 --- a/deeplc/_features.py +++ b/deeplc/_features.py @@ -39,6 +39,7 @@ def encode_peptidoform( dict_aa: dict[str, int] | None = None, dict_index_pos: dict[str, int] | None = None, dict_index: dict[str, int] | None = None, + legacy_positional_deltas: bool = False, ) -> dict[str, np.ndarray]: """ Extract features from a single peptidoform. @@ -57,6 +58,13 @@ def encode_peptidoform( modification on the same residue are indistinguishable. padding_length The maximum length of the sequence after padding. Default is 60. + 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 + offset and to reach only one row. That placement is wrong, but a model trained + against it expects it, so this reproduces it exactly for such models. Affects + modified peptidoforms only; unmodified ones encode identically either way. + Default is False, meaning the corrected placement. positions The positions to consider for feature extraction. Default is DEFAULT_POSITIONS. positions_pos @@ -105,6 +113,7 @@ def encode_peptidoform( dict_index, dict_index_pos, positions, + legacy_positional_deltas, ) _apply_terminal_modifications( std_matrix, @@ -114,6 +123,7 @@ def encode_peptidoform( dict_index, dict_index_pos, positions, + legacy_positional_deltas, ) matrix_all = np.sum(std_matrix, axis=0) @@ -239,6 +249,28 @@ def _positional_rows(i: int, seq_len: int, positions: set[int]) -> list[int]: return rows +def _legacy_positional_rows(i: int, seq_len: int, positions: set[int]) -> list[int]: + """ + Positional rows as written before version 4.0.1. + + Kept because models trained against that encoding expect it. Two differences + from :func:`_positional_rows`, both wrong and both reproduced here: + + * ``pos_mat`` is indexed by ``i`` directly, without subtracting + ``min(positions)``. With the default sets the offset is 4, so a delta at + sequence position 1 landed in the row meaning position -3, and a negative + index wrapped in from the end of the block. + * The two cases were ``if``/``elif``, so a residue occupying both a positive + and a negative row received the delta in only one of them, while + :func:`_fill_pos_matrix` wrote its base composition to both. + """ + if i in positions: + return [i] + if (i - seq_len) in positions: + return [i - seq_len] + return [] + + def _terminal_composition( peptidoform: Peptidoform, dict_index: dict[str, int], @@ -275,15 +307,21 @@ def _apply_composition_to_matrices( dict_index: dict[str, int], dict_index_pos: dict[str, int], positions: set[int], + legacy_positional_deltas: bool = False, ) -> None: """ Apply a composition delta to the standard and positional matrices. Positional rows come from :func:`_positional_rows`, which applies the same offset and the same both-ends handling that :func:`_fill_pos_matrix` uses - for base residue compositions. + for base residue compositions, or from :func:`_legacy_positional_rows` when + reproducing the pre-4.0.1 placement for a model trained against it. """ - rows = _positional_rows(i, seq_len, positions) + rows = ( + _legacy_positional_rows(i, seq_len, positions) + if legacy_positional_deltas + else _positional_rows(i, seq_len, positions) + ) for atom_comp, change in composition.items(): try: mat[i, dict_index[atom_comp]] += change @@ -311,6 +349,7 @@ def _apply_modifications( dict_index: dict[str, int], dict_index_pos: dict[str, int], positions: set[int], + legacy_positional_deltas: bool = False, ) -> None: """Apply modification changes to the matrices.""" for i, token in enumerate(parsed_seq): @@ -332,6 +371,7 @@ def _apply_modifications( dict_index, dict_index_pos, positions, + legacy_positional_deltas, ) @@ -343,6 +383,7 @@ def _apply_terminal_modifications( dict_index: dict[str, int], dict_index_pos: dict[str, int], positions: set[int], + legacy_positional_deltas: bool = False, ) -> None: """Apply N- and C-terminal modification changes to the matrices.""" terminal_mods = [ @@ -370,6 +411,7 @@ def _apply_terminal_modifications( dict_index, dict_index_pos, positions, + legacy_positional_deltas, ) diff --git a/deeplc/_model_ops.py b/deeplc/_model_ops.py index 2f5d777..15e04dc 100644 --- a/deeplc/_model_ops.py +++ b/deeplc/_model_ops.py @@ -1,8 +1,9 @@ """Training, predicting, and evaluating with PyTorch.""" import copy +import inspect import logging -from collections.abc import Callable +from collections.abc import Callable, Sequence from os import PathLike from pathlib import Path @@ -17,7 +18,7 @@ ) from torch.utils.data import DataLoader, Dataset, Subset -from deeplc._architecture import DeepLCModel +from deeplc._architecture import DeepLCModel, FlexCNNMultitaskModel from deeplc.data import DeepLCDataset logger = logging.getLogger(__name__) @@ -34,10 +35,17 @@ def load_model( ) if isinstance(model, (str, PathLike, Path)): + raw = torch.load(model, weights_only=True, map_location=selected_device) + + # Newer checkpoints are a dict describing the model rather than a bare + # state dict, so the architecture and its hyperparameters do not have to + # be guessed from tensor shapes. Older files keep working unchanged. + if isinstance(raw, dict) and "architecture" in raw: + return _load_described_model(raw, selected_device) + # Infer architecture hyperparameters from the saved state dict # Only checks n_heads and final_num_layers; other hyperparameters are set to defaults # May break for models saved with different architectures. - raw = torch.load(model, weights_only=True, map_location=selected_device) n_heads = raw["heads.b2"].shape[0] final_num_layers = sum( 1 for k in raw if k.startswith("shared_trunk.") and k.endswith(".weight") @@ -46,7 +54,7 @@ def load_model( if "adapter.0.weight" in raw: loaded_model.add_adapter(hidden_size=raw["adapter.0.weight"].shape[0]) loaded_model.load_state_dict(raw) - elif isinstance(model, DeepLCModel): + elif isinstance(model, (DeepLCModel, FlexCNNMultitaskModel)): loaded_model = model logger.debug("Using provided PyTorch model instance") elif model is None: @@ -60,6 +68,55 @@ def load_model( return loaded_model +#: Architectures a described checkpoint may name, and the class to build. +_DESCRIBED_ARCHITECTURES = { + "FlexCNNMultitaskModel": FlexCNNMultitaskModel, +} + + +def _load_described_model(blob: dict, device: torch.device | str) -> torch.nn.Module: + """ + Build a model from a checkpoint that describes itself. + + The checkpoint carries the architecture name, its constructor arguments and + the feature specification it was trained against. That last part matters: + the model's first dense layer fixes the width of the global feature vector, + so a model expecting terminal composition cannot be fed the shorter default + vector. Attaching the specification to the returned module lets the caller + build a matching dataset instead of inferring it. + """ + name = blob["architecture"] + try: + cls = _DESCRIBED_ARCHITECTURES[name] + except KeyError: + raise ValueError( + f"Checkpoint names architecture {name!r}, which this version of DeepLC " + f"does not know. Known architectures: " + f"{sorted(_DESCRIBED_ARCHITECTURES)}." + ) from None + + kwargs = dict(blob.get("encoder_kwargs") or {}) + kwargs.update(blob.get("head_kwargs") or {}) + built = cls(n_tasks=blob["n_tasks"], **kwargs) + built.load_state_dict(blob["state_dict"]) + + # Carried on the instance so predict() can build a matching dataset. + built.feature_spec = blob.get("feature_spec") + built.target_units = blob.get("target_units") + built.task_names = blob.get("task_names") + + logger.debug( + "Loaded %s with %d tasks, feature spec %s, targets in %s", + name, + blob["n_tasks"], + (built.feature_spec or {}).get("name", "unspecified"), + built.target_units or "unspecified units", + ) + built.to(device) + built.eval() + return built + + def train( model: DeepLCModel | PathLike | str | None, train_dataset: DeepLCDataset | Subset[DeepLCDataset], @@ -142,7 +199,14 @@ def train( loss_fn = torch.nn.L1Loss() best_model_wts = copy.deepcopy(model.state_dict()) - best_val_loss = float("inf") + + # Score the starting point, so training can never return a model worse than the + # one it began with. With this left at infinity the first epoch always became the + # best, even when it was worse: fine-tuning a small reference set could hand back + # a fit whose predictions had collapsed onto the mean retention time, at ninety + # times the error of the model it started from. + best_val_loss = _validate_epoch(model, val_loader, loss_fn, device) + logger.debug("Validation loss before training: %.4f", best_val_loss) epochs_no_improve = 0 with _create_progress(disable=not show_progress) as progress: @@ -185,16 +249,32 @@ def predict( num_workers: int = 0, num_threads: int | None = None, show_progress: bool = True, + task_idx: Sequence[int] | None = None, ) -> torch.Tensor: """Predict using the model for the given dataset.""" + # ``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, + # so a caller wanting one column should ask for one column. Models whose + # forward does not accept it ignore the argument. 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) + predictions = _predict_epoch( + model, data_loader, device, show_progress=show_progress, task_idx=task_idx + ) return predictions.cpu().detach() +def supports_task_subset(model: torch.nn.Module) -> bool: + """Whether ``model.forward`` accepts a ``task_idx`` argument.""" + try: + return "task_idx" in inspect.signature(model.forward).parameters + except (TypeError, ValueError): + return False + + def evaluate( model: torch.nn.Module | PathLike | str | None, data: Dataset, @@ -265,16 +345,20 @@ def _predict_epoch( data_loader: DataLoader, device: str, show_progress: bool = False, + task_idx: Sequence[int] | None = None, ) -> torch.Tensor: """Predict using the model for one epoch.""" model.eval() + selected = None + if task_idx is not None and supports_task_subset(model): + selected = torch.as_tensor(list(task_idx), dtype=torch.long, device=device) predictions = [] with torch.no_grad(): for features, _ in track( data_loader, description="Predicting...", transient=True, disable=not show_progress ): features = [feature_tensor.to(device) for feature_tensor in features] - outputs = model(*features) + outputs = model(*features) if selected is None else model(*features, task_idx=selected) predictions.append(outputs.cpu()) if not predictions: raise ValueError("Dataset is empty — nothing to predict.") diff --git a/deeplc/core.py b/deeplc/core.py index 1319e54..082857c 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -9,6 +9,7 @@ import numpy as np import torch from psm_utils import PSM, Peptidoform, PSMList +from torch.utils.data import DataLoader from deeplc import _model_ops from deeplc._reference_selection import select_reference_psms @@ -23,6 +24,27 @@ DEEPLC_DIR = Path(__file__).resolve().parent DEFAULT_MODEL = DEEPLC_DIR / "package_data" / "models" / "multitask_model.pt" +#: Below this many reference PSMs, fine-tuning measured worse than calibration on +#: every held-out setup tried, so it is warned about rather than silently attempted. +#: On six unseen LC setups the crossover sat between 300 and 735 reference +#: peptidoforms: at 230 the error went from 1.47 to 91.7 min, at 300 from 0.37 to +#: 0.48, and at 735 and above fine-tuning helped every version. +MIN_FINETUNE_REFERENCE = 500 + +#: A fine-tuned model whose validation error exceeds this fraction of the reference +#: retention-time span has not converged onto the gradient, whatever the loss curve +#: said. Measured failures collapse the output range and predict near the mean, so +#: the error lands at a large fraction of the span: on one 133-minute gradient the +#: adapter path reached 92 minutes, or 69 % of the span, while its correlation stayed +#: above 0.9 because the peptide ordering was never what broke. Ordinary fits sit +#: near 1 %. +MAX_FINETUNE_ERROR_FRACTION = 0.15 + +#: Fused-trunk multitask model, trained across 6,543 LC setups. Not the default: +#: switching would change every prediction, so the choice is left to the caller +#: until the calibration path is adapted to its low-rank head. +FLEXCNN_MULTITASK_MODEL = DEEPLC_DIR / "package_data" / "models" / "multitask_flexcnn_model.pt" + def predict( psm_list: PSMList | list[PSM | Peptidoform | str], @@ -53,10 +75,35 @@ def predict( produces multitask output, in which case shape is ``(n, n_heads)``. """ + # The model is loaded before the dataset is built because the features it + # needs depend on the model. A checkpoint that describes itself carries a + # feature specification, and a model trained on the 67-dimensional global + # vector cannot be fed the 55-dimensional default. + # + # The device is taken from predict_kwargs rather than left to default, so a + # caller asking for CPU does not first have the checkpoint placed on a GPU + # it may not fit on. + kwargs = dict(predict_kwargs or {}) + loaded_model = _model_ops.load_model(model or DEFAULT_MODEL, device=kwargs.get("device")) + feature_spec = getattr(loaded_model, "feature_spec", None) or {} + + # Only one column is wanted unless the caller asked for the matrix. A model + # trained on thousands of LC setups would otherwise materialise a column per + # setup, which at a million peptides is tens of gigabytes of output for a + # result the caller then throws away. + if ( + not return_matrix + and "task_idx" not in kwargs + and _model_ops.supports_task_subset(loaded_model) + ): + kwargs["task_idx"] = [0] + result = _model_ops.predict( - model=model or DEFAULT_MODEL, - data=DeepLCDataset.from_psm_list(_parse_psms(psm_list)), - **(predict_kwargs or {}), + model=loaded_model, + data=DeepLCDataset.from_psm_list( + _parse_psms(psm_list), **_feature_kwargs_from_spec(feature_spec) + ), + **kwargs, ).numpy() if not return_matrix: return result[:, 0] @@ -280,6 +327,59 @@ def finetune_and_predict( return calibrated_rt +def _feature_kwargs_from_spec(spec: dict | None) -> dict: + """ + Feature settings a model expects, taken from the specification it carries. + + A checkpoint that records no specification was written before 4.1.0, which + means it also predates the 4.0.1 correction to positional modification + deltas, so it is fed the encoding it was trained on. Every model DeepLC has + released so far is in that position: all five bundled checkpoints are bare + state dicts. A checkpoint that does record a specification is read literally, + and one written by this version always records the encoding it used. + """ + spec = spec or {} + return { + "add_ccs_features": bool(spec.get("add_ccs_features", False)), + "add_terminal_composition": bool(spec.get("add_terminal_composition", False)), + "padding_length": int(spec.get("padding_length", 60)), + "legacy_positional_deltas": bool(spec.get("legacy_positional_deltas", not spec)), + } + + +def _solve_reference_affine( + model: torch.nn.Module, + dataset, + device: str | None = None, + adapter: bool = False, +) -> bool: + """ + Put a freshly attached output on the right axis before training starts. + + Both adaptation paths end in a layer that is linear in its input, so the values + that best fit the reference data follow in closed form. Solving them first is what + keeps a small reference set from producing a fit that predicts every peptide near + the mean retention time. + + Returns True when the solve was applied. A failure is not fatal: training then + proceeds from the initialised values, which is the previous behaviour. + """ + try: + loader = DataLoader(dataset, batch_size=len(dataset), shuffle=False) + features, targets = next(iter(loader)) + selected = device or ("cuda" if torch.cuda.is_available() else "cpu") + model.to(selected) + moved = tuple(f.to(selected) for f in features) + targets = targets.to(selected).float() + if adapter: + return bool(model.solve_adapter_output(*moved, targets)) + model.solve_new_task_affine(moved, targets) + return True + except Exception: # noqa: BLE001 - never let the anchor break the fit + LOGGER.warning("Could not anchor the output layer; training from the initialised values.") + return False + + def finetune( psm_list_reference: PSMList, psm_list_validation: PSMList | None = None, @@ -299,7 +399,9 @@ def finetune( used. validation_split Fraction of ``psm_list_reference`` to use for validation when ``psm_list_validation`` - is None. + is None. Raised to at least 0.25 when the reference set is smaller than + :data:`MIN_FINETUNE_REFERENCE`, because early stopping cannot work on a + handful of PSMs and an unchecked fit can end up far worse than calibration. model Trained model or path to model file. train_kwargs @@ -312,12 +414,41 @@ def finetune( """ LOGGER.info("Fine-tuning model...") + + # Fine-tuning needs enough reference data to both fit and validate on. The + # default validation split leaves too few PSMs to early-stop against on a small + # reference set, which is how a fit ends up worse than the model it started + # from; the split is widened here so the stopping signal is usable. + n_reference = len(psm_list_reference) + if n_reference < MIN_FINETUNE_REFERENCE: + LOGGER.warning( + "Only %d reference PSMs. Fine-tuning measured worse than calibration " + "below about %d on held-out setups, in the worst case by sixty-fold. " + "Consider predict() with calibrate() instead.", + n_reference, + MIN_FINETUNE_REFERENCE, + ) + validation_split = max(validation_split, 0.25) + LOGGER.info( + "Using a %.0f %% validation split so early stopping has signal.", + validation_split * 100, + ) + if any(psm_list_reference["is_decoy"]): # TODO: Move to reusable validation step? LOGGER.warning("PSM list contains decoy PSMs. These will be used for fine tuning.") - training_data = DeepLCDataset.from_psm_list(psm_list_reference) + # The model is loaded further down, but the datasets have to match its feature + # specification, so peek at it first. + _peek = _model_ops.load_model( + model or DEFAULT_MODEL, device=(train_kwargs or {}).get("device") + ) + _spec = getattr(_peek, "feature_spec", None) or {} + _feature_kwargs = _feature_kwargs_from_spec(_spec) + training_data = DeepLCDataset.from_psm_list(psm_list_reference, **_feature_kwargs) validation_data = ( - DeepLCDataset.from_psm_list(psm_list_validation) if psm_list_validation else None + DeepLCDataset.from_psm_list(psm_list_validation, **_feature_kwargs) + if psm_list_validation + else None ) training_dataset, validation_dataset = split_datasets( training_data, validation_data=validation_data, validation_split=validation_split @@ -327,12 +458,56 @@ def finetune( freeze_epochs = int(train_kwargs_local.pop("freeze_epochs", 5)) train_kwargs_local.setdefault("epochs", 50) - loaded_model = _model_ops.load_model( - model or DEFAULT_MODEL, - device=train_kwargs_local.get("device"), - ) - loaded_model.add_adapter(hidden_size=adapter_hidden_size) - train_kwargs_local["freeze_epochs"] = freeze_epochs + loaded_model = _peek + if hasattr(loaded_model, "add_task_head"): + # A low-rank multitask head is adapted by fitting the new setup's own + # rank + 2 parameters with everything else frozen, rather than by training + # an adapter over the full head vector. There is nothing to unfreeze part + # way through, so freeze_epochs does not apply. + targets = psm_list_reference["retention_time"] + targets = torch.as_tensor( + np.asarray([t for t in targets if t is not None], dtype=np.float32) + ) + n_trainable = loaded_model.add_task_head(targets=targets) + + # Solve the affine part on the reference data before training. Left to the + # optimiser on a small reference set it collapses: on one 133-minute gradient + # with 230 reference peptides the output range shrank to 17 minutes and the + # error reached 91 minutes, with the correlation still above 0.9 because the + # ordering was never what broke. + if _solve_reference_affine(loaded_model, training_data, train_kwargs_local.get("device")): + LOGGER.info("Anchored the new setup's scale and shift by least squares.") + # Sixty-six parameters tolerate, and need, a far larger step than the + # whole-network default: at 1e-3 the fit is still short of its optimum after + # twenty-five epochs (1.32 min against 0.86 on a held-out setup). + train_kwargs_local.setdefault("learning_rate", 0.05) + LOGGER.info( + "Fitting %d parameters for the new setup at lr %.3g; encoder and " + "pretrained setups are frozen.", + n_trainable, + train_kwargs_local["learning_rate"], + ) + elif hasattr(loaded_model, "add_adapter"): + loaded_model.add_adapter(hidden_size=adapter_hidden_size) + train_kwargs_local["freeze_epochs"] = freeze_epochs + + # Put the adapter's output on the right axis before training. From a default + # initialisation on a small reference set the fit can collapse onto the mean + # retention time: on a 133-minute gradient with 230 reference peptides the + # output range shrank to 26 minutes and the error reached 92, with the + # correlation still above 0.9 because only the scale was lost. + if _solve_reference_affine( + loaded_model, + training_data, + train_kwargs_local.get("device"), + adapter=True, + ): + LOGGER.info("Anchored the adapter's output layer by least squares.") + else: + raise NotImplementedError( + f"{type(loaded_model).__name__} supports neither adapter-based " + "fine-tuning nor a low-rank task head." + ) finetuned_model = _model_ops.train( model=loaded_model, @@ -340,9 +515,68 @@ def finetune( validation_dataset=validation_dataset, **train_kwargs_local, ) + + _warn_if_fit_collapsed( + finetuned_model, + validation_dataset, + psm_list_reference, + device=train_kwargs_local.get("device"), + ) return finetuned_model +def _warn_if_fit_collapsed( + model: torch.nn.Module, + validation_dataset, + psm_list_reference: PSMList, + device: str | None = None, +) -> None: + """ + Say so, loudly, when a fine-tuned model is far worse than its own reference data. + + Fine-tuning can converge on a degenerate solution that predicts every peptide + near the mean retention time. The loss curve looks unremarkable and the + correlation stays high, because the ordering is preserved and only the scale is + lost, so nothing in training flags it. Comparing the validation error against the + span of the reference retention times does: a collapsed fit lands at a large + fraction of the span where a working one sits near a hundredth of it. + + This is a report, not a repair. Whether to fall back to calibration is the + caller's decision, and for the adapter path there is no untrained state worth + reverting to. + """ + observed = [t for t in psm_list_reference["retention_time"] if t is not None] + if len(observed) < 3: + return + span = float(np.max(observed) - np.min(observed)) + if span <= 0: + return + + try: + error = _model_ops.evaluate(model, validation_dataset, device=device) + except Exception: # noqa: BLE001 - the check must never break the fit + LOGGER.debug("Could not evaluate the fine-tuned model for the sanity check.") + return + + fraction = error / span + if fraction > MAX_FINETUNE_ERROR_FRACTION: + LOGGER.error( + "Fine-tuned validation error is %.2f, which is %.0f %% of the reference " + "retention-time span of %.1f. A fit this far off has collapsed rather " + "than converged: predictions are probably clustered near the mean " + "retention time. Prefer predict() with calibrate() for this dataset.", + error, + fraction * 100, + span, + ) + else: + LOGGER.info( + "Fine-tuned validation error %.3f, %.1f %% of the reference span.", + error, + fraction * 100, + ) + + def train( psm_list_reference: PSMList, psm_list_validation: PSMList | None = None, @@ -369,9 +603,13 @@ def train( Trained model. """ - training_data = DeepLCDataset.from_psm_list(psm_list_reference) + # A model trained here is new, so it gets the corrected encoding rather than + # the compatibility default the dataset applies for existing checkpoints. + training_data = DeepLCDataset.from_psm_list(psm_list_reference, legacy_positional_deltas=False) validation_data = ( - DeepLCDataset.from_psm_list(psm_list_validation) if psm_list_validation else None + DeepLCDataset.from_psm_list(psm_list_validation, legacy_positional_deltas=False) + if psm_list_validation + else None ) training_dataset, validation_dataset = split_datasets( training_data, validation_data=validation_data, validation_split=validation_split @@ -399,8 +637,17 @@ def save_model(model: torch.nn.Module, path: PathLike | str) -> None: path Destination file path. + Models that can describe themselves are saved with their architecture, + constructor arguments and feature specification alongside the weights, so + that :func:`load_model` can rebuild them. Saving a bare state dict for such a + model produced a file that could not be reloaded, because the loader would + fall back to inferring the architecture from tensor names. + """ - torch.save(model.state_dict(), path) + if hasattr(model, "describe"): + torch.save(model.describe(), path) + else: + torch.save(model.state_dict(), path) def _parse_psms(psm_list: PSMList | list[PSM | Peptidoform | str]) -> PSMList: diff --git a/deeplc/data.py b/deeplc/data.py index 3fff843..07149ce 100644 --- a/deeplc/data.py +++ b/deeplc/data.py @@ -25,6 +25,9 @@ def __init__( peptidoforms: list[Peptidoform | str], target_retention_times: np.ndarray | None = None, add_ccs_features: bool = False, + add_terminal_composition: bool = False, + padding_length: int = 60, + legacy_positional_deltas: bool = True, ): """ Initialize the DeepLCDataset. @@ -39,6 +42,31 @@ def __init__( will be set to NaN. add_ccs_features Whether to include CCS features in the encoded representation. Default is False. + add_terminal_composition + Whether to append the N- and C-terminal group composition to the global + feature vector, lengthening it from 55 to 67. Required by models trained on + that layout; see the ``feature_spec`` recorded in such a model. Default is + False. + padding_length + Length the per-position matrices are padded or truncated to. Must match the + value the model was trained with: the fused-trunk architecture pools rather + than flattens, so a mismatch changes the representation without changing any + shape and would not raise. Default is 60. + legacy_positional_deltas + Whether to place modification deltas in the positional block the way + versions before 4.0.1 did. That placement was wrong and 4.0.1 corrected + it, but every model released against this dataset class was trained on + it, so **the default is True**: a dataset exists to feed a model, and + feeding a model an encoding it was not trained on changes its + predictions on modified peptides without any error. + + Set it to False for a model trained after the correction. + :func:`deeplc.core.predict` and :func:`deeplc.core.finetune` do this + automatically from the ``feature_spec`` a self-describing checkpoint + carries, and :func:`deeplc.core.train` does it for newly trained + 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. Raises ------ @@ -50,6 +78,9 @@ def __init__( self.peptidoforms = peptidoforms self.target_retention_times = target_retention_times self.add_ccs_features = add_ccs_features + self.add_terminal_composition = add_terminal_composition + self.padding_length = padding_length + self.legacy_positional_deltas = legacy_positional_deltas if self.target_retention_times is not None and len(self.target_retention_times) != len( self.peptidoforms ): @@ -67,7 +98,11 @@ def __getitem__(self, idx: int) -> tuple[torch.Tensor, ...]: 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 + 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, ) feature_tuples = ( torch.from_numpy(features["matrix"]).to(dtype=torch.float32), @@ -87,6 +122,9 @@ def from_psm_list( cls, psm_list: PSMList, add_ccs_features: bool = False, + add_terminal_composition: bool = False, + padding_length: int = 60, + legacy_positional_deltas: bool = True, ) -> DeepLCDataset: """ Create a DeepLCDataset from a PSMList. @@ -97,6 +135,29 @@ def from_psm_list( A PSMList containing the peptidoforms and their corresponding retention times. add_ccs_features Whether to include CCS features in the encoded representation. Default is False. + add_terminal_composition + Whether to append the N- and C-terminal group composition to the global + feature vector, lengthening it from 55 to 67. Default is False. + padding_length + Length the per-position matrices are padded or truncated to. Must match the + value the model was trained with: the fused-trunk architecture pools rather + than flattens, so a mismatch changes the representation without changing any + shape and would not raise. Default is 60. + legacy_positional_deltas + Whether to place modification deltas in the positional block the way + versions before 4.0.1 did. That placement was wrong and 4.0.1 corrected + it, but every model released against this dataset class was trained on + it, so **the default is True**: a dataset exists to feed a model, and + feeding a model an encoding it was not trained on changes its + predictions on modified peptides without any error. + + Set it to False for a model trained after the correction. + :func:`deeplc.core.predict` and :func:`deeplc.core.finetune` do this + automatically from the ``feature_spec`` a self-describing checkpoint + carries, and :func:`deeplc.core.train` does it for newly trained + 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. Returns ------- @@ -114,6 +175,9 @@ def from_psm_list( peptidoforms=peptidoforms, target_retention_times=target_retention_times, add_ccs_features=add_ccs_features, + add_terminal_composition=add_terminal_composition, + padding_length=padding_length, + legacy_positional_deltas=legacy_positional_deltas, ) diff --git a/deeplc/package_data/models/multitask_flexcnn_model.pt b/deeplc/package_data/models/multitask_flexcnn_model.pt new file mode 100644 index 0000000..e6dc0c0 Binary files /dev/null and b/deeplc/package_data/models/multitask_flexcnn_model.pt differ diff --git a/tests/test_features.py b/tests/test_features.py index 776c653..f040b9b 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -302,3 +302,374 @@ def test_terminal_composition_is_opt_in_and_separates_terminal_from_side_chain() # the acetyl composition C2H2O appears in the N-terminal block only when terminal assert terminal["matrix_global"][55:61].tolist() == [2, 2, 0, 1, 0, 0] assert side_chain["matrix_global"][55:61].tolist() == [0, 0, 0, 0, 0, 0] + + +# LEGACY ENCODING, FOR MODELS TRAINED BEFORE 4.0.1 + +#: Non-zero entries of ``matrix_global`` as produced by a real v4.0.0 checkout. +#: Captured by running v4.0.0 rather than derived, so this pins the compatibility +#: path to what those models were actually trained against. Only modified +#: peptidoforms are listed, since unmodified ones encode identically either way. +V400_GLOBAL_NONZERO: dict[str, dict[int, float]] = { + "AC[UNIMOD:4]DEK": { + 0: 23, + 1: 37, + 2: 7, + 3: 10, + 4: 1, + 6: 5, + 7: 3, + 8: 5, + 9: 1, + 10: 1, + 11: 1, + 13: 6, + 14: 8, + 15: 2, + 16: 4, + 19: 5, + 20: 7, + 21: 1, + 22: 3, + 25: 6, + 26: 12, + 27: 2, + 28: 1, + 31: 3, + 32: 5, + 33: 1, + 34: 1, + 37: 3, + 38: 5, + 39: 1, + 40: 1, + 41: 1, + 43: 4, + 44: 5, + 45: 1, + 46: 3, + 49: 5, + 50: 7, + 51: 1, + 52: 3, + }, + "[UNIMOD:737]-PEPTIDEK": { + 0: 52, + 1: 83, + 2: 11, + 3: 17, + 6: 8, + 7: 18, + 8: 31, + 9: 3, + 10: 3, + 13: 4, + 14: 5, + 15: 1, + 16: 3, + 19: 5, + 20: 7, + 21: 1, + 22: 3, + 25: 6, + 26: 12, + 27: 2, + 28: 1, + 31: 5, + 32: 7, + 33: 1, + 34: 1, + 37: 5, + 38: 7, + 39: 1, + 40: 3, + 43: 5, + 44: 7, + 45: 1, + 46: 1, + 49: 4, + 50: 7, + 51: 1, + 52: 2, + }, + "PEPTIDEK-[UNIMOD:2]": { + 0: 40, + 1: 64, + 2: 10, + 3: 14, + 6: 8, + 7: 6, + 8: 11, + 9: 1, + 10: 1, + 13: 4, + 14: 5, + 15: 1, + 16: 3, + 19: 5, + 20: 7, + 21: 1, + 22: 3, + 25: 6, + 26: 12, + 27: 2, + 28: 1, + 31: 5, + 32: 7, + 33: 1, + 34: 1, + 37: 5, + 38: 7, + 39: 1, + 40: 3, + 43: 5, + 44: 7, + 45: 1, + 46: 1, + 49: 4, + 50: 8, + 51: 2, + 52: 1, + }, + "M[UNIMOD:35]EEPTIDEK": { + 0: 45, + 1: 72, + 2: 10, + 3: 19, + 4: 1, + 6: 9, + 7: 6, + 8: 11, + 9: 1, + 10: 2, + 13: 4, + 14: 5, + 15: 1, + 16: 3, + 19: 5, + 20: 7, + 21: 1, + 22: 3, + 25: 6, + 26: 12, + 27: 2, + 28: 1, + 31: 5, + 32: 9, + 33: 1, + 34: 1, + 35: 1, + 37: 5, + 38: 7, + 39: 1, + 40: 3, + 43: 5, + 44: 7, + 45: 1, + 46: 3, + 49: 5, + 50: 7, + 51: 1, + 52: 1, + }, + "PEPS[UNIMOD:21]TIDEK": { + 0: 43, + 1: 69, + 2: 10, + 3: 20, + 5: 1, + 6: 9, + 7: 6, + 8: 11, + 9: 1, + 10: 1, + 13: 4, + 14: 5, + 15: 1, + 16: 3, + 19: 5, + 20: 7, + 21: 1, + 22: 3, + 25: 6, + 26: 13, + 27: 2, + 28: 4, + 30: 1, + 31: 5, + 32: 7, + 33: 1, + 34: 1, + 37: 5, + 38: 7, + 39: 1, + 40: 3, + 43: 5, + 44: 7, + 45: 1, + 46: 1, + 49: 3, + 50: 5, + 51: 1, + 52: 2, + }, + "AC[UNIMOD:4]DEKR": { + 0: 29, + 1: 49, + 2: 11, + 3: 11, + 4: 1, + 6: 6, + 7: 4, + 8: 5, + 9: 1, + 10: 3, + 13: 7, + 14: 10, + 15: 2, + 16: 4, + 19: 6, + 20: 12, + 21: 2, + 22: 1, + 25: 6, + 26: 12, + 27: 4, + 28: 1, + 31: 3, + 32: 5, + 33: 1, + 34: 1, + 37: 3, + 38: 5, + 39: 1, + 40: 1, + 41: 1, + 43: 4, + 44: 5, + 45: 1, + 46: 3, + 49: 5, + 50: 7, + 51: 1, + 52: 3, + }, +} + + +@pytest.mark.parametrize("proforma", sorted(V400_GLOBAL_NONZERO)) +def test_legacy_positional_deltas_reproduces_v400(proforma): + """ + The compatibility path must match v4.0.0 exactly, not approximately. + + IM2Deep's CCS models and every DeepLC checkpoint from before 4.0.1 were trained + against the pre-fix placement, so any deviation here silently changes their + predictions on modified peptides. + """ + result = encode_peptidoform(proforma, legacy_positional_deltas=True)["matrix_global"] + expected = np.zeros_like(result) + for index, value in V400_GLOBAL_NONZERO[proforma].items(): + expected[index] = value + np.testing.assert_array_equal(result, expected) + + +@pytest.mark.parametrize("proforma", ["PEPTIDEK", "ACDEK", "LGEYGFQNALIVR", "A" * 70]) +def test_legacy_flag_is_a_no_op_without_modifications(proforma): + """An unmodified peptidoform has no deltas to place, so both paths agree.""" + for key, legacy in encode_peptidoform(proforma, legacy_positional_deltas=True).items(): + np.testing.assert_array_equal(legacy, encode_peptidoform(proforma)[key], err_msg=key) + + +def test_legacy_and_corrected_placement_differ_on_modified_peptidoforms(): + """Guards against the flag being silently wired to nothing.""" + legacy = encode_peptidoform("[Acetyl]-PEPTIDEK", legacy_positional_deltas=True) + corrected = encode_peptidoform("[Acetyl]-PEPTIDEK") + + assert not np.array_equal(legacy["matrix_global"], corrected["matrix_global"]) + # The defect was positional only: the per-residue matrix is unaffected, so the + # difference must be confined to the positional block of matrix_global. + np.testing.assert_array_equal(legacy["matrix"], corrected["matrix"]) + np.testing.assert_array_equal(legacy["matrix_global"][:7], corrected["matrix_global"][:7]) + + +def test_dataset_defaults_to_the_encoding_released_models_expect(): + """ + The default must be the pre-4.0.1 placement, not the corrected one. + + This is what lets a downstream package holding a model trained before the + correction keep working without changing its call. IM2Deep reaches DeepLC only + through ``from_psm_list(psm_list, add_ccs_features=True)``, so that exact call + is what is checked here. + """ + from deeplc.data import DeepLCDataset + + peptidoforms = ["AC[UNIMOD:4]DEK/2", "[UNIMOD:737]-PEPTIDEK/2"] + default = DeepLCDataset.from_psm_list(_psm_list(peptidoforms), add_ccs_features=True) + corrected = DeepLCDataset.from_psm_list( + _psm_list(peptidoforms), add_ccs_features=True, legacy_positional_deltas=False + ) + + for index, proforma in enumerate(peptidoforms): + legacy_expected = encode_peptidoform( + proforma, add_ccs_features=True, legacy_positional_deltas=True + )["matrix_global"].astype(np.float32) + # The dataset stores float32 while matrix_global is float64, so compare at + # float32 precision rather than exactly. + np.testing.assert_array_equal(default[index][0][2].numpy(), legacy_expected) + assert not np.array_equal(default[index][0][2].numpy(), corrected[index][0][2].numpy()), ( + "the corrected path must still be reachable" + ) + + +def _psm_list(peptidoforms): + """Build a PSMList over ``peptidoforms``; from_psm_list does not take strings.""" + from psm_utils import PSM, PSMList + + return PSMList( + psm_list=[ + PSM(peptidoform=Peptidoform(p), spectrum_id=str(i), retention_time=float(i)) + for i, p in enumerate(peptidoforms) + ] + ) + + +def test_undescribed_checkpoint_resolves_to_the_legacy_encoding(): + """ + Every checkpoint DeepLC has released is a bare state dict with no spec. + + Those models predate the correction, so an absent specification has to mean the + old placement or their predictions on modified peptides change silently. + """ + from deeplc.core import _feature_kwargs_from_spec + + for spec in (None, {}): + assert _feature_kwargs_from_spec(spec)["legacy_positional_deltas"] is True + + +def test_described_checkpoint_resolves_to_the_corrected_encoding(): + """A recorded specification is only written by versions that carry the fix.""" + from deeplc.core import _feature_kwargs_from_spec + + resolved = _feature_kwargs_from_spec({"padding_length": 60, "global_dim": 67}) + assert resolved["legacy_positional_deltas"] is False + assert resolved["padding_length"] == 60 + + +def test_described_checkpoint_may_request_the_legacy_encoding(): + """A model trained on the old placement can say so and be believed.""" + from deeplc.core import _feature_kwargs_from_spec + + spec = {"padding_length": 60, "legacy_positional_deltas": True} + assert _feature_kwargs_from_spec(spec)["legacy_positional_deltas"] is True + + +def test_bundled_multitask_model_declares_its_encoding(): + """The shipped model was trained after the correction, so it must say so.""" + import torch + + from deeplc.core import FLEXCNN_MULTITASK_MODEL, _feature_kwargs_from_spec + + blob = torch.load(FLEXCNN_MULTITASK_MODEL, map_location="cpu", weights_only=False) + spec = blob["feature_spec"] + assert spec["legacy_positional_deltas"] is False + assert _feature_kwargs_from_spec(spec)["legacy_positional_deltas"] is False diff --git a/tests/test_flexcnn.py b/tests/test_flexcnn.py new file mode 100644 index 0000000..161cfe0 --- /dev/null +++ b/tests/test_flexcnn.py @@ -0,0 +1,667 @@ +"""Tests for the fused-trunk multitask architecture and self-describing checkpoints.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from deeplc import _model_ops, core +from deeplc._architecture import ( + DeepLCModel, + FactorHead, + FlexCNNMultitaskModel, + InputNorm, +) +from deeplc._features import encode_peptidoform +from deeplc.data import DeepLCDataset + +MAXLEN, N_ATOMS, N_RESIDUES, PAD = 60, 6, 20, 20 +GLOBAL_DIM = 67 + +# A model small enough to build in a test but structurally identical to the shipped +# one: same modules, same forward, only narrower. +SMALL = dict( + global_dim=GLOBAL_DIM, + embed_dim=4, + channels=(8, 8), + kernel_size=5, + stem_channels=6, + stem_layers=2, + width=12, + depth=2, + rank=3, +) + + +def make_batch(lengths, seed=0): + """Structurally valid features for peptides of the given lengths.""" + rng = np.random.RandomState(seed) + batch = len(lengths) + x_atom = np.zeros((batch, MAXLEN, N_ATOMS), dtype=np.float32) + one_hot = np.zeros((batch, MAXLEN, N_RESIDUES), dtype=np.float32) + for i, length in enumerate(lengths): + x_atom[i, :length] = rng.randint(0, 12, size=(length, N_ATOMS)) + residues = rng.randint(0, N_RESIDUES, size=length) + one_hot[i, np.arange(length), residues] = 1.0 + x_global = rng.randn(batch, GLOBAL_DIM).astype(np.float32) + return ( + torch.from_numpy(x_atom), + torch.empty(0), + torch.from_numpy(x_global), + torch.from_numpy(one_hot), + ) + + +# --------------------------------------------------------------------------- # +# architecture +# --------------------------------------------------------------------------- # + + +def test_forward_returns_one_prediction_per_task(): + """Every LC setup gets a prediction for every peptide in the batch.""" + model = FlexCNNMultitaskModel(n_tasks=7, **SMALL).eval() + with torch.no_grad(): + out = model(*make_batch([9, 14, 30])) + assert out.shape == (3, 7) + assert torch.isfinite(out).all() + + +def test_task_subset_matches_full_matrix(): + """Selecting tasks must equal slicing the full output, as calibration relies on it.""" + model = FlexCNNMultitaskModel(n_tasks=9, **SMALL).eval() + batch = make_batch([12, 21]) + idx = torch.tensor([0, 4, 8]) + with torch.no_grad(): + full = model(*batch) + subset = model(*batch, task_idx=idx) + torch.testing.assert_close(full[:, idx], subset, rtol=1e-5, atol=1e-5) + + +def test_padding_does_not_change_prediction(): + """ + A peptide's prediction must not depend on how much padding follows it. + + Masking is central to this architecture, so the two inputs have to differ in + padded length for the test to mean anything: the same ten residues are placed + once in a length-20 array and once in a length-60 one. + """ + model = FlexCNNMultitaskModel(n_tasks=3, **SMALL).eval() + rng = np.random.RandomState(1) + length = 10 + atoms = rng.randint(0, 12, size=(length, N_ATOMS)).astype(np.float32) + residues = rng.randint(0, N_RESIDUES, size=length) + x_global = torch.from_numpy(rng.randn(1, GLOBAL_DIM).astype(np.float32)) + + def build(padded_to): + x_atom = np.zeros((1, padded_to, N_ATOMS), dtype=np.float32) + one_hot = np.zeros((1, padded_to, N_RESIDUES), dtype=np.float32) + x_atom[0, :length] = atoms + one_hot[0, np.arange(length), residues] = 1.0 + return torch.from_numpy(x_atom), torch.from_numpy(one_hot) + + short_atom, short_hot = build(20) + long_atom, long_hot = build(60) + with torch.no_grad(): + a = model(short_atom, torch.empty(0), x_global, short_hot) + b = model(long_atom, torch.empty(0), x_global, long_hot) + torch.testing.assert_close(a, b) + + +def test_length_one_peptide_is_handled(): + """A single residue leaves the max-pool with one valid position, not none.""" + model = FlexCNNMultitaskModel(n_tasks=3, **SMALL).eval() + with torch.no_grad(): + out = model(*make_batch([1])) + assert torch.isfinite(out).all() + + +def test_all_padding_row_does_not_produce_nan(): + """ + An empty peptide masks every position. + + The max over an entirely masked row is -inf before the guard, so this checks + the guard rather than a realistic input. + """ + model = FlexCNNMultitaskModel(n_tasks=3, **SMALL).eval() + x_atom = torch.zeros(1, MAXLEN, N_ATOMS) + one_hot = torch.zeros(1, MAXLEN, N_RESIDUES) + x_global = torch.zeros(1, GLOBAL_DIM) + with torch.no_grad(): + out = model(x_atom, torch.empty(0), x_global, one_hot) + assert torch.isfinite(out).all() + + +def test_x_atom_sum_is_ignored(): + """The fused trunk reads x_atom directly, so the rolling-sum array is unused.""" + model = FlexCNNMultitaskModel(n_tasks=4, **SMALL).eval() + x_atom, _, x_global, one_hot = make_batch([15, 25]) + with torch.no_grad(): + a = model(x_atom, torch.empty(0), x_global, one_hot) + b = model(x_atom, torch.randn(2, 30, N_ATOMS), x_global, one_hot) + torch.testing.assert_close(a, b) + + +def test_wrong_global_width_fails_loudly(): + """ + Feeding the 55-dimensional default vector must raise, not silently rescale. + + This is the failure mode worth protecting: a shape error is recoverable, a + quietly wrong retention time is not. + """ + model = FlexCNNMultitaskModel(n_tasks=3, **SMALL).eval() + x_atom, _, _, one_hot = make_batch([12]) + with pytest.raises(RuntimeError): + model(x_atom, torch.empty(0), torch.zeros(1, 55), one_hot) + + +# --------------------------------------------------------------------------- # +# encoder details +# --------------------------------------------------------------------------- # + + +def test_residue_indices_marks_padding(): + """All-zero one-hot rows are padding; argmax alone would call them residue 0.""" + encoder = FlexCNNMultitaskModel(n_tasks=2, **SMALL).encoder + one_hot = torch.zeros(1, 4, N_RESIDUES) + one_hot[0, 0, 0] = 1.0 # residue 0, genuinely + one_hot[0, 1, 7] = 1.0 + # rows 2 and 3 left empty + idx = encoder.residue_indices(one_hot) + assert idx.tolist() == [[0, 7, PAD, PAD]] + + +def test_residue_counts_ignores_padding(): + """Counts cover the twenty residues and exclude padding positions.""" + encoder = FlexCNNMultitaskModel(n_tasks=2, **SMALL).encoder + idx = torch.tensor([[3, 3, 5, PAD, PAD]]) + counts = encoder.residue_counts(idx) + assert counts.shape == (1, N_RESIDUES) + assert counts[0, 3].item() == 2 + assert counts[0, 5].item() == 1 + assert counts.sum().item() == 3 # padding contributes nothing + + +def test_input_norm_leaves_constant_features_alone(): + """ + A feature with no variance keeps raw units. + + Clamping its standard deviation to a floor would multiply any non-zero test + value by one over that floor, which is how a single phosphate once destroyed + the forward pass for a model trained without phosphorus. + """ + norm = InputNorm(3) + values = torch.tensor([[1.0, 5.0, 0.0], [1.0, 7.0, 0.0], [1.0, 9.0, 0.0]]) + norm.fit(values) + assert norm.std[0].item() == pytest.approx(1.0) + assert norm.std[2].item() == pytest.approx(1.0) + assert norm.std[1].item() > 1.0 + out = norm(torch.tensor([[1.0, 7.0, 1000.0]])) + assert out[0, 2].item() == pytest.approx(1000.0) + + +def test_factor_head_parameter_count(): + """Adding a setup costs rank + 2 parameters, which is the point of the head.""" + head = FactorHead(trunk_dim=16, n_tasks=100, rank=8) + per_task = head.embedding.shape[1] + 2 + assert per_task == 10 + shared = head.proj.weight.numel() + head.proj.bias.numel() + assert shared == 16 * 8 + 8 + + +# --------------------------------------------------------------------------- # +# self-describing checkpoints +# --------------------------------------------------------------------------- # + + +def _write_described(tmp_path, **overrides): + model = FlexCNNMultitaskModel(n_tasks=5, **SMALL) + encoder_kwargs = {k: v for k, v in SMALL.items() if k != "rank"} + blob = { + "state_dict": model.state_dict(), + "architecture": "FlexCNNMultitaskModel", + "encoder_kwargs": encoder_kwargs, + "head_kwargs": {"rank": SMALL["rank"]}, + "n_tasks": 5, + "feature_spec": { + "name": "global67_terminal", + "global_dim": GLOBAL_DIM, + "add_terminal_composition": True, + "add_ccs_features": False, + "padding_length": MAXLEN, + }, + "target_units": "minutes", + "task_names": [f"setup_{i}" for i in range(5)], + } + blob.update(overrides) + path = tmp_path / "described.pt" + torch.save(blob, path) + return model, path + + +def test_described_checkpoint_round_trips(tmp_path): + """A described checkpoint rebuilds the same model and carries its metadata.""" + original, path = _write_described(tmp_path) + loaded = _model_ops.load_model(path, device="cpu") + assert isinstance(loaded, FlexCNNMultitaskModel) + assert loaded.feature_spec["add_terminal_composition"] is True + assert loaded.target_units == "minutes" + assert loaded.task_names[0] == "setup_0" + + batch = make_batch([11, 19], seed=3) + original.eval() + with torch.no_grad(): + torch.testing.assert_close(original(*batch), loaded(*batch)) + + +def test_unknown_architecture_is_rejected(tmp_path): + """An unrecognised architecture name must fail with a clear message.""" + _, path = _write_described(tmp_path, architecture="SomeFutureModel") + with pytest.raises(ValueError, match="does not know"): + _model_ops.load_model(path, device="cpu") + + +def test_bare_state_dict_still_loads(tmp_path): + """The old checkpoint format must keep working.""" + legacy = DeepLCModel(n_heads=3) + path = tmp_path / "legacy.pt" + torch.save(legacy.state_dict(), path) + loaded = _model_ops.load_model(path, device="cpu") + assert isinstance(loaded, DeepLCModel) + assert loaded.heads.b2.shape[0] == 3 + + +# --------------------------------------------------------------------------- # +# features and the prediction path +# --------------------------------------------------------------------------- # + + +def test_terminal_composition_gives_the_expected_width(): + """The terminal block lengthens the global vector from 55 to 67.""" + without = encode_peptidoform("PEPTIDEK")["matrix_global"] + with_terminal = encode_peptidoform("PEPTIDEK", add_terminal_composition=True)["matrix_global"] + assert len(without) == 55 + assert len(with_terminal) == GLOBAL_DIM + # The shorter vector is a prefix of the longer one. + np.testing.assert_allclose(with_terminal[:55], without) + + +def test_dataset_passes_terminal_composition_through(): + """The dataset honours the flag, and still defaults to the short vector.""" + dataset = DeepLCDataset(["PEPTIDEK", "ACDEFGHIK"], add_terminal_composition=True) + features, _ = dataset[0] + assert features[2].shape == (GLOBAL_DIM,) + + default = DeepLCDataset(["PEPTIDEK"]) + features, _ = default[0] + assert features[2].shape == (55,) + + +def test_predict_builds_features_the_model_needs(tmp_path): + """ + ``predict`` must consult the model before encoding. + + The dataset default produces a 55-wide global vector, so a model needing 67 + would fail unless its feature specification is honoured. + """ + _, path = _write_described(tmp_path) + out = core.predict(["PEPTIDEK", "LGEYGFQNALIVR"], model=path, return_matrix=True) + assert out.shape == (2, 5) + assert np.isfinite(out).all() + + single = core.predict(["PEPTIDEK"], model=path) + assert single.shape == (1,) + + +def test_predictions_are_deterministic(tmp_path): + """Repeated calls on the same input return identical values.""" + _, path = _write_described(tmp_path) + first = core.predict(["PEPTIDEK", "ACDEFGHIK"], model=path, return_matrix=True) + second = core.predict(["PEPTIDEK", "ACDEFGHIK"], model=path, return_matrix=True) + np.testing.assert_array_equal(first, second) + + +# --------------------------------------------------------------------------- # +# integration: saving, device handling, and the bundled model +# --------------------------------------------------------------------------- # + + +def test_public_save_model_round_trips(tmp_path): + """ + ``save_model`` must produce a file ``predict`` can read back. + + Saving a bare state dict for this architecture produced a checkpoint the + loader could not rebuild: with no recorded architecture it fell back to + inferring one from tensor names and failed on ``heads.b2``. + """ + _, path = _write_described(tmp_path) + model = _model_ops.load_model(path, device="cpu") + + copy_path = tmp_path / "copy.pt" + core.save_model(model, copy_path) + + reloaded = _model_ops.load_model(copy_path, device="cpu") + assert isinstance(reloaded, FlexCNNMultitaskModel) + assert reloaded.feature_spec["add_terminal_composition"] is True + assert reloaded.task_names == model.task_names + + batch = make_batch([13, 22], seed=5) + with torch.no_grad(): + torch.testing.assert_close(model(*batch), reloaded(*batch)) + + # And through the public API, which is how the failure was first seen. + out = core.predict(["PEPTIDEK"], model=copy_path) + assert out.shape == (1,) + + +def test_single_column_request_does_not_evaluate_every_task(tmp_path): + """ + ``return_matrix=False`` must ask the model for one column, not all of them. + + A model trained on thousands of setups would otherwise build a column per + setup and discard all but one; at a million peptides that intermediate is + tens of gigabytes. + """ + _, path = _write_described(tmp_path) + model = _model_ops.load_model(path, device="cpu") + + seen = {} + original_forward = type(model).forward + + def spy(self, *args, task_idx=None, **kwargs): + seen["task_idx"] = task_idx + return original_forward(self, *args, task_idx=task_idx, **kwargs) + + monkey = type(model) + monkey.forward = spy + try: + single = core.predict(["PEPTIDEK", "ACDEFGHIK"], model=model) + finally: + monkey.forward = original_forward + + assert single.shape == (2,) + assert seen["task_idx"] is not None, "predict should have requested a task subset" + assert len(seen["task_idx"]) == 1 + + +def test_matrix_request_still_returns_every_task(tmp_path): + """Asking for the matrix must not be narrowed by the single-column shortcut.""" + _, path = _write_described(tmp_path) + out = core.predict(["PEPTIDEK"], model=path, return_matrix=True) + assert out.shape == (1, 5) + + +def test_requested_device_is_used_for_loading(tmp_path, monkeypatch): + """ + The device from ``predict_kwargs`` must reach ``load_model``. + + Loading first onto the default device and moving afterwards wastes a copy and + can fail with a GPU out-of-memory error for a caller who explicitly asked for + CPU. + """ + _, path = _write_described(tmp_path) + seen = {} + real_load = _model_ops.load_model + + def spy(model, device=None): + seen["device"] = device + return real_load(model, device=device) + + monkeypatch.setattr(_model_ops, "load_model", spy) + core.predict(["PEPTIDEK"], model=path, predict_kwargs={"device": "cpu"}) + assert seen["device"] == "cpu" + + +def test_add_task_head_trains_only_the_new_setup(tmp_path): + """ + Adapting to a setup must cost rank + 2 parameters and freeze everything else. + + This is the architecture's reason for existing, so the count is asserted rather + than assumed. + """ + _, path = _write_described(tmp_path) + model = _model_ops.load_model(path, device="cpu") + total = sum(p.numel() for p in model.parameters()) + + trainable = model.add_task_head(targets=torch.tensor([5.0, 10.0, 15.0, 20.0])) + assert trainable == SMALL["rank"] + 2 + assert trainable < total + assert all(not p.requires_grad for p in model.encoder.parameters()) + + # Output collapses to one column for the new setup, so the training loop and + # predict() need no special case. + with torch.no_grad(): + out = model(*make_batch([12, 20])) + assert out.shape == (2, 1) + + +def test_finetune_fits_the_low_rank_head(tmp_path): + """ + ``finetune`` adapts a fused-trunk model instead of refusing. + + It previously raised NotImplementedError for this architecture; the low-rank + head is now the adaptation path. + """ + from psm_utils import PSM, PSMList + + _, path = _write_described(tmp_path) + peptides = [ + "PEPTIDEK", + "ACDEFGHIK", + "LGEYGFQNALIVR", + "TVMENFVAFVDK", + "DAFLGSFLYEYSR", + "YICDNQDTISSK", + "SDKPDMAEIEK", + "MNDPKTLLQK", + ] + psms = PSMList( + psm_list=[ + PSM(peptidoform=f"{p}/2", spectrum_id=str(i), retention_time=float(10 + 3 * i)) + for i, p in enumerate(peptides) + ] + ) + tuned = core.finetune( + psms, + model=path, + validation_split=0.25, + train_kwargs={"epochs": 2, "device": "cpu", "show_progress": False, "batch_size": 4}, + ) + assert isinstance(tuned, FlexCNNMultitaskModel) + assert tuned.head.has_new_task + + out = core.predict(peptides[:3], model=tuned) + assert out.shape == (3,) + assert np.isfinite(out).all() + + +def test_new_task_scale_starts_at_a_trained_magnitude(tmp_path): + """ + The affine part must not be seeded from target minutes. + + ``scale`` multiplies a dot product in the normalised space the model was trained + in, roughly 0 to 100, not in minutes. Seeding it with a spread measured in + minutes overshoots by about thirtyfold; on a real setup that put the first + prediction some 600 minutes out. + """ + model = FlexCNNMultitaskModel(n_tasks=5, **SMALL) + with torch.no_grad(): + model.head.scale.fill_(0.8) + model.head.shift.fill_(3.0) + targets = torch.tensor([10.0, 40.0, 70.0, 100.0]) # std about 39 minutes + model.head.add_task(targets=targets) + assert model.head.new_scale.item() == pytest.approx(0.8, abs=1e-6) + assert model.head.new_scale.item() < targets.std().item() / 10 + + +def test_padding_length_is_taken_from_the_feature_spec(tmp_path): + """ + A recorded ``padding_length`` must reach the encoder. + + This architecture pools rather than flattens, so a mismatch changes the + representation without changing any shape and would not raise. + """ + _, path = _write_described( + tmp_path, + feature_spec={ + "name": "global67_terminal", + "global_dim": GLOBAL_DIM, + "add_terminal_composition": True, + "add_ccs_features": False, + "padding_length": 40, + }, + ) + model = _model_ops.load_model(path, device="cpu") + assert model.feature_spec["padding_length"] == 40 + + dataset = DeepLCDataset(["PEPTIDEK"], add_terminal_composition=True, padding_length=40) + features, _ = dataset[0] + assert features[0].shape[0] == 40 + + +def test_bundled_model_predicts_in_minutes(): + """ + The packaged checkpoint must load and predict plausible retention times. + + The synthetic round-trip cannot catch packaging drift: wrong metadata, a + truncated file, or target scaling left in normalised units would all pass the + other tests and fail here. + """ + path = core.FLEXCNN_MULTITASK_MODEL + if not path.exists(): # pragma: no cover - packaged with the distribution + pytest.skip("bundled model not present") + + model = _model_ops.load_model(path, device="cpu") + assert model.n_tasks == 6543 + assert model.task_names is not None + assert len(model.task_names) == 6543 + assert model.target_units == "minutes" + assert model.feature_spec["global_dim"] == GLOBAL_DIM + + peptides = ["LGEYGFQNALIVR", "TVMENFVAFVDK", "PEPTIDEK", "M[UNIMOD:35]NDPKTLLQK"] + out = core.predict(peptides, model=path, return_matrix=True) + assert out.shape == (len(peptides), 6543) + assert np.isfinite(out).all() + + # Retention times in minutes on real gradients: negative values occur on + # indexed scales, but nothing should be near a normalised 0-1 range or + # implausibly late. + assert -100.0 < out.min() < 60.0 + assert 20.0 < out.max() < 1000.0 + + single = core.predict(peptides, model=path) + np.testing.assert_allclose(single, out[:, 0], rtol=1e-5) + + +def test_small_reference_set_warns_and_widens_validation(tmp_path, caplog): + """ + A reference set too small to fine-tune on must say so. + + On held-out setups, fine-tuning below roughly five hundred reference peptides was + worse than calibration every time, once by sixty-fold, because the default + validation split left too few PSMs to early-stop against. + """ + import logging + + from psm_utils import PSM, PSMList + + _, path = _write_described(tmp_path) + peptides = [ + "PEPTIDEK", + "ACDEFGHIK", + "LGEYGFQNALIVR", + "TVMENFVAFVDK", + "DAFLGSFLYEYSR", + "YICDNQDTISSK", + "SDKPDMAEIEK", + "MNDPKTLLQK", + ] + psms = PSMList( + psm_list=[ + PSM(peptidoform=f"{p}/2", spectrum_id=str(i), retention_time=float(10 + 3 * i)) + for i, p in enumerate(peptides) + ] + ) + + with caplog.at_level(logging.WARNING, logger="deeplc.core"): + core.finetune( + psms, + model=path, + train_kwargs={"epochs": 2, "device": "cpu", "show_progress": False, "batch_size": 4}, + ) + assert any("reference PSMs" in r.getMessage() for r in caplog.records) + + +def test_adapter_output_layer_is_anchored(): + """ + Solving the adapter's output layer must put it on the retention-time axis. + + The adapter's ReLU stack is largely dead at its default initialisation, so before + this solve its output was near zero for every peptide, and the activations + reaching the output layer are rank deficient. A CUDA least-squares driver + requires full rank and returns non-finite values on such a system, which is why + the solve has to be done with a rank-tolerant method. + """ + model = DeepLCModel(n_heads=4) + model.add_adapter(hidden_size=32) + model.eval() + + lengths = [8, 12, 20, 31, 44, 7] + x_atom, x_atom_sum, x_global, one_hot = _deeplc_batch(lengths) + targets = torch.tensor([20.0, 35.0, 55.0, 80.0, 110.0, 15.0]) + + with torch.no_grad(): + before = model(x_atom, x_atom_sum, x_global, one_hot).reshape(-1) + assert before.max() - before.min() < 5.0, "expected a near-constant start" + + applied = model.solve_adapter_output(x_atom, x_atom_sum, x_global, one_hot, targets) + assert applied + + with torch.no_grad(): + after = model(x_atom, x_atom_sum, x_global, one_hot).reshape(-1) + # The solve cannot fit dead activations perfectly, but it must at least land on + # the right axis rather than near zero. + assert after.mean().item() > 5.0 + assert abs(after.mean().item() - targets.mean().item()) < 25.0 + + +def _deeplc_batch(lengths): + """Four-branch features for the released architecture.""" + rng = np.random.RandomState(2) + batch = len(lengths) + x_atom = np.zeros((batch, MAXLEN, N_ATOMS), dtype=np.float32) + x_atom_sum = np.zeros((batch, 30, N_ATOMS), dtype=np.float32) + one_hot = np.zeros((batch, MAXLEN, N_RESIDUES), dtype=np.float32) + for i, length in enumerate(lengths): + x_atom[i, :length] = rng.randint(0, 12, size=(length, N_ATOMS)) + x_atom_sum[i, : max(1, length // 2)] = rng.randint( + 0, 20, size=(max(1, length // 2), N_ATOMS) + ) + residues = rng.randint(0, N_RESIDUES, size=length) + one_hot[i, np.arange(length), residues] = 1.0 + # The four-branch model reads the 55-dimensional global vector, not the 67- + # dimensional one the fused trunk needs. + x_global = rng.randn(batch, 55).astype(np.float32) + return ( + torch.from_numpy(x_atom), + torch.from_numpy(x_atom_sum), + torch.from_numpy(x_global), + torch.from_numpy(one_hot), + ) + + +def test_training_scores_its_starting_point(tmp_path): + """ + Training must not return a model worse than the one it started from. + + With the best validation loss left at infinity the first epoch always became the + best, however bad, so a fine-tune on a small reference set could hand back a + collapsed fit at ninety times the error of its own starting point. + """ + import inspect + + from deeplc import _model_ops + + 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