Update code/pdb_featureVector.py

code/pdb_featureVector.py

@@ -240,6 +240,13 @@ def pdb(input_set, mode, impute):
             pdb_file_content = StringIO(pdb_data)
             structure = pdb_parser.get_structure(search, pdb_file_content)
             ppb = PPBuilder()
+            pdb_data_list = pdb_data.split('\n')
+            pdb_data_list = [i for i in pdb_data_list if i.startswith('DBREF')]
+            pdb_data_list = [[list(filter(None,i.split(' '))) for j in i.split(' ') if j == 'UNP'] for i in pdb_data_list]
+            pdb_data_list = [i for i in pdb_data_list if i != []]
+            #pdb_data_list = [[list(filter(None,j)) for j in i] for i in pdb_data_list]
+            pdb_data_list = [[j[2] for j in i] for i in pdb_data_list]
+            pdb_data_list = [i[0] for i in pdb_data_list]
             for model in structure:
                 for pp in ppb.build_peptides(model):
                     sequence = pp.get_sequence()
@@ -251,24 +258,23 @@ def pdb(input_set, mode, impute):
                     # Get the resolution from the PDB header
                     header = structure.header
                     resolution = header.get('resolution', 'N/A')
+                    if chain_id in pdb_data_list:
                     # Print UniProt IDs, chain ID, and resolution for the current model
-                    for i, chain in enumerate(model, start=1):
-                        if chain.dbxrefs == 'UNP':
-                            chain_id = chain.get_id()
-                            st.write(f"---- Information for Chain {chain_id} in Model {i} ----")
-                            st.write(f"UniProt IDs: {', '.join(uniprot_ids)}")
-                            st.write(f"Chain ID: {chain_id}")
-                            st.write(f"PDB ID: {search.upper()}")
-                            st.write(f"Resolution: {resolution}")
-                            st.write(f"Sequence: {sequence}")
-                            pdb_fasta.at[index, 'pdbID'] = search
-                            pdb_fasta.at[index, 'chain'] = chain_id
-                            pdb_fasta.at[index, 'pdbSequence'] = str(sequence)
-                            pdb_info.at[index, 'uniprotID'] = ', '.join(uniprot_ids)
-                            pdb_info.at[index, 'pdbID'] = search
-                            pdb_info.at[index, 'chain'] = chain_id
-                            pdb_info.at[index, 'resolution'] = resolution
-                            index += 1
+                        chain_id = chain.get_id()
+                        st.write(f"---- Information for Chain {chain_id} in Model {i} ----")
+                        st.write(f"UniProt IDs: {', '.join(uniprot_ids)}")
+                        st.write(f"Chain ID: {chain_id}")
+                        st.write(f"PDB ID: {search.upper()}")
+                        st.write(f"Resolution: {resolution}")
+                        st.write(f"Sequence: {sequence}")
+                        pdb_fasta.at[index, 'pdbID'] = search
+                        pdb_fasta.at[index, 'chain'] = chain_id
+                        pdb_fasta.at[index, 'pdbSequence'] = str(sequence)
+                        pdb_info.at[index, 'uniprotID'] = ', '.join(uniprot_ids)
+                        pdb_info.at[index, 'pdbID'] = search
+                        pdb_info.at[index, 'chain'] = chain_id
+                        pdb_info.at[index, 'resolution'] = resolution
+                        index += 1
 
         print()
         st.write()