| import numpy as np
|
| import torch
|
| import torch.nn.functional as F
|
| from pad_feature import parse_fasta_file
|
| import joblib
|
| def cal_dist(x,y):
|
| return np.sqrt(np.sum((x-y)**2))
|
| def get_pdb_xyz(pdb_file):
|
| current_pos = -1000
|
| X = []
|
| current_aa = {}
|
| for line in pdb_file:
|
| if (line[0:4].strip() == "ATOM" and int(line[22:26].strip()) != current_pos) or line[0:4].strip() == "TER":
|
| if current_aa != {}:
|
| X.append(current_aa["CA"])
|
| current_aa = {}
|
| if line[0:4].strip() != "TER":
|
| current_pos = int(line[22:26].strip())
|
|
|
| if line[0:4].strip() == "ATOM":
|
| atom = line[13:16].strip()
|
| if atom in ['N', 'CA', 'C', 'O']:
|
| xyz = np.array([line[30:38].strip(), line[38:46].strip(), line[46:54].strip()]).astype(np.float32)
|
| current_aa[atom] = xyz
|
| return np.array(X)
|
|
|
| def prepare_adj(pdb_id,max_len):
|
| with open('./SC_position/' + pdb_id + '_psepos_SC.pkl', 'rb') as file:
|
| X = joblib.load(file)
|
|
|
|
|
| dis_matrix = []
|
| for radius in X:
|
| dis_radius = []
|
| for i in range(len(X)):
|
| dis_radius.append(cal_dist(radius, X[i]))
|
| dis_matrix.append(dis_radius)
|
| dis_matrix = np.row_stack(dis_matrix)
|
| adjency_matrix = cal_adj_matrix(dis_matrix,14)
|
| adj = normalize(adjency_matrix)
|
| pad_adj = pad_adjacency_matrix(adj,max_len)
|
| torch.save(pad_adj, '../feature/' + f'{pdb_id}_adj.tensor')
|
|
|
| return
|
|
|
| def cal_adj_matrix(dis_matrix,radius):
|
| dist_matrix = dis_matrix
|
| mask = ((dist_matrix >= 0)*(dist_matrix <= radius))
|
| adjency_matrix = mask.astype(np.int32)
|
| return adjency_matrix
|
|
|
| def normalize(mx):
|
| rowsum = np.array(mx.sum(1))
|
| r_inv = (rowsum ** -0.5).flatten()
|
| r_inv[np.isinf(r_inv)] = 0
|
| r_mat_inv = np.diag(r_inv)
|
| result = r_mat_inv @ mx @ r_mat_inv
|
| return result
|
|
|
|
|
| def pad_adjacency_matrix(adj, target_size):
|
| adj = torch.from_numpy(adj).float()
|
| padding_size = target_size - adj.size(0)
|
| padded_adj = F.pad(adj, (0, padding_size, 0, padding_size), "constant", 0)
|
|
|
| return padded_adj
|
|
|
