diff --git a/docs/drevalpy.models.PaccMann.rst b/docs/drevalpy.models.PaccMann.rst new file mode 100644 index 000000000..9609f283b --- /dev/null +++ b/docs/drevalpy.models.PaccMann.rst @@ -0,0 +1,50 @@ +PaccMann +============================= + +PaccMann Model +---------------------------------- + +.. automodule:: drevalpy.models.PaccMann.paccmann + :members: + :undoc-members: + :show-inheritance: + +PaccMannV2 Network +---------------------------------- + +.. automodule:: drevalpy.models.PaccMann.paccmann_network_v2 + :members: + :undoc-members: + :show-inheritance: + +Hyperparameter utils +---------------------------------- + +.. automodule:: drevalpy.models.PaccMann.utils.hyperparams + :members: + :undoc-members: + :show-inheritance: + +Layers +---------------------------------- + +.. automodule:: drevalpy.models.PaccMann.utils.layers + :members: + :undoc-members: + :show-inheritance: + +Loss functions +---------------------------------- + +.. automodule:: drevalpy.models.PaccMann.utils.loss_functions + :members: + :undoc-members: + :show-inheritance: + +Model utils +---------------------------------- + +.. automodule:: drevalpy.models.PaccMann.utils.utils + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/drevalpy.models.rst b/docs/drevalpy.models.rst index 3fda95618..fb72b92fd 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 5fc8ba93e..0ff5dac28 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 54428e3a4..a8ce939dd 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -302,6 +302,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 000000000..524763c24 --- /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 000000000..71941c0a0 --- /dev/null +++ b/drevalpy/models/PaccMann/hyperparameters.yaml @@ -0,0 +1,59 @@ +PaccMann: + gene_list: + - gene_list_paccmann_network_prop + + epochs: + - 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 + + activation_fn: + - relu + + loss_fn: + - mse + + 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/paccmann.py b/drevalpy/models/PaccMann/paccmann.py new file mode 100644 index 000000000..738f20b7a --- /dev/null +++ b/drevalpy/models/PaccMann/paccmann.py @@ -0,0 +1,667 @@ +"""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 +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 .paccmann_network_v2 import PaccMannV2 + +# 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] + + +# 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 PaccMannV2 neural network 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 PaccMannV2 PyTorch model + - keeps the weights of the epoch with the lowest loss on the early stopping set + + The original implementation (https://github.com/PaccMann/paccmann_predictor) trains for a fixed number of + epochs and checkpoints the model whenever the loss on a held-out set improves, using that checkpoint as the + final model. This wrapper follows the same procedure with the early stopping set: training always runs for + the full epoch budget and the weights of the best epoch are restored at the end. There is no patience-based + termination, since the original implementation does not stop early either. + """ + + 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 + padding_idx: index used for padding tokens + unk_idx: index for unknown tokens + smiles_padding_length: sequence length used for padding + number_of_genes: number of gene features + """ + super().__init__() + self.model: PaccMannV2 | 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] = { + "": 0, + "": 1, + } + self.padding_idx = 0 + self.unk_idx = 1 + + 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 PaccMannV2 network 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 _normalize_smiles_array(self, smiles_raw: np.ndarray) -> list[str]: + """Convert SMILES output from FeatureDataset into a list of strings. + + :param smiles_raw: raw SMILES array + :return: list of SMILES strings + """ + smiles_raw = np.asarray(smiles_raw, dtype=object) + + if smiles_raw.ndim == 2 and smiles_raw.shape[1] == 1: + smiles_raw = smiles_raw[:, 0] + + smiles_list = [] + for smile in smiles_raw: + if smile is None: + smiles_list.append("") + else: + smiles_list.append(str(smile)) + return smiles_list + + 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=self.padding_idx, + dtype=np.int64, + ) + + for i, smile in enumerate(smiles_list): + tokens = _tokenize_smiles(smile) + token_ids = [self.smiles_to_idx.get(token, self.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 = self._normalize_smiles_array(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 += self.model.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 for the full epoch budget, evaluating the early stopping set after every epoch + - 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. If None, the weights of the last + epoch are kept. + :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.") + + # Retrieve gene expression features for the training cell lines + gex = cell_line_input.get_feature_matrix("gene_expression", output.cell_line_ids) + + # Retrieve raw SMILES features for the corresponding drugs + smiles_raw = drug_input.get_feature_matrix("smiles", output.drug_ids) + + # Target vector containing the drug response values + y = output.response + + # Convert to numpy arrays with explicit dtypes + gex = np.asarray(gex, dtype=np.float32) + y = np.asarray(y, dtype=np.float32) + + # Convert SMILES to a list of strings + smiles = self._normalize_smiles_array(smiles_raw) + + # Scale gene expression on training data only + gex = self.gene_expression_scaler.fit_transform(gex).astype(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 = { + "": 0, + "": 1, + } + self._build_smiles_vocab(flat_variants) + + # Determine SMILES padding length + 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, + ) + + # Copy hyperparameters and adapt the number of genes to the training data + 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] + + # Build the PaccMann neural network + self.model = PaccMannV2(model_params).to(self.device) + + # Convert all inputs to tensors. The dataset holds the drug index rather than the encoded SMILES, so that + # 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) + + # Create PyTorch dataset and dataloader + 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, + ) + + # Initialize optimizer + 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) + + # The early stopping set is evaluated after every epoch to pick the best epoch + validation_loader = self._build_validation_loader( + output_earlystopping, + cell_line_input, + drug_input, + batch_size, + ) + + best_validation_loss = float("inf") + best_state_dict: dict[str, Any] | None = None + + # Train the model + 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 = self.model.loss(predictions, batch_y) + + loss.backward() + optimizer.step() + + if validation_loader is None: + continue + + # Keep the weights of the epoch with the lowest loss on the early stopping set, which is what the + # original implementation checkpoints. Training itself is never cut short. + 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()) + + # Restore the weights of the best epoch + if best_state_dict is not None: + self.model.load_state_dict(best_state_dict) + + # Optional: save trained model checkpoint + if model_checkpoint_dir is not None: + self.model.save(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.") + + # Scale gene expression and encode SMILES with the training scaler and vocabulary. + # These are 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, + ) + + # Predict drug response values batch-wise + 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 = PaccMannV2(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/PaccMann/paccmann_network_v2.py b/drevalpy/models/PaccMann/paccmann_network_v2.py new file mode 100644 index 000000000..9318d0545 --- /dev/null +++ b/drevalpy/models/PaccMann/paccmann_network_v2.py @@ -0,0 +1,303 @@ +"""Contains the PaccMannV2 model for drug response prediction. + +The model is based on the PaccMann framework for predicting anticancer drug sensitivity +from SMILES strings and gene expression data. + +Original PaccMann repository: +https://github.com/PaccMann/paccmann_predictor +""" + +from collections import OrderedDict + +import torch +import torch.nn as nn + +from .utils.hyperparams import ACTIVATION_FN_FACTORY, LOSS_FN_FACTORY +from .utils.layers import ContextAttentionLayer, convolutional_layer, dense_layer +from .utils.utils import get_device + + +class PaccMannV2(nn.Module): + """PaccMannV2 model for drug response prediction. + + Based on the MCA model in Molecular Pharmaceutics: + https://pubs.acs.org/doi/10.1021/acs.molpharmaceut.9b00520. + + Main idea: + - SMILES strings are embedded and processed with convolutional layers + - gene expression is used as biological context + - context attention connects gene and drug information + - the combined representation is passed through dense layers + - output is a predicted drug sensitivity value + """ + + def __init__(self, params, *args, **kwargs): + """Initialize the PaccMannV2 model. + + :param params: A dictionary containing the parameter to built the dense encoder. + :param args: additional positional arguments passed to nn.Module + :param kwargs: additional keyword arguments passed to nn.Module + + Items in params: + - smiles_padding_length (int): Padding length for SMILES. + - smiles_embedding_size (int): dimension of tokens' embedding. + - smiles_vocabulary_size (int): size of the tokens vocabulary. + - activation_fn (string, optional): Activation function used in all layers + for specification in ACTIVATION_FN_FACTORY. Defaults to 'relu'. + - batch_norm (bool, optional): Whether batch normalization is applied. + Defaults to True. + - dropout (float, optional): Dropout probability in all except parametric layer. Defaults to 0.5. + - filters (list[int], optional): Numbers of filters to learn per + SMILES convolutional layer. Defaults to [64, 64, 64]. + - kernel_sizes (list[list[int]], optional): Sizes of kernels per SMILES convolutional layer. + Defaults to + [[3, params['smiles_embedding_size']], + [5, params['smiles_embedding_size']], + [11, params['smiles_embedding_size']]] + - molecule_heads (list[int], optional): Amount of attentive molecule_heads + per SMILES embedding. Should have len(filters)+1. + Defaults to [4, 4, 4, 4]. + - stacked_dense_hidden_sizes (list[int], optional): Sizes of the + hidden dense layers. Defaults to [1024, 512]. + - smiles_attention_size (int, optional): size of the attentive layer + for the smiles sequence. Defaults to 64. + + :raises ValueError: if the attention head settings or convolution settings do not match + """ + super().__init__(*args, **kwargs) + + # model parameter + self.device = get_device() + self.params = params + + # select loss function + self.loss_fn = LOSS_FN_FACTORY[params.get("loss_fn", "mse")] + + # input sizes + self.smiles_padding_length = params["smiles_padding_length"] + self.number_of_genes = params["number_of_genes"] + + # attention settings + self.smiles_attention_size = params.get("smiles_attention_size", 64) + self.gene_attention_size = params.get("gene_attention_size", 1) + + self.molecule_temperature = params.get("molecule_temperature", 1.0) + self.gene_temperature = params.get("gene_temperature", 1.0) + + # model architecture (hyperparameter) + self.molecule_heads = params.get("molecule_heads", [4, 4, 4, 4]) + self.gene_heads = params.get("gene_heads", [2, 2, 2, 2]) + + if len(self.gene_heads) != len(self.molecule_heads): + raise ValueError("Length of gene and molecule_heads do not match.") + + self.filters = params.get("filters", [64, 64, 64]) + + # size of dense input + self.hidden_sizes = [ + self.molecule_heads[0] * params["smiles_embedding_size"] + + sum([h * f for h, f in zip(self.molecule_heads[1:], self.filters)]) + + sum(self.gene_heads) * self.number_of_genes + ] + params.get("stacked_dense_hidden_sizes", [1024, 512]) + + # general NN settings + self.dropout = params.get("dropout", 0.5) + self.temperature = params.get("temperature", 1.0) + self.act_fn = ACTIVATION_FN_FACTORY[params.get("activation_fn", "relu")] + + # Default convolution kernel size + self.kernel_sizes = params.get( + "kernel_sizes", + [ + [3, params["smiles_embedding_size"]], + [5, params["smiles_embedding_size"]], + [11, params["smiles_embedding_size"]], + ], + ) + if len(self.filters) != len(self.kernel_sizes): + raise ValueError("Length of filter and kernel size lists do not match.") + + if len(self.filters) + 1 != len(self.molecule_heads): + raise ValueError("Length of filter and multihead lists do not match") + + # Build the model + self.smiles_embedding = nn.Embedding( + self.params["smiles_vocabulary_size"], + self.params["smiles_embedding_size"], + scale_grad_by_freq=params.get("embed_scale_grad", False), + ) + + # Convolution layers over embedded SMILES + self.convolutional_layers = nn.Sequential( + OrderedDict( + [ + ( + f"convolutional_{index}", + convolutional_layer( + num_kernel, + kernel_size, + act_fn=self.act_fn, + batch_norm=params.get("batch_norm", False), + dropout=self.dropout, + ).to(self.device), + ) + for index, (num_kernel, kernel_size) in enumerate(zip(self.filters, self.kernel_sizes)) + ] + ) + ) + + # Hidden size of each SMILES representation stage: raw embedding + outputs of conv layer + smiles_hidden_sizes = [params["smiles_embedding_size"]] + self.filters + + # Attention layers: gene context -> SMILES (focus on relevant molecule parts) + self.molecule_attention_layers = nn.Sequential( + OrderedDict( + [ + ( + f"molecule_attention_{layer}_head_{head}", + ContextAttentionLayer( + reference_hidden_size=smiles_hidden_sizes[layer], + reference_sequence_length=self.smiles_padding_length, + context_hidden_size=1, + context_sequence_length=self.number_of_genes, + attention_size=self.smiles_attention_size, + individual_nonlinearity=params.get("context_nonlinearity", nn.Sequential()), + temperature=self.molecule_temperature, + ), + ) + for layer in range(len(self.molecule_heads)) + for head in range(self.molecule_heads[layer]) + ] + ) + ) + + # Attention layers: SMILES -> gene expression (focus on relevant genes) + self.gene_attention_layers = nn.Sequential( + OrderedDict( + [ + ( + f"gene_attention_{layer}_head_{head}", + ContextAttentionLayer( + reference_hidden_size=1, + reference_sequence_length=self.number_of_genes, + context_hidden_size=smiles_hidden_sizes[layer], + context_sequence_length=self.smiles_padding_length, + attention_size=self.gene_attention_size, + individual_nonlinearity=params.get("context_nonlinearity", nn.Sequential()), + temperature=self.gene_temperature, + ), + ) + for layer in range(len(self.molecule_heads)) + for head in range(self.gene_heads[layer]) + ] + ) + ) + + # Batch normalization for the concatenated attention output + # Only applied if params['batch_norm'] = True + self.batch_norm = nn.BatchNorm1d(self.hidden_sizes[0]) + + # Dense layers after attention + self.dense_layers = nn.Sequential( + OrderedDict( + [ + ( + f"dense_{ind}", + dense_layer( + self.hidden_sizes[ind], + self.hidden_sizes[ind + 1], + act_fn=self.act_fn, + dropout=self.dropout, + batch_norm=params.get("batch_norm", True), + ).to(self.device), + ) + for ind in range(len(self.hidden_sizes) - 1) + ] + ) + ) + + # Final output layer + self.final_dense = ( + nn.Linear(self.hidden_sizes[-1], 1) + if not params.get("final_activation", False) + else nn.Sequential( + OrderedDict( + [ + ("projection", nn.Linear(self.hidden_sizes[-1], 1)), + ("sigmoidal", ACTIVATION_FN_FACTORY["sigmoid"]), + ] + ) + ) + ) + + def forward(self, smiles, gep): + """Forward pass through the PaccMannV2. + + :param smiles: tokenized SMILES tensor of shape [bs, smiles_padding_length] + :param gep: gene expression tensor of shape [bs, number_of_genes] + :return: predictions, a tensor of shape [batch_size, 1] + """ + # reshape gene input + gep = torch.unsqueeze(gep, dim=-1) + embedded_smiles = self.smiles_embedding(smiles.to(dtype=torch.int64)) + + # SMILES convolutions, over an input of shape bs x 1 x T x H + encoded_smiles = [embedded_smiles] + [ + self.convolutional_layers[ind](torch.unsqueeze(embedded_smiles, 1)).permute(0, 2, 1) + for ind in range(len(self.convolutional_layers)) + ] + + # Molecule context attention + encodings = [] + for layer in range(len(self.molecule_heads)): + for head in range(self.molecule_heads[layer]): + ind = self.molecule_heads[0] * layer + head + e, _ = self.molecule_attention_layers[ind](encoded_smiles[layer], gep) + encodings.append(e) + + # Gene context attention + for layer in range(len(self.gene_heads)): + for head in range(self.gene_heads[layer]): + ind = self.gene_heads[0] * layer + head + e, _ = self.gene_attention_layers[ind](gep, encoded_smiles[layer], average_seq=False) + encodings.append(e) + + # concat features + encodings = torch.cat(encodings, dim=1) + + # Apply batch normalization if specified + inputs = self.batch_norm(encodings) if self.params.get("batch_norm", False) else encodings + # NOTE: stacking dense layers as a bottleneck + for dl in self.dense_layers: + inputs = dl(inputs) + + # prediction + return self.final_dense(inputs) + + def loss(self, yhat, y): + """Compute the loss between predictions and targets. + + :param yhat: predicted values + :param y: true target values + :return: loss value + """ + return self.loss_fn(yhat, y) + + def load(self, path, *args, **kwargs): + """Load model from path. + + :param path: path to the saved model file + :param args: additional positioinal arguments passed to torch.load + :param kwargs: additional keyword arguments passed to torch.load + """ + weights = torch.load(path, *args, **kwargs) # noqa: S614 + self.load_state_dict(weights) + + def save(self, path, *args, **kwargs): + """Save model to path. + + :param path: path where the model should be saved + :param args: additional positional arguments passed to torch.save + :param kwargs: additional keyword arguments passed to torch.save + """ + torch.save(self.state_dict(), path, *args, **kwargs) diff --git a/drevalpy/models/PaccMann/utils/__init__.py b/drevalpy/models/PaccMann/utils/__init__.py new file mode 100644 index 000000000..2e0a46f6a --- /dev/null +++ b/drevalpy/models/PaccMann/utils/__init__.py @@ -0,0 +1 @@ +"""Utility modules for PaccMann.""" diff --git a/drevalpy/models/PaccMann/utils/hyperparams.py b/drevalpy/models/PaccMann/utils/hyperparams.py new file mode 100644 index 000000000..4c07605ab --- /dev/null +++ b/drevalpy/models/PaccMann/utils/hyperparams.py @@ -0,0 +1,25 @@ +"""Customizable model hyperparameters.""" + +import torch.nn as nn + +from drevalpy.models.PaccMann.utils.loss_functions import ( + correlation_coefficient_loss, + mse_cc_loss, +) + +LOSS_FN_FACTORY = { + "mse": nn.MSELoss(), + "l1": nn.L1Loss(), + "mse_and_pearson": mse_cc_loss, + "pearson": correlation_coefficient_loss, + "binary_cross_entropy": nn.BCELoss(), +} + +ACTIVATION_FN_FACTORY = { + "relu": nn.ReLU(), + "sigmoid": nn.Sigmoid(), + "selu": nn.SELU(), + "tanh": nn.Tanh(), + "lrelu": nn.LeakyReLU(), + "elu": nn.ELU(), +} diff --git a/drevalpy/models/PaccMann/utils/layers.py b/drevalpy/models/PaccMann/utils/layers.py new file mode 100644 index 000000000..cf4ed11ed --- /dev/null +++ b/drevalpy/models/PaccMann/utils/layers.py @@ -0,0 +1,217 @@ +"""Custom layers implementation.""" + +from collections import OrderedDict + +import torch +import torch.nn as nn + +from .utils import Squeeze, Temperature + + +def dense_layer( + input_size, + hidden_size, + act_fn=None, + batch_norm=False, + dropout=0.0, +): + """Build a dense layer block. + + :param input_size: Input feature size + :param hidden_size: Output feature size + :param act_fn: Activation module + :param batch_norm: whether batch normalization is applied + :param dropout: Dropout probability + :return: Sequential dense layer block + """ + if act_fn is None: + act_fn = nn.ReLU() + + return nn.Sequential( + OrderedDict( + [ + ("projection", nn.Linear(input_size, hidden_size)), + ( + "batch_norm", + nn.BatchNorm1d(hidden_size) if batch_norm else nn.Identity(), + ), + ("act_fn", act_fn), + ("dropout", nn.Dropout(p=dropout)), + ] + ) + ) + + +def convolutional_layer( + num_kernel, + kernel_size, + act_fn=None, + batch_norm=False, + dropout=0.0, + input_channels=1, +): + """Convolutional layer. + + :param num_kernel: number of convolution kernels + :param kernel_size: size of the convolution kernels + :param act_fn: activation module + :param batch_norm: whether batch normalization is applied + :param dropout: dropout probability + :param input_channels: number of input channels + :return: sequential convolutional layer block + """ + if act_fn is None: + act_fn = nn.ReLU() + + return nn.Sequential( + OrderedDict( + [ + ( + "convolve", + torch.nn.Conv2d( + input_channels, # channel_in + num_kernel, # channel_out + kernel_size, # kernel_size + padding=[kernel_size[0] // 2, 0], # pad for valid conv. + ), + ), + ("squeeze", Squeeze()), + ("act_fn", act_fn), + ("dropout", nn.Dropout(p=dropout)), + ( + "batch_norm", + nn.BatchNorm1d(num_kernel) if batch_norm else nn.Identity(), + ), + ] + ) + ) + + +class ContextAttentionLayer(nn.Module): + """Context attention layer used in the PaccMann architecture. + + It implements context attention as described in the PaccMann paper and + supports an optional hidden size in the context representation. + """ + + def __init__( + self, + reference_hidden_size: int, + reference_sequence_length: int, + context_hidden_size: int, + context_sequence_length: int = 1, + attention_size: int = 16, + individual_nonlinearity=None, + temperature: float = 1.0, + ): + """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 + :param context_sequence_length: sequence length of the context + :param attention_size: size of the attention space + :param individual_nonlinearity: optional activation module applied to each projection + :param temperature: temperature used for the softmax + """ + super().__init__() + + if individual_nonlinearity is None: + individual_nonlinearity = nn.Sequential() + + self.reference_sequence_length = reference_sequence_length + self.reference_hidden_size = reference_hidden_size + self.context_sequence_length = context_sequence_length + self.context_hidden_size = context_hidden_size + self.attention_size = attention_size + self.individual_nonlinearity = individual_nonlinearity + self.temperature = temperature + + # Project the reference into the attention space + self.reference_projection = nn.Sequential( + OrderedDict( + [ + ( + "projection", + nn.Linear(reference_hidden_size, attention_size), + ), + ("act_fn", individual_nonlinearity), + ] + ) + ) + + # Project the context into the attention space + self.context_projection = nn.Sequential( + OrderedDict( + [ + ( + "projection", + nn.Linear(context_hidden_size, attention_size), + ), + ("act_fn", individual_nonlinearity), + ] + ) + ) + + # Optionally reduce the hidden size in context + if context_sequence_length > 1: + self.context_hidden_projection = nn.Sequential( + OrderedDict( + [ + ( + "projection", + nn.Linear( + context_sequence_length, + reference_sequence_length, + ), + ), + ("act_fn", individual_nonlinearity), + ] + ) + ) + else: + self.context_hidden_projection = nn.Sequential() + + self.alpha_projection = nn.Sequential( + OrderedDict( + [ + ("projection", nn.Linear(attention_size, 1, bias=False)), + ("squeeze", Squeeze()), + ("temperature", Temperature(self.temperature)), + ("softmax", nn.Softmax(dim=1)), + ] + ) + ) + + def forward( + self, + reference: torch.Tensor, + context: torch.Tensor, + average_seq: bool = True, + ): + """Forward pass through a context attention layer. + + :param reference: reference tensor of shape 'bs x ref_seq_length x ref_hidden_size' + :param context: context tensor of shape 'bs x context_seq_length x context_hidden_size' + :param average_seq: whether to average over the sequence length + :return: Tuple (output, attention_weights) + :raises ValueError: If reference or context is not 3-dimensional. + """ + if len(reference.shape) != 3: + raise ValueError("Reference tensor needs to be 3D") + + if len(context.shape) != 3: + raise ValueError("Context tensor needs to be 3D") + + 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)) + + output = reference * torch.unsqueeze(alphas, -1) + # Squeeze only the last dimension. A bare torch.squeeze would also drop the batch dimension + # for a batch of size 1, which breaks the concatenation of the encodings further downstream. + output = torch.sum(output, 1) if average_seq else torch.squeeze(output, -1) + + return output, alphas diff --git a/drevalpy/models/PaccMann/utils/loss_functions.py b/drevalpy/models/PaccMann/utils/loss_functions.py new file mode 100644 index 000000000..45a7e35da --- /dev/null +++ b/drevalpy/models/PaccMann/utils/loss_functions.py @@ -0,0 +1,61 @@ +"""Loss function definitions for PaccMann.""" + +import torch +import torch.nn as nn + + +def pearsonr(x, y): + """Compute Pearson correlation. + + :param x: 1D vector + :param y: 1D vector of the same size as x + :return: Pearson correlation coefficient + :raises TypeError: if inputs are not torch.Tensors + :raises ValueError: if inputs are not 1D, have different lengths, have length < 2 or are constant + """ + if not isinstance(x, torch.Tensor) or not isinstance(y, torch.Tensor): + raise TypeError("Function expects torch Tensors.") + + if len(x.shape) > 1 or len(y.shape) > 1: + raise ValueError("x and y must be 1D Tensors.") + + if len(x) != len(y): + raise ValueError("x and y must have the same length.") + + if len(x) < 2: + raise ValueError("x and y must have length at least 2.") + + # If an input is constant, the correlation coefficient is not defined. + if bool((x == x[0]).all()) or bool((y == y[0]).all()): + raise ValueError("Constant input, r is not defined.") + + mx = x - torch.mean(x) + my = y - torch.mean(y) + cost = torch.sum(mx * my) / (torch.sqrt(torch.sum(mx**2)) * torch.sqrt(torch.sum(my**2))) + return torch.clamp(cost, min=-1.0, max=1.0) + + +def correlation_coefficient_loss(labels, predictions): + """Compute loss based on Pearson correlation. + + :param labels: reference values + :param predictions: predicted values + :return: Loss value defined as 1 - r(labels, predictions)^2 + """ + return 1 - pearsonr(labels, predictions) ** 2 + + +def mse_cc_loss(labels, predictions): + """Compute loss based on MSE and Pearson correlation. + + The main assumption is that MSE lies in [0,1] range, i.e.: range is + comparable with Pearson correlation-based loss. + + :param labels: reference values + :param predictions: predicted values + :return: Loss defined as mse(labels, predictions) + 1 - r(labels, predictions)^2 + """ + mse_loss_fn = nn.MSELoss() + mse_loss = mse_loss_fn(predictions, labels) + cc_loss = correlation_coefficient_loss(labels, predictions) + return mse_loss + cc_loss diff --git a/drevalpy/models/PaccMann/utils/utils.py b/drevalpy/models/PaccMann/utils/utils.py new file mode 100644 index 000000000..f6627c2ab --- /dev/null +++ b/drevalpy/models/PaccMann/utils/utils.py @@ -0,0 +1,52 @@ +"""Utility functions.""" + +import torch +import torch.nn as nn + + +def get_device(): + """Return the active torch device. + + :return: torch.device("cuda") if cuda is available otherwise torch.device("cpu") + """ + return torch.device("cuda" if cuda() else "cpu") + + +def cuda(): + """Check whether cuda is available. + + :return: True if cuda is available otherwise False. + """ + return torch.cuda.is_available() + + +class Squeeze(nn.Module): + """Squeeze wrapper for nn.Sequential.""" + + def forward(self, data): + """Squeeze the last dimension of the input tensor. + + :param data: input tensor + :return: squeezed tensor + """ + return torch.squeeze(data, -1) + + +class Temperature(nn.Module): + """Temperature wrapper for nn.Sequential.""" + + def __init__(self, temperature): + """Initialize the temperature wrapper. + + :param temperature: Temperature value used for scaling. + """ + super().__init__() + self.temperature = temperature + + def forward(self, data): + """Scale the input tensor by the temperature value. + + :param data: input tensor + :return: scaled tensor + """ + return data / self.temperature diff --git a/drevalpy/models/__init__.py b/drevalpy/models/__init__.py index 334d407cb..d5fafd2d1 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 b01698f4c..f19898cdd 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 71ef279ec..a87cbe093 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 c97942527..c5950b804 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 ae633acd1..296d58abb 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