core.models.gemnet_gp.layers.interaction_block

core.models.gemnet_gp.layers.interaction_block#

Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes#

InteractionBlockTripletsOnly

Interaction block for GemNet-T/dT.

TripletInteraction

Triplet-based message passing block.

Module Contents#

class core.models.gemnet_gp.layers.interaction_block.InteractionBlockTripletsOnly(emb_size_atom: int, emb_size_edge: int, emb_size_trip: int, emb_size_rbf: int, emb_size_cbf: int, emb_size_bil_trip: int, num_before_skip: int, num_after_skip: int, num_concat: int, num_atom: int, activation: str | None = None, name: str = 'Interaction')#

Bases: torch.nn.Module

Interaction block for GemNet-T/dT.

Parameters:

  • emb_size_atom (int) – Embedding size of the atoms.

  • emb_size_edge (int) – Embedding size of the edges.

  • emb_size_trip (int) – (Down-projected) Embedding size in the triplet message passing block.

  • emb_size_rbf (int) – Embedding size of the radial basis transformation.

  • emb_size_cbf (int) – Embedding size of the circular basis transformation (one angle).

  • emb_size_bil_trip (int) – Embedding size of the edge embeddings in the triplet-based message passing block after the bilinear layer.

  • num_before_skip (int) – Number of residual blocks before the first skip connection.

  • num_after_skip (int) – Number of residual blocks after the first skip connection.

  • num_concat (int) – Number of residual blocks after the concatenation.

  • num_atom (int) – Number of residual blocks in the atom embedding blocks.

  • activation (str) – Name of the activation function to use in the dense layers except for the final dense layer.

name#
block_nr#
dense_ca#
trip_interaction#
layers_before_skip#
layers_after_skip#
atom_update#
concat_layer#
residual_m#
inv_sqrt_2#
forward(h: torch.Tensor, m: torch.Tensor, rbf3, cbf3, id3_ragged_idx, id_swap, id3_ba, id3_ca, rbf_h, idx_s, idx_t, edge_offset, Kmax, nAtoms)#
Returns:

  • h (torch.Tensor, shape=(nEdges, emb_size_atom)) – Atom embeddings.

  • m (torch.Tensor, shape=(nEdges, emb_size_edge)) – Edge embeddings (c->a).

  • Node (h)

  • Edge (m, rbf3, id_swap, rbf_h, idx_s, idx_t, cbf3[0], cbf3[1] (dense))

  • Triplet (id3_ragged_idx, id3_ba, id3_ca)

class core.models.gemnet_gp.layers.interaction_block.TripletInteraction(emb_size_edge: int, emb_size_trip: int, emb_size_bilinear: int, emb_size_rbf: int, emb_size_cbf: int, activation: str | None = None, name: str = 'TripletInteraction', **kwargs)#

Bases: torch.nn.Module

Triplet-based message passing block.

Parameters:

  • emb_size_edge (int) – Embedding size of the edges.

  • emb_size_trip (int) – (Down-projected) Embedding size of the edge embeddings after the hadamard product with rbf.

  • emb_size_bilinear (int) – Embedding size of the edge embeddings after the bilinear layer.

  • emb_size_rbf (int) – Embedding size of the radial basis transformation.

  • emb_size_cbf (int) – Embedding size of the circular basis transformation (one angle).

  • activation (str) – Name of the activation function to use in the dense layers except for the final dense layer.

name#
dense_ba#
mlp_rbf#
scale_rbf#
mlp_cbf#
scale_cbf_sum#
down_projection#
up_projection_ca#
up_projection_ac#
inv_sqrt_2#
forward(m: torch.Tensor, rbf3, cbf3, id3_ragged_idx, id_swap, id3_ba, id3_ca, edge_offset, Kmax)#
Returns:

m – Edge embeddings (c->a).

Return type:

torch.Tensor, shape=(nEdges, emb_size_edge)

Contents