Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import torch | |
| # ============================================================ | |
| def get_pair_dist(a, b): | |
| """calculate pair distances between two sets of points | |
| Parameters | |
| ---------- | |
| a,b : pytorch tensors of shape [batch,nres,3] | |
| store Cartesian coordinates of two sets of atoms | |
| Returns | |
| ------- | |
| dist : pytorch tensor of shape [batch,nres,nres] | |
| stores paitwise distances between atoms in a and b | |
| """ | |
| dist = torch.cdist(a, b, p=2) | |
| return dist | |
| # ============================================================ | |
| def get_ang(a, b, c): | |
| """calculate planar angles for all consecutive triples (a[i],b[i],c[i]) | |
| from Cartesian coordinates of three sets of atoms a,b,c | |
| Parameters | |
| ---------- | |
| a,b,c : pytorch tensors of shape [batch,nres,3] | |
| store Cartesian coordinates of three sets of atoms | |
| Returns | |
| ------- | |
| ang : pytorch tensor of shape [batch,nres] | |
| stores resulting planar angles | |
| """ | |
| v = a - b | |
| w = c - b | |
| v = v / torch.norm(v, dim=-1, keepdim=True) | |
| w = w / torch.norm(w, dim=-1, keepdim=True) | |
| # this is not stable at the poles | |
| #vw = torch.sum(v*w, dim=-1) | |
| #ang = torch.acos(vw) | |
| # this is better | |
| # https://math.stackexchange.com/questions/1143354/numerically-stable-method-for-angle-between-3d-vectors/1782769 | |
| y = torch.norm(v-w,dim=-1) | |
| x = torch.norm(v+w,dim=-1) | |
| ang = 2*torch.atan2(y, x) | |
| return ang | |
| # ============================================================ | |
| def get_dih(a, b, c, d): | |
| """calculate dihedral angles for all consecutive quadruples (a[i],b[i],c[i],d[i]) | |
| given Cartesian coordinates of four sets of atoms a,b,c,d | |
| Parameters | |
| ---------- | |
| a,b,c,d : pytorch tensors of shape [batch,nres,3] | |
| store Cartesian coordinates of four sets of atoms | |
| Returns | |
| ------- | |
| dih : pytorch tensor of shape [batch,nres] | |
| stores resulting dihedrals | |
| """ | |
| b0 = a - b | |
| b1r = c - b | |
| b2 = d - c | |
| b1 = b1r/torch.norm(b1r, dim=-1, keepdim=True) | |
| v = b0 - torch.sum(b0*b1, dim=-1, keepdim=True)*b1 | |
| w = b2 - torch.sum(b2*b1, dim=-1, keepdim=True)*b1 | |
| x = torch.sum(v*w, dim=-1) | |
| y = torch.sum(torch.cross(b1,v,dim=-1)*w, dim=-1) | |
| ang = torch.atan2(y, x) | |
| return ang | |
| # ============================================================ | |
| def xyz_to_c6d(xyz, params): | |
| """convert cartesian coordinates into 2d distance | |
| and orientation maps | |
| Parameters | |
| ---------- | |
| xyz : pytorch tensor of shape [batch,3,nres,3] | |
| stores Cartesian coordinates of backbone N,Ca,C atoms | |
| Returns | |
| ------- | |
| c6d : pytorch tensor of shape [batch,nres,nres,4] | |
| stores stacked dist,omega,theta,phi 2D maps | |
| """ | |
| batch = xyz.shape[0] | |
| nres = xyz.shape[2] | |
| # 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 = torch.cross(b, c, dim=-1) | |
| Cb = -0.58273431*a + 0.56802827*b - 0.54067466*c + Ca | |
| # 6d coordinates order: (dist,omega,theta,phi) | |
| c6d = torch.zeros([batch,nres,nres,4],dtype=xyz.dtype,device=xyz.device) | |
| dist = get_pair_dist(Cb,Cb) | |
| dist[torch.isnan(dist)] = 999.9 | |
| c6d[...,0] = dist + 999.9*torch.eye(nres,device=xyz.device)[None,...] | |
| b,i,j = torch.where(c6d[...,0]<params['DMAX']) | |
| c6d[b,i,j,torch.full_like(b,1)] = get_dih(Ca[b,i], Cb[b,i], Cb[b,j], Ca[b,j]) | |
| c6d[b,i,j,torch.full_like(b,2)] = get_dih(N[b,i], Ca[b,i], Cb[b,i], Cb[b,j]) | |
| c6d[b,i,j,torch.full_like(b,3)] = get_ang(Ca[b,i], Cb[b,i], Cb[b,j]) | |
| # fix long-range distances | |
| c6d[...,0][c6d[...,0]>=params['DMAX']] = 999.9 | |
| return c6d | |
| # ============================================================ | |
| def c6d_to_bins(c6d,params): | |
| """bin 2d distance and orientation maps | |
| """ | |
| dstep = (params['DMAX'] - params['DMIN']) / params['DBINS'] | |
| astep = 2.0*np.pi / params['ABINS'] | |
| dbins = torch.linspace(params['DMIN']+dstep, params['DMAX'], params['DBINS'],dtype=c6d.dtype,device=c6d.device) | |
| ab360 = torch.linspace(-np.pi+astep, np.pi, params['ABINS'],dtype=c6d.dtype,device=c6d.device) | |
| ab180 = torch.linspace(astep, np.pi, params['ABINS']//2,dtype=c6d.dtype,device=c6d.device) | |
| db = torch.bucketize(c6d[...,0].contiguous(),dbins) | |
| ob = torch.bucketize(c6d[...,1].contiguous(),ab360) | |
| tb = torch.bucketize(c6d[...,2].contiguous(),ab360) | |
| pb = torch.bucketize(c6d[...,3].contiguous(),ab180) | |
| ob[db==params['DBINS']] = params['ABINS'] | |
| tb[db==params['DBINS']] = params['ABINS'] | |
| pb[db==params['DBINS']] = params['ABINS']//2 | |
| return torch.stack([db,ob,tb,pb],axis=-1).to(torch.uint8) | |
| # ============================================================ | |
| def dist_to_bins(dist,params): | |
| """bin 2d distance maps | |
| """ | |
| dstep = (params['DMAX'] - params['DMIN']) / params['DBINS'] | |
| db = torch.round((dist-params['DMIN']-dstep/2)/dstep) | |
| db[db<0] = 0 | |
| db[db>params['DBINS']] = params['DBINS'] | |
| return db.long() | |
| # ============================================================ | |
| def c6d_to_bins2(c6d,params): | |
| """bin 2d distance and orientation maps | |
| (alternative slightly simpler version) | |
| """ | |
| dstep = (params['DMAX'] - params['DMIN']) / params['DBINS'] | |
| astep = 2.0*np.pi / params['ABINS'] | |
| db = torch.round((c6d[...,0]-params['DMIN']-dstep/2)/dstep) | |
| ob = torch.round((c6d[...,1]+np.pi-astep/2)/astep) | |
| tb = torch.round((c6d[...,2]+np.pi-astep/2)/astep) | |
| pb = torch.round((c6d[...,3]-astep/2)/astep) | |
| # put all d<dmin into one bin | |
| db[db<0] = 0 | |
| # synchronize no-contact bins | |
| db[db>params['DBINS']] = params['DBINS'] | |
| ob[db==params['DBINS']] = params['ABINS'] | |
| tb[db==params['DBINS']] = params['ABINS'] | |
| pb[db==params['DBINS']] = params['ABINS']//2 | |
| return torch.stack([db,ob,tb,pb],axis=-1).long() | |
| # ============================================================ | |
| def get_cb(N,Ca,C): | |
| """recreate Cb given N,Ca,C""" | |
| b = Ca - N | |
| c = C - Ca | |
| a = torch.cross(b, c, dim=-1) | |
| Cb = -0.58273431*a + 0.56802827*b - 0.54067466*c + Ca | |
| return Cb | |