MoleculeNet Classification Data : Modularization, dynamic splits, Clean Code

.vscode/settings.json

@@ -2,7 +2,7 @@
     "python.testing.unittestArgs": [
         "-v",
         "-s",
-        "./tests",
+        "./tests/unit",
         "-p",
         "test*.py"
     ],

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,16 +64,15 @@ 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,
         label_filter: Optional[int] = None,
         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(

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,13 +52,15 @@ 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,
         augment_smiles: bool = False,
         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)