| |
| import numpy as np |
| from collections import defaultdict |
| import prody as pr |
| import os |
|
|
| from datasets.constants import chi, atom_order, aa_long2short, aa_short2aa_idx, aa_idx2aa_short |
|
|
|
|
| def get_dihedral_indices(resname, chi_num): |
| """Return the atom indices for the specified dihedral angle. |
| """ |
| if resname not in chi: |
| return np.array([np.nan]*4) |
| if chi_num not in chi[resname]: |
| return np.array([np.nan]*4) |
| return np.array([atom_order[resname].index(x) for x in chi[resname][chi_num]]) |
|
|
|
|
| dihedral_indices = defaultdict(list) |
| for aa in atom_order.keys(): |
| for i in range(1, 5): |
| inds = get_dihedral_indices(aa, i) |
| dihedral_indices[aa].append(inds) |
| dihedral_indices[aa] = np.array(dihedral_indices[aa]) |
|
|
|
|
| def vector_batch(a, b): |
| return a - b |
|
|
|
|
| def unit_vector_batch(v): |
| return v / np.linalg.norm(v, axis=1, keepdims=True) |
|
|
|
|
| def dihedral_angle_batch(p): |
| b0 = vector_batch(p[:, 0], p[:, 1]) |
| b1 = vector_batch(p[:, 1], p[:, 2]) |
| b2 = vector_batch(p[:, 2], p[:, 3]) |
| |
| n1 = np.cross(b0, b1) |
| n2 = np.cross(b1, b2) |
| |
| m1 = np.cross(n1, b1 / np.linalg.norm(b1, axis=1, keepdims=True)) |
| |
| x = np.sum(n1 * n2, axis=1) |
| y = np.sum(m1 * n2, axis=1) |
| |
| deg = np.degrees(np.arctan2(y, x)) |
|
|
| deg[deg < 0] += 360 |
|
|
| return deg |
|
|
|
|
| def batch_compute_dihedral_angles(sidechains): |
| sidechains_np = np.array(sidechains) |
| dihedral_angles = dihedral_angle_batch(sidechains_np) |
| return dihedral_angles |
|
|
|
|
| def get_coords(prody_pdb): |
| resindices = sorted(set(prody_pdb.ca.getResindices())) |
| coords = np.full((len(resindices), 14, 3), np.nan) |
| for i, resind in enumerate(resindices): |
| sel = prody_pdb.select(f'resindex {resind}') |
| resname = sel.getResnames()[0] |
| for j, name in enumerate(atom_order[aa_long2short[resname] if resname in aa_long2short else 'X']): |
| sel_resnum_name = sel.select(f'name {name}') |
| if sel_resnum_name is not None: |
| coords[i, j, :] = sel_resnum_name.getCoords()[0] |
| else: |
| coords[i, j, :] = [np.nan, np.nan, np.nan] |
| return coords |
|
|
|
|
| def get_onehot_sequence(seq): |
| onehot = np.zeros((len(seq), 20)) |
| for i, aa in enumerate(seq): |
| idx = aa_short2aa_idx[aa] if aa in aa_short2aa_idx else 7 |
| onehot[i, idx] = 1 |
| return onehot |
|
|
|
|
| def get_dihedral_indices(onehot_sequence): |
| return np.array([dihedral_indices[aa_idx2aa_short[aa_idx]] for aa_idx in np.where(onehot_sequence)[1]]) |
|
|
|
|
| def _get_chi_angles(coords, indices): |
| X = coords |
| Y = indices.astype(int) |
| N = coords.shape[0] |
| mask = np.isnan(indices) |
| Y[mask] = 0 |
| Z = X[np.arange(N)[:, None, None], Y, :] |
| Z[mask] = np.nan |
| chi_angles = batch_compute_dihedral_angles(Z.reshape(-1, 4, 3)).reshape(N, 4) |
| return chi_angles |
|
|
|
|
| def get_chi_angles(coords, seq, return_onehot=False): |
| """ |
| |
| Parameters |
| ---------- |
| prody_pdb : prody.AtomGroup |
| prody pdb object or selection |
| return_coords : bool, optional |
| return coordinates of prody_pdb in (N, 14, 3) array format, by default False |
| return_onehot : bool, optional |
| return one-hot sequence of prody_pdb, by default False |
| |
| Returns |
| ------- |
| numpy array of shape (N, 4) |
| Array contains chi angles of sidechains in row-order of residue indices in prody_pdb. |
| If a chi angle is not defined for a residue, due to missing atoms or GLY / ALA, it is set to np.nan. |
| """ |
| onehot = get_onehot_sequence(seq) |
| dihedral_indices = get_dihedral_indices(onehot) |
| if return_onehot: |
| return _get_chi_angles(coords, dihedral_indices), onehot |
| return _get_chi_angles(coords, dihedral_indices) |
|
|
|
|
| def test_get_chi_angles(print_chi_angles=False): |
| |
| pdb = pr.parsePDB('6w70') |
| prody_pdb = pdb.select('chain A') |
| chi_angles = get_chi_angles(prody_pdb) |
| assert chi_angles.shape == (prody_pdb.ca.numAtoms(), 4) |
| assert chi_angles[0,0] < 56.0 and chi_angles[0,0] > 55.0 |
| print('test_get_chi_angles passed') |
| try: |
| os.remove('6w70.pdb.gz') |
| except: |
| pass |
| if print_chi_angles: |
| print(chi_angles) |
| return True |
|
|
|
|
| if __name__ == '__main__': |
| test_get_chi_angles(print_chi_angles=True) |
|
|
|
|
|
|
