(self, msa, pair, R_in, T_in, xyz, state, idx, motif_mask, top_k=64, eps=1e-5)
| 235 | |
| 236 | @torch.cuda.amp.autocast(enabled=False) |
| 237 | def forward(self, msa, pair, R_in, T_in, xyz, state, idx, motif_mask, top_k=64, eps=1e-5): |
| 238 | B, N, L = msa.shape[:3] |
| 239 | |
| 240 | if motif_mask is None: |
| 241 | motif_mask = torch.zeros(L).bool() |
| 242 | |
| 243 | # process msa & pair features |
| 244 | node = self.norm_msa(msa[:,0]) |
| 245 | pair = self.norm_pair(pair) |
| 246 | state = self.norm_state(state) |
| 247 | |
| 248 | node = torch.cat((node, state), dim=-1) |
| 249 | node = self.norm_node(self.embed_x(node)) |
| 250 | pair = self.norm_edge1(self.embed_e1(pair)) |
| 251 | |
| 252 | neighbor = get_seqsep(idx) |
| 253 | rbf_feat = rbf(torch.cdist(xyz[:,:,1], xyz[:,:,1])) |
| 254 | pair = torch.cat((pair, rbf_feat, neighbor), dim=-1) |
| 255 | pair = self.norm_edge2(self.embed_e2(pair)) |
| 256 | |
| 257 | # define graph |
| 258 | if top_k != 0: |
| 259 | G, edge_feats = make_topk_graph(xyz[:,:,1,:], pair, idx, top_k=top_k) |
| 260 | else: |
| 261 | G, edge_feats = make_full_graph(xyz[:,:,1,:], pair, idx, top_k=top_k) |
| 262 | l1_feats = xyz - xyz[:,:,1,:].unsqueeze(2) |
| 263 | l1_feats = l1_feats.reshape(B*L, -1, 3) |
| 264 | |
| 265 | # apply SE(3) Transformer & update coordinates |
| 266 | shift = self.se3(G, node.reshape(B*L, -1, 1), l1_feats, edge_feats) |
| 267 | |
| 268 | state = shift['0'].reshape(B, L, -1) # (B, L, C) |
| 269 | |
| 270 | offset = shift['1'].reshape(B, L, 2, 3) |
| 271 | offset[:,motif_mask,...] = 0 # NOTE: motif mask is all zeros if not freeezing the motif |
| 272 | |
| 273 | delTi = offset[:,:,0,:] / 10.0 # translation |
| 274 | R = offset[:,:,1,:] / 100.0 # rotation |
| 275 | |
| 276 | Qnorm = torch.sqrt( 1 + torch.sum(R*R, dim=-1) ) |
| 277 | qA, qB, qC, qD = 1/Qnorm, R[:,:,0]/Qnorm, R[:,:,1]/Qnorm, R[:,:,2]/Qnorm |
| 278 | |
| 279 | delRi = torch.zeros((B,L,3,3), device=xyz.device) |
| 280 | delRi[:,:,0,0] = qA*qA+qB*qB-qC*qC-qD*qD |
| 281 | delRi[:,:,0,1] = 2*qB*qC - 2*qA*qD |
| 282 | delRi[:,:,0,2] = 2*qB*qD + 2*qA*qC |
| 283 | delRi[:,:,1,0] = 2*qB*qC + 2*qA*qD |
| 284 | delRi[:,:,1,1] = qA*qA-qB*qB+qC*qC-qD*qD |
| 285 | delRi[:,:,1,2] = 2*qC*qD - 2*qA*qB |
| 286 | delRi[:,:,2,0] = 2*qB*qD - 2*qA*qC |
| 287 | delRi[:,:,2,1] = 2*qC*qD + 2*qA*qB |
| 288 | delRi[:,:,2,2] = qA*qA-qB*qB-qC*qC+qD*qD |
| 289 | |
| 290 | Ri = einsum('bnij,bnjk->bnik', delRi, R_in) |
| 291 | Ti = delTi + T_in #einsum('bnij,bnj->bni', delRi, T_in) + delTi |
| 292 | |
| 293 | alpha = self.sc_predictor(msa[:,0], state) |
| 294 | return Ri, Ti, state, alpha |
nothing calls this directly
no test coverage detected