diff --git a/.vscode/settings.json b/.vscode/settings.json index dbebc3c5..f81b15ba 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -2,7 +2,7 @@ "python.testing.unittestArgs": [ "-v", "-s", - "./tests", + "./tests/unit", "-p", "test*.py" ], diff --git a/chebai/callbacks/epoch_metrics.py b/chebai/callbacks/epoch_metrics.py index 76d6a8fd..3edb20b6 100644 --- a/chebai/callbacks/epoch_metrics.py +++ b/chebai/callbacks/epoch_metrics.py @@ -1,5 +1,8 @@ +import warnings + import torch import torchmetrics +from sklearn.metrics import roc_auc_score def custom_reduce_fx(input: torch.Tensor) -> torch.Tensor: @@ -179,3 +182,83 @@ def compute(self) -> torch.Tensor: balanced_acc = (tpr + tnr) / 2 return torch.mean(balanced_acc) + + +class HiMolMacroAUROC(torchmetrics.Metric): + """ + Macro-averaged multilabel AUROC that exactly replicates HiMol's eval(): + - missing labels (== ignore_index) are excluded per-task before scoring + - any task that, after exclusion, has only one class present is + dropped from BOTH the sum and the divisor of the macro average + (instead of being scored as 0.0 and diluting the mean, which is + torchmetrics' default MultilabelAUROC behavior) + + Assumes preds/target are shape (N, num_labels), target values in {0, 1}, + with `ignore_index` marking missing/unlabeled entries. + + References: + https://github.com/ZangXuan/HiMol/blob/ffdcb247b361a1f85ddb741862cff25e4a3b3341/finetune/optimization.py#L95-L104 + """ + + full_state_update = False + is_differentiable = False + higher_is_better = True + + def __init__(self, num_labels: int, ignore_index: int = 0, **kwargs): + super().__init__(**kwargs) + self.num_labels = num_labels + self.ignore_index = ignore_index + + self.add_state("preds", default=[], dist_reduce_fx="cat") + self.add_state("target", default=[], dist_reduce_fx="cat") + + def update(self, preds: torch.Tensor, target: torch.Tensor) -> None: + if preds.shape != target.shape: + raise ValueError( + f"preds/target shape mismatch: {preds.shape} vs {target.shape}" + ) + if preds.ndim != 2 or preds.shape[1] != self.num_labels: + raise ValueError( + f"expected shape (N, {self.num_labels}), got {tuple(preds.shape)}" + ) + + self.preds.append(preds.detach().cpu()) + self.target.append(target.detach().cpu()) + + def compute(self) -> torch.Tensor: + preds = torch.cat(self.preds, dim=0).cpu().numpy() + target = torch.cat(self.target, dim=0).cpu().numpy() + + roc_list = [] + n_skipped = 0 + + for i in range(self.num_labels): + col_target = target[:, i] + col_preds = preds[:, i] + + valid = col_target != self.ignore_index + col_target = col_target[valid] + col_preds = col_preds[valid] + + # need at least one of each class to define AUC + if ( + len(col_target) == 0 + or (col_target == 0).sum() == 0 + or (col_target == 1).sum() == 0 + ): + n_skipped += 1 + continue + + roc_list.append(roc_auc_score(col_target, col_preds)) + + if n_skipped > 0: + warnings.warn( + f"{n_skipped}/{self.num_labels} labels skipped (missing or single-class " + f"after masking). Macro AUROC computed over {len(roc_list)} labels.", + stacklevel=2, + ) + + if len(roc_list) == 0: + return torch.tensor(float("nan")) + + return torch.tensor(sum(roc_list) / len(roc_list)) diff --git a/chebai/cli.py b/chebai/cli.py index d65dd51e..df355a98 100644 --- a/chebai/cli.py +++ b/chebai/cli.py @@ -75,6 +75,8 @@ def call_data_methods(data: Type[XYBaseDataModule]): "mse", "rmse", "r2", + "pr-auc", + "himol-marco-roc-auc", ): # When using lightning > 2.5.1 then need to uncomment all metrics that are not used # for average in ("mse", "rmse","r2"): # for regression diff --git a/chebai/loss/asymmetric_loss.py b/chebai/loss/asymmetric_loss.py index 9ed6f170..8b795907 100644 --- a/chebai/loss/asymmetric_loss.py +++ b/chebai/loss/asymmetric_loss.py @@ -9,10 +9,10 @@ class AsymmetricLoss(nn.Module): Asymmetric Loss for multi-label and single-label classification tasks. Implementation adapted from: https://github.com/Alibaba-MIIL/ASL - + Asymmetric Loss from: "Asymmetric Loss For Multi-Label Classification" https://openaccess.thecvf.com/content/ICCV2021/papers/Ben-Baruch_Asymmetric_Loss_For_Multi-Label_Classification_ICCV_2021_paper.pdf - + Args: gamma_neg (float): Negative focusing parameter. Default is 1. gamma_pos (float): Positive focusing parameter. Default is 1. @@ -47,12 +47,12 @@ def __init__( def forward(self, inputs, targets, **kwargs): """ Forward pass to compute the Asymmetric Loss. - + Args: inputs: Predictions (logits) from the model. targets: Ground truth labels. **kwargs: Additional keyword arguments (for compatibility with training framework). - + Returns: Loss tensor with the reduction option applied. """ @@ -68,7 +68,7 @@ def forward(self, inputs, targets, **kwargs): def _multi_label_asymmetric_loss(self, x, y): """ Standard asymmetric loss for multi-label classification. - + Parameters ---------- x: input logits @@ -112,12 +112,14 @@ def _multi_label_asymmetric_loss(self, x, y): def _single_label_asymmetric_loss(self, inputs, target): """ Asymmetric loss for single-label classification problems. - + "input" dimensions: - (batch_size, number_classes) "target" dimensions: - (batch_size) """ log_preds = self.logsoftmax(inputs) - self.targets_classes = torch.zeros_like(inputs).scatter_(1, target.long().unsqueeze(1), 1) + self.targets_classes = torch.zeros_like(inputs).scatter_( + 1, target.long().unsqueeze(1), 1 + ) # ASL weights targets = self.targets_classes @@ -125,13 +127,15 @@ def _single_label_asymmetric_loss(self, inputs, target): xs_pos = torch.exp(log_preds) xs_neg = 1 - xs_pos grad_ctx = ( - torch.no_grad() if self.disable_torch_grad_focal_loss else nullcontext() - ) + torch.no_grad() if self.disable_torch_grad_focal_loss else nullcontext() + ) with grad_ctx: xs_pos = xs_pos * targets xs_neg = xs_neg * anti_targets - asymmetric_w = torch.pow(1 - xs_pos - xs_neg, - self.gamma_pos * targets + self.gamma_neg * anti_targets) + asymmetric_w = torch.pow( + 1 - xs_pos - xs_neg, + self.gamma_pos * targets + self.gamma_neg * anti_targets, + ) log_preds = log_preds * asymmetric_w # loss calculation diff --git a/chebai/loss/base.py b/chebai/loss/base.py new file mode 100644 index 00000000..94834eec --- /dev/null +++ b/chebai/loss/base.py @@ -0,0 +1,29 @@ +from typing import Optional + +import torch + + +class BCELogitLossWithValidLabels(torch.nn.BCEWithLogitsLoss): + def __init__(self, **kwargs): + kwargs["reduction"] = "none" + super().__init__(**kwargs) + + def forward( + self, + input: torch.Tensor, + target: torch.Tensor, + valid_label_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + loss_mat = super().forward(input, target) + + if valid_label_mask is None: + return loss_mat.mean() + + loss_mat = torch.where( + valid_label_mask, + loss_mat, + torch.zeros_like(loss_mat), + ) + + return loss_mat.sum() / valid_label_mask.sum().clamp_min(1) diff --git a/chebai/loss/bce_weighted.py b/chebai/loss/bce_weighted.py index 440495cf..7453321b 100644 --- a/chebai/loss/bce_weighted.py +++ b/chebai/loss/bce_weighted.py @@ -3,11 +3,12 @@ import torch +from chebai.loss.base import BCELogitLossWithValidLabels from chebai.preprocessing.datasets.base import XYBaseDataModule from chebai.preprocessing.datasets.chebi import _ChEBIDataExtractor -class BCEWeighted(torch.nn.BCEWithLogitsLoss): +class BCEWeighted(BCELogitLossWithValidLabels): """ BCEWithLogitsLoss with weights automatically computed according to the beta parameter. @@ -104,10 +105,9 @@ def forward( torch.Tensor: The computed loss. """ self.set_pos_weight(input) - return super().forward(input, target) + return super().forward(input, target, **kwargs) -class UnWeightedBCEWithLogitsLoss(torch.nn.BCEWithLogitsLoss): +class UnWeightedBCEWithLogitsLoss(BCELogitLossWithValidLabels): def forward(self, input, target, **kwargs): - # As the custom passed kwargs are not used in BCEWithLogitsLoss, we can ignore them - return super().forward(input, target) + return super().forward(input, target, **kwargs) diff --git a/chebai/models/base.py b/chebai/models/base.py index 58d387a9..48412b16 100644 --- a/chebai/models/base.py +++ b/chebai/models/base.py @@ -295,6 +295,9 @@ def _execute( model_output, labels, data.get("loss_kwargs", dict()) ) loss_kwargs = dict() + loss_kwargs["valid_label_mask"] = loss_kwargs_candidates[ + "valid_label_mask" + ] if self.pass_loss_kwargs: loss_kwargs = loss_kwargs_candidates loss_kwargs["current_epoch"] = self.trainer.current_epoch diff --git a/chebai/preprocessing/collate.py b/chebai/preprocessing/collate.py index 06b4ff34..3bfb3d70 100644 --- a/chebai/preprocessing/collate.py +++ b/chebai/preprocessing/collate.py @@ -61,10 +61,13 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: """ Collate ragged data samples (i.e., samples of unequal size, such as molecular sequences) into a batch. - Handles both fully and partially labeled data, where some samples may have `None` as their label. The indices - of non-null labels are stored in the `non_null_labels` field, which is used to filter out predictions for - unlabeled data during evaluation (e.g., F1, MSE). For models supporting partially labeled data, this method - ensures alignment between features and labels. Missing labels are passed as a loss keyword. + Handles both fully and partially labeled data by use of the following fields in the returned XYData: + + `non_null_labels`: Stores batch row indices of samples where the whole `labels` field is not None, like [0, 2]. + - Example: [[True, False], None, [False, None]] would result in `non_null_labels` = [0, 2]. + - This is used to filter out predictions for unlabeled samples during evaluation. + + `valid_label_mask`: Stores a per-sample, per-label-position boolean mask for valid entries inside a label row, like the True in [1, None, 0]. Args: data (List[Union[Dict, Tuple]]): List of ragged data samples. Each sample can be a dictionary or tuple @@ -81,30 +84,25 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: if isinstance(data[0], tuple): # For legacy data x, y, idents = zip(*data) - missing_labels = None else: x, y, idents = zip( *((d["features"], d["labels"], d.get("ident")) for d in data) ) - missing_labels = [ - d.get( - "missing_labels", - [False for _ in y[0]] if y[0] is not None else [False], - ) - for d in data - ] + valid_label_mask = self._get_valid_label_mask(y) + + # Typical y: ([True, False], None, [True, None], [True]) if any(x is not None for x in y): - # If any label is not None: (None, None, `1`, None) + # If any label is not None: (None, None, `[True, None]`, None) if any(x is None for x in y): - # If any label is None: (`None`, `None`, 1, `None`) + # If any label is None: (`None`, [True, False], [True], [False]) non_null_labels = [i for i, r in enumerate(y) if r is not None] y = self.process_label_rows( tuple(ye for i, ye in enumerate(y) if i in non_null_labels) ) loss_kwargs["non_null_labels"] = non_null_labels else: - # If all labels are not None: (`0`, `2`, `1`, `3`) + # If all labels are not None: (`[True, False]`, `[False, True, True]`, `[False]`, `[True]`) y = self.process_label_rows(y) else: @@ -112,7 +110,7 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData: y = None loss_kwargs["non_null_labels"] = [] - loss_kwargs["missing_labels"] = torch.tensor(missing_labels) + loss_kwargs["valid_label_mask"] = valid_label_mask # Calculate the lengths of each sequence, create a binary mask for valid (non-padded) positions lens = torch.tensor(list(map(len, x))) model_kwargs["mask"] = torch.arange(max(lens))[None, :] < lens[:, None] @@ -146,3 +144,26 @@ def process_label_rows(self, labels: Tuple) -> torch.Tensor: ], batch_first=True, ) + + def _get_valid_label_mask(self, y: Tuple) -> torch.Tensor | None: + # Compute the per-sample, per-label-position boolean mask for unknown entries + # (e.g., the None in [1, None, 0]) on the *original* labels, before any + # filtering/padding is applied to `y`. Rows whose entire label is None are + # represented as all-False rows of the maximum label length. + if any(labels is not None for labels in y): + max_label_len = max(len(labels) for labels in y if labels is not None) + valid_label_mask = pad_sequence( + [ + torch.tensor([label is not None for label in labels]) + if labels is not None + else torch.zeros(max_label_len, dtype=torch.bool) + for labels in y + ], + batch_first=True, + ) + if (~valid_label_mask).sum() != 0: + # If there are any invalid labels, return the valid_label_mask + # Else, return None to indicate that all labels are valid (no None entries). + return valid_label_mask + + return None diff --git a/chebai/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index b3110783..0ea7f453 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -2,23 +2,20 @@ import random from abc import ABC, abstractmethod from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, List, Optional, Tuple, Union import lightning as pl import numpy as np import pandas as pd -from rdkit import Chem import torch import tqdm from lightning.pytorch.core.datamodule import LightningDataModule from lightning_utilities.core.rank_zero import rank_zero_info +from rdkit import Chem from torch.utils.data import DataLoader from chebai.preprocessing import reader as dr -if TYPE_CHECKING: - import networkx as nx - class XYBaseDataModule(LightningDataModule): """ @@ -37,7 +34,6 @@ class XYBaseDataModule(LightningDataModule): label_filter (Optional[int]): The index of the label to filter. Default is None. balance_after_filter (Optional[float]): The ratio of negative samples to positive samples after filtering. Default is None. num_workers (int): The number of worker processes for data loading. Default is 1. - chebi_version (int): The version of ChEBI to use. Default is 200. inner_k_folds (int): The number of folds for inner cross-validation. Use -1 to disable inner cross-validation. Default is -1. fold_index (Optional[int]): The index of the fold to use for training and validation. Default is None. base_dir (Optional[str]): The base directory for storing processed and raw data. Default is None. @@ -54,7 +50,6 @@ class XYBaseDataModule(LightningDataModule): label_filter (Optional[int]): The index of the label to filter. balance_after_filter (Optional[float]): The ratio of negative samples to positive samples after filtering. num_workers (int): The number of worker processes for data loading. - chebi_version (int): The version of ChEBI to use. inner_k_folds (int): The number of folds for inner cross-validation. If it is less than to, no cross-validation will be performed. fold_index (Optional[int]): The index of the fold to use for training and validation (only relevant for cross-validation). _base_dir (Optional[str]): The base directory for storing processed and raw data. @@ -70,8 +65,8 @@ class XYBaseDataModule(LightningDataModule): def __init__( self, batch_size: int = 1, - test_split: Optional[float] = 0.1, - validation_split: Optional[float] = 0.05, + test_split: float = 0.1, + validation_split: float = 0.05, reader_kwargs: Optional[dict] = None, prediction_kind: str = "test", data_limit: Optional[int] = None, @@ -79,7 +74,6 @@ def __init__( balance_after_filter: Optional[float] = None, num_workers: int = 1, persistent_workers: bool = True, - chebi_version: int = 200, inner_k_folds: int = -1, # use inner cross-validation if > 1 fold_index: Optional[int] = None, base_dir: Optional[str] = None, @@ -103,7 +97,6 @@ def __init__( self.balance_after_filter = balance_after_filter self.num_workers = num_workers self.persistent_workers: bool = bool(persistent_workers) - self.chebi_version = chebi_version assert type(inner_k_folds) is int self.inner_k_folds = inner_k_folds self.use_inner_cross_validation = ( @@ -284,6 +277,12 @@ def dataloader(self, kind: str, **kwargs) -> DataLoader: random.shuffle(dataset) if self.data_limit is not None: dataset = dataset[: self.data_limit] + + if len(dataset) == 0: + raise ValueError( + f"Dataset is empty for {kind} data.\nPlease check the data preparation and filtering steps.", + ) + return DataLoader( dataset, collate_fn=self.reader.collator, @@ -342,6 +341,13 @@ def _load_data_from_file(self, path: str) -> List[Dict[str, Any]]: if d["features"] is not None ] + number_of_unique_ids = len(set(d["ident"] for d in data)) + assert len(data) == number_of_unique_ids, ( + "Duplicate entries found in the dataset. " + f"Unique entries {number_of_unique_ids}. " + f"Total entries {len(data)}. " + ) + data = [val for val in data if self._filter_to_token_limit(val)] return data @@ -642,6 +648,14 @@ def classes_txt_file_path(self) -> Optional[str]: # - chebai/cli.py: to link this property to `model.init_args.classes_txt_file_path` return None + @property + def data_type(self) -> str: + """ + Returns the type of data (e.g., chebi, protein, HIV, Tox21, etc.) that the dataset represents. + This property is used to create a separate tokens directory for each data type. + """ + raise NotImplementedError + class MergedDataset(XYBaseDataModule): MERGED = [] @@ -816,7 +830,7 @@ class _DynamicDataset(XYBaseDataModule, ABC): apply_id_filter (Optional[str]): Path to a data.pt file for ID filtering. """ - # ---- Index for columns of processed `data.pkl` (should be derived from `_graph_to_raw_dataset` method) ------ + # ---- Index for columns of processed `data.pkl` (should be derived from `_preprocess_data_into_dataframe` method) ------ _ID_IDX: int = None _DATA_REPRESENTATION_IDX: int = None _LABELS_START_IDX: int = None @@ -910,10 +924,7 @@ def _perform_data_preparation(self, *args: Any, **kwargs: Any) -> None: print(f"Missing processed data file (`{processed_name}` file)") os.makedirs(self.processed_dir_main, exist_ok=True) data_path = self._download_required_data() - from chebi_utils import build_chebi_graph - - g = build_chebi_graph(data_path) - data_df = self._graph_to_raw_dataset(g) + data_df = self._preprocess_data_into_dataframe(data_path) self.save_processed(data_df, processed_name) @abstractmethod @@ -927,17 +938,15 @@ def _download_required_data(self) -> str: pass @abstractmethod - def _graph_to_raw_dataset(self, graph: "nx.DiGraph") -> pd.DataFrame: + def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: """ - Converts the graph to a raw dataset. - Uses the graph created by chebi_utils to extract the - raw data in Dataframe format with additional columns corresponding to each multi-label class. + Preprocesses the raw data into a DataFrame. Args: - graph (nx.DiGraph): The class hierarchy graph. + raw_data_path (str): Path to the raw data. Returns: - pd.DataFrame: The raw dataset. + pd.DataFrame: The preprocessed data as a DataFrame. """ pass @@ -949,7 +958,8 @@ def save_processed(self, data: pd.DataFrame, filename: str) -> None: data (pd.DataFrame): The processed dataset to be saved. filename (str): The filename for the pickle file. """ - pd.to_pickle(data, open(os.path.join(self.processed_dir_main, filename), "wb")) + if data is not None and not data.empty: + data.to_pickle(os.path.join(self.processed_dir_main, filename)) def get_processed_pickled_df_file(self, filename: str) -> Optional[pd.DataFrame]: """ @@ -972,7 +982,7 @@ def setup_processed(self) -> None: Transforms `data.pkl` into a model input data format (`data.pt`), ensuring that the data is in a format compatible for input to the model. The transformed data contains the following keys: `ident`, `features`, `labels`, and `group`. - This method uses a subclass of Data Reader to perform the transformation. + This method uses assigned subclass of `DataReader` to perform the transformation. Returns: None @@ -1107,6 +1117,7 @@ def _retrieve_splits_from_csv(self) -> None: splits.csv to reconstruct the train, validation, and test splits. """ print(f"\nLoading splits from {self.splits_file_path}...") + assert self.splits_file_path is not None, "splits_file_path should not be None" splits_df = pd.read_csv(self.splits_file_path) filename = self.processed_file_names_dict["data"] @@ -1144,6 +1155,15 @@ def _retrieve_splits_from_csv(self) -> None: self._dynamic_df_train = df_data[df_data["ident"].isin(train_ids)] self._dynamic_df_val = df_data[df_data["ident"].isin(validation_ids)] self._dynamic_df_test = df_data[df_data["ident"].isin(test_ids)] + assert len(self._dynamic_df_train) > 0, ( + "No training data found after applying splits" + ) + assert len(self._dynamic_df_val) > 0, ( + "No validation data found after applying splits" + ) + assert len(self._dynamic_df_test) > 0, ( + "No test data found after applying splits" + ) # ------------------------------ Phase: DataLoaders ----------------------------------- def load_processed_data( diff --git a/chebai/preprocessing/datasets/chebi.py b/chebai/preprocessing/datasets/chebi.py index d152f4f1..7c1709a4 100644 --- a/chebai/preprocessing/datasets/chebi.py +++ b/chebai/preprocessing/datasets/chebi.py @@ -9,19 +9,20 @@ from itertools import cycle, permutations, product from typing import TYPE_CHECKING, Any, Generator, List, Literal, Optional -from networkx import DiGraph import numpy as np import pandas as pd +from networkx import DiGraph from rdkit import Chem from chebai.preprocessing import reader as dr from chebai.preprocessing.datasets.base import _DynamicDataset +from chebai.preprocessing.splitters import MultiLabelSplitter if TYPE_CHECKING: import networkx as nx -class _ChEBIDataExtractor(_DynamicDataset, ABC): +class _ChEBIDataExtractor(MultiLabelSplitter, _DynamicDataset, ABC): """ A class for extracting and processing data from the ChEBI dataset. @@ -51,6 +52,7 @@ class _ChEBIDataExtractor(_DynamicDataset, ABC): def __init__( self, + chebi_version: int = 241, chebi_version_train: Optional[int] = None, single_class: Optional[int] = None, subset: Optional[Literal["2_STAR", "3_STAR"]] = None, @@ -58,6 +60,7 @@ def __init__( aug_smiles_variations: Optional[int] = None, **kwargs, ): + self.chebi_version = chebi_version if bool(augment_smiles): assert int(aug_smiles_variations) > 0, ( "Number of variations must be greater than 0" @@ -80,6 +83,7 @@ def __init__( self.subset = subset super(_ChEBIDataExtractor, self).__init__(**kwargs) + # use different version of chebi for training and validation (if not None) # (still uses self.chebi_version for test set) self.chebi_version_train = chebi_version_train @@ -150,6 +154,21 @@ def _download_required_data(self) -> str: self._load_sdf() return self._load_chebi() + def _preprocess_data_into_dataframe(self, raw_data_path: str) -> pd.DataFrame: + """ + Preprocesses the raw data into a DataFrame. + + Args: + raw_data_path (str): Path to the raw data. + + Returns: + pd.DataFrame: The preprocessed data as a DataFrame. + """ + from chebi_utils import build_chebi_graph + + g = build_chebi_graph(raw_data_path) + return self._graph_to_raw_dataset(g) + def _load_chebi(self, version: Optional[int] = None) -> str: """ Load the ChEBI ontology file. @@ -381,27 +400,6 @@ def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, No for feat, labels, ident in zip(features, all_labels, idents): yield dict(features=feat, labels=labels, ident=ident) - def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: - """ - Loads encoded/transformed data and generates training, validation, and test splits. - """ - - filename = self.processed_file_names_dict["data"] - data = self.load_processed_data_from_file(filename) - df_data = pd.DataFrame(data) - - from chebi_utils import create_multilabel_splits - - splits = create_multilabel_splits( - df_data, - self._LABELS_START_IDX, - 1 - self.validation_split - self.test_split, - self.validation_split, - self.test_split, - self.dynamic_data_split_seed, - ) - return splits["train"], splits["validation"], splits["test"] - def _setup_pruned_test_set( self, df_test_chebi_version: pd.DataFrame ) -> pd.DataFrame: @@ -539,6 +537,14 @@ def classes_txt_file_path(self) -> str: # - chebai/cli.py: to link this property to `model.init_args.classes_txt_file_path` return os.path.join(self.processed_dir_main, "classes.txt") + @property + def data_type(self) -> str: + """ + Returns the type of data (e.g., chebi, protein, HIV, Tox21, etc.) that the dataset represents. + This property is used to create a separate tokens directory for each data type. + """ + return "chebi" + class ChEBIFromList(_ChEBIDataExtractor): """ @@ -657,6 +663,21 @@ class ChEBIOver50(ChEBIOverX): THRESHOLD: int = 50 +class ChEBIOver50_ChemDataReader(ChEBIOverX): + """ + A class for extracting data from the ChEBI dataset with a threshold of 50 for selecting classes. + + Inherits from ChEBIOverX. + + Attributes: + THRESHOLD (int): The threshold for selecting classes (50). + """ + + READER = dr.ChemDataReader + + THRESHOLD: int = 50 + + class ChEBIOver25(ChEBIOverX): """ A class for extracting data from the ChEBI dataset with a threshold of 25 for selecting classes. @@ -753,12 +774,12 @@ def _graph_to_raw_dataset(self, g: "nx.DiGraph") -> pd.DataFrame: """ # Extract mol objects from SDF using chebi-utils + import networkx as nx from chebi_utils import ( build_labeled_dataset, extract_molecules, get_hierarchy_subgraph, ) - import networkx as nx sdf_path = os.path.join(self.raw_dir, self.raw_file_names_dict["sdf"]) mol_df = extract_molecules(sdf_path) diff --git a/chebai/preprocessing/datasets/molecule_classification.py b/chebai/preprocessing/datasets/molecule_classification.py deleted file mode 100644 index c2916675..00000000 --- a/chebai/preprocessing/datasets/molecule_classification.py +++ /dev/null @@ -1,1052 +0,0 @@ -import csv -import gzip -import os -import shutil -from tempfile import NamedTemporaryFile -from typing import Dict, List -from urllib import request - -import numpy as np -import torch -from sklearn.model_selection import GroupShuffleSplit, train_test_split - -from chebai.preprocessing import reader as dr -from chebai.preprocessing.datasets.base import XYBaseDataModule - - -class ClinTox(XYBaseDataModule): - """Data module for ClinTox MoleculeNet dataset.""" - - HEADERS = [ - "FDA_APPROVED", - "CT_TOX", - ] - - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "ClinTox" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 2 - - @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["clintox.csv"] - - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: - """Downloads and extracts the dataset.""" - with NamedTemporaryFile("rb") as gout: - request.urlretrieve( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/clintox.csv.gz", - gout.name, - ) - with gzip.open(gout.name) as gfile: - with open(os.path.join(self.raw_dir, "clintox.csv"), "wt") as fout: - fout.write(gfile.read().decode()) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list( - self._load_data_from_file(os.path.join(self.raw_dir, "clintox.csv")) - ) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d for d in (data[temp_split_index[i]] for i in test_split_index) - ] - validation_split = [ - d for d in (data[temp_split_index[i]] for i in validation_split_index) - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["smiles"] - labels = [ - bool(int(label)) if label else None - for label in (row[k] for k in self.HEADERS) - ] - # group = int(row["group"]) - yield dict( - features=smiles, - labels=labels, - ident=i, - # group=group - ) - # yield dict(features=smiles, labels=labels, ident=i) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class BBBP(XYBaseDataModule): - """Data module for ClinTox MoleculeNet dataset.""" - - HEADERS = [ - "p_np", - ] - - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "BBBP" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 1 - - @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["bbbp.csv"] - - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: - """Downloads and extracts the dataset.""" - with open(os.path.join(self.raw_dir, "bbbp.csv"), "ab") as dst: - with request.urlopen( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/BBBP.csv", - ) as src: - shutil.copyfileobj(src, dst) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "bbbp.csv"))) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - print("Group shuffled") - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d - for d in (data[temp_split_index[i]] for i in test_split_index) - # if d["original"] - ] - validation_split = [ - d - for d in (data[temp_split_index[i]] for i in validation_split_index) - # if d["original"] - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["smiles"] - labels = [int(row["p_np"])] - # group = int(row["group"]) - yield dict( - features=smiles, - labels=labels, - ident=i, - # , group=group - ) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class Sider(XYBaseDataModule): - """Data module for ClinTox MoleculeNet dataset.""" - - HEADERS = [ - "Hepatobiliary disorders", - "Metabolism and nutrition disorders", - "Product issues", - "Eye disorders", - "Investigations", - "Musculoskeletal and connective tissue disorders", - "Gastrointestinal disorders", - "Social circumstances", - "Immune system disorders", - "Reproductive system and breast disorders", - "Neoplasms benign, malignant and unspecified (incl cysts and polyps)", - "General disorders and administration site conditions", - "Endocrine disorders", - "Surgical and medical procedures", - "Vascular disorders", - "Blood and lymphatic system disorders", - "Skin and subcutaneous tissue disorders", - "Congenital, familial and genetic disorders", - "Infections and infestations", - "Respiratory, thoracic and mediastinal disorders", - "Psychiatric disorders", - "Renal and urinary disorders", - "Pregnancy, puerperium and perinatal conditions", - "Ear and labyrinth disorders", - "Cardiac disorders", - "Nervous system disorders", - "Injury, poisoning and procedural complications", - ] - - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "Sider" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 27 - - @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["sider.csv"] - - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: - """Downloads and extracts the dataset.""" - with NamedTemporaryFile("rb") as gout: - request.urlretrieve( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/sider.csv.gz", - gout.name, - ) - with gzip.open(gout.name) as gfile: - with open(os.path.join(self.raw_dir, "sider.csv"), "wt") as fout: - fout.write(gfile.read().decode()) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "sider.csv"))) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d - for d in (data[temp_split_index[i]] for i in test_split_index) - # if d["original"] - ] - validation_split = [ - d - for d in (data[temp_split_index[i]] for i in validation_split_index) - # if d["original"] - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["smiles"] - labels = [ - bool(int(label)) if label else None - for label in (row[k] for k in self.HEADERS) - ] - # group = row["group"] - yield dict( - features=smiles, - labels=labels, - ident=i, - # , group=group - ) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class Bace(XYBaseDataModule): - """Data module for ClinTox MoleculeNet dataset.""" - - HEADERS = [ - "class", - ] - - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "Bace" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 1 - - @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["bace.csv"] - - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: - """Downloads and extracts the dataset.""" - with open(os.path.join(self.raw_dir, "bace.csv"), "ab") as dst: - with request.urlopen( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/bace.csv", - ) as src: - shutil.copyfileobj(src, dst) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "bace.csv"))) - # groups = np.array([d.get("group") for d in data]) - - # if not all(g is None for g in groups): - # split_size = int(len(set(groups)) * (1 - self.test_split - self.validation_split)) - # os.makedirs(self.processed_dir, exist_ok=True) - # splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - # train_split_index, temp_split_index = next( - # splitter.split(data, groups=groups) - # ) - - # split_groups = groups[temp_split_index] - - # splitter = GroupShuffleSplit( - # train_size=int(len(set(split_groups)) * (1 - self.test_split - self.validation_split)), n_splits=1 - # ) - # test_split_index, validation_split_index = next( - # splitter.split(temp_split_index, groups=split_groups) - # ) - # train_split = [data[i] for i in train_split_index] - # test_split = [ - # d - # for d in (data[temp_split_index[i]] for i in test_split_index) - # ] - # validation_split = [ - # d - # for d in (data[temp_split_index[i]] for i in validation_split_index) - # ] - # else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["mol"] - labels = [int(row["Class"])] - # group = row["group"] - yield dict(features=smiles, labels=labels, ident=i) # , group=group - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class HIV(XYBaseDataModule): - """Data module for ClinTox MoleculeNet dataset.""" - - HEADERS = [ - "HIV_active", - ] - - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "HIV" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 1 - - @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["hiv.csv"] - - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: - """Downloads and extracts the dataset.""" - with open(os.path.join(self.raw_dir, "hiv.csv"), "ab") as dst: - with request.urlopen( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/HIV.csv", - ) as src: - shutil.copyfileobj(src, dst) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "hiv.csv"))) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - print("Group shuffled") - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d for d in (data[temp_split_index[i]] for i in test_split_index) - ] - validation_split = [ - d for d in (data[temp_split_index[i]] for i in validation_split_index) - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - if len(row) > 1: - i += 1 - smiles = row["smiles"] - labels = [int(row["HIV_active"])] - # group = int(row["group"]) - yield dict( - features=smiles, - labels=labels, - ident=i, - # , group=group - ) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class MUV(XYBaseDataModule): - """Data module for ClinTox MoleculeNet dataset.""" - - HEADERS = [ - "MUV-466", - "MUV-548", - "MUV-600", - "MUV-644", - "MUV-652", - "MUV-689", - "MUV-692", - "MUV-712", - "MUV-713", - "MUV-733", - "MUV-737", - "MUV-810", - "MUV-832", - "MUV-846", - "MUV-852", - "MUV-858", - "MUV-859", - ] - - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "MUV" - - @property - def label_number(self) -> int: - """Returns the number of labels.""" - return 17 - - @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["muv.csv"] - - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: - """Downloads and extracts the dataset.""" - with NamedTemporaryFile("rb") as gout: - request.urlretrieve( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/muv.csv.gz", - gout.name, - ) - with gzip.open(gout.name) as gfile: - with open(os.path.join(self.raw_dir, "muv.csv"), "wt") as fout: - fout.write(gfile.read().decode()) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "muv.csv"))) - groups = np.array([d["group"] for d in data]) - if not all(g is None for g in groups): - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d - for d in (data[temp_split_index[i]] for i in test_split_index) - # if d["original"] - ] - validation_split = [ - d - for d in (data[temp_split_index[i]] for i in validation_split_index) - # if d["original"] - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - self._after_setup() - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - i = 0 - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - i += 1 - smiles = row["smiles"] - labels = [ - bool(int(label)) if label else None - for label in (row[k] for k in self.HEADERS) - ] - # group = row["group"] - yield dict(features=smiles, labels=labels, ident=i) # , group=group) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=i)) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - -class BaceChem(Bace): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class SiderChem(Sider): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class BBBPChem(BBBP): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class ClinToxChem(ClinTox): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class HIVChem(HIV): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader - - -class MUVChem(MUV): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader diff --git a/chebai/preprocessing/datasets/molecule_net_classification.py b/chebai/preprocessing/datasets/molecule_net_classification.py new file mode 100644 index 00000000..0eb0ac0b --- /dev/null +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -0,0 +1,374 @@ +import os +from abc import ABC, abstractmethod +from typing import Any, Generator + +import deepchem as dc +import pandas as pd +from deepchem.data import DiskDataset + +from chebai.preprocessing import reader as dr +from chebai.preprocessing.datasets.base import _DynamicDataset + + +class MoleculeNetDataExtractor(_DynamicDataset, ABC): + """ + Base class for MoleculeNet dataset extraction and preprocessing. + + Reference: + - https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html + - Zhenqin Wu, Bharath Ramsundar, Evan N. Feinberg, Joseph Gomes, Caleb Geniesse, + Aneesh S. Pappu, Karl Leswing, Vijay Pande; MoleculeNet: a benchmark for molecular + machine learning. Chem. Sci. 2018; 9 (2): 513–530. https://doi.org/10.1039/c7sc02664a + """ + + def __init__( + self, + test_split: float | None = None, + validation_split: float | None = None, + **kwargs, + ): + if test_split is not None or validation_split is not None: + raise ValueError( + "Custom splits are not supported for MoleculeNet datasets. " + "Please use the predefined splits provided by the deepchem community" + "by using `--splits_file_path=`" + ) + super().__init__( + test_split=test_split, validation_split=validation_split, **kwargs + ) + + READER = dr.ChemDataReader + + def _preprocess_data_into_dataframe(self, raw_data_path: str) -> None: + pass + + def _download_required_data(self) -> None: + pass + + def save_processed(self, data: pd.DataFrame, filename: str) -> None: + """ + Save the processed dataset to a pickle file. + + Args: + data (pd.DataFrame): The processed dataset to be saved. + filename (str): The filename for the pickle file. + """ + if data is not None: + data.to_pickle(os.path.join(self.processed_dir_main, filename)) + + def _get_data_size(self, input_file_path: str) -> None: + pass + + def _load_dict(self, input_file_path: str) -> Generator[dict[str, Any], None, None]: + """Loads data from a CSV file. + + Args: + input_file_path (str): Path to the CSV file. + + Returns: + List[Dict]: List of data dictionaries. + """ + splits = [] + train, valid, test = self._deep_chem_data_loader_api() + idx = 0 + for split_name, data in [ + ("train", train), + ("validation", valid), + ("test", test), + ]: + for mol, labels, w, smiles in data.itersamples(): + # https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html + # Note that the “w” matrix represents the weight of each sample. Some assays may have missing values, in which case the weight is 0. + # Otherwise, the weight is 1. This is when `transformers` are set to `[]` in the deepchem data loader API. + # By default `transformers` is set to ['balancing'], which doesn't hand you the raw 0/1 weight matrix. + # Instead meaning DeepChem automatically applies a BalancingTransformer before returning the dataset. + # That transformer reweights the observed labels per task so positive and negative examples end up with equal total weight — it upweights the rarer class. + # Currently, transformers are set to `[]` in the deepchem data loader API, so we can get the raw 0/1 weight matrix. + labels = [ + bool(label) if int(weight) != 0 else None + for label, weight in zip(labels, w) + ] + yield dict(features=mol, labels=labels, ident=idx) + splits.append({"id": idx, "split": split_name}) + idx += 1 + splits_file_path = os.path.join(self.processed_dir_main, "splits.csv") + if not os.path.exists(splits_file_path): + splits_df = pd.DataFrame(splits) + splits_df.to_csv(splits_file_path, index=False) + + @abstractmethod + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + pass + + def _get_data_splits(self) -> None: + pass + + def _generate_dynamic_splits(self) -> None: + raise ValueError( + "Custom splits are not supported for MoleculeNet datasets. " + "Please use the predefined splits provided by the deepchem community" + "by using `--splits_file_path=`" + ) + + @property + def base_dir(self) -> str: + """ + Return the base directory path for data. + + Returns: + str: The base directory path for data. + """ + return os.path.join("data", f"{self._name}:MNClassification") + + @property + def raw_file_names_dict(self) -> None: + """Returns a dictionary of raw file names.""" + pass + + +class ClinTox(MoleculeNetDataExtractor): + """Data module for ClinTox MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_clintox( + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "clin_tox" + + @property + def _name(self) -> str: + return "ClinTox" + + +class BBBP(MoleculeNetDataExtractor): + """Data module for BBBP MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Scaffold splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_bbbp( + featurizer="Raw", + splitter="scaffold", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "bbbp" + + @property + def _name(self) -> str: + return "BBBP" + + +class SIDER(MoleculeNetDataExtractor): + """Data module for Sider MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_sider( + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "sider" + + @property + def _name(self) -> str: + return "SIDER" + + +class BACE(MoleculeNetDataExtractor): + """Data module for Bace MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Scaffold splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_bace_classification( + featurizer="Raw", + splitter="scaffold", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "bace" + + @property + def _name(self) -> str: + return "BACE" + + +class HIV(MoleculeNetDataExtractor): + """Data module for HIV MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Scaffold splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_hiv( + featurizer="Raw", + splitter="scaffold", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "hiv" + + @property + def _name(self) -> str: + return "HIV" + + +class MUV(MoleculeNetDataExtractor): + """Data module for MUV MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Scaffold splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_muv( + featurizer="Raw", + splitter="scaffold", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "muv" + + @property + def _name(self) -> str: + return "MUV" + + +class Tox21(MoleculeNetDataExtractor): + """Data module for Tox21MolNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_tox21( + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "tox21" + + @property + def _name(self) -> str: + return "Tox21" + + +class ToxCast(MoleculeNetDataExtractor): + """Data module for ToxCast MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_toxcast( + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "toxcast" + + @property + def _name(self) -> str: + return "ToxCast" + + +class PCBA(MoleculeNetDataExtractor): + """Data module for PCBA MoleculeNet dataset.""" + + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + # Random splitting is recommended for this dataset. + tasks, datasets, transformers = dc.molnet.load_pcba( + featurizer="Raw", + splitter="random", + data_dir=self.raw_dir, + save_dir=self.processed_dir_main, + transformers=[], + ) + return datasets + + @property + def data_type(self) -> str: + return "pcba" + + @property + def _name(self) -> str: + return "PCBA" + + +if __name__ == "__main__": + # Example usage + for dataset_class in [ + ClinTox, + BBBP, + SIDER, + BACE, + Tox21, + ToxCast, + # PCBA, + # HIV, + # MUV, + ]: + dataset = dataset_class() + dataset.prepare_data() + dataset.setup() + # dataset = SIDER() + # dataset.prepare_data() + # dataset.setup() diff --git a/chebai/preprocessing/datasets/pubchem.py b/chebai/preprocessing/datasets/pubchem.py index ea5e8978..5ac33439 100644 --- a/chebai/preprocessing/datasets/pubchem.py +++ b/chebai/preprocessing/datasets/pubchem.py @@ -138,11 +138,6 @@ def _download_required_data(self) -> str: self.download() return self._raw_data_source_path - def _graph_to_raw_dataset(self, graph): - raise NotImplementedError( - "PubChem does not use a graph-based data preparation pipeline." - ) - def download(self): """ Downloads PubChem data based on `_k` parameter. diff --git a/chebai/preprocessing/datasets/tox21.py b/chebai/preprocessing/datasets/tox21.py index f6298293..da478c14 100644 --- a/chebai/preprocessing/datasets/tox21.py +++ b/chebai/preprocessing/datasets/tox21.py @@ -1,5 +1,4 @@ import csv -import gzip import os import shutil import zipfile @@ -7,192 +6,13 @@ from typing import Dict, Generator, List, Optional from urllib import request -import numpy as np import torch from rdkit import Chem -from sklearn.model_selection import GroupShuffleSplit, train_test_split from chebai.preprocessing import reader as dr from chebai.preprocessing.datasets.base import XYBaseDataModule -class Tox21MolNet(XYBaseDataModule): - """Data module for Tox21MolNet dataset.""" - - HEADERS = [ - "NR-AR", - "NR-AR-LBD", - "NR-AhR", - "NR-Aromatase", - "NR-ER", - "NR-ER-LBD", - "NR-PPAR-gamma", - "SR-ARE", - "SR-ATAD5", - "SR-HSE", - "SR-MMP", - "SR-p53", - ] - - @property - def _name(self) -> str: - """Returns the name of the dataset.""" - return "Tox21MN" - - @property - def raw_file_names(self) -> List[str]: - """Returns a list of raw file names.""" - return ["tox21.csv"] - - # @property - # def processed_file_names(self) -> List[str]: - # """Returns a list of processed file names.""" - # return ["test.pt", "train.pt", "validation.pt"] - - @property - def processed_file_names_dict(self) -> dict: - return { - "test": "test.pt", - "train": "train.pt", - "validation": "validation.pt", - } - - def download(self) -> None: - """Downloads and extracts the dataset.""" - with NamedTemporaryFile("rb") as gout: - request.urlretrieve( - "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/tox21.csv.gz", - gout.name, - ) - with gzip.open(gout.name) as gfile: - with open(os.path.join(self.raw_dir, "tox21.csv"), "wt") as fout: - fout.write(gfile.read().decode()) - - def setup_processed(self) -> None: - """Processes and splits the dataset.""" - print("Create splits") - data = list(self._load_data_from_file(os.path.join(self.raw_dir, "tox21.csv"))) - groups = np.array([d.get("group") for d in data]) - - if not all(g is None for g in groups): - split_size = int( - len(set(groups)) * (1 - self.test_split - self.validation_split) - ) - os.makedirs(self.processed_dir, exist_ok=True) - splitter = GroupShuffleSplit(train_size=split_size, n_splits=1) - - train_split_index, temp_split_index = next( - splitter.split(data, groups=groups) - ) - - split_groups = groups[temp_split_index] - - splitter = GroupShuffleSplit( - train_size=int( - len(set(split_groups)) - * (1 - self.test_split - self.validation_split) - ), - n_splits=1, - ) - test_split_index, validation_split_index = next( - splitter.split(temp_split_index, groups=split_groups) - ) - train_split = [data[i] for i in train_split_index] - test_split = [ - d - for d in (data[temp_split_index[i]] for i in test_split_index) - # if d["original"] - ] - validation_split = [ - d - for d in (data[temp_split_index[i]] for i in validation_split_index) - # if d["original"] - ] - else: - train_split, test_split = train_test_split( - data, test_size=self.test_split, shuffle=True - ) - train_split, validation_split = train_test_split( - train_split, test_size=self.validation_split, shuffle=True - ) - - for k, split in [ - ("test", test_split), - ("train", train_split), - ("validation", validation_split), - ]: - print("transform", k) - torch.save( - split, - os.path.join(self.processed_dir, f"{k}.pt"), - ) - - def setup(self, **kwargs) -> None: - """Sets up the dataset by downloading and processing if necessary.""" - if self._setup_data_flag != 1: - return - - self._setup_data_flag += 1 - if any( - not os.path.isfile(os.path.join(self.raw_dir, f)) - for f in self.raw_file_names - ): - self.download() - if any( - not os.path.isfile(os.path.join(self.processed_dir, f)) - for f in self.processed_file_names - ): - self.setup_processed() - - # self._set_processed_data_props() - self._after_setup() - - def _load_dict(self, input_file_path: str) -> List[Dict]: - """Loads data from a CSV file. - - Args: - input_file_path (str): Path to the CSV file. - - Returns: - List[Dict]: List of data dictionaries. - """ - with open(input_file_path, "r") as input_file: - reader = csv.DictReader(input_file) - for row in reader: - smiles = row["smiles"] - labels = [ - bool(int(float(label))) if len(label) >= 1 else None - for label in (row[k] for k in self.HEADERS) - ] - # group = int(row["group"]) - yield dict( - features=smiles, - labels=labels, - ident=row["mol_id"], - # group=group - ) - # yield self.reader.to_data(dict(features=smiles, labels=labels, ident=row["mol_id"])) - - def _set_processed_data_props(self): - """ - Load processed data and extract metadata. - - Sets: - - self._num_of_labels: Number of target labels in the dataset. - - self._feature_vector_size: Maximum feature vector length across all data points. - """ - pt_file_path = os.path.join( - self.processed_dir, self.processed_file_names_dict["train"] - ) - data_pt = torch.load(pt_file_path, weights_only=False) - - self._num_of_labels = len(data_pt[0]["labels"]) - self._feature_vector_size = max(len(d["features"]) for d in data_pt) - - def _perform_data_preparation(self, *args, **kwargs) -> None: - pass - - class Tox21Challenge(XYBaseDataModule): """Data module for Tox21Challenge dataset.""" @@ -381,9 +201,3 @@ class Tox21ChallengeChem(Tox21Challenge): """Chemical data reader for Tox21Challenge dataset.""" READER = dr.ChemDataReader - - -class Tox21MolNetChem(Tox21MolNet): - """Chemical data reader for Tox21MolNet dataset.""" - - READER = dr.ChemDataReader diff --git a/chebai/preprocessing/reader.py b/chebai/preprocessing/reader.py index 664a8d8f..33ad3c16 100644 --- a/chebai/preprocessing/reader.py +++ b/chebai/preprocessing/reader.py @@ -99,11 +99,6 @@ def _read_components(self, row: Dict[str, Any]) -> Dict[str, Any]: under the additional `missing_labels` keyword.""" labels = self._get_raw_label(row) additional_kwargs = self._get_additional_kwargs(row) - if labels is not None: - if any(label is None for label in labels): - additional_kwargs["missing_labels"] = [ - label is None for label in labels - ] return dict( features=self._get_raw_data(row), labels=labels, @@ -207,8 +202,12 @@ def _read_data(self, raw_data: str | Chem.Mol) -> Optional[List[int]]: try: if isinstance(raw_data, str): mol = smiles_or_inchi_to_mol(raw_data.strip()) - else: + elif isinstance(raw_data, Chem.Mol): mol = raw_data + else: + raise ValueError( + f"Invalid input type: {type(raw_data)}. Expected str or Chem.Mol." + ) if mol is None: raise ValueError(f"Invalid input: {raw_data}") except ValueError as e: diff --git a/chebai/preprocessing/splitters/__init__.py b/chebai/preprocessing/splitters/__init__.py new file mode 100644 index 00000000..49b0395e --- /dev/null +++ b/chebai/preprocessing/splitters/__init__.py @@ -0,0 +1,5 @@ +from .group import GroupSplitter +from .multilabel import MultiLabelSplitter +from .random import RandomSplitter + +__all__ = ["GroupSplitter", "MultiLabelSplitter", "RandomSplitter"] diff --git a/chebai/preprocessing/splitters/group.py b/chebai/preprocessing/splitters/group.py new file mode 100644 index 00000000..7659b1ac --- /dev/null +++ b/chebai/preprocessing/splitters/group.py @@ -0,0 +1,138 @@ +"""Generate group-based train/validation/test splits from DataFrames.""" + +from __future__ import annotations + +from abc import ABC + +import pandas as pd +from sklearn.model_selection import GroupShuffleSplit + +from chebai.preprocessing.datasets.base import _DynamicDataset + + +class GroupSplitter(_DynamicDataset, ABC): + def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """ + Loads encoded/transformed data and generates training, validation, and test splits. + """ + + filename = self.processed_file_names_dict["data"] + data = self.load_processed_data_from_file(filename) + df_data = pd.DataFrame(data) + + splits = create_group_splits( + df_data, + self._LABELS_START_IDX, + 1 - self.validation_split - self.test_split, + self.validation_split, + self.test_split, + self.dynamic_data_split_seed, + ) + return splits["train"], splits["validation"], splits["test"] + + +def create_group_splits( + df: pd.DataFrame, + label_start_col: int = 2, + train_ratio: float = 0.8, + val_ratio: float = 0.1, + test_ratio: float = 0.1, + seed: int | None = 42, +) -> dict[str, pd.DataFrame]: + """Create group-based train/validation/test splits for DataFrames. + + Splitting is done with ``GroupShuffleSplit`` using the ``group`` column, + so that all rows sharing the same group value are assigned to the same + split (no group leaks across train/val/test). This is **not** a + stratified split: label balance across splits is not guaranteed, even + though label columns are used to build the ``y`` array passed to the + splitter (``GroupShuffleSplit`` ignores label values and only inspects + the ``groups`` argument). + + Parameters + ---------- + df : pd.DataFrame + Input data. Columns ``0`` to ``label_start_col - 1`` are treated as + feature/metadata columns; all remaining columns are boolean label + columns. A typical ChEBI DataFrame has columns + ``["chebi_id", "mol", "label1", "label2", ...]``. A ``group`` column + must also be present and is used to keep related rows together. + label_start_col : int + Index of the first label column (default 2). + train_ratio : float + Fraction of data for training (default 0.8). + val_ratio : float + Fraction of data for validation (default 0.1). + test_ratio : float + Fraction of data for testing (default 0.1). + seed : int or None + Random seed for reproducibility. + + Returns + ------- + dict + Dictionary with keys ``'train'``, ``'validation'``, ``'test'``, each + containing a DataFrame. + + Raises + ------ + ValueError + If the ratios do not sum to 1, any ratio is outside ``[0, 1]``, + *label_start_col* is out of range, the ``group`` column is missing, + or fewer than 2 unique groups are present. + """ + if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6: + raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") + if any(r < 0 or r > 1 for r in [train_ratio, val_ratio, test_ratio]): + raise ValueError("All ratios must be between 0 and 1") + if label_start_col >= len(df.columns): + raise ValueError( + f"label_start_col={label_start_col} is out of range for a DataFrame " + f"with {len(df.columns)} columns" + ) + + if "group" not in df.columns: + raise ValueError( + "Input DataFrame must contain a 'group' column for group split" + ) + + if len(df["group"].unique()) < 2: + raise ValueError( + "Input DataFrame must contain at least 2 unique groups for group split" + ) + + y = df.iloc[:, label_start_col:].values + # StratifiedShuffleSplit requires a 1-D label array + + df_reset = df.reset_index(drop=True) + + # ── Step 1: carve out the test set ────────────────────────────────────── + test_splitter = GroupShuffleSplit( + n_splits=1, test_size=test_ratio, random_state=seed + ) + + train_val_idx, test_idx = next(test_splitter.split(y, y, groups=df_reset["group"])) + + df_test = df_reset.iloc[test_idx] + df_trainval = df_reset.iloc[train_val_idx] + + # ── Step 2: split train/val from the remaining data ───────────────────── + y_trainval = y[train_val_idx] + val_ratio_adjusted = val_ratio / (1.0 - test_ratio) + + val_splitter = GroupShuffleSplit( + n_splits=1, test_size=val_ratio_adjusted, random_state=seed + ) + + train_idx_inner, val_idx_inner = next( + val_splitter.split(y_trainval, y_trainval, groups=df_trainval["group"]) + ) + + df_train = df_trainval.iloc[train_idx_inner] + df_val = df_trainval.iloc[val_idx_inner] + + return { + "train": df_train.reset_index(drop=True), + "validation": df_val.reset_index(drop=True), + "test": df_test.reset_index(drop=True), + } diff --git a/chebai/preprocessing/splitters/multilabel.py b/chebai/preprocessing/splitters/multilabel.py new file mode 100644 index 00000000..a7c0ced2 --- /dev/null +++ b/chebai/preprocessing/splitters/multilabel.py @@ -0,0 +1,32 @@ +"""Generate stratified train/validation/test splits from ChEBI DataFrames.""" + +from __future__ import annotations + +from abc import ABC + +import pandas as pd + +from chebai.preprocessing.datasets.base import _DynamicDataset + + +class MultiLabelSplitter(_DynamicDataset, ABC): + def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """ + Loads encoded/transformed data and generates training, validation, and test splits. + """ + + filename = self.processed_file_names_dict["data"] + data = self.load_processed_data_from_file(filename) + df_data = pd.DataFrame(data) + + from chebi_utils import create_multilabel_splits + + splits = create_multilabel_splits( + df_data, + self._LABELS_START_IDX, + 1 - self.validation_split - self.test_split, + self.validation_split, + self.test_split, + self.dynamic_data_split_seed, + ) + return splits["train"], splits["validation"], splits["test"] diff --git a/chebai/preprocessing/splitters/random.py b/chebai/preprocessing/splitters/random.py new file mode 100644 index 00000000..b84490ba --- /dev/null +++ b/chebai/preprocessing/splitters/random.py @@ -0,0 +1,99 @@ +"""Generate random (non-stratified) train/validation/test splits from DataFrames.""" + +from __future__ import annotations + +from abc import ABC + +import pandas as pd +from sklearn.model_selection import train_test_split + +from chebai.preprocessing.datasets.base import _DynamicDataset + + +class RandomSplitter(_DynamicDataset, ABC): + def _get_data_splits(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """ + Loads encoded/transformed data and generates training, validation, and test splits. + """ + + filename = self.processed_file_names_dict["data"] + data = self.load_processed_data_from_file(filename) + df_data = pd.DataFrame(data) + + splits = create_random_splits( + df_data, + 1 - self.validation_split - self.test_split, + self.validation_split, + self.test_split, + self.dynamic_data_split_seed, + ) + return splits["train"], splits["validation"], splits["test"] + + +def create_random_splits( + df: pd.DataFrame, + train_ratio: float = 0.8, + val_ratio: float = 0.1, + test_ratio: float = 0.1, + seed: int | None = 42, +) -> dict[str, pd.DataFrame]: + """Create random (non-stratified) train/validation/test splits. + + Rows are split purely at random using ``train_test_split`` from + scikit-learn, with no regard to label distribution or grouping. + + Parameters + ---------- + df : pd.DataFrame + Input data. + label_start_col : int + Index of the first label column (default 2). Unused by this + function; retained for consistency with related split functions. + train_ratio : float + Fraction of data for training (default 0.8). + val_ratio : float + Fraction of data for validation (default 0.1). + test_ratio : float + Fraction of data for testing (default 0.1). + seed : int or None + Random seed for reproducibility. + + Returns + ------- + dict + Dictionary with keys ``'train'``, ``'validation'``, ``'test'``, each + containing a DataFrame. + + Raises + ------ + ValueError + If the ratios do not sum to 1, any ratio is outside ``[0, 1]``, or + *label_start_col* is out of range. + """ + if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6: + raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") + if any(r < 0 or r > 1 for r in [train_ratio, val_ratio, test_ratio]): + raise ValueError("All ratios must be between 0 and 1") + + df_reset = df.reset_index(drop=True) + + # ── Step 1: carve out the test set ────────────────────────────────────── + df_trainval, df_test = train_test_split( + df_reset, test_size=test_ratio, shuffle=True, random_state=seed + ) + + # ── Step 2: split train/val from the remaining data ───────────────────── + val_ratio_adjusted = val_ratio / (1.0 - test_ratio) + + df_train, df_val = train_test_split( + df_trainval, + test_size=val_ratio_adjusted, + shuffle=True, + random_state=seed, + ) + + return { + "train": df_train.reset_index(drop=True), + "validation": df_val.reset_index(drop=True), + "test": df_test.reset_index(drop=True), + } diff --git a/chebai/result/compute_avg_performance.py b/chebai/result/compute_avg_performance.py new file mode 100644 index 00000000..2fb63b5a --- /dev/null +++ b/chebai/result/compute_avg_performance.py @@ -0,0 +1,287 @@ +""" +For each of one or more local W&B run files (run-*.wandb), find the step +with the best validation macro-F1 score, print that score along with the +corresponding validation micro-F1 (and any other metrics logged at that +same step). At the end, print the average +/- sample standard deviation +(ddof=1, the standard convention for reporting results across seeds) of +the best macro-F1 and its corresponding metrics across all the given files. + +Usage: + python find_best_f1.py 2cb51q4o 0nwo7wrt s4w2w2cx + python find_best_f1.py *.wandb --macro-metric val/macro_f1 --micro-metric val/micro_f1 + python find_best_f1.py *.wandb --max-epoch 200 + +If --macro-metric / --micro-metric aren't given, the script auto-detects +them per file (case-insensitive match on "f1"+"macro" / "f1"+"micro", +preferring keys that also mention "val"/"eval"/"test"). +""" + +import argparse +import json +import sys +from pathlib import Path +from statistics import mean, stdev + +try: + # Newer wandb versions (>=0.16 or so) + from wandb.sdk.internal.datastore import DataStore +except ImportError: + # Older wandb versions + from wandb.old.datastore import DataStore + +from wandb.proto import wandb_internal_pb2 as pb + +INTERNAL_KEY_PREFIXES = ("_",) # e.g. _step, _timestamp, _runtime + + +def iter_history_rows(wandb_path: str): + """ + Yields dicts of {key: value} for every 'history' record logged in the run, + parsed straight out of the binary .wandb file (no network / API calls). + """ + ds = DataStore() + ds.open_for_scan(wandb_path) + + while True: + data = ds.scan_data() + if data is None: + break + record = pb.Record() + record.ParseFromString(data) + + if record.WhichOneof("record_type") == "history": + row = {} + for item in record.history.item: + key = item.key if item.key else ".".join(item.nested_key) + try: + val = json.loads(item.value_json) + except Exception: + val = item.value_json + row[key] = val + yield row + + +def detect_metric_key(all_keys, must_contain, explicit=None): + """Find a logged key matching all substrings in must_contain (case-insensitive), + preferring ones that also look like validation metrics.""" + if explicit: + return explicit + + candidates = [k for k in all_keys if all(s in k.lower() for s in must_contain)] + val_candidates = [ + k + for k in candidates + if any(tag in k.lower() for tag in ("val", "eval", "test")) + ] + chosen = val_candidates if val_candidates else candidates + return chosen[0] if chosen else None + + +def get_epoch(row, epoch_key=None): + if epoch_key: + return row.get(epoch_key) + if "epoch" in row: + return row["epoch"] + return row.get("_step") + + +def process_file(path, macro_metric, micro_metric, epoch_key, max_epoch): + rows = list(iter_history_rows(str(path))) + if not rows: + raise ValueError(f" [!] No history records found in {path}, skipping.") + + all_keys = set() + for row in rows: + all_keys.update(row.keys()) + + macro_key = detect_metric_key(all_keys, ("f1", "macro"), macro_metric) + if macro_key is None: + raise ValueError( + f" [!] Could not find a macro-F1 metric in {path}. Available keys: {', '.join(sorted(all_keys))}" + ) + + micro_key = detect_metric_key(all_keys, ("f1", "micro"), micro_metric) + + best_row = None + best_epoch = None + best_val = None + for row in rows: + if macro_key not in row or row[macro_key] is None: + continue + epoch = get_epoch(row, epoch_key) + if epoch is not None and max_epoch is not None and epoch > max_epoch: + continue + try: + val = float(row[macro_key]) + except (TypeError, ValueError): + continue + if best_val is None or val > best_val: + best_val = val + best_row = row + best_epoch = epoch + + if best_row is None: + raise ValueError( + f" [!] No numeric values for '{macro_key}' within epoch <= {max_epoch} in {path}." + ) + + return { + "file": str(path), + "macro_key": macro_key, + "micro_key": micro_key, + "epoch": best_epoch, + "row": best_row, + } + + +def format_mean_std(vals): + """Mean +/- sample standard deviation (ddof=1), the convention used in + research for reporting performance across seeds/runs. Falls back to + 'no tolerance' when only one value is available (stdev is undefined).""" + m = mean(vals) + if len(vals) > 1: + s = stdev(vals) # sample std (n-1 denominator) + return f"{m:.4f} \u00b1 {s:.4f}" + return f"{m:.4f} (n=1, no std)" + + +def print_result(result): + row = result["row"] + macro_key = result["macro_key"] + micro_key = result["micro_key"] + + print(f"File: {result['file']}") + print(f" Best epoch/step: {result['epoch']}") + print(f" {macro_key}: {row[macro_key]:.4f}") + + if micro_key and micro_key in row and row[micro_key] is not None: + print(f" {micro_key} (corresponding): {row[micro_key]:.4f}") + elif micro_key: + print(f" {micro_key} (corresponding): N/A") + else: + raise ValueError(f" [!] No micro-F1 metric found in {result['file']}.") + + shown = {macro_key, micro_key, "epoch", "_step"} + other_keys = sorted( + k + for k in row.keys() + if k not in shown and not k.startswith(INTERNAL_KEY_PREFIXES) + ) + if other_keys: + print(" Other metrics at this step:") + for k in other_keys: + v = row[k] + if isinstance(v, float): + print(f" {k}: {v:.4f}") + else: + print(f" {k}: {v}") + print() + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "wandb_run_ids", + nargs="+", + help="Identifiers of local W&B run files (run-*.wandb) to process", + ) + parser.add_argument( + "--macro-metric", default=None, help="Exact key for macro-F1 (skip auto-detect)" + ) + parser.add_argument( + "--micro-metric", default=None, help="Exact key for micro-F1 (skip auto-detect)" + ) + parser.add_argument( + "--epoch-key", + default=None, + help="Key to use as the epoch/step number (default: auto epoch/_step)", + ) + parser.add_argument( + "--max-epoch", + type=int, + default=200, + help="Only consider steps/epochs up to this value", + ) + args = parser.parse_args() + + results = [] + for wandb_id in args.wandb_run_ids: + file_name = f"run-{wandb_id}.wandb" + matches = list(Path(".").rglob(file_name)) + + if len(matches) == 0: + raise FileNotFoundError(f"Could not find {file_name}") + + if len(matches) > 1: + raise RuntimeError( + f"Found multiple files named {file_name}:\n" + + "\n".join(str(p.resolve()) for p in matches) + ) + file_path = matches[0].resolve() + path = Path(file_path) + if not path.exists(): + raise FileNotFoundError(f" [!] File not found: {path}.") + + print(f"Processing {path.name} ...") + result = process_file( + path, args.macro_metric, args.micro_metric, args.epoch_key, args.max_epoch + ) + if result: + print_result(result) + results.append(result) + + if not results: + sys.exit("No valid results across the given files.") + + if len(results) == len(args.wandb_run_ids): + print( + f"({len(results)}/{len(args.wandb_run_ids)} files produced a valid result)\n" + ) + else: + raise ValueError( + f"({len(results)} and {len(args.wandb_run_ids)} do not match)\n" + ) + + print("=" * 50) + print(f"Average across {len(results)} file(s):") + + # Average (+/- sample std) the macro-F1 across files + macro_vals = [r["row"][r["macro_key"]] for r in results] + print(f" Best macro-F1: {format_mean_std(macro_vals)} (n={len(macro_vals)})") + + # Average (+/- sample std) the corresponding micro-F1 across files (where present) + micro_vals = [ + r["row"][r["micro_key"]] + for r in results + if r["micro_key"] + and r["micro_key"] in r["row"] + and r["row"][r["micro_key"]] is not None + ] + if micro_vals: + print( + f" Corresponding micro-F1: {format_mean_std(micro_vals)} (n={len(micro_vals)}/{len(results)})" + ) + + # Average every other numeric key found in the best rows (union across files) + shown = ( + {"epoch", "_step"} + | {r["macro_key"] for r in results} + | {r["micro_key"] for r in results if r["micro_key"]} + ) + other_key_values = {} + for r in results: + for k, v in r["row"].items(): + if k in shown or k.startswith(INTERNAL_KEY_PREFIXES): + continue + if isinstance(v, (int, float)): + other_key_values.setdefault(k, []).append(v) + + for k in sorted(other_key_values): + vals = other_key_values[k] + print(f" {k}: {format_mean_std(vals)} (n={len(vals)}/{len(results)})") + + +if __name__ == "__main__": + main() diff --git a/configs/data/moleculenet/bace_moleculenet.yml b/configs/data/moleculenet/bace_moleculenet.yml index bd6c04a8..da9736c5 100644 --- a/configs/data/moleculenet/bace_moleculenet.yml +++ b/configs/data/moleculenet/bace_moleculenet.yml @@ -1,5 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.BaceChem +class_path: chebai.preprocessing.datasets.molecule_net_classification.BACE init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/bbbp_moleculenet.yml b/configs/data/moleculenet/bbbp_moleculenet.yml index 01479443..8ad18668 100644 --- a/configs/data/moleculenet/bbbp_moleculenet.yml +++ b/configs/data/moleculenet/bbbp_moleculenet.yml @@ -1,5 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.BBBPChem +class_path: chebai.preprocessing.datasets.molecule_net_classification.BBBP init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/clintox_moleculenet.yml b/configs/data/moleculenet/clintox_moleculenet.yml index d7b7c3be..a389f725 100644 --- a/configs/data/moleculenet/clintox_moleculenet.yml +++ b/configs/data/moleculenet/clintox_moleculenet.yml @@ -1,5 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.ClinToxChem +class_path: chebai.preprocessing.datasets.molecule_net_classification.ClinTox init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/hiv_moleculenet.yml b/configs/data/moleculenet/hiv_moleculenet.yml index 3bef06b2..d1febe83 100644 --- a/configs/data/moleculenet/hiv_moleculenet.yml +++ b/configs/data/moleculenet/hiv_moleculenet.yml @@ -1,5 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.HIVChem +class_path: chebai.preprocessing.datasets.molecule_net_classification.HIV init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/muv_moleculenet.yml b/configs/data/moleculenet/muv_moleculenet.yml index d7498305..9e02496d 100644 --- a/configs/data/moleculenet/muv_moleculenet.yml +++ b/configs/data/moleculenet/muv_moleculenet.yml @@ -1,5 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.MUVChem +class_path: chebai.preprocessing.datasets.molecule_net_classification.MUV init_args: batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/sider_moleculenet.yml b/configs/data/moleculenet/sider_moleculenet.yml index 1a1d81ee..2ae64c2e 100644 --- a/configs/data/moleculenet/sider_moleculenet.yml +++ b/configs/data/moleculenet/sider_moleculenet.yml @@ -1,5 +1,3 @@ -class_path: chebai.preprocessing.datasets.molecule_classification.SiderChem +class_path: chebai.preprocessing.datasets.molecule_net_classification.SIDER init_args: batch_size: 10 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/data/moleculenet/tox21_moleculenet.yml b/configs/data/moleculenet/tox21_moleculenet.yml new file mode 100644 index 00000000..8ce308e9 --- /dev/null +++ b/configs/data/moleculenet/tox21_moleculenet.yml @@ -0,0 +1,3 @@ +class_path: chebai.preprocessing.datasets.molecule_net_classification.Tox21 +init_args: + batch_size: 32 diff --git a/configs/data/tox21/tox21_moleculenet.yml b/configs/data/tox21/tox21_moleculenet.yml deleted file mode 100644 index 1e8af70f..00000000 --- a/configs/data/tox21/tox21_moleculenet.yml +++ /dev/null @@ -1,5 +0,0 @@ -class_path: chebai.preprocessing.datasets.tox21.Tox21MolNetChem -init_args: - batch_size: 32 - validation_split: 0.05 - test_split: 0.15 diff --git a/configs/metrics/binary-f1-roc-auc.yml b/configs/metrics/binary-f1-roc-auc.yml index 05834343..d87bb04f 100644 --- a/configs/metrics/binary-f1-roc-auc.yml +++ b/configs/metrics/binary-f1-roc-auc.yml @@ -1,6 +1,6 @@ class_path: torchmetrics.MetricCollection init_args: - metrics: + metrics: # Use this for: BACE, BBBP, HIV f1: class_path: torchmetrics.classification.BinaryF1Score roc-auc: diff --git a/configs/metrics/micro-macro-f1-roc-auc.yml b/configs/metrics/micro-macro-f1-roc-auc.yml index c659b877..88bdcad3 100644 --- a/configs/metrics/micro-macro-f1-roc-auc.yml +++ b/configs/metrics/micro-macro-f1-roc-auc.yml @@ -1,6 +1,6 @@ class_path: torchmetrics.MetricCollection init_args: - metrics: + metrics: # Use this for: SIDER, ClinTox, Tox21, ToxCast micro-f1: class_path: torchmetrics.classification.MultilabelF1Score init_args: @@ -9,3 +9,13 @@ init_args: class_path: chebai.callbacks.epoch_metrics.MacroF1 roc-auc: class_path: torchmetrics.classification.MultilabelAUROC + init_args: + ignore_index: -1 + pr-auc: # Especially used for MUV, PCBA + class_path: torchmetrics.classification.MultilabelAveragePrecision + init_args: + ignore_index: -1 + himol-marco-roc-auc: + class_path: chebai.callbacks.epoch_metrics.HiMolMacroAUROC + init_args: + ignore_index: -1 diff --git a/configs/training/binary_callbacks.yml b/configs/training/binary_callbacks.yml index 5f6369ac..9bad0262 100644 --- a/configs/training/binary_callbacks.yml +++ b/configs/training/binary_callbacks.yml @@ -4,16 +4,15 @@ mode: 'max' filename: 'best_f1_{epoch:02d}_{val_loss:.4f}_{val_f1:.4f}_{val_roc-auc:.4f}' every_n_epochs: 1 - save_top_k: 1 + save_top_k: 3 - class_path: chebai.callbacks.model_checkpoint.CustomModelCheckpoint init_args: monitor: val_roc-auc mode: 'max' filename: 'best_roc-auc_{epoch:02d}_{val_loss:.4f}_{val_f1:.4f}_{val_roc-auc:.4f}' every_n_epochs: 1 - save_top_k: 1 + save_top_k: 3 - class_path: chebai.callbacks.model_checkpoint.CustomModelCheckpoint init_args: filename: 'per_{epoch:02d}_{val_loss:.4f}_{val_f1:.4f}_{val_roc-auc:.4f}' - every_n_epochs: 25 - save_top_k: -1 + save_top_k: 1 diff --git a/configs/training/default_callbacks.yml b/configs/training/default_callbacks.yml index 7635b88c..a20b9315 100644 --- a/configs/training/default_callbacks.yml +++ b/configs/training/default_callbacks.yml @@ -1,12 +1,11 @@ - class_path: chebai.callbacks.model_checkpoint.CustomModelCheckpoint init_args: - monitor: val_micro-f1 + monitor: val_macro-f1 mode: 'max' - filename: 'best_micro_f1_{epoch:02d}_{val_loss:.4f}_{val_macro-f1:.4f}_{val_micro-f1:.4f}' + filename: 'best_macro_f1_{epoch:02d}_{val_loss:.4f}_{val_macro-f1:.4f}_{val_micro-f1:.4f}' every_n_epochs: 1 save_top_k: 3 - class_path: chebai.callbacks.model_checkpoint.CustomModelCheckpoint init_args: filename: 'per_{epoch:02d}_{val_loss:.4f}_{val_macro-f1:.4f}_{val_micro-f1:.4f}' - every_n_epochs: 25 - save_top_k: -1 + save_top_k: 1 diff --git a/pyproject.toml b/pyproject.toml index f0594fc8..5143888e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,10 @@ dev = [ "deepsmiles", "torchmetrics", "chebi-utils>=0.4", + # In case of urllib.error.URLError: None: - """ - Set up Tox21 dataset and compute overlaps between data splits. - """ - cls.tox21 = Tox21MolNetChem() - cls.getDataSplitsOverlaps() - - @classmethod - def getDataSplitsOverlaps(cls) -> None: - """ - Get the overlap between data splits based on SMILES features and IDs. - """ - processed_path = os.path.join(os.getcwd(), cls.tox21.processed_dir) - print(f"Checking Data from - {processed_path}") - - train_set = torch.load( - os.path.join(processed_path, "train.pt"), weights_only=False - ) - val_set = torch.load( - os.path.join(processed_path, "validation.pt"), weights_only=False - ) - test_set = torch.load( - os.path.join(processed_path, "test.pt"), weights_only=False - ) - - train_smiles, train_smiles_ids = cls.get_features_ids(train_set) - val_smiles, val_smiles_ids = cls.get_features_ids(val_set) - test_smiles, test_smiles_ids = cls.get_features_ids(test_set) - - # Get overlaps based on SMILES features - cls.overlaps_train_val = cls.get_overlaps(train_smiles, val_smiles) - cls.overlaps_train_test = cls.get_overlaps(train_smiles, test_smiles) - cls.overlaps_val_test = cls.get_overlaps(val_smiles, test_smiles) - - # Get overlaps based on SMILES IDs - cls.overlaps_train_val_ids = cls.get_overlaps(train_smiles_ids, val_smiles_ids) - cls.overlaps_train_test_ids = cls.get_overlaps( - train_smiles_ids, test_smiles_ids - ) - cls.overlaps_val_test_ids = cls.get_overlaps(val_smiles_ids, test_smiles_ids) - - @staticmethod - def get_features_ids(data_split: List[Dict]) -> Tuple[List, List]: - """ - Returns SMILES features/tokens and SMILES IDs from the data. - - Args: - data_split (List[Dict]): List of dictionaries containing SMILES features and IDs. - - Returns: - Tuple[List, List]: Tuple of lists containing SMILES features and SMILES IDs. - """ - smiles_features, smiles_ids = [], [] - for entry in data_split: - smiles_features.append(entry["features"]) - smiles_ids.append(entry["ident"]) - - return smiles_features, smiles_ids - - @staticmethod - def get_overlaps(list_1: List, list_2: List) -> List: - """ - Get overlaps between two lists. - - Args: - list_1 (List): First list. - list_2 (List): Second list. - - Returns: - List: List of elements common to both input lists. - """ - overlap = [] - for element in list_1: - if element in list_2: - overlap.append(element) - return overlap - - def test_train_val_overlap_based_on_smiles(self) -> None: - """ - Check that train-val splits are performed correctly based on SMILES features. - """ - self.assertEqual( - len(self.overlaps_train_val), - 0, - "Duplicate entities present in Train and Validation set based on SMILES", - ) - - def test_train_test_overlap_based_on_smiles(self) -> None: - """ - Check that train-test splits are performed correctly based on SMILES features. - """ - self.assertEqual( - len(self.overlaps_train_test), - 0, - "Duplicate entities present in Train and Test set based on SMILES", - ) - - def test_val_test_overlap_based_on_smiles(self) -> None: - """ - Check that val-test splits are performed correctly based on SMILES features. - """ - self.assertEqual( - len(self.overlaps_val_test), - 0, - "Duplicate entities present in Validation and Test set based on SMILES", - ) - - def test_train_val_overlap_based_on_ids(self) -> None: - """ - Check that train-val splits are performed correctly based on SMILES IDs. - """ - self.assertEqual( - len(self.overlaps_train_val_ids), - 0, - "Duplicate entities present in Train and Validation set based on IDs", - ) - - def test_train_test_overlap_based_on_ids(self) -> None: - """ - Check that train-test splits are performed correctly based on SMILES IDs. - """ - self.assertEqual( - len(self.overlaps_train_test_ids), - 0, - "Duplicate entities present in Train and Test set based on IDs", - ) - - def test_val_test_overlap_based_on_ids(self) -> None: - """ - Check that val-test splits are performed correctly based on SMILES IDs. - """ - self.assertEqual( - len(self.overlaps_val_test_ids), - 0, - "Duplicate entities present in Validation and Test set based on IDs", - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/collators/testRaggedCollator.py b/tests/unit/collators/testRaggedCollator.py index d9ab2b1d..aa0258cd 100644 --- a/tests/unit/collators/testRaggedCollator.py +++ b/tests/unit/collators/testRaggedCollator.py @@ -73,20 +73,50 @@ def test_call_with_missing_entire_labels(self) -> None: data: List[Dict] = [ {"features": [1, 2], "labels": [True, False], "ident": "sample1"}, {"features": [3, 4, 5], "labels": None, "ident": "sample2"}, - {"features": [6], "labels": [True], "ident": "sample3"}, + {"features": [7], "labels": [True, None], "ident": "sample3"}, + {"features": [6], "labels": [True], "ident": "sample4"}, + {"features": [8, 9], "labels": [None, None], "ident": "sample5"}, ] result: XYData = self.collator(data) # https://github.com/ChEB-AI/python-chebai/pull/48#issuecomment-2324393829 - expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0]]) + expected_x = torch.tensor( + [ + [1, 2, 0], + [3, 4, 5], + [7, 0, 0], + [6, 0, 0], + [8, 9, 0], + ] + ) expected_y = torch.tensor( - [[True, False], [True, False]] + [ + [True, False], + [True, False], + [True, False], + [False, False], + ] ) # True -> 1, False -> 0 expected_mask_for_x = torch.tensor( - [[True, True, False], [True, True, True], [True, False, False]] + [ + [True, True, False], + [True, True, True], + [True, False, False], + [True, False, False], + [True, True, False], + ] + ) + expected_lens_for_x = torch.tensor([2, 3, 1, 1, 2]) + expected_valid_label_mask = torch.tensor( + [ + [True, True], # sample1 has no missing labels + [False, False], # sample2 has no missing labels (entire label is None) + [True, False], # sample3 has a missing label at index 1 + [True, False], # sample4 has no missing labels + [False, False], # sample5 has missing labels at both indices + ] ) - expected_lens_for_x = torch.tensor([2, 3, 1]) self.assertTrue( torch.equal(result.x, expected_x), @@ -110,19 +140,26 @@ def test_call_with_missing_entire_labels(self) -> None: ) self.assertEqual( result.additional_fields["loss_kwargs"]["non_null_labels"], - [0, 2], + [0, 2, 3, 4], "The non-null labels list does not match the expected output.", ) self.assertEqual( len(result.additional_fields["loss_kwargs"]["non_null_labels"]), - result.y.shape[1], + result.y.shape[0], "The length of non null labels list must match with target label variable size", ) self.assertEqual( result.additional_fields["idents"], - ("sample1", "sample2", "sample3"), + ("sample1", "sample2", "sample3", "sample4", "sample5"), "The identifiers do not match the expected output when labels are missing.", ) + self.assertTrue( + torch.equal( + result.additional_fields["loss_kwargs"]["valid_label_mask"], + expected_valid_label_mask, + ), + "The valid label mask tensor does not match the expected output when labels are missing.", + ) def test_call_with_none_in_labels(self) -> None: """ @@ -132,18 +169,32 @@ def test_call_with_none_in_labels(self) -> None: {"features": [1, 2], "labels": [None, True], "ident": "sample1"}, {"features": [3, 4, 5], "labels": [True, False], "ident": "sample2"}, {"features": [6], "labels": [True], "ident": "sample3"}, + {"features": [7, 8], "labels": [None, None], "ident": "sample4"}, ] result: XYData = self.collator(data) - expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0]]) + expected_x = torch.tensor([[1, 2, 0], [3, 4, 5], [6, 0, 0], [7, 8, 0]]) expected_y = torch.tensor( - [[False, True], [True, False], [True, False]] + [[False, True], [True, False], [True, False], [False, False]] ) # None -> False expected_mask_for_x = torch.tensor( - [[True, True, False], [True, True, True], [True, False, False]] + [ + [True, True, False], + [True, True, True], + [True, False, False], + [True, True, False], + ] + ) + expected_lens_for_x = torch.tensor([2, 3, 1, 2]) + expected_valid_label_mask = torch.tensor( + [ + [False, True], # sample1 has a missing label at index 0 + [True, True], # sample2 has no missing labels + [True, False], # sample3 has no missing labels + [False, False], # sample4 has missing labels at both indices + ] ) - expected_lens_for_x = torch.tensor([2, 3, 1]) self.assertTrue( torch.equal(result.x, expected_x), @@ -167,9 +218,16 @@ def test_call_with_none_in_labels(self) -> None: ) self.assertEqual( result.additional_fields["idents"], - ("sample1", "sample2", "sample3"), + ("sample1", "sample2", "sample3", "sample4"), "The identifiers do not match the expected output when labels contain None.", ) + self.assertTrue( + torch.equal( + result.additional_fields["loss_kwargs"]["valid_label_mask"], + expected_valid_label_mask, + ), + "The valid label mask tensor does not match the expected output when labels contain None.", + ) def test_call_with_empty_data(self) -> None: """ diff --git a/tests/unit/dataset_classes/testTox21MolNet.py b/tests/unit/dataset_classes/testTox21MolNet.py deleted file mode 100644 index 30383524..00000000 --- a/tests/unit/dataset_classes/testTox21MolNet.py +++ /dev/null @@ -1,185 +0,0 @@ -import unittest -from typing import List -from unittest.mock import MagicMock, mock_open, patch - -import torch - -from chebai.preprocessing.datasets.tox21 import Tox21MolNet -from chebai.preprocessing.reader import ChemDataReader -from tests.unit.mock_data.tox_mock_data import Tox21MolNetMockData - - -class TestTox21MolNet(unittest.TestCase): - @classmethod - @patch("os.makedirs", return_value=None) - def setUpClass(cls, mock_makedirs: MagicMock) -> None: - """ - Initialize a Tox21MolNet instance for testing. - - Args: - mock_makedirs (MagicMock): Mocked `os.makedirs` function. - """ - Tox21MolNet.READER = ChemDataReader - cls.data_module = Tox21MolNet() - - @patch( - "builtins.open", - new_callable=mock_open, - read_data=Tox21MolNetMockData.get_raw_data(), - ) - def test_load_data_from_file(self, mock_open_file: mock_open) -> None: - """ - Test the `_load_data_from_file` method for correct output. - - Args: - mock_open_file (mock_open): Mocked open function to simulate file reading. - """ - actual_data: list = self.data_module._load_data_from_file("fake/file/path.csv") - - first_instance = actual_data[0] - - # Check for required keys - required_keys = ["features", "labels", "ident"] - for key in required_keys: - self.assertIn( - key, first_instance, f"'{key}' key is missing in the output data." - ) - - self.assertTrue( - all(isinstance(feature, int) for feature in first_instance["features"]), - "Not all elements in 'features' are integers.", - ) - - # Check that 'features' can be converted to a tensor - features = first_instance["features"] - try: - tensor_features = torch.tensor(features) - self.assertTrue( - tensor_features.ndim > 0, - "'features' should be convertible to a non-empty tensor.", - ) - except Exception as e: - self.fail(f"'features' cannot be converted to a tensor: {str(e)}") - - @patch( - "builtins.open", - new_callable=mock_open, - read_data=Tox21MolNetMockData.get_raw_data(), - ) - @patch("torch.save") - def test_setup_processed_simple_split( - self, - mock_torch_save: MagicMock, - mock_open_file: mock_open, - ) -> None: - """ - Test the `setup_processed` method for basic data splitting and saving. - - Args: - mock_torch_save (MagicMock): Mocked `torch.save` function to avoid actual file writes. - mock_open_file (mock_open): Mocked `open` function to simulate file reading. - """ - self.data_module.setup_processed() - - # Verify if torch.save was called for each split (train, test, validation) - self.assertEqual( - mock_torch_save.call_count, 3, "Expected torch.save to be called 3 times." - ) - call_args_list = mock_torch_save.call_args_list - self.assertIn("test", call_args_list[0][0][1], "Missing 'test' split.") - self.assertIn("train", call_args_list[1][0][1], "Missing 'train' split.") - self.assertIn( - "validation", call_args_list[2][0][1], "Missing 'validation' split." - ) - - # Check for non-overlap between train, test, and validation splits - test_split: List[str] = [d["ident"] for d in call_args_list[0][0][0]] - train_split: List[str] = [d["ident"] for d in call_args_list[1][0][0]] - validation_split: List[str] = [d["ident"] for d in call_args_list[2][0][0]] - - self.assertTrue( - set(train_split).isdisjoint(test_split), - "Overlap detected between the train and test splits.", - ) - self.assertTrue( - set(train_split).isdisjoint(validation_split), - "Overlap detected between the train and validation splits.", - ) - self.assertTrue( - set(test_split).isdisjoint(validation_split), - "Overlap detected between the test and validation splits.", - ) - - @patch.object( - Tox21MolNet, - "_load_data_from_file", - return_value=Tox21MolNetMockData.get_processed_grouped_data(), - ) - @patch("torch.save") - def test_setup_processed_with_group_split( - self, mock_torch_save: MagicMock, mock_load_file: MagicMock - ) -> None: - """ - Test the `setup_processed` method for group-based splitting and saving. - - Args: - mock_torch_save (MagicMock): Mocked `torch.save` function to avoid actual file writes. - mock_load_file (MagicMock): Mocked `_load_data_from_file` to provide custom data. - """ - # self.data_module.train_split = 0.5 - # To get the train split as 50%, set test and validation splits to 25% each - # Refer: https://github.com/ChEB-AI/python-chebai/pull/102 - self.data_module.test_split = 0.25 - self.data_module.validation_split = 0.25 - self.data_module.setup_processed() - - # Verify if torch.save was called for each split - self.assertEqual( - mock_torch_save.call_count, 3, "Expected torch.save to be called 3 times." - ) - call_args_list = mock_torch_save.call_args_list - self.assertIn("test", call_args_list[0][0][1], "Missing 'test' split.") - self.assertIn("train", call_args_list[1][0][1], "Missing 'train' split.") - self.assertIn( - "validation", call_args_list[2][0][1], "Missing 'validation' split." - ) - - # Check for non-overlap between train, test, and validation splits (based on 'ident') - test_split: List[str] = [d["ident"] for d in call_args_list[0][0][0]] - train_split: List[str] = [d["ident"] for d in call_args_list[1][0][0]] - validation_split: List[str] = [d["ident"] for d in call_args_list[2][0][0]] - - self.assertTrue( - set(train_split).isdisjoint(test_split), - "Overlap detected between the train and test splits (based on 'ident').", - ) - self.assertTrue( - set(train_split).isdisjoint(validation_split), - "Overlap detected between the train and validation splits (based on 'ident').", - ) - self.assertTrue( - set(test_split).isdisjoint(validation_split), - "Overlap detected between the test and validation splits (based on 'ident').", - ) - - # Check for non-overlap between train, test, and validation splits (based on 'group') - test_split_grp: List[str] = [d["group"] for d in call_args_list[0][0][0]] - train_split_grp: List[str] = [d["group"] for d in call_args_list[1][0][0]] - validation_split_grp: List[str] = [d["group"] for d in call_args_list[2][0][0]] - - self.assertTrue( - set(train_split_grp).isdisjoint(test_split_grp), - "Overlap detected between the train and test splits (based on 'group').", - ) - self.assertTrue( - set(train_split_grp).isdisjoint(validation_split_grp), - "Overlap detected between the train and validation splits (based on 'group').", - ) - self.assertTrue( - set(test_split_grp).isdisjoint(validation_split_grp), - "Overlap detected between the test and validation splits (based on 'group').", - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/mock_data/tox_mock_data.py b/tests/unit/mock_data/tox_mock_data.py index fcf5633f..7567d6b2 100644 --- a/tests/unit/mock_data/tox_mock_data.py +++ b/tests/unit/mock_data/tox_mock_data.py @@ -394,10 +394,6 @@ def data_in_dict_format() -> List[Dict]: for dict_ in data_list: dict_["features"] = Tox21ChallengeMockData.FEATURE_OF_SMILES dict_["group"] = None - if any(label is None for label in dict_["labels"]): - dict_["missing_labels"] = [ - True if label is None else False for label in dict_["labels"] - ] return data_list @@ -509,9 +505,5 @@ def get_setup_processed_output_data() -> List[Dict]: "group": None, } ) - if any(label is None for label in dict_["labels"]): - complete_list[-1]["missing_labels"] = [ - True if label is None else False for label in dict_["labels"] - ] return complete_list