models/datasets/pdb_data.py

import pathlib
from typing import Callable, Dict, List, Literal, Optional, Tuple, Union

import pandas as pd
import torch
from loguru import logger
from torch.utils.data import Dataset
from torch_geometric.data import Data
from tqdm import tqdm

from onescience.utils.openfold.np.residue_constants import resname_to_idx
from .base_data import BaseLightningDataModule
from models.utils.cluster_utils import (
    cluster_sequences,
    df_to_fasta,
    expand_cluster_splits,
    fasta_to_df,
    read_cluster_tsv,
    setup_clustering_file_paths,
    split_dataframe,
)
from models.utils.constants import PDB_TO_OPENFOLD_INDEX_TENSOR

from graphein.ml.datasets import PDBManager
from graphein.protein.tensor.io import protein_to_pyg
from graphein.protein.utils import download_pdb_multiprocessing


class PDBDataSelector:
    def __init__(
        self,
        data_dir: str,
        fraction: float = 1.0,
        min_length: int = None,
        max_length: int = None,
        molecule_type: str = None,
        experiment_types: List[str] = None,
        oligomeric_min: int = None,
        oligomeric_max: int = None,
        best_resolution: float = None,
        worst_resolution: float = None,
        has_ligands: List[str] = None,
        remove_ligands: List[str] = None,
        remove_non_standard_residues: bool = True,
        remove_pdb_unavailable: bool = True,
        labels: Optional[List[Literal["uniprot_id", "cath_code", "ec_number"]]] = None,
        remove_cath_unavailable: bool = False,
        exclude_ids: List[str] = None,
        exclude_ids_from_file: str = None,
        num_workers: int = 32,
    ):
        """
        Initialize the PDBDataSelector with the specified parameters.

        Args:
            data_dir (str): Directory path where the data is stored.
            fraction (float): Fraction of the data to be selected.
            min_length (int): Minimum length of the sequences to be included.
            max_length (int): Maximum length of the sequences to be included.
            molecule_type (str): Type of the molecule (e.g., "protein", "DNA", "RNA").
            experiment_types (List[str]): List of experiment types to be included.
            oligomeric_min (int): Minimum oligomeric state of the structures to be included.
            oligomeric_max (int): Maximum oligomeric state of the structures to be included.
            best_resolution (float): Best resolution threshold for the structures to be included.
            worst_resolution (float): Worst resolution threshold for the structures to be included.
            has_ligands (List[str]): List of ligands that must be present in the structures.
            remove_ligands (List[str]): List of ligands to be removed from the structures.
            remove_non_standard_residues (bool): Whether to remove non-standard residues from the structures.
            remove_pdb_unavailable (bool): Whether to remove structures that are not available in the PDB.
            labels (Optional[List[Literal["uniprot_id", "cath_code", "ec_number"]]], optional): A list of names corresponding to metadata labels that should be included in PDB manager dataframe.
                Defaults to ``None``.
            remove_cath_unavailable (bool): Whether to remove structures that don't have CATH labels.
            exclude_ids (List[str]): List of PDB IDs to be excluded from the selection.
            exclude_ids_from_file (str, optional): Path to a txt file containing IDs to be excluded from the selection.
            num_workers (int): Number of workers for parallel processing. Defaults to 32.

        Raises:
            AssertionError: If the sum of train_val_test fractions is not equal to 1.
            TypeError: If split_time_frames does not contain valid dates for np.datetime64 format.
        """
        self.database = "pdb"
        self.data_dir = pathlib.Path(data_dir)
        self.raw_dir = self.data_dir / "raw"
        self.fraction = fraction
        self.molecule_type = molecule_type
        self.experiment_types = experiment_types
        self.oligomeric_min = oligomeric_min
        self.oligomeric_max = oligomeric_max
        self.best_resolution = best_resolution
        self.worst_resolution = worst_resolution
        self.has_ligands = has_ligands
        self.remove_ligands = remove_ligands
        self.remove_non_standard_residues = remove_non_standard_residues
        self.remove_pdb_unavailable = remove_pdb_unavailable
        self.min_length = min_length
        self.max_length = max_length
        self.exclude_ids = exclude_ids
        self.exclude_ids_from_file = exclude_ids_from_file
        self.labels = labels
        self.remove_cath_unavailable = remove_cath_unavailable
        self.num_workers = num_workers
        self.df_data = None

    def create_dataset(self) -> pd.DataFrame:
        """Filter PDB data based on metadata and constraints and return a dataframe with the selected datapoints.

        Returns:
            pd.DataFrame: dataframe containing all datapoints that satisfy the criteria.
        """
        # lazy init
        if self.df_data:
            return self.df_data

        self.data_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"Initializing PDBManager in {self.data_dir}...")
        pdb_manager = PDBManager(root_dir=self.data_dir, labels=self.labels)

        num_chains = len(pdb_manager.df)
        logger.info(f"Starting with: {num_chains} chains")

        # subsample dataframe based on provided fraction
        if self.fraction != 1.0:
            logger.info(f"Subsampling data to {self.fraction} fraction")
            pdb_manager.df = pdb_manager.df.sample(frac=self.fraction)
            logger.info(f"{len(pdb_manager.df)} chains remaining")

        if self.experiment_types:
            logger.info(
                f"Removing chains that are not in one of the following experiment types: {self.experiment_types}"
            )
            pdb_manager.experiment_types(self.experiment_types, update=True)

        if self.max_length:
            logger.info(f"Removing chains longer than {self.max_length}...")
            pdb_manager.length_shorter_than(self.max_length, update=True)
            logger.info(f"{len(pdb_manager.df)} chains remaining")

        if self.min_length:
            logger.info(f"Removing chains shorter than {self.min_length}...")
            pdb_manager.length_longer_than(self.min_length, update=True)
            logger.info(f"{len(pdb_manager.df)} chains remaining")

        if self.molecule_type:
            logger.info(
                f"Removing chains molecule types not in selection: {self.molecule_type}..."
            )
            pdb_manager.molecule_type(self.molecule_type, update=True)
            logger.info(f"{len(pdb_manager.df)} chains remaining")

        logger.info(
            f"Removing chains oligomeric state not in selection: {self.oligomeric_min} - {self.oligomeric_max}..."
        )
        if self.oligomeric_min:
            pdb_manager.oligomeric(self.oligomeric_min, "greater", update=True)
        if self.oligomeric_max:
            pdb_manager.oligomeric(self.oligomeric_max, "less", update=True)
        logger.info(f"{len(pdb_manager.df)} chains remaining")

        logger.info(
            f"Removing chains with resolution not in selection: {self.best_resolution} - {self.worst_resolution}..."
        )
        if self.worst_resolution:
            pdb_manager.resolution_better_than_or_equal_to(
                self.worst_resolution, update=True
            )
        if self.best_resolution:
            pdb_manager.resolution_worse_than_or_equal_to(
                self.best_resolution, update=True
            )
        logger.info(f"{len(pdb_manager.df)} chains remaining")

        if self.remove_ligands:
            logger.info(
                f"Removing chains with ligands in selection: {self.remove_ligands}..."
            )
            pdb_manager.has_ligands(self.remove_ligands, inverse=True, update=True)
            logger.info(f"{len(pdb_manager.df)} chains remaining")

        if self.has_ligands:
            logger.info(
                f"Removing chains without ligands in selection: {self.has_ligands}..."
            )
            pdb_manager.has_ligands(self.has_ligands, update=True)
            logger.info(f"{len(pdb_manager.df)} chains remaining")

        if self.remove_non_standard_residues:
            logger.info("Removing chains with non-standard residues...")
            pdb_manager.remove_non_standard_alphabet_sequences(update=True)
            logger.info(f"{len(pdb_manager.df)} chains remaining")
        if self.remove_pdb_unavailable:
            logger.info("Removing chains with PDB unavailable...")
            pdb_manager.remove_unavailable_pdbs(update=True)
            logger.info(f"{len(pdb_manager.df)} chains remaining")
        if self.remove_cath_unavailable:
            logger.info("Removing chains with cath code unavailable...")
            mask = ~pdb_manager.df["cath_code"].isna()
            pdb_manager.df = pdb_manager.df[mask]
            logger.info(f"{len(pdb_manager.df)} chains remaining")

        all_exclude_ids = set()
        # Add IDs from direct list if present
        if self.exclude_ids:
            all_exclude_ids.update(self.exclude_ids)
            # Add IDs from file if present
        if self.exclude_ids_from_file:
            with open(self.exclude_ids_from_file, "r") as f:
                file_ids = {line.strip() for line in f if line.strip()}
            all_exclude_ids.update(file_ids)

        logger.info(f"Removing excluded chains ({len(all_exclude_ids)} gathered)")

        mask = ~pdb_manager.df["id"].isin(all_exclude_ids)
        pdb_manager.df = pdb_manager.df[mask]
        logger.info(f"{len(pdb_manager.df)} chains remaining")
        self.df_data = pdb_manager.df
        return self.df_data


class PDBDataSplitter:
    def __init__(
        self,
        df_data: pd.DataFrame = None,
        data_dir: str = None,
        train_val_test: List[float] = [0.8, 0.15, 0.05],
        split_type: Literal["random", "sequence_similarity"] = "random",
        split_sequence_similarity: Optional[int] = None,
        overwrite_sequence_clusters: Optional[bool] = False,
    ) -> None:
        """Initialise DataSplitter object for splitting data based on arguments into train, val and test set.

        Args:
            df_data (pd.DataFrame, optional): DataFrame containing the sample IDs and metadata. Defaults to None.
            data_dir (str, optional): directory contain the sample files. Defaults to None.
            train_val_test (List[float], optional): proportion of train, validation and test set. Defaults to [0.8, 0.15, 0.05].
            split_type (Literal["random", "sequence_similarity"], optional): If the dataset should be
                split randomly into train, val and test or via sequence similarity clustering.
                Defaults to "random".
            split_sequence_similarity (Optional[float], optional): if split_type == "sequence_similarity",
                which sequence similarity threshold should be chosen (0.3 means 30% sequence similarity
                clustering). Defaults to None.
            overwrite_sequence_clusters (Optional[bool], optional): if split_type == "sequence_similarity", if previously
                generated clusters (if present with the same sequence similarity threshold) should be overwritten
                or reused. Defaults to False (reuse).
        """

        self.df_data = df_data
        self.data_dir = data_dir
        self.train_val_test = train_val_test
        self.split_type = split_type
        self.split_sequence_similarity = split_sequence_similarity
        self.overwrite_sequence_clusters = overwrite_sequence_clusters
        self.splits = ["train", "val", "test"]
        self.dfs_splits = None
        self.clusterid_to_seqid_mappings = None

    def split_data(self, df_data: pd.DataFrame, file_identifier: str) -> Dict:
        """
        Splits the dataframe into train, val and test splits based on the split type and sampling mode.

        Args:
            df_data (pd.DataFrame): dataframe containing the data to be split

        Returns:
            dfs_splits (Dict): dictionary containing the train/val/test splits of the dataframe.
        """
        if self.split_type == "random":
            logger.info(
                f"Splitting dataset via random split into {self.train_val_test}..."
            )
            self.dfs_splits = split_dataframe(df_data, self.splits, self.train_val_test)
            self.clusterid_to_seqid_mappings = None

        elif self.split_type == "sequence_similarity":
            logger.info(
                f"Splitting dataset via sequence-similarity split into {self.train_val_test}..."
            )
            logger.info(
                f"Using {self.split_sequence_similarity} sequence similarity for split"
            )
            input_fasta_filepath, cluster_fasta_filepath, cluster_tsv_filepath = (
                setup_clustering_file_paths(
                    self.data_dir,
                    file_identifier,
                    self.split_sequence_similarity,
                )
            )

            if not input_fasta_filepath.exists() or self.overwrite_sequence_clusters:
                logger.info("Retrieving sequences and writing them to fasta file...")
                df_to_fasta(df=df_data, output_file=input_fasta_filepath)

            if not cluster_fasta_filepath.exists() or self.overwrite_sequence_clusters:
                logger.info("Clustering sequences via mmseqs2...")
                cluster_sequences(
                    fasta_input_filepath=input_fasta_filepath,
                    cluster_output_filepath=cluster_fasta_filepath,
                    min_seq_id=self.split_sequence_similarity,
                    overwrite=self.overwrite_sequence_clusters,
                )
            # read representative sequences in
            df_cluster_reps = fasta_to_df(cluster_fasta_filepath)
            seq_ids = df_cluster_reps["id"].to_numpy().tolist()
            # only select sequence representatives from original df to generate random splits of clusters
            df_sequences_reps = df_data.loc[df_data.id.isin(seq_ids)]
            splits = split_dataframe(
                df_sequences_reps, self.splits, self.train_val_test
            )
            # construct cluster_dict to map from cluster representative to all sequence ids in cluster
            clusterid_to_seqid_mapping = read_cluster_tsv(cluster_tsv_filepath)
            # use cluster dict to extend splits from cluster representatives to all sequence ids included in these clusters
            self.dfs_splits, self.clusterid_to_seqid_mappings = expand_cluster_splits(
                cluster_rep_splits=splits,
                clusterid_to_seqid_mapping=clusterid_to_seqid_mapping,
            )
        return (
            self.dfs_splits,
            self.clusterid_to_seqid_mappings,
        )


class PDBDataset(Dataset):
    def __init__(
        self,
        pdb_codes: List[str],
        chains: Optional[List[str]] = None,
        data_dir: Optional[str] = None,
        transform: Optional[Callable] = None,
        format: Literal["mmtf", "pdb", "cif", "ent"] = "cif",
        in_memory: bool = False,
        file_names: Optional[List[str]] = None,
        num_workers: int = 64,
    ):
        """
        Args:
            pdb_codes (List[str]): List of PDB codes or identifiers specific to your
                filenames for the structures to load.
            chains (List[str], optional): List of chains to load for each PDB code.
                Defaults to None.
            data_dir (str, optional): Path to the data directory. Defaults to None.
            transform (Callable, optional): Transform to apply to each
                example. Defaults to None.
            format (str, optional): Format to save structures in. Can be one of
                "mmtf", "pdb", "cif" or "ent". Defaults to "cif".
            in_memory (bool, optional): Whether to load data into memory.
                Defaults to False.
            file_names (List[str], optional): How to name the processed data files. By default '{pdb_code}.pt'.
            num_workers (int, optional): How many workers to use for pdb data downloads.
                Defaults to 8.
        """
        self.database = "pdb"
        self.pdb_codes = [pdb.lower() for pdb in pdb_codes]
        self.chains = chains
        self.format = format
        self.data_dir = pathlib.Path(data_dir)
        self.processed_dir = self.data_dir / "processed"
        self.in_memory = in_memory
        self.file_names = file_names
        self.num_workers = num_workers
        self.transform = transform
        self.sequence_id_to_idx = None

        if self.in_memory:
            logger.info("Reading data into memory")
            self.data = [torch.load(self.processed_dir / f) for f in tqdm(file_names)]

    def __len__(self):
        return len(self.file_names)

    def __getitem__(self, idx: int) -> Data:
        """Return PyTorch Geometric Data object for a given index.

        Args:
            idx (int): Index to retrieve.

        Returns:
            Data: PyTorch Geometric Data object.
        """
        if self.in_memory:
            graph = self.data[idx]
        else:
            if self.file_names is not None:
                fname = f"{self.file_names[idx]}.pt"
            elif self.chains is not None:
                fname = f"{self.pdb_codes[idx]}_{self.chains[idx]}.pt"
            else:
                fname = f"{self.pdb_codes[idx]}.pt"

            graph = torch.load(self.data_dir / "processed" / fname, weights_only=False)

        # reorder coords to be in OpenFold and not PDB convention
        graph.coords = graph.coords[:, PDB_TO_OPENFOLD_INDEX_TENSOR, :]
        graph.coord_mask = graph.coord_mask[:, PDB_TO_OPENFOLD_INDEX_TENSOR]

        if self.transform:
            graph = self.transform(graph)

        return graph


class PDBLightningDataModule(BaseLightningDataModule):
    def __init__(
        self,
        data_dir: Optional[str] = None,
        dataselector: Optional[PDBDataSelector] = None,
        datasplitter: Optional[PDBDataSplitter] = None,
        in_memory: bool = False,
        format: Literal["mmtf", "pdb", "cif", "ent"] = "cif",
        overwrite: bool = False,
        store_het: bool = False,
        store_bfactor: bool = True,
        # arguments for BaseLightningDataModule
        batch_padding: bool = True,
        sampling_mode: Literal["random", "cluster-random", "cluster-reps"] = "random",
        transforms: Optional[List[Callable]] = None,
        pre_transforms: Optional[List[Callable]] = None,
        pre_filters: Optional[List[Callable]] = None,
        batch_size: int = 32,
        num_workers: int = 32,
        pin_memory: bool = False,
        **kwargs,
    ):
        """Initializes the PDBLightningDataModule.

        Args:
            data_dir (str, optional): directory where PDB data should be stored.
                Default to None.
            dataselector (PDBDataSelector, optional): Selector for PDB data.
                Defaults to None.
            datasplitter (PDBDataSplitter, optional): Splitter for PDB data
                to create train/val/test splits. Defaults to None.
            in_memory (bool, optional): Whether to load the entire dataset into
                memory. Defaults to False.
            format (str, optional): Format to save structures in. Can be one of
                "mmtf", "pdb", "cif" or "ent". Defaults to "cif".
            overwrite (bool, optional): Whether to overwrite existing processed
                data. Defaults to False.
            store_het (bool, optional): Whether to store heteroatoms in the processed data.
                Defaults to False.
            store_bfactor (bool, optional): Whether to store B factors in the processed data.
                Defaults to True.
            batch_padding (bool, optional): Whether batches should be padded to a dense representation
                with the length being either a pre-specified max length or the maximum length of the
                sample in the batch (base PyTorch batch) or whether a sparse representation should be
                used (PyG batch). Defaults to True (base PyTorch batch).
            sampling_mode (Literal["random", "cluster-random", "cluster-reps"], optional): How the data should be
                sampled from the dataset later on:
                - "random": Select a random sequence and ignore clusters.
                - "cluster-random": Select a random sequence from each cluster. Keep all samples for each cluster.
                - "cluster-reps": Select the cluster representative from each cluster. Only keep the representative for each cluster.
                  Defaults to "random".
            transforms (List[Callable]): List of transforms applied to each example.
            pre_transforms (Callable): Transform applied to each example before processing.
            pre_filters (Callable): Filter applied to each example before processing.
            batch_size (int, optional): Batch size used for dataloaders. Defaults to 32.
            num_workers (int, optional): Number of workers used for dataloading. Defaults to 32.
            pin_memory (bool, optional): Whether memory should be pinned. Defaults to False.
        """
        super().__init__(
            batch_padding=batch_padding,
            sampling_mode=sampling_mode,
            transforms=transforms,
            pre_transforms=pre_transforms,
            pre_filters=pre_filters,
            batch_size=batch_size,
            num_workers=num_workers,
            pin_memory=pin_memory,
            **kwargs,
        )
        self.data_dir = pathlib.Path(data_dir)
        self.raw_dir = self.data_dir / "raw"
        self.processed_dir = self.data_dir / "processed"
        self.raw_dir.mkdir(parents=True, exist_ok=True)
        self.processed_dir.mkdir(parents=True, exist_ok=True)
        self.dataselector = dataselector
        self.datasplitter = datasplitter
        self.sampling_mode = sampling_mode
        self.format = format
        self.overwrite = overwrite
        self.in_memory = in_memory
        self.store_het = store_het
        self.store_bfactor = store_bfactor
        self.df_data = None
        self.dfs_splits = None
        self.clusterid_to_seqid_mappings = None
        self.file_names = None

    def prepare_data(self):
        if self.dataselector:
            file_identifier = self._get_file_identifier(self.dataselector)
            df_data_name = f"{file_identifier}.csv"
            if not self.overwrite and (self.data_dir / df_data_name).exists():
                logger.info(
                    f"{df_data_name} already exists, skipping data selection and processing stage."
                )
            else:
                logger.info(f"{df_data_name} does not exist yet, creating dataset now.")
                df_data = self.dataselector.create_dataset()
                logger.info(
                    f"Dataset created with {len(df_data)} entries. Now downloading structure data..."
                )
                self._download_structure_data(df_data["pdb"].tolist())
                # process pdb files into seperate chains and save processed objects as .pt files
                self._process_structure_data(
                    df_data["pdb"].tolist(), df_data["chain"].tolist()
                )

                # save df_data to disk for later use (in splitting, dataloading etc)
                logger.info(f"Saving dataset csv to {df_data_name}")
                df_data.to_csv(self.data_dir / df_data_name, index=False)

        else:  # user-provided dataset
            df_data_name = f"{self.data_dir.name}.csv"
            if not self.overwrite and (self.data_dir / df_data_name).exists():
                logger.info(
                    f"{df_data_name} already exists, skipping data selection and processing stage."
                )
            else:
                logger.info(f"{df_data_name} does not exist yet, creating dataset now.")
                df_data = self._load_pdb_folder_data(self.raw_dir)
                # process pdb files into seperate chains and save processed objects as .pt files
                self._process_structure_data(
                    pdb_codes=df_data["pdb"].tolist(),
                    chains=None,
                )
                # save df_data to disk for later use (in splitting, dataloading etc)
                logger.info(f"Saving dataset csv to {df_data_name}")
                df_data.to_csv(self.data_dir / df_data_name, index=False)
            
    def _load_pdb_folder_data(self, data_dir: pathlib.Path) -> pd.DataFrame:
        """
        Load PDB files from a folder and create a DataFrame with filenames.
        
        Args:
            data_dir (pathlib.Path): Path to the directory containing PDB files
            
        Returns:
            pd.DataFrame: DataFrame with 'pdb' column containing filenames
        """
        # Get all files with the specified format extension
        pdb_files = list(data_dir.glob(f"*.{self.format}"))
        
        # Create DataFrame with filenames
        df_data = pd.DataFrame({
            'pdb': [pdb_file.stem for pdb_file in pdb_files],
            'id': [pdb_file.stem for pdb_file in pdb_files],
        })
        
        if len(df_data) == 0:
            raise ValueError(f"No files with extension .{self.format} found in {data_dir}")
            
        logger.info(f"Found {len(df_data)} {self.format} files in {data_dir}")
        
        return df_data

    def _get_file_identifier(self, ds):
        file_identifier = (
            f"df_pdb_f{ds.fraction}_minl{ds.min_length}_maxl{ds.max_length}_mt{ds.molecule_type}"
            f"_et{''.join(ds.experiment_types) if ds.experiment_types else ''}"
            f"_mino{ds.oligomeric_min}_maxo{ds.oligomeric_max}"
            f"_minr{ds.best_resolution}_maxr{ds.worst_resolution}"
            f"_hl{''.join(ds.has_ligands) if ds.has_ligands else ''}"
            f"_rl{''.join(ds.remove_ligands) if ds.remove_ligands else ''}"
            f"_rnsr{ds.remove_non_standard_residues}_rpu{ds.remove_pdb_unavailable}"
            f"_l{''.join(ds.labels) if ds.labels else ''}"
            f"_rcu{ds.remove_cath_unavailable}"
        )
        return file_identifier

    def setup(self, stage: Optional[str] = None):
        """Split data into train, val and test sets and create dataset objects.

        Args:
            stage (Optional[str], optional): Which dataset should be created (train, val or test). Defaults to None.
        """
        # load dataframe with metadata from disk
        if not self.df_data:
            if self.dataselector:
                file_identifier = self._get_file_identifier(self.dataselector)
            else:
                file_identifier = self.data_dir.name

            df_data_name = f"{file_identifier}.csv"
            logger.info(f"Loading dataset csv from {df_data_name}")
            self.df_data = pd.read_csv(self.data_dir / df_data_name)

        # split the dataset into train, val and test and set attributes that are used for dataset creation
        (self.dfs_splits, self.clusterid_to_seqid_mappings) = (
            self.datasplitter.split_data(self.df_data, file_identifier)
        )

        # create appropriate datasets based on the selected stage
        if stage == "fit" or stage is None:
            self.train_ds = self.train_dataset()
            self.val_ds = self.val_dataset()
        elif stage == "test":
            self.test_ds = self.test_dataset()

    def _process_structure_data(self, pdb_codes, chains):
        """Process raw data sequentially instead of using multiprocessing."""
        if chains is not None:
            index_pdb_tuples = [
                (i, pdb, chains[i])
                for i, pdb in enumerate(pdb_codes)
                if not (self.processed_dir / f"{pdb}_{chains[i]}.pt").exists()
            ]
        else:
            index_pdb_tuples = [
                (i, pdb)
                for i, pdb in enumerate(pdb_codes)
                if not (self.processed_dir / f"{pdb}.pt").exists()
            ]

        file_names = []
        for tuple_ in tqdm(index_pdb_tuples, desc="Processing structures", unit="file"):
            result = self._load_and_process_pdb(tuple_)
            if result is not None:
                file_names.append(result)
        
        logger.info("Completed processing.")
        return file_names

    def _load_and_process_pdb(
        self, index_pdb_tuple: Union[Tuple[int, str], Tuple[int, str, str]]
    ) -> Optional[str]:
        """
        Load and process a PDB file, converting it to a PyTorch Geometric graph.

        This function takes a tuple containing an index and a PDB code (and optionally a chain),
        loads the corresponding PDB file, processes it into a graph, and saves the result.

        Args:
            index_pdb_tuple (Union[Tuple[int, str], Tuple[int, str, str]]): A tuple containing:
                - index (int): The index of the PDB file in the list.
                - pdb (str): The PDB code.
                - chains (str, optional): The chains to process. If not provided, all chains are processed.

        Returns:
            Optional[str]: The filename of the saved processed graph, or None if processing failed.

        Raises:
            FileNotFoundError: If the PDB file is not found in the raw directory.
        """
        try:
            if len(index_pdb_tuple) == 3:
                i, pdb, chains = index_pdb_tuple
            elif len(index_pdb_tuple) == 2:
                i, pdb = index_pdb_tuple
                chains = "all"
            else:
                raise ValueError("index_pdb_tuple must have 2 or 3 elements")

            path = self.raw_dir / f"{pdb}.{self.format}"
            if path.exists():
                path = str(path)
            elif path.with_suffix("." + self.format + ".gz").exists():
                path = str(path.with_suffix("." + self.format + ".gz"))
            else:
                raise FileNotFoundError(
                    f"{pdb} not found in raw directory. Are you sure it's downloaded and has the format {self.format}?"
                )

            fill_value_coords = 1e-5
            graph = protein_to_pyg(
                path=path,
                chain_selection=chains,
                keep_insertions=True,
                store_het=self.store_het,
                store_bfactor=self.store_bfactor,
                fill_value_coords=fill_value_coords,
            )

        except Exception as e:
            logger.warning(f"Error processing {pdb} {chains}: {e}")
            return None
        fname = f"{pdb}.pt" if chains == "all" else f"{pdb}_{chains}.pt"

        graph.id = fname.split(".")[0]
        coord_mask = graph.coords != fill_value_coords
        graph.coord_mask = coord_mask[..., 0]
        graph.residue_type = torch.tensor(
            [resname_to_idx[residue] for residue in graph.residues]
        ).long()
        graph.database = "pdb"
        graph.bfactor_avg = torch.mean(graph.bfactor, dim=-1)
        graph.residue_pdb_idx = torch.tensor(
            [int(s.split(":")[2]) for s in graph.residue_id], dtype=torch.long
        )
        graph.seq_pos = torch.arange(graph.coords.shape[0]).unsqueeze(-1)

        if self.pre_transform:
            graph = self.pre_transform(graph)

        if self.pre_filter:
            if self.pre_filter(graph) is not True:
                return None

        torch.save(graph, self.processed_dir / fname)
        return fname

    def _download_structure_data(self, pdb_codes) -> None:
        if pdb_codes is not None:
            to_download = (
                pdb_codes
                if self.overwrite
                else [
                    pdb
                    for pdb in pdb_codes
                    if not (
                        (self.raw_dir / f"{pdb}.{self.format}").exists()
                        or (self.raw_dir / f"{pdb}.{self.format}.gz").exists()
                    )
                ]
            )
            to_download = list(set(to_download))
            # Determine whether to download raw structures
            if to_download:
                logger.info(
                    f"Downloading {len(to_download)} structures to {self.processed_dir}"
                )
                file_format = (
                    self.format[:-3] if self.format.endswith(".gz") else self.format
                )
                # calculate number of downloads per worker
                chunksize = (
                    len(to_download) // self.num_workers + 1
                )  # +1 handles edge case where num_workers > len(to_download)
                download_pdb_multiprocessing(
                    to_download,
                    self.raw_dir,
                    format=file_format,
                    max_workers=self.num_workers,
                    chunksize=chunksize,
                )
            else:
                logger.info(
                    f"No structures to download, all {len(pdb_codes)} structure files already present"
                )

    def _get_dataset(self, split: Literal["train", "val", "test"]) -> PDBDataset:
        """Initialises a dataset for a given split.

        Args:
            split Literal["train", "val", "test"]: Split to initialise.

        Returns:
            PDBCompDataset: initialised dataset for one split
        """
        df_split = self.dfs_splits[split]
        self.clusterid_to_seqid_mappings = self.clusterid_to_seqid_mappings
        pdb_codes = df_split["pdb"].tolist()
        # Check if 'chain' column exists in the DataFrame
        if 'chain' in df_split.columns:
            chains = df_split["chain"].tolist()
            file_names = [f"{pdb}_{chain}" for pdb, chain in zip(pdb_codes, chains)]
        else:
            chains = None
            file_names = [f"{pdb}" for pdb in pdb_codes]

        return PDBDataset(
            pdb_codes=pdb_codes,
            chains=chains,
            data_dir=self.data_dir,
            transform=self.transform,
            format=self.format,
            in_memory=self.in_memory,
            file_names=file_names,
            num_workers=self.num_workers,
        )