Protein Close Residues Lifting (Graph to Hypergraph)

configs/datasets/QM9.yaml

@@ -0,0 +1,14 @@
+data_domain: graph
+data_type: QM9
+data_name: QM9
+data_dir: datasets/${data_domain}/${data_type}
+#data_split_dir: ${oc.env:PROJECT_ROOT}/datasets/data_splits/${data_name}
+
+# Dataset parameters
+num_features: 11
+num_classes: 1
+task: regression
+loss_type: mse
+monitor_metric: mae
+task_level: graph
+

configs/datasets/UniProt.yaml

@@ -0,0 +1,22 @@
+data_domain: graph
+data_type: UniProt
+data_name: UniProt
+data_dir: datasets/${data_domain}/${data_type}
+#data_split_dir: ${oc.env:PROJECT_ROOT}/datasets/data_splits/${data_name}
+
+# Some parameters to do the query
+query: "length:[95 TO 155]" # number of residues per protein
+format: "tsv"
+fields: "accession,length"
+size: 100 # number of proteins to load
+
+threshold: 6.0 # distance between proteins to create the initial graph
+
+# Dataset parameters
+num_features: 20
+num_classes: 1
+task: regression
+loss_type: mse
+monitor_metric: mae
+task_level: graph
+

configs/datasets/manual_prot.yaml

@@ -0,0 +1,12 @@
+data_domain: graph
+data_type: toy_dataset
+data_name: manual_prot
+data_dir: datasets/${data_domain}/${data_type}
+
+# Dataset parameters
+num_features: 1
+num_classes: 2
+task: classification
+loss_type: cross_entropy
+monitor_metric: accuracy
+task_level: node

configs/transforms/liftings/graph2hypergraph/close_lifting.yaml

@@ -0,0 +1,5 @@
+transform_type: 'lifting'
+transform_name: "HypergraphCloseLifting"
+feature_lifting: ProjectionSum
+
+distance: 6.0

loaders.py

@@ -0,0 +1,227 @@
+import os
+
+import numpy as np
+import rootutils
+import torch_geometric
+from omegaconf import DictConfig
+
+# silent RDKit warnings
+from rdkit import Chem, RDLogger
+
+from modules.data.load.base import AbstractLoader
+from modules.data.utils.concat2geometric_dataset import ConcatToGeometricDataset
+from modules.data.utils.custom_dataset import CustomDataset
+from modules.data.utils.utils import (
+    load_cell_complex_dataset,
+    load_hypergraph_pickle_dataset,
+    load_manual_graph,
+    load_simplicial_dataset,
+)
+
+RDLogger.DisableLog("rdApp.*")
+
+
+class GraphLoader(AbstractLoader):
+    r"""Loader for graph datasets.
+
+    Parameters
+    ----------
+    parameters : DictConfig
+        Configuration parameters.
+    """
+
+    def __init__(self, parameters: DictConfig):
+        super().__init__(parameters)
+        self.parameters = parameters
+
+    def is_valid_smiles(self, smiles):
+        """Check if a SMILES string is valid using RDKit."""
+        mol = Chem.MolFromSmiles(smiles)
+        return mol is not None
+
+    def filter_qm9_dataset(self, dataset):
+        """Filter the QM9 dataset to remove invalid SMILES strings."""
+        return [data for data in dataset if self.is_valid_smiles(data.smiles)]
+
+    def load(self) -> torch_geometric.data.Dataset:
+        r"""Load graph dataset.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        torch_geometric.data.Dataset
+            torch_geometric.data.Dataset object containing the loaded data.
+        """
+        # Define the path to the data directory
+        root_folder = rootutils.find_root()
+        root_data_dir = os.path.join(root_folder, self.parameters["data_dir"])
+
+        self.data_dir = os.path.join(root_data_dir, self.parameters["data_name"])
+        if (
+            self.parameters.data_name.lower() in ["cora", "citeseer", "pubmed"]
+            and self.parameters.data_type == "cocitation"
+        ):
+            dataset = torch_geometric.datasets.Planetoid(
+                root=root_data_dir,
+                name=self.parameters["data_name"],
+            )
+
+        elif self.parameters.data_name in [
+            "MUTAG",
+            "ENZYMES",
+            "PROTEINS",
+            "COLLAB",
+            "IMDB-BINARY",
+            "IMDB-MULTI",
+            "REDDIT-BINARY",
+            "NCI1",
+            "NCI109",
+        ]:
+            dataset = torch_geometric.datasets.TUDataset(
+                root=root_data_dir,
+                name=self.parameters["data_name"],
+                use_node_attr=False,
+            )
+
+        elif self.parameters.data_name in ["ZINC", "AQSOL"]:
+            datasets = []
+            for split in ["train", "val", "test"]:
+                if self.parameters.data_name == "ZINC":
+                    datasets.append(
+                        torch_geometric.datasets.ZINC(
+                            root=root_data_dir,
+                            subset=True,
+                            split=split,
+                        )
+                    )
+                elif self.parameters.data_name == "AQSOL":
+                    datasets.append(
+                        torch_geometric.datasets.AQSOL(
+                            root=root_data_dir,
+                            split=split,
+                        )
+                    )
+            # The splits are predefined
+            # Extract and prepare split_idx
+            split_idx = {"train": np.arange(len(datasets[0]))}
+            split_idx["valid"] = np.arange(
+                len(datasets[0]), len(datasets[0]) + len(datasets[1])
+            )
+            split_idx["test"] = np.arange(
+                len(datasets[0]) + len(datasets[1]),
+                len(datasets[0]) + len(datasets[1]) + len(datasets[2]),
+            )
+            # Join dataset to process it
+            dataset = datasets[0] + datasets[1] + datasets[2]
+            dataset = ConcatToGeometricDataset(dataset)
+
+        elif self.parameters.data_name == "QM9":
+            dataset = torch_geometric.datasets.QM9(root=root_data_dir)
+            # Filter the QM9 dataset to remove invalid SMILES strings
+            valid_dataset = self.filter_qm9_dataset(dataset)
+            # dataset = ConcatToGeometricDataset(valid_dataset)
+            dataset = CustomDataset(valid_dataset, self.data_dir)
+
+        elif self.parameters.data_name in ["manual"]:
+            data = load_manual_graph()
+            dataset = CustomDataset([data], self.data_dir)
+
+        else:
+            raise NotImplementedError(
+                f"Dataset {self.parameters.data_name} not implemented"
+            )
+
+        return dataset
+
+
+class CellComplexLoader(AbstractLoader):
+    r"""Loader for cell complex datasets.
+
+    Parameters
+    ----------
+    parameters : DictConfig
+        Configuration parameters.
+    """
+
+    def __init__(self, parameters: DictConfig):
+        super().__init__(parameters)
+        self.parameters = parameters
+
+    def load(
+        self,
+    ) -> torch_geometric.data.Dataset:
+        r"""Load cell complex dataset.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        torch_geometric.data.Dataset
+            torch_geometric.data.Dataset object containing the loaded data.
+        """
+        return load_cell_complex_dataset(self.parameters)
+
+
+class SimplicialLoader(AbstractLoader):
+    r"""Loader for simplicial datasets.
+
+    Parameters
+    ----------
+    parameters : DictConfig
+        Configuration parameters.
+    """
+
+    def __init__(self, parameters: DictConfig):
+        super().__init__(parameters)
+        self.parameters = parameters
+
+    def load(
+        self,
+    ) -> torch_geometric.data.Dataset:
+        r"""Load simplicial dataset.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        torch_geometric.data.Dataset
+            torch_geometric.data.Dataset object containing the loaded data.
+        """
+        return load_simplicial_dataset(self.parameters)
+
+
+class HypergraphLoader(AbstractLoader):
+    r"""Loader for hypergraph datasets.
+
+    Parameters
+    ----------
+    parameters : DictConfig
+        Configuration parameters.
+    """
+
+    def __init__(self, parameters: DictConfig):
+        super().__init__(parameters)
+        self.parameters = parameters
+
+    def load(
+        self,
+    ) -> torch_geometric.data.Dataset:
+        r"""Load hypergraph dataset.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        torch_geometric.data.Dataset
+            torch_geometric.data.Dataset object containing the loaded data.
+        """
+        return load_hypergraph_pickle_dataset(self.parameters)