2026 model update

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

README.md

@@ -51,6 +51,17 @@ To use the multi-output prediction model (requires optional dependencies):
 im2deep <path/to/peptide_file.csv> --calibration-precursors <path/to/calibration_file.csv> --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 <path/to/peptide_file.csv> --calibration-precursors <path/to/calibration_file.csv> --model-name maldi
+```
+In Python, pass `model=im2deep.constants.MALDI_MODEL` to `predict()` or
+`predict_and_calibrate()`.
+
 **Output options:**
 ```bash
 im2deep <path/to/peptide_file.csv> --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

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

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!")

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))