"""Pytorch Geometric Dataset classes for Protein Graphs."""
# Graphein
# Author: Arian Jamasb <arian@jamasb.io>
# License: MIT
# Project Website: https://github.com/a-r-j/graphein
# Code Repository: https://github.com/a-r-j/graphein
from __future__ import annotations
import os
from pathlib import Path
from typing import Callable, Dict, Generator, List, Optional
import networkx as nx
from loguru import logger as log
from tqdm.auto import tqdm
from graphein.ml.conversion import GraphFormatConvertor
from graphein.protein.config import ProteinGraphConfig
from graphein.protein.graphs import construct_graphs_mp
from graphein.protein.utils import (
download_alphafold_structure,
download_pdb,
download_pdb_multiprocessing,
)
from graphein.utils.dependencies import import_message
try:
import torch
from torch_geometric.data import Data, Dataset, InMemoryDataset
except ImportError:
import_message(
"graphein.ml.datasets.torch_geometric_dataset",
"torch_geometric",
conda_channel="pyg",
pip_install=True,
)
[docs]
class InMemoryProteinGraphDataset(InMemoryDataset):
[docs]
def __init__(
self,
root: str,
name: str,
paths: Optional[List[str]] = None,
pdb_codes: Optional[List[str]] = None,
uniprot_ids: Optional[List[str]] = None,
graph_label_map: Optional[Dict[str, torch.Tensor]] = None,
node_label_map: Optional[Dict[str, torch.Tensor]] = None,
chain_selection_map: Optional[Dict[str, List[str]]] = None,
graphein_config: ProteinGraphConfig = ProteinGraphConfig(),
graph_format_convertor: GraphFormatConvertor = GraphFormatConvertor(
src_format="nx", dst_format="pyg"
),
graph_transformation_funcs: Optional[List[Callable]] = None,
transform: Optional[Callable] = None,
pdb_transform: Optional[List[Callable]] = None,
pre_transform: Optional[Callable] = None,
pre_filter: Optional[Callable] = None,
num_cores: int = 16,
af_version: int = 6,
):
"""In Memory dataset for protein graphs.
Dataset base class for creating graph datasets which easily fit
into CPU memory. Inherits from
:class:`torch_geometric.data.InMemoryDataset`, which inherits from
:class:`torch_geometric.data.Dataset`.
See `here <https://pytorch-geometric.readthedocs.io/en/latest/notes/
create_dataset.html#creating-in-memory-datasets>`__ for the accompanying
tutorial.
:param root: Root directory where the dataset should be saved.
:type root: str
:param name: Name of the dataset. Will be saved to ``data_$name.pt``.
:type name: str
:param paths: List of full path of PDB or MMTF files to load. Defaults
to ``None``.
:type paths: Optional[List[str]], optional
:param pdb_codes: List of PDB codes to download and parse from the PDB.
Defaults to None.
:type pdb_codes: Optional[List[str]], optional
:param uniprot_ids: List of Uniprot IDs to download and parse from
Alphafold Database. Defaults to ``None``.
:type uniprot_ids: Optional[List[str]], optional
:param graph_label_map: Dictionary mapping PDB/Uniprot IDs to
graph-level labels. Defaults to ``None``.
:type graph_label_map: Optional[Dict[str, Tensor]], optional
:param node_label_map: Dictionary mapping PDB/Uniprot IDs to node-level
labels. Defaults to ``None``.
:type node_label_map: Optional[Dict[str, torch.Tensor]], optional
:param chain_selection_map: Dictionary mapping, defaults to ``None``.
:type chain_selection_map: Optional[Dict[str, List[str]]], optional
:param graphein_config: Protein graph construction config, defaults to
``ProteinGraphConfig()``.
:type graphein_config: ProteinGraphConfig, optional
:param graph_format_convertor: Conversion handler for graphs, defaults
to ``GraphFormatConvertor(src_format="nx", dst_format="pyg")``.
:type graph_format_convertor: GraphFormatConvertor, optional
:param pdb_transform: List of functions that consume a list of paths to
the downloaded structures. This provides an entry point to apply
pre-processing from bioinformatics tools of your choosing. Defaults
to ``None``.
:type pdb_transform: Optional[List[Callable]], optional
:param graph_transformation_funcs: List of functions that consume a
``nx.Graph`` and return a ``nx.Graph``. Applied to graphs after
construction but before conversion to pyg. Defaults to ``None``.
:type graph_transformation_funcs: Optional[List[Callable]], optional
:param transform: A function/transform that takes in a
``torch_geometric.data.Data`` object and returns a transformed
version. The data object will be transformed before every access.
Defaults to ``None``.
:type transform: Optional[Callable], optional
:param pre_transform: A function/transform that takes in an
``torch_geometric.data.Data`` object and returns a transformed
version. The data object will be transformed before being saved to
disk. Defaults to ``None``.
:type pre_transform: Optional[Callable], optional
:param pre_filter: A function that takes in a
``torch_geometric.data.Data`` object and returns a boolean value,
indicating whether the data object should be included in the final
dataset. Optional, defaults to ``None``.
:type pre_filter: Optional[Callable], optional
:param num_cores: Number of cores to use for multiprocessing of graph
construction, defaults to ``16``.
:type num_cores: int, optional
:param af_version: Version of AlphaFoldDB structures to use,
defaults to ``6``.
:type af_version: int, optional
"""
self.name = name
self.pdb_codes = (
[pdb.lower() for pdb in pdb_codes]
if pdb_codes is not None
else None
)
self.uniprot_ids = (
[up.upper() for up in uniprot_ids]
if uniprot_ids is not None
else None
)
self.paths = paths
if self.paths is None:
if self.pdb_codes and self.uniprot_ids:
self.structures = self.pdb_codes + self.uniprot_ids
elif self.pdb_codes:
self.structures = pdb_codes
elif self.uniprot_ids:
self.structures = uniprot_ids
# Use local saved pdb_files instead of download or move them to
# self.root/raw dir
else:
if isinstance(self.paths, list):
self.structures = [
os.path.splitext(os.path.split(path)[-1])[0]
for path in self.paths
]
self.path, _ = os.path.split(self.paths[0])
if self.pdb_codes and self.uniprot_ids:
self.structures = self.pdb_codes + self.uniprot_ids
elif self.pdb_codes:
self.structures = pdb_codes
elif self.uniprot_ids:
self.structures = uniprot_ids
self.af_version = af_version
self.bad_pdbs: List[
str
] = [] # list of pdb codes that failed to download
# Labels & Chains
self.graph_label_map = graph_label_map
self.node_label_map = node_label_map
self.chain_selection_map = chain_selection_map
# Configs
self.config = graphein_config
self.graph_format_convertor = graph_format_convertor
self.graph_transformation_funcs = graph_transformation_funcs
self.pdb_transform = pdb_transform
self.num_cores = num_cores
self.af_version = af_version
super().__init__(
root,
transform=transform,
pre_transform=pre_transform,
pre_filter=pre_filter,
)
self.config.pdb_dir = Path(self.raw_dir)
self.data, self.slices = torch.load(
self.processed_paths[0], weights_only=False
)
@property
def raw_file_names(self) -> List[str]:
"""Name of the raw files in the dataset."""
return [f"{pdb}.pdb" for pdb in self.structures]
@property
def processed_file_names(self) -> List[str]:
"""Name of the processed file."""
return [f"data_{self.name}.pt"]
@property
def raw_dir(self) -> str:
if self.paths is not None:
return self.path # replace raw dir with user local path
else:
return os.path.join(self.root, "raw")
[docs]
def download(self):
"""Download the PDB files from RCSB or Alphafold."""
self.config.pdb_dir = Path(self.raw_dir)
if self.pdb_codes:
# Only download PDBs that are not already downloaded
to_download = [
pdb
for pdb in set(self.pdb_codes)
if not os.path.exists(Path(self.raw_dir) / f"{pdb}.pdb")
]
download_pdb_multiprocessing(
to_download,
self.raw_dir,
max_workers=self.num_cores,
strict=False,
)
self.bad_pdbs = self.bad_pdbs + [
pdb
for pdb in set(self.pdb_codes)
if not os.path.exists(Path(self.raw_dir) / f"{pdb}.pdb")
]
if self.uniprot_ids:
[
download_alphafold_structure(
uniprot,
out_dir=self.raw_dir,
version=self.af_version,
aligned_score=False,
)
for uniprot in tqdm(self.uniprot_ids)
]
def __len__(self) -> int:
return len(self.structures)
[docs]
def process(self):
"""Process structures into PyG format and save to disk."""
# Read data into huge `Data` list.
structure_files = [
f"{self.raw_dir}/{pdb}.pdb" for pdb in self.structures
]
# Apply transformations to raw PDB files.
if self.pdb_transform is not None:
print("Transforming PDB files...")
self.transform_pdbs()
if self.chain_selection_map:
chain_selections = [
(
self.chain_selection_map[pdb]
if pdb in self.chain_selection_map
else "all"
)
for pdb in self.structures
]
else:
chain_selections = None
# Create graph objects
print("Constructing Graphs...")
graphs = construct_graphs_mp(
path_it=structure_files,
config=self.config,
chain_selections=chain_selections,
return_dict=True,
num_cores=self.num_cores,
)
# Transform graphs
if self.graph_transformation_funcs is not None:
print("Transforming Nx Graphs...")
for func in self.graph_transformation_funcs:
graphs = {k: func(v) for k, v in graphs.items()}
# Convert to PyTorch Geometric Data
print("Converting Graphs...")
graphs = {k: self.graph_format_convertor(v) for k, v in graphs.items()}
graphs = dict(zip(self.structures, graphs.values()))
# Assign labels
if self.graph_label_map:
print("Assigning graph Labels...")
for k, v in self.graph_label_map.items():
try:
graphs[k].graph_y = v
except KeyError:
print(f"{k} not found in graphs. Skipping.")
if self.node_label_map:
print("Assigning node Labels...")
for k, v in self.node_label_map.items():
try:
graphs[k].node_y = v
except KeyError:
print(f"{k} not found in graphs. Skipping.")
data_list = list(graphs.values())
del graphs
if self.pre_filter is not None:
print("Pre-filtering Data...")
data_list = [g for g in data_list if self.pre_filter(g)]
if self.pre_transform is not None:
print("Pre-transforming data...")
data_list = [self.pre_transform(data) for data in data_list]
print("Saving Data...")
data, slices = self.collate(data_list)
torch.save((data, slices), self.processed_paths[0])
print("Done!")
[docs]
class ProteinGraphDataset(Dataset):
[docs]
def __init__(
self,
root: str,
paths: Optional[List[str]] = None,
pdb_codes: Optional[List[str]] = None,
uniprot_ids: Optional[List[str]] = None,
# graph_label_map: Optional[Dict[str, int]] = None,
graph_labels: Optional[List[torch.Tensor]] = None,
node_labels: Optional[List[torch.Tensor]] = None,
chain_selections: Optional[List[str]] = None,
# node_label_map: Optional[Dict[str, int]] = None,
# chain_selection_map: Optional[Dict[str, List[str]]] = None,
graphein_config: ProteinGraphConfig = ProteinGraphConfig(),
graph_format_convertor: GraphFormatConvertor = GraphFormatConvertor(
src_format="nx", dst_format="pyg"
),
graph_transformation_funcs: Optional[List[Callable]] = None,
pdb_transform: Optional[List[Callable]] = None,
transform: Optional[Callable] = None,
pre_transform: Optional[Callable] = None,
pre_filter: Optional[Callable] = None,
num_cores: int = 16,
af_version: int = 6,
):
"""Dataset class for protein graphs.
Dataset base class for creating graph datasets.
See `here <https://pytorch-geometric.readthedocs.io/en/latest/notes/
create_dataset.html>`__ for the accompanying tutorial.
:param root: Root directory where the dataset should be saved.
:type root: str
:param paths: List of full path of PDB or MMTF files to load. Defaults
to ``None``.
:type paths: Optional[List[str]], optional
:param pdb_codes: List of PDB codes to download and parse from the PDB.
Defaults to ``None``.
:type pdb_codes: Optional[List[str]], optional
:param uniprot_ids: List of Uniprot IDs to download and parse from
Alphafold Database. Defaults to ``None``.
:type uniprot_ids: Optional[List[str]], optional
:param graph_label_map: Dictionary mapping PDB/Uniprot IDs to
graph-level labels. Defaults to ``None``.
:type graph_label_map: Optional[Dict[str, Tensor]], optional
:param node_label_map: Dictionary mapping PDB/Uniprot IDs to node-level
labels. Defaults to ``None``.
:type node_label_map: Optional[Dict[str, torch.Tensor]], optional
:param chain_selection_map: Dictionary mapping, defaults to ``None``.
:type chain_selection_map: Optional[Dict[str, List[str]]], optional
:param graphein_config: Protein graph construction config, defaults to
``ProteinGraphConfig()``.
:type graphein_config: ProteinGraphConfig, optional
:param graph_format_convertor: Conversion handler for graphs, defaults
to ``GraphFormatConvertor(src_format="nx", dst_format="pyg")``.
:type graph_format_convertor: GraphFormatConvertor, optional
:param graph_transformation_funcs: List of functions that consume a
``nx.Graph`` and return a ``nx.Graph``. Applied to graphs after
construction but before conversion to pyg. Defaults to ``None``.
:type graph_transformation_funcs: Optional[List[Callable]], optional
:param pdb_transform: List of functions that consume a list of paths to
the downloaded structures. This provides an entry point to apply
pre-processing from bioinformatics tools of your choosing. Defaults
to ``None``.
:type pdb_transform: Optional[List[Callable]], optional
:param transform: A function/transform that takes in a
``torch_geometric.data.Data`` object and returns a transformed
version. The data object will be transformed before every access.
Defaults to ``None``.
:type transform: Optional[Callable], optional
:param pre_transform: A function/transform that takes in an
``torch_geometric.data.Data`` object and returns a transformed
version. The data object will be transformed before being saved to
disk. Defaults to ``None``.
:type pre_transform: Optional[Callable], optional
:param pre_filter: A function that takes in a
``torch_geometric.data.Data`` object and returns a boolean value,
indicating whether the data object should be included in the final
dataset. Optional, defaults to ``None``.
:type pre_filter: Optional[Callable], optional
:param num_cores: Number of cores to use for multiprocessing of graph
construction, defaults to ``16``.
:type num_cores: int, optional
:param af_version: Version of AlphaFoldDB structures to use,
defaults to ``6``.
:type af_version: int, optional
"""
self.pdb_codes = (
[pdb.lower() for pdb in pdb_codes]
if pdb_codes is not None
else None
)
self.uniprot_ids = (
[up.upper() for up in uniprot_ids]
if uniprot_ids is not None
else None
)
self.paths = paths
if self.paths is None:
if self.pdb_codes and self.uniprot_ids:
self.structures = self.pdb_codes + self.uniprot_ids
elif self.pdb_codes:
self.structures = pdb_codes
elif self.uniprot_ids:
self.structures = uniprot_ids
# Use local saved pdb_files instead of download or move them to
# self.root/raw dir
else:
if isinstance(self.paths, list):
self.structures = [
os.path.splitext(os.path.split(path)[-1])[0]
for path in self.paths
]
self.path, _ = os.path.split(self.paths[0])
# Labels & Chains
self.examples: Dict[int, str] = dict(enumerate(self.structures))
if graph_labels is not None:
self.graph_label_map = dict(enumerate(graph_labels))
else:
self.graph_label_map = None
if node_labels is not None:
self.node_label_map = dict(enumerate(node_labels))
else:
self.node_label_map = None
if chain_selections is not None:
self.chain_selection_map = dict(enumerate(chain_selections))
else:
self.chain_selection_map = None
self.validate_input()
self.bad_pdbs: List[str] = []
# Configs
self.config = graphein_config
self.graph_format_convertor = graph_format_convertor
self.num_cores = num_cores
self.pdb_transform = pdb_transform
self.graph_transformation_funcs = graph_transformation_funcs
self.af_version = af_version
super().__init__(
root,
transform=transform,
pre_transform=pre_transform,
pre_filter=pre_filter,
)
self.config.pdb_dir = Path(self.raw_dir)
@property
def raw_file_names(self) -> List[str]:
"""Names of raw files in the dataset."""
return [f"{pdb}.pdb" for pdb in self.structures]
@property
def processed_file_names(self) -> List[str]:
"""Names of processed files to look for"""
if self.chain_selection_map is not None:
return [
f"{pdb}_{chain}.pt"
for pdb, chain in zip(
self.structures, self.chain_selection_map.values()
)
]
else:
return [f"{pdb}.pt" for pdb in self.structures]
@property
def raw_dir(self) -> str:
if self.paths is not None:
return self.path # replace raw dir with user local path
else:
return os.path.join(self.root, "raw")
def validate_input(self):
if self.graph_label_map is not None:
assert len(self.structures) == len(self.graph_label_map), (
"Number of proteins and graph labels must match"
)
if self.node_label_map is not None:
assert len(self.structures) == len(self.node_label_map), (
"Number of proteins and node labels must match"
)
if self.chain_selection_map is not None:
assert len(self.structures) == len(self.chain_selection_map), (
"Number of proteins and chain selections must match"
)
assert len(
{
f"{pdb}_{chain}"
for pdb, chain in zip(
self.structures, self.chain_selection_map
)
}
) == len(self.structures), "Duplicate protein/chain combinations"
[docs]
def download(self):
"""Download the PDB files from RCSB or Alphafold."""
self.config.pdb_dir = Path(self.raw_dir)
if self.pdb_codes:
# Only download undownloaded PDBs
to_download = [
pdb
for pdb in set(self.pdb_codes)
if not os.path.exists(Path(self.raw_dir) / f"{pdb}.pdb")
]
download_pdb_multiprocessing(
to_download,
self.raw_dir,
max_workers=self.num_cores,
strict=False,
)
self.bad_pdbs = self.bad_pdbs + [
pdb
for pdb in set(self.pdb_codes)
if not os.path.exists(Path(self.raw_dir) / f"{pdb}.pdb")
]
if self.uniprot_ids:
to_download = [
uniprot
for uniprot in set(self.uniprot_ids)
if not os.path.exists(Path(self.raw_dir) / f"{uniprot}.pdb")
]
[
download_alphafold_structure(
uniprot,
out_dir=self.raw_dir,
version=self.af_version,
aligned_score=False,
)
for uniprot in tqdm(to_download)
]
[docs]
def len(self) -> int:
"""Returns length of data set (number of structures)."""
return len(self.structures)
def transform_graphein_graphs(self, graph: nx.Graph):
for func in self.graph_transformation_funcs:
graph = func(graph)
return graph
[docs]
def process(self):
"""Processes structures from files into PyTorch Geometric Data."""
# Preprocess PDB files
if self.pdb_transform:
self.transform_pdbs()
# Chunk dataset for parallel processing
chunk_size = 128
def divide_chunks(l: List[str], n: int = 2) -> Generator:
for i in range(0, len(l), n):
yield l[i : i + n]
# chunks = list(divide_chunks(self.structures, chunk_size))
chunks: List[int] = list(
divide_chunks(list(self.examples.keys()), chunk_size)
)
for chunk in tqdm(chunks):
pdbs = [self.examples[idx] for idx in chunk]
# Get chain selections
if self.chain_selection_map is not None:
chain_selections = [
self.chain_selection_map[idx] for idx in chunk
]
else:
chain_selections = ["all"] * len(chunk)
# Create graph objects
file_names = [f"{self.raw_dir}/{pdb}.pdb" for pdb in pdbs]
graphs = construct_graphs_mp(
path_it=file_names,
config=self.config,
chain_selections=chain_selections,
return_dict=False,
)
if self.graph_transformation_funcs is not None:
graphs = [self.transform_graphein_graphs(g) for g in graphs]
# Convert to PyTorch Geometric Data
graphs = [self.graph_format_convertor(g) for g in graphs]
# Assign labels
if self.graph_label_map:
labels = [self.graph_label_map[idx] for idx in chunk]
for i, _ in enumerate(chunk):
graphs[i].graph_y = labels[i]
if self.node_label_map:
labels = [self.node_label_map[idx] for idx in chunk]
for i, _ in enumerate(chunk):
graphs[i].graph_y = labels[i]
data_list = graphs
del graphs
if self.pre_filter is not None:
data_list = [g for g in data_list if self.pre_filter(g)]
if self.pre_transform is not None:
data_list = [self.pre_transform(data) for data in data_list]
for i, (pdb, chain) in enumerate(zip(pdbs, chain_selections)):
if self.chain_selection_map is None:
torch.save(
data_list[i],
os.path.join(self.processed_dir, f"{pdb}.pt"),
)
else:
torch.save(
data_list[i],
os.path.join(self.processed_dir, f"{pdb}_{chain}.pt"),
)
[docs]
def get(self, idx: int):
"""
Returns PyTorch Geometric Data object for a given index.
:param idx: Index to retrieve.
:type idx: int
:return: PyTorch Geometric Data object.
"""
if self.chain_selection_map is not None:
return torch.load(
os.path.join(
self.processed_dir,
f"{self.structures[idx]}_\
{self.chain_selection_map[idx]}.pt",
),
weights_only=False,
)
else:
return torch.load(
os.path.join(self.processed_dir, f"{self.structures[idx]}.pt"),
weights_only=False,
)
[docs]
class ProteinGraphListDataset(InMemoryDataset):
[docs]
def __init__(
self, root: str, data_list: List[Data], name: str, transform=None
):
"""Creates a dataset from a list of PyTorch Geometric Data objects.
:param root: Root directory where the dataset is stored.
:type root: str
:param data_list: List of protein graphs as PyTorch Geometric Data
objects.
:type data_list: List[Data]
:param name: Name of dataset. Data will be saved as ``data_{name}.pt``.
:type name: str
:param transform: A function/transform that takes in an
:obj:`torch_geometric.data.Data` object and returns a transformed
version. The data object will be transformed before every access.
(default: :obj:`None`)
:type transform: Optional[Callable], optional
"""
self.data_list = data_list
self.name = name
super().__init__(root, transform)
self.data, self.slices = torch.load(
self.processed_paths[0], weights_only=False
)
@property
def processed_file_names(self):
"""The name of the files in the :obj:`self.processed_dir` folder that
must be present in order to skip processing."""
return f"data_{self.name}.pt"
[docs]
def process(self):
"""Saves data files to disk."""
torch.save(self.collate(self.data_list), self.processed_paths[0])