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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.33
+Version: 0.0.35
 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.33 → rcsb_embedding_model-0.0.35}/pyproject.toml

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

{rcsb_embedding_model-0.0.33 → rcsb_embedding_model-0.0.35}/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, 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)
+            try:
+                atom_ch = filter_residues(atom_ch)
+                atom_ch = rename_atom_attr(atom_ch)
+                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-0.0.33 → rcsb_embedding_model-0.0.35}/src/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
 
@@ -54,24 +56,16 @@ def get_assemblies(structure, structure_format="mmcif"):
         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):
@@ -85,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,