LLMGM/pocket_features.py at main · simmzx/LLMGM · GitHub

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
"""
Protein Pocket Feature Extraction
"""

import numpy as np
from typing import Dict
from collections import Counter


AMINO_ACID_PROPERTIES = {
    'ALA': {'type': 'hydrophobic'}, 'VAL': {'type': 'hydrophobic'}, 'LEU': {'type': 'hydrophobic'},
    'ILE': {'type': 'hydrophobic'}, 'MET': {'type': 'hydrophobic'}, 'PHE': {'type': 'hydrophobic', 'aromatic': True},
    'TRP': {'type': 'hydrophobic', 'aromatic': True}, 'PRO': {'type': 'hydrophobic'},
    'SER': {'type': 'polar', 'hb_donor': True, 'hb_acceptor': True},
    'THR': {'type': 'polar', 'hb_donor': True, 'hb_acceptor': True},
    'CYS': {'type': 'polar', 'hb_donor': True},
    'TYR': {'type': 'polar', 'aromatic': True, 'hb_donor': True, 'hb_acceptor': True},
    'ASN': {'type': 'polar', 'hb_donor': True, 'hb_acceptor': True},
    'GLN': {'type': 'polar', 'hb_donor': True, 'hb_acceptor': True},
    'LYS': {'type': 'charged', 'charge': '+', 'hb_donor': True},
    'ARG': {'type': 'charged', 'charge': '+', 'hb_donor': True},
    'HIS': {'type': 'charged', 'charge': '+', 'aromatic': True, 'hb_donor': True},
    'ASP': {'type': 'charged', 'charge': '-', 'hb_acceptor': True},
    'GLU': {'type': 'charged', 'charge': '-', 'hb_acceptor': True},
    'GLY': {'type': 'special'},
}


def extract_pocket_features_from_pdb(pdb_file: str) -> Dict:
    residues = []
    coords = []
    residue_dict = {}

    with open(pdb_file, 'r') as f:
        for line in f:
            if not (line.startswith('ATOM') or line.startswith('HETATM')):
                continue
            try:
                res_name = line[17:20].strip()
                chain_id = line[21].strip() or 'A'
                res_num = line[22:26].strip()
                x, y, z = float(line[30:38]), float(line[38:46]), float(line[46:54])

                key = (chain_id, res_num, res_name)
                if key not in residue_dict:
                    residue_dict[key] = []
                residue_dict[key].append([x, y, z])
            except:
                continue

    for (chain, num, name), coord_list in residue_dict.items():
        residues.append(f"{name}{num}:{chain}")
        coords.extend(coord_list)

    coords = np.array(coords)
    if len(coords) == 0:
        return create_mock_pocket_features()

    # Volume and depth
    volume = np.prod(coords.max(axis=0) - coords.min(axis=0))
    center = coords.mean(axis=0)
    depth = np.linalg.norm(coords - center, axis=1).max()

    # Residue analysis
    unique_residues = list(set(residues))
    res_types = [res.split(':')[0][:3] for res in unique_residues]
    total = len(res_types)

    hydrophobic = sum(1 for r in res_types if AMINO_ACID_PROPERTIES.get(r, {}).get('type') == 'hydrophobic')
    polar = sum(1 for r in res_types if AMINO_ACID_PROPERTIES.get(r, {}).get('type') == 'polar')
    charged = sum(1 for r in res_types if AMINO_ACID_PROPERTIES.get(r, {}).get('type') == 'charged')

    hb_donors = [r for r in unique_residues if AMINO_ACID_PROPERTIES.get(r.split(':')[0][:3], {}).get('hb_donor')]
    hb_acceptors = [r for r in unique_residues if AMINO_ACID_PROPERTIES.get(r.split(':')[0][:3], {}).get('hb_acceptor')]
    aromatic = [r for r in unique_residues if AMINO_ACID_PROPERTIES.get(r.split(':')[0][:3], {}).get('aromatic')]

    return {
        'volume': round(volume, 1),
        'depth': round(depth, 1),
        'num_residues': len(unique_residues),
        'hydrophobic_ratio': hydrophobic / total if total > 0 else 0,
        'polar_ratio': polar / total if total > 0 else 0,
        'charged_ratio': charged / total if total > 0 else 0,
        'key_residues': unique_residues[:15],
        'hb_donors': hb_donors[:8],
        'hb_acceptors': hb_acceptors[:8],
        'aromatic': aromatic[:6],
    }


def create_mock_pocket_features() -> Dict:
    return {
        'volume': 450.0,
        'depth': 12.5,
        'num_residues': 25,
        'hydrophobic_ratio': 0.48,
        'polar_ratio': 0.32,
        'charged_ratio': 0.20,
        'key_residues': ['LEU83:A', 'VAL111:A', 'ASP189:A', 'PHE200:A', 'TYR245:A'],
        'hb_donors': ['SER195:A', 'THR224:A', 'TYR245:A', 'ARG105:A'],
        'hb_acceptors': ['ASP189:A', 'GLU166:A', 'SER195:A'],
        'aromatic': ['PHE200:A', 'TYR245:A', 'TRP215:A'],
    }