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=9)
| 107 | return neigh.unsqueeze(-1) |
| 108 | |
| 109 | def make_full_graph(xyz, pair, idx, top_k=64, kmin=9): |
| 110 | ''' |
| 111 | Input: |
| 112 | - xyz: current backbone cooordinates (B, L, 3, 3) |
| 113 | - pair: pair features from Trunk (B, L, L, E) |
| 114 | - idx: residue index from ground truth pdb |
| 115 | Output: |
| 116 | - G: defined graph |
| 117 | ''' |
| 118 | |
| 119 | B, L = xyz.shape[:2] |
| 120 | device = xyz.device |
| 121 | |
| 122 | # seq sep |
| 123 | sep = idx[:,None,:] - idx[:,:,None] |
| 124 | b,i,j = torch.where(sep.abs() > 0) |
| 125 | |
| 126 | src = b*L+i |
| 127 | tgt = b*L+j |
| 128 | G = dgl.graph((src, tgt), num_nodes=B*L).to(device) |
| 129 | G.edata['rel_pos'] = (xyz[b,j,:] - xyz[b,i,:]).detach() # no gradient through basis function |
| 130 | |
| 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 | ''' |