File size: 5,638 Bytes
80a72c3 | 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 155 156 157 158 159 160 | import logging
import os
from typing import List
import Bio.PDB
import numpy as np
import torch
from scipy.spatial import distance_matrix
from model import constants
from model.constants import _3to1
from model.utils.cif2pdb import cif2pdb
from model.utils.secondary_structure import renum_pdb_file, \
calculate_ss, make_ss_matrix
LOG = logging.getLogger(__name__)
def get_model_structure(structure_path) -> Bio.PDB.Structure:
"""
Returns the Bio.PDB.Structure object for a given PDB or MMCIF file
"""
structure_id = os.path.split(structure_path)[-1].split('.')[0]
if structure_path.endswith('.pdb'):
structure = Bio.PDB.PDBParser().get_structure(structure_id, structure_path)
elif structure_path.endswith('.cif'):
structure = Bio.PDB.MMCIFParser().get_structure(structure_id, structure_path)
else:
raise ValueError(f'Unrecognized file extension: {structure_path}')
model = structure[0]
return model
class Residue:
def __init__(self, index: int, res_label: str, aa: str):
self.index = int(index)
self.res_label = str(res_label)
self.aa = str(aa)
def get_model_structure_residues(structure_model: Bio.PDB.Structure, chain='A') -> List[Residue]:
"""
Returns a list of residues from a given PDB or MMCIF structure
"""
residues = []
res_index = 1
for biores in structure_model[chain].child_list:
res_num = biores.id[1]
res_ins = biores.id[2]
res_label = str(res_num)
if res_ins != ' ':
res_label += str(res_ins)
aa3 = biores.get_resname()
if aa3 not in _3to1:
continue
aa = _3to1[aa3]
res = Residue(res_index, res_label, aa)
residues.append(res)
# increment the residue index after we have filtered out non-standard amino acids
res_index += 1
return residues
def inference_time_create_features(pdb_path, feature_config, chain="A", *,
model_structure: Bio.PDB.Structure=None,
renumber_pdbs=True, stride_path=constants.STRIDE_EXE,
):
if pdb_path.endswith(".cif"):
pdb_path = cif2pdb(pdb_path)
if not model_structure:
model_structure = get_model_structure(pdb_path)
dist_matrix = get_distance(model_structure, chain=chain)
if renumber_pdbs:
output_pdb_path = pdb_path.replace('.pdb', '_renum.pdb').replace('.cif', '_renum.cif')
renum_pdb_file(pdb_path, output_pdb_path)
else:
output_pdb_path = pdb_path
ss_filepath = pdb_path + '_ss.txt'
calculate_ss(output_pdb_path, chain, stride_path, ssfile=ss_filepath)
helix, strand = make_ss_matrix(ss_filepath, nres=dist_matrix.shape[-1])
if feature_config['ss_bounds']:
end_res_val = -1 if feature_config['negative_ss_end'] else 1
helix_boundaries = make_boundary_matrix(helix, end_res_val=end_res_val)
strand_boundaries = make_boundary_matrix(strand, end_res_val=end_res_val)
if renumber_pdbs:
os.remove(output_pdb_path)
os.remove(ss_filepath)
LOG.info(f"Distance matrix shape: {dist_matrix.shape}, SS matrix shape: {helix.shape}")
if feature_config['ss_bounds']:
if feature_config['same_channel_boundaries_and_ss']:
helix_boundaries[helix == 1] = 1
strand_boundaries[strand == 1] = 1
stacked_features = np.stack((dist_matrix, helix_boundaries, strand_boundaries), axis=0)
else:
stacked_features = np.stack((dist_matrix, helix, strand, helix_boundaries, strand_boundaries), axis=0)
else:
stacked_features = np.stack((dist_matrix, helix, strand), axis=0)
stacked_features = stacked_features[None] # add batch dimension
return torch.Tensor(stacked_features)
def distance_matrix(x):
"""Compute the distance matrix.
Returns the matrix of all pair-wise distances.
Parameters
----------
x : (M, K) array_like
Matrix of M vectors in K dimensions.
p : float, 1 <= p <= infinity
Which Minkowski p-norm to use.
Returns
-------
result : (M, M) ndarray
Matrix containing the distance from every vector in `x` to every vector
in `x`.
"""
x1 = x[:, np.newaxis, :] # Expand x to 3D for broadcasting, shape (M, 1, K)
x2 = x[np.newaxis, :, :] # Expand x to 3D for broadcasting, shape (1, M, K)
distance_matrix = np.sum(np.abs(x2 - x1) ** 2, axis=-1) ** 0.5
return distance_matrix.astype(np.float16)
def get_distance(structure_model: Bio.PDB.Structure, chain='A'):
alpha_coords = np.array([residue['CA'].get_coord() for residue in \
structure_model[chain].get_residues() if Bio.PDB.is_aa(residue) and \
'CA' in residue and residue.get_resname() in _3to1], dtype=np.float16)
x = distance_matrix(alpha_coords)
return x
def make_boundary_matrix(ss, end_res_val=1):
"""
makes a matrix where the boundary residues
of the sec struct component are 1
"""
ss_lines = np.zeros_like(ss)
diag = np.diag(ss)
if max(diag) == 0:
return ss_lines
padded_diag = np.zeros(len(diag) + 2)
padded_diag[1:-1] = diag
diff_before = diag - padded_diag[:-2]
diff_after = diag - padded_diag[2:]
start_res = np.where(diff_before == 1)[0]
end_res = np.where(diff_after == 1)[0]
ss_lines[start_res, :] = 1
ss_lines[:, start_res] = 1
ss_lines[end_res, :] = end_res_val
ss_lines[:, end_res] = end_res_val
return ss_lines
|