rcsb-embedding-model 0.0.32__py3-none-any.whl → 0.0.34__py3-none-any.whl

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, check_all_hetero, remove_hetero
+from rcsb_embedding_model.utils.structure_parser import rename_atom_attr,filter_residues
 from rcsb_embedding_model.utils.structure_provider import StructureProvider
 
 
@@ -72,13 +72,17 @@ 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}")
-            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)
+            atom_ch = filter_residues(atom_ch)
+            atom_ch = rename_atom_attr(atom_ch)
+            try:
+                protein_chain = ProteinChain.from_atomarray(atom_ch)
+                protein_chain = ESMProtein.from_protein_chain(protein_chain)
+            except Exception as e:
+                raise IOError(f"Error while creating ESMProtein from structure chain {src_name}.{chain_id}: {e}")
+
             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/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/utils/structure_parser.py

@@ -1,4 +1,6 @@
-from biotite.structure import filter_amino_acids, filter_polymer, chain_iter, get_chains, get_residues, AtomArray
+import numpy as np
+from biotite.structure import filter_amino_acids, filter_polymer, chain_iter, get_chains, get_residues, \
+    filter_peptide_backbone, residue_iter, array
 from biotite.structure.io.pdb import PDBFile, get_structure as get_pdb_structure, get_assembly as get_pdb_assembly, list_assemblies as list_pdb_assemblies
 from biotite.structure.io.pdbx import CIFFile, get_structure, get_assembly, BinaryCIFFile, list_assemblies
 
@@ -9,17 +11,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,34 +43,29 @@ 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))
-    for idx, atom in enumerate(atom_ch):
-        atom.chain_id = ch
-        renamed_atom_ch[idx] = atom
-    return renamed_atom_ch
+def rename_atom_attr(atom_ch):
+    return array([__rename_atom(a) for a in 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 filter_residues(atom_ch):
+    return atom_ch[filter_amino_acids(atom_ch)]
 
 
-def check_all_hetero(atom_ch):
-    return sum(atom_ch.hetero) == len(atom_ch)
+def get_backbone_atoms(atom_ch):
+    return np.array([(lambda x: [a.coord for a in x])(r) for r in  residue_iter(atom_ch[filter_peptide_backbone(atom_ch)])])
 
 
 def __get_pdb_structure(pdb_file, assembly_id=None):
@@ -79,6 +79,14 @@ def __get_pdb_structure(pdb_file, assembly_id=None):
     )
 
 
+def __rename_atom(atom):
+    atom.chain_id = "A"
+    atom.hetero = False
+    if len(atom.res_name) > 3:
+        atom.res_name = 'UNK'
+    return atom
+
+
 def __get_structure(cif_file, assembly_id=None):
     return get_structure(
         cif_file,

{rcsb_embedding_model-0.0.32.dist-info → rcsb_embedding_model-0.0.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.32
+Version: 0.0.34
 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.32.dist-info → rcsb_embedding_model-0.0.34.dist-info}/RECORD

@@ -2,7 +2,7 @@ rcsb_embedding_model/__init__.py,sha256=7YfYO-V-u__19eAZfQ3t5Gf2qrhd_gwQB8rHO0J0
 rcsb_embedding_model/rcsb_structure_embedding.py,sha256=dKp9hXQO0JAnO4SEfjJ_mG_jHu3UxAPguv6jkOjp-BI,4487
 rcsb_embedding_model/cli/args_utils.py,sha256=7nP2q8pL5dWK_U7opxtWmoFcYVwasky6elHk-dASFaI,165
 rcsb_embedding_model/cli/inference.py,sha256=67_Tr3LWeA3T4KS5mkjq6tw77Ypy0R8IwMxEG2FwVqQ,19901
-rcsb_embedding_model/dataset/esm_prot_from_chain.py,sha256=HuKII1zOnO0JF6wZXyqBykiGunSFXmhzBimsSpbEC1I,4291
+rcsb_embedding_model/dataset/esm_prot_from_chain.py,sha256=p0IqLjGj1v-0sKGbrpWfkyYZXSurF3OOVkVk0g3zxyE,4410
 rcsb_embedding_model/dataset/esm_prot_from_structure.py,sha256=3HzXCCc-UqmZNbJaeXHyUsSIZZxMc2erbxAPGIxSmfE,2621
 rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py,sha256=69h1VkrIXesHZi1cG3BOMMytSDeRzcBBP0_Z3Xz3dM8,2869
 rcsb_embedding_model/dataset/residue_assembly_embedding_from_tensor_file.py,sha256=Hd9oH-IVgY6d7Dxy5VfiwHvSaK-Wwhk6ccUBgOwl0TU,3740
@@ -18,13 +18,13 @@ rcsb_embedding_model/modules/chain_module.py,sha256=KsZw2uagO4rpAKWv6ivqEMxIEzgt
 rcsb_embedding_model/modules/esm_module.py,sha256=otJRbCb319nCCob_4E1W_UClhkex9eDqcCyzWQO-vIs,740
 rcsb_embedding_model/modules/structure_module.py,sha256=4js02XzKvhc_G26ELsGhJ9SCi_wlvtVolObxfWt3BhE,1077
 rcsb_embedding_model/types/api_types.py,sha256=SCwALwvEb0KRKaoWKbuN7JyfOH-1whsI0Z4ki41dht8,1235
-rcsb_embedding_model/utils/data.py,sha256=BOjYdIRHrFqk8qFuKGrgCtVyfDupzgOVmH_0C-ecMvg,3813
+rcsb_embedding_model/utils/data.py,sha256=ThrcYycIizsV_Ycn6PPxF12JRr1m2K-v8TsIaVqx10A,3816
 rcsb_embedding_model/utils/model.py,sha256=xr3p02ohOgJ5UInwdIupN68Oq4yvNFhxobZRacS1adg,953
-rcsb_embedding_model/utils/structure_parser.py,sha256=z_Aid5QgGUAJkLL-fMzG2Uh2wPrHXmBdjhLn_YkPl-k,3035
+rcsb_embedding_model/utils/structure_parser.py,sha256=fSIbq_a_aEigCWY_1dUcW9d9Law0ZDOcZAxJlZL0Rt8,3377
 rcsb_embedding_model/utils/structure_provider.py,sha256=eWtxjkPpmRfmil_DKR1J6miaXR3lQ28DF5O0qrqSgGA,786
 rcsb_embedding_model/writer/batch_writer.py,sha256=rTFNasB0Xp4-XCNTXKeEWZxSrb7lvZytoRldJUWn9Jg,3312
-rcsb_embedding_model-0.0.32.dist-info/METADATA,sha256=JWSb2a_wZ2Y8YWzyM3vk_QG--ukrIGsI5Q1-R3bApjs,5351
-rcsb_embedding_model-0.0.32.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-rcsb_embedding_model-0.0.32.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
-rcsb_embedding_model-0.0.32.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
-rcsb_embedding_model-0.0.32.dist-info/RECORD,,
+rcsb_embedding_model-0.0.34.dist-info/METADATA,sha256=zL59HhwjvhQAOrh_hJAMQ_997sXCnVzboStFC-j5qog,5351
+rcsb_embedding_model-0.0.34.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rcsb_embedding_model-0.0.34.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
+rcsb_embedding_model-0.0.34.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
+rcsb_embedding_model-0.0.34.dist-info/RECORD,,