Repository: a-r-j/graphein (language: python)

Issue:
Add Multiprocessing to Protein Graph Construction
Implements #72

Implements multiprocessing of a list of PDB codes/paths to speed up data preparation.

#### Tests
Rough sanity check tests implemented. Would be great to hear any suggestions for additional tests :)

—

Created via [Raycast](https://www.raycast.com?ref=signatureGithub)


Discussion / hints:

Fix patch:
diff --git a/graphein/datasets/__init__.py b/graphein/datasets/__init__.py
new file mode 100644
index 00000000..b9804882
--- /dev/null
+++ b/graphein/datasets/__init__.py
@@ -0,0 +1,10 @@
+# Graphein
+# Author: Arian Jamasb <arian@jamasb.io>
+# License: MIT
+# Code Repository: https://github.com/a-r-j/graphein
+
+from .proteins_ligands import PROTEINS_LIGANDS
+from .proteins_metal import PROTEINS_METAL
+from .proteins_nucleic import PROTEINS_NUCLEIC
+from .proteins_nucleotides import PROTEINS_NUCLEOTIDES
+from .pscdb import PSCDB
diff --git a/graphein/datasets/base.py b/graphein/datasets/base.py
new file mode 100644
index 00000000..9706ab13
--- /dev/null
+++ b/graphein/datasets/base.py
@@ -0,0 +1,234 @@
+"""Base class for working with the PROTEINS_X datasets"""
+import logging
+import multiprocessing
+# Graphein
+# Author: Arian Jamasb <arian@jamasb.io>
+# License: MIT
+# Code Repository: https://github.com/a-r-j/graphein
+import traceback
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import List, Optional, Tuple
+
+import networkx as nx
+import numpy as np
+import pandas as pd
+from Bio.PDB import PDBList
+from sklearn.model_selection import train_test_split
+from tqdm import tqdm
+
+from graphein.protein.config import ProteinGraphConfig
+from graphein.protein.graphs import construct_graph
+
+logging.basicConfig(level="DEBUG")
+log = logging.getLogger(__name__)
+
+
+class AbstractClassificationDataset(ABC):
+    def __init__(
+        self,
+        df: Optional[pd.DataFrame] = None,
+        protein_graph_config: Optional[
+            ProteinGraphConfig
+        ] = ProteinGraphConfig(),
+        num_cores: int = 16,
+    ):
+        self.df = self._load_dataset() if df is None else df
+        self._num_cores: int = num_cores
+        self.config = protein_graph_config
+
+        self.sequences: List[str] = self.get_sequences(self.df)
+        self.graph_labels: List[str] = self.get_graph_labels(self.df)
+        self.graph_classes: List[str] = list(set(self.graph_labels))
+        self.residue_labels: List[str] = self.get_residue_labels(self.df)
+        self.pdb_list: List[str] = self.get_pdb_list(self.df)
+        self.chain_list: List[str] = self.get_chain_list(self.df)
+        self.chain_length: List[str] = self.get_chain_lengths(self.df)
+        self.node_labels: List[np.ndarray] = [
+            self.encode_interactions(ex) for ex in self.residue_labels
+        ]
+        if protein_graph_config is not None:
+            self.graphs: List[nx.Graph] = self.construct_graphs()
+        self.bad_pdbs: List[str] = []
+
+    @abstractmethod
+    def _load_dataset(self):
+        pass
+
+    @abstractmethod
+    def split_data(self):
+        pass
+
+    def _split_data(
+        self,
+        strategy: str = "random",
+        training_size: float = 0.7,
+        validation_size: float = 0.2,
+        test_size: float = 0.1,
+        shuffle: bool = True,
+    ):
+        if strategy == "random":
+            train_df, test_df = train_test_split(
+                self.df,
+                train_size=training_size,
+                test_size=test_size + validation_size,
+                shuffle=shuffle,
+            )
+            val_df, test_df = train_test_split(
+                test_df,
+                train_size=validation_size / (validation_size + test_size),
+                test_size=test_size / (test_size + validation_size),
+                shuffle=shuffle,
+            )
+        elif strategy == "stratified":
+            y = self.df[["interactor"]]
+            train_df, temp_df, _, _ = train_test_split(
+                self.df,
+                y,
+                stratify=y,
+                train_size=training_size,
+                test_size=test_size + validation_size,
+            )
+            y = temp_df[["interactor"]]
+            val_df, test_df, _, _ = train_test_split(
+                temp_df,
+                y,
+                stratify=y,
+                train_size=validation_size / (validation_size + test_size),
+                test_size=test_size / (validation_size + test_size),
+            )
+        return train_df, val_df, test_df
+
+    @staticmethod
+    def get_sequences(df) -> List[str]:
+        """
+        Returns the list of protein sequences in the dataset
+
+        :return: List of protein sequences
+        :rtype: List[str]
+        """
+        return list(df["sequence"])
+
+    @staticmethod
+    def get_residue_labels(df) -> List[str]:
+        """
+        Returns residue labels for sequences in the dataset (in string form).
+
+        '-' denotes a residue is non-interacting
+        '+' denotes a residue does interact
+        :return: List of residue labels for the dataset
+        :rtype: List[str]
+        """
+        return list(df["interacting_residues"])
+
+    @staticmethod
+    def get_pdb_list(df) -> List[str]:
+        """
+        Returns list of PDBs in the dataset
+
+        :return: List of PDB accession codes
+        :rtype: List[str]
+        """
+        return list(df["PDB"])
+
+    @staticmethod
+    def get_graph_labels(df) -> List[str]:
+        return list(df["interactor"])
+
+    @staticmethod
+    def get_chain_list(df) -> List[str]:
+        """Returns list of proteins chains in dataset
+
+        :return: List of protein chains
+        :rtype: List[str]
+        """
+        return list(df["chain"])
+
+    @staticmethod
+    def get_chain_lengths(df) -> List[int]:
+        """Returns list of chain lengths
+
+        :return: List of chain lengths
+        :rtype: List[int]
+        """
+        return list(df["length"])
+
+    @staticmethod
+    def encode_interactions(s: str) -> np.ndarray:
+        """
+        Encodes interacting residue string as a binary numpy array
+
+        :param s:  String of interaction status (e.g '----++----+-')
+        :type s: str
+        :return: binary numpy array representation of interacting residues (nodes)
+        :rtype: np.array
+        """
+        return np.array([1 if c == "+" else 0 for c in s])
+
+    def construct_graphs(self) -> List[nx.Graph]:
+        """
+        Constructs graphs for protein chains in the dataset.
+
+        :param config: Config specifying protein graph construction parameters
+        :type config: graphein.protein.ProteinGraphConfig
+        :return: List of protein structure graphs
+        :rtype: List[nx.Graph]
+        """
+        pool = multiprocessing.Pool(self._num_cores)
+        graph_list = list(
+            pool.map(
+                self._graph_constructor,
+                [
+                    (pdb, self.chain_list[i])
+                    for i, pdb in enumerate(self.pdb_list)
+                ],
+            )
+        )
+        pool.close()
+        pool.join()
+
+        # for i, pdb in enumerate(pdbs):
+        #    assert g.number_of_nodes() == len(
+        #        self.residue_labels[i]
+        #    ), f"Lengths do not match: Nodes: {g.number_of_nodes()} Labels: {len(self.residue_labels[i])}"
+
+        return graph_list
+
+    def _graph_constructor(self, args: Tuple[str, str]):
+        """
+        Partialed graph constructor for multiprocessing
+
+        :param args: Tuple of pdb code and chain to build graph of
+        :type args: Tuple[str, str]
+        :return: Protein structure graph
+        :rtype: nx.Graph
+        """
+        log.info(
+            f"Constructing graph for: {args[0]}. Chain selection: {args[1]}"
+        )
+        func = partial(construct_graph, config=self.config)
+        try:
+            return func(pdb_code=args[0], chain_selection=args[1])
+        except Exception:
+            log.info(
+                f"Graph construction error (PDB={args[0]})! {traceback.format_exc()}"
+            )
+            self.bad_pdbs.append(args[0])
+
+    def download_pdbs(self, path: str):
+        """
+        Downloads dataset PDBs to a specified directories.
+
+        :param path: Path to desired output location
+        :type path: str
+        """
+        pdbl = PDBList()
+        pdbl.download_pdb_files(pdb_codes=self.pdb_list, pdir=path)
+
+    def __len__(self) -> int:
+        """Returns length of the dataset.
+
+        :returns: Dataset length
+        :rtype: int
+        """
+        return len(self.df)
diff --git a/graphein/datasets/dataloader.py b/graphein/datasets/dataloader.py
new file mode 100644
index 00000000..ba716ccc
--- /dev/null
+++ b/graphein/datasets/dataloader.py
@@ -0,0 +1,433 @@
+"""Dataloader class for working with protein structure graphs."""
+# Graphein
+# Author: Arian Jamasb <arian@jamasb.io>
+# License: MIT
+# Code Repository: https://github.com/a-r-j/graphein
+from __future__ import annotations
+
+import logging
+import os
+from typing import List, Optional, Tuple, Union
+
+import networkx as nx
+import numpy as np
+from Bio.PDB import PDBList
+
+from graphein.ml.conversion import GraphFormatConvertor
+from graphein.protein.config import ProteinGraphConfig
+from graphein.protein.graphs import construct_graphs_mp
+from graphein.utils.utils import import_message
+
+logging.basicConfig(level="DEBUG")
+log = logging.getLogger(__name__)
+
+try:
+    import torch
+except ImportError:
+    message = import_message(
+        "graphein.datasets.dataloader", "torch", pip_install=True
+    )
+    log.warning(message)
+    log.warning("PyTorch not installed. Skipping torch-related functionality.")
+
+try:
+    from torch_geometric.loader import DataLoader
+except ImportError:
+    message = import_message(
+        "graphein.datasets.dataloader", "torch_geometric", pip_install=True
+    )
+    log.warning(message)
+
+
+class GrapheinDataLoader:
+    def __init__(
+        self,
+        pdb_file_list: Optional[List[str]] = None,
+        pdb_code_list: Optional[List[str]] = None,
+        pdb_in_dir: Optional[str] = None,
+        pdb_out_dir: str = "/tmp/",
+        graph_labels: Optional[List[Union[np.ndarray, torch.Tensor]]] = None,
+        node_labels: Optional[List[Union[np.ndarray, torch.Tensor]]] = None,
+        chain_list: Optional[List[str]] = None,
+        protein_graph_config: Optional[
+            ProteinGraphConfig
+        ] = ProteinGraphConfig(),
+        num_cores: int = 16,
+        graph_convertor: Optional[GraphFormatConvertor] = None,
+        train_idxs: Optional[List[int]] = None,
+        valid_idxs: Optional[List[int]] = None,
+        test_idxs: Optional[List[int]] = None,
+        batch_size: int = 32,
+        shuffle_batches: bool = True,
+        drop_last: bool = True,
+    ):
+        """Dataloader class for working with protein structure graphs.
+
+        :param pdb_file_list: List of paths to PDB files to use, defaults to None
+        :type pdb_file_list: Optional[List[str]], optional
+        :param pdb_code_list: List of PDB codes to use, defaults to None
+        :type pdb_code_list: Optional[List[str]], optional
+        :param pdb_in_dir: Path to direcotry containing pdb files to use, defaults to None
+        :type pdb_in_dir: Optional[str], optional
+        :param pdb_out_dir: Download location for PDB files, defaults to "/tmp/"
+        :type pdb_out_dir: str, optional
+        :param graph_labels: Labels for protein graphs, defaults to None
+        :type graph_labels: Optional[List[Union[np.ndarray, torch.Tensor]]], optional
+        :param node_labels: Labels for nodes in protein graphs, defaults to None
+        :type node_labels: Optional[List[Union[np.ndarray, torch.Tensor]]], optional
+        :param chain_list: List of chain selections for each structures, defaults to None
+        :type chain_list: Optional[List[str]], optional
+        :param protein_graph_config: Configuration object for protein graph construction, defaults to ProteinGraphConfig()
+        :type protein_graph_config: Optional[ ProteinGraphConfig ], optional
+        :param num_cores: Number of cores to use for multiprocessing of graph construction, defaults to 16
+        :type num_cores: int, optional
+        :param graph_convertor: GraphFormatConversion object to handle nc.Graph -> PyG Data conversion. Also specifies which data are kept from the source nx.Graph, defaults to None
+        :type graph_convertor: Optional[GraphFormatConvertor], optional
+        :param train_idxs: List of indices corresponding to training examples, defaults to None
+        :type train_idxs: Optional[List[int]], optional
+        :param valid_idxs: List of indices corresponding to validation examples, defaults to None
+        :type valid_idxs: Optional[List[int]], optional
+        :param test_idxs: List of indices corresponding to test examples, defaults to None
+        :type test_idxs: Optional[List[int]], optional
+        :param batch_size: Batch size for dataloader, defaults to 32
+        :type batch_size: int, optional
+        :param shuffle_batches: Whether or not to shuffle batches, defaults to True
+        :type shuffle_batches: bool, optional
+        :param drop_last: Whether or not to drop last batch, defaults to True
+        :type drop_last: bool, optional
+        """
+        # Graphein objects for constructing and converting graphs
+        self.protein_graph_config = protein_graph_config
+        self.graph_convertor = graph_convertor
+
+        # PDB inputs
+        self.pdb_file_list = pdb_file_list
+        self.pdb_code_list = pdb_code_list
+        self.pdb_in_dir = pdb_in_dir
+        self.pdb_out_dir = pdb_out_dir
+        self.chain_list = chain_list
+
+        # Indexes for train/test/split
+        self.train_idxs = train_idxs
+        self.valid_idxs = valid_idxs
+        self.test_idxs = test_idxs
+
+        # Get PDBs from input
+        if self.pdb_code_list:
+            self._validate_pdb_codes()
+            self.pdb_source = "code"
+            self.pdb_list = self.pdb_code_list
+        if self.pdb_file_list:
+            self._validate_pdb_files()
+            self.pdb_source = "file"
+            self.pdb_list = self.pdb_file_list
+        if self.pdb_in_dir:
+            self.pdb_file_list = os.listdir(self.pdb_in_dir)
+            self.pdb_file_list = [
+                f for f in self.pdb_file_list if f.endswith(".pdb")
+            ]
+            self.pdb_source = "file"
+            self.pdb_list = self.pdb_file_list
+        if self.chain_list is None:
+            self.chain_list = ["all"] * len(self.pdb_list)
+
+        # Graph labels
+        self.graph_labels = graph_labels
+        self.node_labels = node_labels
+
+        self._validate_input_data()
+
+        self._num_cores = num_cores
+
+        # Data loader params
+        self.batch_size = batch_size
+        self.shuffle_batches = shuffle_batches
+        self.drop_last = drop_last
+
+    def split_data(
+        self,
+    ) -> Tuple[GrapheinDataLoader, GrapheinDataLoader, GrapheinDataLoader]:
+        """Splits data based on provided indices.
+
+        :returns: Tuple of GrapheinDataLoaders for train, validation, and test
+        :rtype: Tuple[GrapheinDataLoader, GrapheinDataLoader, GrapheinDataLoader]
+        """
+        loaders: List[GrapheinDataLoader] = []
+        if self.train_idxs:
+            train = GrapheinDataLoader(
+                pdb_file_list=[
+                    self.pdb_file_list[idx] for idx in self.train_idxs
+                ]
+                if self.pdb_file_list is not None
+                else None,
+                pdb_code_list=[
+                    self.pdb_code_list[idx] for idx in self.train_idxs
+                ]
+                if self.pdb_code_list is not None
+                else None,
+                graph_labels=[
+                    self.graph_labels[idx] for idx in self.train_idxs
+                ]
+                if self.graph_labels is not None
+                else None,
+                node_labels=[self.node_labels[idx] for idx in self.train_idxs]
+                if self.node_labels is not None
+                else None,
+                chain_list=[self.chain_list[idx] for idx in self.train_idxs]
+                if self.node_labels is not None
+                else None,
+                pdb_in_dir=self.pdb_in_dir,
+                pdb_out_dir=self.pdb_out_dir,
+                protein_graph_config=self.protein_graph_config,
+                graph_convertor=self.graph_convertor,
+                num_cores=self._num_cores,
+                batch_size=self.batch_size,
+                shuffle_batches=self.shuffle_batches,
+                drop_last=self.drop_last,
+            )
+            self.train = train
+            loaders.append(train)
+        if self.valid_idxs:
+            valid = GrapheinDataLoader(
+                pdb_file_list=[
+                    self.pdb_file_list[idx] for idx in self.valid_idxs
+                ]
+                if self.pdb_file_list is not None
+                else None,
+                pdb_code_list=[
+                    self.pdb_code_list[idx] for idx in self.valid_idxs
+                ]
+                if self.pdb_code_list is not None
+                else None,
+                graph_labels=[
+                    self.graph_labels[idx] for idx in self.valid_idxs
+                ]
+                if self.graph_labels is not None
+                else None,
+                node_labels=[self.node_labels[idx] for idx in self.valid_idxs]
+                if self.node_labels is not None
+                else None,
+                chain_list=[self.chain_list[idx] for idx in self.valid_idxs]
+                if self.node_labels is not None
+                else None,
+                pdb_in_dir=self.pdb_in_dir,
+                pdb_out_dir=self.pdb_out_dir,
+                protein_graph_config=self.protein_graph_config,
+                graph_convertor=self.graph_convertor,
+                num_cores=self._num_cores,
+                batch_size=self.batch_size,
+                shuffle_batches=self.shuffle_batches,
+                drop_last=self.drop_last,
+            )
+            self.valid = valid
+            loaders.append(valid)
+        if self.test_idxs:
+            test = GrapheinDataLoader(
+                pdb_file_list=[
+                    self.pdb_file_list[idx] for idx in self.test_idxs
+                ]
+                if self.pdb_file_list is not None
+                else None,
+                pdb_code_list=[
+                    self.pdb_code_list[idx] for idx in self.test_idxs
+                ]
+                if self.pdb_code_list is not None
+                else None,
+                graph_labels=[self.graph_labels[idx] for idx in self.test_idxs]
+                if self.graph_labels is not None
+                else None,
+                node_labels=[self.node_labels[idx] for idx in self.test_idxs]
+                if self.node_labels is not None
+                else None,
+                chain_list=[self.chain_list[idx] for idx in self.test_idxs]
+                if self.node_labels is not None
+                else None,
+                pdb_in_dir=self.pdb_in_dir,
+                pdb_out_dir=self.pdb_out_dir,
+                protein_graph_config=self.protein_graph_config,
+                graph_convertor=self.graph_convertor,
+                num_cores=self._num_cores,
+                batch_size=self.batch_size,
+                shuffle_batches=self.shuffle_batches,
+                drop_last=self.drop_last,
+            )
+            self.test = test
+            loaders.append(test)
+
+        return tuple(loaders)
+
+    def construct_graphs(self) -> List[nx.Graph]:
+        """
+        Constructs graphs for protein chains in the dataset.
+
+        :param config: Config specifying protein graph construction parameters
+        :type config: graphein.protein.ProteinGraphConfig
+        :return: List of protein structure graphs
+        :rtype: List[nx.Graph]
+        """
+        graph_list = construct_graphs_mp(
+            pdb_code_it=self.pdb_code_list,
+            pdb_path_it=self.pdb_file_list,
+            chain_selections=self.chain_list,
+            num_cores=self._num_cores,
+            return_dict=False,
+        )
+
+        if self.graph_convertor:
+            graph_list = self.convert_graphs(graph_list, self.graph_convertor)
+
+        self.graphs = graph_list
+        return graph_list
+
+    def convert_graphs(
+        self, graph_list: List[nx.Graph], convertor: GraphFormatConvertor
+    ) -> List[Data]:
+        """Converts graphs to desired format.
+
+        :param graph_list: List of graphs to conert
+        :type graph_list: List[nx.Graph]
+        :param convertor: Graph format convertor
+        :type convertor: GraphFormatConvertor
+        :return: List of graphs converted to desired format.
+        :rtype: List[Data]
+        """
+        graph_list = [convertor(g) for g in graph_list]
+        if self.node_labels:
+            self._validate_node_labels(graph_list)
+            for i, g in enumerate(graph_list):
+                g.node_y = self.node_labels[i]
+        if self.graph_labels:
+            for i, g in enumerate(graph_list):
+                g.graph_y = self.graph_labels[i]
+        return graph_list
+
+    def dataloader(self) -> DataLoader:
+        """Returns a PyTorch Geometric Dataloader for the dataset.
+
+        :return: Pytorch Geometric Dataloader
+        :rtype: DataLoader
+        """
+        loader = DataLoader(
+            self.graphs,
+            batch_size=self.batch_size,
+            shuffle=self.shuffle_batches,
+            drop_last=self.drop_last,
+        )
+        self.loader = loader
+        return loader
+
+    def download_pdbs(self):
+        """
+        Downloads dataset PDBs to a specified directories.
+
+        :param path: Path to desired output location
+        :type path: str
+        """
+        pdbl = PDBList()
+        pdbl.download_pdb_files(
+            pdb_codes=self.pdb_code_list, pdir=self.pdb_out_dir
+        )
+
+    def _validate_pdb_codes(self):
+        """Validates provided PDB codes."""
+        for code in self.pdb_code_list:
+            assert (
+                len(code) == 4
+            ), f"PDB codes {code} must be 4 characters long."
+
+    def _validate_pdb_files(self):
+        """Validates provided PDB files by asserting that they exist."""
+        for file in self.pdb_file_list:
+            assert os.path.isfile(file), f"PDB file {file} does not exist."
+
+    def _validate_node_labels(self, graph_list: List[nx.Graph]):
+        """Checks length of node labels matches number of nodes in graph."""
+        for i, g in enumerate(graph_list):
+            assert (
+                g.number_of_nodes() == self.node_labels[i].shape[0]
+            ), f"Node label lengths ({self.node_labels[i].shape[0]}) do not match number of noes ({g.number_of_nodes()})"
+
+    def _validate_input_data(self):
+        """Checks length of provided inputs match with number of structures provided."""
+        if self.node_labels:
+            assert len(self.node_labels) == len(
+                self.pdb_list
+            ), f"Node labels length ({len(self.node_labels)}) does not match PDB list length ({len(self.pdb_list)})"
+        if self.graph_labels:
+            assert len(self.graph_labels) == len(
+                self.pdb_list
+            ), f"Graph labels length ({len(self.graph_labels)}) does not match PDB list length ({len(self.pdb_list)})"
+        if self.chain_list:
+            assert len(self.chain_list) == len(
+                self.pdb_list
+            ), f"Chain list length ({len(self.chain_list)}) does not match PDB list length ({len(self.pdb_list)})"
+
+    def __repr__(self) -> str:
+        """Print summary for dataloader."""
+        summary = (
+            f"Graphein Dataloader containing: {len(self.pdb_list)} PDBs. \n"
+        )
+        if self.graph_labels:
+            summary += "Using graph labels.\n"
+        if self.node_labels:
+            summary += "Using node labels.\n"
+        return summary
+
+    def __len__(self) -> int:
+        return len(self.pdb_list) if not self.graphs else len(self.graphs)
+
+
+if __name__ == "__main__":
+    import numpy as np
+
+    pdb_list = [
+        "3eiy",
+        "4hhb",
+        "1a0a",
+        "1a0b",
+        "1a0c",
+        "1a0d",
+        "1a0e",
+        "1a0f",
+        "1a0g",
+        "1a0h",
+        "1a0i",
+        "1a0j",
+        "1a0k",
+        "1a0l",
+        "1a0m",
+        "1a0n",
+        "1a0o",
+        "1a0p",
+        "1a0q",
+        "1a0r",
+        "1a0s",
+        "1a0t",
+        "1a0u",
+        "1a0v",
+        "1a0w",
+        "1a0x",
+        "1a0y",
+    ]
+    train_idxs = list(range(0, 10))
+    valid_idx = list(range(10, 20))
+    test_idx = list(range(20, 25))
+
+    format_convertor = GraphFormatConvertor(
+        "nx", "pyg", verbose="gnn", columns=None
+    )
+
+    dl = GrapheinDataLoader(
+        pdb_code_list=pdb_list,
+        train_idxs=train_idxs,
+        valid_idxs=valid_idx,
+        test_idxs=test_idx,
+        graph_convertor=format_convertor,
+        batch_size=2,
+        graph_labels=[np.random.randint(1)] * 27,
+    )
+
+    dl.split_data()
+    dl.train.construct_graphs()
+    dl.valid.construct_graphs()
+    dl.test.construct_graphs()
diff --git a/graphein/datasets/proteins_ligands.py b/graphein/datasets/proteins_ligands.py
new file mode 100644
index 00000000..776a7715
--- /dev/null
+++ b/graphein/datasets/proteins_ligands.py
@@ -0,0 +1,45 @@
+"""Class for working with the PROTEINS_LIGANDS dataset"""
+# Graphein
+# Author: Arian Jamasb <arian@jamasb.io>
+# License: MIT
+# Code Repository: https://github.com/a-r-j/graphein
+
+from functools import lru_cache
+from pathlib import Path
+
+import pandas as pd
+
+from graphein.datasets.base import AbstractClassificationDataset
+
+
+class PROTEINS_LIGANDS(AbstractClassificationDataset):
+    @lru_cache
+    def _load_dataset(self=None) -> pd.DataFrame:
+        file_path = (
+            Path(__file__).parent.parent.parent.resolve()
+            / "datasets"
+            / "proteins_ligands"
+            / "PROTEINS_LIGANDS.csv"
+        )
+        return pd.read_csv(file_path)
+
+    def split_data(
+        self,
+        strategy: str = "random",
+        training_size: float = 0.7,
+        validation_size: float = 0.2,
+        test_size: float = 0.1,
+        shuffle: bool = True,
+    ):
+
+        train_df, val_df, test_df = self._split_data(
+            strategy, training_size, validation_size, test_size, shuffle
+        )
+
+        self.train: PROTEINS_LIGANDS = PROTEINS_LIGANDS(train_df)
+        self.validation: PROTEINS_LIGANDS = PROTEINS_LIGANDS(val_df)
+        self.test: PROTEINS_LIGANDS = PROTEINS_LIGANDS(test_df)
+
+
+if __name__ == "__main__":
+    c = PROTEINS_LIGANDS()
diff --git a/graphein/datasets/proteins_metal.py b/graphein/datasets/proteins_metal.py
new file mode 100644
index 00000000..fe965983
--- /dev/null
+++ b/graphein/datasets/proteins_metal.py
@@ -0,0 +1,51 @@
+"""Class for working with the PROTEINS_METAL dataset"""
+# Graphein
+# Author: Arian Jamasb <arian@jamasb.io>
+# License: MIT
+# Code Repository: https://github.com/a-r-j/graphein
+
+from functools import lru_cache
+from pathlib import Path
+
+import pandas as pd
+
+from graphein.datasets.base import AbstractClassificationDataset
+from graphein.protein.config import ProteinGraphConfig
+
+
+class PROTEINS_METAL(AbstractClassificationDataset):
+    @lru_cache
+    def _load_dataset(self=None) -> pd.DataFrame:
+        file_path = (
+            Path(__file__).parent.parent.parent.resolve()
+            / "datasets"
+            / "proteins_metal"
+            / "PROTEINS_METAL.csv"
+        )
+        return pd.read_csv(file_path)
+
+    def split_data(
+        self,
+        strategy: str = "random",
+        training_size: float = 0.7,
+        validation_size: float = 0.2,
+        test_size: float = 0.1,
+        shuffle: bool = True,
+    ):
+
+        train_df, val_df, test_df = self._split_data(
+            strategy, training_size, validation_size, test_size, shuffle
+        )
+
+        self.train: PROTEINS_METAL = PROTEINS_METAL(train_df)
+        self.validation: PROTEINS_METAL = PROTEINS_METAL(val_df)
+        self.test: PROTEINS_METAL = PROTEINS_METAL(test_df)
+
+
+if __name__ == "__main__":
+    config = ProteinGraphConfig(keep_hets=True)
+    c = PROTEINS_METAL()
+    c.split_data()
+    c.train.construct_graphs(config)
+
+    print(c.df)
diff --git a/graphein/datasets/proteins_nucleic.py b/graphein/datasets/proteins_nucleic.py
new file mode 100644
index 00000000..08b12979
--- /dev/null
+++ b/graphein/datasets/proteins_nucleic.py
@@ -0,0 +1,45 @@
+"""Class for working with the PROTEINS_NUCLEIC dataset"""
+
+
+# License: MIT
+# Code Repository: https://github.com/a-r-j/graphein
+from functools import lru_cache
+# Graphein
+# Author: Arian Jamasb <arian@jamasb.io>
+from pathlib import Path
+
+import pandas as pd
+
+from graphein.datasets.base import AbstractClassificationDataset
+
+
+class PROTEINS_NUCLEIC(AbstractClassificationDataset):
+    @lru_cache
+    def _load_dataset(self=None) -> pd.DataFrame:
+        file_path = (
+            Path(__file__).parent.parent.parent.resolve()
+            / "datasets"
+            / "proteins_nucleic"
+            / "PROTEINS_NUCLEIC.csv"
+        )
+        return pd.read_csv(file_path)
+
+    def split_data(
+        self,
+        strategy: str = "random",
+        training_size: float = 0.7,
+        validation_size: float = 0.2,
+        test_size: float = 0.1,
+        shuffle: bool = True,
+    ):
+        train_df, val_df, test_df = self._split_data(
+            strategy, training_size, validation_size, test_size, shuffle
+        )
+
+        self.train: PROTEINS_NUCLEIC = PROTEINS_NUCLEIC(train_df)
+        self.validation: PROTEINS_NUCLEIC = PROTEINS_NUCLEIC(val_df)
+        self.test: PROTEINS_NUCLEIC = PROTEINS_NUCLEIC(test_df)
+
+
+if __name__ == "__main__":
+    c = PROTEINS_NUCLEIC()
diff --git a/graphein/datasets/proteins_nucleotides.py b/graphein/datasets/proteins_nucleotides.py
new file mode 100644
index 00000000..96650204
--- /dev/null
+++ b/graphein/datasets/proteins_nucleotides.py
@@ -0,0 +1,44 @@
+"""Class for working with the PROTEINS_NUCLEOTIDES dataset"""
+# Graphein
+# Author: Arian Jamasb <arian@jamasb.io>
+# License: MIT
+# Code Repository: https://github.com/a-r-j/graphein
+
+from functools import lru_cache
+from pathlib import Path
+
+import pandas as pd
+
+from graphein.datasets.base import AbstractClassificationDataset
+
+
+class PROTEINS_NUCLEOTIDES(AbstractClassificationDataset):
+    @lru_cache
+    def _load_dataset(self=None) -> pd.DataFrame:
+        file_path = (
+            Path(__file__).parent.parent.parent.resolve()
+            / "datasets"
+            / "proteins_nucleotides"
+            / "PROTEINS_NUCLEOTIDES.csv"
+        )
+        return pd.read_csv(file_path)
+
+    def split_data(
+        self,
+        strategy: str = "random",
+        training_size: float = 0.7,
+        validation_size: float = 0.2,
+        test_size: float = 0.1,
+        shuffle: bool = True,
+    ):
+        train_df, val_df, test_df = self._split_data(
+            strategy, training_size, validation_size, test_size, shuffle
+        )
+
+        self.train = PROTEINS_NUCLEOTIDES(train_df)
+        self.validation = PROTEINS_NUCLEOTIDES(val_df)
+        self.test = PROTEINS_NUCLEOTIDES(test_df)
+
+
+if __name__ == "__main__":
+    c = PROTEINS_NUCLEOTIDES()
diff --git a/graphein/datasets/pscdb.py b/graphein/datasets/pscdb.py
new file mode 100644
index 00000000..2c3586d8
--- /dev/null
+++ b/graphein/datasets/pscdb.py
@@ -0,0 +1,181 @@
+"""Class for working with the PSCDB dataset"""
+import logging
+import multiprocessing
+import time
+# Graphein
+# Author: Arian Jamasb <arian@jamasb.io>
+# License: MIT
+# Code Repository: https://github.com/a-r-j/graphein
+import traceback
+from functools import partial
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import networkx as nx
+import numpy as np
+import pandas as pd
+from Bio.PDB import PDBList
+from tqdm import tqdm
+
+from graphein.protein.config import ProteinGraphConfig
+from graphein.protein.graphs import construct_graph
+
+logging.basicConfig(level="DEBUG")
+log = logging.getLogger(__name__)
+
+
+class PSCDB:
+    def __init__(
+        self,
+        compute_bound_graphs: bool = False,
+        compute_free_graphs: bool = False,
+        protein_graph_config: Optional[
+            ProteinGraphConfig
+        ] = ProteinGraphConfig(),
+        num_cores: int = 8,
+    ):
+        """
+        Initialises the PSCDB Dataset class
+
+        :param compute_bound_graphs: Whether or not to compute (ligand) bound protein graphs
+        :type compute_bound_graphs:  bool
+        :param compute_free_graphs: Whether or not to compute unbound protein graphs
+        :type compute_free_graphs: bool
+        :param protein_graph_config: Config object to be used for protein graph construction
+        :type protein_graph_config: graphein.protein.config.ProteinGraphConfig
+        :param num_cores: Number of cores to use for graph construction. The more the merrier :)
+        :type num_cores: int, defaults to 8.
+        """
+        self.df = self._load_dataset()
+        self.bound_pdbs: List[str] = self.get_bound_pdb()
+        self.bound_chains: List[str] = self.get_bound_chains()
+        self.free_pdbs: List[str] = self.get_free_pdb()
+        self.free_chains: List[str] = self.get_free_chains()
+        self.protein_names: List[str] = self.get_protein_names()
+        self.ligands: List[str] = self.get_ligands()
+        self.classification: List[str] = self.get_classification()
+        self.motion_type: List[str] = self.get_motion_type()
+        self.PSCID: List[str] = self.get_pscid()
+        self.protein_names: List[str] = self.get_protein_names()
+
+        self._num_cores: int = num_cores
+        self.bad_pdbs: List[str] = []
+
+        if compute_bound_graphs:
+            self.config = protein_graph_config
+            self.bound_graphs: List[nx.Graph] = self.construct_graphs(
+                pdbs=self.bound_pdbs,
+                chains=self.bound_chains,
+            )
+        if compute_free_graphs:
+            self.config = protein_graph_config
+            self.free_graphs: List[nx.Graph] = self.construct_graphs(
+                pdbs=self.free_pdbs,
+                chains=self.free_chains,
+            )
+
+    @staticmethod
+    def _load_dataset() -> pd.DataFrame:
+        file_path = Path(
+            "../../datasets/pscdb/structural_rearrangement_data.csv"
+        )
+        return pd.read_csv(file_path)
+
+    def get_pscid(self) -> List[str]:
+        """
+        Returns list of PSCID (PSCDB identifiers)
+
+        :return: List of PSCDB Identifiers
+        :rtype: List[str]
+        """
+        return list(self.df["PSCID"])
+
+    def get_protein_names(self) -> List[str]:
+        return list(self.df["Protein Name"])
+
+    def get_ligands(self) -> List[str]:
+        return list(self.df["Ligands"])
+
+    def get_classification(self) -> List[str]:
+        return list(self.df["Classification(?)"])
+
+    def get_motion_type(self) -> List[str]:
+        return list(self.df["motion_type"])
+
+    def get_bound_pdb(self):
+        return list(self.df["Bound PDB"])
+
+    def get_bound_chains(self):
+        return list(self.df["Bound Chains"])
+
+    def get_free_pdb(self):
+        return list(self.df["Free PDB"])
+
+    def get_free_chains(self):
+        return list(self.df["Free Chains"])
+
+    def construct_graphs(
+        self, pdbs: List[str], chains: List[str]
+    ) -> List[nx.Graph]:
+        pool = multiprocessing.Pool(self._num_cores)
+        graph_list = list(
+            pool.map(
+                self._graph_constructor,
+                [(pdb, chains[i]) for i, pdb in enumerate(pdbs)],
+            )
+        )
+        pool.close()
+        pool.join()
+        return graph_list
+
+    def _graph_constructor(self, args: Tuple[str, str]):
+        log.info(
+            f"Constructing graph for: {args[0]}. Chain selection: {args[1]}"
+        )
+        func = partial(construct_graph, config=self.config)
+        try:
+            result = func(pdb_code=args[0], chain_selection=args[1])
+            return result
+        except Exception:
+            log.info(
+                f"Graph construction error (PDB={args[0]})! {traceback.format_exc()}"
+            )
+            self.bad_pdbs.append(args[0])
+
+    def split_dataset(
+        self,
+        strategy: str,
+        train_size: float = 0.7,
+        val_size: float = 0.2,
+        test_size: float = 0.1,
+    ):
+        assert (
+            train_size + val_size + test_size == 1.0
+        ), "Train, test & val sizes must sum to 1"
+
+        if strategy == "random":
+            pass
+
+    def download_pdbs(self, path: str):
+        """
+        Downloads dataset PDBs to a specified directories
+
+        :param path: Path to desired output location
+        :type path: str
+        """
+        pdbl = PDBList()
+        pdbl.download_pdb_files(pdb_codes=self.pdb_list, pdir=path)
+
+    def __len__(self):
+        """Returns length of the dataset
+
+        :returns: Dataset length
+        :rtype: int
+        """
+        return len(self.df)
+
+
+if __name__ == "__main__":
+    ts = time.time()
+    c = PSCDB(compute_free_graphs=True, num_cores=16)
+    print(time.time() - ts)
diff --git a/graphein/grn/features/node_features.py b/graphein/grn/features/node_features.py
index 1991d737..32764442 100644
--- a/graphein/grn/features/node_features.py
+++ b/graphein/grn/features/node_features.py
@@ -1,14 +1,19 @@
+import logging
+
 from graphein.utils.utils import import_message
 
+log = logging.getLogger(__name__)
+
 try:
     from bioservices import HGNC, UniProt
 except ImportError:
-    import_message(
+    message = import_message(
         submodule="graphein.grn.features.node_features",
         package="bioservices",
         conda_channel="bioconda",
         pip_install=True,
     )
+    log.warning(message)
 
 
 def add_sequence_to_nodes(n, d):
diff --git a/graphein/ppi/features/node_features.py b/graphein/ppi/features/node_features.py
index 7d7dde6a..5162c93a 100644
--- a/graphein/ppi/features/node_features.py
+++ b/graphein/ppi/features/node_features.py
@@ -1,4 +1,6 @@
 """Functions for adding nodes features to a PPI Graph"""
+import logging
+
 # %%
 # Graphein
 # Author: Ramon Vinas, Arian Jamasb <arian@jamasb.io>
@@ -11,15 +13,19 @@
 
 from graphein.utils.utils import import_message
 
+log = logging.getLogger(__name__)
+
+
 try:
     from bioservices import HGNC, UniProt
 except ImportError:
-    import_message(
+    message = import_message(
         submodule="graphein.ppi.features.nodes_features",
         package="bioservices",
         conda_channel="bioconda",
         pip_install=True,
     )
+    log.warning(message)
 
 
 def add_sequence_to_nodes(n: str, d: Dict[str, Any]):
diff --git a/graphein/protein/features/nodes/aaindex.py b/graphein/protein/features/nodes/aaindex.py
index f21c5ee8..8e8e9aae 100644
--- a/graphein/protein/features/nodes/aaindex.py
+++ b/graphein/protein/features/nodes/aaindex.py
@@ -1,17 +1,22 @@
+import logging
 from typing import Dict, Tuple
 
 import networkx as nx
 
 from graphein.utils.utils import import_message, protein_letters_3to1_all_caps
 
+log = logging.getLogger(__name__)
+
+
 try:
     from pyaaisc import Aaindex
 except ImportError:
-    import_message(
+    message = import_message(
         submodule="graphein.protein.features.nodes.aaindex",
         package="pyaaisc",
         pip_install=True,
     )
+    log.warning(message)
 
 
 def fetch_AAIndex(accession: str) -> Tuple[str, Dict[str, float]]:
diff --git a/graphein/protein/features/sequence/embeddings.py b/graphein/protein/features/sequence/embeddings.py
index 5ffa8edb..878b6013 100644
--- a/graphein/protein/features/sequence/embeddings.py
+++ b/graphein/protein/features/sequence/embeddings.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import os
 from functools import lru_cache, partial
 from pathlib import Path
@@ -12,24 +13,29 @@
 )
 from graphein.utils.utils import import_message
 
+log = logging.getLogger(__name__)
+
+
 try:
     import torch
 except ImportError:
-    import_message(
+    message = import_message(
         submodule="graphein.protein.features.sequence.embeddings",
         package="torch",
         pip_install=True,
         conda_channel="pytorch",
     )
+    log.warning(message)
 
 try:
     import biovec
 except ImportError:
-    import_message(
+    message = import_message(
         submodule="graphein.protein.features.sequence.embeddings",
         package="biovec",
         pip_install=True,
     )
+    log.warning(message)
 
 
 @lru_cache()
diff --git a/graphein/protein/graphs.py b/graphein/protein/graphs.py
index 71ca8384..7d3cb4f3 100644
--- a/graphein/protein/graphs.py
+++ b/graphein/protein/graphs.py
@@ -8,7 +8,10 @@
 from __future__ import annotations
 
 import logging
-from typing import Callable, List, Optional
+import multiprocessing
+import traceback
+from functools import partial
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import networkx as nx
 import numpy as np
@@ -623,6 +626,106 @@ def construct_graph(
     return g
 
 
+def _mp_graph_constructor(
+    args: Tuple[str, str], use_pdb_code: bool, config: ProteinGraphConfig
+) -> nx.Graph:
+    """
+    Protein graph constructor for use in multiprocessing several protein structure graphs.
+
+    :param args: Tuple of pdb code/path and the chain selection for that PDB
+    :type args: Tuple[str, str]
+    :param use_pdb_code: Whether or not we are using pdb codes or paths
+    :type use_pdb_code: bool
+    :param config: Protein structure graph construction config
+    :type config: ProteinGraphConfig
+    :return: Protein structure graph
+    :rtype: nx.Graph
+    """
+    log.info(f"Constructing graph for: {args[0]}. Chain selection: {args[1]}")
+    func = partial(construct_graph, config=config)
+    if use_pdb_code:
+        try:
+            result = func(pdb_code=args[0], chain_selection=args[1])
+            return result
+        except Exception as ex:
+            log.info(
+                f"Graph construction error (PDB={args[0]})! {traceback.format_exc()}"
+            )
+            log.info(ex)
+            return None
+    elif not use_pdb_code:
+        try:
+            result = func(pdb_path=args[0], chain_selection=args[1])
+            return result
+        except Exception as ex:
+            log.info(
+                f"Graph construction error (PDB={args[0]})! {traceback.format_exc()}"
+            )
+            log.info(ex)
+            return None
+
+
+def construct_graphs_mp(
+    pdb_code_it: Optional[List[str]] = None,
+    pdb_path_it: Optional[List[str]] = None,
+    chain_selections: Optional[list[str]] = None,
+    config: ProteinGraphConfig = ProteinGraphConfig(),
+    num_cores: int = 16,
+    return_dict: bool = True,
+) -> Union[List[nx.Graph], Dict[str, nx.Graph]]:
+    """
+    Constructs protein graphs for a list of pdb codes or pdb paths using multiprocessing.
+
+    :param pdb_code_it: List of pdb codes to use for protein graph construction
+    :type pdb_code_it: Optional[List[str]], defaults to None
+    :param pdb_path_it: List of paths to PDB files to use for protein graph construction
+    :type pdb_path_it: Optional[List[str]], defaults to None
+    :param chain_selections: List of chains to select from the protein structures (e.g. ["ABC", "A", "L", "CD"...])
+    :type chain_selections: Optional[List[str]], defaults to None
+    :param config: ProteinGraphConfig to use.
+    :type config: graphein.protein.config.ProteinGraphConfig, defaults to default config params
+    :param num_cores: Number of cores to use for multiprocessing. The more the merrier
+    :type num_cores: int, defaults to 16
+    :param return_dict: Whether or not to return a dictionary (indexed by pdb codes/paths) or a list of graphs.
+    :type return_dict: bool, default to True
+    :return: Iterable of protein graphs. None values indicate there was a problem in constructing the graph for this particular pdb
+    :rtype: Union[List[nx.Graph], Dict[str, nx.Graph]]
+    """
+    assert (
+        pdb_code_it is not None or pdb_path_it is not None
+    ), "Iterable of pdb codes OR pdb paths required."
+
+    if pdb_code_it is not None:
+        pdbs = pdb_code_it
+        use_pdb_code = True
+
+    if pdb_path_it is not None:
+        pdbs = pdb_path_it
+        use_pdb_code = False
+
+    if chain_selections is None:
+        chain_selections = ["all"] * len(pdbs)
+
+    constructor = partial(
+        _mp_graph_constructor, use_pdb_code=use_pdb_code, config=config
+    )
+
+    pool = multiprocessing.Pool(num_cores)
+    graphs = list(
+        pool.map(
+            constructor,
+            [(pdb, chain_selections[i]) for i, pdb in enumerate(pdbs)],
+        )
+    )
+    pool.close()
+    pool.join()
+
+    if return_dict:
+        graphs = {pdb: graphs[i] for i, pdb in enumerate(pdbs)}
+
+    return graphs
+
+
 if __name__ == "__main__":
     from functools import partial
 
@@ -643,9 +746,29 @@ def construct_graph(
         partial(add_k_nn_edges, k=3, long_interaction_threshold=0)
     ]
     # Test High-level API
-    g = construct_graph(
-        config=config,
-        pdb_path="../examples/pdbs/3eiy.pdb",
+    # g = construct_graph(
+    #    config=config,
+    #    pdb_path="../examples/pdbs/3eiy.pdb",
+    # )
+
+    # Test multiprocessing
+    graph_list = [
+        "2olg",
+        "1bjq",
+        "1omr",
+        "1a4g",
+        "2je9",
+        "3vm5",
+        "1el1",
+        "3fzo",
+        "1mn1",
+        "1ff5",
+        "1fic",
+        "3a47",
+        "1bir",
+    ] * 5
+    g = construct_graphs_mp(
+        pdb_code_it=graph_list, config=ProteinGraphConfig(), return_dict=True
     )
 
     """
diff --git a/graphein/protein/meshes.py b/graphein/protein/meshes.py
index a77da1e5..4b59b15c 100644
--- a/graphein/protein/meshes.py
+++ b/graphein/protein/meshes.py
@@ -14,18 +14,19 @@
 from graphein.utils.pymol import MolViewer
 from graphein.utils.utils import import_message
 
+log = logging.getLogger(__name__)
+
+
 try:
     from pytorch3d.structures import Meshes
 except ImportError:
-    import_message(
+    message = import_message(
         submodule="graphein.protein.meshes",
         package="pytorch3d",
         conda_channel="pytorch3d",
         pip_install=True,
     )
-
-
-log = logging.getLogger(__name__)
+    log.warning(message)
 
 
 def check_for_pymol_installation():
diff --git a/graphein/protein/visualisation.py b/graphein/protein/visualisation.py
index 50f3dcdc..f47e3736 100644
--- a/graphein/protein/visualisation.py
+++ b/graphein/protein/visualisation.py
@@ -6,6 +6,7 @@
 # Code Repository: https://github.com/a-r-j/graphein
 from __future__ import annotations
 
+import logging
 from itertools import count
 from typing import List, Optional, Tuple
 
@@ -18,14 +19,18 @@
 
 from graphein.utils.utils import import_message
 
+log = logging.getLogger(__name__)
+
+
 try:
     from pytorch3d.ops import sample_points_from_meshes
 except ImportError:
-    import_message(
+    message = import_message(
         submodule="graphein.protein.visualisation",
         package="pytorch3d",
         conda_channel="pytorch3d",
     )
+    log.warning(message)
 
 
 def plot_pointcloud(mesh: Meshes, title: str = "") -> Axes3D:
diff --git a/graphein/utils/utils.py b/graphein/utils/utils.py
index 8c7aa289..2dd77e57 100644
--- a/graphein/utils/utils.py
+++ b/graphein/utils/utils.py
@@ -1,4 +1,4 @@
-"""Utilities for working with graph objects"""
+"""Utilities for working with graph objects."""
 # Graphein
 # Author: Arian Jamasb <arian@jamasb.io>, Eric Ma
 # License: MIT
@@ -318,7 +318,7 @@ def import_message(
     package: str,
     conda_channel: Optional[str] = None,
     pip_install: bool = False,
-):
+) -> str:
     """
     Return warning if package is not found.
     Generic message for indicating to the user when a function relies on an
@@ -348,17 +348,12 @@ def import_message(
         installable = False
         installation = f"{package} cannot be installed via pip"
 
-    print(
-        f"To use the Graphein submodule {submodule}, you need to install "
-        f"{package}."
-    )
-    print()
+    message = f"To use the Graphein submodule {submodule}, you need to install: {package} "
     if installable:
-        print("To do so, use the following command:")
-        print()
-        print(f"    {installation}")
+        message += f"\nTo do so, use the following command: {installation}"
     else:
-        print(f"{installation}")
+        message += f"\n{installation}"
+    return message
 
 
 def parse_config(path: Path) -> ProteinGraphConfig:


Test patch:
diff --git a/tests/protein/test_graphs.py b/tests/protein/test_graphs.py
index f56866a5..ba534255 100644
--- a/tests/protein/test_graphs.py
+++ b/tests/protein/test_graphs.py
@@ -38,7 +38,11 @@
     esm_sequence_embedding,
 )
 from graphein.protein.features.sequence.sequence import molecular_weight
-from graphein.protein.graphs import construct_graph, read_pdb_to_dataframe
+from graphein.protein.graphs import (
+    construct_graph,
+    construct_graphs_mp,
+    read_pdb_to_dataframe,
+)
 
 DATA_PATH = Path(__file__).resolve().parent / "test_data" / "4hhb.pdb"
 
@@ -91,6 +95,37 @@ def test_construct_graph():
     assert len(peptide_bond_edges) == 570
 
 
+def test_construct_graphs_mp():
+    graph_list = [
+        "2olg",
+        "1bjq",
+        "1omr",
+        "1a4g",
+        "2je9",
+        "3vm5",
+        "1el1",
+        "3fzo",
+        "1mn1",
+        "1ff5",
+        "1fic",
+        "3a47",
+        "1bir",
+    ] * 5
+
+    g = construct_graphs_mp(
+        pdb_code_it=graph_list, config=ProteinGraphConfig(), return_dict=True
+    )
+    assert isinstance(g, dict)
+    assert len(g.keys()) == len(graph_list) / 5
+    for k, v in g.items():
+        assert isinstance(v, (nx.Graph, None))
+    g = construct_graphs_mp(
+        pdb_code_it=graph_list, config=ProteinGraphConfig(), return_dict=False
+    )
+    assert isinstance(g, list)
+    assert len(g) == len(graph_list)
+
+
 def test_chain_selection():
     """Example-based test that chain selection works correctly.