Spaces:
Running
on
T4
Running
on
T4
| import numpy as np | |
| import scipy | |
| import scipy.spatial | |
| # calculate dihedral angles defined by 4 sets of points | |
| def get_dihedrals(a, b, c, d): | |
| b0 = -1.0*(b - a) | |
| b1 = c - b | |
| b2 = d - c | |
| b1 /= np.linalg.norm(b1, axis=-1)[:,None] | |
| v = b0 - np.sum(b0*b1, axis=-1)[:,None]*b1 | |
| w = b2 - np.sum(b2*b1, axis=-1)[:,None]*b1 | |
| x = np.sum(v*w, axis=-1) | |
| y = np.sum(np.cross(b1, v)*w, axis=-1) | |
| return np.arctan2(y, x) | |
| # calculate planar angles defined by 3 sets of points | |
| def get_angles(a, b, c): | |
| v = a - b | |
| v /= np.linalg.norm(v, axis=-1)[:,None] | |
| w = c - b | |
| w /= np.linalg.norm(w, axis=-1)[:,None] | |
| x = np.sum(v*w, axis=1) | |
| #return np.arccos(x) | |
| return np.arccos(np.clip(x, -1.0, 1.0)) | |
| # get 6d coordinates from x,y,z coords of N,Ca,C atoms | |
| def get_coords6d(xyz, dmax): | |
| nres = xyz.shape[1] | |
| # three anchor atoms | |
| N = xyz[0] | |
| Ca = xyz[1] | |
| C = xyz[2] | |
| # recreate Cb given N,Ca,C | |
| b = Ca - N | |
| c = C - Ca | |
| a = np.cross(b, c) | |
| Cb = -0.58273431*a + 0.56802827*b - 0.54067466*c + Ca | |
| # fast neighbors search to collect all | |
| # Cb-Cb pairs within dmax | |
| kdCb = scipy.spatial.cKDTree(Cb) | |
| indices = kdCb.query_ball_tree(kdCb, dmax) | |
| # indices of contacting residues | |
| idx = np.array([[i,j] for i in range(len(indices)) for j in indices[i] if i != j]).T | |
| idx0 = idx[0] | |
| idx1 = idx[1] | |
| # Cb-Cb distance matrix | |
| dist6d = np.full((nres, nres),999.9, dtype=np.float32) | |
| dist6d[idx0,idx1] = np.linalg.norm(Cb[idx1]-Cb[idx0], axis=-1) | |
| # matrix of Ca-Cb-Cb-Ca dihedrals | |
| omega6d = np.zeros((nres, nres), dtype=np.float32) | |
| omega6d[idx0,idx1] = get_dihedrals(Ca[idx0], Cb[idx0], Cb[idx1], Ca[idx1]) | |
| # matrix of polar coord theta | |
| theta6d = np.zeros((nres, nres), dtype=np.float32) | |
| theta6d[idx0,idx1] = get_dihedrals(N[idx0], Ca[idx0], Cb[idx0], Cb[idx1]) | |
| # matrix of polar coord phi | |
| phi6d = np.zeros((nres, nres), dtype=np.float32) | |
| phi6d[idx0,idx1] = get_angles(Ca[idx0], Cb[idx0], Cb[idx1]) | |
| mask = np.zeros((nres, nres), dtype=np.float32) | |
| mask[idx0, idx1] = 1.0 | |
| return dist6d, omega6d, theta6d, phi6d, mask | |