rcsb-embedding-model 0.0.31__tar.gz → 0.0.33__tar.gz

{rcsb_embedding_model-0.0.31 → rcsb_embedding_model-0.0.33}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.31
+Version: 0.0.33
 Summary: Protein Embedding Model for Structure Search
 Project-URL: Homepage, https://github.com/rcsb/rcsb-embedding-model
 Project-URL: Issues, https://github.com/rcsb/rcsb-embedding-model/issues

{rcsb_embedding_model-0.0.31 → rcsb_embedding_model-0.0.33}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.31"
+version = "0.0.33"
 authors = [
     { name="Joan Segura", email="joan.segura@rcsb.org" },
 ]

{rcsb_embedding_model-0.0.31 → rcsb_embedding_model-0.0.33}/src/rcsb_embedding_model/dataset/esm_prot_from_chain.py

@@ -11,7 +11,7 @@ import pandas as pd
 
 from rcsb_embedding_model.types.api_types import StructureFormat, StructureLocation, SrcLocation
 from rcsb_embedding_model.utils.data import stringio_from_url
-from rcsb_embedding_model.utils.structure_parser import rename_atom_ch
+from rcsb_embedding_model.utils.structure_parser import rename_atom_ch, check_all_hetero, remove_hetero
 from rcsb_embedding_model.utils.structure_provider import StructureProvider
 
 
@@ -72,8 +72,14 @@ class EsmProtFromChain(Dataset):
         for atom_ch in chain_iter(structure):
             if len(atom_ch) == 0:
                 raise IOError(f"No atoms were found in structure chain {src_name}.{chain_id}")
-            protein_chain = ProteinChain.from_atomarray(rename_atom_ch(atom_ch))
-            return ESMProtein.from_protein_chain(protein_chain), item_name
+            if check_all_hetero(atom_ch):
+                atom_ch = remove_hetero(atom_ch)
+            atom_ch = rename_atom_ch(atom_ch)
+            protein_chain = ProteinChain.from_atomarray(atom_ch)
+            protein_chain = ESMProtein.from_protein_chain(protein_chain)
+            if len(protein_chain) == 0:
+                raise IOError(f"No atoms were found in structure chain {src_name}.{chain_id}")
+            return protein_chain, item_name
         raise IOError(f"No atoms were found in structure chain {src_name}.{chain_id}")
 
 

{rcsb_embedding_model-0.0.31 → rcsb_embedding_model-0.0.33}/src/rcsb_embedding_model/utils/data.py

@@ -4,6 +4,7 @@ import gzip
 from io import StringIO, BytesIO
 
 import torch
+from requests import RequestException
 
 
 def collate_seq_embeddings(batch_list):
@@ -52,13 +53,10 @@ def stringio_from_url(url):
                 return StringIO(f.read())
         else:
             return StringIO(response.text)
-    except requests.exceptions.RequestException as e:
-        print(f"Error fetching URL: {e}")
-        return None
+    except RequestException as e:
+        raise RuntimeError(f"Error fetching URL: {e}")
     except (OSError, gzip.BadGzipFile) as e:
-        print(f"Error decompressing gzip file: {e}")
-        return None
-
+        raise RuntimeError(f"Error decompressing gzip file: {e}")
 
 
 def concatenate_tensors(file_list, max_residues, dim=0):

{rcsb_embedding_model-0.0.31 → rcsb_embedding_model-0.0.33}/src/rcsb_embedding_model/utils/structure_parser.py

@@ -9,17 +9,20 @@ def get_structure_from_src(
         chain_id=None,
         assembly_id=None
 ):
-    if structure_format == "pdb":
-        pdb_file = PDBFile.read(src_structure)
-        structure = __get_pdb_structure(pdb_file, assembly_id)
-    elif structure_format == "mmcif":
-        cif_file = CIFFile.read(src_structure)
-        structure = __get_structure(cif_file, assembly_id)
-    elif structure_format == "binarycif":
-        cif_file = BinaryCIFFile.read(src_structure)
-        structure = __get_structure(cif_file, assembly_id)
-    else:
-        raise RuntimeError(f"Unknown file format {structure_format}")
+    try:
+        if structure_format == "pdb":
+            pdb_file = PDBFile.read(src_structure)
+            structure = __get_pdb_structure(pdb_file, assembly_id)
+        elif structure_format == "mmcif":
+            cif_file = CIFFile.read(src_structure)
+            structure = __get_structure(cif_file, assembly_id)
+        elif structure_format == "binarycif":
+            cif_file = BinaryCIFFile.read(src_structure)
+            structure = __get_structure(cif_file, assembly_id)
+        else:
+            raise RuntimeError(f"Unknown file format {structure_format}")
+    except Exception as e:
+        raise RuntimeError(f"Error reading structure from {src_structure}: {e}")
 
     if chain_id is not None:
         return structure[structure.chain_id == chain_id]
@@ -38,26 +41,39 @@ def get_protein_chains(structure, min_res_n=0):
 
 
 def get_assemblies(structure, structure_format="mmcif"):
-    if structure_format == "pdb":
-        return tuple(list_pdb_assemblies(PDBFile.read(structure)))
-    elif structure_format == "mmcif":
-        return tuple(list_assemblies(CIFFile.read(structure)).keys())
-    elif structure_format == "binarycif":
-        return tuple(list_assemblies(BinaryCIFFile.read(structure)))
-    else:
-        raise RuntimeError(f"Unknown file format {structure_format}")
+    try:
+        if structure_format == "pdb":
+            return tuple(list_pdb_assemblies(PDBFile.read(structure)))
+        elif structure_format == "mmcif":
+            return tuple(list_assemblies(CIFFile.read(structure)).keys())
+        elif structure_format == "binarycif":
+            return tuple(list_assemblies(BinaryCIFFile.read(structure)))
+        else:
+            raise RuntimeError(f"Unknown file format {structure_format}")
+    except Exception as e:
+        raise RuntimeError(f"Error reading assemblies from {structure}: {e}")
 
 
 def rename_atom_ch(atom_ch, ch="A"):
     renamed_atom_ch = AtomArray(len(atom_ch))
-    n = 0
-    for atom in atom_ch:
+    for idx, atom in enumerate(atom_ch):
         atom.chain_id = ch
-        renamed_atom_ch[n] = atom
-        n += 1
+        renamed_atom_ch[idx] = atom
     return renamed_atom_ch
 
 
+def remove_hetero(atom_ch):
+    renamed_atom_ch = AtomArray(len(atom_ch))
+    for idx, atom in enumerate(atom_ch):
+        atom.hetero = False
+        renamed_atom_ch[idx] = atom
+    return renamed_atom_ch
+
+
+def check_all_hetero(atom_ch):
+    return sum(atom_ch.hetero) == len(atom_ch)
+
+
 def __get_pdb_structure(pdb_file, assembly_id=None):
     return get_pdb_structure(
         pdb_file,