diff --git a/docs/drevalpy.models.PaccMann.rst b/docs/drevalpy.models.PaccMann.rst new file mode 100644 index 00000000..43d64cac --- /dev/null +++ b/docs/drevalpy.models.PaccMann.rst @@ -0,0 +1,18 @@ +PaccMann +============================= + +PaccMann Model +---------------------------------- + +.. automodule:: drevalpy.models.PaccMann.paccmann + :members: + :undoc-members: + :show-inheritance: + +PaccMann Network +---------------------------------- + +.. automodule:: drevalpy.models.PaccMann.network + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/drevalpy.models.rst b/docs/drevalpy.models.rst index 3fda9561..fb72b92f 100644 --- a/docs/drevalpy.models.rst +++ b/docs/drevalpy.models.rst @@ -27,6 +27,7 @@ Implemented models drevalpy.models.DIPK drevalpy.models.DrugGNN drevalpy.models.MOLIR + drevalpy.models.PaccMann drevalpy.models.PharmaFormer drevalpy.models.Precily drevalpy.models.SRMF diff --git a/docs/installation.rst b/docs/installation.rst index 5fc8ba93..0ff5dac2 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -65,6 +65,10 @@ default ``pip install drevalpy``. They are provided as optional `extras`: * - ``xgboost`` - The ``MultiViewXGBoost`` baseline model - ``xgboost`` + * - ``paccmann`` + - SMILES augmentation for the ``PaccMann`` model. Without it, PaccMann trains on the + unaugmented SMILES and warns. + - ``rdkit`` * - ``multiprocessing`` - Parallelized cross-validation / tuning via Ray - ``ray`` (and ``pydantic``, usually already present) diff --git a/docs/usage.rst b/docs/usage.rst index f19dc8f7..9c96c158 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -310,6 +310,8 @@ See the sklearn model :ref:`flexible-inputs` or the SimpleNeuralNetwork :ref:`fl +---------------------------------+----------------------------+--------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | Precily | Published Model | Multi-Drug Model | `Precily `_ from Chawla et al. Uses GSVA pathway-activity scores with SMILESVec drug embeddings. Features are concatenated and passed through multiple linear layers with ReLU and Dropout. | +---------------------------------+----------------------------+--------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| PaccMann | Published Model | Multi-Drug Model | `PaccMann `_ from Manica et al. Embeds tokenized drug SMILES and encodes them with multi-scale convolutional layers, while cell line gene expression of a curated gene panel serves as biological context. Contextual attention layers connect the gene and molecule representations, which are concatenated and passed through stacked dense layers to predict the response. | ++---------------------------------+----------------------------+--------------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ Available Datasets diff --git a/drevalpy/models/PaccMann/__init__.py b/drevalpy/models/PaccMann/__init__.py new file mode 100644 index 00000000..524763c2 --- /dev/null +++ b/drevalpy/models/PaccMann/__init__.py @@ -0,0 +1,5 @@ +"""Module for the Paccmann model.""" + +from .paccmann import PaccMann + +__all__ = ["PaccMann"] diff --git a/drevalpy/models/PaccMann/hyperparameters.yaml b/drevalpy/models/PaccMann/hyperparameters.yaml new file mode 100644 index 00000000..02b05048 --- /dev/null +++ b/drevalpy/models/PaccMann/hyperparameters.yaml @@ -0,0 +1,41 @@ +PaccMann: + gene_list: + - gene_list_paccmann_network_prop + epochs: + - 10 + patience: + - 10 + batch_size: + - 64 + learning_rate: + - 0.001 + weight_decay: + - 0.0 + smiles_embedding_size: + - 8 + filters: + - [16, 16, 16] + molecule_heads: + - [2, 2, 2, 2] + gene_heads: + - [2, 2, 2, 2] + smiles_padding_length: + - 512 + # Train each drug on several equivalent SMILES strings. Requires rdkit; without it + # training falls back to the unaugmented SMILES and warns. + augment_smiles: + - true + dropout: + - 0.5 + batch_norm: + - true + smiles_attention_size: + - 64 + gene_attention_size: + - 1 + molecule_temperature: + - 1.0 + gene_temperature: + - 1.0 + stacked_dense_hidden_sizes: + - [512, 256] diff --git a/drevalpy/models/PaccMann/network.py b/drevalpy/models/PaccMann/network.py new file mode 100644 index 00000000..c7659a59 --- /dev/null +++ b/drevalpy/models/PaccMann/network.py @@ -0,0 +1,270 @@ +"""PaccMann network: context attention between SMILES convolutions and gene expression. + +Reimplementation of the MCA architecture from Manica et al., Molecular Pharmaceutics 2019 +(https://pubs.acs.org/doi/10.1021/acs.molpharmaceut.9b00520), as released at +https://github.com/PaccMann/paccmann_predictor. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +from torch import nn + + +@dataclass +class PaccMannConfig: + """Resolved, typed configuration for PaccMannNetwork. + + smiles_vocabulary_size, smiles_padding_length, and number_of_genes are derived from the training data; + every other field is a model hyperparameter, see PaccMann/hyperparameters.yaml for their defaults. + """ + + smiles_vocabulary_size: int + smiles_embedding_size: int + smiles_padding_length: int + number_of_genes: int + molecule_heads: list[int] = field(default_factory=lambda: [4, 4, 4, 4]) + gene_heads: list[int] = field(default_factory=lambda: [2, 2, 2, 2]) + filters: list[int] = field(default_factory=lambda: [64, 64, 64]) + kernel_sizes: list[tuple[int, int]] | None = None + dropout: float = 0.5 + batch_norm: bool = False + smiles_attention_size: int = 64 + gene_attention_size: int = 1 + molecule_temperature: float = 1.0 + gene_temperature: float = 1.0 + stacked_dense_hidden_sizes: list[int] = field(default_factory=lambda: [1024, 512]) + + @property + def resolved_kernel_sizes(self) -> list[tuple[int, int]]: + """Kernel sizes, defaulting to token windows of 3, 5, and 11 sized to the embedding dimension. + + :return: one (token window, embedding dimension) pair per convolution + """ + if self.kernel_sizes is not None: + return self.kernel_sizes + return [(window, self.smiles_embedding_size) for window in (3, 5, 11)] + + def __post_init__(self) -> None: + """Validate that head counts, filters, and kernel sizes line up. + + :raises ValueError: if the molecule/gene head counts, filters, or kernel sizes are inconsistent + """ + if len(self.gene_heads) != len(self.molecule_heads): + raise ValueError("gene_heads and molecule_heads must have the same length.") + if len(self.filters) != len(self.resolved_kernel_sizes): + raise ValueError("filters and kernel_sizes must have the same length.") + if len(self.filters) + 1 != len(self.molecule_heads): + raise ValueError("molecule_heads must have exactly one more entry than filters.") + + @classmethod + def from_hyperparameters(cls, params: dict[str, Any]) -> PaccMannConfig: + """Resolve a config from a raw hyperparameter dictionary. + + :param params: hyperparameters, see PaccMann/hyperparameters.yaml for the available keys, plus + smiles_vocabulary_size, smiles_padding_length, and number_of_genes, which are derived from the + training data rather than configured + :return: resolved config + """ + field_names = {f for f in cls.__dataclass_fields__ if f in params} + return cls(**{name: params[name] for name in field_names}) + + +class ContextAttentionLayer(nn.Module): + """Context attention layer: lets one modality attend over another (PaccMann paper, Fig. 2C).""" + + def __init__( + self, + reference_hidden_size: int, + reference_sequence_length: int, + context_hidden_size: int, + context_sequence_length: int, + attention_size: int, + temperature: float, + ) -> None: + """Initialize the context attention layer. + + :param reference_hidden_size: hidden size of the reference input + :param reference_sequence_length: sequence length of the reference input + :param context_hidden_size: hidden size (or feature count) of the context input + :param context_sequence_length: sequence length of the context input + :param attention_size: size of the shared attention space + :param temperature: softmax temperature; below 1 sharpens the attention, above 1 smooths it + """ + super().__init__() + self.reference_projection = nn.Linear(reference_hidden_size, attention_size) + self.context_projection = nn.Linear(context_hidden_size, attention_size) + self.context_hidden_projection = ( + nn.Linear(context_sequence_length, reference_sequence_length) + if context_sequence_length > 1 + else nn.Identity() + ) + self.alpha_projection = nn.Linear(attention_size, 1, bias=False) + self.temperature = temperature + + def forward( + self, reference: torch.Tensor, context: torch.Tensor, average_seq: bool = True + ) -> tuple[torch.Tensor, torch.Tensor]: + """Attend over ``reference`` using ``context``. + + :param reference: tensor of shape (batch, reference_sequence_length, reference_hidden_size) + :param context: tensor of shape (batch, context_sequence_length, context_hidden_size) + :param average_seq: sum the attended reference over its sequence dimension + :return: attended output and attention weights + """ + reference_attention = self.reference_projection(reference) + context_attention = self.context_hidden_projection(self.context_projection(context).permute(0, 2, 1)).permute( + 0, 2, 1 + ) + alphas = self.alpha_projection(torch.tanh(reference_attention + context_attention)).squeeze(-1) + alphas = torch.softmax(alphas / self.temperature, dim=1) + + output = reference * alphas.unsqueeze(-1) + # Squeeze only the trailing dimension: a batch of size 1 must keep its batch dimension. + return (output.sum(dim=1) if average_seq else output.squeeze(-1)), alphas + + +class _ConvolutionBlock(nn.Module): + """Convolution over embedded SMILES tokens, followed by activation, dropout, and batch norm.""" + + def __init__(self, num_kernel: int, kernel_size: tuple[int, int], dropout: float, batch_norm: bool) -> None: + """Initialize the convolution block. + + :param num_kernel: number of convolution kernels, i.e. output channels + :param kernel_size: (token window, embedding dimension) size of the convolution kernel + :param dropout: dropout probability + :param batch_norm: whether to apply batch normalization + """ + super().__init__() + self.convolve = nn.Conv2d(1, num_kernel, kernel_size, padding=(kernel_size[0] // 2, 0)) + self.dropout = nn.Dropout(dropout) + self.batch_norm = nn.BatchNorm1d(num_kernel) if batch_norm else nn.Identity() + + def forward(self, embedded_smiles: torch.Tensor) -> torch.Tensor: + """Convolve embedded SMILES tokens. + + :param embedded_smiles: tensor of shape (batch, 1, smiles_padding_length, embedding_size) + :return: tensor of shape (batch, num_kernel, smiles_padding_length) + """ + activated = torch.relu(self.convolve(embedded_smiles).squeeze(-1)) + return self.batch_norm(self.dropout(activated)) + + +def _dense_block(input_size: int, hidden_size: int, dropout: float, batch_norm: bool) -> nn.Sequential: + """Build a linear layer followed by batch norm, ReLU, and dropout. + + :param input_size: input feature size + :param hidden_size: output feature size + :param dropout: dropout probability + :param batch_norm: whether to apply batch normalization + :return: sequential dense block + """ + return nn.Sequential( + nn.Linear(input_size, hidden_size), + nn.BatchNorm1d(hidden_size) if batch_norm else nn.Identity(), + nn.ReLU(), + nn.Dropout(dropout), + ) + + +class PaccMannNetwork(nn.Module): + """PaccMann drug-response network. + + SMILES tokens are embedded and convolved at several kernel sizes; gene expression attends over each + resulting representation and vice versa; the concatenated attention outputs are passed through dense + layers down to a single drug-response prediction. + """ + + def __init__(self, config: PaccMannConfig) -> None: + """Build the network from a resolved config. + + :param config: resolved network configuration + """ + super().__init__() + + self.smiles_embedding = nn.Embedding(config.smiles_vocabulary_size, config.smiles_embedding_size) + self.convolutions = nn.ModuleList( + [ + _ConvolutionBlock(num_kernel, kernel_size, config.dropout, config.batch_norm) + for num_kernel, kernel_size in zip(config.filters, config.resolved_kernel_sizes) + ] + ) + + # Flat lists of attention heads, grouped back into per-layer chunks in forward() via self._molecule_heads + # and self._gene_heads: nn.ModuleList cannot be nested and stay iterable under mypy's torch stubs. + smiles_hidden_sizes = [config.smiles_embedding_size] + config.filters + self._molecule_heads = config.molecule_heads + self.molecule_attentions = nn.ModuleList( + [ + ContextAttentionLayer( + reference_hidden_size=smiles_hidden_sizes[layer], + reference_sequence_length=config.smiles_padding_length, + context_hidden_size=1, + context_sequence_length=config.number_of_genes, + attention_size=config.smiles_attention_size, + temperature=config.molecule_temperature, + ) + for layer, heads in enumerate(config.molecule_heads) + for _ in range(heads) + ] + ) + self._gene_heads = config.gene_heads + self.gene_attentions = nn.ModuleList( + [ + ContextAttentionLayer( + reference_hidden_size=1, + reference_sequence_length=config.number_of_genes, + context_hidden_size=smiles_hidden_sizes[layer], + context_sequence_length=config.smiles_padding_length, + attention_size=config.gene_attention_size, + temperature=config.gene_temperature, + ) + for layer, heads in enumerate(config.gene_heads) + for _ in range(heads) + ] + ) + + attention_output_size = ( + config.molecule_heads[0] * config.smiles_embedding_size + + sum(heads * num_filters for heads, num_filters in zip(config.molecule_heads[1:], config.filters)) + + sum(config.gene_heads) * config.number_of_genes + ) + hidden_sizes = [attention_output_size, *config.stacked_dense_hidden_sizes] + + self.batch_norm = nn.BatchNorm1d(hidden_sizes[0]) if config.batch_norm else nn.Identity() + self.dense_layers = nn.ModuleList( + [ + _dense_block(hidden_sizes[i], hidden_sizes[i + 1], config.dropout, config.batch_norm) + for i in range(len(hidden_sizes) - 1) + ] + ) + self.output_layer = nn.Linear(hidden_sizes[-1], 1) + + def forward(self, smiles: torch.Tensor, gene_expression: torch.Tensor) -> torch.Tensor: + """Predict a drug-response score from tokenized SMILES and gene expression. + + :param smiles: token ids of shape (batch, smiles_padding_length) + :param gene_expression: gene expression of shape (batch, number_of_genes) + :return: predicted response of shape (batch, 1) + """ + gene_expression = gene_expression.unsqueeze(-1) + embedded_smiles = self.smiles_embedding(smiles) + smiles_encodings = [embedded_smiles] + [ + conv(embedded_smiles.unsqueeze(1)).permute(0, 2, 1) for conv in self.convolutions + ] + + attended = [] + molecule_attentions = iter(self.molecule_attentions) + for heads, encoding in zip(self._molecule_heads, smiles_encodings): + attended += [next(molecule_attentions)(encoding, gene_expression)[0] for _ in range(heads)] + gene_attentions = iter(self.gene_attentions) + for heads, encoding in zip(self._gene_heads, smiles_encodings): + attended += [next(gene_attentions)(gene_expression, encoding, average_seq=False)[0] for _ in range(heads)] + + hidden = self.batch_norm(torch.cat(attended, dim=1)) + for dense in self.dense_layers: + hidden = dense(hidden) + return self.output_layer(hidden) diff --git a/drevalpy/models/PaccMann/paccmann.py b/drevalpy/models/PaccMann/paccmann.py new file mode 100644 index 00000000..4ff4e64e --- /dev/null +++ b/drevalpy/models/PaccMann/paccmann.py @@ -0,0 +1,629 @@ +"""PaccMann model.""" + +from __future__ import annotations + +import copy +import json +import os +import re +import warnings +from typing import Any + +import joblib +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as functional +from sklearn.preprocessing import StandardScaler +from torch.utils.data import DataLoader, TensorDataset + +from drevalpy.datasets.dataset import DrugResponseDataset, FeatureDataset +from drevalpy.models.drp_model import DRPModel +from drevalpy.models.utils import load_and_select_gene_features + +from .network import PaccMannConfig, PaccMannNetwork + +# Atom-level SMILES tokenizer, copied verbatim from pytoda.smiles.processing.SMILES_TOKENIZER, which is what +# the original implementation tokenizes with. Splitting SMILES by character instead would break multi-character +# atoms: "Cl" and "Br" would collide with chlorine/bromine-free molecules that contain carbon or boron, and +# bracket atoms such as "[C@@H]" or "[Pt+2]" would fall apart into their individual characters. +SMILES_TOKENIZER = re.compile( + r"(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|" r"-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])" +) + + +def _tokenize_smiles(smiles: str) -> list[str]: + """Split a SMILES string into atom-level tokens. + + :param smiles: SMILES string + :return: list of tokens + """ + return [token for token in SMILES_TOKENIZER.split(smiles) if token] + + +def _smiles_column_to_list(smiles_matrix: np.ndarray) -> list[str]: + """Flatten the single-column SMILES feature matrix into a list of strings. + + :param smiles_matrix: SMILES feature matrix of shape (n_drugs, 1), as returned by load_drug_features + :return: list of SMILES strings + """ + return smiles_matrix.reshape(-1).astype(str).tolist() + + +PAD_IDX = 0 +UNK_IDX = 1 +BASE_SMILES_VOCAB = {"": PAD_IDX, "": UNK_IDX} + + +# Number of SMILES variants generated per drug. The original implementation re-randomizes on every access, which +# would mean an RDKit call per sample per epoch. Since a dataset holds only a few hundred distinct drugs, a fixed +# bank of variants is built once and sampled from during training instead. +N_SMILES_VARIANTS = 20 + +# Seed for building the variant bank, so that repeated runs on the same data produce the same variants +SMILES_AUGMENTATION_SEED = 42 + + +def _randomize_smiles(smiles_list: list[str], n_variants: int) -> list[list[str]]: + """Generate alternative SMILES strings for the same molecules. + + Every molecule is re-serialized from a shuffled atom order, which yields a different SMILES string that + describes the exact same molecule. This is the augmentation the original implementation applies to the drug + modality. Molecules that RDKit cannot parse keep their original SMILES. + + :param smiles_list: list of distinct SMILES strings + :param n_variants: number of variants to generate per molecule, including the original SMILES + :return: list holding the variants of each molecule, the original SMILES first + :raises ImportError: if RDKit is not installed + """ + try: + from rdkit import Chem, RDLogger + except ImportError as e: # pragma: no cover - depends on the environment + raise ImportError("Please install rdkit to augment SMILES for PaccMann: pip install rdkit") from e + + RDLogger.DisableLog("rdApp.*") # RDKit warns loudly about SMILES it can still parse + rng = np.random.default_rng(SMILES_AUGMENTATION_SEED) + + variants = [] + for smiles in smiles_list: + molecule = Chem.MolFromSmiles(smiles) + if molecule is None: + variants.append([smiles] * n_variants) + continue + + atom_order = list(range(molecule.GetNumAtoms())) + molecule_variants = [smiles] + for _ in range(n_variants - 1): + rng.shuffle(atom_order) + renumbered = Chem.RenumberAtoms(molecule, atom_order) + molecule_variants.append(Chem.MolToSmiles(renumbered, canonical=False)) + variants.append(molecule_variants) + + return variants + + +class PaccMann(DRPModel): + """PaccMann model for drug response prediction. + + This DrEval wrapper combines cell line gene expression features and tokenized SMILES representations of drugs + and uses the PaccMannNetwork to predict drug response values. + + This wrapper: + - loads gene expression features for cell lines + - loads SMILES strings for drugs + - tokenizes SMILES into padded integer sequences + - augments the drugs with equivalent SMILES strings, unless augment_smiles is disabled + - scales gene expression on training data only + - trains a PaccMannNetwork PyTorch model, stopping early once the early stopping set stops improving + - keeps the weights of the epoch with the lowest loss on the early stopping set + """ + + early_stopping = True + is_single_drug_model = False + + cell_line_views = ["gene_expression"] + drug_views = ["smiles"] + + def __init__(self) -> None: + """Initialize the PaccMann model wrapper. + + Initialized attributes: + model: stores the PaccMann neural network + hyperparameters: stores the passed hyperparameters + device: CPU or GPU device + gene_expression_scaler: scaler fitted on training gene expression + smiles_to_idx: SMILES vocabulary + smiles_padding_length: sequence length used for padding + number_of_genes: number of gene features + """ + super().__init__() + self.model: PaccMannNetwork | None = None + self.hyperparameters = {} + + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.gene_expression_scaler = StandardScaler() + + self.smiles_to_idx: dict[str, int] = dict(BASE_SMILES_VOCAB) + + self.smiles_padding_length: int | None = None + self.number_of_genes: int | None = None + + @classmethod + def get_model_name(cls) -> str: + """Return the model name. + + :return: model name + """ + return "PaccMann" + + def load_cell_line_features(self, data_path: str, dataset_name: str) -> FeatureDataset: + """Load gene expression features. + + :param data_path: path to the data directory + :param dataset_name: name of the dataset + :return: FeatureDataset containing gene expression features + """ + return load_and_select_gene_features( + feature_type="gene_expression", + data_path=data_path, + dataset_name=dataset_name, + gene_list=self.hyperparameters.get("gene_list", "gene_list_paccmann_network_prop"), + ) + + def load_drug_features(self, data_path: str, dataset_name: str) -> FeatureDataset: + """Load raw SMILES features. + + Only the id and SMILES columns are read from the csv. The fingerprint columns + (cactvs_fingerprint, fingerprint) are huge decimal bit-strings that are not needed + here, and are skipped rather than loaded and dropped, since pandas' numeric type + inference on them can raise an OverflowError. + + :param data_path: path to the data directory + :param dataset_name: name of the dataset + :return: FeatureDataset containing SMILES features + """ + id_column = "pubchem_id" + smiles_column = "canonical_smiles" + + data = pd.read_csv( + f"{data_path}/{dataset_name}/drug_smiles.csv", + usecols=[id_column, smiles_column], + dtype={id_column: str, smiles_column: str}, + ) + data = data.drop_duplicates(subset=id_column, keep="first") + + features = { + str(pubchem_id): {"smiles": np.array([smiles], dtype=object)} + for pubchem_id, smiles in zip(data[id_column], data[smiles_column], strict=True) + } + + return FeatureDataset(features=features, meta_info={"smiles": [smiles_column]}) + + def build_model(self, hyperparameters: dict[str, Any]) -> None: + """Store hyperparameters for later model initialization. + + The actual PaccMannNetwork is initialized in train(), + because the number of genes depends on the loaded training data. + + :param hyperparameters: dictionary containing model hyperparameters + """ + self.log_hyperparameters(hyperparameters) + self.hyperparameters = hyperparameters + + def _build_smiles_vocab(self, smiles_list: list[str]) -> None: + """Build a token vocabulary from training SMILES strings. + + :param smiles_list: list of SMILES strings + """ + for smile in smiles_list: # Build vocabulary: "Cl", "C", "=" ... -> {"Cl": 2, "C": 3, "=": 4} + for token in _tokenize_smiles(smile): + if token not in self.smiles_to_idx: + self.smiles_to_idx[token] = len(self.smiles_to_idx) + + def _encode_smiles(self, smiles_list: list[str]) -> np.ndarray: + """Encode SMILES strings as padded integer sequences. + + :param smiles_list: list of SMILES strings + :return: encoded SMILES array + :raises ValueError: if smiles_padding_length is not set + """ + if self.smiles_padding_length is None: + raise ValueError("smiles_padding_length is not set.") + + encoded = np.full( + (len(smiles_list), self.smiles_padding_length), + fill_value=PAD_IDX, + dtype=np.int64, + ) + + for i, smile in enumerate(smiles_list): + tokens = _tokenize_smiles(smile) + token_ids = [self.smiles_to_idx.get(token, UNK_IDX) for token in tokens] # "CCO" -> [2, 2, 3] + token_ids = token_ids[: self.smiles_padding_length] + encoded[i, : len(token_ids)] = token_ids # Padding: [2,2,2] -> [2,2,3,0,0,0,...] + + return encoded + + def _encode_inputs( + self, + cell_line_input: FeatureDataset, + drug_input: FeatureDataset, + cell_line_ids: np.ndarray, + drug_ids: np.ndarray, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Turn cell line and drug features into model input tensors. + + Gene expression is scaled with the scaler fitted on the training data and SMILES are encoded with the + vocabulary built from the training data, so this may only be called after train() has fitted both. + + :param cell_line_input: FeatureDataset containing cell line features + :param drug_input: FeatureDataset containing drug features + :param cell_line_ids: array of cell line identifiers + :param drug_ids: array of drug identifiers + :return: tuple of the encoded SMILES tensor and the scaled gene expression tensor + """ + gex = cell_line_input.get_feature_matrix("gene_expression", cell_line_ids) + gex = np.asarray(gex, dtype=np.float32) + gex = self.gene_expression_scaler.transform(gex).astype(np.float32) + + smiles = _smiles_column_to_list(drug_input.get_feature_matrix("smiles", drug_ids)) + smiles_encoded = self._encode_smiles(smiles) + + return ( + torch.tensor(smiles_encoded, dtype=torch.long), + torch.tensor(gex, dtype=torch.float32), + ) + + def _build_smiles_variants(self, unique_smiles: list[str]) -> list[list[str]]: + """Build the SMILES variants each drug is trained on. + + With augmentation enabled every drug gets several equivalent SMILES strings; without it each drug keeps + its single original SMILES. If RDKit is missing, augmentation is skipped with a warning rather than + failing, so the model stays usable without the optional dependency. + + :param unique_smiles: list of distinct SMILES strings + :return: list holding the variants of each molecule, the original SMILES first + """ + if not self.hyperparameters.get("augment_smiles", True): + return [[smiles] for smiles in unique_smiles] + + try: + return _randomize_smiles(unique_smiles, N_SMILES_VARIANTS) + except ImportError as e: # pragma: no cover - depends on the environment + warnings.warn( + f"{e} Training PaccMann without SMILES augmentation.", + stacklevel=2, + ) + return [[smiles] for smiles in unique_smiles] + + def _build_validation_loader( + self, + output_earlystopping: DrugResponseDataset | None, + cell_line_input: FeatureDataset, + drug_input: FeatureDataset, + batch_size: int, + ) -> DataLoader | None: + """Build a loader over the early stopping set, used to pick the best epoch. + + :param output_earlystopping: early stopping dataset, may be None + :param cell_line_input: FeatureDataset containing cell line features + :param drug_input: FeatureDataset containing drug features + :param batch_size: batch size to use + :return: DataLoader over the early stopping set, or None if there is nothing to evaluate + """ + if output_earlystopping is None or len(output_earlystopping) == 0: + return None + + smiles_tensor, gex_tensor = self._encode_inputs( + cell_line_input, + drug_input, + output_earlystopping.cell_line_ids, + output_earlystopping.drug_ids, + ) + y_tensor = torch.tensor(np.asarray(output_earlystopping.response, dtype=np.float32)).view(-1, 1) + + return DataLoader( + TensorDataset(smiles_tensor, gex_tensor, y_tensor), + batch_size=batch_size, + shuffle=False, + ) + + def _validation_loss(self, validation_loader: DataLoader) -> float: + """Compute the mean loss over the early stopping set. + + :param validation_loader: DataLoader over the early stopping set + :return: mean loss per batch + :raises ValueError: if the model has not been built yet + """ + if self.model is None: + raise ValueError("Model has not been built yet.") + + self.model.eval() + total_loss = 0.0 + with torch.no_grad(): + for batch_smiles, batch_gex, batch_y in validation_loader: + batch_smiles = batch_smiles.to(self.device) + batch_gex = batch_gex.to(self.device) + batch_y = batch_y.to(self.device) + + predictions = self.model(batch_smiles, batch_gex) + total_loss += functional.mse_loss(predictions, batch_y).item() + + return total_loss / len(validation_loader) + + def train( + self, + output: DrugResponseDataset, + cell_line_input: FeatureDataset, + drug_input: FeatureDataset | None = None, + output_earlystopping: DrugResponseDataset | None = None, + model_checkpoint_dir: str | None = None, + ) -> None: + """Train the PaccMann model on gene DrEval data. + + Procedure: + - get gene expression data for all the cell lines + - get raw SMILES for the drugs + - scale gene expression features + - build a SMILES vocabulary + - encode and pad the SMILES strings + - initialize the PaccMann network + - convert both inputs to tensors + - train the network, stopping early once the early stopping set has not improved for `patience` epochs + - restore the weights of the epoch with the lowest loss on the early stopping set + + :param output: training dataset containing response values, cell line ids, and drug ids + :param cell_line_input: FeatureDataset containing cell line features + :param drug_input: FeatureDataset containing drug features + :param output_earlystopping: dataset used to select the best epoch and to trigger early stopping. + If None, the weights of the last epoch are kept and training runs for the full epoch budget. + :param model_checkpoint_dir: optional directory to save a model checkpoint + :raises ValueError: if drug_input is None + :raises ValueError: if the model has not been built yet + """ + if drug_input is None: + raise ValueError("drug_input (SMILES) is required for PaccMann.") + + if self.hyperparameters is None: + raise ValueError("Model has not been built yet. Call build_model first.") + + gex = cell_line_input.get_feature_matrix("gene_expression", output.cell_line_ids) + gex = np.asarray(gex, dtype=np.float32) + gex = self.gene_expression_scaler.fit_transform(gex).astype(np.float32) + + smiles = _smiles_column_to_list(drug_input.get_feature_matrix("smiles", output.drug_ids)) + y = np.asarray(output.response, dtype=np.float32) + + # Each drug is trained on several equivalent SMILES strings, so the variants are built per distinct drug + # and every response row only needs to remember which drug it belongs to. + unique_smiles, row_to_drug = np.unique(np.asarray(smiles, dtype=object), return_inverse=True) + smiles_variants = self._build_smiles_variants(list(unique_smiles)) + flat_variants = [variant for molecule_variants in smiles_variants for variant in molecule_variants] + + # Build SMILES vocabulary from training data only. It has to cover the augmented variants as well, + # otherwise their tokens would be encoded as unknown. + self.smiles_to_idx = dict(BASE_SMILES_VOCAB) + self._build_smiles_vocab(flat_variants) + + if "smiles_padding_length" in self.hyperparameters: + self.smiles_padding_length = int(self.hyperparameters["smiles_padding_length"]) + else: + self.smiles_padding_length = max(len(_tokenize_smiles(variant)) for variant in flat_variants) + + # Encode the variants into a bank of shape (drugs, variants, padding length), sampled from per batch + n_variants = len(smiles_variants[0]) + smiles_bank = torch.tensor( + self._encode_smiles(flat_variants).reshape(len(smiles_variants), n_variants, -1), + dtype=torch.long, + ) + + model_params = dict(self.hyperparameters) + model_params["number_of_genes"] = gex.shape[1] + model_params["smiles_padding_length"] = self.smiles_padding_length + model_params["smiles_vocabulary_size"] = len(self.smiles_to_idx) + self.number_of_genes = gex.shape[1] + + self.model = PaccMannNetwork(PaccMannConfig.from_hyperparameters(model_params)).to(self.device) + + # The dataset holds the drug index rather than the encoded SMILES, so a fresh variant can be drawn for + # every row in every epoch without materializing one tensor per epoch. + drug_tensor = torch.tensor(row_to_drug, dtype=torch.long) + gex_tensor = torch.tensor(gex, dtype=torch.float32) + y_tensor = torch.tensor(y, dtype=torch.float32).view(-1, 1) + + dataset = TensorDataset(drug_tensor, gex_tensor, y_tensor) + batch_size = model_params.get("batch_size", 64) + + # The batch norm layers cannot process a batch that holds a single sample, so a trailing batch of size 1 + # has to be dropped. The original implementation always drops the last batch; dropping it only when it + # would contain a single sample keeps training sets smaller than one batch usable. + drop_last = len(dataset) > batch_size and len(dataset) % batch_size == 1 + + train_loader = DataLoader( + dataset, + batch_size=batch_size, + shuffle=True, + drop_last=drop_last, + ) + + optimizer = torch.optim.Adam( + self.model.parameters(), + lr=model_params.get("learning_rate", 1e-3), + weight_decay=model_params.get("weight_decay", 0.0), + ) + + epochs = model_params.get("epochs", 20) + + validation_loader = self._build_validation_loader( + output_earlystopping, + cell_line_input, + drug_input, + batch_size, + ) + + patience = model_params.get("patience", 10) + best_validation_loss = float("inf") + best_state_dict: dict[str, Any] | None = None + epochs_without_improvement = 0 + + for epoch in range(epochs): + self.model.train() + for batch_drugs, batch_gex, batch_y in train_loader: + # Draw one of the equivalent SMILES strings per row, so a drug is seen through a different + # SMILES string in every epoch + variant = torch.randint(0, n_variants, (batch_drugs.shape[0],)) + batch_smiles = smiles_bank[batch_drugs, variant].to(self.device) + batch_gex = batch_gex.to(self.device) + batch_y = batch_y.to(self.device) + + optimizer.zero_grad() + + predictions = self.model(batch_smiles, batch_gex) + loss = functional.mse_loss(predictions, batch_y) + + loss.backward() + optimizer.step() + + if validation_loader is None: + continue + + validation_loss = self._validation_loss(validation_loader) + self.log_metrics({"validation_loss": validation_loss}, step=epoch) + + if validation_loss < best_validation_loss: + best_validation_loss = validation_loss + best_state_dict = copy.deepcopy(self.model.state_dict()) + epochs_without_improvement = 0 + else: + epochs_without_improvement += 1 + if epochs_without_improvement >= patience: + break + + if best_state_dict is not None: + self.model.load_state_dict(best_state_dict) + + if model_checkpoint_dir is not None: + torch.save(self.model.state_dict(), f"{model_checkpoint_dir}/paccmann.pt") + + def predict( + self, + cell_line_ids: np.ndarray, + drug_ids: np.ndarray, + cell_line_input: FeatureDataset, + drug_input: FeatureDataset | None = None, + ) -> np.ndarray: + """Predict drug response values. + + Procedure: + - load appropriate cell lines features and drug features + - scale gene expression features + - encode and pad SMILES strings + - convert inputs to tensors + - run the trained PaccMann model in evaluation mode + - return predicted drug response values + + :param cell_line_ids: array of cell line identifiers + :param drug_ids: array of drug identifiers + :param cell_line_input: FeatureDataset containing cell line features + :param drug_input: FeatureDataset containing drug features + :return: predicted drug response values + :raises ValueError: if drug_input is None + :raises ValueError: if the model has not been trained yet + """ + if drug_input is None: + raise ValueError("drug_input (SMILES) is required for PaccMann.") + + if self.model is None: + raise ValueError("Model has not been trained yet.") + + # _encode_inputs returns CPU tensors; batches are moved to the device one at a time below. + smiles_tensor, gex_tensor = self._encode_inputs(cell_line_input, drug_input, cell_line_ids, drug_ids) + + dataset = TensorDataset(smiles_tensor, gex_tensor) + predict_loader = DataLoader( + dataset, + batch_size=self.hyperparameters.get("batch_size", 64), + shuffle=False, + ) + + self.model.eval() + predictions_list = [] + with torch.no_grad(): + for batch_smiles, batch_gex in predict_loader: + batch_smiles = batch_smiles.to(self.device) + batch_gex = batch_gex.to(self.device) + batch_predictions = self.model(batch_smiles, batch_gex) + predictions_list.append(batch_predictions.cpu()) + + predictions = torch.cat(predictions_list, dim=0) + return predictions.numpy().reshape(-1) + + def save(self, path: str) -> None: + """Save the trained PaccMann wrapper. + + Saved files: + - model.pt: trained model weights + - config.json: model hyperparameters + - scaler.pkl: fitted gene expression scaler + - vocab.json: SMILES vocabulary + - meta.json: additional metadata needed for loading + + :param path: directory where the model should be saved + :raises ValueError: if no model is available + """ + os.makedirs(path, exist_ok=True) + + if self.model is None: + raise ValueError("No model to save.") + + torch.save(self.model.state_dict(), f"{path}/model.pt") + + with open(f"{path}/config.json", "w") as f: + json.dump(self.hyperparameters, f) + + joblib.dump(self.gene_expression_scaler, f"{path}/scaler.pkl") + + with open(f"{path}/vocab.json", "w") as f: + json.dump(self.smiles_to_idx, f) + + with open(f"{path}/meta.json", "w") as f: + json.dump( + { + "padding_length": self.smiles_padding_length, + "num_genes": self.number_of_genes, + }, + f, + ) + + @classmethod + def load(cls, path: str) -> PaccMann: + """Load a trained PaccMann wrapper. + + :param path: directory containing the saved model files + :return: loaded PaccMann instance + """ + instance = cls() + + with open(f"{path}/config.json") as f: + instance.hyperparameters = json.load(f) + + instance.gene_expression_scaler = joblib.load(f"{path}/scaler.pkl") + + with open(f"{path}/vocab.json") as f: + instance.smiles_to_idx = json.load(f) + + with open(f"{path}/meta.json") as f: + meta = json.load(f) + instance.smiles_padding_length = meta["padding_length"] + instance.number_of_genes = meta["num_genes"] + + params = dict(instance.hyperparameters) + params["smiles_padding_length"] = instance.smiles_padding_length + params["smiles_vocabulary_size"] = len(instance.smiles_to_idx) + params["number_of_genes"] = instance.number_of_genes + + instance.model = PaccMannNetwork(PaccMannConfig.from_hyperparameters(params)).to(instance.device) + instance.model.load_state_dict(torch.load(f"{path}/model.pt", map_location=instance.device)) # noqa: S614 + instance.model.eval() + + return instance diff --git a/drevalpy/models/__init__.py b/drevalpy/models/__init__.py index 334d407c..d5fafd2d 100644 --- a/drevalpy/models/__init__.py +++ b/drevalpy/models/__init__.py @@ -32,6 +32,7 @@ "MultiViewXGBoost", "MultiViewLightGBM", "SparseGO", + "PaccMann", ] from .baselines.multi_view_lightgbm import MultiViewLightGBM @@ -59,6 +60,7 @@ from .drp_model import DRPModel from .DrugGNN import DrugGNN from .MOLIR.molir import MOLIR +from .PaccMann.paccmann import PaccMann from .PharmaFormer.pharmaformer import PharmaFormerModel from .Precily import PrecilyModel from .SimpleNeuralNetwork.multi_view_neural_network import MultiViewNeuralNetwork @@ -105,6 +107,7 @@ "SRMF": SRMF, "Precily": PrecilyModel, "SparseGO": SparseGOModel, + "PaccMann": PaccMann, } # MODEL_FACTORY is used in the pipeline! diff --git a/noxfile.py b/noxfile.py index b01698f4..f19898cd 100644 --- a/noxfile.py +++ b/noxfile.py @@ -144,7 +144,7 @@ def tests(session: Session) -> None: :param session: The Session object. """ - session.install(".[xgboost,precily,sparsego]") + session.install(".[xgboost,precily,sparsego,paccmann]") session.install("coverage[toml]", "pytest", "pygments") try: session.run( @@ -188,7 +188,7 @@ def typeguard(session: Session) -> None: :param session: The Session object. """ - session.install(".[xgboost,precily,sparsego]") + session.install(".[xgboost,precily,sparsego,paccmann]") session.install("pytest", "typeguard", "pygments") session.run( diff --git a/poetry.lock b/poetry.lock index 71ef279e..a87cbe09 100644 --- a/poetry.lock +++ b/poetry.lock @@ -4870,6 +4870,41 @@ serve-grpc = ["aiohttp (>=3.13.3)", "aiohttp_cors", "colorful", "fastapi (>=0.13 train = ["fsspec", "pandas", "pyarrow (>=17.0.0)", "pydantic (>=2.13.0,<3) ; python_version >= \"3.14\"", "pydantic (>=2.5.0,<3) ; python_version < \"3.14\"", "requests", "tensorboardX (>=1.9)"] tune = ["fsspec", "pandas", "pyarrow (>=17.0.0)", "pydantic (>=2.13.0,<3) ; python_version >= \"3.14\"", "pydantic (>=2.5.0,<3) ; python_version < \"3.14\"", "requests", "tensorboardX (>=1.9)"] +[[package]] +name = "rdkit" +version = "2026.3.5" +description = "A collection of chemoinformatics and machine-learning software written in C++ and Python" +optional = true +python-versions = "*" +groups = ["main"] +markers = "extra == \"paccmann\"" +files = [ + {file = "rdkit-2026.3.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ff2103456dd726fbb3f67952ae72c42026cdba46f565352f7f585741b02ba681"}, + {file = "rdkit-2026.3.5-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:4ed70419877db5dd3f47dc10d3d4b5ab317853cfb9218499b0634854b36ed003"}, + {file = "rdkit-2026.3.5-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:3b299c2d1e2c4da14b38fae62c650ea315468cb54b2b83a42e1eeed212f5027d"}, + {file = "rdkit-2026.3.5-cp310-cp310-win_amd64.whl", hash = "sha256:33b7f6e604ee29e4dc426627062ecf64179acc1e7767ad00b738a4d262734f0c"}, + {file = "rdkit-2026.3.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f3b0102b8d8a2f45faf039d31440e7fcf0dcd423d1416def321e2d772b270f41"}, + {file = "rdkit-2026.3.5-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:ed230d8d085a85aa45e02509689ea181aad3a1ddcc9aed7f03d03fede8bbee7f"}, + {file = "rdkit-2026.3.5-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:989323fee0059fa468f408f86012e3bbd3252fa9d0355382cc22c29d1cb0cd78"}, + {file = "rdkit-2026.3.5-cp311-cp311-win_amd64.whl", hash = "sha256:47ae231c5e8aa03359b91e9639025501cc966b17cc9d8f4ea9f0fcf14bfdb469"}, + {file = "rdkit-2026.3.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:74e621083f26360ae3128b2c283def72e1729c114577c58115294b1cefe0200b"}, + {file = "rdkit-2026.3.5-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:d6c6b167b4c795468cdd273d35a323226dddbeb204aa34350257ad26d5dfd024"}, + {file = "rdkit-2026.3.5-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b75944ba959d908e97b4d68754e5950216ac08aa81faf67cfd1d7a3cb5b2bad7"}, + {file = "rdkit-2026.3.5-cp312-cp312-win_amd64.whl", hash = "sha256:b60a2b6e8e2cecd89f775c6a3e691d3dacc5ae05cf521154822e7e5a54602825"}, + {file = "rdkit-2026.3.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f56eb842b8c0716348b31fc97fe0c6581fc39d32567085f19f96bd7f51f0a96c"}, + {file = "rdkit-2026.3.5-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:91358e266e5189c2402cc8c7df1f34688ec9e25a7235f2191b829f92fadc56df"}, + {file = "rdkit-2026.3.5-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:04a75a048cd61ef934fd2fd474ff426f40cf22e83fe8cd038d1563d695f7e314"}, + {file = "rdkit-2026.3.5-cp313-cp313-win_amd64.whl", hash = "sha256:91b31d4ce9f380a09fb263a882d4d8a97eee3dca54acca5b9c568d9bc859dc96"}, + {file = "rdkit-2026.3.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:74c5b98da1d52e83f42ceffcc9dc91fcb0b7615eab82e9a1e736aeb4a6c1cb0e"}, + {file = "rdkit-2026.3.5-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:7fa98cb6ad79238c7a9cf0a7b42c8abbfed659e9ab5f7b2cba9b17db2915653c"}, + {file = "rdkit-2026.3.5-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:b79c185602aa21bdfd1f01daae7171598090ee4c4d984d5544905b5318d29288"}, + {file = "rdkit-2026.3.5-cp314-cp314-win_amd64.whl", hash = "sha256:f0096fdc40ab259151a04933e8201d877f5d274966680c922cd63d3cffb330aa"}, +] + +[package.dependencies] +numpy = "*" +Pillow = "*" + [[package]] name = "referencing" version = "0.37.0" @@ -7000,6 +7035,7 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [extras] multiprocessing = ["pydantic", "ray"] +paccmann = ["rdkit"] precily = ["gseapy"] sparsego = ["mygene", "obonet"] xgboost = ["xgboost"] @@ -7007,4 +7043,4 @@ xgboost = ["xgboost"] [metadata] lock-version = "2.1" python-versions = ">=3.11,<3.14" -content-hash = "8b4114488e9751e2bfa258bd9d4bf2fdf1964ebac8b25ce3b7520d68382f6cb4" +content-hash = "f22f494d706a5d6aa4ba9c3398cafa6d656d4dc418087d0ad8534d6c26c478ff" diff --git a/pyproject.toml b/pyproject.toml index c9794252..c5950b80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ rich = ">=15.0.0" gseapy = { version = ">=1.1.0", optional = true } mygene = { version = "*", optional = true } obonet = { version = "*", optional = true } +rdkit = { version = ">=2022.9", optional = true } [tool.poetry.requires-plugins] poetry-plugin-export = ">=1.8" @@ -65,6 +66,7 @@ multiprocessing = ["ray", "pydantic"] xgboost = ["xgboost"] precily = ["gseapy"] sparsego = ["mygene", "obonet"] +paccmann = ["rdkit"] [tool.poetry.dependencies.ray] extras = ["tune"] diff --git a/tests/models/test_global_models.py b/tests/models/test_global_models.py index ae633acd..296d58ab 100644 --- a/tests/models/test_global_models.py +++ b/tests/models/test_global_models.py @@ -25,6 +25,7 @@ "SimpleNeuralNetwork[chemberta]", "MultiViewNeuralNetwork", "PharmaFormer", + "PaccMann", "Precily", "SparseGO", ], @@ -83,6 +84,9 @@ def test_global_models( elif model_name == "SparseGO": hpam_combi["epochs"] = 1 hpam_combi["batch_size"] = 32 + elif model_name == "PaccMann": + hpam_combi["epochs"] = 1 + hpam_combi["gene_list"] = None elif model_name == "AdaBoostDecisionTree": hpam_combi["max_depth"] = 2 hpam_combi["min_samples_split"] = 2