File size: 7,829 Bytes
ad9fbbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | import numpy as np
import torch
import torch.nn.functional as F
import joblib
import pandas as pd
import os
import subprocess
from pad_feature import parse_fasta_file
def def_atom_features():
A = {'N':[0,1,0], 'CA':[0,1,0], 'C':[0,0,0], 'O':[0,0,0], 'CB':[0,3,0]}
V = {'N':[0,1,0], 'CA':[0,1,0], 'C':[0,0,0], 'O':[0,0,0], 'CB':[0,1,0], 'CG1':[0,3,0], 'CG2':[0,3,0]}
F = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0],'CB':[0,2,0],
'CG':[0,0,1], 'CD1':[0,1,1], 'CD2':[0,1,1], 'CE1':[0,1,1], 'CE2':[0,1,1], 'CZ':[0,1,1] }
P = {'N': [0, 0, 1], 'CA': [0, 1, 1], 'C': [0, 0, 0], 'O': [0, 0, 0],'CB':[0,2,1], 'CG':[0,2,1], 'CD':[0,2,1]}
L = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'CG':[0,1,0], 'CD1':[0,3,0], 'CD2':[0,3,0]}
I = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,1,0], 'CG1':[0,2,0], 'CG2':[0,3,0], 'CD1':[0,3,0]}
R = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0],
'CG':[0,2,0], 'CD':[0,2,0], 'NE':[0,1,0], 'CZ':[1,0,0], 'NH1':[0,2,0], 'NH2':[0,2,0] }
D = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'CG':[-1,0,0], 'OD1':[-1,0,0], 'OD2':[-1,0,0]}
E = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'CG':[0,2,0], 'CD':[-1,0,0], 'OE1':[-1,0,0], 'OE2':[-1,0,0]}
S = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'OG':[0,1,0]}
T = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,1,0], 'OG1':[0,1,0], 'CG2':[0,3,0]}
C = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'SG':[-1,1,0]}
N = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'CG':[0,0,0], 'OD1':[0,0,0], 'ND2':[0,2,0]}
Q = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'CG':[0,2,0], 'CD':[0,0,0], 'OE1':[0,0,0], 'NE2':[0,2,0]}
H = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0],
'CG':[0,0,1], 'ND1':[-1,1,1], 'CD2':[0,1,1], 'CE1':[0,1,1], 'NE2':[-1,1,1]}
K = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'CG':[0,2,0], 'CD':[0,2,0], 'CE':[0,2,0], 'NZ':[0,3,1]}
Y = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0],
'CG':[0,0,1], 'CD1':[0,1,1], 'CD2':[0,1,1], 'CE1':[0,1,1], 'CE2':[0,1,1], 'CZ':[0,0,1], 'OH':[-1,1,0]}
M = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0], 'CG':[0,2,0], 'SD':[0,0,0], 'CE':[0,3,0]}
W = {'N': [0, 1, 0], 'CA': [0, 1, 0], 'C': [0, 0, 0], 'O': [0, 0, 0], 'CB':[0,2,0],
'CG':[0,0,1], 'CD1':[0,1,1], 'CD2':[0,0,1], 'NE1':[0,1,1], 'CE2':[0,0,1], 'CE3':[0,1,1], 'CZ2':[0,1,1], 'CZ3':[0,1,1], 'CH2':[0,1,1]}
G = {'N': [0, 1, 0], 'CA': [0, 2, 0], 'C': [0, 0, 0], 'O': [0, 0, 0]}
atom_features = {'A': A, 'V': V, 'F': F, 'P': P, 'L': L, 'I': I, 'R': R, 'D': D, 'E': E, 'S': S,
'T': T, 'C': C, 'N': N, 'Q': Q, 'H': H, 'K': K, 'Y': Y, 'M': M, 'W': W, 'G': G}
for atom_fea in atom_features.values():
for i in atom_fea.keys():
i_fea = atom_fea[i]
atom_fea[i] = [i_fea[0]/2+0.5,i_fea[1]/3,i_fea[2]]
return atom_features
def get_pdb_DF(file_path):
atom_fea_dict = def_atom_features()
res_dict = {'GLY': 'G', 'ALA': 'A', 'VAL': 'V', 'ILE': 'I', 'LEU': 'L', 'PHE': 'F', 'PRO': 'P', 'MET': 'M',
'TRP': 'W', 'CYS': 'C', 'SER': 'S', 'THR': 'T', 'ASN': 'N', 'GLN': 'Q', 'TYR': 'Y', 'HIS': 'H',
'ASP': 'D', 'GLU': 'E', 'LYS': 'K', 'ARG': 'R'}
Relative_atomic_mass = {'H': 1, 'C': 12, 'O': 16, 'N': 14, 'S': 32, 'FE': 56, 'P': 31, 'BR': 80, 'F': 19,
'CO': 59, 'V': 51, 'I': 127, 'CL': 35.5, 'CA': 40, 'B': 10.8, 'ZN': 65.5, 'MG': 24.3,
'NA': 23, 'HG': 200.6, 'MN': 55, 'K': 39.1, 'AP': 31, 'AC': 227, 'AL': 27, 'W': 183.9,
'SE': 79, 'NI': 58.7}
atom_count = -1
res_count = -1
res_id_list = []
before_res_pdb_id = None
atoms = [] # List to store atom data
with open(file_path, 'r') as pdb_file:
for line in pdb_file:
if line.startswith('ATOM'):
atom_type = line[76:78].strip()
if atom_type not in Relative_atomic_mass:
continue
atom_count += 1
res_pdb_id = int(line[22:26].strip())
if res_pdb_id != before_res_pdb_id:
res_count += 1
before_res_pdb_id = res_pdb_id
atom_name = line[12:16].strip()
is_sidechain = 0 if atom_name in ['N', 'CA', 'C', 'O', 'H'] else 1
res = res_dict.get(line[17:20].strip(), 'X') # Default to 'X' if unknown
xyz = [float(line[30:38]), float(line[38:46]), float(line[46:54])]
try:
atom_fea = atom_fea_dict[res][atom_name]
except KeyError:
atom_fea = [0.5, 0.5, 0.5]
try:
bfactor = float(line[60:66])
except ValueError:
bfactor = 0.5
atoms.append({
'ID': atom_count,
'atom': atom_name,
'atom_type': atom_type,
'res': res,
'res_id': res_pdb_id,
'xyz': xyz,
'B_factor': bfactor,
'mass': Relative_atomic_mass[atom_type],
'is_sidechain': is_sidechain,
'charge': atom_fea[0],
'num_H': atom_fea[1],
'ring': atom_fea[2]
})
if len(res_id_list) == 0 or res_id_list[-1] != res_pdb_id:
res_id_list.append(res_pdb_id)
elif line.startswith('TER'):
break
return atoms, res_id_list
def PDBFeature(query_id, PDB_chain_dir, results_dir):
pdb_path = f"{PDB_chain_dir}/{query_id}.pdb"
atoms, res_id_list = get_pdb_DF(pdb_path)
# with open(f"{results_dir}/{query_id}.df", 'wb') as f:
# joblib.dump({'atoms': atoms, 'res_id_list': res_id_list}, f)
res_sidechain_centroid = []
res_types = []
for res_id in res_id_list:
res_atoms = [atom for atom in atoms if atom['res_id'] == res_id]
if not res_atoms:
continue
res_type = res_atoms[0]['res']
res_types.append(res_type)
xyz = np.array([atom['xyz'] for atom in res_atoms])
masses = np.array([atom['mass'] for atom in res_atoms]).reshape(-1, 1)
centroid = np.sum(masses * xyz, axis=0) / np.sum(masses)
sidechain_atoms = [atom for atom in res_atoms if atom['is_sidechain'] == 1]
if not sidechain_atoms:
res_sidechain_centroid.append(centroid)
else:
xyz = np.array([atom['xyz'] for atom in sidechain_atoms])
masses = np.array([atom['mass'] for atom in sidechain_atoms]).reshape(-1, 1)
sidechain_centroid = np.sum(masses * xyz, axis=0) / np.sum(masses)
res_sidechain_centroid.append(sidechain_centroid)
res_sidechain_centroid = np.array(res_sidechain_centroid)
with open(f"{results_dir}/{query_id}_psepos_SC.pkl", 'wb') as f:
joblib.dump(res_sidechain_centroid, f)
sequence = ''.join(res_types)
with open(f"{results_dir}/{query_id}.seq", 'w') as f:
f.write(f">{query_id}\n")
f.write(sequence)
|