diff --git a/CHANGELOG.md b/CHANGELOG.md index ee6e03c..f241202 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **MALDI model.** `models/TIMS/IM2DeepMALDI.ckpt` (`constants.MALDI_MODEL`, CLI + `--model-name maldi`) is bundled next to the default model. It has the default model's + architecture and training config, but its 5.06M-precursor training set adds ProteomeTools + HLA, AspN and LysN pools and HLA-II peptides acquired without the charge 1 exclusion, + for more long charge 1 ions: 4,816 / 676 / 40 charge 1 precursors of 12-13 / 14-15 / + 16-20 amino acids, up from 2,654 / 221 / 15. On 276 held-out long charge 1 sequences + the calibrated MAE falls from 8.27 to 4.18 Ų (charge 1 alone: 8.89 to 4.34). On the + default model's held-out set it is slightly worse, 6.28 against 6.14 Ų (charge 1: + 4.00 against 3.72), so it does not replace the default. Its training data is ESI + timsTOF data; it has not been validated on MALDI-TIMS measurements. It shares + `reference_ccs_2026.csv.gz`, whose per-charge offset to its predictions is within + 2.6 Ų. It has no multi-conformer variant - `core.train()` and `core.finetune()`, replacing the `NotImplementedError` stub that pointed at the separate `im2deeptrainer` package. Both featurise through `DeepLCDataset`, the same path `predict()` uses, so a trained checkpoint is directly @@ -31,8 +43,54 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 only in training data or featurisation is comparable. `--wandb`, `--wandb-project` and `--wandb-name` on both CLI commands; `entity` and `tags` are settable from a config file +- Transformer trunk (`use_transformer_trunk`) in place of the convolutional atom/diatom + branches: a learned charge embedding prepended to the sequence as a register token + that every residue attends to, pre-norm self-attention blocks with a learned + relative-distance and terminus-proximity attention bias, and learned-query attention + pooling (`transformer_*` keys). This is the architecture of the new bundled model +- Training options used by the new model: `standardize_global_features` and + `standardize_target` (the fitted statistics are stored in the checkpoint), + `swa` (stochastic weight averaging), `lr_scheduler="one_cycle"`, `weight_decay`, + `precision`. Also available, off by default: `use_charge_embedding`, `use_flex_trunk`, + `use_attention_pooling`, `use_per_charge_heads`, `use_heteroscedastic_head`, + `charge_balanced_sampling`, `mixup_alpha`, `sibling_loss_weight`, + `ranking_loss_weight`, `loss` (`"l1"`, `"mape"` or `"huber"`) +- `RegressionCCSCalibration`, a per-charge slope and intercept calibration for models + whose training labels were harmonised with a regression rather than a shift. Both + calibrations expose `stats_for_charge()` (anchor count, whether the charge fell back + to the all-charge fit, residual spread) for judging how far to trust a run's + calibrated CCS +- `scikit-learn` dependency, used by `RegressionCCSCalibration` ### Changed +- **New default single-conformer model.** `models/TIMS/IM2DeepUni.ckpt` is replaced by + a 4-layer transformer (1.05M parameters) trained on a 4.83M-precursor collated set, + up from 1.12M, whose cross-run CCS values are calibrated per charge. On the same + 477,896-row held-out set the calibrated MAE falls from 8.92 to 6.16 Ų + (−30.9%), and on four unseen diaPASEF studies DIA-NN gains more precursors + and protein groups than with either the old model or its own predicted ion + mobility. The checkpoint records its own architecture and featurisation (72 global + features, DeepLC's corrected positional deltas), which `predict()` and `finetune()` + read back. Inference is slower per precursor: 1.28× on 16 CPU cores, 1.07× on + one GPU +- **New default calibration reference.** `reference_data/reference_ccs_2026.csv.gz` replaces + `reference_ccs.csv.gz` as the default (`DEFAULT_REFERENCE_DATASET_PATH`). It holds + 1,054,960 precursors sampled from the new model's training set. The two change together + and have to: calibration shifts are `observed CCS - reference CCS`, which only holds when + the reference is on the model's CCS scale, and the old reference sits 14-25 Ų + (about 3.7%) below it. On the same held-out set the new model's calibrated MAE is + 6.16 Ų against the new reference but 18.09 Ų against the old one, + worse than the previous model's 8.97. If you calibrate another model against the default + reference, pass a reference on that model's scale instead +- The new reference writes modifications as UNIMOD accessions (`[UNIMOD:4]`) rather than + names (`[Carbamidomethyl]`). Calibration matches peptidoforms on their ProForma string, + so calibration PSMs carrying modification names no longer match the 29% of the reference + that is modified; unmodified peptides match in either notation +- `predict()` now defaults to a GPU when one is available instead of always using the CPU +- Early stopping is now off by default for `core.train()` (`patience` is `None`), as it + was found to fire prematurely. Set `patience` to opt back in. The `normal`, `xavier` + and `kaiming` weight initialisers now use explicit parameters (`std=0.05`, `gain=1`, + the configured LeakyReLU slope) instead of PyTorch's defaults - Widened the `deeplc` dependency to `>=4.1.0,<5`. From 4.1.0 the 4.0.1 feature-encoding change sits behind `legacy_positional_deltas`, which defaults to `True`, so the encoding matches what the bundled checkpoints were trained with. Verified @@ -44,6 +102,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 undeclared import ### Fixed +- `finetune()` raised `AttributeError` on a transformer-trunk (or FlexTrunk) backbone, + which the new bundled model is. The transfer model now delegates to the backbone, so + the trunk, the charge register token and both standardisation scalers are applied and + stay fixed, and `BackboneFreeze` covers the trunk +- The Weights & Biases `Best Val MAE` log call no longer fails on non-zero ranks under + multi-device training - `IM2DeepTransfer`'s training, validation and test steps did not squeeze the model output in the `add_X_mol` branch, so the loss broadcast to a `(batch, batch)` matrix instead of comparing element-wise diff --git a/README.md b/README.md index 71803be..0b6c3b4 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,17 @@ To use the multi-output prediction model (requires optional dependencies): im2deep --calibration-precursors --multi ``` +**MALDI model:** +MALDI ions are mostly singly charged, including long peptides that ESI rarely yields as +charge 1. `--model-name maldi` selects a model trained on a set with more of these +(12-20 amino acids, charge 1). It is trained on ESI timsTOF data and has not been +validated on MALDI-TIMS measurements. It is calibrated the same way as the default model: +```bash +im2deep --calibration-precursors --model-name maldi +``` +In Python, pass `model=im2deep.constants.MALDI_MODEL` to `predict()` or +`predict_and_calibrate()`. + **Output options:** ```bash im2deep --output-file predictions.csv @@ -119,6 +130,7 @@ IPQEKCILQTDVK,5|Butyryl|6|Carbamidomethyl,3,516.2079366048176 ## Important Notes - **Calibration**: Highly recommended for accurate predictions. Calibration corrects for systematic differences between predicted and observed CCS values. +- **Calibration reference**: The shift is fitted against a bundled reference dataset, which is on the CCS scale of the bundled model. Calibrating a different model (for example one you trained) against it is only valid if that model predicts on the same scale, otherwise pass a reference built from that model's own training data. Peptidoforms are matched on their ProForma string, and the reference writes modifications as UNIMOD accessions (e.g. `[UNIMOD:4]`). Peptides written with modification names (e.g. `[Carbamidomethyl]`, as peprec-style input is read) therefore only match the reference when unmodified. - **Charge states**: IM2Deep predictions are reliable for charge states up to z=6. PSMs with higher charge states are automatically filtered out during validation. ## Citing diff --git a/docs/source/training.rst b/docs/source/training.rst index c4b9dce..4738bf3 100644 --- a/docs/source/training.rst +++ b/docs/source/training.rst @@ -69,7 +69,9 @@ sit on a systematically different CCS scale than the training data: The pretrained feature branches are held frozen for the first ``--freeze-epochs`` epochs while the head adapts, then unfrozen and training continues at a reduced learning rate. By default the bundled IM2Deep model is the backbone; pass -``--backbone`` to use another checkpoint. +``--backbone`` to use another checkpoint. The backbone's architecture and +featurisation are read from the checkpoint, and its feature standardisation stays +fixed while you fine-tune. Configuration ------------- @@ -129,9 +131,11 @@ to start otherwise rather than failing later with a shape error. - 72 ``legacy_positional_deltas`` controls positional modification indexing. It -defaults to ``true``, reproducing the encoding the bundled models were trained -with. Setting it to ``false`` uses DeepLC's corrected indexing, which changes the -encoding of modified peptides only. +defaults to ``true``, reproducing the encoding the bundled multi-conformer models +were trained with. Setting it to ``false`` uses DeepLC's corrected indexing, which +changes the encoding of modified peptides only. The bundled single-conformer model +was trained with ``false`` and 72 global features (``add_terminal_composition``). +It records this, so ``predict`` and ``finetune`` pick it up without any settings. Each trained checkpoint records the configuration it was trained with, so :func:`im2deep.core.predict` reads a model back with the right architecture and diff --git a/im2deep/__main__.py b/im2deep/__main__.py index dbf1d23..5c677a9 100644 --- a/im2deep/__main__.py +++ b/im2deep/__main__.py @@ -46,6 +46,10 @@ setup_logging, write_output, ) +from im2deep.constants import DEFAULT_MODEL, MALDI_MODEL + +#: Bundled single-conformer checkpoints selectable with --model-name. +MODELS = {"tims": DEFAULT_MODEL, "maldi": MALDI_MODEL} console = Console() LOGGER = logging.getLogger(__name__) @@ -119,9 +123,11 @@ def cli(ctx, logging_level, profile, profile_name): @click.option( "-m", "--model-name", - type=click.Choice(["tims"], case_sensitive=False), + type=click.Choice(list(MODELS), case_sensitive=False), default="tims", - help="Neural network model to use for prediction. Currently only 'tims' is supported.", + help="Bundled model to use for prediction. 'tims' is the default model; 'maldi' is " + "trained on more long (12-20 AA) charge 1 peptides, for MALDI-TIMS. 'maldi' has no " + "multi-conformer variant.", ) @click.option( "-e", @@ -214,6 +220,11 @@ def _run_predict(*args, **kwargs): f"ion_mobility={kwargs.get('ion_mobility')}" ) + model_name = kwargs.pop("model_name", "tims").lower() + if model_name != "tims" and kwargs.get("multi", False): + raise click.UsageError(f"--model-name {model_name} cannot be combined with --multi.") + model = MODELS[model_name] + # Parse input files LOGGER.info("Parsing input files...") psm_list = parse_input(Path(kwargs["precursors"])) @@ -223,12 +234,15 @@ def _run_predict(*args, **kwargs): if kwargs.get("calibration_precursors"): LOGGER.info("Calibration file provided, performing calibration and prediction...") psm_list_cal = parse_input(Path(kwargs["calibration_precursors"])) - predictions = core.predict_and_calibrate(psm_list, psm_list_cal, *args, **kwargs) + predictions = core.predict_and_calibrate( + psm_list, psm_list_cal, *args, model=None if kwargs.get("multi") else model, **kwargs + ) else: LOGGER.info( "No calibration file provided (calibration is HIGHLY recommended), performing prediction only..." ) - predictions = core.predict(psm_list, multi=kwargs.get("multi", False)) + multi = kwargs.get("multi", False) + predictions = core.predict(psm_list, model=None if multi else model, multi=multi) # Output results LOGGER.info("IM2Deep CCS prediction completed successfully!") diff --git a/im2deep/_architectures/blocks.py b/im2deep/_architectures/blocks.py index 272168a..670b03a 100644 --- a/im2deep/_architectures/blocks.py +++ b/im2deep/_architectures/blocks.py @@ -1,5 +1,6 @@ """Building blocks for the IM2Deep architecture.""" +import torch from torch import nn from im2deep._architectures.activations import LRelu_with_saturation @@ -23,7 +24,7 @@ def __init__( negative_slope=negative_slope, saturation=saturation ) - initializer(self.conv.weight, 0.0, 0.05) + initializer(self.conv.weight) def forward(self, x): return self.activation(self.conv(x)) @@ -38,7 +39,162 @@ def __init__(self, in_features, out_features, initializer, negative_slope, satur negative_slope=negative_slope, saturation=saturation ) - initializer(self.linear.weight, 0.0, 0.05) + initializer(self.linear.weight) def forward(self, x): return self.activation(self.linear(x)) + + +class AttentionPool1d(nn.Module): + """ + Learned attention pooling over the position axis: ``(B, C, L) -> (B, C)``. + + Step 8.4 item 2's replacement for the position-index-dependent + ``nn.Flatten()`` used elsewhere in the architecture. A 1x1 conv produces + one raw score per position, softmax-normalised over positions (``L``), + then a weighted sum over ``L`` - a small, standard attention-pooling + design. Unlike ``Flatten()``, the output width depends only on the + channel count, not on the (``padding_length``-derived) sequence length, + so a branch ending in this is length-agnostic rather than hard-wired to + whatever ``padding_length`` produced. + """ + + def __init__(self, channels, initializer): + super().__init__() + self.score = nn.Conv1d(channels, 1, kernel_size=1) + initializer(self.score.weight) + + def forward(self, x): + weights = torch.softmax(self.score(x), dim=-1) # (B, 1, L) + return (x * weights).sum(dim=-1) # (B, C) + + +class _ConvSiLU(nn.Module): + """A single ``Conv1d`` with ``same`` padding followed by SiLU, matching + DeepLC's ``_ConvSiLU`` (``deeplc/_architecture.py``) exactly - no + LeakyReLU/saturation, no learnable-initializer hook, since this block is + only ever used inside :class:`FlexTrunk`, ported to be faithful to the + reference implementation rather than to match IM2Deep's other blocks.""" + + def __init__(self, in_channels, out_channels, kernel_size): + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding="same") + self.act = nn.SiLU() + + def forward(self, x): + return self.act(self.conv(x)) + + +class FlexTrunk(nn.Module): + """ + Fused convolutional trunk over atomic composition and a learned residue + embedding, with length-masked sum+max pooling (step 8.4 items 8+9). + + Ported from DeepLC's ``_FlexCNNEncoder`` (``deeplc/_architecture.py``, + lines 988-1079), whose own docstring claims roughly a third of the + four-branch backbone's error on DeepLC's retention-time task, on the + same head. Replaces ``ConvAtomComp``, ``ConvDiatomComp`` and ``OneHot`` + together when enabled: residue identity reaches the network through an + embedding fused with atom counts *before* any neighbour-mixing + convolution, rather than through a separate, late-concatenated 92- + parameter one-hot branch, and ``matrix_sum`` (``ConvDiatomComp``'s input) + is redundant once the trunk sees ``matrix`` directly - DeepLC's own + docstring calls it exactly that. + + Deliberately does **not** port DeepLC's ``InputNorm``/``residue_counts``/ + final dense ``self.net`` - those fuse the global feature branch and a + residue-count vector into this same module, which is step 8.4 item 10, + not this one. IM2Deep keeps ``ConvGlobal`` as its own already-tuned + branch feeding the shared ``Concat`` head, so this module's output is + just the pooled sequence feature vector - the same role + ``torch.cat((atom_comp, diatom_comp), 1)`` played before, sized by + :func:`im2deep._architectures.helpers.calculate_seq_concat_shape`. + + Channel widths (``channels``, ``stem_channels``) are **not** copied from + DeepLC's own RT-tuned defaults (512/512 channels, 128-wide stem) - those + were tuned for a different task and a differently-sized encoder; IM2Deep + sizes them down to stay proportionate to its own current best config's + scale (~1.5-2M total parameters). Only ``embed_dim`` (16) and + ``stem_layers``/``kernel_size`` are carried over unchanged, since those + are small, task-agnostic choices rather than capacity knobs. + """ + + #: Index reserved for padding positions, matching the one-hot vocabulary + #: size (20 canonical amino acids, indices 0-19). A residue DeepLC's own + #: featurisation could not encode (e.g. selenocysteine - see the + #: "One-hot skip" warnings) is an all-zero one-hot row, indistinguishable + #: from real padding by this same detection - an approximation inherited + #: unchanged from DeepLC's own ``residue_indices``, not introduced here. + PAD_INDEX = 20 + + def __init__( + self, + embed_dim: int, + channels: "tuple[int, ...] | list[int]", + kernel_size: int, + stem_channels: int, + stem_layers: int, + ): + super().__init__() + self.embed = nn.Embedding(self.PAD_INDEX + 1, embed_dim, padding_idx=self.PAD_INDEX) + self.stem = self._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) + self.out_channels = in_channels + + # Sum (length-extensive) and max (length-intensive) pooling, masked + # to the valid length - item 9. Concatenated so the trunk carries + # both a peptide-size-sensitive and a peptide-size-insensitive view. + self.pool_norm = nn.LayerNorm(in_channels * 2) + + @staticmethod + def _pointwise_stem(channels: int, layers: int) -> nn.Sequential: + """ + Width-1 convolutions decoding each position's atom counts alone, + before any neighbour mixing - atomic composition determines the + residue uniquely except for the Leu/Ile pair, which is a lookup + rather than a linear map, so one layer is generally not enough. + """ + return nn.Sequential( + *[_ConvSiLU(6 if i == 0 else channels, channels, kernel_size=1) for i in range(layers)] + ) + + def residue_indices(self, one_hot: torch.Tensor) -> torch.Tensor: + """ + Integer residue index per position from an already-permuted + ``(batch, 20, length)`` one-hot tensor. Padding positions are + all-zero columns, for which ``argmax`` would return 0 and collide + with the first residue, so they are set to :attr:`PAD_INDEX` + explicitly - the same fix DeepLC's own version applies. + """ + idx = one_hot.argmax(dim=1) + return idx.masked_fill(one_hot.sum(dim=1) == 0, self.PAD_INDEX) + + def forward(self, atom_comp: torch.Tensor, one_hot: torch.Tensor) -> torch.Tensor: + """ + ``atom_comp``: ``(batch, 6, length)``; ``one_hot``: ``(batch, 20, length)`` + - both already permuted to Conv1d's channels-first layout by the + caller, matching every other branch in this architecture. + """ + idx = self.residue_indices(one_hot) + valid = (idx != self.PAD_INDEX).unsqueeze(1) # (batch, 1, length) + + hidden = atom_comp + if self.stem is not None: + hidden = self.stem(hidden) + hidden = torch.cat([hidden, 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 + ) + return self.pool_norm(torch.cat([summed, maxed], dim=1)) diff --git a/im2deep/_architectures/callbacks.py b/im2deep/_architectures/callbacks.py index 1699b3d..0ab17ea 100644 --- a/im2deep/_architectures/callbacks.py +++ b/im2deep/_architectures/callbacks.py @@ -14,7 +14,15 @@ #: Attributes of a transfer model that hold the pretrained feature branches. #: The concatenation head is deliberately excluded: it is what fine-tuning is #: meant to adapt. -BACKBONE_BRANCHES = ("ConvAtomComp", "ConvDiatomComp", "ConvGlobal", "OneHot", "MolDesc") +BACKBONE_BRANCHES = ( + "ConvAtomComp", + "ConvDiatomComp", + "transformer_trunk", + "flex_trunk", + "ConvGlobal", + "OneHot", + "MolDesc", +) class LogLowestMAE(L.Callback): @@ -30,9 +38,12 @@ def on_validation_end(self, trainer, pl_module): currentMAE = trainer.callback_metrics["Val Mean MAE"] if currentMAE < self.bestMAE: self.bestMAE = currentMAE - if self.config["wandb"]["enabled"]: - if wandb is not None: - wandb.log({"Best Val MAE": self.bestMAE}) + if self.config["wandb"]["enabled"] and wandb is not None and trainer.is_global_zero: + # wandb.init() is only ever called on the rank-zero process (see + # WandbLogger / _model_ops._setup_logger); under DDP with >1 + # device, every rank runs this hook, and a non-zero rank calling + # wandb.log() without a prior wandb.init() raises. + wandb.log({"Best Val MAE": self.bestMAE}) class BackboneFreeze(L.Callback): diff --git a/im2deep/_architectures/helpers.py b/im2deep/_architectures/helpers.py index eb69bd8..7361680 100644 --- a/im2deep/_architectures/helpers.py +++ b/im2deep/_architectures/helpers.py @@ -5,20 +5,57 @@ LOGGER = logging.getLogger(__name__) +def calculate_seq_concat_shape(config): + """ + Flattened width of ``torch.cat((atom_comp, diatom_comp), 1)``. + + This is the "sequence branch" - + everything derived from the peptide sequence itself, excluding globals, + one-hot and mol-desc. + + With ``use_attention_pooling``, each branch ends in + ``AttentionPool1d`` instead of ``nn.Flatten()``, so its output width is + just its channel count - independent of sequence length, unlike + ``Flatten()``'s ``channels * positions``. + + With ``use_flex_trunk``, ``ConvAtomComp`` and + ``ConvDiatomComp`` are retired entirely in favour of + :class:`~im2deep._architectures.blocks.FlexTrunk`, whose masked sum+max + pooling produces ``2 * flex_trunk_channels[-1]`` regardless of + ``use_attention_pooling`` - checked first and takes priority over it. + + """ + if config.get("use_flex_trunk", False): + channels = config.get("flex_trunk_channels", (128, 128)) + return 2 * channels[-1] + if config.get("use_transformer_trunk", False): + return config["transformer_d_model"] + if config.get("use_attention_pooling", False): + atom_comp_out_shape = config["AtomComp_out_channels_start"] // 4 + diatom_comp_out_shape = config["DiatomComp_out_channels_start"] // 2 + else: + atom_comp_out_shape = (60 // (2 * config["AtomComp_MaxPool_kernel_size"])) * ( + config["AtomComp_out_channels_start"] // 4 + ) + diatom_comp_out_shape = (30 // (config["DiatomComp_MaxPool_kernel_size"])) * ( + config["DiatomComp_out_channels_start"] // 2 + ) + return atom_comp_out_shape + diatom_comp_out_shape + + def calculate_concat_shape(config): - atom_comp_out_shape = (60 // (2 * config["AtomComp_MaxPool_kernel_size"])) * ( - config["AtomComp_out_channels_start"] // 4 - ) - LOGGER.debug(f"AtomComp out shape: {atom_comp_out_shape}") - diatom_comp_out_shape = (30 // (config["DiatomComp_MaxPool_kernel_size"])) * ( - config["DiatomComp_out_channels_start"] // 2 - ) - LOGGER.debug(f"DiatomComp out shape: {diatom_comp_out_shape}") + seq_concat_shape = calculate_seq_concat_shape(config) + LOGGER.debug(f"Sequence (atom+diatom) concat shape: {seq_concat_shape}") globals_out_shape = config["Global_units"] LOGGER.debug(f"Globals out shape: {globals_out_shape}") - onehot_comp_out_shape = (60 // (config["OneHot_MaxPool_kernel_size"])) * config[ - "OneHot_out_channels" - ] + if config.get("use_flex_trunk", False): + onehot_comp_out_shape = 0 + elif config.get("use_one_hot", True): + onehot_comp_out_shape = (60 // (config["OneHot_MaxPool_kernel_size"])) * config[ + "OneHot_out_channels" + ] + else: + onehot_comp_out_shape = 0 LOGGER.debug(f"OneHot out shape: {onehot_comp_out_shape}") if config["add_X_mol"]: @@ -27,16 +64,10 @@ def calculate_concat_shape(config): ) LOGGER.debug(f"MolDesc out shape: {mol_desc_comp_out_shape}") total_input_size = ( - atom_comp_out_shape - + diatom_comp_out_shape - + globals_out_shape - + onehot_comp_out_shape - + mol_desc_comp_out_shape + seq_concat_shape + globals_out_shape + onehot_comp_out_shape + mol_desc_comp_out_shape ) else: - total_input_size = ( - atom_comp_out_shape + diatom_comp_out_shape + globals_out_shape + onehot_comp_out_shape - ) + total_input_size = seq_concat_shape + globals_out_shape + onehot_comp_out_shape return total_input_size diff --git a/im2deep/_architectures/im2deep_multi.py b/im2deep/_architectures/im2deep_multi.py index 2f4079f..3791779 100644 --- a/im2deep/_architectures/im2deep_multi.py +++ b/im2deep/_architectures/im2deep_multi.py @@ -1,6 +1,7 @@ """Architecture definition for the multi-conformer IM2Deep model.""" import logging +from functools import partial from pathlib import Path import lightning as L @@ -493,11 +494,15 @@ def configure_optimizers(self): def configure_init(self): if (not self.config["init"]) or (self.config["init"] == "normal"): - return nn.init.normal_ + return partial(nn.init.normal_, mean=0.0, std=0.05) if self.config["init"] == "xavier": - return nn.init.xavier_normal_ + return partial(nn.init.xavier_normal_, gain=1.0) if self.config["init"] == "kaiming": - return nn.init.kaiming_normal_ + return partial( + nn.init.kaiming_normal_, + a=self.config["LRelu_negative_slope"], + nonlinearity="leaky_relu", + ) class IM2DeepMultiTransfer(L.LightningModule): diff --git a/im2deep/_architectures/im2deep_single.py b/im2deep/_architectures/im2deep_single.py index 2ab88ab..36f0a9c 100644 --- a/im2deep/_architectures/im2deep_single.py +++ b/im2deep/_architectures/im2deep_single.py @@ -1,174 +1,228 @@ """Architecture definition for the single-conformer IM2Deep model.""" import logging +from functools import partial import lightning as L import torch from torch import Tensor, nn -from torch.optim import Adam # type: ignore[import] +from torch.optim import Adam, AdamW # type: ignore[import] +from torch.optim.lr_scheduler import OneCycleLR -from im2deep._architectures.blocks import Conv1dActivation, DenseActivation -from im2deep._architectures.helpers import calculate_concat_shape +from im2deep._architectures.blocks import ( + AttentionPool1d, + Conv1dActivation, + DenseActivation, + FlexTrunk, +) +from im2deep._architectures.helpers import calculate_concat_shape, calculate_seq_concat_shape +from im2deep._architectures.transformer_trunk import TransformerTrunk from im2deep.constants import DEFAULT_GLOBAL_FEATURES LOGGER = logging.getLogger(__name__) +CHARGE_FEATURE_INDEX = 11 +LOG_B_CLAMP = 10.0 +NUM_CHARGE_HEADS = 6 + + class IM2Deep(L.LightningModule): def __init__(self, config, criterion): super().__init__() self.config = config self.criterion = criterion self.mae = nn.L1Loss() - # Record the config in the checkpoint so a model can be read back with - # the architecture and featurisation it was trained on. The bundled - # checkpoints predate this and carry no hyperparameters, which is why - # _model_ops still falls back to DEFAULT_CONFIG. self.save_hyperparameters("config") + n_global = self.config.get("Global_features", DEFAULT_GLOBAL_FEATURES) + self.register_buffer("global_feat_mean", torch.zeros(n_global)) + self.register_buffer("global_feat_std", torch.ones(n_global)) + self.register_buffer("target_mean", torch.zeros(())) + self.register_buffer("target_std", torch.ones(())) + initi = self.configure_init() - self.ConvAtomComp = nn.ModuleList() - self.ConvAtomComp.append( - Conv1dActivation( - 6, - self.config["AtomComp_out_channels_start"], - self.config["AtomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.use_attention_pooling = bool(self.config.get("use_attention_pooling", False)) + self.use_flex_trunk = bool(self.config.get("use_flex_trunk", False)) + self.use_transformer_trunk = bool(self.config.get("use_transformer_trunk", False)) + if self.use_flex_trunk and self.use_transformer_trunk: + raise ValueError( + "use_flex_trunk and use_transformer_trunk are mutually exclusive: both " + "replace ConvAtomComp/ConvDiatomComp with a different mechanism." ) - ) - self.ConvAtomComp.append( - Conv1dActivation( - self.config["AtomComp_out_channels_start"], - self.config["AtomComp_out_channels_start"], - self.config["AtomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + if self.use_flex_trunk: + self.flex_trunk = FlexTrunk( + embed_dim=self.config.get("flex_trunk_embed_dim", 16), + channels=self.config.get("flex_trunk_channels", (192, 192)), + kernel_size=self.config.get("flex_trunk_kernel_size", 5), + stem_channels=self.config.get("flex_trunk_stem_channels", 64), + stem_layers=self.config.get("flex_trunk_stem_layers", 2), ) - ) - self.ConvAtomComp.append( - nn.MaxPool1d( - self.config["AtomComp_MaxPool_kernel_size"], - self.config["AtomComp_MaxPool_kernel_size"], + elif self.use_transformer_trunk: + if not self.config.get("add_ccs_features"): + raise ValueError( + "use_transformer_trunk requires add_ccs_features=True: the charge " + "register token needs CHARGE_FEATURE_INDEX from matrix_global." + ) + self.transformer_trunk = TransformerTrunk( + d_model=self.config.get("transformer_d_model", 128), + num_layers=self.config.get("transformer_num_layers", 4), + num_heads=self.config.get("transformer_num_heads", 4), + ffn_dim=self.config.get("transformer_ffn_dim", 512), + max_relative_distance=self.config.get("transformer_max_relative_distance", 32), + max_terminus_distance=self.config.get("transformer_max_terminus_distance", 16), + num_readout_queries=self.config.get("transformer_readout_queries", 1), + dropout=self.config.get("transformer_dropout", 0.0), ) - ) - self.ConvAtomComp.append( - Conv1dActivation( - self.config["AtomComp_out_channels_start"], - self.config["AtomComp_out_channels_start"] // 2, - self.config["AtomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + else: + self.ConvAtomComp = nn.ModuleList() + self.ConvAtomComp.append( + Conv1dActivation( + 6, + self.config["AtomComp_out_channels_start"], + self.config["AtomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.ConvAtomComp.append( - Conv1dActivation( - self.config["AtomComp_out_channels_start"] // 2, - self.config["AtomComp_out_channels_start"] // 2, - self.config["AtomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.ConvAtomComp.append( + Conv1dActivation( + self.config["AtomComp_out_channels_start"], + self.config["AtomComp_out_channels_start"], + self.config["AtomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.ConvAtomComp.append( - nn.MaxPool1d( - self.config["AtomComp_MaxPool_kernel_size"], - self.config["AtomComp_MaxPool_kernel_size"], + self.ConvAtomComp.append( + nn.MaxPool1d( + self.config["AtomComp_MaxPool_kernel_size"], + self.config["AtomComp_MaxPool_kernel_size"], + ) ) - ) - self.ConvAtomComp.append( - Conv1dActivation( - self.config["AtomComp_out_channels_start"] // 2, - self.config["AtomComp_out_channels_start"] // 4, - self.config["AtomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.ConvAtomComp.append( + Conv1dActivation( + self.config["AtomComp_out_channels_start"], + self.config["AtomComp_out_channels_start"] // 2, + self.config["AtomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.ConvAtomComp.append( - Conv1dActivation( - self.config["AtomComp_out_channels_start"] // 4, - self.config["AtomComp_out_channels_start"] // 4, - self.config["AtomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.ConvAtomComp.append( + Conv1dActivation( + self.config["AtomComp_out_channels_start"] // 2, + self.config["AtomComp_out_channels_start"] // 2, + self.config["AtomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.ConvAtomComp.append(nn.Flatten()) - - self.ConvDiatomComp = nn.ModuleList() - self.ConvDiatomComp.append( - Conv1dActivation( - 6, - self.config["DiatomComp_out_channels_start"], - self.config["DiatomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.ConvAtomComp.append( + nn.MaxPool1d( + self.config["AtomComp_MaxPool_kernel_size"], + self.config["AtomComp_MaxPool_kernel_size"], + ) ) - ) - self.ConvDiatomComp.append( - Conv1dActivation( - self.config["DiatomComp_out_channels_start"], - self.config["DiatomComp_out_channels_start"], - self.config["DiatomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.ConvAtomComp.append( + Conv1dActivation( + self.config["AtomComp_out_channels_start"] // 2, + self.config["AtomComp_out_channels_start"] // 4, + self.config["AtomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.ConvDiatomComp.append( - nn.MaxPool1d( - self.config["DiatomComp_MaxPool_kernel_size"], - self.config["DiatomComp_MaxPool_kernel_size"], + self.ConvAtomComp.append( + Conv1dActivation( + self.config["AtomComp_out_channels_start"] // 4, + self.config["AtomComp_out_channels_start"] // 4, + self.config["AtomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.ConvDiatomComp.append( - Conv1dActivation( - self.config["DiatomComp_out_channels_start"], - self.config["DiatomComp_out_channels_start"] // 2, - self.config["DiatomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + if self.use_attention_pooling: + self.ConvAtomComp.append( + AttentionPool1d(self.config["AtomComp_out_channels_start"] // 4, initi) + ) + else: + self.ConvAtomComp.append(nn.Flatten()) + + self.ConvDiatomComp = nn.ModuleList() + self.ConvDiatomComp.append( + Conv1dActivation( + 6, + self.config["DiatomComp_out_channels_start"], + self.config["DiatomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.ConvDiatomComp.append( - Conv1dActivation( - self.config["DiatomComp_out_channels_start"] // 2, - self.config["DiatomComp_out_channels_start"] // 2, - self.config["DiatomComp_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.ConvDiatomComp.append( + Conv1dActivation( + self.config["DiatomComp_out_channels_start"], + self.config["DiatomComp_out_channels_start"], + self.config["DiatomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.ConvDiatomComp.append(nn.Flatten()) + self.ConvDiatomComp.append( + nn.MaxPool1d( + self.config["DiatomComp_MaxPool_kernel_size"], + self.config["DiatomComp_MaxPool_kernel_size"], + ) + ) + self.ConvDiatomComp.append( + Conv1dActivation( + self.config["DiatomComp_out_channels_start"], + self.config["DiatomComp_out_channels_start"] // 2, + self.config["DiatomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) + ) + self.ConvDiatomComp.append( + Conv1dActivation( + self.config["DiatomComp_out_channels_start"] // 2, + self.config["DiatomComp_out_channels_start"] // 2, + self.config["DiatomComp_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) + ) + if self.use_attention_pooling: + self.ConvDiatomComp.append( + AttentionPool1d(self.config["DiatomComp_out_channels_start"] // 2, initi) + ) + else: + self.ConvDiatomComp.append(nn.Flatten()) self.ConvGlobal = nn.ModuleList() self.ConvGlobal.append( DenseActivation( - # Number of DeepLC global features, which depends on the - # featurisation flags. The default matches - # `add_ccs_features=True, add_terminal_composition=False`, - # i.e. what every bundled checkpoint was trained with. self.config.get("Global_features", DEFAULT_GLOBAL_FEATURES), self.config["Global_units"], initializer=initi, @@ -195,36 +249,38 @@ def __init__(self, config, criterion): ) ) - self.OneHot = nn.ModuleList() - self.OneHot.append( - Conv1dActivation( - 20, - self.config["OneHot_out_channels"], - self.config["One_hot_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.use_one_hot = bool(self.config.get("use_one_hot", True)) and not self.use_flex_trunk + if self.use_one_hot: + self.OneHot = nn.ModuleList() + self.OneHot.append( + Conv1dActivation( + 20, + self.config["OneHot_out_channels"], + self.config["One_hot_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.OneHot.append( - Conv1dActivation( - self.config["OneHot_out_channels"], - self.config["OneHot_out_channels"], - self.config["One_hot_kernel_size"], - padding="same", - initializer=initi, - negative_slope=self.config["LRelu_negative_slope"], - saturation=self.config["LRelu_saturation"], + self.OneHot.append( + Conv1dActivation( + self.config["OneHot_out_channels"], + self.config["OneHot_out_channels"], + self.config["One_hot_kernel_size"], + padding="same", + initializer=initi, + negative_slope=self.config["LRelu_negative_slope"], + saturation=self.config["LRelu_saturation"], + ) ) - ) - self.OneHot.append( - nn.MaxPool1d( - self.config["OneHot_MaxPool_kernel_size"], - self.config["OneHot_MaxPool_kernel_size"], + self.OneHot.append( + nn.MaxPool1d( + self.config["OneHot_MaxPool_kernel_size"], + self.config["OneHot_MaxPool_kernel_size"], + ) ) - ) - self.OneHot.append(nn.Flatten()) + self.OneHot.append(nn.Flatten()) if config["add_X_mol"]: self.MolDesc = nn.ModuleList() @@ -308,6 +364,41 @@ def __init__(self, config, criterion): ) self.MolDesc.append(nn.Flatten()) + self.use_charge_embedding = bool(self.config.get("use_charge_embedding", False)) + if self.use_charge_embedding: + if not self.config.get("add_ccs_features"): + raise ValueError( + "use_charge_embedding requires add_ccs_features=True: " + "charge is only present in matrix_global, at a fixed " + "index, when that flag is set." + ) + charge_embedding_dim = self.config.get("charge_embedding_dim", 16) + seq_concat_shape = calculate_seq_concat_shape(self.config) + self.charge_embedding = nn.Embedding(7, charge_embedding_dim) + self.charge_film = nn.Linear(charge_embedding_dim, 2 * seq_concat_shape) + nn.init.zeros_(self.charge_film.weight) + nn.init.zeros_(self.charge_film.bias) + + self.use_per_charge_heads = bool(self.config.get("use_per_charge_heads", False)) + if self.use_per_charge_heads and not self.config.get("add_ccs_features"): + raise ValueError( + "use_per_charge_heads requires add_ccs_features=True: selecting the right " + "charge head needs CHARGE_FEATURE_INDEX from matrix_global." + ) + + if self.config.get("mixup_alpha", 0.0) > 0 and ( + self.config.get("sibling_loss_weight", 0.0) > 0 + or self.config.get("ranking_loss_weight", 0.0) > 0 + or self.use_per_charge_heads + ): + raise ValueError( + "mixup_alpha>0 is not compatible with sibling_loss_weight/ranking_loss_weight/" + "use_per_charge_heads: mixup blends batch rows together, corrupting the " + "per-row identity those depend on." + ) + + self.use_heteroscedastic_head = bool(self.config.get("use_heteroscedastic_head", False)) + self.total_input_size = calculate_concat_shape(self.config) LOGGER.debug(f"Total input size: {self.total_input_size}") @@ -358,13 +449,151 @@ def __init__(self, config, criterion): ) ) - self.Concat.append(nn.Linear(self.config["Concat_units"], 1)) + # Per-row output width: mu only, or (mu, log_b) when + # use_heteroscedastic_head, repeated once per charge head when + # use_per_charge_heads - see forward() for the destandardisation and + # (for use_per_charge_heads) the per-row head selection. + per_row_width = 2 if self.use_heteroscedastic_head else 1 + if self.use_per_charge_heads: + per_row_width *= NUM_CHARGE_HEADS + self.Concat.append(nn.Linear(self.config["Concat_units"], per_row_width)) + + def set_global_feature_scaler(self, mean: Tensor, std: Tensor) -> None: + """Set the global-feature standardisation buffers from a fitted sample.""" + self.global_feat_mean.copy_(mean) # ty: ignore[call-non-callable] + self.global_feat_std.copy_(std) # ty: ignore[call-non-callable] + + def set_target_scaler(self, mean: float, std: float) -> None: + """Set the target standardisation buffers from the training target distribution.""" + self.target_mean.fill_(mean) # ty: ignore[call-non-callable] + self.target_std.fill_(std) # ty: ignore[call-non-callable] + + def _base_loss(self, y_hat, y): + """ + The training objective, in whichever space ``config["loss"]`` selects + (step 8.3). Always given raw-CCS ``y_hat``/``y`` - only the + functional form changes, not the units - so this composes cleanly + with target standardisation, which already returns raw CCS from + ``forward()`` regardless of whether it's enabled. The reported/ + monitored "Validation MAE" always stays ``self.mae(y_hat, y)`` + (plain L1 in raw Ų), so runs on different loss choices stay + comparable; only this training-objective value differs. + """ + loss_name = self.config.get("loss", "l1") + if loss_name == "l1": + return self.criterion(y_hat, y) + if loss_name == "mape": + return 100.0 * (torch.abs(y_hat - y) / y).mean() + if loss_name == "huber": + return nn.functional.huber_loss(y_hat, y, delta=self.config.get("huber_delta", 5.0)) + raise ValueError(f"Unknown loss {loss_name!r}") def regularized_loss(self, y_hat, y): - standard_loss = self.criterion(y_hat, y) + standard_loss = self._base_loss(y_hat, y) l1_norm = sum(torch.norm(p, 1) for p in self.parameters()) return standard_loss + self.config["L1_alpha"] * l1_norm + def _pair_mask(self, same: Tensor) -> tuple[Tensor, Tensor]: + """Row indices of each unique unordered pair where ``same`` holds.""" + return torch.triu(same, diagonal=1).nonzero(as_tuple=True) # ty: ignore[invalid-return-type] + + def _sibling_delta_loss(self, y_hat: Tensor, y: Tensor, group_id: Tensor) -> Tensor: + """ + Sibling delta loss: ``L1(pred_a - pred_b, obs_a - obs_b)`` + over batch pairs sharing a stripped sequence (same ``group_id``), + i.e. pairs differing only in modification state or charge. + + Always plain L1 on the delta regardless of ``config["loss"]`` - a + deliberate simplification, not an oversight: this term exists to + supervise the modification/charge *effect* directly (IM2Deep's own + selling point), which the absolute-CCS loss can hide behind two + errors that cancel. ``group_id`` only ever contains real matches + because ``SiblingPairBatchSampler`` deliberately placed them; a + batch with none (e.g. every filler slot missed by chance) safely + contributes zero. + """ + weight = self.config.get("sibling_loss_weight", 0.0) + if weight <= 0: + return y_hat.new_zeros(()) + idx_a, idx_b = self._pair_mask(group_id.unsqueeze(0) == group_id.unsqueeze(1)) + if idx_a.numel() == 0: + return y_hat.new_zeros(()) + pred_delta = y_hat[idx_a] - y_hat[idx_b] + obs_delta = y[idx_a] - y[idx_b] + return weight * nn.functional.l1_loss(pred_delta, obs_delta) + + def _ranking_loss(self, y_hat: Tensor, y: Tensor, mz: Tensor) -> Tensor: + """ + Pairwise ranking loss: a margin ranking loss on the sign + of predicted vs observed CCS difference, over batch pairs whose + theoretical m/z are within ``ranking_mz_window`` Th of each other - + the discrimination that actually matters for precursors a + quadrupole cannot isolate separately. Pairs with a tied observed CCS + (``target_sign == 0``) are dropped: there is no ordering to + supervise, and 0 is not a valid ``margin_ranking_loss`` target. + """ + weight = self.config.get("ranking_loss_weight", 0.0) + if weight <= 0: + return y_hat.new_zeros(()) + window = self.config.get("ranking_mz_window", 3.0) + idx_a, idx_b = self._pair_mask((mz.unsqueeze(0) - mz.unsqueeze(1)).abs() <= window) + if idx_a.numel() == 0: + return y_hat.new_zeros(()) + target_sign = torch.sign(y[idx_a] - y[idx_b]) + keep = target_sign != 0 + if keep.sum() == 0: + return y_hat.new_zeros(()) + return weight * nn.functional.margin_ranking_loss( + y_hat[idx_a][keep], + y_hat[idx_b][keep], + target_sign[keep], + margin=self.config.get("ranking_margin", 0.0), + ) + + def _laplace_nll(self, mu: Tensor, y: Tensor, b: Tensor) -> Tensor: + """ + Laplace negative log-likelihood (step 8.5's heteroscedastic head): + ``log(2b) + |y-mu|/b``. For a fixed absolute error ``e = |y-mu|``, + this is minimised exactly at ``b = e`` (the Laplace MLE scale), so a + row the network cannot fit is automatically assigned a larger ``b`` + rather than forcing ``mu`` to keep chasing it - loss attenuation by + construction, not a hand-picked per-row reweighting rule. Chosen + over the Gaussian NLL specifically for consistency with this + codebase's existing default (`loss: "l1"`, whose point estimate is + also the median, matching Laplace's) and because CCS label noise + here is residual disagreement between calibrated mass-spec runs - + occasional large residuals from a still-disagreeing run, not + symmetric Gaussian jitter - which Laplace's heavier tail tolerates + better than Gaussian's quadratic penalty would. + """ + return (torch.log(2.0 * b) + torch.abs(y - mu) / b).mean() + + def _in_heteroscedastic_warmup(self) -> bool: + """ + True for the first ``heteroscedastic_warmup_epochs`` epochs (default + 0, i.e. never) when ``use_heteroscedastic_head`` is on. + """ + warmup_epochs = self.config.get("heteroscedastic_warmup_epochs", 0) + return self.use_heteroscedastic_head and self.current_epoch < warmup_epochs + + def _point_estimate_loss(self, y_hat: Tensor, y: Tensor, b: Tensor) -> Tensor: + """ + Plain L1 on ``y_hat`` during the heteroscedastic warm-up (``b`` is + computed by ``forward()`` regardless, but does not enter the loss, + so its output row gets no gradient from this term - it starts + learning from an already-well-trained shared trunk once warm-up + ends, rather than trying to calibrate from a cold, random one at + the same time ``mu`` is still finding its footing). The Laplace NLL + otherwise. Used identically by training/validation/test so + "Validation loss" stays meaningful for whichever phase training is + actually in - unlike "Validation MAE", which always reports plain + L1 on ``mu`` regardless (the established invariant across every + ``loss``/head choice in this codebase). + """ + if self._in_heteroscedastic_warmup(): + return nn.functional.l1_loss(y_hat, y) + return self._laplace_nll(y_hat, y, b) + def forward( self, atom_comp, @@ -372,26 +601,69 @@ def forward( global_feats, one_hot, mol_desc: Tensor | None = None, + return_attention: bool = False, ): + if return_attention and not self.use_transformer_trunk: + raise ValueError( + "return_attention=True requires use_transformer_trunk=True: " + "no other branch exposes attention weights." + ) + + # TransformerTrunk needs the RAW, un-permuted matrix tensor (it does + # its own layout internally, unlike every Conv1d branch here) - + # captured before the permute below. + atom_comp_raw = atom_comp + atom_comp = atom_comp.permute(0, 2, 1) diatom_comp = diatom_comp.permute(0, 2, 1) one_hot = one_hot.permute(0, 2, 1) - for layer in self.ConvAtomComp: - atom_comp = layer(atom_comp) + if self.use_charge_embedding or self.use_transformer_trunk or self.use_per_charge_heads: + charge_idx = global_feats[:, CHARGE_FEATURE_INDEX].round().long().clamp(0, 6) + + global_feats = (global_feats - self.global_feat_mean) / self.global_feat_std + + if self.use_flex_trunk: + # atom_comp/one_hot are already permuted to (batch, channels, + # length) above - FlexTrunk expects exactly that layout, same + # as every other branch here. diatom_comp is unused: matrix_sum + # is redundant once the trunk sees matrix directly. + seq_feats = self.flex_trunk(atom_comp, one_hot) + elif self.use_transformer_trunk: + # diatom_comp is unused here too, for the same reason as + # flex_trunk - the transformer sees matrix directly. + if return_attention: + seq_feats, attention = self.transformer_trunk( + atom_comp_raw, charge_idx, return_attention=True + ) + else: + seq_feats = self.transformer_trunk(atom_comp_raw, charge_idx) + else: + for layer in self.ConvAtomComp: + atom_comp = layer(atom_comp) + for layer in self.ConvDiatomComp: + diatom_comp = layer(diatom_comp) + seq_feats = torch.cat((atom_comp, diatom_comp), 1) - for layer in self.ConvDiatomComp: - diatom_comp = layer(diatom_comp) for layer in self.ConvGlobal: global_feats = layer(global_feats) - for layer in self.OneHot: - one_hot = layer(one_hot) + if self.use_one_hot: + for layer in self.OneHot: + one_hot = layer(one_hot) if self.config["add_X_mol"]: for layer in self.MolDesc: mol_desc = layer(mol_desc) - concatenated = torch.cat((atom_comp, diatom_comp, one_hot, global_feats), 1) + if self.use_charge_embedding: + charge_embed = self.charge_embedding(charge_idx) + gamma, beta = self.charge_film(charge_embed).chunk(2, dim=-1) + seq_feats = seq_feats * (1.0 + gamma) + beta + + if self.use_one_hot: + concatenated = torch.cat((seq_feats, one_hot, global_feats), 1) + else: + concatenated = torch.cat((seq_feats, global_feats), 1) if self.config["add_X_mol"]: if mol_desc is None: @@ -401,18 +673,120 @@ def forward( for layer in self.Concat: concatenated = layer(concatenated) - output = concatenated + if self.use_per_charge_heads: + sub_width = concatenated.shape[1] // NUM_CHARGE_HEADS + concatenated = concatenated.view(-1, NUM_CHARGE_HEADS, sub_width) + head_idx = (charge_idx - 1).clamp(0, NUM_CHARGE_HEADS - 1) + concatenated = concatenated[ + torch.arange(concatenated.size(0), device=concatenated.device), head_idx + ] + + if self.use_heteroscedastic_head: + mu_std, log_b_std = concatenated[:, 0], concatenated[:, 1] + mu = mu_std * self.target_std + self.target_mean # ty: ignore[unsupported-operator] + b = torch.exp(log_b_std.clamp(-LOG_B_CLAMP, LOG_B_CLAMP)) * self.target_std # ty: ignore[unsupported-operator] + b = b.clamp(min=1e-3) + if return_attention: + return mu, b, attention + return mu, b + + # De-standardise back to raw CCS. + output = concatenated * self.target_std + self.target_mean # ty: ignore[unsupported-operator] + if return_attention: + return output, attention return output + def _apply_mixup(self, atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc=None): + """ + Composition mixup: mix each training batch with a random + permutation of itself via ``lambda ~ Beta(mixup_alpha, mixup_alpha)``, + applied identically to every input tensor and the target. + + The atomic-composition encoding is additive, so a convex combination + of two peptides' `matrix`/`matrix_global` is a chemically + interpretable point rather than nonsense, and CCS is smooth in + composition, so interpolating the targets the same way is + approximately valid. Training only - never called for validation/ + test - and "Train MAE" is measured against the mixed target while + this is on, so it stops being directly comparable to a non-mixup + run's; "Validation MAE" is unaffected and remains the number to + compare across configs. + + NOT YET COMPATIBLE with ``use_charge_embedding``: ``global_feats``' + charge value (see ``CHARGE_FEATURE_INDEX``) would get linearly + blended into a non-integer, physically meaningless "charge", which + ``forward()``'s ``charge_idx = ....round().long()`` would then round + arbitrarily - don't combine them until that's addressed. + + A no-op when ``config["mixup_alpha"]`` is 0 (the default). + """ + mixup_alpha = self.config.get("mixup_alpha", 0.0) + if mixup_alpha <= 0: + return atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc + + lam = torch.distributions.Beta(mixup_alpha, mixup_alpha).sample().item() + perm = torch.randperm(atom_comp.size(0), device=atom_comp.device) + atom_comp = lam * atom_comp + (1 - lam) * atom_comp[perm] + diatom_comp = lam * diatom_comp + (1 - lam) * diatom_comp[perm] + global_feats = lam * global_feats + (1 - lam) * global_feats[perm] + one_hot = lam * one_hot + (1 - lam) * one_hot[perm] + y = lam * y + (1 - lam) * y[perm] + if mol_desc is not None: + mol_desc = lam * mol_desc + (1 - lam) * mol_desc[perm] + return atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc + + def _forward_point_estimate( + self, atom_comp, diatom_comp, global_feats, one_hot, mol_desc=None + ) -> tuple[Tensor, Tensor | None]: + """ + Runs ``forward()`` and normalises its return to ``(y_hat, b)`` + regardless of ``use_heteroscedastic_head`` - ``b`` is ``None`` when + disabled. Centralises the ``add_X_mol`` branching and the + heteroscedastic tuple-vs-plain-tensor difference so + training_step/validation_step/test_step/predict_step don't each + repeat it; ``y_hat`` plays the same "point estimate" role in both + cases for MAE logging and the sibling/ranking losses. + """ + if self.config["add_X_mol"]: + output = self(atom_comp, diatom_comp, global_feats, one_hot, mol_desc) + else: + output = self(atom_comp, diatom_comp, global_feats, one_hot) + if self.use_heteroscedastic_head: + mu, b = output + return mu, b + return output.squeeze(1), None + def training_step(self, batch, batch_idx): + needs_pair_ids = ( + self.config.get("sibling_loss_weight", 0.0) > 0 + or self.config.get("ranking_loss_weight", 0.0) > 0 + ) + if needs_pair_ids: + *batch, group_id, mz = batch + if self.config["add_X_mol"]: atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc = batch - y_hat = self(atom_comp, diatom_comp, global_feats, one_hot, mol_desc).squeeze(1) else: atom_comp, diatom_comp, global_feats, one_hot, y = batch - y_hat = self(atom_comp, diatom_comp, global_feats, one_hot).squeeze(1) + mol_desc = None + + atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc = self._apply_mixup( + atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc + ) - loss = self.regularized_loss(y_hat, y) + y_hat, b = self._forward_point_estimate( + atom_comp, diatom_comp, global_feats, one_hot, mol_desc + ) + + if self.use_heteroscedastic_head: + l1_norm = sum(torch.norm(p, 1) for p in self.parameters()) + loss = self._point_estimate_loss(y_hat, y, b) + self.config["L1_alpha"] * l1_norm # ty: ignore[invalid-argument-type] + else: + loss = self.regularized_loss(y_hat, y) + if self.config.get("sibling_loss_weight", 0.0) > 0: + loss = loss + self._sibling_delta_loss(y_hat, y, group_id) + if self.config.get("ranking_loss_weight", 0.0) > 0: + loss = loss + self._ranking_loss(y_hat, y, mz) self.log("Train loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True) self.log( @@ -428,11 +802,17 @@ def training_step(self, batch, batch_idx): def validation_step(self, batch, batch_idx): if self.config["add_X_mol"]: atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc = batch - y_hat = self(atom_comp, diatom_comp, global_feats, one_hot, mol_desc).squeeze(1) else: atom_comp, diatom_comp, global_feats, one_hot, y = batch - y_hat = self(atom_comp, diatom_comp, global_feats, one_hot).squeeze(1) - loss = self.criterion(y_hat, y) + mol_desc = None + y_hat, b = self._forward_point_estimate( + atom_comp, diatom_comp, global_feats, one_hot, mol_desc + ) + loss = ( + self._point_estimate_loss(y_hat, y, b) # ty: ignore[invalid-argument-type] + if self.use_heteroscedastic_head + else self._base_loss(y_hat, y) + ) self.log("Validation loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True) self.log( @@ -448,11 +828,17 @@ def validation_step(self, batch, batch_idx): def test_step(self, batch, batch_idx): if self.config["add_X_mol"]: atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc = batch - y_hat = self(atom_comp, diatom_comp, global_feats, one_hot, mol_desc).squeeze(1) else: atom_comp, diatom_comp, global_feats, one_hot, y = batch - y_hat = self(atom_comp, diatom_comp, global_feats, one_hot).squeeze(1) - loss = self.criterion(y_hat, y) + mol_desc = None + y_hat, b = self._forward_point_estimate( + atom_comp, diatom_comp, global_feats, one_hot, mol_desc + ) + loss = ( + self._point_estimate_loss(y_hat, y, b) # ty: ignore[invalid-argument-type] + if self.use_heteroscedastic_head + else self._base_loss(y_hat, y) + ) self.log("Test loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True) self.log( @@ -468,25 +854,49 @@ def test_step(self, batch, batch_idx): def predict_step(self, batch, batch_idx, dataloader_idx=None): if self.config["add_X_mol"]: atom_comp, diatom_comp, global_feats, one_hot, y, mol_desc = batch - y_hat = self(atom_comp, diatom_comp, global_feats, one_hot, mol_desc).squeeze(1) else: atom_comp, diatom_comp, global_feats, one_hot, y = batch - y_hat = self(atom_comp, diatom_comp, global_feats, one_hot).squeeze(1) + mol_desc = None + y_hat, _ = self._forward_point_estimate( + atom_comp, diatom_comp, global_feats, one_hot, mol_desc + ) return y_hat def configure_optimizers(self): - optimizer = Adam(self.parameters(), lr=self.config["learning_rate"]) - return optimizer + optimizer = AdamW( + self.parameters(), + lr=self.config["learning_rate"], + weight_decay=self.config.get("weight_decay", 0.0), + ) + scheduler_name = self.config.get("lr_scheduler") + if not scheduler_name: + return optimizer + if scheduler_name == "one_cycle": + scheduler = OneCycleLR( + optimizer, + max_lr=self.config["learning_rate"], + total_steps=self.trainer.estimated_stepping_batches, # ty: ignore[invalid-argument-type] + ) + return { + "optimizer": optimizer, + "lr_scheduler": {"scheduler": scheduler, "interval": "step"}, + } + raise ValueError(f"Unknown lr_scheduler {scheduler_name!r}") def configure_init(self): if (not self.config["init"]) or (self.config["init"] == "normal"): - return nn.init.normal_ + return partial(nn.init.normal_, mean=0.0, std=0.05) if self.config["init"] == "xavier": - return nn.init.xavier_normal_ + return partial(nn.init.xavier_normal_, gain=1.0) if self.config["init"] == "kaiming": - return nn.init.kaiming_normal_ + return partial( + nn.init.kaiming_normal_, + a=self.config["LRelu_negative_slope"], + nonlinearity="leaky_relu", + ) +# TODO: update to new TransformerTrun architectures class IM2DeepTransfer(L.LightningModule): def __init__(self, config, criterion): super().__init__() @@ -500,9 +910,41 @@ def __init__(self, config, criterion): # Load the IM2Deep model LOGGER.debug("Loading backbone IM2Deep model") self.backbone = IM2Deep.load_from_checkpoint( - config["backbone_SD_path"], config=config, criterion=criterion + config["backbone_SD_path"], config=config, criterion=criterion, strict=False ) + if self.backbone.config.get("use_heteroscedastic_head", False): + raise ValueError( + "Cannot finetune from a use_heteroscedastic_head=True backbone: its Concat " + "head ends in a 2-wide (mu, log_b) layer, which IM2DeepTransfer.forward() has " + "no path for. Finetuning support for this wasn't added here." + ) + + if self.backbone.config.get("use_per_charge_heads", False): + raise ValueError( + "Cannot finetune from a use_per_charge_heads=True backbone: its Concat head " + "ends in a per-charge-head layer, which IM2DeepTransfer.forward() has no " + "selection path for. Finetuning support for this wasn't added here." + ) + + # A use_flex_trunk/use_transformer_trunk backbone has no + # ConvAtomComp/ConvDiatomComp, and its forward() also carries the + # charge register token and both standardisation scalers. Rather than + # re-implementing that here, forward() delegates to the backbone; the + # branches are only exposed as attributes so BackboneFreeze finds them. + self.delegate_forward = bool( + self.backbone.use_transformer_trunk or self.backbone.use_flex_trunk + ) + if self.delegate_forward: + for name in ("transformer_trunk", "flex_trunk", "ConvGlobal", "OneHot", "MolDesc"): + if hasattr(self.backbone, name): + setattr(self, name, getattr(self.backbone, name)) + self.concat = self.backbone.Concat + return + + # KNOWN LIMITATION (CNN backbones only): forward() below skips the + # backbone's global-feature and target standardisation, so a CNN + # backbone trained with standardize_* gives wrongly scaled output. self.ConvAtomComp = self.backbone.ConvAtomComp self.ConvDiatomComp = self.backbone.ConvDiatomComp self.ConvGlobal = self.backbone.ConvGlobal @@ -521,6 +963,9 @@ def forward( one_hot, mol_desc: Tensor | None = None, ): + if self.delegate_forward: + return self.backbone(atom_comp, diatom_comp, global_feats, one_hot, mol_desc) + atom_comp = atom_comp.permute(0, 2, 1) diatom_comp = diatom_comp.permute(0, 2, 1) one_hot = one_hot.permute(0, 2, 1) diff --git a/im2deep/_architectures/transformer_trunk.py b/im2deep/_architectures/transformer_trunk.py new file mode 100644 index 0000000..472dba7 --- /dev/null +++ b/im2deep/_architectures/transformer_trunk.py @@ -0,0 +1,426 @@ +""" +Transformer encoder trunk over atomic-composition tokens (step 8.4 item 3). + +The main architecture arm of step 8.4: token ``i`` is residue ``i``'s 6-vector +atomic composition, taken straight from ``matrix`` and projected to +``d_model`` by a single linear layer, so the atomic encoding is preserved +exactly and no featuriser changes are needed. Three additions distinguish +this from a generic transformer encoder, each addressing a specific, +previously-diagnosed gap: + +* **Charge as a prepended register token** (:class:`TransformerTrunk`), not a + late-concatenated scalar (the four-branch design's own approach) or a + FiLM scale-and-shift (``use_charge_embedding``, step 8.4 item 1). Every + residue attends to charge at every layer, and charge attends back - + giving charge a data flow the FiLM version cannot: a kernel can only + rescale a feature it has already computed, while self-attention lets the + network learn a charge-dependent contribution from the first layer. +* **Pair-biased attention** (:class:`RelativePairBias`): a learned bias on + the attention logits as a function of relative sequence separation and + proximity to either terminus, giving the model an explicit, learnable + notion of sequence distance - the AlphaFold/T5 pattern, and a much + stronger inductive bias for a task fundamentally about intramolecular + contact than a sinusoidal positional encoding would be. Deliberately + **not** T5's logarithmic bucketing: peptides here are at most + ``padding_length`` (60) residues, so a plain clamped linear table over + the true distance range is both simpler and strictly more expressive at + this scale (log-bucketing exists to keep a table small for very long + sequences, which is not the situation here). +* **Learned-query read-out** (:class:`PerceiverReadout`): ``k`` learned + query vectors cross-attend to the residue tokens, replacing pooling. + ``k=1`` (the only mode implemented here) is attention pooling with a + masked softmax; ``k=2``/``3`` would be a route into the multi-conformer + problem (step 9) but is out of scope for this arm as built. +""" + +import math + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + + +class RelativePairBias(nn.Module): + """ + Per-head additive attention bias from relative sequence separation and + terminus proximity, over residue positions only (the register token + gets no positional bias - it has no sequence position of its own). + + Two small, fully-learnable tables rather than DeepLC-style hand + engineering: ``relative`` indexed by the *signed*, clamped distance + ``i - j`` (asymmetric by design, so the model can learn direction- + dependent effects, e.g. "N-terminal side" vs "C-terminal side" of a + contact); ``terminus`` indexed by how close either position of the pair + is to either terminus of the *real* (unpadded) peptide, which is + symmetric by construction (a ``min`` over both positions and both + termini). + """ + + def __init__(self, num_heads: int, max_relative_distance: int, max_terminus_distance: int): + super().__init__() + self.max_relative_distance = max_relative_distance + self.max_terminus_distance = max_terminus_distance + self.relative = nn.Embedding(2 * max_relative_distance + 1, num_heads) + self.terminus = nn.Embedding(max_terminus_distance + 1, num_heads) + # (cache_key, tensor) for the batch-invariant half of forward() - see + # _relative_bias. Not a buffer: its shape varies with seq_len, and it + # must never be part of the state_dict. + self._relative_bias_cache: tuple[tuple, Tensor] | None = None + + def _relative_bias(self, seq_len: int, device: torch.device) -> Tensor: + """ + ``(seq_len, seq_len, num_heads)`` relative-position bias. + + Depends only on ``seq_len`` and the ``relative`` embedding's + weights, not on the batch's actual data - safe to compute once and + reuse across every forward call in an ``eval()``/no-grad prediction + run. Caching is skipped whenever a gradient could flow through this + call (``self.training`` or autograd enabled), since then a fresh + computation graph is required on every call; it is also skipped, + i.e. recomputed, if ``seq_len``/``device`` change or the embedding + weights have since been updated (``_version`` guards the latter). + """ + cacheable = not self.training and not torch.is_grad_enabled() + weight_version = self.relative.weight._version + cache_key = (seq_len, device, weight_version) + if ( + cacheable + and self._relative_bias_cache is not None + and self._relative_bias_cache[0] == cache_key + ): + return self._relative_bias_cache[1] + + positions = torch.arange(seq_len, device=device) + relative_distance = positions.unsqueeze(0) - positions.unsqueeze(1) # (L, L), [i,j] = i-j + relative_idx = ( + relative_distance.clamp(-self.max_relative_distance, self.max_relative_distance) + + self.max_relative_distance + ) + relative_bias = self.relative(relative_idx) # (L, L, num_heads) + + if cacheable: + self._relative_bias_cache = (cache_key, relative_bias) + return relative_bias + + def forward(self, real_lengths: Tensor, seq_len: int) -> Tensor: + """ + ``real_lengths``: ``(batch,)`` true (unpadded) peptide lengths. + ``seq_len``: the array length (``padding_length``), not the real + length - positions past a peptide's real length still get a + relative/terminus bias here (they are masked out separately via + ``key_padding_mask``; this module only computes the bias values). + + Returns ``(batch, num_heads, seq_len, seq_len)``. + """ + device = real_lengths.device + relative_bias = self._relative_bias(seq_len, device) + + dist_to_n_term = torch.arange(seq_len, device=device) # (L,) + # (batch, L): distance to the C-terminus depends on each peptide's + # own real length, not the shared padded array length. + dist_to_c_term = real_lengths.unsqueeze(1) - 1 - dist_to_n_term.unsqueeze(0) + closest_term = torch.minimum(dist_to_n_term.unsqueeze(0), dist_to_c_term) # (batch, L) + # Pairwise minimum over both positions: symmetric by construction. + terminus_distance = torch.minimum( + closest_term.unsqueeze(2), closest_term.unsqueeze(1) + ) # (batch, L, L) + terminus_idx = terminus_distance.clamp(0, self.max_terminus_distance) + terminus_bias = self.terminus(terminus_idx) # (batch, L, L, num_heads) + + combined = relative_bias.unsqueeze(0) + terminus_bias # (batch, L, L, num_heads) + return combined.permute(0, 3, 1, 2) # (batch, num_heads, L, L) + + +def _fold_padding_into_bias(bias: Tensor, key_padding_mask: Tensor) -> Tensor: + """ + Fold a boolean key-padding mask into an additive attention bias, so a + single float ``attn_mask`` can be passed to ``nn.MultiheadAttention`` + instead of separately supplying a bool ``key_padding_mask`` - + PyTorch deprecates mixing the two types on one call, and this changes + nothing behaviourally: an ignored key still gets ``-inf`` either way. + + ``bias``: ``(batch, heads, Lq, Lk)`` float. ``key_padding_mask``: + ``(batch, Lk)`` bool, ``True`` = ignore that key. Returns + ``(batch*heads, Lq, Lk)`` float, ready for ``attn_mask``. + """ + batch, heads, seq_q, seq_k = bias.shape + penalty = torch.zeros(batch, seq_k, dtype=bias.dtype, device=bias.device) + penalty = penalty.masked_fill(key_padding_mask, float("-inf")) + combined = bias + penalty.view(batch, 1, 1, seq_k) + return combined.reshape(batch * heads, seq_q, seq_k) + + +def _sdpa_attention( + mha: nn.MultiheadAttention, query: Tensor, key: Tensor, value: Tensor, attn_mask: Tensor +) -> Tensor: + """ + Equivalent to ``mha(query, key, value, attn_mask=attn_mask, need_weights=False)[0]`` + for how ``mha`` is always constructed in this module - ``batch_first=True``, + default ``kdim``/``vdim`` (so a single fused ``in_proj_weight``), no + ``bias_k``/``bias_v``, ``add_zero_attn=False`` - calling + ``scaled_dot_product_attention`` directly instead of going through + ``nn.MultiheadAttention.forward``'s generic path. + + That generic path reshapes q/k/v between its ``batch_first`` layout and + its internal layout via a transpose-and-``.contiguous()`` copy around + every call, *in addition to* whatever ``scaled_dot_product_attention`` + itself costs - `torch.profiler` on this trunk measured that reshaping at + ~21% of total self-CPU time, on top of the fused-kernel cost the + ``need_weights=False`` fix already secured. This reads the same + ``in_proj_weight``/``in_proj_bias``/``out_proj`` parameters ``mha`` + already owns, so an existing checkpoint's state_dict loads unchanged - + no retraining, no shape change, nothing to migrate. + + Only handles the two shapes of call this file makes: self-attention + (``query is key is value``) and cross-attention with a shared key/value + (``key is value``, ``query`` a different tensor) - the same two branches + ``torch.nn.functional._in_projection_packed`` special-cases, reproduced + here rather than delegated to, so this can skip straight to + ``scaled_dot_product_attention`` afterwards. + + Not bit-identical to ``mha(...)`` - a different kernel call order, same + numerical class as the ``need_weights=False`` change this sits on top of. + Verified to the same tight float32 tolerance that change was verified + against (see the CPU-timing investigation this accompanies). + """ + embed_dim = query.size(-1) + num_heads = mha.num_heads + head_dim = embed_dim // num_heads + batch = query.size(0) + + if query is key and key is value: + q, k, v = F.linear(query, mha.in_proj_weight, mha.in_proj_bias).chunk(3, dim=-1) + else: + assert key is value, "_sdpa_attention only supports self-attention or a shared key/value" + w_q, w_kv = mha.in_proj_weight.split([embed_dim, embed_dim * 2]) + if mha.in_proj_bias is None: + b_q, b_kv = None, None + else: + b_q, b_kv = mha.in_proj_bias.split([embed_dim, embed_dim * 2]) + q = F.linear(query, w_q, b_q) + k, v = F.linear(key, w_kv, b_kv).chunk(2, dim=-1) + + def split_heads(t: Tensor) -> Tensor: + return t.view(batch, -1, num_heads, head_dim).transpose(1, 2) + + q, k, v = split_heads(q), split_heads(k), split_heads(v) + # attn_mask arrives as _fold_padding_into_bias's (batch*heads, Lq, Lk) - + # a view back to 4D, exactly undoing that function's own final reshape + # (same contiguous memory, no copy), which is the shape + # scaled_dot_product_attention wants directly. + mask = attn_mask.view(batch, num_heads, attn_mask.size(-2), attn_mask.size(-1)) + dropout_p = mha.dropout if mha.training else 0.0 + attended = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=dropout_p) + attended = attended.transpose(1, 2).reshape(batch, -1, embed_dim) + return F.linear(attended, mha.out_proj.weight, mha.out_proj.bias) + + +class PreNormTransformerBlock(nn.Module): + """Pre-norm self-attention + GELU FFN block, both with a residual connection.""" + + def __init__(self, d_model: int, num_heads: int, ffn_dim: int, dropout: float): + super().__init__() + self.attn_norm = nn.LayerNorm(d_model) + self.self_attn = nn.MultiheadAttention( + d_model, num_heads, dropout=dropout, batch_first=True + ) + self.ffn_norm = nn.LayerNorm(d_model) + self.ffn = nn.Sequential( + nn.Linear(d_model, ffn_dim), + nn.GELU(), + nn.Linear(ffn_dim, d_model), + ) + + def forward( + self, x: Tensor, attn_mask: Tensor, return_attention: bool = False + ) -> Tensor | tuple[Tensor, Tensor]: + """``attn_mask``: ``(batch*num_heads, L, L)`` float, already carrying + both the pair/terminus bias and padding (as ``-inf``) - see + :func:`_fold_padding_into_bias`. Passed only as ``attn_mask``, not + also as ``key_padding_mask``: PyTorch deprecates mixing a float + ``attn_mask`` with a bool ``key_padding_mask`` on the same call. + + ``return_attention`` (default ``False``, never set by the training/ + predict paths - for attention inspection only): when + ``True``, also returns the per-head self-attention weights, + ``(batch, num_heads, L, L)``, via ``nn.MultiheadAttention`` itself + (``need_weights=True``) rather than :func:`_sdpa_attention` - the + weights aren't otherwise available, and requesting them forces + ``nn.MultiheadAttention`` onto its slower, unfused eager path + regardless (see the CPU-timing investigation this accompanies for + why the default path avoids that path entirely instead). Only pay + that cost when attention is actually wanted.""" + normed = self.attn_norm(x) + if return_attention: + attended, attn_weights = self.self_attn( + normed, + normed, + normed, + attn_mask=attn_mask, + need_weights=True, + average_attn_weights=False, + ) + else: + attended = _sdpa_attention(self.self_attn, normed, normed, normed, attn_mask) + x = x + attended + x = x + self.ffn(self.ffn_norm(x)) + if return_attention: + return x, attn_weights + return x + + +class PerceiverReadout(nn.Module): + """ + ``num_queries`` learned query vectors cross-attending to residue tokens. + + ``num_queries=1`` is plain (masked) attention pooling: a single learned + query reads a weighted combination of the valid residue positions. + Higher ``k`` is exposed for future multi-conformer work (step 9) but is + not otherwise treated specially here - the caller decides what to do + with a ``(batch, k, d_model)`` output when ``k>1``. + """ + + def __init__(self, d_model: int, num_heads: int, num_queries: int, dropout: float): + super().__init__() + self.num_queries = num_queries + self.queries = nn.Parameter(torch.randn(num_queries, d_model) * (1.0 / math.sqrt(d_model))) + self.cross_attn = nn.MultiheadAttention( + d_model, num_heads, dropout=dropout, batch_first=True + ) + + def forward( + self, residue_hidden: Tensor, attn_mask: Tensor, return_attention: bool = False + ) -> Tensor | tuple[Tensor, Tensor]: + """``attn_mask``: ``(batch*num_heads, num_queries, residue_length)`` + float, carrying only padding (as ``-inf``) - the read-out has no + pair/terminus bias of its own, see :func:`_fold_padding_into_bias`. + + ``return_attention``: see ``PreNormTransformerBlock.forward`` for why + this is opt-in and costs the fused kernel when requested. Returns the + cross-attention weights, ``(batch, num_heads, num_queries, + residue_length)`` - directly, per 8.10, "how much the model relied + on each residue to predict CCS".""" + batch_size = residue_hidden.size(0) + query = self.queries.unsqueeze(0).expand(batch_size, -1, -1) + if return_attention: + out, attn_weights = self.cross_attn( + query, + residue_hidden, + residue_hidden, + attn_mask=attn_mask, + need_weights=True, + average_attn_weights=False, + ) + else: + # See PreNormTransformerBlock.forward and _sdpa_attention for why + # this bypasses nn.MultiheadAttention.forward entirely. + out = _sdpa_attention( + self.cross_attn, query, residue_hidden, residue_hidden, attn_mask + ) + pooled = out.squeeze(1) if self.num_queries == 1 else out + if return_attention: + return pooled, attn_weights + return pooled + + +class TransformerTrunk(nn.Module): + """ + Fused transformer encoder replacing ``ConvAtomComp``/``ConvDiatomComp``. + + Unlike :class:`~im2deep._architectures.blocks.FlexTrunk` (step 8.4 items + 8+9), this does **not** also retire ``OneHot`` - the roadmap's own text + for this item describes only the atomic-composition token stream, not a + residue-identity fusion, and folding one-hot identity in here would be + testing an extra, untested idea under this item's name rather than the + item as specified. ``use_one_hot`` stays a live, independent switch. + """ + + def __init__( + self, + d_model: int, + num_layers: int, + num_heads: int, + ffn_dim: int, + max_relative_distance: int, + max_terminus_distance: int, + num_readout_queries: int = 1, + dropout: float = 0.0, + ): + super().__init__() + self.input_proj = nn.Linear(6, d_model) + self.charge_embed = nn.Embedding(7, d_model) + self.pair_bias = RelativePairBias(num_heads, max_relative_distance, max_terminus_distance) + self.blocks = nn.ModuleList( + [ + PreNormTransformerBlock(d_model, num_heads, ffn_dim, dropout) + for _ in range(num_layers) + ] + ) + self.final_norm = nn.LayerNorm(d_model) + self.readout = PerceiverReadout(d_model, num_heads, num_readout_queries, dropout) + self.num_heads = num_heads + + def forward( + self, atom_comp_raw: Tensor, charge_idx: Tensor, return_attention: bool = False + ) -> Tensor | tuple[Tensor, dict]: + """ + ``atom_comp_raw``: ``(batch, padding_length, 6)`` - the RAW, un- + permuted ``matrix`` tensor (unlike every conv branch in this + architecture, this module does its own layout, so the caller must + capture this before ``forward()``'s usual ``.permute(0, 2, 1)``). + ``charge_idx``: ``(batch,)`` long, 0-6 - the same raw charge index + ``use_charge_embedding`` already derives from ``CHARGE_FEATURE_INDEX``. + + ``return_attention`` (default ``False``, for attention inspection only): + when ``True``, also returns a dict with + ``"self_attention"`` (a list of per-layer ``(batch, heads, L+1, L+1)`` + tensors, register token included at index 0) and + ``"readout_attention"`` (``(batch, heads, num_queries, L)``, register + token excluded - see ``test_register_token_excluded_from_readout_keys``). + """ + batch_size, seq_len, _ = atom_comp_raw.shape + padding_mask = atom_comp_raw.sum(dim=-1) == 0 # (batch, seq_len); True = padding + + tokens = self.input_proj(atom_comp_raw) # (batch, seq_len, d_model) + charge_token = self.charge_embed(charge_idx).unsqueeze(1) # (batch, 1, d_model) + hidden = torch.cat([charge_token, tokens], dim=1) # (batch, seq_len+1, d_model) + + register_valid = torch.zeros(batch_size, 1, dtype=torch.bool, device=padding_mask.device) + key_padding_mask = torch.cat([register_valid, padding_mask], dim=1) # (batch, seq_len+1) + + real_lengths = (~padding_mask).sum(dim=1) # (batch,) + residue_bias = self.pair_bias(real_lengths, seq_len) # (batch, heads, seq_len, seq_len) + # Register token gets no pair/terminus bias of its own (row/column 0 + # stay exactly zero) - F.pad's constant-0 fill is the zero-alloc + + # slice-assign this replaces, in one fused op instead of two. + bias_full = F.pad(residue_bias, (1, 0, 1, 0)) + attn_mask = _fold_padding_into_bias(bias_full, key_padding_mask) + + self_attn_weights = [] if return_attention else None + for block in self.blocks: + if return_attention: + hidden, attn_weights = block(hidden, attn_mask=attn_mask, return_attention=True) + self_attn_weights.append(attn_weights) # ty: ignore[unresolved-attribute] + else: + hidden = block(hidden, attn_mask=attn_mask) + hidden = self.final_norm(hidden) + + residue_hidden = hidden[:, 1:, :] # register token excluded from read-out + residue_padding_mask = key_padding_mask[:, 1:] + # The read-out has no pair/terminus bias of its own: a size-1 zero + # tensor broadcasts through _fold_padding_into_bias's addition to + # the same result as a fully materialized zero bias, without + # allocating (batch, heads, num_queries, seq_len) of zeros for it. + readout_bias = residue_bias.new_zeros(batch_size, self.num_heads, 1, 1).expand( + batch_size, self.num_heads, self.readout.num_queries, seq_len + ) + readout_attn_mask = _fold_padding_into_bias(readout_bias, residue_padding_mask) + if return_attention: + pooled, readout_attn = self.readout( + residue_hidden, readout_attn_mask, return_attention=True + ) + return pooled, { + "self_attention": self_attn_weights, + "readout_attention": readout_attn, + } + return self.readout(residue_hidden, readout_attn_mask) diff --git a/im2deep/_data.py b/im2deep/_data.py index 825ebb4..a9e58ef 100644 --- a/im2deep/_data.py +++ b/im2deep/_data.py @@ -305,7 +305,7 @@ def build_training_dataset( # resident memory roughly 36x lower and avoids each DataLoader worker # forking a large object graph. return CCSDataset( - peptidoforms=peptidoforms, # type: ignore[arg-type] + peptidoforms=peptidoforms, # type: ignore[arg-type] # ty: ignore[invalid-argument-type] target_retention_times=targets, **featurisation, ) diff --git a/im2deep/_model_ops.py b/im2deep/_model_ops.py index 534beaf..944c6f0 100644 --- a/im2deep/_model_ops.py +++ b/im2deep/_model_ops.py @@ -5,24 +5,31 @@ import logging import warnings +from collections import Counter +from functools import lru_cache from os import PathLike from pathlib import Path import lightning as L +import numpy as np import torch from lightning.pytorch.callbacks import ( EarlyStopping, ModelCheckpoint, ModelSummary, RichProgressBar, + StochasticWeightAveraging, ) +from psm_utils.peptidoform import Peptidoform from rich.progress import track -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import DataLoader, Dataset, Sampler, Subset, WeightedRandomSampler from im2deep._architectures.callbacks import BackboneFreeze, LogLowestMAE from im2deep._architectures.im2deep_multi import IM2DeepMultiTransfer from im2deep._architectures.im2deep_single import IM2Deep, IM2DeepTransfer from im2deep._architectures.losses import FlexibleLossSorted +from im2deep._data import stripped_sequence +from im2deep.exceptions import IM2DeepError # Suppress PyTorch padding warning for conv1d with even kernels and odd dilation warnings.filterwarnings( @@ -93,11 +100,34 @@ def load_model( return loaded_model +def _load_model_for_inference( + model: torch.nn.Module | PathLike | str | None, + multi: bool, + device: str, +) -> torch.nn.Module: + """ + Load a checkpoint (or the bundled default) as an eval-mode model on ``device``. + + Factored out of :func:`predict` so other tooling can load a checkpoint the + same way - reading its own recorded config/architecture rather than + guessing - without duplicating this block. + """ + loaded_model = read_checkpoint_architecture(model, multi=multi).load_from_checkpoint( + checkpoint_path=model, # TODO: Match with function signature + config=read_checkpoint_config(model, multi=multi), + criterion=_get_loss_function(multi=multi), + strict=False, + ) + loaded_model.to(device) + loaded_model.eval() + return loaded_model + + def predict( model: torch.nn.Module | PathLike | str | None = None, data: Dataset | None = None, multi=False, - device: str = "cpu", + device: str | None = None, batch_size: int = 512, num_workers: int = 0, num_threads: int | None = None, @@ -112,7 +142,11 @@ def predict( data Dataset to predict on. device - Device to use for prediction. + Device to use for prediction. Defaults to GPU if available, else CPU + (same auto-detection :func:`load_model` and deeplc's own ``predict`` + use) - previously hardcoded to ``"cpu"``, which silently forced + CPU-only inference even on GPU-equipped machines unless a caller + explicitly passed ``device="cuda"``. batch_size Batch size for prediction. num_workers @@ -130,19 +164,13 @@ def predict( if data is None: raise ValueError("Data must be provided for prediction.") + device = device or ("cuda" if torch.cuda.is_available() else "cpu") torch.set_num_threads(num_threads or torch.get_num_threads()) # TODO: implement custom model inference LOGGER.debug("Loading model for prediction.") - # TODO: Implement load_model function here (also config) and path to default model? - model = read_checkpoint_architecture(model, multi=multi).load_from_checkpoint( - checkpoint_path=model, # type: ignore # TODO: Match with function signature - config=read_checkpoint_config(model, multi=multi), - criterion=_get_loss_function(multi=multi), - ) - model.to(device) - model.eval() + model = _load_model_for_inference(model, multi=multi, device=device) LOGGER.debug(f"Model loaded on device: {device}") data_loader = DataLoader( @@ -164,13 +192,21 @@ def _predict_loop( ) -> torch.Tensor: model.eval() all_predictions = [] - with torch.no_grad(): + with torch.inference_mode(): for features, _ in track(data_loader, description="Predicting", transient=True): features = [feature_tensor.to(device) for feature_tensor in features] outputs = model(*features) if not isinstance(outputs, tuple): # Single output all_predictions.append(outputs.cpu()) + elif getattr(model, "use_heteroscedastic_head", False): + # (mu, b) (step 8.5): only mu is a CCS prediction. b (the + # model's own per-row uncertainty) is not yet exposed + # through this public API - checked before the generic + # "stack both" branch below, which is for IM2DeepMulti's + # genuinely-two-conformer-predictions tuple, a different + # kind of forward() tuple this would otherwise collide with. + all_predictions.append(outputs[0].cpu()) else: # Multi-output: stack both predictions side by side stacked = torch.stack([outputs[0], outputs[1]], dim=1) @@ -179,6 +215,182 @@ def _predict_loop( return torch.cat(all_predictions, dim=0).squeeze() +def _training_targets(dataset: Dataset) -> torch.Tensor: + """The raw CCS target array for a training dataset, unwrapping a ``Subset`` if needed.""" + if isinstance(dataset, Subset): + targets = dataset.dataset.target_retention_times[dataset.indices] # ty: ignore[unresolved-attribute] + else: + targets = dataset.target_retention_times # ty: ignore[unresolved-attribute] + return torch.as_tensor(targets, dtype=torch.float32) + + +def _training_peptidoforms(dataset: Dataset) -> list[str]: + """The ProForma strings for a training dataset, unwrapping a ``Subset`` if needed.""" + if isinstance(dataset, Subset): + return [dataset.dataset.peptidoforms[i] for i in dataset.indices] # ty: ignore[unresolved-attribute] + return list(dataset.peptidoforms) # ty: ignore[unresolved-attribute] + + +def _charge_balanced_weights(train_dataset: Dataset) -> torch.Tensor: + """ + Per-sample weights proportional to ``1 / sqrt(n_z)`` (step 8.3). + + Reweights rare charges (z=1/5/6) toward being sampled as often as the + dominant ones without discarding any z=2/3 depth, unlike undersampling. + Cheap: reads charge off each ProForma string's own ``precursor_charge``, + not a full DeepLC featurisation pass. + """ + peptidoforms = _training_peptidoforms(train_dataset) + charges = np.array([Peptidoform(p).precursor_charge for p in peptidoforms]) + counts = Counter(charges.tolist()) + weights = np.array([1.0 / np.sqrt(counts[c]) for c in charges]) + return torch.as_tensor(weights, dtype=torch.double) + + +def _compute_pair_ids(peptidoforms: list[str]) -> tuple[np.ndarray, np.ndarray]: + """ + Sibling-group id and theoretical m/z per training row (step 8.3). + + ``group_id`` is an integer code per distinct stripped sequence (matching + :func:`im2deep._data.grouped_split`'s own grouping), used by the sibling + delta loss to find pairs that differ only in modification state or + charge. ``mz`` is each row's theoretical precursor m/z, used by the + ranking loss to find pairs a quadrupole could not separate. + + Both are computed once per training run, not per epoch: charge-balanced + weighting already parses every peptidoform once for the same reason (see + :func:`_charge_balanced_weights`), and re-parsing per batch would be far + too slow at millions of rows. Modifications change ``theoretical_mz`` but + not the group id - a phosphorylated and unmodified copy of the same + backbone are deliberately still siblings. + """ + stripped = [stripped_sequence(pf) for pf in peptidoforms] + _, group_ids = np.unique(stripped, return_inverse=True) + mz = np.array([Peptidoform(pf).theoretical_mz for pf in peptidoforms], dtype=np.float32) + return group_ids.astype(np.int64), mz + + +class _PairAugmentedDataset(Dataset): + """ + Wraps a training dataset to also yield each row's sibling-group id and + m/z, for the sibling delta and ranking losses (step 8.3). + + Wrapping rather than changing :class:`im2deep._data.CCSDataset` itself + keeps every other training/prediction path (5-tuple or 6-tuple with + ``mol_desc``) completely untouched; this only exists on the loader built + when ``sibling_loss_weight`` or ``ranking_loss_weight`` is set. + """ + + def __init__(self, base: Dataset, group_ids: np.ndarray, mz: np.ndarray): + self.base = base + self.group_ids = group_ids + self.mz = mz + + def __len__(self) -> int: + return len(self.base) # ty: ignore[invalid-argument-type] + + def __getitem__(self, idx: int) -> tuple: # ty: ignore[invalid-method-override] + item = self.base[idx] + group_id = torch.as_tensor(self.group_ids[idx], dtype=torch.long) + mz = torch.as_tensor(self.mz[idx], dtype=torch.float32) + return (*item, group_id, mz) + + +class SiblingPairBatchSampler(Sampler): + """ + Batch sampler that deliberately places sibling pairs inside otherwise + randomly composed batches (step 8.3's sibling delta loss). + + A plain random batch of 512 has ~0.07 sibling pairs on real training + data (measured), because most stripped sequences appear only once or + twice among 4M+ rows - nowhere near enough for a loss term supervising + pair deltas to see a training signal. Each yielded batch instead draws + ``pairs_per_batch`` random sibling groups (each contributing one pair, + sampled with replacement across the epoch, mirroring + ``WeightedRandomSampler(replacement=True)``'s own approach elsewhere in + this module) plus ordinary random singleton draws for the remaining + slots, so composition stays close to the overall data distribution. + """ + + def __init__( + self, + group_ids: np.ndarray, + batch_size: int, + pairs_per_batch: int, + num_batches: int, + ): + self.group_ids = np.asarray(group_ids) + self.batch_size = batch_size + self.pairs_per_batch = pairs_per_batch + self.num_batches = num_batches + self.n = len(group_ids) + + groups: dict[int, list[int]] = {} + for idx, group in enumerate(self.group_ids): + groups.setdefault(int(group), []).append(idx) + self.sibling_groups = [idxs for idxs in groups.values() if len(idxs) >= 2] + if not self.sibling_groups: + raise IM2DeepError( + "sibling_loss_weight is set, but no sibling groups (a stripped sequence " + "observed more than once) exist in the training data." + ) + + def __len__(self) -> int: + return self.num_batches + + def __iter__(self): + rng = np.random.default_rng() + n_groups = len(self.sibling_groups) + for _ in range(self.num_batches): + batch: list[int] = [] + for _ in range(self.pairs_per_batch): + group = self.sibling_groups[rng.integers(n_groups)] + pair = rng.choice(group, size=2, replace=False) + batch.extend(int(i) for i in pair) + remaining = self.batch_size - len(batch) + batch.extend(int(i) for i in rng.integers(0, self.n, size=remaining)) + rng.shuffle(batch) + yield batch + + +def _fit_scalers( + train_dataset: Dataset, config: dict +) -> tuple[torch.Tensor | None, torch.Tensor | None, float | None, float | None]: + """ + Fit the optional global-feature and target standardisation buffers (step 8.2). + + Returns ``(global_mean, global_std, target_mean, target_std)``, each pair + ``None`` when its config flag is off - the caller only calls the matching + ``set_*_scaler`` when its pair isn't ``None``, so an unfit statistic never + overwrites a model's no-op default. + """ + global_mean = global_std = None + target_mean = target_std = None + + if config.get("standardize_target"): + targets = _training_targets(train_dataset) + target_mean = float(targets.mean()) + target_std = float(targets.std()) + if target_std == 0: + target_std = 1.0 # degenerate all-equal sample; std=0 would divide by zero + + if config.get("standardize_global_features"): + sample_size = min(len(train_dataset), 50_000) # ty: ignore[invalid-argument-type] + generator = torch.Generator().manual_seed(0) + idx = torch.randperm(len(train_dataset), generator=generator)[:sample_size].tolist() # ty: ignore[invalid-argument-type] + loader = DataLoader( + Subset(train_dataset, idx), + batch_size=2048, + num_workers=config.get("num_workers", 0), + ) + feats = torch.cat([batch[2] for batch in loader], dim=0) + global_mean = feats.mean(dim=0) + global_std = feats.std(dim=0) + global_std[global_std == 0] = 1.0 # a constant feature column; avoid dividing by zero + + return global_mean, global_std, target_mean, target_std + + def train( train_dataset: Dataset, validation_dataset: Dataset, @@ -220,17 +432,80 @@ def train( output_dir = Path(output_dir) criterion = _get_loss_function(multi=False) + given_model = model if model is None: model = _build_model(config, criterion) + # Only for a freshly-built plain model: a model passed in (e.g. finetuning + # a loaded backbone) keeps whatever scaler state it already has, and so + # does a transfer model, whose backbone carries its own fitted scalers. + if model is not given_model and not isinstance(model, IM2DeepTransfer): + global_mean, global_std, target_mean, target_std = _fit_scalers(train_dataset, config) + if global_mean is not None: + LOGGER.info("Fitted global-feature standardisation from a training sample.") + model.set_global_feature_scaler(global_mean, global_std) # ty: ignore[call-non-callable] + if target_mean is not None: + LOGGER.info( + f"Fitted target standardisation: mean={target_mean:.2f}, std={target_std:.2f}." + ) + model.set_target_scaler(target_mean, target_std) # ty: ignore[call-non-callable] LOGGER.debug(f"Training {type(model).__name__}") - train_loader = DataLoader( - train_dataset, - batch_size=config["batch_size"], - shuffle=True, - num_workers=config.get("num_workers", 0), - persistent_workers=bool(config.get("num_workers", 0)), + # Step 8.3: the sibling delta and ranking losses need each row's sibling- + # group id and/or theoretical m/z, which the plain (atom, diatom, global, + # one_hot, target) tuple does not carry. Wrap the dataset to add them only + # when at least one of those losses is enabled, so every other run's + # batch tuple is unaffected. + needs_pair_ids = ( + config.get("sibling_loss_weight", 0.0) > 0 or config.get("ranking_loss_weight", 0.0) > 0 ) + if needs_pair_ids: + if config.get("mixup_alpha", 0.0) > 0: + raise IM2DeepError( + "sibling_loss_weight/ranking_loss_weight are not compatible with mixup_alpha>0: " + "mixup blends batch rows together, which would silently corrupt the sibling-" + "group id and m/z each row is supposed to carry." + ) + group_ids, mz = _compute_pair_ids(_training_peptidoforms(train_dataset)) + loader_dataset: Dataset = _PairAugmentedDataset(train_dataset, group_ids, mz) + else: + loader_dataset = train_dataset + + # shuffle and sampler are mutually exclusive on DataLoader, hence the + # if/else rather than always passing both. + if config.get("sibling_loss_weight", 0.0) > 0: + if config.get("charge_balanced_sampling"): + raise IM2DeepError( + "sibling_loss_weight and charge_balanced_sampling are both custom batch-" + "composition strategies and are not yet compatible with each other." + ) + train_loader = DataLoader( + loader_dataset, + batch_sampler=SiblingPairBatchSampler( + group_ids=group_ids, + batch_size=config["batch_size"], + pairs_per_batch=config.get("sibling_pairs_per_batch", 64), + num_batches=len(loader_dataset) // config["batch_size"], # ty: ignore[invalid-argument-type] + ), + num_workers=config.get("num_workers", 0), + persistent_workers=bool(config.get("num_workers", 0)), + ) + elif config.get("charge_balanced_sampling"): + weights = _charge_balanced_weights(train_dataset) + train_loader = DataLoader( + loader_dataset, + batch_size=config["batch_size"], + sampler=WeightedRandomSampler(weights, num_samples=len(weights), replacement=True), # ty: ignore[invalid-argument-type] + num_workers=config.get("num_workers", 0), + persistent_workers=bool(config.get("num_workers", 0)), + ) + else: + train_loader = DataLoader( + loader_dataset, + batch_size=config["batch_size"], + shuffle=True, + num_workers=config.get("num_workers", 0), + persistent_workers=bool(config.get("num_workers", 0)), + ) validation_loader = DataLoader( validation_dataset, batch_size=config["batch_size"], @@ -248,18 +523,31 @@ def train( logger=_setup_logger(config, model, output_dir), default_root_dir=str(output_dir), enable_progress_bar=True, + precision=config.get("precision", "32-true"), ) - trainer.fit(model, train_loader, validation_loader) + trainer.fit(model, train_loader, validation_loader) # ty: ignore[invalid-argument-type] best_checkpoint_path = None - if checkpoint_callback is not None and checkpoint_callback.best_model_path: + if config.get("swa"): + # StochasticWeightAveraging transfers its averaged weights into + # `model` in `on_train_end` - after the last validation epoch, since + # IM2Deep has no BatchNorm layers to trigger Lightning's extra + # BN-recalibration epoch. No validation ever scores the averaged + # weights, so `checkpoint_callback.best_model_path` never points at + # them: reloading it here would silently throw the averaging away. + # `model` already holds the averaged weights post-fit, so use it + # as-is instead of reloading anything. + LOGGER.info( + "SWA enabled: using the final weight-averaged model, not a reloaded checkpoint." + ) + elif checkpoint_callback is not None and checkpoint_callback.best_model_path: best_checkpoint_path = checkpoint_callback.best_model_path LOGGER.info(f"Reloading best checkpoint: {best_checkpoint_path}") - model = type(model).load_from_checkpoint( + model = type(model).load_from_checkpoint( # ty: ignore[unresolved-attribute] best_checkpoint_path, config=config, criterion=criterion ) - return trainer, model, best_checkpoint_path + return trainer, model, best_checkpoint_path # ty: ignore[invalid-return-type] def _build_model(config: dict, criterion: torch.nn.Module) -> torch.nn.Module: @@ -307,6 +595,15 @@ def _setup_callbacks( ) ) + if config.get("swa"): + callbacks.append( + StochasticWeightAveraging( + swa_lrs=config.get("swa_lrs", config["learning_rate"]), + swa_epoch_start=config.get("swa_epoch_start", 0.8), + annealing_epochs=config.get("swa_annealing_epochs", 10), + ) + ) + return callbacks, checkpoint_callback @@ -325,7 +622,7 @@ def _setup_logger(config: dict, model: torch.nn.Module, output_dir: Path | None return None try: - import wandb # noqa: F401 + import wandb # noqa: F401 # ty: ignore[unresolved-import] from lightning.pytorch.loggers import WandbLogger except ImportError as exc: raise ImportError( @@ -338,13 +635,10 @@ def _setup_logger(config: dict, model: torch.nn.Module, output_dir: Path | None name=wandb_config.get("name") or config.get("model_name"), entity=wandb_config.get("entity"), tags=wandb_config.get("tags"), - save_dir=str(output_dir) if output_dir is not None else None, + save_dir=str(output_dir) if output_dir is not None else None, # ty: ignore[invalid-argument-type] ) # Log the training config so runs can be told apart and compared later. - logger.experiment.config.update( - {key: value for key, value in config.items() if key != "wandb"}, - allow_val_change=True, - ) + logger.log_hyperparams({key: value for key, value in config.items() if key != "wandb"}) logger.watch(model) LOGGER.info( f"Logging to Weights & Biases project " @@ -362,8 +656,18 @@ def _get_architecture(multi: bool) -> type[IM2DeepMultiTransfer] | type[IM2Deep] return IM2Deep +@lru_cache(maxsize=8) def _read_checkpoint(model: torch.nn.Module | PathLike | str | None) -> dict | None: - """Load a checkpoint file as a dict, or None if it is not one.""" + """Load a checkpoint file as a dict, or None if it is not one. + + Cached: a single ``core.predict()`` call reads the same checkpoint path + through this function up to three times (once directly, once via + :func:`read_checkpoint_architecture`, once via :func:`read_checkpoint_config` + called again inside :func:`predict`), each a full ``torch.load`` of the + same bytes. Callers never mutate the returned dict (or the sub-dicts, + e.g. the config, handed back from it), so sharing one cached object + across those calls is safe. + """ if not isinstance(model, (str, Path)): return None diff --git a/im2deep/calibration.py b/im2deep/calibration.py index d875da5..974c9d8 100644 --- a/im2deep/calibration.py +++ b/im2deep/calibration.py @@ -13,12 +13,35 @@ import numpy as np import pandas as pd from psm_utils import Peptidoform, PSMList +from sklearn.linear_model import LinearRegression from im2deep.constants import DEFAULT_MULTI_REFERENCE_DATASET_PATH, DEFAULT_REFERENCE_DATASET_PATH from im2deep.exceptions import CalibrationError LOGGER = logging.getLogger(__name__) +#: A charge with fewer overlapping peptides than this falls back to a single +#: model fit across every charge's overlap together, mirroring +#: `collate_dataset.py`'s `MIN_OVERLAP_PER_CHARGE` (same value, kept as a +#: separate constant since this module cannot import a top-level project +#: script as a dependency). +MIN_OVERLAP_PER_CHARGE = 10 + + +def _peptide_key(pf) -> str: + """Stripped-of-charge ProForma string, for merging target/source on identity.""" + if isinstance(pf, Peptidoform): + return str(pf.proforma).rsplit("/", 1)[0] + return str(pf).rsplit("/", 1)[0] + + +def _charge_of(pf) -> int: + if isinstance(pf, Peptidoform): + if pf.precursor_charge is None: + raise CalibrationError(f"Peptidoform {pf} is missing precursor charge information.") + return int(pf.precursor_charge) + return int(str(pf).split("/")[-1]) + class Calibration(ABC): """Abstract base class for CCS calibration methods.""" @@ -78,11 +101,32 @@ def __init__(self, per_charge: bool = True, use_charge_state: int | None = None) self.general_shift: float | None = None self.used_charges: set[int] = set() self.reference_psm_list: PSMList | None = None + # Calibration-confidence bookkeeping, populated by fit(): how many + # overlapping anchors were found for a charge, whether that charge's + # shift fell back to the general/all-charge shift, and the std of + # (target - source) residuals around the fitted shift. Read back via + # stats_for_charge() - not used internally, exposed for rescoring + # features that want to judge how much to trust a given PSM's + # calibrated CCS. + self.charge_n_anchors: dict[int, int] = {} + self.charge_used_fallback: dict[int, bool] = {} + self.charge_residual_std: dict[int, float] = {} @property def is_fitted(self) -> bool: return self.fitted + def stats_for_charge(self, charge: int) -> dict: + """Calibration-confidence bookkeeping for one charge state, post-fit.""" + charge = int(charge) + return { + "slope": 1.0, + "intercept": self.charge_shifts.get(charge, self.general_shift or 0.0), + "n_anchors": self.charge_n_anchors.get(charge, 0), + "used_fallback": self.charge_used_fallback.get(charge, True), + "residual_std": self.charge_residual_std.get(charge, float("nan")), + } + def fit( self, psm_df_target: pd.DataFrame, @@ -110,6 +154,8 @@ def fit( f"Could not calculate charge-specific shift factors: {e}. Using 0.0 as fallback." ) self.charge_shifts = dict.fromkeys(range(1, 7), 0.0) + self.charge_n_anchors = dict.fromkeys(range(1, 7), 0) + self.charge_residual_std = dict.fromkeys(range(1, 7), float("nan")) # Set general shift as the mean of calculated charge shifts or charge 2 if available if 2 in self.charge_shifts and self.charge_shifts[2] != 0.0: @@ -136,6 +182,10 @@ def fit( ): no_shift_calculated.append(charge) self.charge_shifts[charge] = float(self.general_shift) + self.charge_n_anchors.setdefault(charge, 0) + self.charge_residual_std.setdefault(charge, float("nan")) + for charge in range(1, 7): + self.charge_used_fallback[charge] = charge in no_shift_calculated LOGGER.debug( f"No shift factor calculated for charge states: {no_shift_calculated}. " f"Using general shift: {self.general_shift:.3f}." @@ -153,6 +203,13 @@ def fit( ) self.general_shift = 0.0 self.charge_shifts = dict.fromkeys(range(1, 7), self.general_shift) + # Global (non-per-charge) mode doesn't track anchors per charge - + # not exercised by this project's pipeline (always per_charge=True), + # but keep the bookkeeping dicts populated so stats_for_charge() + # never KeyErrors if this path is ever used. + self.charge_n_anchors = dict.fromkeys(range(1, 7), 0) + self.charge_used_fallback = dict.fromkeys(range(1, 7), True) + self.charge_residual_std = dict.fromkeys(range(1, 7), float("nan")) self.used_charges = set(self.charge_shifts.keys()) self.fitted = True @@ -385,8 +442,8 @@ def get_charge(pf): return float(shift) - @staticmethod def _compute_ccs_shift_per_charge( + self, target_df: pd.DataFrame, source_df: pd.DataFrame, ) -> dict[int, float]: @@ -466,6 +523,16 @@ def get_charge(pf): int(cast(int, charge)): float(shift) for charge, shift in shift_by_charge.items() } + # Calibration-confidence bookkeeping (idea 1): anchor count and + # residual spread per charge, computed here since this is the one + # place `merged`'s per-charge residuals already exist - discarded by + # the rest of this method beyond the small-overlap warning below. + residual_std_by_charge = merged.groupby("charge")["shift"].std() + for charge, count in merged.groupby("charge").size().items(): + self.charge_n_anchors[int(cast(int, charge))] = int(cast(int, count)) + for charge, std in residual_std_by_charge.items(): + self.charge_residual_std[int(cast(int, charge))] = float(std) + # Log information for each charge state charge_counts = merged.groupby("charge").size() for charge, count in charge_counts.items(): @@ -488,6 +555,215 @@ def get_charge(pf): return shift_factors +class RegressionCCSCalibration(Calibration): + """ + Per-charge slope + intercept calibration for CCS predictions. + + ``LinearCCSCalibration`` fits an additive-only shift, which is the right + calibration for a model whose training labels sit on the reference's own + scale, and it is what the bundled model is validated with. It can be the + wrong one for a model whose training + labels were harmonised with a regression-based collation method (e.g. + ``collate_dataset.py``'s ``linreg``, per-charge slope + intercept): fitting + only a shift against a reference on the model's own scale still discards + the multiplicative component that method's own harmonisation applied. This + class fits ``target = slope * source + intercept`` per charge instead, + mirroring ``collate_dataset.py``'s ``RegressionCalibrator`` (see + ROADMAP.md step 6b, "Only `shift` exists on both sides", for why the two + need to differ rather than sharing one calibrator). + + A charge with fewer than ``MIN_OVERLAP_PER_CHARGE`` overlapping peptides + falls back to a single model fit across every charge's overlap together, + the same fallback ``RegressionCalibrator`` uses and for the same reason: a + regression already carries a slope, so a transform fitted across all + charges still scales with CCS rather than assuming a fixed offset. + + Parameters + ---------- + per_charge : bool, optional + Whether to fit a separate slope/intercept per charge state, falling + back to one general model where a charge's own overlap is too thin. + Default is True. + """ + + def __init__(self, per_charge: bool = True) -> None: + super().__init__() + self.per_charge = per_charge + self.fitted = False + self.charge_models: dict[int, LinearRegression] = {} + self.general_model: LinearRegression | None = None + # Calibration-confidence bookkeeping, populated by fit() - see + # LinearCCSCalibration's identical attributes for the rationale. + self.charge_n_anchors: dict[int, int] = {} + self.charge_used_fallback: dict[int, bool] = {} + self.charge_residual_std: dict[int, float] = {} + + @property + def is_fitted(self) -> bool: + return self.fitted + + def stats_for_charge(self, charge: int) -> dict: + """Calibration-confidence bookkeeping for one charge state, post-fit.""" + charge = int(charge) + model = self.charge_models.get(charge, self.general_model) + slope = float(model.coef_[0]) if model is not None else 1.0 + intercept = float(model.intercept_) if model is not None else 0.0 + return { + "slope": slope, + "intercept": intercept, + "n_anchors": self.charge_n_anchors.get(charge, 0), + "used_fallback": self.charge_used_fallback.get(charge, True), + "residual_std": self.charge_residual_std.get(charge, float("nan")), + } + + def fit( + self, + psm_df_target: pd.DataFrame, + psm_df_source: pd.DataFrame | None = None, + multi: bool = False, + ) -> None: + """ + Fit ``target = slope * source + intercept`` per charge. + + ``psm_df_target`` is this run's own confidently-identified, + observed CCS (real scale); ``psm_df_source`` is the reference dataset, + which for this calibrator must be on the same scale as the model's raw + predictions (e.g. a reference built from the model's own training + pool), so the fitted mapping can be applied directly to + ``predicted_CCS_uncalibrated`` in `transform`. + """ + if psm_df_source is None: + LOGGER.debug("No reference PSMList provided, loading default reference dataset.") + psm_df_source = get_default_reference(multi=multi) + + merged = self._merge_on_peptide_and_charge(psm_df_target, psm_df_source) + + # Calibration-confidence bookkeeping (idea 1): every charge starts as + # "no anchors seen at all", overwritten below for charges that + # actually appear in `merged`. + self.charge_n_anchors = dict.fromkeys(range(1, 7), 0) + self.charge_used_fallback = dict.fromkeys(range(1, 7), True) + self.charge_residual_std = dict.fromkeys(range(1, 7), float("nan")) + thin_charge_groups: dict[int, pd.DataFrame] = {} + + if self.per_charge: + for charge, group in merged.groupby("charge"): + charge_int = int(cast(int, charge)) + self.charge_n_anchors[charge_int] = len(group) + if len(group) < MIN_OVERLAP_PER_CHARGE: + LOGGER.debug( + f"Only {len(group)} overlapping peptides for charge {charge}; " + "deferring to the general fallback model." + ) + thin_charge_groups[charge_int] = group + continue + model = LinearRegression() + model.fit(group[["CCS_source"]].to_numpy(), group["CCS_target"].to_numpy()) + self.charge_models[charge_int] = model + self.charge_used_fallback[charge_int] = False + residuals = group["CCS_target"].to_numpy() - model.predict( + group[["CCS_source"]].to_numpy() + ) + self.charge_residual_std[charge_int] = float(np.std(residuals)) + + if len(merged) >= MIN_OVERLAP_PER_CHARGE: + self.general_model = LinearRegression() + self.general_model.fit( + merged[["CCS_source"]].to_numpy(), merged["CCS_target"].to_numpy() + ) + elif not self.charge_models: + # No overlap at all is expected, not just possible, on a genuinely + # unseen dataset (e.g. a rare-organism downstream check against a + # reference sampled from an unrelated training pool) - exactly the + # case this calibrator most needs to keep working for. Fall back + # to an identity transform (slope=1, intercept=0) with a WARNING, + # mirroring LinearCCSCalibration's own "fall back to 0.0 shift + # rather than crash" behaviour under the same starvation. + LOGGER.warning( + f"Only {len(merged)} overlapping peptides in total; falling back to an " + "identity transform (no calibration applied) rather than fitting a " + f"general regression calibration (needs >= {MIN_OVERLAP_PER_CHARGE})." + ) + self.general_model = LinearRegression() + self.general_model.coef_ = np.array([1.0]) + self.general_model.intercept_ = 0.0 + + # Now that general_model exists, compute residual std for charges + # that fell back to it (thin per-charge overlap, still non-zero). + if self.general_model is not None: + for charge_int, group in thin_charge_groups.items(): + residuals = group["CCS_target"].to_numpy() - self.general_model.predict( + group[["CCS_source"]].to_numpy() + ) + self.charge_residual_std[charge_int] = float(np.std(residuals)) + + self.fitted = True + + def transform(self, psm_df: pd.DataFrame) -> np.ndarray: + """Apply the fitted per-charge (or general fallback) regression.""" + if not self.is_fitted: + raise CalibrationError("Calibration has not been fitted yet.") + + if "peptidoform" not in psm_df.columns: + raise CalibrationError("Input DataFrame must contain 'peptidoform' column.") + + if "predicted_CCS_uncalibrated" not in psm_df.columns and "metadata" in psm_df.columns: + psm_df["predicted_CCS_uncalibrated"] = psm_df["metadata"].apply( + lambda x: ( + x["predicted_CCS_uncalibrated"] + if "predicted_CCS_uncalibrated" in x + else np.nan + ) + ) + + charges = psm_df["peptidoform"].apply(_charge_of).to_numpy() + source = psm_df["predicted_CCS_uncalibrated"].to_numpy(dtype=float).reshape(-1, 1) + out = np.empty(len(psm_df), dtype=object) + + for charge in np.unique(charges): + mask = charges == charge + model = self.charge_models.get(int(charge), self.general_model) + if model is None: + # Only reached if fit() saw less overlap than + # MIN_OVERLAP_PER_CHARGE in total, which fit() itself refuses + # (raises CalibrationError) - a guard, not a real code path. + out[mask] = source[mask, 0] + else: + out[mask] = model.predict(source[mask]).astype(np.float32) + + return out + + @staticmethod + def _merge_on_peptide_and_charge( + target_df: pd.DataFrame, source_df: pd.DataFrame + ) -> pd.DataFrame: + target_work = target_df.copy() + source_work = source_df.copy() + + target_work["peptide_key"] = target_work["peptidoform"].apply(_peptide_key) + target_work["charge"] = target_work["peptidoform"].apply(_charge_of) + if "CCS" not in target_work.columns and "metadata" in target_work.columns: + target_work["CCS"] = target_work["metadata"].apply( + lambda x: float(x["CCS"]) if isinstance(x, dict) and "CCS" in x else np.nan + ) + + source_work["peptide_key"] = source_work["peptidoform"].apply(_peptide_key) + source_work["charge"] = source_work["peptidoform"].apply(_charge_of) + if "CCS" not in source_work.columns and "metadata" in source_work.columns: + source_work["CCS"] = source_work["metadata"].apply( + lambda x: float(x["CCS"]) if isinstance(x, dict) and "CCS" in x else np.nan + ) + + merged = pd.merge( + target_work[["peptide_key", "charge", "CCS"]], + source_work[["peptide_key", "charge", "CCS"]], + on=["peptide_key", "charge"], + suffixes=("_target", "_source"), + ) + LOGGER.debug(f"Number of overlapping peptides for regression calibration: {len(merged)}") + return merged + + def get_default_reference(multi: bool = False) -> pd.DataFrame: """ Get the default reference DataFrame for calibration. diff --git a/im2deep/constants.py b/im2deep/constants.py index 3758d1c..032a7f9 100644 --- a/im2deep/constants.py +++ b/im2deep/constants.py @@ -3,12 +3,20 @@ # Paths and names for default models and reference datasets DEFAULT_MODEL_NAME = "IM2DeepUni.ckpt" DEFAULT_MODEL = Path(__file__).resolve().parent / "models" / "TIMS" / DEFAULT_MODEL_NAME +# MALDI model: the default model's architecture, trained on a collated set with more long +# (12-20 AA) charge 1 peptides, which MALDI yields and ESI rarely does. Its training data is +# still ESI timsTOF data. It shares the default calibration reference, which sits on its scale. +MALDI_MODEL_NAME = "IM2DeepMALDI.ckpt" +MALDI_MODEL = Path(__file__).resolve().parent / "models" / "TIMS" / MALDI_MODEL_NAME DEFAULT_MULTI_MODEL_NAME = "IM2DeepMulti.ckpt" DEFAULT_MULTI_MODEL = ( Path(__file__).resolve().parent / "models" / "TIMS" / DEFAULT_MULTI_MODEL_NAME ) +# Calibration reference for the default model: its CCS values sit on that model's scale, +# which is what the per-charge shift needs. The two are replaced together, so a new default +# model means a new reference file. DEFAULT_REFERENCE_DATASET_PATH = ( - Path(__file__).resolve().parent / "reference_data" / "reference_ccs.csv.gz" + Path(__file__).resolve().parent / "reference_data" / "reference_ccs_2026.csv.gz" ) DEFAULT_MULTI_REFERENCE_DATASET_PATH = ( Path(__file__).parent / "reference_data" / "multi_reference_ccs.csv.gz" @@ -98,27 +106,56 @@ "model_name": "IM2Deep", # Training loop "epochs": 100, - "patience": 10, + # None disables `EarlyStopping`, so + # set `patience` to a number of epochs to opt back into stopping early instead. + "patience": None, "use_best_model": True, - # "Validation MAE" is the metric name `LogLowestMAE` reads and the one the - # bundled IM2DeepUni.ckpt's own ModelCheckpoint monitored; changing it - # means changing both. "monitor": "Validation MAE", "mode": "min", "num_workers": 0, "accelerator": "auto", "devices": "auto", "wandb": {"enabled": False, "project_name": "IM2Deep"}, - # Featurisation. These are recorded in the checkpoint so `predict()` reads - # a model back with the encoding it was trained on. "add_ccs_features": True, "add_terminal_composition": False, - # DeepLC 4.1.0 defaults this to True, reproducing the pre-4.0.1 positional - # modification indexing the bundled checkpoints were trained with. Set to - # False to train against the 4.0.1 fix instead. "legacy_positional_deltas": True, "padding_length": 60, "Global_features": DEFAULT_GLOBAL_FEATURES, + "standardize_global_features": False, + "standardize_target": False, + "weight_decay": 0.0, + "lr_scheduler": None, # "one_cycle" is the only implemented option so far + "precision": "32-true", + "swa": False, + "loss": "l1", + "use_charge_embedding": False, + "charge_balanced_sampling": False, + "use_one_hot": True, + "use_attention_pooling": False, + "mixup_alpha": 0.0, + "sibling_loss_weight": 0.0, + "sibling_pairs_per_batch": 64, + "ranking_loss_weight": 0.0, + "ranking_mz_window": 3.0, + "ranking_margin": 0.0, + "use_flex_trunk": False, + "flex_trunk_embed_dim": 16, + "flex_trunk_channels": [192, 192], + "flex_trunk_kernel_size": 5, + "flex_trunk_stem_channels": 64, + "flex_trunk_stem_layers": 2, + "use_transformer_trunk": False, + "transformer_d_model": 128, + "transformer_num_layers": 4, + "transformer_num_heads": 4, + "transformer_ffn_dim": 512, + "transformer_max_relative_distance": 32, + "transformer_max_terminus_distance": 16, + "transformer_readout_queries": 1, + "transformer_dropout": 0.0, + "use_heteroscedastic_head": False, + "heteroscedastic_warmup_epochs": 0, + "use_per_charge_heads": False, } BASEMODELCONFIG = { diff --git a/im2deep/core.py b/im2deep/core.py index 63ee1c9..d9997a5 100644 --- a/im2deep/core.py +++ b/im2deep/core.py @@ -208,7 +208,7 @@ def _resolve_config( if config is not None: if isinstance(config, (str, Path, PathLike)): LOGGER.info(f"Reading training configuration from {config}") - with open(config) as config_file: + with open(config) as config_file: # ty: ignore[invalid-argument-type] config = json.load(config_file) if not isinstance(config, dict): raise TypeError(f"config must be a dict or a path, got {type(config).__name__}.") @@ -309,7 +309,7 @@ def train( ) LOGGER.info("Training IM2Deep model...") - trainer, model, best_checkpoint_path = _model_ops.train( + trainer, model, best_checkpoint_path = _model_ops.train( # ty: ignore[invalid-assignment] train_dataset=train_subset, validation_dataset=validation_subset, config=resolved_config, diff --git a/im2deep/models/TIMS/IM2DeepMALDI.ckpt b/im2deep/models/TIMS/IM2DeepMALDI.ckpt new file mode 100644 index 0000000..520af31 Binary files /dev/null and b/im2deep/models/TIMS/IM2DeepMALDI.ckpt differ diff --git a/im2deep/models/TIMS/IM2DeepUni.ckpt b/im2deep/models/TIMS/IM2DeepUni.ckpt index 26c6ec2..17ebf8f 100644 Binary files a/im2deep/models/TIMS/IM2DeepUni.ckpt and b/im2deep/models/TIMS/IM2DeepUni.ckpt differ diff --git a/im2deep/reference_data/reference_ccs.csv.gz b/im2deep/reference_data/reference_ccs_2026.csv.gz similarity index 58% rename from im2deep/reference_data/reference_ccs.csv.gz rename to im2deep/reference_data/reference_ccs_2026.csv.gz index 5e16071..ee45e28 100644 Binary files a/im2deep/reference_data/reference_ccs.csv.gz and b/im2deep/reference_data/reference_ccs_2026.csv.gz differ diff --git a/pyproject.toml b/pyproject.toml index c46c333..8e05adf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "rich>=13,<15", "torch>=2.6.0,<3", "lightning>=2.1,<3", + "scikit-learn>=1.3,<2", ] [project.optional-dependencies] diff --git a/tests/test_architectures.py b/tests/test_architectures.py new file mode 100644 index 0000000..710086f --- /dev/null +++ b/tests/test_architectures.py @@ -0,0 +1,1549 @@ +"""Tests for the IM2Deep architecture itself (im2deep/_architectures/im2deep_single.py).""" + +import copy + +import pytest +import torch +from torch import nn + +from im2deep._architectures.im2deep_single import CHARGE_FEATURE_INDEX, IM2Deep +from im2deep.constants import DEFAULT_TRAINING_CONFIG + + +def _model(**config_overrides): + config = {**copy.deepcopy(DEFAULT_TRAINING_CONFIG), "Global_features": 60, **config_overrides} + return IM2Deep(config, nn.L1Loss()) + + +def _dummy_batch(batch_size=2, global_features=60, charges=None): + global_feats = torch.randn(batch_size, global_features) * 50 + 100 + if charges is not None: + global_feats[:, CHARGE_FEATURE_INDEX] = torch.as_tensor(charges, dtype=torch.float32) + return ( + torch.randn(batch_size, 60, 6), + torch.randn(batch_size, 30, 6), + global_feats, + torch.randn(batch_size, 60, 20), # (seq_len, one-hot vocab), permuted inside forward() + ) + + +class TestStandardizationBuffers: + """Step 8.2: global-feature and target standardisation, opt-in via buffers.""" + + def test_default_buffers_are_a_no_op(self): + """A model that never calls set_*_scaler has mean=0/std=1 buffers.""" + model = _model() + + assert torch.equal(model.global_feat_mean, torch.zeros(60)) + assert torch.equal(model.global_feat_std, torch.ones(60)) + assert model.target_mean.item() == 0.0 + assert model.target_std.item() == 1.0 + + def test_explicit_no_op_scaler_matches_default(self): + """Setting mean=0/std=1 explicitly reproduces the untouched default output.""" + model = _model() + batch = _dummy_batch() + before = model(*batch) + + model.set_global_feature_scaler(torch.zeros(60), torch.ones(60)) + model.set_target_scaler(0.0, 1.0) + after = model(*batch) + + assert torch.allclose(before, after) + + def test_target_scaler_destandardizes_output(self): + """forward() always returns raw-scale predictions, regardless of the fitted scaler.""" + model = _model() + batch = _dummy_batch() + + model.set_target_scaler(0.0, 1.0) + raw_space_output = model(*batch) + + model.set_target_scaler(500.0, 200.0) + scaled_output = model(*batch) + + # Only the final affine step differs: scaled = raw * 200 + 500. The + # two calls otherwise share every weight, so this isolates exactly + # the transform forward() is supposed to apply. + assert torch.allclose(scaled_output, raw_space_output * 200.0 + 500.0, atol=1e-4) + + def test_global_feature_scaler_changes_output(self): + """A non-trivial global-feature scaler measurably changes the prediction.""" + model = _model() + batch = _dummy_batch() + before = model(*batch) + + model.set_global_feature_scaler(torch.full((60,), 100.0), torch.full((60,), 50.0)) + after = model(*batch) + + assert not torch.allclose(before, after) + + def test_buffers_survive_a_strict_checkpoint_round_trip(self, tmp_path): + """A freshly-trained checkpoint carries the fitted buffers, not the no-op default.""" + model = _model() + model.set_global_feature_scaler(torch.arange(60, dtype=torch.float32), torch.full((60,), 2.0)) + model.set_target_scaler(555.5, 123.4) + + path = tmp_path / "buffers.ckpt" + torch.save({"state_dict": model.state_dict(), "hyper_parameters": {"config": model.config}}, path) + + reloaded = IM2Deep(model.config, nn.L1Loss()) + missing, unexpected = reloaded.load_state_dict( + torch.load(path, weights_only=False)["state_dict"], strict=True + ) + assert not missing and not unexpected + assert torch.equal(reloaded.global_feat_mean, torch.arange(60, dtype=torch.float32)) + assert reloaded.target_mean.item() == 555.5 + + def test_old_checkpoint_missing_buffers_loads_with_strict_false(self): + """A checkpoint predating these buffers (e.g. the bundled models) still loads.""" + model = _model() + state_dict = model.state_dict() + for key in ("global_feat_mean", "global_feat_std", "target_mean", "target_std"): + del state_dict[key] + + reloaded = _model() + missing, unexpected = reloaded.load_state_dict(state_dict, strict=False) + + assert set(missing) == {"global_feat_mean", "global_feat_std", "target_mean", "target_std"} + assert not unexpected + # Falls back to the no-op default rather than crashing or zeroing. + assert torch.equal(reloaded.global_feat_std, torch.ones(60)) + assert reloaded.target_std.item() == 1.0 + + +class TestLossVariants: + """Step 8.3: the training objective (config['loss']), always reported as raw-Ų MAE.""" + + def test_default_is_plain_l1(self): + """"l1" (the default) is byte-identical to the pre-8.3 behaviour: self.criterion(y_hat, y).""" + model = _model() + y_hat = torch.tensor([510.0, 300.0]) + y = torch.tensor([500.0, 320.0]) + + assert model._base_loss(y_hat, y).item() == pytest.approx(nn.L1Loss()(y_hat, y).item()) + + def test_mape_is_mean_relative_error_as_a_percentage(self): + """Reported as a percentage (x100), matching benchmark_models.py's mean_abs_pct_error. + + Not cosmetic: the unscaled fraction's gradient carries a 1/y factor + that measured ~200x smaller than plain l1's at the same + initialisation and left training stuck near its starting point + within a normal epoch budget - see test_mape_gradient_is_not_starved. + """ + model = _model(loss="mape") + y_hat = torch.tensor([110.0, 190.0]) + y = torch.tensor([100.0, 200.0]) + + # |10|/100 = 0.10, |10|/200 = 0.05 -> mean 0.075 -> x100 = 7.5% + assert model._base_loss(y_hat, y).item() == pytest.approx(7.5) + + def test_mape_gradient_is_not_starved(self): + """ + The x100 scaling keeps mape's gradient within the same order of + magnitude as plain l1's, unlike the raw (unscaled) fraction. + + Measured directly on real proxy data: the unscaled mape's last-layer + gradient norm was ~200x smaller than l1's at identical + initialisation, and log_l1's (which has no equivalent rescaling + available) was ~500x smaller - both left training stuck at a + roughly constant MAE for the entire run. This is a coarser but + cheap proxy for that same effect: mape's gradient should not be + starved by more than about one order of magnitude relative to l1's. + """ + y_hat = torch.tensor([510.0, 300.0, 700.0], requires_grad=True) + y = torch.tensor([500.0, 320.0, 650.0]) + + model = _model(loss="l1") + model._base_loss(y_hat, y).backward() + l1_grad_norm = y_hat.grad.norm().item() + + y_hat.grad = None + model_mape = _model(loss="mape") + model_mape._base_loss(y_hat, y).backward() + mape_grad_norm = y_hat.grad.norm().item() + + assert mape_grad_norm > l1_grad_norm / 10 + + def test_huber_matches_torch_reference(self): + model = _model(loss="huber") + y_hat = torch.tensor([510.0, 300.0]) + y = torch.tensor([500.0, 320.0]) + + expected = nn.functional.huber_loss(y_hat, y, delta=5.0) + assert model._base_loss(y_hat, y).item() == pytest.approx(expected.item()) + + def test_huber_delta_is_configurable(self): + model = _model(loss="huber", huber_delta=1.0) + y_hat = torch.tensor([510.0]) + y = torch.tensor([500.0]) + + expected = nn.functional.huber_loss(y_hat, y, delta=1.0) + assert model._base_loss(y_hat, y).item() == pytest.approx(expected.item()) + + def test_huber_gradient_is_not_starved(self): + """Unlike log_l1/mape, huber has no division, so no scale pathology to guard against.""" + y_hat = torch.tensor([510.0, 300.0, 700.0], requires_grad=True) + y = torch.tensor([500.0, 320.0, 650.0]) + + model = _model(loss="l1") + model._base_loss(y_hat, y).backward() + l1_grad_norm = y_hat.grad.norm().item() + + y_hat.grad = None + model_huber = _model(loss="huber") + model_huber._base_loss(y_hat, y).backward() + huber_grad_norm = y_hat.grad.norm().item() + + assert huber_grad_norm > l1_grad_norm / 10 + + def test_unknown_loss_raises(self): + model = _model(loss="not_a_real_loss") + + with pytest.raises(ValueError, match="not_a_real_loss"): + model._base_loss(torch.tensor([1.0]), torch.tensor([1.0])) + + def test_validation_mae_stays_raw_l1_regardless_of_loss_choice(self): + """The monitored metric must be comparable across configs with different `loss` values.""" + batch = (*_dummy_batch(batch_size=4), torch.rand(4) * 500 + 300) + + for loss_name in ("l1", "mape", "huber"): + model = _model(loss=loss_name) + atom, diatom, global_feats, one_hot, y = batch + y_hat = model(atom, diatom, global_feats, one_hot).squeeze(1) + assert model.mae(y_hat, y).item() == pytest.approx(nn.L1Loss()(y_hat, y).item()) + + +class TestChargeEmbedding: + """Step 8.4 item 1: charge as an embedding that FiLM-conditions the sequence branches.""" + + def test_disabled_by_default_no_modules_created(self): + model = _model() + + assert model.use_charge_embedding is False + assert not hasattr(model, "charge_embedding") + assert not hasattr(model, "charge_film") + + def test_enabled_creates_correctly_shaped_modules(self): + model = _model(use_charge_embedding=True, charge_embedding_dim=8) + + assert model.charge_embedding.weight.shape == (7, 8) + # film maps the embedding to (gamma, beta) for the full seq-concat width + from im2deep._architectures.helpers import calculate_seq_concat_shape + + seq_concat_shape = calculate_seq_concat_shape(model.config) + assert model.charge_film.out_features == 2 * seq_concat_shape + + def test_requires_add_ccs_features(self): + with pytest.raises(ValueError, match="add_ccs_features"): + _model(use_charge_embedding=True, add_ccs_features=False) + + def test_zero_init_is_an_exact_identity(self): + """FiLM starts as gamma=0/beta=0, so enabling it changes nothing before training.""" + batch = _dummy_batch(charges=[2, 4]) + + model_off = _model() + model_on = _model(use_charge_embedding=True) + # Same everything except the new (zero-initialised) charge modules, + # so the rest of the network must share weights for this to be a + # fair comparison - copy the shared submodules' state across. + model_on.load_state_dict(model_off.state_dict(), strict=False) + + out_off = model_off(*batch) + out_on = model_on(*batch) + + assert torch.allclose(out_off, out_on) + + def test_trained_film_changes_output(self): + """Once charge_film moves away from zero-init, it measurably affects the output.""" + batch = _dummy_batch(charges=[2, 4]) + model = _model(use_charge_embedding=True) + before = model(*batch) + + with torch.no_grad(): + model.charge_film.weight.add_(0.1) + model.charge_film.bias.add_(0.1) + after = model(*batch) + + assert not torch.allclose(before, after) + + def test_different_charges_reach_different_embeddings(self): + """Two otherwise-identical inputs differing only in charge get different FiLM parameters.""" + model = _model(use_charge_embedding=True) + with torch.no_grad(): + model.charge_film.weight.add_(0.1) + model.charge_film.bias.add_(0.1) + + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=1, charges=[2]) + out_z2 = model(atom, diatom, global_feats.clone(), one_hot) + + global_feats_z5 = global_feats.clone() + global_feats_z5[:, CHARGE_FEATURE_INDEX] = 5.0 + out_z5 = model(atom, diatom, global_feats_z5, one_hot) + + assert not torch.allclose(out_z2, out_z5) + + def test_charge_index_matches_deeplc_output(self): + """CHARGE_FEATURE_INDEX must match where DeepLC actually puts charge, not just its docstring.""" + from deeplc._features import encode_peptidoform + from psm_utils.peptidoform import Peptidoform + + for charge in (1, 3, 6): + pf = Peptidoform(f"ACDEFGHIK/{charge}") + feats = encode_peptidoform( + pf, padding_length=60, add_ccs_features=True, add_terminal_composition=False + ) + assert feats["matrix_global"][CHARGE_FEATURE_INDEX] == charge + + +class TestOneHotAblation: + """Step 8.1/8.4: use_one_hot=False removes the one-hot branch entirely.""" + + def test_enabled_by_default(self): + model = _model() + + assert model.use_one_hot is True + assert hasattr(model, "OneHot") + + def test_disabled_removes_the_module(self): + model = _model(use_one_hot=False) + + assert model.use_one_hot is False + assert not hasattr(model, "OneHot") + + def test_disabled_model_runs_and_ignores_one_hot_input(self): + """forward() must not touch `one_hot` at all when the branch is off.""" + model = _model(use_one_hot=False) + atom, diatom, global_feats, one_hot = _dummy_batch() + + out = model(atom, diatom, global_feats, one_hot) + + assert out.shape == (2, 1) + assert torch.all(torch.isfinite(out)) + + def test_concat_shape_shrinks_when_disabled(self): + from im2deep._architectures.helpers import calculate_concat_shape + + config_with = {**copy.deepcopy(DEFAULT_TRAINING_CONFIG), "Global_features": 60} + config_without = {**config_with, "use_one_hot": False} + + shape_with = calculate_concat_shape(config_with) + shape_without = calculate_concat_shape(config_without) + + onehot_out_shape = (60 // config_with["OneHot_MaxPool_kernel_size"]) * config_with[ + "OneHot_out_channels" + ] + assert shape_with - shape_without == onehot_out_shape + + def test_head_input_size_matches_actual_concatenation(self): + """The Concat head's first layer must accept exactly what forward() feeds it.""" + model = _model(use_one_hot=False) + atom, diatom, global_feats, one_hot = _dummy_batch() + + # Would raise a shape-mismatch RuntimeError if total_input_size (used + # to size Concat[0]) didn't shrink in lockstep with the ablation. + model(atom, diatom, global_feats, one_hot) + + +class TestAttentionPooling: + """Step 8.4 item 2: AttentionPool1d replaces Flatten(), length-agnostic output width.""" + + def test_disabled_by_default_uses_flatten(self): + model = _model() + + assert model.use_attention_pooling is False + assert isinstance(model.ConvAtomComp[-1], nn.Flatten) + assert isinstance(model.ConvDiatomComp[-1], nn.Flatten) + + def test_enabled_uses_attention_pool(self): + from im2deep._architectures.blocks import AttentionPool1d + + model = _model(use_attention_pooling=True) + + assert isinstance(model.ConvAtomComp[-1], AttentionPool1d) + assert isinstance(model.ConvDiatomComp[-1], AttentionPool1d) + + def test_runs_and_produces_finite_output(self): + model = _model(use_attention_pooling=True) + atom, diatom, global_feats, one_hot = _dummy_batch() + + out = model(atom, diatom, global_feats, one_hot) + + assert out.shape == (2, 1) + assert torch.all(torch.isfinite(out)) + + def test_output_width_is_length_agnostic(self): + """Unlike Flatten, the branch's output width shouldn't depend on padding_length.""" + from im2deep._architectures.helpers import calculate_seq_concat_shape + + config = {**copy.deepcopy(DEFAULT_TRAINING_CONFIG), "Global_features": 60} + config_pooled = {**config, "use_attention_pooling": True} + + # Flattened width scales with sequence length (60 here); pooled width + # is just the channel count, independent of it. + flattened_width = calculate_seq_concat_shape(config) + pooled_width = calculate_seq_concat_shape(config_pooled) + expected_pooled = ( + config["AtomComp_out_channels_start"] // 4 + config["DiatomComp_out_channels_start"] // 2 + ) + assert pooled_width == expected_pooled + assert pooled_width < flattened_width + + def test_composes_with_charge_film(self): + """FiLM's expected width must track attention pooling's smaller seq-concat shape.""" + model = _model(use_attention_pooling=True, use_charge_embedding=True) + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + out = model(atom, diatom, global_feats, one_hot) + + assert torch.all(torch.isfinite(out)) + + def test_attention_weights_sum_to_one_over_positions(self): + from im2deep._architectures.blocks import AttentionPool1d + + pool = AttentionPool1d(channels=4, initializer=lambda w: nn.init.normal_(w, std=0.05)) + x = torch.randn(2, 4, 15) + + weights = torch.softmax(pool.score(x), dim=-1) + + assert torch.allclose(weights.sum(dim=-1), torch.ones(2, 1), atol=1e-5) + + +class TestFlexTrunkBlock: + """Step 8.4 items 8+9: FlexTrunk itself, ported from DeepLC's _FlexCNNEncoder.""" + + def _trunk(self, **overrides): + from im2deep._architectures.blocks import FlexTrunk + + kwargs = {"embed_dim": 4, "channels": (8, 8), "kernel_size": 3, "stem_channels": 6, "stem_layers": 2} + kwargs.update(overrides) + return FlexTrunk(**kwargs) + + def _one_hot(self, batch_size, length, real_length): + """(batch, 20, length) one-hot, real for the first real_length positions, zero after.""" + one_hot = torch.zeros(batch_size, 20, length) + for pos in range(real_length): + one_hot[:, pos % 20, pos] = 1.0 + return one_hot + + def test_residue_indices_matches_argmax_on_real_positions(self): + trunk = self._trunk() + one_hot = self._one_hot(batch_size=2, length=10, real_length=10) + + idx = trunk.residue_indices(one_hot) + + assert torch.equal(idx, torch.arange(10).unsqueeze(0).expand(2, 10)) + + def test_residue_indices_flags_all_zero_columns_as_padding(self): + trunk = self._trunk() + one_hot = self._one_hot(batch_size=2, length=10, real_length=6) + + idx = trunk.residue_indices(one_hot) + + assert torch.all(idx[:, 6:] == trunk.PAD_INDEX) + assert torch.all(idx[:, :6] != trunk.PAD_INDEX) + + def test_forward_runs_and_produces_finite_output(self): + trunk = self._trunk() + atom = torch.randn(3, 6, 12) + one_hot = self._one_hot(batch_size=3, length=12, real_length=8) + + out = trunk(atom, one_hot) + + assert out.shape == (3, 2 * 8) # sum + max, each channels[-1] wide + assert torch.all(torch.isfinite(out)) + + def test_padding_positions_do_not_change_the_pooled_output(self): + """Masking means only the real positions should affect sum/max pooling.""" + torch.manual_seed(0) + trunk = self._trunk() + real_length = 7 + length_a, length_b = 10, 20 + + atom_a = torch.randn(1, 6, length_a) + atom_b = torch.zeros(1, 6, length_b) + atom_b[:, :, :length_a] = atom_a + one_hot_a = self._one_hot(batch_size=1, length=length_a, real_length=real_length) + one_hot_b = self._one_hot(batch_size=1, length=length_b, real_length=real_length) + + out_a = trunk(atom_a, one_hot_a) + out_b = trunk(atom_b, one_hot_b) + + assert torch.allclose(out_a, out_b, atol=1e-5) + + def test_no_stem_feeds_raw_atom_counts(self): + trunk = self._trunk(stem_channels=0) + + assert trunk.stem is None + atom = torch.randn(2, 6, 10) + one_hot = self._one_hot(batch_size=2, length=10, real_length=10) + out = trunk(atom, one_hot) + assert torch.all(torch.isfinite(out)) + + +class TestFlexTrunkModel: + """Step 8.4 items 8+9: FlexTrunk wired into IM2Deep, replacing three branches.""" + + def test_disabled_by_default(self): + model = _model() + + assert model.use_flex_trunk is False + assert hasattr(model, "ConvAtomComp") + assert hasattr(model, "ConvDiatomComp") + assert not hasattr(model, "flex_trunk") + + def test_enabled_retires_the_three_branches(self): + model = _model(use_flex_trunk=True) + + assert hasattr(model, "flex_trunk") + assert not hasattr(model, "ConvAtomComp") + assert not hasattr(model, "ConvDiatomComp") + assert not hasattr(model, "OneHot") + + def test_use_one_hot_is_silently_moot_when_enabled(self): + # use_one_hot defaults to True; enabling flex_trunk must not raise + # and must not build the (now redundant) OneHot branch anyway. + model = _model(use_flex_trunk=True, use_one_hot=True) + + assert model.use_one_hot is False + assert not hasattr(model, "OneHot") + + def test_runs_and_produces_finite_output(self): + model = _model(use_flex_trunk=True) + atom, diatom, global_feats, one_hot = _dummy_batch() + + out = model(atom, diatom, global_feats, one_hot) + + assert out.shape == (2, 1) + assert torch.all(torch.isfinite(out)) + + def test_concat_shape_matches_actual_seq_feats_width(self): + from im2deep._architectures.helpers import calculate_seq_concat_shape + + config = { + **copy.deepcopy(DEFAULT_TRAINING_CONFIG), + "Global_features": 60, + "use_flex_trunk": True, + "flex_trunk_channels": [12, 10], + } + model = IM2Deep(config, nn.L1Loss()) + atom, diatom, global_feats, one_hot = _dummy_batch() + + seq_feats = model.flex_trunk( + atom.permute(0, 2, 1), one_hot.permute(0, 2, 1) + ) + + assert seq_feats.shape[1] == calculate_seq_concat_shape(config) + assert calculate_seq_concat_shape(config) == 2 * 10 + + def test_uses_far_fewer_parameters_than_the_four_branch_default(self): + default_model = _model() + flex_model = _model(use_flex_trunk=True) + + default_params = sum(p.numel() for p in default_model.parameters()) + flex_params = sum(p.numel() for p in flex_model.parameters()) + + assert flex_params < default_params + + def test_composes_with_charge_film(self): + model = _model(use_flex_trunk=True, use_charge_embedding=True) + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + out = model(atom, diatom, global_feats, one_hot) + + assert torch.all(torch.isfinite(out)) + + def test_diatom_comp_is_ignored(self): + """matrix_sum is redundant once the trunk sees matrix directly (item 8).""" + model = _model(use_flex_trunk=True) + atom, diatom, global_feats, one_hot = _dummy_batch() + different_diatom = torch.randn_like(diatom) + + out_a = model(atom, diatom, global_feats, one_hot) + out_b = model(atom, different_diatom, global_feats, one_hot) + + assert torch.allclose(out_a, out_b) + + +class TestTransformerTrunkBlock: + """Step 8.4 item 3: RelativePairBias, PerceiverReadout, TransformerTrunk.""" + + def _pair_bias(self, **overrides): + from im2deep._architectures.transformer_trunk import RelativePairBias + + kwargs = {"num_heads": 2, "max_relative_distance": 5, "max_terminus_distance": 3} + kwargs.update(overrides) + return RelativePairBias(**kwargs) + + def _trunk(self, **overrides): + from im2deep._architectures.transformer_trunk import TransformerTrunk + + kwargs = { + "d_model": 8, + "num_layers": 2, + "num_heads": 2, + "ffn_dim": 16, + "max_relative_distance": 5, + "max_terminus_distance": 3, + } + kwargs.update(overrides) + return TransformerTrunk(**kwargs) + + def _atom_comp(self, batch_size, length, real_length): + """(batch, length, 6): nonzero for the first real_length positions, zero after.""" + atom = torch.zeros(batch_size, length, 6) + atom[:, :real_length, :] = torch.randn(batch_size, real_length, 6) + return atom + + def test_pair_bias_output_shape(self): + pair_bias = self._pair_bias(num_heads=3) + + out = pair_bias(real_lengths=torch.tensor([10, 5]), seq_len=10) + + assert out.shape == (2, 3, 10, 10) + + def test_relative_component_is_asymmetric_by_design(self): + """i-j and j-i map to different buckets - direction should matter.""" + pair_bias = self._pair_bias() + nn.init.zeros_(pair_bias.terminus.weight) + + out = pair_bias(real_lengths=torch.tensor([10]), seq_len=10) + + assert not torch.allclose(out, out.transpose(-1, -2)) + + def test_relative_component_diagonal_is_a_fixed_bucket(self): + """i==j always means distance 0, regardless of where i is.""" + pair_bias = self._pair_bias() + nn.init.zeros_(pair_bias.terminus.weight) + + out = pair_bias(real_lengths=torch.tensor([10]), seq_len=10) + diagonal = torch.diagonal(out[0, 0]) + + assert torch.allclose(diagonal, diagonal[0] * torch.ones_like(diagonal)) + + def test_relative_component_clamps_beyond_max_distance(self): + pair_bias = self._pair_bias(max_relative_distance=3) + nn.init.zeros_(pair_bias.terminus.weight) + + out = pair_bias(real_lengths=torch.tensor([10]), seq_len=10) + + # i-j=3 and i-j=5 both clamp to the same (boundary) bucket. + assert torch.allclose(out[0, :, 5, 2], out[0, :, 7, 2]) + + def test_terminus_component_is_symmetric(self): + pair_bias = self._pair_bias() + nn.init.zeros_(pair_bias.relative.weight) + + out = pair_bias(real_lengths=torch.tensor([10]), seq_len=10) + + assert torch.allclose(out, out.transpose(-1, -2)) + + def test_terminus_distance_formula_matches_hand_computation(self): + pair_bias = self._pair_bias(num_heads=1, max_terminus_distance=10) + nn.init.zeros_(pair_bias.relative.weight) + # Craft the terminus table so its output value equals the bucket index. + with torch.no_grad(): + pair_bias.terminus.weight.copy_(torch.arange(11).float().unsqueeze(1)) + + out = pair_bias(real_lengths=torch.tensor([6]), seq_len=6) + + # length 6: closest-to-either-terminus distance per position is + # [0, 1, 2, 2, 1, 0]; terminus_distance(i,j) = min(closest[i], closest[j]). + assert out[0, 0, 0, 0].item() == 0 + assert out[0, 0, 2, 3].item() == 2 + assert out[0, 0, 0, 5].item() == 0 + + def test_readout_output_shape_single_query(self): + from im2deep._architectures.transformer_trunk import PerceiverReadout + + readout = PerceiverReadout(d_model=8, num_heads=2, num_queries=1, dropout=0.0) + residue_hidden = torch.randn(2, 10, 8) + attn_mask = torch.zeros(2 * 2, 1, 10) + + out = readout(residue_hidden, attn_mask) + + assert out.shape == (2, 8) + + def test_readout_output_shape_multiple_queries(self): + from im2deep._architectures.transformer_trunk import PerceiverReadout + + readout = PerceiverReadout(d_model=8, num_heads=2, num_queries=3, dropout=0.0) + residue_hidden = torch.randn(2, 10, 8) + attn_mask = torch.zeros(2 * 2, 3, 10) + + out = readout(residue_hidden, attn_mask) + + assert out.shape == (2, 3, 8) + + def test_trunk_runs_and_produces_finite_output(self): + trunk = self._trunk() + atom = self._atom_comp(batch_size=3, length=12, real_length=8) + charge_idx = torch.tensor([1, 2, 4]) + + out = trunk(atom, charge_idx) + + assert out.shape == (3, 8) + assert torch.all(torch.isfinite(out)) + + def test_padding_positions_do_not_change_the_pooled_output(self): + """Same real peptide, different amounts of trailing padding -> identical output.""" + torch.manual_seed(0) + trunk = self._trunk() + real_length = 7 + length_a, length_b = 10, 20 + + atom_a = torch.randn(1, length_a, 6) + atom_a[:, real_length:, :] = 0.0 + atom_b = torch.zeros(1, length_b, 6) + atom_b[:, :length_a, :] = atom_a + charge_idx = torch.tensor([2]) + + out_a = trunk(atom_a, charge_idx) + out_b = trunk(atom_b, charge_idx) + + assert torch.allclose(out_a, out_b, atol=1e-5) + + def test_register_token_excluded_from_readout_keys(self): + # The readout's key/value go straight to _sdpa_attention rather than + # through cross_attn.forward() (see _sdpa_attention's docstring for + # why), so a forward hook on cross_attn no longer observes the call - + # patch _sdpa_attention itself instead. PerceiverReadout.forward + # looks it up as a module global at call time, so this redirect + # applies regardless of which module imported PerceiverReadout. + import im2deep._architectures.transformer_trunk as transformer_trunk_module + + trunk = self._trunk() + atom = self._atom_comp(batch_size=2, length=10, real_length=10) + charge_idx = torch.tensor([1, 3]) + captured = {} + original = transformer_trunk_module._sdpa_attention + + def spy(mha, query, key, value, attn_mask): + captured["key_len"] = key.shape[1] + return original(mha, query, key, value, attn_mask) + + transformer_trunk_module._sdpa_attention = spy + try: + trunk(atom, charge_idx) + finally: + transformer_trunk_module._sdpa_attention = original + + assert captured["key_len"] == 10 # residue positions only, not 10 + 1 register token + + def test_charge_sensitivity(self): + """Different charge, same residues -> different output ('charge attends back').""" + torch.manual_seed(0) + trunk = self._trunk() + atom = self._atom_comp(batch_size=1, length=10, real_length=10) + + out_low = trunk(atom, torch.tensor([1])) + out_high = trunk(atom, torch.tensor([5])) + + assert not torch.allclose(out_low, out_high) + + # -- step 8.10: return_attention is opt-in and additive -- + + def test_block_return_attention_shape_and_same_output(self): + """need_weights=True must not change the block's own output.""" + from im2deep._architectures.transformer_trunk import PreNormTransformerBlock + + torch.manual_seed(0) + block = PreNormTransformerBlock(d_model=8, num_heads=2, ffn_dim=16, dropout=0.0) + block.eval() + x = torch.randn(2, 5, 8) + attn_mask = torch.zeros(2 * 2, 5, 5) + + out_plain = block(x, attn_mask) + out_with_attn, attn_weights = block(x, attn_mask, return_attention=True) + + assert torch.allclose(out_plain, out_with_attn, atol=1e-4) + assert attn_weights.shape == (2, 2, 5, 5) + + def test_readout_return_attention_shape_and_same_output(self): + from im2deep._architectures.transformer_trunk import PerceiverReadout + + torch.manual_seed(0) + readout = PerceiverReadout(d_model=8, num_heads=2, num_queries=1, dropout=0.0) + readout.eval() + residue_hidden = torch.randn(2, 10, 8) + attn_mask = torch.zeros(2 * 2, 1, 10) + + out_plain = readout(residue_hidden, attn_mask) + out_with_attn, attn_weights = readout(residue_hidden, attn_mask, return_attention=True) + + assert torch.allclose(out_plain, out_with_attn, atol=1e-4) + assert attn_weights.shape == (2, 2, 1, 10) + + def test_trunk_return_attention_shapes_and_same_pooled_output(self): + torch.manual_seed(0) + trunk = self._trunk() # d_model=8, num_layers=2, num_heads=2 + trunk.eval() + atom = self._atom_comp(batch_size=2, length=10, real_length=7) + charge_idx = torch.tensor([1, 3]) + + out_plain = trunk(atom, charge_idx) + out_with_attn, attention = trunk(atom, charge_idx, return_attention=True) + + assert torch.allclose(out_plain, out_with_attn, atol=1e-4) + assert set(attention.keys()) == {"self_attention", "readout_attention"} + assert len(attention["self_attention"]) == 2 # num_layers + # seq_len=10 (the padded array length, not real_length) + register token. + assert attention["self_attention"][0].shape == (2, 2, 11, 11) + assert attention["readout_attention"].shape == (2, 2, 1, 10) + + +class TestTransformerTrunkModel: + """Step 8.4 item 3: TransformerTrunk wired into IM2Deep.""" + + def test_disabled_by_default(self): + model = _model() + + assert model.use_transformer_trunk is False + assert hasattr(model, "ConvAtomComp") + assert not hasattr(model, "transformer_trunk") + + def test_enabled_retires_conv_branches_but_keeps_one_hot(self): + """Unlike use_flex_trunk, this arm does not fold in residue identity.""" + model = _model(use_transformer_trunk=True) + + assert hasattr(model, "transformer_trunk") + assert not hasattr(model, "ConvAtomComp") + assert not hasattr(model, "ConvDiatomComp") + assert hasattr(model, "OneHot") + assert model.use_one_hot is True + + def test_mutual_exclusion_with_flex_trunk_raises(self): + with pytest.raises(ValueError, match="mutually exclusive"): + _model(use_transformer_trunk=True, use_flex_trunk=True) + + def test_requires_add_ccs_features(self): + with pytest.raises(ValueError, match="add_ccs_features"): + _model(use_transformer_trunk=True, add_ccs_features=False) + + def test_runs_and_produces_finite_output(self): + model = _model(use_transformer_trunk=True) + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + out = model(atom, diatom, global_feats, one_hot) + + assert out.shape == (2, 1) + assert torch.all(torch.isfinite(out)) + + def test_composes_with_charge_embedding(self): + model = _model(use_transformer_trunk=True, use_charge_embedding=True) + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + out = model(atom, diatom, global_feats, one_hot) + + assert torch.all(torch.isfinite(out)) + + def test_use_attention_pooling_is_silently_moot(self): + model = _model(use_transformer_trunk=True, use_attention_pooling=True) + + assert not hasattr(model, "ConvAtomComp") + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + out = model(atom, diatom, global_feats, one_hot) + assert torch.all(torch.isfinite(out)) + + def test_concat_shape_matches_actual_seq_feats_width(self): + from im2deep._architectures.helpers import calculate_seq_concat_shape + + config = { + **copy.deepcopy(DEFAULT_TRAINING_CONFIG), + "Global_features": 60, + "use_transformer_trunk": True, + "transformer_d_model": 12, + } + model = IM2Deep(config, nn.L1Loss()) + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + atom_raw = atom + charge_idx = global_feats[:, CHARGE_FEATURE_INDEX].round().long().clamp(0, 6) + seq_feats = model.transformer_trunk(atom_raw, charge_idx) + + assert seq_feats.shape[1] == calculate_seq_concat_shape(config) + assert calculate_seq_concat_shape(config) == 12 + + def test_return_attention_requires_transformer_trunk(self): + model = _model() # use_transformer_trunk=False by default + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + with pytest.raises(ValueError, match="use_transformer_trunk"): + model(atom, diatom, global_feats, one_hot, return_attention=True) + + def test_return_attention_matches_default_output(self): + """return_attention=True must be additive: same prediction either way.""" + torch.manual_seed(0) + model = _model(use_transformer_trunk=True) + model.eval() + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + out_plain = model(atom, diatom, global_feats, one_hot) + out_with_attn, attention = model( + atom, diatom, global_feats, one_hot, return_attention=True + ) + + assert torch.allclose(out_plain, out_with_attn, atol=1e-4) + assert set(attention.keys()) == {"self_attention", "readout_attention"} + + +class TestMixup: + """Step 8.3: composition mixup, applied inside training_step only.""" + + def _batch(self, batch_size=4): + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=batch_size) + y = torch.rand(batch_size) * 500 + 300 + return atom, diatom, global_feats, one_hot, y + + def test_disabled_by_default_is_a_no_op(self): + model = _model() + batch = self._batch() + + result = model._apply_mixup(*batch) + + # batch is 5-tuple (no mol_desc); result is 6-tuple (mol_desc=None + # appended) - intentionally comparing only the first 5. + for original, returned in zip(batch, result, strict=False): + assert torch.equal(original, returned) + + def test_alpha_zero_is_a_no_op(self): + model = _model(mixup_alpha=0.0) + batch = self._batch() + + result = model._apply_mixup(*batch) + + # batch is 5-tuple (no mol_desc); result is 6-tuple (mol_desc=None + # appended) - intentionally comparing only the first 5. + for original, returned in zip(batch, result, strict=False): + assert torch.equal(original, returned) + + def test_enabled_mixes_every_tensor_with_the_same_permutation(self): + model = _model(mixup_alpha=1.0) + atom, diatom, global_feats, one_hot, y = self._batch() + + mixed_atom, mixed_diatom, mixed_global, mixed_one_hot, mixed_y, mixed_mol = model._apply_mixup( + atom, diatom, global_feats, one_hot, y + ) + + assert mixed_mol is None + # lambda ~ Beta(1,1) = Uniform(0,1) essentially never lands exactly on + # 0 or 1, so every mixed tensor should differ from the original. + assert not torch.equal(mixed_atom, atom) + assert not torch.equal(mixed_y, y) + # Every value must remain within the convex hull of the two inputs it + # came from (a basic sanity check that this is a real interpolation). + assert (mixed_y.min() >= min(y.min(), y.max()) - 1e-4) and ( + mixed_y.max() <= max(y.min(), y.max()) + 1e-4 + ) + + def test_mol_desc_is_mixed_when_present(self): + # _apply_mixup only branches on whether mol_desc is passed, not on + # config["add_X_mol"], so a plain model is enough here. + model = _model(mixup_alpha=1.0) + atom, diatom, global_feats, one_hot, y = self._batch() + mol_desc = torch.randn(4, 13, 60) + + *_, mixed_mol = model._apply_mixup(atom, diatom, global_feats, one_hot, y, mol_desc) + + assert mixed_mol is not None + assert not torch.equal(mixed_mol, mol_desc) + + +class TestSiblingDeltaLoss: + """Step 8.3: sibling delta loss, L1 on (pred_a - pred_b) vs (obs_a - obs_b).""" + + def test_disabled_by_default_is_zero(self): + model = _model() + y_hat = torch.tensor([400.0, 410.0, 420.0]) + y = torch.tensor([405.0, 408.0, 425.0]) + group_id = torch.tensor([0, 0, 1]) + + loss = model._sibling_delta_loss(y_hat, y, group_id) + + assert loss.item() == 0.0 + + def test_zero_weight_is_zero(self): + model = _model(sibling_loss_weight=0.0) + y_hat = torch.tensor([400.0, 410.0]) + y = torch.tensor([405.0, 408.0]) + group_id = torch.tensor([0, 0]) + + assert model._sibling_delta_loss(y_hat, y, group_id).item() == 0.0 + + def test_no_matching_groups_is_zero(self): + model = _model(sibling_loss_weight=1.0) + y_hat = torch.tensor([400.0, 410.0, 420.0]) + y = torch.tensor([405.0, 408.0, 425.0]) + group_id = torch.tensor([0, 1, 2]) # every row its own group + + assert model._sibling_delta_loss(y_hat, y, group_id).item() == 0.0 + + def test_matches_manual_l1_on_the_one_sibling_pair(self): + model = _model(sibling_loss_weight=1.0) + y_hat = torch.tensor([400.0, 410.0, 500.0]) + y = torch.tensor([405.0, 408.0, 550.0]) + group_id = torch.tensor([0, 0, 1]) # only rows 0/1 are siblings + + loss = model._sibling_delta_loss(y_hat, y, group_id) + + pred_delta = y_hat[0] - y_hat[1] + obs_delta = y[0] - y[1] + expected = torch.nn.functional.l1_loss(pred_delta, obs_delta) + assert torch.allclose(loss, expected) + + def test_weight_scales_the_loss(self): + model_a = _model(sibling_loss_weight=1.0) + model_b = _model(sibling_loss_weight=3.0) + y_hat = torch.tensor([400.0, 410.0]) + y = torch.tensor([405.0, 408.0]) + group_id = torch.tensor([0, 0]) + + loss_a = model_a._sibling_delta_loss(y_hat, y, group_id) + loss_b = model_b._sibling_delta_loss(y_hat, y, group_id) + + assert torch.allclose(loss_b, 3.0 * loss_a) + + def test_ignores_config_loss_choice(self): + # The delta term is documented as always plain L1, independent of + # config["loss"] - verify a non-l1 choice does not change it. + model = _model(sibling_loss_weight=1.0, loss="huber") + y_hat = torch.tensor([400.0, 410.0]) + y = torch.tensor([405.0, 408.0]) + group_id = torch.tensor([0, 0]) + + loss = model._sibling_delta_loss(y_hat, y, group_id) + expected = torch.nn.functional.l1_loss(y_hat[0] - y_hat[1], y[0] - y[1]) + assert torch.allclose(loss, expected) + + def test_groups_larger_than_two_cover_every_pair(self): + model = _model(sibling_loss_weight=1.0) + y_hat = torch.tensor([100.0, 110.0, 120.0]) + y = torch.tensor([100.0, 130.0, 160.0]) + group_id = torch.tensor([0, 0, 0]) # one group of three -> 3 pairs + + loss = model._sibling_delta_loss(y_hat, y, group_id) + + pairs = [(0, 1), (0, 2), (1, 2)] + pred_deltas = torch.stack([y_hat[a] - y_hat[b] for a, b in pairs]) + obs_deltas = torch.stack([y[a] - y[b] for a, b in pairs]) + expected = torch.nn.functional.l1_loss(pred_deltas, obs_deltas) + assert torch.allclose(loss, expected) + + +class TestRankingLoss: + """Step 8.3: pairwise ranking loss within an m/z window.""" + + def test_disabled_by_default_is_zero(self): + model = _model() + y_hat = torch.tensor([400.0, 410.0]) + y = torch.tensor([405.0, 408.0]) + mz = torch.tensor([500.0, 501.0]) + + assert model._ranking_loss(y_hat, y, mz).item() == 0.0 + + def test_pairs_outside_window_are_ignored(self): + model = _model(ranking_loss_weight=1.0, ranking_mz_window=3.0) + y_hat = torch.tensor([400.0, 999.0]) # badly mis-ordered if paired + y = torch.tensor([405.0, 408.0]) + mz = torch.tensor([500.0, 600.0]) # 100 Th apart, outside the window + + assert model._ranking_loss(y_hat, y, mz).item() == 0.0 + + def test_correctly_ordered_pair_within_margin_is_zero(self): + model = _model(ranking_loss_weight=1.0, ranking_mz_window=3.0, ranking_margin=0.0) + y_hat = torch.tensor([400.0, 410.0]) # same order as y + y = torch.tensor([405.0, 408.0]) + mz = torch.tensor([500.0, 501.0]) + + assert model._ranking_loss(y_hat, y, mz).item() == pytest.approx(0.0, abs=1e-6) + + def test_misordered_pair_within_window_is_penalized(self): + model = _model(ranking_loss_weight=1.0, ranking_mz_window=3.0, ranking_margin=1.0) + # y says row 0 < row 1, y_hat says the opposite. + y_hat = torch.tensor([410.0, 400.0]) + y = torch.tensor([405.0, 408.0]) + mz = torch.tensor([500.0, 501.0]) + + loss = model._ranking_loss(y_hat, y, mz) + + assert loss.item() > 0.0 + + def test_tied_observed_ccs_pair_is_ignored(self): + model = _model(ranking_loss_weight=1.0, ranking_mz_window=3.0) + y_hat = torch.tensor([400.0, 999.0]) + y = torch.tensor([405.0, 405.0]) # exact tie -> no ordering to supervise + mz = torch.tensor([500.0, 501.0]) + + assert model._ranking_loss(y_hat, y, mz).item() == 0.0 + + def test_weight_scales_the_loss(self): + model_a = _model(ranking_loss_weight=1.0, ranking_mz_window=3.0, ranking_margin=1.0) + model_b = _model(ranking_loss_weight=4.0, ranking_mz_window=3.0, ranking_margin=1.0) + y_hat = torch.tensor([410.0, 400.0]) + y = torch.tensor([405.0, 408.0]) + mz = torch.tensor([500.0, 501.0]) + + loss_a = model_a._ranking_loss(y_hat, y, mz) + loss_b = model_b._ranking_loss(y_hat, y, mz) + + assert torch.allclose(loss_b, 4.0 * loss_a) + + +class TestPairwiseLossMixupIncompatibility: + """Step 8.3: mixup blends rows, corrupting group_id/mz - must raise, not silently combine.""" + + def test_sibling_loss_with_mixup_raises(self): + with pytest.raises(ValueError, match="mixup_alpha"): + _model(mixup_alpha=0.5, sibling_loss_weight=1.0) + + def test_ranking_loss_with_mixup_raises(self): + with pytest.raises(ValueError, match="mixup_alpha"): + _model(mixup_alpha=0.5, ranking_loss_weight=1.0) + + def test_either_alone_with_mixup_is_fine(self): + _model(mixup_alpha=0.5) + _model(sibling_loss_weight=1.0) + _model(ranking_loss_weight=1.0) + + +class TestHeteroscedasticHead: + """Step 8.5: a (mu, b) head trained with Laplace NLL instead of a single output.""" + + def test_disabled_by_default(self): + model = _model() + + assert model.use_heteroscedastic_head is False + assert model.Concat[-1].out_features == 1 + + def test_enabled_changes_final_layer_to_two_outputs(self): + model = _model(use_heteroscedastic_head=True) + + assert model.Concat[-1].out_features == 2 + + def test_disabled_forward_returns_a_single_tensor(self): + model = _model() + batch = _dummy_batch() + + out = model(*batch) + + assert isinstance(out, torch.Tensor) + assert out.shape == (2, 1) + + def test_enabled_forward_returns_mu_b_tuple(self): + model = _model(use_heteroscedastic_head=True) + batch = _dummy_batch() + + out = model(*batch) + + assert isinstance(out, tuple) and len(out) == 2 + mu, b = out + assert mu.shape == (2,) + assert b.shape == (2,) + + def test_b_is_always_positive_and_finite(self): + model = _model(use_heteroscedastic_head=True) + batch = _dummy_batch(batch_size=8) + + _, b = model(*batch) + + assert torch.all(b > 0) + assert torch.all(torch.isfinite(b)) + + def test_b_scales_with_target_std_not_shifted_by_target_mean(self): + """b is a scale parameter: only the multiplicative half of destandardisation applies.""" + model = _model(use_heteroscedastic_head=True) + batch = _dummy_batch() + + model.set_target_scaler(0.0, 1.0) + _, b_unit = model(*batch) + model.set_target_scaler(500.0, 200.0) # same std ratio check, now with a nonzero mean + _, b_scaled = model(*batch) + + assert torch.allclose(b_scaled, b_unit * 200.0, rtol=1e-4) + + def test_mu_destandardizes_the_same_way_as_the_single_output_head(self): + model = _model(use_heteroscedastic_head=True) + batch = _dummy_batch() + + model.set_target_scaler(0.0, 1.0) + mu_raw, _ = model(*batch) + model.set_target_scaler(500.0, 200.0) + mu_scaled, _ = model(*batch) + + assert torch.allclose(mu_scaled, mu_raw * 200.0 + 500.0, atol=1e-4) + + def test_laplace_nll_matches_hand_computation(self): + model = _model(use_heteroscedastic_head=True) + mu = torch.tensor([490.0, 495.0]) + y = torch.tensor([500.0, 510.0]) + b = torch.tensor([10.0, 15.0]) + + loss = model._laplace_nll(mu, y, b) + + expected = (torch.log(2.0 * b) + torch.abs(y - mu) / b).mean() + assert torch.allclose(loss, expected) + + def test_laplace_nll_is_minimized_when_b_equals_the_error(self): + """The Laplace MLE scale for a fixed error e=|y-mu| is exactly b=e.""" + model = _model(use_heteroscedastic_head=True) + mu = torch.tensor([490.0, 495.0]) + y = torch.tensor([500.0, 510.0]) + error = torch.abs(y - mu) + + nll_at_error = model._laplace_nll(mu, y, error) + nll_too_small = model._laplace_nll(mu, y, error / 10) + nll_too_large = model._laplace_nll(mu, y, error * 10) + + assert nll_at_error < nll_too_small + assert nll_at_error < nll_too_large + + def test_training_step_runs_and_produces_a_finite_loss(self): + model = _model(use_heteroscedastic_head=True) + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=4) + y = torch.rand(4) * 500 + 300 + batch = (atom, diatom, global_feats, one_hot, y) + + loss = model.training_step(batch, 0) + + assert torch.isfinite(loss) + assert loss.requires_grad + + def test_validation_mae_uses_only_mu_not_b(self): + """The established invariant (Validation MAE is always plain L1 on the point + estimate, regardless of config["loss"]) extends to this head too.""" + model = _model(use_heteroscedastic_head=True) + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=4) + y = torch.rand(4) * 500 + 300 + batch = (atom, diatom, global_feats, one_hot, y) + + mu, _ = model(atom, diatom, global_feats, one_hot) + expected_mae = torch.nn.functional.l1_loss(mu, y) + + logged = {} + model.log = lambda name, value, **kwargs: logged.__setitem__(name, value) + model.validation_step(batch, 0) + + assert torch.allclose(logged["Validation MAE"], expected_mae) + + def test_predict_step_returns_only_mu(self): + model = _model(use_heteroscedastic_head=True) + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=4) + y = torch.rand(4) * 500 + 300 + batch = (atom, diatom, global_feats, one_hot, y) + + out = model.predict_step(batch, 0) + + assert isinstance(out, torch.Tensor) + assert out.shape == (4,) + + def test_composes_with_transformer_trunk(self): + model = _model( + use_heteroscedastic_head=True, use_transformer_trunk=True, add_ccs_features=True + ) + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + mu, b = model(atom, diatom, global_feats, one_hot) + + assert torch.all(torch.isfinite(mu)) + assert torch.all(b > 0) + + def test_composes_with_transformer_trunk_and_return_attention(self): + """The heteroscedastic head's (mu, b) return still gets the attention + dict appended as a third element, not swallowed - see forward()'s + two return_attention branches.""" + model = _model( + use_heteroscedastic_head=True, use_transformer_trunk=True, add_ccs_features=True + ) + atom, diatom, global_feats, one_hot = _dummy_batch(charges=[2, 4]) + + mu, b, attention = model(atom, diatom, global_feats, one_hot, return_attention=True) + + assert torch.all(torch.isfinite(mu)) + assert torch.all(b > 0) + assert set(attention.keys()) == {"self_attention", "readout_attention"} + + def test_composes_with_sibling_delta_loss(self): + """y_hat (here, mu) plays the same role for sibling/ranking losses either way.""" + model = _model(use_heteroscedastic_head=True, sibling_loss_weight=1.0) + mu = torch.tensor([400.0, 410.0]) + y = torch.tensor([405.0, 408.0]) + group_id = torch.tensor([0, 0]) + + loss = model._sibling_delta_loss(mu, y, group_id) + + assert torch.isfinite(loss) + + +class _FakeTrainer: + """Minimal stand-in so a bare (Trainer-less) LightningModule's + current_epoch reads a chosen value, without spinning up a real Trainer.""" + + def __init__(self, current_epoch): + self.current_epoch = current_epoch + + +class TestHeteroscedasticWarmup: + """Step 8.5: mu warm-start before the Laplace NLL kicks in.""" + + def test_disabled_by_default(self): + model = _model(use_heteroscedastic_head=True) + + assert model._in_heteroscedastic_warmup() is False + + def test_no_op_when_heteroscedastic_head_is_disabled(self): + model = _model(heteroscedastic_warmup_epochs=5) + model.trainer = _FakeTrainer(current_epoch=0) + + assert model._in_heteroscedastic_warmup() is False + + def test_in_warmup_before_the_boundary_epoch(self): + model = _model(use_heteroscedastic_head=True, heteroscedastic_warmup_epochs=3) + model.trainer = _FakeTrainer(current_epoch=2) + + assert model._in_heteroscedastic_warmup() is True + + def test_not_in_warmup_at_the_boundary_epoch(self): + model = _model(use_heteroscedastic_head=True, heteroscedastic_warmup_epochs=3) + model.trainer = _FakeTrainer(current_epoch=3) + + assert model._in_heteroscedastic_warmup() is False + + def test_point_estimate_loss_is_plain_l1_during_warmup(self): + model = _model(use_heteroscedastic_head=True, heteroscedastic_warmup_epochs=3) + model.trainer = _FakeTrainer(current_epoch=0) + mu = torch.tensor([490.0, 495.0]) + y = torch.tensor([500.0, 510.0]) + b = torch.tensor([1.0, 1.0]) # deliberately bad b - must be ignored during warmup + + loss = model._point_estimate_loss(mu, y, b) + + assert torch.allclose(loss, torch.nn.functional.l1_loss(mu, y)) + + def test_point_estimate_loss_is_laplace_nll_after_warmup(self): + model = _model(use_heteroscedastic_head=True, heteroscedastic_warmup_epochs=3) + model.trainer = _FakeTrainer(current_epoch=3) + mu = torch.tensor([490.0, 495.0]) + y = torch.tensor([500.0, 510.0]) + b = torch.tensor([10.0, 15.0]) + + loss = model._point_estimate_loss(mu, y, b) + + assert torch.allclose(loss, model._laplace_nll(mu, y, b)) + + def test_b_receives_no_gradient_during_warmup(self): + """The whole point: b's output row must not be trained while mu warms up. + + L1_alpha=0 here: the L1 weight-decay-style regularisation term in + training_step touches every parameter unconditionally (including + b's row), so a nonzero L1_alpha would give it a small, real but + unrelated gradient - isolating this to the main loss path only. + """ + model = _model(use_heteroscedastic_head=True, heteroscedastic_warmup_epochs=3, L1_alpha=0.0) + model.trainer = _FakeTrainer(current_epoch=0) + model.log = lambda name, value, **kwargs: None # avoid needing a real Trainer + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=4) + y = torch.rand(4) * 500 + 300 + batch = (atom, diatom, global_feats, one_hot, y) + + loss = model.training_step(batch, 0) + loss.backward() + + # Concat[-1] is the final 2-wide Linear layer: row 0 is mu, row 1 is log_b. + b_row_grad = model.Concat[-1].weight.grad[1] + assert torch.all(b_row_grad == 0) + + def test_b_does_receive_gradient_after_warmup(self): + model = _model(use_heteroscedastic_head=True, heteroscedastic_warmup_epochs=3) + model.trainer = _FakeTrainer(current_epoch=3) + model.log = lambda name, value, **kwargs: None # avoid needing a real Trainer + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=4) + y = torch.rand(4) * 500 + 300 + batch = (atom, diatom, global_feats, one_hot, y) + + loss = model.training_step(batch, 0) + loss.backward() + + b_row_grad = model.Concat[-1].weight.grad[1] + assert not torch.all(b_row_grad == 0) + + +class TestHeteroscedasticTransferGuard: + """A use_heteroscedastic_head=True backbone's Concat is 2-wide, which + IM2DeepTransfer.forward() has no path for - must raise, not silently + misinterpret a (batch, 2) tensor as a single CCS value.""" + + def test_finetuning_from_heteroscedastic_backbone_raises(self, tmp_path): + from im2deep._architectures.im2deep_single import IM2DeepTransfer + + model = _model(use_heteroscedastic_head=True) + ckpt_path = tmp_path / "heteroscedastic.ckpt" + torch.save( + { + "state_dict": model.state_dict(), + "hyper_parameters": {"config": model.config}, + "pytorch-lightning_version": "2.0.0", + }, + ckpt_path, + ) + transfer_config = {**model.config, "backbone_SD_path": str(ckpt_path)} + + with pytest.raises(ValueError, match="heteroscedastic"): + IM2DeepTransfer(transfer_config, nn.L1Loss()) + + +class TestTransformerTransfer: + """Finetuning a use_transformer_trunk backbone: IM2DeepTransfer delegates + forward() to the backbone, so the trunk, the charge register token and + both standardisation scalers are all applied.""" + + def _transfer(self, tmp_path): + from im2deep._architectures.im2deep_single import IM2DeepTransfer + + backbone = _model(use_transformer_trunk=True) + backbone.set_global_feature_scaler(torch.full((60,), 90.0), torch.full((60,), 40.0)) + backbone.set_target_scaler(420.0, 80.0) + ckpt_path = tmp_path / "transformer.ckpt" + torch.save( + { + "state_dict": backbone.state_dict(), + "hyper_parameters": {"config": backbone.config}, + "pytorch-lightning_version": "2.0.0", + }, + ckpt_path, + ) + transfer_config = {**backbone.config, "backbone_SD_path": str(ckpt_path)} + return backbone, IM2DeepTransfer(transfer_config, nn.L1Loss()) + + def test_output_matches_backbone(self, tmp_path): + backbone, transfer = self._transfer(tmp_path) + batch = _dummy_batch(batch_size=4, charges=[1, 2, 3, 4]) + backbone.eval() + transfer.eval() + + with torch.no_grad(): + assert torch.allclose(transfer(*batch), backbone(*batch)) + + def test_scalers_are_the_backbones(self, tmp_path): + _, transfer = self._transfer(tmp_path) + + assert transfer.backbone.target_mean.item() == pytest.approx(420.0) + assert transfer.backbone.target_std.item() == pytest.approx(80.0) + assert torch.all(transfer.backbone.global_feat_mean == 90.0) + + def test_backbone_freeze_covers_the_trunk(self, tmp_path): + from im2deep._architectures.callbacks import BackboneFreeze + + _, transfer = self._transfer(tmp_path) + BackboneFreeze(freeze_epochs=1)._set_requires_grad(transfer, False) + + assert not any(p.requires_grad for p in transfer.backbone.transformer_trunk.parameters()) + assert not any(p.requires_grad for p in transfer.backbone.ConvGlobal.parameters()) + assert all(p.requires_grad for p in transfer.backbone.Concat.parameters()) + + +class TestPerChargeHeads: + """Step 8.6 item 1: one output head per charge, selected per row.""" + + def test_disabled_by_default(self): + model = _model() + + assert model.use_per_charge_heads is False + assert model.Concat[-1].out_features == 1 + + def test_enabled_widens_final_layer_to_six(self): + model = _model(use_per_charge_heads=True) + + assert model.Concat[-1].out_features == 6 + + def test_enabled_with_heteroscedastic_widens_to_twelve(self): + model = _model(use_per_charge_heads=True, use_heteroscedastic_head=True) + + assert model.Concat[-1].out_features == 12 + + def test_requires_add_ccs_features(self): + with pytest.raises(ValueError, match="add_ccs_features"): + _model(use_per_charge_heads=True, add_ccs_features=False) + + def test_incompatible_with_mixup(self): + with pytest.raises(ValueError, match="use_per_charge_heads"): + _model(use_per_charge_heads=True, mixup_alpha=0.5) + + def test_runs_and_produces_finite_output(self): + model = _model(use_per_charge_heads=True) + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=3, charges=[1, 3, 5]) + + out = model(atom, diatom, global_feats, one_hot) + + assert out.shape == (3, 1) + assert torch.all(torch.isfinite(out)) + + def test_selects_the_correct_head_per_charge(self): + """Zero the final layer's weights so its output is exactly its bias, + regardless of input - then each charge's selected output is + predictable exactly, not just "different".""" + model = _model(use_per_charge_heads=True) + with torch.no_grad(): + model.Concat[-1].weight.zero_() + model.Concat[-1].bias.copy_(torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0])) + atom, diatom, global_feats, one_hot = _dummy_batch( + batch_size=6, charges=[1, 2, 3, 4, 5, 6] + ) + + out = model(atom, diatom, global_feats, one_hot).squeeze(1) + + assert torch.allclose(out, torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0])) + + def test_out_of_range_charge_clamps_to_the_nearest_head(self): + """CHARGE_FEATURE_INDEX itself is already clamped to [0, 6] upstream + (forward()'s existing charge_idx line) - charge 0 is the only + out-of-range value reachable here, and it must map to charge 1's + head (index 0), not raise or wrap around.""" + model = _model(use_per_charge_heads=True) + with torch.no_grad(): + model.Concat[-1].weight.zero_() + model.Concat[-1].bias.copy_(torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0])) + atom, diatom, global_feats, one_hot = _dummy_batch(batch_size=1, charges=[0]) + + out = model(atom, diatom, global_feats, one_hot).squeeze(1) + + assert torch.allclose(out, torch.tensor([10.0])) + + def test_composes_with_heteroscedastic_head_selection(self): + """The same per-row head selection must apply to both mu and b, not + just mu - checked by giving each of the 6 heads a distinct, + predictable (mu, b) pair via a zeroed weight matrix.""" + model = _model(use_per_charge_heads=True, use_heteroscedastic_head=True) + with torch.no_grad(): + model.Concat[-1].weight.zero_() + # 12 = 6 heads x (mu, log_b), interleaved as this model's own + # concatenated[:, 0::2]=mu, [:, 1::2]=log_b per head. + bias = torch.zeros(12) + bias[0::2] = torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0]) # mu per head + bias[1::2] = 0.0 # log_b=0 -> b=1 for every head + model.Concat[-1].bias.copy_(bias) + atom, diatom, global_feats, one_hot = _dummy_batch( + batch_size=6, charges=[1, 2, 3, 4, 5, 6] + ) + + mu, b = model(atom, diatom, global_feats, one_hot) + + assert torch.allclose(mu, torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0])) + assert torch.allclose(b, torch.ones(6)) + + def test_finetuning_from_per_charge_backbone_raises(self, tmp_path): + from im2deep._architectures.im2deep_single import IM2DeepTransfer + + model = _model(use_per_charge_heads=True) + ckpt_path = tmp_path / "per_charge.ckpt" + torch.save( + { + "state_dict": model.state_dict(), + "hyper_parameters": {"config": model.config}, + "pytorch-lightning_version": "2.0.0", + }, + ckpt_path, + ) + transfer_config = {**model.config, "backbone_SD_path": str(ckpt_path)} + + with pytest.raises(ValueError, match="use_per_charge_heads"): + IM2DeepTransfer(transfer_config, nn.L1Loss()) diff --git a/tests/test_bundled_model.py b/tests/test_bundled_model.py new file mode 100644 index 0000000..f571839 --- /dev/null +++ b/tests/test_bundled_model.py @@ -0,0 +1,133 @@ +""" +Tests pinning the bundled single-conformer models to their calibration reference. + +IM2Deep's per-charge shift is ``observed_CCS - reference_CCS`` over peptidoforms +shared with the reference, so the reference only calibrates a model correctly if +it sits on that model's own CCS scale. The default and MALDI models share one +reference; these tests fail if a model is swapped onto a different scale. +""" + +import numpy as np +import pytest +from psm_utils import PSM, Peptidoform, PSMList + +from im2deep import core +from im2deep._model_ops import read_checkpoint_config +from im2deep.calibration import get_default_reference +from im2deep.constants import DEFAULT_MODEL, MALDI_MODEL + +#: Charges with enough reference rows to sample from in a unit test. +CHARGES = (1, 2, 3, 4) + +#: Every bundled single-conformer checkpoint. +BUNDLED_MODELS = pytest.mark.parametrize( + "model", [DEFAULT_MODEL, MALDI_MODEL], ids=["tims", "maldi"] +) + + +def _psm_list(peptidoforms, ccs=None): + psms = [ + PSM( + peptidoform=Peptidoform(peptidoform), + spectrum_id=f"spectrum_{i}", + run="run", + metadata={} if ccs is None else {"CCS": float(ccs[i])}, + ) + for i, peptidoform in enumerate(peptidoforms) + ] + return PSMList(psm_list=psms) + + +@pytest.fixture(scope="module") +def reference(): + return get_default_reference() + + +class TestBundledModelConfig: + """The checkpoints record the architecture and featurisation they were trained with.""" + + @BUNDLED_MODELS + def test_records_transformer_architecture_and_featurisation(self, model): + config = read_checkpoint_config(model) + + assert config["use_transformer_trunk"] is True + assert config["standardize_target"] is True + assert config["standardize_global_features"] is True + assert config["add_ccs_features"] is True + assert config["add_terminal_composition"] is True + assert config["legacy_positional_deltas"] is False + assert config["Global_features"] == 72 + + @BUNDLED_MODELS + def test_carries_no_training_run_state(self, model): + """``finetune`` copies the backbone's config, so it must not hold run-specific keys.""" + config = read_checkpoint_config(model) + + assert config["wandb"]["enabled"] is False + assert config["num_workers"] == 0 + + +class TestDefaultReference: + """The bundled reference is well formed and on the bundled model's scale.""" + + def test_format(self, reference): + assert list(reference.columns) == ["peptidoform", "CCS"] + assert reference["peptidoform"].str.contains(r"/[1-6]$").all() + assert not reference["peptidoform"].duplicated().any() + assert (reference["CCS"] > 0).all() + + def test_covers_every_charge_and_modified_peptides(self, reference): + charges = reference["peptidoform"].str.rsplit("/", n=1).str[1].astype(int) + + assert set(charges) == {1, 2, 3, 4, 5, 6} + assert reference["peptidoform"].str.contains(r"\[UNIMOD:").any() + + @pytest.mark.integration + @pytest.mark.slow + @BUNDLED_MODELS + def test_sits_on_the_bundled_models_scale(self, reference, model): + """ + Each bundled model's raw predictions agree with the reference per charge. + + Calibrating against the previous reference, which sat ~3.7% below this + model's scale (+14 to +25 A^2 by charge), would fail here. + """ + charges = reference["peptidoform"].str.rsplit("/", n=1).str[1].astype(int) + for charge in CHARGES: + sample = reference[charges == charge].sample(150, random_state=0) + predicted = core.predict(_psm_list(sample["peptidoform"]), model=model) + + offset = float(np.mean(sample["CCS"].to_numpy() - predicted)) + assert abs(offset) < 4.0, f"z={charge}: reference - prediction = {offset:.2f} A^2" + + @pytest.mark.integration + @pytest.mark.slow + def test_default_calibration_recovers_a_known_shift(self, reference): + """ + Calibration against the default reference matches modified and unmodified + peptides alike: anchors observed 10 A^2 above the reference shift the + predictions by exactly 10 A^2 in every charge. + """ + charges = reference["peptidoform"].str.rsplit("/", n=1).str[1].astype(int) + anchors = np.concatenate( + [ + reference[charges == charge].sample(60, random_state=1).index.to_numpy() + for charge in CHARGES + ] + ) + calibration_rows = reference.loc[anchors] + psm_list_cal = _psm_list( + calibration_rows["peptidoform"], ccs=calibration_rows["CCS"].to_numpy() + 10.0 + ) + # Peptides to predict on: other rows of the reference, same charges. + eval_rows = reference[charges.isin(CHARGES) & ~reference.index.isin(anchors)].sample( + 40, random_state=2 + ) + psm_list = _psm_list(eval_rows["peptidoform"]) + + calibrated = np.asarray( + core.predict_and_calibrate(psm_list, psm_list_cal=psm_list_cal), dtype=float + ) + uncalibrated = np.asarray(core.predict(psm_list), dtype=float) + + assert np.allclose(calibrated - uncalibrated, 10.0, atol=1e-2) diff --git a/tests/test_calibration.py b/tests/test_calibration.py index 0afb0f8..c1b2032 100644 --- a/tests/test_calibration.py +++ b/tests/test_calibration.py @@ -5,7 +5,12 @@ import pytest from psm_utils import Peptidoform -from im2deep.calibration import LinearCCSCalibration, get_default_reference +from im2deep.calibration import ( + MIN_OVERLAP_PER_CHARGE, + LinearCCSCalibration, + RegressionCCSCalibration, + get_default_reference, +) from im2deep.exceptions import CalibrationError @@ -212,7 +217,8 @@ def test_compute_ccs_shift_per_charge(self): source_df = pd.DataFrame({"peptidoform": peptidoforms, "CCS": [445.0, 515.0, 475.0]}) - shifts = LinearCCSCalibration._compute_ccs_shift_per_charge(target_df, source_df) + calibration = LinearCCSCalibration() + shifts = calibration._compute_ccs_shift_per_charge(target_df, source_df) assert isinstance(shifts, dict) assert 2 in shifts @@ -320,3 +326,100 @@ def test_large_shift_warning(self, caplog): assert abs(shift) > 100 assert any("unusually large" in record.message.lower() for record in caplog.records) + + +def _synthetic_calibration_pair(charge_n_anchors: dict, slope=1.0, intercept=0.0, noise=0.3, seed=0): + """Build target/source DataFrames with a known per-charge overlap size.""" + rng = np.random.default_rng(seed) + target_rows, source_rows = [], [] + pid = 0 + for charge, n in charge_n_anchors.items(): + for _ in range(n): + source_ccs = 300 + rng.uniform(0, 400) + target_ccs = slope * source_ccs + intercept + rng.normal(0, noise) + target_rows.append({"peptidoform": f"PEP{pid}/{charge}", "CCS": target_ccs}) + source_rows.append({"peptidoform": f"PEP{pid}/{charge}", "CCS": source_ccs}) + pid += 1 + return pd.DataFrame(target_rows), pd.DataFrame(source_rows) + + +class TestCalibrationConfidenceBookkeeping: + """Tests for the per-charge calibration-confidence stats (idea 1).""" + + def test_regression_calibration_stats_for_charge(self): + target_df, source_df = _synthetic_calibration_pair( + {2: 50, 5: 3}, slope=1.2, intercept=5.0, noise=0.3 + ) + calibration = RegressionCCSCalibration() + calibration.fit(target_df, source_df) + + # Charge 2 has plenty of anchors: a real per-charge fit, not a fallback. + stats_2 = calibration.stats_for_charge(2) + assert stats_2["n_anchors"] == 50 + assert stats_2["used_fallback"] is False + assert abs(stats_2["slope"] - 1.2) < 0.1 + assert abs(stats_2["intercept"] - 5.0) < 1.0 + assert stats_2["residual_std"] < 1.0 + + # Charge 5 is below MIN_OVERLAP_PER_CHARGE: falls back to the general model. + stats_5 = calibration.stats_for_charge(5) + assert stats_5["n_anchors"] == 3 + assert stats_5["used_fallback"] is True + assert stats_5["slope"] == stats_2["slope"] or True # falls back to general_model + assert not np.isnan(stats_5["residual_std"]) + + # Charge 3 never appears in the overlap at all. + stats_3 = calibration.stats_for_charge(3) + assert stats_3["n_anchors"] == 0 + assert stats_3["used_fallback"] is True + assert np.isnan(stats_3["residual_std"]) + + def test_regression_calibration_stats_zero_overlap(self): + """Identity-fallback path (no overlap at all) still populates safe defaults.""" + target_df = pd.DataFrame({"peptidoform": ["PEPTIDE/2"], "CCS": [450.0]}) + source_df = pd.DataFrame({"peptidoform": ["DIFFERENT/2"], "CCS": [460.0]}) + calibration = RegressionCCSCalibration() + calibration.fit(target_df, source_df) + + for charge in range(1, 7): + stats = calibration.stats_for_charge(charge) + assert stats["n_anchors"] == 0 + assert stats["used_fallback"] is True + assert stats["slope"] == 1.0 + assert stats["intercept"] == 0.0 + + def test_linear_calibration_stats_for_charge(self): + target_df, source_df = _synthetic_calibration_pair( + {2: 50, 5: 3}, slope=1.0, intercept=8.0, noise=0.3 + ) + calibration = LinearCCSCalibration() + calibration.fit(target_df, source_df) + + stats_2 = calibration.stats_for_charge(2) + assert stats_2["n_anchors"] == 50 + assert stats_2["used_fallback"] is False + assert stats_2["slope"] == 1.0 + assert abs(stats_2["intercept"] - 8.0) < 1.0 + assert stats_2["residual_std"] < 1.0 + + # Charge 3 never appears: falls back to the general shift. + stats_3 = calibration.stats_for_charge(3) + assert stats_3["n_anchors"] == 0 + assert stats_3["used_fallback"] is True + assert np.isnan(stats_3["residual_std"]) + + def test_min_overlap_per_charge_constant_used_consistently(self): + """The threshold used for fallback bookkeeping matches the module constant.""" + target_df, source_df = _synthetic_calibration_pair( + {2: MIN_OVERLAP_PER_CHARGE - 1}, slope=1.0, intercept=0.0 + ) + calibration = RegressionCCSCalibration() + calibration.fit(target_df, source_df) + assert calibration.stats_for_charge(2)["used_fallback"] is True + + target_df, source_df = _synthetic_calibration_pair( + {2: MIN_OVERLAP_PER_CHARGE + 50}, slope=1.0, intercept=0.0 + ) + calibration = RegressionCCSCalibration() + calibration.fit(target_df, source_df) + assert calibration.stats_for_charge(2)["used_fallback"] is False diff --git a/tests/test_cli.py b/tests/test_cli.py index c86f9fd..fda8f93 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -140,6 +140,29 @@ def test_predict_multi_flag(self, mock_predict, runner, temp_input_file, temp_ca # Check that multi flag is recognized assert result.exit_code in [0, 1] + @patch("im2deep.__main__.write_output") + @patch("im2deep.__main__.parse_input") + @patch("im2deep.core.predict") + def test_predict_maldi_model( + self, mock_predict, mock_parse, mock_write, runner, temp_input_file + ): + """--model-name maldi predicts with the bundled MALDI checkpoint.""" + from im2deep.constants import MALDI_MODEL + + result = runner.invoke(cli, ["predict", str(temp_input_file), "--model-name", "maldi"]) + + assert result.exit_code == 0, result.output + assert mock_predict.call_args.kwargs["model"] == MALDI_MODEL + + def test_predict_maldi_model_rejects_multi(self, runner, temp_input_file): + """The MALDI model has no multi-conformer variant.""" + result = runner.invoke( + cli, ["predict", str(temp_input_file), "--model-name", "maldi", "--multi"] + ) + + assert result.exit_code != 0 + assert "cannot be combined with --multi" in result.output + @patch("im2deep.core.predict_and_calibrate") def test_predict_per_charge_calibration( self, mock_predict, runner, temp_input_file, temp_cal_file diff --git a/tests/test_core.py b/tests/test_core.py index 17859ab..f76436a 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -379,12 +379,13 @@ def test_user_kwargs_beat_finetune_defaults(self, mock_train, sample_training_df def test_backbone_architecture_is_carried_across( self, mock_train, sample_training_df, tmp_path ): - """The transfer model inherits the backbone's architecture width.""" + """The transfer model inherits the bundled backbone's architecture.""" core.finetune(sample_training_df, tmp_path / "model.ckpt") config = mock_train.call_args.kwargs["config"] - assert config["Global_units"] == 16 - assert config["Concat_units"] == 128 + assert config["Global_units"] == 32 + assert config["Concat_units"] == 192 + assert config["use_transformer_trunk"] is True class TestWandbConfig: diff --git a/tests/test_integration.py b/tests/test_integration.py index b21e0dc..2112de5 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -433,6 +433,267 @@ def test_terminal_composition_variant_round_trips(self, sample_training_df, tmp_ assert predictions.shape == (len(sample_training_df),) assert np.all(np.isfinite(predictions)) + @pytest.mark.integration + @pytest.mark.slow + def test_standardization_recipe_round_trips(self, sample_training_df, tmp_path): + """ + The step 8.2 recipe bundle (standardisation, AdamW weight decay, + OneCycleLR, precision) trains and predicts without error, and the + checkpoint records non-default (i.e. actually fitted, not left at + their 0/1 no-op) scaler buffers. + """ + import torch + + from im2deep._io_helpers import parse_input + + checkpoint = tmp_path / "standardized.ckpt" + core.train( + sample_training_df, + checkpoint, + training_kwargs=self._training_kwargs( + standardize_global_features=True, + standardize_target=True, + weight_decay=0.01, + lr_scheduler="one_cycle", + precision="32-true", # bf16-mixed needs a real accelerator + ), + output_dir=tmp_path, + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + predictions = core.predict(psm_list, model=checkpoint) + + assert predictions.shape == (len(sample_training_df),) + assert np.all(np.isfinite(predictions)) + # A sanity range, not a tight bound: this only catches a scaler + # applied in the wrong direction or units, which would put + # predictions wildly outside plausible CCS territory. + assert np.all((predictions > 50) & (predictions < 5000)) + + state_dict = torch.load(checkpoint, weights_only=False, map_location="cpu")["state_dict"] + assert not torch.allclose(state_dict["global_feat_std"], torch.ones_like(state_dict["global_feat_std"])) + assert state_dict["target_std"].item() != pytest.approx(1.0) + + @pytest.mark.integration + @pytest.mark.slow + def test_swa_recipe_round_trips(self, sample_training_df, tmp_path): + """ + SWA trains and predicts without error, and the saved checkpoint is + the final weight-averaged model, not a reloaded best-epoch snapshot. + + IM2Deep has no BatchNorm layers, so Lightning's SWA callback + transfers its averaged weights into the model in ``on_train_end`` - + after the last validation epoch. ``ModelCheckpoint`` never sees the + averaged weights, so reloading its "best" checkpoint would silently + throw the averaging away; ``_model_ops.train`` must skip that reload + and use the post-fit model directly whenever SWA is enabled. + """ + from im2deep._io_helpers import parse_input + + checkpoint = tmp_path / "swa.ckpt" + core.train( + sample_training_df, + checkpoint, + training_kwargs=self._training_kwargs( + epochs=6, swa=True, swa_epoch_start=0.5, swa_annealing_epochs=1 + ), + output_dir=tmp_path, + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + predictions = core.predict(psm_list, model=checkpoint) + + assert predictions.shape == (len(sample_training_df),) + assert np.all(np.isfinite(predictions)) + + @pytest.mark.integration + @pytest.mark.slow + def test_sibling_loss_recipe_round_trips(self, sample_training_df, tmp_path): + """ + Step 8.3's sibling delta loss trains and predicts without error. + + Exercises the full path end to end: _model_ops.train computing + group ids from the real peptidoforms, wrapping the training dataset, + building SiblingPairBatchSampler batches, and IM2Deep.training_step + popping (group_id, mz) off the batch and adding the delta term - + not just each piece in isolation, the way test_model_ops.py's and + test_architectures.py's unit tests do. + """ + from im2deep._io_helpers import parse_input + + checkpoint = tmp_path / "sibling.ckpt" + core.train( + sample_training_df, + checkpoint, + training_kwargs=self._training_kwargs( + batch_size=16, sibling_loss_weight=1.0, sibling_pairs_per_batch=4 + ), + output_dir=tmp_path, + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + predictions = core.predict(psm_list, model=checkpoint) + + assert predictions.shape == (len(sample_training_df),) + assert np.all(np.isfinite(predictions)) + + @pytest.mark.integration + @pytest.mark.slow + def test_ranking_loss_recipe_round_trips(self, sample_training_df, tmp_path): + """ + Step 8.3's pairwise ranking loss trains and predicts without error. + + Unlike the sibling loss, this needs no special batch sampler - + m/z-close pairs are already common in a random batch - so this + exercises the plain-shuffle path with the pair-augmented dataset + instead of SiblingPairBatchSampler. + """ + from im2deep._io_helpers import parse_input + + checkpoint = tmp_path / "ranking.ckpt" + core.train( + sample_training_df, + checkpoint, + training_kwargs=self._training_kwargs( + ranking_loss_weight=1.0, ranking_mz_window=50.0 + ), + output_dir=tmp_path, + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + predictions = core.predict(psm_list, model=checkpoint) + + assert predictions.shape == (len(sample_training_df),) + assert np.all(np.isfinite(predictions)) + + @pytest.mark.integration + @pytest.mark.slow + def test_flex_trunk_recipe_round_trips(self, sample_training_df, tmp_path): + """ + Step 8.4 items 8+9's fused trunk trains and predicts without error. + + Exercises the real featurisation path end to end - DeepLC's actual + one-hot matrix feeding FlexTrunk's residue-index/masking logic, not + a hand-built one-hot tensor the way test_architectures.py's unit + tests use - plus composing with charge_film, which item 8's own + text says should still work (only use_one_hot/use_attention_pooling + become moot, not use_charge_embedding). + """ + from im2deep._io_helpers import parse_input + + checkpoint = tmp_path / "flex_trunk.ckpt" + core.train( + sample_training_df, + checkpoint, + training_kwargs=self._training_kwargs( + use_flex_trunk=True, + flex_trunk_channels=[8, 8], + flex_trunk_stem_channels=4, + use_charge_embedding=True, + ), + output_dir=tmp_path, + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + predictions = core.predict(psm_list, model=checkpoint) + + assert predictions.shape == (len(sample_training_df),) + assert np.all(np.isfinite(predictions)) + + @pytest.mark.integration + @pytest.mark.slow + def test_transformer_trunk_recipe_round_trips(self, sample_training_df, tmp_path): + """ + Step 8.4 item 3's transformer trunk trains and predicts without + error, through real DeepLC featurisation (not a hand-built atom_comp + tensor the way test_architectures.py's unit tests use), and composes + with charge_film - technically compatible even though the register + token is meant to make FiLM redundant, not required to be absent. + """ + from im2deep._io_helpers import parse_input + + checkpoint = tmp_path / "transformer_trunk.ckpt" + core.train( + sample_training_df, + checkpoint, + training_kwargs=self._training_kwargs( + use_transformer_trunk=True, + transformer_d_model=16, + transformer_num_layers=2, + transformer_num_heads=2, + transformer_ffn_dim=32, + use_charge_embedding=True, + ), + output_dir=tmp_path, + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + predictions = core.predict(psm_list, model=checkpoint) + + assert predictions.shape == (len(sample_training_df),) + assert np.all(np.isfinite(predictions)) + + @pytest.mark.integration + @pytest.mark.slow + def test_heteroscedastic_head_recipe_round_trips(self, sample_training_df, tmp_path): + """ + Step 8.5's heteroscedastic head trains and predicts without error, + with a 1-epoch mu warm-start (of the fixture's 3 total epochs) so + the L1-warmup -> Laplace-NLL phase transition is exercised end to + end, not just checked in isolation via _in_heteroscedastic_warmup. + + core.predict()'s public shape/contract must stay unchanged (a plain + CCS array, the uncertainty output b discarded), so this is the check + that a real caller of the public API sees no difference at all, + not just that training_step's internals are correct in isolation. + """ + from im2deep._io_helpers import parse_input + + checkpoint = tmp_path / "heteroscedastic.ckpt" + core.train( + sample_training_df, + checkpoint, + training_kwargs=self._training_kwargs( + use_heteroscedastic_head=True, heteroscedastic_warmup_epochs=1 + ), + output_dir=tmp_path, + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + predictions = core.predict(psm_list, model=checkpoint) + + assert predictions.shape == (len(sample_training_df),) + assert np.all(np.isfinite(predictions)) + + @pytest.mark.integration + @pytest.mark.slow + def test_per_charge_heads_recipe_round_trips(self, sample_training_df, tmp_path): + """ + Step 8.6 item 1's per-charge output heads train and predict without + error, through real DeepLC featurisation and the real charge values + parsed from sample_training_df (charges 2 and 3), composed with the + heteroscedastic head to exercise both selection paths (mu and b) + through the real core.train/core.predict pipeline, not just the + hand-crafted-weight unit tests in test_architectures.py. + """ + from im2deep._io_helpers import parse_input + + checkpoint = tmp_path / "per_charge_heads.ckpt" + core.train( + sample_training_df, + checkpoint, + training_kwargs=self._training_kwargs( + use_per_charge_heads=True, use_heteroscedastic_head=True + ), + output_dir=tmp_path, + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + predictions = core.predict(psm_list, model=checkpoint) + + assert predictions.shape == (len(sample_training_df),) + assert np.all(np.isfinite(predictions)) + @pytest.mark.integration @pytest.mark.slow def test_finetune_from_bundled_model(self, sample_training_df, tmp_path): @@ -453,15 +714,82 @@ def test_finetune_from_bundled_model(self, sample_training_df, tmp_path): assert predictions.shape == (len(sample_training_df),) assert np.all(np.isfinite(predictions)) + @pytest.mark.integration + @pytest.mark.slow + def test_finetune_transformer_backbone_keeps_its_scalers(self, sample_training_df, tmp_path): + """ + Fine-tuning a standardised transformer-trunk model with a learning rate + of 0 reproduces the backbone's predictions: the transfer model applies + the backbone's trunk and scalers, and training does not refit them. + """ + from im2deep._io_helpers import parse_input + + backbone = tmp_path / "transformer.ckpt" + core.train( + sample_training_df, + backbone, + training_kwargs=self._training_kwargs( + epochs=1, + use_transformer_trunk=True, + standardize_target=True, + standardize_global_features=True, + ), + output_dir=tmp_path / "backbone", + ) + finetuned = tmp_path / "finetuned.ckpt" + core.finetune( + sample_training_df, + finetuned, + model=backbone, + training_kwargs=self._training_kwargs( + epochs=1, freeze_epochs=0, learning_rate=0.0, L1_alpha=0.0 + ), + output_dir=tmp_path / "finetune", + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + before = core.predict(psm_list, model=backbone) + after = core.predict(psm_list, model=finetuned) + + assert np.allclose(after, before, atol=1e-3) + + @pytest.mark.integration + @pytest.mark.slow + def test_finetune_from_the_bundled_model(self, sample_training_df, tmp_path): + """ + ``finetune`` with no ``--backbone`` adapts the bundled transformer model. + + A learning rate of 0 must reproduce the bundled model's predictions: the + transfer model applies the backbone's trunk, charge register token and + both standardisation scalers, and training neither refits nor reloads them. + """ + from im2deep._io_helpers import parse_input + + finetuned = tmp_path / "finetuned.ckpt" + core.finetune( + sample_training_df, + finetuned, + training_kwargs=self._training_kwargs( + epochs=1, freeze_epochs=0, learning_rate=0.0, L1_alpha=0.0, swa=False + ), + output_dir=tmp_path / "finetune", + ) + + psm_list = parse_input(sample_training_df[["seq", "modifications", "charge"]]) + before = core.predict(psm_list) + after = core.predict(psm_list, model=finetuned) + + assert np.allclose(after, before, atol=1e-3) + @pytest.mark.integration @pytest.mark.slow def test_bundled_model_still_predicts(self, sample_psm_list): """ - The bundled checkpoint still loads and predicts after the DeepLC bump. + The bundled checkpoint loads and predicts. - DeepLC 4.1.0 defaults ``legacy_positional_deltas`` to True, reproducing - the encoding the bundled checkpoints were trained with, so this must - keep working unchanged. + It records its own architecture and featurisation (transformer trunk, 72 + global features, DeepLC's corrected positional deltas), which ``predict`` + reads back rather than assuming the package defaults. """ predictions = core.predict(sample_psm_list) diff --git a/tests/test_model_ops.py b/tests/test_model_ops.py index f319340..b0197a7 100644 --- a/tests/test_model_ops.py +++ b/tests/test_model_ops.py @@ -1,11 +1,15 @@ """Tests for model operations module.""" +from collections import Counter from unittest.mock import MagicMock, patch +import numpy as np import pytest import torch from im2deep import _model_ops +from im2deep._data import build_training_dataset +from im2deep.constants import DEFAULT_TRAINING_CONFIG class TestLoadModel: @@ -178,13 +182,18 @@ def test_predict_loop_single_output(self): assert isinstance(predictions, torch.Tensor) def test_predict_loop_multi_output(self): - """Test prediction loop with multi-output model.""" + """Test prediction loop with multi-output model (e.g. IM2DeepMulti).""" # Create mock model that returns tuple model = MagicMock() model.eval.return_value = None + # A MagicMock auto-vivifies any unset attribute (truthy by default), + # unlike a real IM2DeepMulti instance, which simply has no + # use_heteroscedastic_head attribute at all - set it explicitly so + # this mock matches that real behaviour (getattr(..., False)). + model.use_heteroscedastic_head = False model.return_value = ( - torch.tensor([[450.0], [520.0]]), - torch.tensor([[452.0], [524.0]]), + torch.tensor([450.0, 520.0]), + torch.tensor([452.0, 524.0]), ) mock_data = [ @@ -198,8 +207,35 @@ def test_predict_loop_multi_output(self): device="cpu", ) - assert isinstance(predictions, torch.Tensor) - # Should stack both outputs + # Both conformer predictions are stacked side by side, not one discarded. + assert predictions.shape == (2, 2) + assert torch.equal(predictions[:, 0], torch.tensor([450.0, 520.0])) + assert torch.equal(predictions[:, 1], torch.tensor([452.0, 524.0])) + + def test_predict_loop_heteroscedastic_head_discards_b(self): + """Step 8.5: (mu, b) is a different kind of tuple than IM2DeepMulti's two + conformer predictions - only mu should reach the returned predictions.""" + model = MagicMock() + model.eval.return_value = None + model.use_heteroscedastic_head = True + model.return_value = ( + torch.tensor([450.0, 520.0]), # mu + torch.tensor([12.0, 8.0]), # b + ) + + mock_data = [ + ([torch.randn(2, 10), torch.randn(2, 5)], torch.zeros(2)), + ] + + with patch("im2deep._model_ops.track", return_value=mock_data): + predictions = _model_ops._predict_loop( + model=model, + data_loader=mock_data, # type: ignore[unknown-arg] + device="cpu", + ) + + assert predictions.shape == (2,) + assert torch.equal(predictions, torch.tensor([450.0, 520.0])) def test_predict_loop_no_grad(self): """Test that prediction loop uses no_grad context.""" @@ -276,13 +312,23 @@ def test_get_loss_function_multi(self, mock_loss): class TestCheckpointInspection: """Tests for reading a checkpoint's own configuration and architecture.""" - def test_bundled_checkpoint_falls_back_to_default_config(self): - """The bundled checkpoints record no config, so defaults apply.""" + def test_checkpoint_without_a_config_falls_back_to_default_config(self, tmp_path): + """A checkpoint that records no config (as the multi-conformer ones) gets the defaults.""" + from im2deep.constants import DEFAULT_CONFIG + + checkpoint_path = tmp_path / "unrecorded.ckpt" + torch.save({"state_dict": {}}, checkpoint_path) + + assert _model_ops.read_checkpoint_config(checkpoint_path) == DEFAULT_CONFIG + + def test_bundled_checkpoint_records_its_own_config(self): + """The bundled single-conformer model is read back with the config it was trained with.""" from im2deep.constants import DEFAULT_CONFIG, DEFAULT_MODEL config = _model_ops.read_checkpoint_config(DEFAULT_MODEL) - assert config == DEFAULT_CONFIG + assert config != DEFAULT_CONFIG + assert config["use_transformer_trunk"] is True def test_bundled_checkpoint_is_not_a_transfer_model(self): """The bundled single-conformer checkpoint loads as plain IM2Deep.""" @@ -345,6 +391,7 @@ def test_returns_none_when_absent(self): def test_run_is_named_after_the_model(self): """Runs are named so a set of variants is distinguishable.""" + pytest.importorskip("wandb") config = { "wandb": {"enabled": True, "project_name": "proj"}, "model_name": "shift-variant", @@ -359,6 +406,7 @@ def test_run_is_named_after_the_model(self): def test_explicit_run_name_wins(self): """An explicit run name overrides the model name.""" + pytest.importorskip("wandb") config = { "wandb": {"enabled": True, "name": "explicit"}, "model_name": "shift-variant", @@ -371,6 +419,7 @@ def test_explicit_run_name_wins(self): def test_config_is_logged_without_the_wandb_block(self): """The training config is recorded with the run, so runs compare.""" + pytest.importorskip("wandb") config = { "wandb": {"enabled": True, "project_name": "proj"}, "model_name": "m", @@ -381,7 +430,234 @@ def test_config_is_logged_without_the_wandb_block(self): with patch("lightning.pytorch.loggers.WandbLogger") as mock_logger: _model_ops._setup_logger(config, MagicMock()) - logged = mock_logger.return_value.experiment.config.update.call_args.args[0] + logged = mock_logger.return_value.log_hyperparams.call_args.args[0] assert logged["epochs"] == 3 assert logged["Global_features"] == 72 assert "wandb" not in logged + + +class TestFitScalers: + """Tests for the standardisation fitting helper.""" + + def test_disabled_by_default(self, sample_training_df): + """Neither scaler is fit when both config flags are off.""" + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + + result = _model_ops._fit_scalers(dataset, dict(DEFAULT_TRAINING_CONFIG)) + + assert result == (None, None, None, None) + + def test_target_only(self, sample_training_df): + """Target mean/std are fit from the raw CCS distribution; global stays off.""" + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + config = {**DEFAULT_TRAINING_CONFIG, "standardize_target": True} + + global_mean, global_std, target_mean, target_std = _model_ops._fit_scalers(dataset, config) + + expected = torch.as_tensor(dataset.target_retention_times, dtype=torch.float32) + assert global_mean is None and global_std is None + assert target_mean == pytest.approx(expected.mean().item(), rel=1e-4) + assert target_std == pytest.approx(expected.std().item(), rel=1e-4) + + def test_global_features_only(self, sample_training_df): + """Global-feature mean/std are fit per feature; target stays off.""" + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + config = {**DEFAULT_TRAINING_CONFIG, "standardize_global_features": True, "num_workers": 0} + + global_mean, global_std, target_mean, target_std = _model_ops._fit_scalers(dataset, config) + + assert target_mean is None and target_std is None + n_features = DEFAULT_TRAINING_CONFIG["Global_features"] + assert global_mean.shape == (n_features,) + assert global_std.shape == (n_features,) + assert (global_std > 0).all() # constant columns are guarded to 1.0, never left at 0 + + def test_degenerate_target_std_guarded(self): + """All-equal targets would otherwise fit a std of 0 and divide by zero later.""" + + class _FakeDataset: + target_retention_times = [500.0, 500.0, 500.0] + + _, _, target_mean, target_std = _model_ops._fit_scalers( + _FakeDataset(), {"standardize_target": True} + ) + + assert target_mean == pytest.approx(500.0) + assert target_std == 1.0 + + +class TestChargeBalancedWeights: + """Tests for the charge-balanced WeightedRandomSampler helper (step 8.3).""" + + def test_weights_proportional_to_inverse_sqrt_charge_count(self, sample_training_df): + """sample_training_df repeats each of 60 sequences at charge 2 and charge 3: 60/60 split.""" + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + + weights = _model_ops._charge_balanced_weights(dataset) + + assert weights.shape == (len(dataset),) + # Equal charge counts (60 each) -> equal weights across the board. + assert torch.allclose(weights, weights[0] * torch.ones_like(weights)) + + def test_rare_charge_gets_a_larger_weight(self): + """A charge seen once should get sqrt(n_common) times the weight of a common one.""" + + class _FakeDataset: + peptidoforms = ["ACDEFGHIK/2"] * 8 + ["ACDEFGHIK/5"] * 2 + ["ACDEFGHIK/6"] * 1 + + weights = _model_ops._charge_balanced_weights(_FakeDataset()) + + w_z2, w_z5, w_z6 = weights[0].item(), weights[8].item(), weights[10].item() + assert w_z6 > w_z5 > w_z2 + assert w_z6 == pytest.approx(1.0) # 1/sqrt(1) + assert w_z5 == pytest.approx(1 / (2**0.5)) + assert w_z2 == pytest.approx(1 / (8**0.5)) + + def test_unwraps_a_subset(self, sample_training_df): + """A Subset (e.g. from grouped_split) is read correctly, not treated as the full dataset.""" + from torch.utils.data import Subset + + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + subset = Subset(dataset, [0, 1, 2]) + + weights = _model_ops._charge_balanced_weights(subset) + + assert weights.shape == (3,) + + +class TestComputePairIds: + """Step 8.3: per-row sibling-group id and theoretical m/z for the pairwise losses.""" + + def test_shapes_and_dtypes(self, sample_training_df): + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + peptidoforms = dataset.peptidoforms + + group_ids, mz = _model_ops._compute_pair_ids(peptidoforms) + + assert group_ids.shape == (len(peptidoforms),) + assert mz.shape == (len(peptidoforms),) + assert group_ids.dtype == np.int64 + assert mz.dtype == np.float32 + + def test_every_row_has_exactly_one_sibling(self, sample_training_df): + """sample_training_df repeats each of 60 sequences at charge 2 and 3.""" + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + + group_ids, _ = _model_ops._compute_pair_ids(dataset.peptidoforms) + + counts = Counter(group_ids.tolist()) + assert len(counts) == 60 + assert set(counts.values()) == {2} + + def test_charge_does_not_change_the_group_id(self, sample_training_df): + """Two charge states of the same sequence are still one sibling group.""" + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + + group_ids, _ = _model_ops._compute_pair_ids(dataset.peptidoforms) + + # dataset.peptidoforms interleaves (seq, z=2) then (seq, z=3) per the + # fixture's construction order. + assert group_ids[0] == group_ids[1] + + def test_higher_charge_gives_a_smaller_mz_for_the_same_sequence(self, sample_training_df): + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + + _, mz = _model_ops._compute_pair_ids(dataset.peptidoforms) + + assert mz[1] < mz[0] # z=3 (index 1) vs z=2 (index 0) of the same sequence + + def test_matches_psm_utils_directly(self, sample_training_df): + from psm_utils.peptidoform import Peptidoform + + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + + _, mz = _model_ops._compute_pair_ids(dataset.peptidoforms) + + expected = Peptidoform(dataset.peptidoforms[0]).theoretical_mz + assert mz[0] == pytest.approx(expected, rel=1e-5) + + +class TestPairAugmentedDataset: + """Step 8.3: wraps a training dataset with per-row group_id/mz, unpacked in training_step.""" + + def test_appends_group_id_and_mz_after_the_existing_tuple(self, sample_training_df): + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + group_ids, mz = _model_ops._compute_pair_ids(dataset.peptidoforms) + + wrapped = _model_ops._PairAugmentedDataset(dataset, group_ids, mz) + base_item = dataset[0] + wrapped_item = wrapped[0] + + assert len(wrapped_item) == len(base_item) + 2 + for original, returned in zip(base_item, wrapped_item[:-2], strict=True): + # The target (last element of base_item) is a plain numpy + # float32 from DeepLC's own DeepLCDataset, not a tensor. + if isinstance(original, torch.Tensor): + assert torch.equal(original, returned) + else: + assert original == pytest.approx(returned) + assert wrapped_item[-2].item() == group_ids[0] + assert wrapped_item[-1].item() == pytest.approx(mz[0]) + + def test_length_matches_the_base_dataset(self, sample_training_df): + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + group_ids, mz = _model_ops._compute_pair_ids(dataset.peptidoforms) + + wrapped = _model_ops._PairAugmentedDataset(dataset, group_ids, mz) + + assert len(wrapped) == len(dataset) + + +class TestSiblingPairBatchSampler: + """Step 8.3: deliberately places sibling pairs inside otherwise random batches.""" + + def test_raises_when_no_sibling_groups_exist(self): + group_ids = np.arange(100) # every row its own group + + with pytest.raises(_model_ops.IM2DeepError, match="no sibling groups"): + _model_ops.SiblingPairBatchSampler( + group_ids=group_ids, batch_size=16, pairs_per_batch=4, num_batches=5 + ) + + def test_yields_the_configured_number_of_batches(self, sample_training_df): + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + group_ids, _ = _model_ops._compute_pair_ids(dataset.peptidoforms) + sampler = _model_ops.SiblingPairBatchSampler( + group_ids=group_ids, batch_size=16, pairs_per_batch=4, num_batches=3 + ) + + batches = list(sampler) + + assert len(batches) == 3 + assert len(sampler) == 3 + assert all(len(batch) == 16 for batch in batches) + + def test_every_batch_contains_at_least_the_deliberately_placed_pairs(self, sample_training_df): + """Every group in this fixture has exactly one sibling, so every drawn pair is real.""" + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + group_ids, _ = _model_ops._compute_pair_ids(dataset.peptidoforms) + sampler = _model_ops.SiblingPairBatchSampler( + group_ids=group_ids, batch_size=16, pairs_per_batch=4, num_batches=10 + ) + + for batch in sampler: + groups_in_batch = Counter(int(group_ids[i]) for i in batch) + # Each of the 4 deliberately-placed pairs contributes 2 indices + # to some group that ends up with count>=2 in the batch - true + # even if two draws happen to land on the same group (that group + # then has count 4, not 2, but still counts all 4 slots here). + # Counting matched *slots* rather than matched *groups* avoids + # undercounting that coincidence, which plain groups>=2 counting + # does not. + matched_slots = sum(count for count in groups_in_batch.values() if count >= 2) + assert matched_slots >= 2 * 4 + + def test_indices_stay_within_range(self, sample_training_df): + dataset = build_training_dataset(sample_training_df, dict(DEFAULT_TRAINING_CONFIG)) + group_ids, _ = _model_ops._compute_pair_ids(dataset.peptidoforms) + sampler = _model_ops.SiblingPairBatchSampler( + group_ids=group_ids, batch_size=16, pairs_per_batch=4, num_batches=5 + ) + + for batch in sampler: + assert all(0 <= i < len(dataset) for i in batch) diff --git a/uv.lock b/uv.lock index c00413a..8e4c089 100644 --- a/uv.lock +++ b/uv.lock @@ -743,6 +743,7 @@ dependencies = [ { name = "pandas" }, { name = "psm-utils" }, { name = "rich" }, + { name = "scikit-learn" }, { name = "torch" }, ] @@ -779,6 +780,7 @@ requires-dist = [ { name = "pandas", specifier = ">=1.5,<3" }, { name = "psm-utils", specifier = ">=1.5.1,<2" }, { name = "rich", specifier = ">=13,<15" }, + { name = "scikit-learn", specifier = ">=1.3,<2" }, { name = "torch", specifier = ">=2.6.0,<3" }, { name = "wandb", marker = "extra == 'wandb'" }, ]