concatenate

concatenate#

Concatenation utilities.

prxteinmpnn.utils.concatenate

prxteinmpnn.utils.concatenate.concatenate_neighbor_nodes(node_features, edge_features, neighbor_indices)[source]#

Concatenate node features with neighbor edge features.

Parameters:

  • node_features (Int[Array, 'num_atoms num_features']) – (L, C_V) node features

  • edge_features (Float[Array, 'num_atoms num_neighbors num_features']) – (L, K, C_E) edge features

  • neighbor_indices (Int[Array, 'num_atoms num_neighbors']) – (L, K) neighbor indices

Return type:

Float[Array, 'num_atoms num_neighbors num_features']

Returns:

(L, K, C_V + C_E) concatenated features for neighbors

Contents