diff --git a/deeplc/_composition_architecture.py b/deeplc/_composition_architecture.py new file mode 100644 index 0000000..0372379 --- /dev/null +++ b/deeplc/_composition_architecture.py @@ -0,0 +1,215 @@ +""" +Composition-encoder multitask architecture. + +An alternative to :class:`deeplc._architecture.MultitaskDeepLCModel`, differing +in three ways that were each measured: + +``pointwise stem`` + A kernel-1 convolution decodes each position's six atom counts on its own, + before any neighbour mixing. The main convolutions use kernel 5 and would + otherwise have to disentangle residue identity and sequence context at the + same time. +``low-rank read-out`` + Each LC setup is a 64-dimensional vector dotted with a projected trunk, + rather than an independent head. A new setup costs 66 parameters, and a + calibration fitted in that space converges to the trained values rather + than approximating them. +``corrected global features`` + Requires ``matrix_global`` with the terminal-composition blocks, that is + ``encode_peptidoform(..., add_terminal_composition=True)``. Hence + :attr:`CompositionMultitaskModel.requires_terminal_composition`. + +``matrix_sum`` is accepted for signature compatibility and ignored: it equals +``matrix.reshape(30, 2, 6).sum(1)`` exactly, so it carries nothing the atom +matrix does not already provide. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +from torch import nn + +PAD_INDEX = 20 +N_ATOMS = 6 + + +class InputNorm(nn.Module): + """Standardise dense inputs with statistics fitted on the training set.""" + + def __init__(self, n_features: int): + super().__init__() + self.register_buffer("mean", torch.zeros(n_features)) + self.register_buffer("std", torch.ones(n_features)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Standardise ``x`` with the stored statistics.""" + return (x - self.mean) / self.std + + +class ConvBlock(nn.Module): + """One convolution stage: convolution then 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: + """Apply the convolution and activation.""" + return self.act(self.conv(x)) + + +class FactorHead(nn.Module): + """ + Low-rank read-out: a per-setup vector dotted with a projected trunk. + + ``prediction[:, j] = (proj(trunk) . embedding[j]) * scale[j] + shift[j]`` + """ + + def __init__(self, trunk_dim: int, n_tasks: int, rank: int = 64): + super().__init__() + 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 forward(self, trunk: torch.Tensor) -> torch.Tensor: + """Return one prediction per LC setup.""" + return self.proj(trunk) @ self.embedding.t() * self.scale + self.shift + + +class CompositionEncoder(nn.Module): + """Peptide encoder: pointwise stem, convolution stack, masked pooling, dense trunk.""" + + def __init__( + self, + global_dim: int = 67, + embed_dim: int = 16, + channels: Sequence[int] = (512, 512), + kernel_size: int = 5, + stem_pointwise: int = 128, + stem_layers: int = 2, + width: int = 256, + depth: int = 3, + ): + super().__init__() + self.embed = nn.Embedding(PAD_INDEX + 1, embed_dim, padding_idx=PAD_INDEX) + self.stem = nn.Sequential( + *[ + ConvBlock(N_ATOMS if i == 0 else stem_pointwise, stem_pointwise, 1) + for i in range(stem_layers) + ] + ) + + in_channels = stem_pointwise + embed_dim + blocks = [] + for out_channels in channels: + blocks.append(ConvBlock(in_channels, out_channels, kernel_size)) + in_channels = out_channels + self.blocks = nn.ModuleList(blocks) + + pooled_dim = in_channels * 2 # sum and max pooling + self.pool_norm = nn.LayerNorm(pooled_dim) + dense_dim = global_dim + PAD_INDEX # global features and residue counts + self.norm = InputNorm(dense_dim) + + layers: list[nn.Module] = [] + previous = dense_dim + pooled_dim + for _ in range(depth): + layers.extend([nn.Linear(previous, width), nn.SiLU()]) + previous = width + self.net = nn.Sequential(*layers) + self.trunk_dim = width + + def forward( + self, x_atom: torch.Tensor, x_global: torch.Tensor, residue_index: torch.Tensor + ) -> torch.Tensor: + """Encode a batch of peptides into trunk vectors.""" + valid = (residue_index != PAD_INDEX).unsqueeze(1) + + hidden = torch.cat( + [self.stem(x_atom.transpose(1, 2)), self.embed(residue_index).transpose(1, 2)], dim=1 + ) + for block in self.blocks: + hidden = block(hidden) + hidden = hidden * valid + + pooled = torch.cat( + [ + hidden.sum(dim=2), + torch.nan_to_num( + hidden.masked_fill(~valid, float("-inf")).max(dim=2).values, neginf=0.0 + ), + ], + dim=1, + ) + pooled = self.pool_norm(pooled) + + counts = ( + nn.functional.one_hot(residue_index.clamp(0, PAD_INDEX), PAD_INDEX + 1) + .sum(1)[:, :PAD_INDEX] + .float() + ) + dense = self.norm(torch.cat([x_global.float(), counts], dim=1)) + return self.net(torch.cat([dense, pooled], dim=1)) + + +class CompositionMultitaskModel(nn.Module): + """ + Composition encoder with a low-rank multitask read-out. + + Takes the same four feature tensors as :class:`DeepLCModel` so it is a drop-in + for the prediction path, and returns ``(batch, n_tasks)`` predictions. + """ + + #: ``matrix_global`` must carry the terminal-composition blocks for this model. + requires_terminal_composition = True + + def __init__( + self, n_tasks: int = 1025, rank: int = 64, global_dim: int = 67, **encoder_kwargs + ): + super().__init__() + self.encoder = CompositionEncoder(global_dim=global_dim, **encoder_kwargs) + self.head = FactorHead(self.encoder.trunk_dim, n_tasks, rank=rank) + + def forward( + self, + x_atom: torch.Tensor, + x_atom_sum: torch.Tensor, # noqa: ARG002 - accepted for signature compatibility + x_global: torch.Tensor, + x_one_hot: torch.Tensor, + ) -> torch.Tensor: + """Predict retention time for every LC setup the model was trained on.""" + residue_index = x_one_hot.argmax(dim=2).long() + residue_index = residue_index.masked_fill(x_one_hot.sum(dim=2) == 0, PAD_INDEX) + return self.head(self.encoder(x_atom, x_global, residue_index)) + + @classmethod + def from_state_dict(cls, state: dict) -> CompositionMultitaskModel: + """Rebuild the architecture from the shapes stored in a state dict.""" + channels = [ + state[key].shape[0] + for key in state + if key.startswith("encoder.blocks.") and key.endswith(".conv.weight") + ] + stem_layers = sum( + 1 for key in state if key.startswith("encoder.stem.") and key.endswith(".conv.weight") + ) + model = cls( + n_tasks=state["head.embedding"].shape[0], + rank=state["head.embedding"].shape[1], + global_dim=state["encoder.norm.mean"].shape[0] - PAD_INDEX, + embed_dim=state["encoder.embed.weight"].shape[1], + channels=channels, + kernel_size=state["encoder.blocks.0.conv.weight"].shape[2], + stem_pointwise=state["encoder.stem.0.conv.weight"].shape[0], + stem_layers=stem_layers, + width=state["encoder.net.0.weight"].shape[0], + depth=sum( + 1 for key in state if key.startswith("encoder.net.") and key.endswith(".weight") + ), + ) + model.load_state_dict(state) + return model diff --git a/deeplc/_model_ops.py b/deeplc/_model_ops.py index 2f5d777..f719b17 100644 --- a/deeplc/_model_ops.py +++ b/deeplc/_model_ops.py @@ -18,6 +18,7 @@ from torch.utils.data import DataLoader, Dataset, Subset from deeplc._architecture import DeepLCModel +from deeplc._composition_architecture import CompositionMultitaskModel from deeplc.data import DeepLCDataset logger = logging.getLogger(__name__) @@ -38,6 +39,11 @@ def load_model( # 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) + if "head.embedding" in raw: + # Composition encoder with a low-rank read-out; shapes are self-describing. + loaded_model = CompositionMultitaskModel.from_state_dict(raw) + loaded_model.to(selected_device) + return loaded_model 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,14 +52,16 @@ 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, torch.nn.Module): + # Any nn.Module is accepted so alternative architectures, such as + # CompositionMultitaskModel, can be passed in already loaded. loaded_model = model logger.debug("Using provided PyTorch model instance") elif model is None: loaded_model = DeepLCModel(n_heads=1) logger.debug("Initialized new DeepLCModel with default architecture") else: - raise TypeError(f"Expected a DeepLCModel or a file path, got {type(model)} instead.") + raise TypeError(f"Expected a torch.nn.Module or a file path, got {type(model)} instead.") loaded_model.to(selected_device) diff --git a/deeplc/core.py b/deeplc/core.py index 1319e54..5426f98 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -53,9 +53,14 @@ def predict( produces multitask output, in which case shape is ``(n, n_heads)``. """ + # Load first: the feature encoding a model needs is a property of the model. + loaded_model = _model_ops.load_model(model or DEFAULT_MODEL) result = _model_ops.predict( - model=model or DEFAULT_MODEL, - data=DeepLCDataset.from_psm_list(_parse_psms(psm_list)), + model=loaded_model, + data=DeepLCDataset.from_psm_list( + _parse_psms(psm_list), + add_terminal_composition=getattr(loaded_model, "requires_terminal_composition", False), + ), **(predict_kwargs or {}), ).numpy() if not return_matrix: diff --git a/deeplc/data.py b/deeplc/data.py index 3fff843..6e4a70b 100644 --- a/deeplc/data.py +++ b/deeplc/data.py @@ -25,6 +25,7 @@ def __init__( peptidoforms: list[Peptidoform | str], target_retention_times: np.ndarray | None = None, add_ccs_features: bool = False, + add_terminal_composition: bool = False, ): """ Initialize the DeepLCDataset. @@ -39,6 +40,10 @@ 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 compositions to ``matrix_global``. + Required by models whose ``requires_terminal_composition`` attribute is True. + Default is False. Raises ------ @@ -50,6 +55,7 @@ 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 if self.target_retention_times is not None and len(self.target_retention_times) != len( self.peptidoforms ): @@ -67,7 +73,9 @@ 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, ) feature_tuples = ( torch.from_numpy(features["matrix"]).to(dtype=torch.float32), @@ -87,6 +95,7 @@ def from_psm_list( cls, psm_list: PSMList, add_ccs_features: bool = False, + add_terminal_composition: bool = False, ) -> DeepLCDataset: """ Create a DeepLCDataset from a PSMList. @@ -97,6 +106,9 @@ 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 terminal group compositions to ``matrix_global``. + Default is False. Returns ------- @@ -114,6 +126,7 @@ def from_psm_list( peptidoforms=peptidoforms, target_retention_times=target_retention_times, add_ccs_features=add_ccs_features, + add_terminal_composition=add_terminal_composition, ) diff --git a/deeplc/package_data/models/composition_multitask_model.pt b/deeplc/package_data/models/composition_multitask_model.pt new file mode 100644 index 0000000..26ad3fb Binary files /dev/null and b/deeplc/package_data/models/composition_multitask_model.pt differ diff --git a/tests/test_composition_architecture.py b/tests/test_composition_architecture.py new file mode 100644 index 0000000..2d26a60 --- /dev/null +++ b/tests/test_composition_architecture.py @@ -0,0 +1,73 @@ +"""Tests for the composition-encoder multitask model.""" + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from deeplc._composition_architecture import CompositionMultitaskModel +from deeplc._model_ops import load_model +from deeplc.core import predict + +MODEL_PATH = Path(__file__).parent.parent / "deeplc" / "package_data" / "models" / ( + "composition_multitask_model.pt" +) + + +def test_model_declares_its_feature_requirement(): + """The model needs terminal composition, and says so.""" + model = CompositionMultitaskModel(n_tasks=4) + assert model.requires_terminal_composition is True + + +def test_forward_shape_and_ignored_sum_matrix(): + """Output is one prediction per LC setup; matrix_sum does not affect it.""" + model = CompositionMultitaskModel(n_tasks=7).eval() + x_atom = torch.rand(3, 60, 6) + x_global = torch.rand(3, 67) + one_hot = torch.zeros(3, 60, 20) + one_hot[:, :9, 5] = 1.0 + + with torch.no_grad(): + first = model(x_atom, torch.zeros(3, 30, 6), x_global, one_hot) + second = model(x_atom, torch.rand(3, 30, 6), x_global, one_hot) + + assert first.shape == (3, 7) + assert torch.allclose(first, second), "matrix_sum must not change the prediction" + + +@pytest.mark.skipif(not MODEL_PATH.exists(), reason="packaged model not available") +def test_packaged_model_loads_and_predicts(): + """The shipped checkpoint round-trips through load_model and predict.""" + model = load_model(MODEL_PATH, device="cpu") + assert isinstance(model, CompositionMultitaskModel) + assert model.requires_terminal_composition is True + + peptides = ["PEPTIDEK", "[Acetyl]-PEPTIDEK", "ELVISLIVESK", "M[Oxidation]ACGHTR"] + predictions = predict( + peptides, + model=MODEL_PATH, + return_matrix=True, + predict_kwargs={"show_progress": False, "device": "cpu"}, + ) + + assert predictions.shape[0] == len(peptides) + assert predictions.shape[1] > 1, "multitask model should return one column per LC setup" + assert np.isfinite(predictions).all() + + +@pytest.mark.skipif(not MODEL_PATH.exists(), reason="packaged model not available") +def test_terminal_modification_changes_the_prediction(): + """An N-terminal acetyl must not predict identically to the unmodified peptide. + + This is what the terminal-composition features exist for; without them the two + encode the same and the model cannot tell them apart. + """ + predictions = predict( + ["PEPTIDEK", "[Acetyl]-PEPTIDEK"], + model=MODEL_PATH, + return_matrix=True, + predict_kwargs={"show_progress": False, "device": "cpu"}, + ) + assert not np.allclose(predictions[0], predictions[1])