diff --git a/CHANGELOG.md b/CHANGELOG.md index 269b58b..e1bad6b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to ### Added +- Bundled multitask pretrained model as the new default, trained across multiple LC setups. Automatic head selection in `calibrate()` based on Pearson correlation ensures the best fitting setup is used for predictions. Fine-tuning uses adapter-based transfer learning: a small MLP is attached on top of the multi-setup prediction heads. - NiceGUI-based web interface, launchable as a browser app (`deeplc gui`) or native desktop window (`deeplc gui --native`) - `[gui]` optional dependency group (nicegui, plotly, pywebview) for desktop use - `[web]` optional dependency group (nicegui, plotly) for server/Docker use @@ -18,6 +19,7 @@ and this project adheres to - `predict_and_calibrate()` core function combining prediction and calibration in one call, with optional automatic reference PSM selection - `finetune_and_predict()` core function for transfer learning followed by calibrated prediction - Automatic calibration reference selection from input PSMs using q-value filtering or top-scoring fraction +- `Calibration.selected_model_head` field to record which model output head a calibration was fitted to - Publish workflow with Windows installer build, Docker image build, and dry-run mode for CI testing without publishing ### Changed diff --git a/deeplc/__main__.py b/deeplc/__main__.py index 1db4624..4053618 100644 --- a/deeplc/__main__.py +++ b/deeplc/__main__.py @@ -179,6 +179,7 @@ def gui(native): def main(): + """Entry point for the DeepLC CLI.""" cli() diff --git a/deeplc/_architecture.py b/deeplc/_architecture.py index a21a1cc..f6e1110 100644 --- a/deeplc/_architecture.py +++ b/deeplc/_architecture.py @@ -1,10 +1,12 @@ """ PyTorch architecture definitions for DeepLC. -This module contains the neural network architectures used by DeepLC for -predicting peptide retention times based on atomic composition and other features. +This module contains the neural network architectures used by DeepLC for predicting peptide +retention times based on atomic composition and other features. """ +from __future__ import annotations + import torch import torch.nn as nn @@ -13,8 +15,8 @@ class LeakyReLUSaturation(nn.Module): """ Leaky ReLU activation with saturation (max value clipping). - This custom activation function applies leaky ReLU followed by clamping - to a maximum value, matching the original TensorFlow implementation's behavior. + This custom activation function applies leaky ReLU followed by clamping to a maximum value, + matching the original TensorFlow implementation's behavior. Parameters ---------- @@ -22,6 +24,7 @@ class LeakyReLUSaturation(nn.Module): Negative slope coefficient for leaky ReLU (default: 0.1) max_value Maximum output value for saturation (default: 20.0) + """ def __init__(self, negative_slope: float = 0.1, max_value: float = 20.0): @@ -53,6 +56,7 @@ class ConvBlock(nn.Module): Size of the max pooling window (default: 2) regularizer_val L1 regularization coefficient (default: 0.000005) + """ def __init__( @@ -86,6 +90,7 @@ def __init__( ) self.activation2 = LeakyReLUSaturation() + self.pool: nn.MaxPool1d | None = None if use_pooling: self.pool = nn.MaxPool1d(kernel_size=pool_size, stride=pool_size) @@ -97,7 +102,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.activation1(x) x = self.conv2(x) x = self.activation2(x) - if self.use_pooling: + if self.pool is not None: x = self.pool(x) return x @@ -116,6 +121,7 @@ class GlobalFeatureBranch(nn.Module): Layer size for each hidden layer regularizer_val L1 regularization coefficient (default: 0.000005) + """ def __init__( @@ -147,8 +153,8 @@ class OneHotBranch(nn.Module): """ Convolutional branch for processing one-hot encoded amino acid sequences. - This branch uses tanh activation instead of leaky ReLU and processes - one-hot encoded amino acid features. + This branch uses tanh activation instead of leaky ReLU and processes one-hot encoded amino + acid features. Parameters ---------- @@ -158,6 +164,7 @@ class OneHotBranch(nn.Module): Length of the input sequence kernel_size Size of the convolutional kernel (default: 2) + """ def __init__( @@ -201,21 +208,56 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x +class BatchedHeads(nn.Module): + """ + Parallel output heads sharing a hidden projection. + + Each head maps the shared trunk output to a scalar via a two-step computation: a batched + linear projection followed by a per-head dot product with a learned weight vector. + + Parameters + ---------- + input_size + Size of the input feature vector (output of shared trunk). + n_heads + Number of parallel output heads. + hidden + Hidden dimension per head (default: 32). + + """ + + def __init__(self, input_size: int, n_heads: int, hidden: int = 32): + super().__init__() + self.layer1 = nn.Linear(input_size, n_heads * hidden) + self.w2 = nn.Parameter(torch.zeros(n_heads, hidden)) + self.b2 = nn.Parameter(torch.zeros(n_heads)) + nn.init.normal_(self.w2, std=0.05) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.layer1(x) + n_heads = self.b2.shape[0] + h = torch.relu(h.view(h.shape[0], n_heads, h.shape[1] // n_heads)) + return (h * self.w2.unsqueeze(0)).sum(dim=-1) + self.b2 # (batch, n_heads) + + class DeepLCModel(nn.Module): """ - Complete DeepLC model for peptide retention time prediction. + DeepLC model for peptide retention time prediction. + + Four parallel input branches — per-position atomic composition CNN, summed atomic composition + CNN, global feature dense network, and one-hot amino acid CNN — are concatenated and passed + through a shared dense trunk. Outputs are projected by :class:`BatchedHeads` to ``[batch, + n_heads]``. - This model consists of multiple branches processing different feature types: - - Atomic composition CNN (per-position atomic features) - - Summed atomic composition CNN (aggregated atomic features) - - Global feature branch (peptide-level features) - - One-hot encoding branch (amino acid sequence) + When ``n_heads > 1`` the model is a multitask backbone trained across multiple LC setups. + Call :meth:`add_adapter` to attach a fine-tuning MLP that maps the head vector to a single + RT value ``[batch, 1]``. - The outputs are concatenated and passed through a deep fully connected network - for final prediction. Parameters ---------- + n_heads + Number of parallel output heads (default: 1). atom_sequence_length Length of the atomic composition sequence (default: 60) atom_channels @@ -254,10 +296,12 @@ class DeepLCModel(nn.Module): Number of final dense layers (default: 5) regularizer_val L1 regularization coefficient (default: 0.000005) + """ def __init__( self, + n_heads: int = 1, atom_sequence_length: int = 60, atom_channels: int = 6, atom_sum_sequence_length: int = 30, @@ -275,13 +319,13 @@ def __init__( global_layer_size: int = 64, global_num_layers: int = 4, final_layer_size: int = 128, - final_num_layers: int = 5, + final_num_layers: int = 4, regularizer_val: float = 0.000005, ): super().__init__() # Branch A: Atomic composition CNN - a_layers = [] + a_layers: list[nn.Module] = [] in_channels = atom_channels for block_idx in range(atom_cnn_blocks): out_channels = int(atom_cnn_filters_start / (2**block_idx)) @@ -297,11 +341,10 @@ def __init__( ) ) in_channels = out_channels - self.branch_a = nn.Sequential(*a_layers, nn.Flatten()) # Branch B: Summed atomic composition CNN - b_layers = [] + b_layers: list[nn.Module] = [] in_channels = atom_channels for block_idx in range(sum_cnn_blocks): out_channels = int(sum_cnn_filters_start / (2**block_idx)) @@ -317,7 +360,6 @@ def __init__( ) ) in_channels = out_channels - self.branch_b = nn.Sequential(*b_layers, nn.Flatten()) # Branch C: Global features @@ -330,48 +372,61 @@ def __init__( # Branch D: One-hot encoding self.branch_d = OneHotBranch( - one_hot_channels, - one_hot_sequence_length, - kernel_size=one_hot_kernel_size, + one_hot_channels, one_hot_sequence_length, kernel_size=one_hot_kernel_size ) - # Calculate concatenated feature size - # Need to compute output sizes after convolutions and pooling + # Compute concatenated feature size via a dummy forward pass with torch.no_grad(): - dummy_a = torch.zeros(1, atom_channels, atom_sequence_length) - dummy_b = torch.zeros(1, atom_channels, atom_sum_sequence_length) - dummy_c = torch.zeros(1, global_feature_size) - - out_a = self.branch_a(dummy_a) - out_b = self.branch_b(dummy_b) - out_c = self.branch_c(dummy_c) - - dummy_d = torch.zeros(1, one_hot_channels, one_hot_sequence_length) - out_d = self.branch_d(dummy_d) - - concat_size = out_a.shape[1] + out_b.shape[1] + out_c.shape[1] + out_d.shape[1] + concat_size = ( + self.branch_a(torch.zeros(1, atom_channels, atom_sequence_length)).shape[1] + + self.branch_b(torch.zeros(1, atom_channels, atom_sum_sequence_length)).shape[1] + + self.branch_c(torch.zeros(1, global_feature_size)).shape[1] + + self.branch_d(torch.zeros(1, one_hot_channels, one_hot_sequence_length)).shape[1] + ) - # Final dense layers - final_layers = [] + # Shared trunk: dense layers without output linear + trunk_layers: list[nn.Module] = [] for i in range(final_num_layers): in_features = concat_size if i == 0 else final_layer_size - final_layers.extend( - [ - nn.Linear(in_features, final_layer_size), - LeakyReLUSaturation(), - ] - ) + trunk_layers.extend([nn.Linear(in_features, final_layer_size), LeakyReLUSaturation()]) + self.shared_trunk = nn.Sequential(*trunk_layers) - # Output layer - final_layers.append(nn.Linear(final_layer_size, 1)) + # Output heads + self.heads = BatchedHeads(final_layer_size, n_heads) - self.final_network = nn.Sequential(*final_layers) + # Optional fine-tuning adapter (None until add_adapter() is called) + self.adapter: nn.Module | None = None - # Initialize weights self._initialize_weights() - def _initialize_weights(self): - """Initialize weights using normal distribution (matching TensorFlow's RandomNormal).""" + def add_adapter(self, hidden_size: int = 256) -> None: + """ + Attach a fine-tuning adapter mapping the head vector to one RT output. + + Adapter parameters are left at PyTorch default initialization and trained from scratch + during fine-tuning. + """ + n_heads = self.heads.b2.shape[0] + self.adapter = nn.Sequential( + nn.Linear(n_heads, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, max(1, hidden_size // 2)), + nn.ReLU(), + nn.Linear(max(1, hidden_size // 2), 1), + ) + self.adapter.to(self.heads.b2.device) + + def freeze_backbone(self) -> None: + """Freeze all parameters except the adapter.""" + for name, param in self.named_parameters(): + param.requires_grad = name.startswith("adapter.") + + def unfreeze_backbone(self) -> None: + """Unfreeze all parameters.""" + for param in self.parameters(): + param.requires_grad = True + + def _initialize_weights(self) -> None: for module in self.modules(): if isinstance(module, (nn.Conv1d, nn.Linear)): nn.init.normal_(module.weight, mean=0.0, std=0.05) @@ -386,40 +441,28 @@ def forward( x_one_hot: torch.Tensor, ) -> torch.Tensor: """ - Forward pass through the DeepLC model. - - Parameters - ---------- - x_atom - Atomic composition features [batch, sequence_length, channels] - x_atom_sum - Summed atomic composition features [batch, sequence_length, channels] - x_global - Global peptide features [batch, feature_size] - x_one_hot - One-hot encoded amino acid features [batch, sequence_length, channels] + Forward pass. Returns ------- torch.Tensor - Predicted retention times [batch, 1] + Shape ``[batch, n_heads]``, or ``[batch, 1]`` when an adapter is attached. """ - # Transpose to Conv1D format: (batch, channels, length) x_atom = x_atom.transpose(1, 2) x_atom_sum = x_atom_sum.transpose(1, 2) x_one_hot = x_one_hot.transpose(1, 2) - - # Process each branch - out_a = self.branch_a(x_atom) - out_b = self.branch_b(x_atom_sum) - out_c = self.branch_c(x_global) - out_d = self.branch_d(x_one_hot) - - # Concatenate features - concatenated = torch.cat([out_a, out_b, out_c, out_d], dim=1) - - # Final prediction - output = self.final_network(concatenated) - - return output + concatenated = torch.cat( + [ + self.branch_a(x_atom), + self.branch_b(x_atom_sum), + self.branch_c(x_global), + self.branch_d(x_one_hot), + ], + dim=1, + ) + out = self.heads(self.shared_trunk(concatenated)) # [batch, n_heads] + adapter = getattr(self, "adapter", None) + if adapter is not None: + return adapter(out) # [batch, 1] + return out diff --git a/deeplc/_model_ops.py b/deeplc/_model_ops.py index b3245cc..2f5d777 100644 --- a/deeplc/_model_ops.py +++ b/deeplc/_model_ops.py @@ -26,32 +26,42 @@ def load_model( model: torch.nn.Module | PathLike | str | None = None, device: str | None = None, -) -> torch.nn.Module: +) -> DeepLCModel: """Load a model from a file or return a randomly initialized model if none is provided.""" # If device is not specified, use the default device (GPU if available, else CPU) - selected_device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") + selected_device: torch.device | str = device or torch.device( + "cuda" if torch.cuda.is_available() else "cpu" + ) - # Load model from file if a path is provided if isinstance(model, (str, PathLike, Path)): - loaded_model = torch.load(model, weights_only=False, map_location=selected_device) - elif isinstance(model, torch.nn.Module): + # Infer architecture hyperparameters from the saved state dict + # Only checks n_heads and final_num_layers; other hyperparameters are set to defaults + # May break for models saved with different architectures. + raw = torch.load(model, weights_only=True, map_location=selected_device) + n_heads = raw["heads.b2"].shape[0] + final_num_layers = sum( + 1 for k in raw if k.startswith("shared_trunk.") and k.endswith(".weight") + ) + loaded_model = DeepLCModel(n_heads=n_heads, final_num_layers=final_num_layers) + if "adapter.0.weight" in raw: + loaded_model.add_adapter(hidden_size=raw["adapter.0.weight"].shape[0]) + loaded_model.load_state_dict(raw) + elif isinstance(model, DeepLCModel): loaded_model = model logger.debug("Using provided PyTorch model instance") elif model is None: - # Initialize a new model with default architecture - loaded_model = DeepLCModel() + loaded_model = DeepLCModel(n_heads=1) logger.debug("Initialized new DeepLCModel with default architecture") else: - raise TypeError(f"Expected a PyTorch Module or a file path, got {type(model)} instead.") + raise TypeError(f"Expected a DeepLCModel or a file path, got {type(model)} instead.") - # Ensure the model is on the specified device loaded_model.to(selected_device) return loaded_model def train( - model: torch.nn.Module | PathLike | str | None, + model: DeepLCModel | PathLike | str | None, train_dataset: DeepLCDataset | Subset[DeepLCDataset], validation_dataset: DeepLCDataset | Subset[DeepLCDataset], device: str | None = None, @@ -61,6 +71,8 @@ def train( epochs: int = 25, batch_size: int = 512, patience: int = 10, + freeze_epochs: int = 0, + unfreeze_lr_scale: float = 0.1, show_progress: bool = True, ) -> torch.nn.Module: """ @@ -88,6 +100,10 @@ def train( Batch size for training and validation. patience Number of epochs with no improvement before early stopping. + freeze_epochs + Number of initial epochs to train with backbone frozen (adapter only). + unfreeze_lr_scale + Learning rate multiplier applied after unfreezing the backbone. show_progress If True, display a Rich progress bar during training. If False, run silently. @@ -119,6 +135,9 @@ def train( "Validation data loader is empty. Adjust validation data or validation_split." ) + has_freeze = hasattr(model, "freeze_backbone") and hasattr(model, "unfreeze_backbone") + if has_freeze and freeze_epochs > 0: + model.freeze_backbone() optimizer = _get_optimizer(model, learning_rate) loss_fn = torch.nn.L1Loss() @@ -129,7 +148,11 @@ def train( with _create_progress(disable=not show_progress) as progress: epoch_task = progress.add_task("Epochs", total=epochs, status="") - for _epoch in range(epochs): + for epoch in range(epochs): + if has_freeze and freeze_epochs > 0 and epoch == freeze_epochs: + model.unfreeze_backbone() + optimizer = _get_optimizer(model, learning_rate * unfreeze_lr_scale) + avg_loss = _train_epoch(model, train_loader, optimizer, loss_fn, device) avg_val_loss = _validate_epoch(model, val_loader, loss_fn, device) @@ -254,8 +277,8 @@ def _predict_epoch( outputs = model(*features) predictions.append(outputs.cpu()) if not predictions: - return torch.empty(0, dtype=torch.float32) - return torch.cat(predictions, dim=0).squeeze() + raise ValueError("Dataset is empty — nothing to predict.") + return torch.cat(predictions, dim=0) def _create_progress(disable: bool = False) -> Progress: diff --git a/deeplc/_reference_selection.py b/deeplc/_reference_selection.py index 4f6bd87..db6e995 100644 --- a/deeplc/_reference_selection.py +++ b/deeplc/_reference_selection.py @@ -29,8 +29,8 @@ def select_reference_psms(psm_list: PSMList) -> PSMList: Parameters ---------- psm_list - PSMList to select reference PSMs from. Must contain PSMs with observed - retention times and either q-values, scores + target-decoy labels, or scores. + PSMList to select reference PSMs from. Must contain PSMs with observed retention + times and either q-values, scores + target-decoy labels, or scores. Returns ------- @@ -106,7 +106,8 @@ def _select_by_qvalue(candidates: PSMList) -> PSMList: def _select_by_computed_qvalue(psm_list: PSMList) -> PSMList: - """Compute q-values from target-decoy competition, then select by threshold. + """ + Compute q-values from target-decoy competition, then select by threshold. Note: this modifies ``psm_list`` in-place by assigning computed q-values. """ diff --git a/deeplc/calibration.py b/deeplc/calibration.py index b89137c..4e451f0 100644 --- a/deeplc/calibration.py +++ b/deeplc/calibration.py @@ -19,8 +19,11 @@ class Calibration(ABC): """Abstract base class for calibration.""" + selected_model_head: int | None = None + @abstractmethod def __init__(self, *args, **kwargs): + """Initialize the calibration model.""" super().__init__() @property @@ -45,16 +48,21 @@ class IdentityCalibration(Calibration): @property def is_fitted(self) -> bool: + """Always fitted; identity calibration requires no fitting.""" return True def fit(self, target: np.ndarray, source: np.ndarray) -> None: # noqa: ARG002 + """No-op; identity calibration does not fit.""" return None def transform(self, source: np.ndarray) -> np.ndarray: + """Return source unchanged.""" return source class PiecewiseLinearCalibration(Calibration): + """Piece-wise linear calibration based on per-split anchors.""" + def __init__( self, number_of_splits: int = 10, @@ -79,6 +87,7 @@ def __init__( Minimum number of samples required for a segment to contribute an anchor. Segments with fewer samples are skipped, which helps avoid unstable anchors in sparse regions when using many splits. + """ super().__init__() self.number_of_splits = int(number_of_splits) @@ -97,6 +106,7 @@ def __init__( @property def is_fitted(self) -> bool: + """True if the calibration model has been fitted.""" return ( self._calibrate_min is not None and self._calibrate_max is not None @@ -107,10 +117,12 @@ def is_fitted(self) -> bool: @property def calibrate_min(self) -> float | None: + """Minimum source value seen during fitting.""" return self._calibrate_min @property def calibrate_max(self) -> float | None: + """Maximum source value seen during fitting.""" return self._calibrate_max def fit(self, target: np.ndarray, source: np.ndarray) -> None: @@ -228,7 +240,10 @@ def get_calibration_curve(self) -> tuple[np.ndarray, np.ndarray]: class SplineTransformerCalibration(Calibration): + """Spline-based calibration using sklearn's SplineTransformer.""" + def __init__(self) -> None: + """Initialize SplineTransformerCalibration.""" super().__init__() self._calibrate_min: float | None = None self._calibrate_max: float | None = None @@ -238,6 +253,7 @@ def __init__(self) -> None: @property def is_fitted(self) -> bool: + """True if the calibration model has been fitted.""" return ( self._calibrate_min is not None and self._calibrate_max is not None @@ -287,6 +303,8 @@ def transform(self, source: np.ndarray) -> np.ndarray: model_main = cast(Pipeline | LinearRegression, self._model_main) model_left = cast(LinearRegression, self._model_left) model_right = cast(LinearRegression, self._model_right) + calibrate_min = cast(float, self._calibrate_min) + calibrate_max = cast(float, self._calibrate_max) if source.shape[0] == 0: return np.array([]) @@ -294,15 +312,15 @@ def transform(self, source: np.ndarray) -> np.ndarray: y_pred_spline = model_main.predict(source.reshape(-1, 1)) y_pred_left = model_left.predict(source.reshape(-1, 1)) y_pred_right = model_right.predict(source.reshape(-1, 1)) - within_range = (source >= self._calibrate_min) & (source <= self._calibrate_max) + within_range = (source >= calibrate_min) & (source <= calibrate_max) within_range = within_range.ravel() cal_preds = np.copy(y_pred_spline) - cal_preds[~within_range & (source.ravel() < self._calibrate_min)] = y_pred_left[ - ~within_range & (source.ravel() < self._calibrate_min) + cal_preds[~within_range & (source.ravel() < calibrate_min)] = y_pred_left[ + ~within_range & (source.ravel() < calibrate_min) ] - cal_preds[~within_range & (source.ravel() > self._calibrate_max)] = y_pred_right[ - ~within_range & (source.ravel() > self._calibrate_max) + cal_preds[~within_range & (source.ravel() > calibrate_max)] = y_pred_right[ + ~within_range & (source.ravel() > calibrate_max) ] return np.array(cal_preds) diff --git a/deeplc/core.py b/deeplc/core.py index 319f4e0..a4b440c 100644 --- a/deeplc/core.py +++ b/deeplc/core.py @@ -21,14 +21,14 @@ LOGGER = logging.getLogger(__name__) DEEPLC_DIR = Path(__file__).resolve().parent -DEFAULT_MODEL_NAME = "full_hc_PXD005573_pub_1fd8363d9af9dcad3be7553c39396960.pt" -DEFAULT_MODEL = DEEPLC_DIR / "package_data" / "models" / DEFAULT_MODEL_NAME +DEFAULT_MODEL = DEEPLC_DIR / "package_data" / "models" / "multitask_model.pt" def predict( psm_list: PSMList | list[PSM | Peptidoform | str], model: torch.nn.Module | PathLike | str | None = None, predict_kwargs: dict | None = None, + return_matrix: bool = False, ) -> np.ndarray: """ Predict retention times for a list of PSMs using a trained model. @@ -41,18 +41,26 @@ def predict( Trained model or path to model file. If None, the default DeepLC model is used. predict_kwargs Additional keyword arguments to pass to the prediction function. + return_matrix + If True, return the full prediction matrix of shape ``(n, n_heads)`` when using a + multitask model. If False (default), return a 1D array of shape ``(n,)`` using + head 0 when model output is 2D. Returns ------- np.ndarray - Retention time predictions. + Retention time predictions. Shape ``(n,)`` unless ``return_matrix=True`` and model + produces multitask output, in which case shape is ``(n, n_heads)``. """ - return _model_ops.predict( + result = _model_ops.predict( model=model or DEFAULT_MODEL, data=DeepLCDataset.from_psm_list(_parse_psms(psm_list)), **(predict_kwargs or {}), ).numpy() + if not return_matrix: + return result[:, 0] + return result def calibrate( @@ -106,11 +114,18 @@ def calibrate( psm_list=psm_list_reference, model=model, predict_kwargs=predict_kwargs, + return_matrix=True, ) # Fit calibration LOGGER.debug("Fitting calibration...") target_rt_cal = np.array(psm_list_reference["retention_time"], dtype=np.float32) + + # Select the best head for calibration if the model predicts for multiple LC setups + if source_rt_cal.shape[1] > 1: + calibration.selected_model_head = _best_correlating_head(source_rt_cal, target_rt_cal) + source_rt_cal = source_rt_cal[:, calibration.selected_model_head or 0] + calibration.fit(target=target_rt_cal, source=source_rt_cal) return calibration @@ -161,6 +176,7 @@ def predict_and_calibrate( psm_list=parsed_psm_list, model=model, predict_kwargs=predict_kwargs, + return_matrix=True, ) if calibration is not None and not isinstance(calibration, Calibration): @@ -179,6 +195,17 @@ def predict_and_calibrate( else: LOGGER.info("Calibration is already fitted, skipping fitting step.") + if predicted_rt.shape[1] > 1: + if calibration.selected_model_head is None: + raise ValueError( + "Calibration has no selected_model_head. Either use calibrate() to fit it, " + "or set calibration.selected_model_head manually before calling " + "predict_and_calibrate() with a multitask model." + ) + predicted_rt = predicted_rt[:, calibration.selected_model_head] + else: + predicted_rt = predicted_rt[:, 0] + # Apply calibration to predictions calibrated_rt = calibration.transform(predicted_rt) @@ -270,6 +297,9 @@ def finetune( psm_list_validation List of PSMs to use for validation during fine-tuning. If None, a split from psm_list is used. + validation_split + Fraction of ``psm_list_reference`` to use for validation when ``psm_list_validation`` + is None. model Trained model or path to model file. train_kwargs @@ -292,11 +322,22 @@ def finetune( training_dataset, validation_dataset = split_datasets( training_data, validation_data=validation_data, validation_split=validation_split ) + train_kwargs_local = dict(train_kwargs or {}) + adapter_hidden_size = int(train_kwargs_local.pop("adapter_hidden_size", 256)) + freeze_epochs = int(train_kwargs_local.pop("freeze_epochs", 5)) + + loaded_model = _model_ops.load_model( + model or DEFAULT_MODEL, + device=train_kwargs_local.get("device"), + ) + loaded_model.add_adapter(hidden_size=adapter_hidden_size) + train_kwargs_local["freeze_epochs"] = freeze_epochs + finetuned_model = _model_ops.train( - model=model or DEFAULT_MODEL, + model=loaded_model, train_dataset=training_dataset, validation_dataset=validation_dataset, - **(train_kwargs or {}), + **train_kwargs_local, ) return finetuned_model @@ -344,6 +385,23 @@ def train( return trained_model +def save_model(model: torch.nn.Module, path: PathLike | str) -> None: + """ + Save a model's state dict to a file. + + Use :func:`load_model` (via :func:`predict`) to reload the saved checkpoint. + + Parameters + ---------- + model + Trained model instance to save. + path + Destination file path. + + """ + torch.save(model.state_dict(), path) + + def _parse_psms(psm_list: PSMList | list[PSM | Peptidoform | str]) -> PSMList: """ Parse a list of PSMs, Peptidoforms, or strings into a PSMList. @@ -368,3 +426,25 @@ def _parse_psms(psm_list: PSMList | list[PSM | Peptidoform | str]) -> PSMList: raise ValueError("List must contain either PSMs, Peptidoforms, or strings.") else: raise ValueError("Input must be a PSMList or a list of PSMs, Peptidoforms, or strings.") + + +def _best_correlating_head(predictions: np.ndarray, targets: np.ndarray) -> int: + """Return the head index with highest valid Pearson correlation to targets.""" + best_idx = 0 + best_corr = float("-inf") + + for idx in range(predictions.shape[1]): + pred_col = predictions[:, idx] + mask = np.isfinite(pred_col) & np.isfinite(targets) + if mask.sum() < 3: + continue + pred_masked = pred_col[mask] + target_masked = targets[mask] + if np.std(pred_masked) < 1e-8 or np.std(target_masked) < 1e-8: + continue + corr = np.corrcoef(pred_masked, target_masked)[0, 1] + if np.isfinite(corr) and corr > best_corr: + best_corr = corr + best_idx = idx + + return best_idx diff --git a/deeplc/data.py b/deeplc/data.py index a2807ce..3fff843 100644 --- a/deeplc/data.py +++ b/deeplc/data.py @@ -1,3 +1,5 @@ +"""Dataset classes and utilities for DeepLC.""" + from __future__ import annotations import logging @@ -56,10 +58,12 @@ def __init__( f"match length of peptidoforms ({len(self.peptidoforms)})" ) - def __len__(self): + def __len__(self) -> int: + """Return number of peptidoforms in the dataset.""" return len(self.peptidoforms) - def __getitem__(self, idx) -> tuple: + def __getitem__(self, idx: int) -> tuple[torch.Tensor, ...]: + """Return encoded features and target RT for peptidoform at index.""" if not isinstance(idx, int): raise TypeError(f"Index must be an integer, got {type(idx)} instead.") features = encode_peptidoform( diff --git a/deeplc/exceptions.py b/deeplc/exceptions.py index 9b81629..1289fed 100644 --- a/deeplc/exceptions.py +++ b/deeplc/exceptions.py @@ -2,12 +2,12 @@ class DeepLCError(Exception): - pass + """Base class for DeepLC exceptions.""" class CalibrationError(DeepLCError): - pass + """Raised when calibration fails.""" class ReferenceSelectionError(DeepLCError): - pass + """Raised when reference PSM selection fails.""" diff --git a/deeplc/gui.py b/deeplc/gui.py index c0c3b1d..b4270b7 100644 --- a/deeplc/gui.py +++ b/deeplc/gui.py @@ -33,7 +33,6 @@ def create_app(): """Create and configure the NiceGUI app.""" - _logo_ref = importlib.resources.files("deeplc.package_data.gui_images").joinpath( "deeplc_logo.svg" ) diff --git a/deeplc/package_data/__init__.py b/deeplc/package_data/__init__.py index e69de29..90c6b89 100644 --- a/deeplc/package_data/__init__.py +++ b/deeplc/package_data/__init__.py @@ -0,0 +1 @@ +"""DeepLC package data.""" diff --git a/deeplc/package_data/baseline_performance/__init__.py b/deeplc/package_data/baseline_performance/__init__.py index e69de29..76db23c 100644 --- a/deeplc/package_data/baseline_performance/__init__.py +++ b/deeplc/package_data/baseline_performance/__init__.py @@ -0,0 +1 @@ +"""Baseline performance data for DeepLC.""" diff --git a/deeplc/package_data/gui_images/__init__.py b/deeplc/package_data/gui_images/__init__.py index e69de29..cb0894e 100644 --- a/deeplc/package_data/gui_images/__init__.py +++ b/deeplc/package_data/gui_images/__init__.py @@ -0,0 +1 @@ +"""GUI image assets for DeepLC.""" diff --git a/deeplc/package_data/models/__init__.py b/deeplc/package_data/models/__init__.py index e69de29..ad935b8 100644 --- a/deeplc/package_data/models/__init__.py +++ b/deeplc/package_data/models/__init__.py @@ -0,0 +1 @@ +"""Pre-trained model files for DeepLC.""" diff --git a/deeplc/package_data/models/multitask_model.pt b/deeplc/package_data/models/multitask_model.pt new file mode 100644 index 0000000..5ba394b Binary files /dev/null and b/deeplc/package_data/models/multitask_model.pt differ diff --git a/deeplc/plot.py b/deeplc/plot.py index b4e7272..eb22974 100644 --- a/deeplc/plot.py +++ b/deeplc/plot.py @@ -1,3 +1,5 @@ +"""Plotting utilities for DeepLC results.""" + # TODO: Use package data; potentially move to MS²Rescore? from pathlib import Path @@ -70,8 +72,7 @@ def distribution_baseline( observed_column: str = "Observed retention time", ) -> go.Figure: """ - Plot a distribution plot of the relative mean absolute error of the current - DeepLC performance compared to the baseline performance. + Plot relative MAE distribution comparing current DeepLC performance to baseline. Parameters ---------- diff --git a/pyproject.toml b/pyproject.toml index f3fc1a0..4732c76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ authors = [ { name = "Robbin Bouwmeester", email = "robbin.bouwmeester@ugent.be" }, { name = "Ralf Gabriels" }, { name = "Kevin Velghe" }, - { name = "Robbe Devreese"}, + { name = "Robbe Devreese" }, { name = "Arthur Declercq" }, { name = "Alireza Nameni" }, { name = "Ralf Gabriels" }, @@ -57,11 +57,7 @@ PyPi = "https://pypi.org/project/deeplc/" CompOmics = "https://www.compomics.com" [dependency-groups] -dev = [ - "ruff", - "pytest", - "pyinstaller>=6.21.0", -] +dev = ["ruff", "pytest", "pyinstaller>=6.21.0"] docs = [ "numpydoc>=1,<2", "sphinx_rtd_theme", @@ -85,7 +81,8 @@ line-length = 99 target-version = 'py311' [tool.ruff.lint] -select = ["E", "W", "F", "UP", "B", "SIM", "I"] +select = ["E", "W", "F", "UP", "B", "SIM", "I", "D"] +ignore = ["D203", "D212"] [tool.ruff.format] docstring-code-format = true diff --git a/tests/test_core.py b/tests/test_core.py index 578c870..51994a0 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -64,6 +64,13 @@ def test_predict_and_calibrate_returns_calibrated_values_differ_from_raw(): assert not np.allclose(raw, calibrated) +def test_predict_returns_matrix_when_flag_set(): + result = deeplc.core.predict(_make_psm_list(_PEPTIDES), return_matrix=True) + assert result.ndim == 2 + assert result.shape[0] == len(_PEPTIDES) + assert result.shape[1] > 1 + + def test_predict_and_calibrate_auto_selects_reference(): # 200 PSMs cycling through _PEPTIDES; 100 with qvalue<=0.01, 100 with qvalue=1.0. # auto-selection picks the 100 low-qvalue PSMs as reference. diff --git a/tests/test_model_ops.py b/tests/test_model_ops.py index 38306a0..0480eae 100644 --- a/tests/test_model_ops.py +++ b/tests/test_model_ops.py @@ -1,10 +1,15 @@ from __future__ import annotations +import sys + import pytest import torch -from torch.utils.data import Dataset from deeplc import _model_ops +from torch.utils.data import Dataset + +from deeplc._architecture import DeepLCModel +from deeplc.core import DEFAULT_MODEL from deeplc.data import split_datasets @@ -26,17 +31,10 @@ def __getitem__(self, index: int): return features, target -class _DummyModel(torch.nn.Module): - def forward(self, matrix, matrix_sum, matrix_global, matrix_hc): # noqa: ARG002 - batch_size = matrix.shape[0] - return torch.zeros((batch_size, 1), dtype=torch.float32) - - -def test_predict_returns_empty_tensor_for_empty_dataset(): +def test_predict_raises_for_empty_dataset(): empty_data = _TinyDeepLCDataset(length=0) - preds = _model_ops.predict(model=_DummyModel(), data=empty_data, show_progress=False) - assert isinstance(preds, torch.Tensor) - assert preds.numel() == 0 + with pytest.raises(ValueError, match="empty"): + _model_ops.predict(model=DeepLCModel(n_heads=1), data=empty_data, show_progress=False) def test_split_datasets_rejects_too_small_dataset_without_validation_data(): @@ -51,10 +49,35 @@ def test_split_datasets_rejects_too_small_dataset_without_validation_data(): def test_train_rejects_empty_validation_loader(): with pytest.raises(ValueError, match="Validation data loader is empty"): _model_ops.train( - model=_DummyModel(), + model=DeepLCModel(n_heads=1), train_dataset=_TinyDeepLCDataset(length=2), validation_dataset=_TinyDeepLCDataset(length=0), epochs=1, batch_size=2, show_progress=False, ) + + +@pytest.mark.skipif( + not DEFAULT_MODEL.exists(), + reason="multitask model not bundled", +) +def test_load_multitask_model_without_prior_shim(): + """multitask_model.pt must load even when the legacy module is not pre-registered.""" + # Remove any previously registered shim so the test is self-contained. + sys.modules.pop("multitask_model", None) + + model = _model_ops.load_model(DEFAULT_MODEL, device="cpu") + + assert isinstance(model, DeepLCModel) + + x_atom = torch.zeros(2, 60, 6) + x_sum = torch.zeros(2, 30, 6) + x_global = torch.zeros(2, 55) + x_hc = torch.zeros(2, 60, 20) + with torch.no_grad(): + out = model(x_atom, x_sum, x_global, x_hc) + + assert out.ndim == 2 + assert out.shape[0] == 2 + assert out.shape[1] > 1 # multiple heads