Spaces:
Runtime error
Runtime error
File size: 3,766 Bytes
753e275 |
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 |
import torch
from Bio.PDB import Selection
from Bio.PDB.Residue import Residue
from easydict import EasyDict
from .constants import (
AA, max_num_heavyatoms,
restype_to_heavyatom_names,
BBHeavyAtom
)
class ParsingException(Exception):
pass
def _get_residue_heavyatom_info(res: Residue):
pos_heavyatom = torch.zeros([max_num_heavyatoms, 3], dtype=torch.float)
mask_heavyatom = torch.zeros([max_num_heavyatoms, ], dtype=torch.bool)
restype = AA(res.get_resname())
for idx, atom_name in enumerate(restype_to_heavyatom_names[restype]):
if atom_name == '': continue
if atom_name in res:
pos_heavyatom[idx] = torch.tensor(res[atom_name].get_coord().tolist(), dtype=pos_heavyatom.dtype)
mask_heavyatom[idx] = True
return pos_heavyatom, mask_heavyatom
def parse_biopython_structure(entity, unknown_threshold=1.0, max_resseq=None):
chains = Selection.unfold_entities(entity, 'C')
chains.sort(key=lambda c: c.get_id())
data = EasyDict({
'chain_id': [],
'resseq': [], 'icode': [], 'res_nb': [],
'aa': [],
'pos_heavyatom': [], 'mask_heavyatom': [],
})
tensor_types = {
'resseq': torch.LongTensor,
'res_nb': torch.LongTensor,
'aa': torch.LongTensor,
'pos_heavyatom': torch.stack,
'mask_heavyatom': torch.stack,
}
count_aa, count_unk = 0, 0
for i, chain in enumerate(chains):
seq_this = 0 # Renumbering residues
residues = Selection.unfold_entities(chain, 'R')
residues.sort(key=lambda res: (res.get_id()[1], res.get_id()[2])) # Sort residues by resseq-icode
for _, res in enumerate(residues):
resseq_this = int(res.get_id()[1])
if max_resseq is not None and resseq_this > max_resseq:
continue
resname = res.get_resname()
if not AA.is_aa(resname): continue
if not (res.has_id('CA') and res.has_id('C') and res.has_id('N')): continue
restype = AA(resname)
count_aa += 1
if restype == AA.UNK:
count_unk += 1
continue
# Chain info
data.chain_id.append(chain.get_id())
# Residue types
data.aa.append(restype) # Will be automatically cast to torch.long
# Heavy atoms
pos_heavyatom, mask_heavyatom = _get_residue_heavyatom_info(res)
data.pos_heavyatom.append(pos_heavyatom)
data.mask_heavyatom.append(mask_heavyatom)
# Sequential number
resseq_this = int(res.get_id()[1])
icode_this = res.get_id()[2]
if seq_this == 0:
seq_this = 1
else:
d_CA_CA = torch.linalg.norm(data.pos_heavyatom[-2][BBHeavyAtom.CA] - data.pos_heavyatom[-1][BBHeavyAtom.CA], ord=2).item()
if d_CA_CA <= 4.0:
seq_this += 1
else:
d_resseq = resseq_this - data.resseq[-1]
seq_this += max(2, d_resseq)
data.resseq.append(resseq_this)
data.icode.append(icode_this)
data.res_nb.append(seq_this)
if len(data.aa) == 0:
raise ParsingException('No parsed residues.')
if (count_unk / count_aa) >= unknown_threshold:
raise ParsingException(
f'Too many unknown residues, threshold {unknown_threshold:.2f}.'
)
seq_map = {}
for i, (chain_id, resseq, icode) in enumerate(zip(data.chain_id, data.resseq, data.icode)):
seq_map[(chain_id, resseq, icode)] = i
for key, convert_fn in tensor_types.items():
data[key] = convert_fn(data[key])
return data, seq_map
|