Input: - xyz: current backbone cooordinates (B, L, 3, 3) - pair: pair features from Trunk (B, L, L, E) - idx: residue index from ground truth pdb Output: - G: defined graph
(xyz, pair, idx, top_k=64, kmin=32, eps=1e-6)
| 131 | return G, pair[b,i,j][...,None] |
| 132 | |
| 133 | def make_topk_graph(xyz, pair, idx, top_k=64, kmin=32, eps=1e-6): |
| 134 | ''' |
| 135 | Input: |
| 136 | - xyz: current backbone cooordinates (B, L, 3, 3) |
| 137 | - pair: pair features from Trunk (B, L, L, E) |
| 138 | - idx: residue index from ground truth pdb |
| 139 | Output: |
| 140 | - G: defined graph |
| 141 | ''' |
| 142 | |
| 143 | B, L = xyz.shape[:2] |
| 144 | device = xyz.device |
| 145 | |
| 146 | # distance map from current CA coordinates |
| 147 | D = torch.cdist(xyz, xyz) + torch.eye(L, device=device).unsqueeze(0)*999.9 # (B, L, L) |
| 148 | # seq sep |
| 149 | sep = idx[:,None,:] - idx[:,:,None] |
| 150 | sep = sep.abs() + torch.eye(L, device=device).unsqueeze(0)*999.9 |
| 151 | D = D + sep*eps |
| 152 | |
| 153 | # get top_k neighbors |
| 154 | D_neigh, E_idx = torch.topk(D, min(top_k, L), largest=False) # shape of E_idx: (B, L, top_k) |
| 155 | topk_matrix = torch.zeros((B, L, L), device=device) |
| 156 | topk_matrix.scatter_(2, E_idx, 1.0) |
| 157 | |
| 158 | # put an edge if any of the 3 conditions are met: |
| 159 | # 1) |i-j| <= kmin (connect sequentially adjacent residues) |
| 160 | # 2) top_k neighbors |
| 161 | cond = torch.logical_or(topk_matrix > 0.0, sep < kmin) |
| 162 | b,i,j = torch.where(cond) |
| 163 | |
| 164 | src = b*L+i |
| 165 | tgt = b*L+j |
| 166 | G = dgl.graph((src, tgt), num_nodes=B*L).to(device) |
| 167 | G.edata['rel_pos'] = (xyz[b,j,:] - xyz[b,i,:]).detach() # no gradient through basis function |
| 168 | |
| 169 | return G, pair[b,i,j][...,None] |
| 170 | |
| 171 | def make_rotX(angs, eps=1e-6): |
| 172 | B,L = angs.shape[:2] |