terminator.models.layers.utils.merge_duplicate_pairE_geometric

terminator.models.layers.utils.merge_duplicate_pairE_geometric(h_E, edge_index)[source]

Sparse method to average pair energy tables across bidirectional edges with Torch Geometric.

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)

This function assumes edge_index is sorted by columns, and will fail if this is not the case.

Parameters:

  • h_E (torch.Tensor) – Pair energies in Torch Geometric sparse form Shape : n_edge x 400

  • E_idx (torch.LongTensor) – Torch Geometric sparse edge indices Shape : 2 x n_edge

Returns:

Pair energies with merged energies for bidirectional edges Shape : n_edge x 400

Return type:

torch.Tensor