terminator.models.layers.utils.merge_duplicate_pairE_dense

terminator.models.layers.utils.merge_duplicate_pairE_dense(h_E, E_idx)

Dense method to average pair energy tables across bidirectional edges.

TERMinator edges are represented as two bidirectional edges, and to allow for communication between these edges we average the embeddings. In the case for pair energies, we transpose the tables to ensure that the pair energy table is symmetric upon inverse (e.g. the pair energy between i and j should be the same as the pair energy between j and i)

Parameters:

• h_E (torch.Tensor) – Pair energies in kNN sparse form Shape : n_batch x n_res x k x n_aa x n_aa

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

Returns:

Pair energies with merged energies for bidirectional edges Shape : n_batch x n_res x k x n_aa x n_aa

Return type:

torch.Tensor