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/preprocessing/datasets/base.py b/chebai/preprocessing/datasets/base.py index 24655b0d..f036ff86 100644 --- a/chebai/preprocessing/datasets/base.py +++ b/chebai/preprocessing/datasets/base.py @@ -2,7 +2,7 @@ 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 @@ -15,9 +15,6 @@ from chebai.preprocessing import reader as dr -if TYPE_CHECKING: - import networkx as nx - class XYBaseDataModule(LightningDataModule): """ @@ -36,7 +33,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. @@ -53,7 +49,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. @@ -69,8 +64,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, @@ -78,7 +73,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, @@ -102,7 +96,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 = ( @@ -283,6 +276,13 @@ 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.", + "Please check the data preparation and filtering steps.", + ) + return DataLoader( dataset, collate_fn=self.reader.collator, @@ -617,6 +617,7 @@ def raw_file_names(self) -> List[str]: return list(self.raw_file_names_dict.values()) @property + @abstractmethod def raw_file_names_dict(self) -> dict: """ Returns the dictionary of raw file names (i.e., files that are directly obtained from an external source). @@ -815,7 +816,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 @@ -909,10 +910,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 @@ -926,17 +924,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 @@ -948,7 +944,7 @@ 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")) + data.to_pickle(open(os.path.join(self.processed_dir_main, filename), "wb")) def get_processed_pickled_df_file(self, filename: str) -> Optional[pd.DataFrame]: """ @@ -971,7 +967,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 @@ -1106,6 +1102,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"] @@ -1142,6 +1139,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 a0af3ac8..6845d00d 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. @@ -374,27 +393,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: @@ -746,12 +744,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..06311ee0 --- /dev/null +++ b/chebai/preprocessing/datasets/molecule_net_classification.py @@ -0,0 +1,257 @@ +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 + """ + + READER = dr.ChemDataReader + + @property + def _name(self) -> str: + """Returns the name of the dataset.""" + return str(self.__class__.__name__) + + 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(open(os.path.join(self.processed_dir_main, filename), "wb")) + + 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() + for split_name, data in [ + ("train", train), + ("valid", valid), + ("test", test), + ]: + for idx, (mol, labels, wi, smiles) in enumerate(data.itersamples()): + yield dict( + features=mol, + labels=labels, + ident=idx, + ) + splits.append( + { + "id": idx, + "split": split_name, + } + ) + splits_df = pd.DataFrame(splits) + splits_df.to_csv( + os.path.join(self.processed_dir_main, "splits.csv"), index=False + ) + + @abstractmethod + def _deep_chem_data_loader_api( + self, + ) -> tuple[DiskDataset, DiskDataset, DiskDataset]: + pass + + def _get_data_splits(self) -> None: + pass + + @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, + ) + return datasets + + +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, + ) + return datasets + + +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, + ) + return datasets + + +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, + ) + return datasets + + +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, + ) + return datasets + + +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, + ) + return datasets + + +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, + ) + return datasets + + +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, + ) + return datasets + + +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, + ) + return datasets + + +if __name__ == "__main__": + # Example usage + dataset = BBBP() + 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 0dae39cd..0c73a7ef 100644 --- a/chebai/preprocessing/reader.py +++ b/chebai/preprocessing/reader.py @@ -205,8 +205,12 @@ def _read_data(self, raw_data: str | Chem.Mol) -> Optional[List[int]]: try: if isinstance(raw_data, str): mol = Chem.MolFromSmiles(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..e4ddf118 --- /dev/null +++ b/chebai/preprocessing/splitters/group.py @@ -0,0 +1,138 @@ +"""Generate stratified train/validation/test splits from ChEBI 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..bf407fdc --- /dev/null +++ b/chebai/preprocessing/splitters/random.py @@ -0,0 +1,99 @@ +"""Generate stratified train/validation/test splits from ChEBI 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/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..9e7b0d45 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,5 @@ init_args: class_path: chebai.callbacks.epoch_metrics.MacroF1 roc-auc: class_path: torchmetrics.classification.MultilabelAUROC + pr-auc: # Especially used for MUV, PCBA + class_path: torchmetrics.classification.MultilabelAveragePrecision diff --git a/pyproject.toml b/pyproject.toml index 6c00552f..26abb251 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,9 @@ dev = [ "deepsmiles", "torchmetrics", "chebi-utils>=0.3", + # 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/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()