terminator.models.layers.utils.gather_pairEs

terminator.models.layers.utils.gather_pairEs(pairEs, neighbor_idx)[source]

Gather the pair energies features of the nearest neighbors.

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

Parameters:

  • pairEs (torch.Tensor) – The pair energies in dense form Shape: n_batch x n_res x n_res x n_aa x n_aa

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

Returns:

pairE_features – The gathered pair energies Shape : n_batch x n_res x k x n_aa x n_aa

Return type:

torch.Tensor