| 100 | res2surf = res2surf + (surf_graph.num_cum_nodes - surf_graph.num_nodes)[graph.residue2graph].unsqueeze(-1) |
| 101 | |
| 102 | def surface_feature_init(self, graph, surf_graph, input): |
| 103 | # Residue -> surface graph correspondence |
| 104 | surf2res, dist = surface.knn_atoms(surf_graph.node_position, graph.node_position, k=self.k, batch_x=surf_graph.node2graph, batch_y=graph.node2graph) |
| 105 | surf2res = surf2res[:, :self.k] |
| 106 | dist = dist[:, :self.k].sqrt() |
| 107 | |
| 108 | # Knn graph for surface points |
| 109 | surf_edge_index = knn_graph(surf_graph.node_position, k=self.num_surf_graph_neighbor, batch=surf_graph.node2graph) |
| 110 | surf_node_in, surf_node_out = surf_edge_index |
| 111 | surf_pos_in, surf_pos_out = surf_graph.node_position[surf_node_in], surf_graph.node_position[surf_node_out] |
| 112 | surf_vec_edge = (surf_pos_in - surf_pos_out).unsqueeze(-2) # [n_edge, 1, 3] |
| 113 | h_surf_edge = rbf((surf_pos_out - surf_pos_in).norm(dim=-1), dim=self.surf_rbf_dim), surf_vec_edge |
| 114 | |
| 115 | # Inherit node features from residues with distance |
| 116 | h_surf_node = torch.cat((input[surf2res.flatten()], dist.view(-1, 1)), dim=-1) |
| 117 | h_surf_node = h_surf_node.view(surf_graph.num_node, self.k, -1) |
| 118 | h_surf_node = self.surf_in_linear(h_surf_node) |
| 119 | h_surf_node = h_surf_node.mean(dim=1) # Average pooling of K neighbors on the residue graph |
| 120 | h_surf_node = self.surf_in_mlp(torch.cat((h_surf_node, surf_graph.node_feature), dim=-1)) |
| 121 | |
| 122 | return surf_edge_index, h_surf_node, h_surf_edge |
| 123 | |
| 124 | def forward(self, graph, input, surf_graph, all_loss=None, metric=None): |
| 125 | # Input features |