Fix ml dataset

graphein/ml/datasets/torch_geometric_dataset.py

@@ -9,6 +9,7 @@
 import os
 from pathlib import Path
 from typing import Callable, Dict, Generator, List, Optional
+from urllib.error import HTTPError
 
 import networkx as nx
 from loguru import logger as log
@@ -167,6 +168,9 @@ def __init__(
         self.bad_pdbs: List[
             str
         ] = []  # list of pdb codes that failed to download
+        self.bad_uniprot_ids: List[
+            str
+        ] = []  # list of uniprot ids that failed to download
 
         # Labels & Chains
         self.graph_label_map = graph_label_map
@@ -282,6 +286,9 @@ def process(self):
             return_dict=True,
             num_cores=self.num_cores,
         )
+        # Keep only graphs that were successfully constructed 
+        graphs = [g for g in graphs if g is not None]
+
         # Transform graphs
         if self.graph_transformation_funcs is not None:
             print("Transforming Nx Graphs...")
@@ -412,6 +419,8 @@ def __init__(
             defaults to ``2``.
         :type af_version: int, optional
         """
+
+
         self.pdb_codes = (
             [pdb.lower() for pdb in pdb_codes]
             if pdb_codes is not None
@@ -460,7 +469,7 @@ def __init__(
             self.chain_selection_map = None
         self.validate_input()
         self.bad_pdbs: List[str] = []
-
+        self.bad_uniprot_ids: List[str] = []
         # Configs
         self.config = graphein_config
         self.graph_format_convertor = graph_format_convertor
@@ -533,6 +542,7 @@ def download(self):
                 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,
@@ -545,15 +555,34 @@ def download(self):
                 if not os.path.exists(Path(self.raw_dir) / f"{pdb}.pdb")
             ]
         if self.uniprot_ids:
-            [
-                download_alphafold_structure(
+
+            # Only download undownloaded Uniprot IDs
+            to_download = [
+                uniprot
+                for uniprot in set(self.uniprot_ids)
+                if not os.path.exists(Path(self.raw_dir) / f"{uniprot}.pdb")
+            ]
+
+            for uniprot in tqdm(to_download):
+
+
+                fn = download_alphafold_structure(
                     uniprot,
                     out_dir=self.raw_dir,
                     version=self.af_version,
                     aligned_score=False,
+                    rename=True,
                 )
-                for uniprot in tqdm(self.uniprot_ids)
+
+            self.bad_uniprot_ids = self.bad_uniprot_ids + [
+                uniprot
+                for uniprot in set(self.uniprot_ids)
+                if not os.path.exists(Path(self.raw_dir) / f"{uniprot}.pdb")
             ]
+
+
+
+        # TODO: remove bad uniprot / pdb ids from self.structures
 
     def len(self) -> int:
         """Returns length of data set (number of structures)."""
@@ -580,7 +609,7 @@ def process(self):
         if self.pdb_transform:
             self.transform_pdbs()
 
-        idx = 0
+        idx = 0 
         # Chunk dataset for parallel processing
         chunk_size = 128
 
@@ -612,6 +641,12 @@ def divide_chunks(l: List[str], n: int = 2) -> Generator:
                 chain_selections=chain_selections,
                 return_dict=False,
             )
+            graphs = [
+                g 
+                for g in graphs 
+                if g is not None
+            ]
+
             if self.graph_transformation_funcs is not None:
                 graphs = [self.transform_graphein_graphs(g) for g in graphs]
 

graphein/protein/edges/distance.py

@@ -1263,7 +1263,7 @@ def node_coords(G: nx.Graph, n: str) -> Tuple[float, float, float]:
     :return: Tuple of coordinates ``(x, y, z)``
     :rtype: Tuple[float, float, float]
     """
-    x, y, z = tuple(G.nodes[n]["coords"])
+    (x, y, z) = tuple(G.nodes[n]["coords"])
     return x, y, z