terminator.models.layers.utils.cat_neighbors_nodes

terminator.models.layers.utils.cat_neighbors_nodes(h_nodes, h_neighbors, E_idx)[source]

Concatenate node features onto the ends of gathered edge features given kNN sparse edge indices

From https://github.com/jingraham/neurips19-graph-protein-design

Parameters:

  • h_nodes (torch.Tensor) – The node features for all nodes Shape: n_batch x n_res x n_hidden

  • h_neighbors (torch.Tensor) – The gathered edge features Shape: n_batch x n_res x k x n_hidden

  • E_idx (torch.LongTensor) – kNN sparse edge indices Shape : n_batch x n_res x k

Returns:

h_nn – The gathered concatenated node and edge features Shape : n_batch x n_res x k x n_hidden

Return type:

torch.Tensor