core.models.gemnet_oc.layers.embedding_block

core.models.gemnet_oc.layers.embedding_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#

AtomEmbedding

Initial atom embeddings based on the atom type

EdgeEmbedding

Edge embedding based on the concatenation of atom embeddings

Module Contents#

class core.models.gemnet_oc.layers.embedding_block.AtomEmbedding(emb_size: int, num_elements: int)#

Bases: torch.nn.Module

Initial atom embeddings based on the atom type

Parameters:

emb_size (int) – Atom embeddings size

emb_size#
embeddings#
forward(Z) → torch.Tensor#
Returns:

h – Atom embeddings.

Return type:

torch.Tensor, shape=(nAtoms, emb_size)

class core.models.gemnet_oc.layers.embedding_block.EdgeEmbedding(atom_features: int, edge_features: int, out_features: int, activation: str | None = None)#

Bases: torch.nn.Module

Edge embedding based on the concatenation of atom embeddings and a subsequent dense layer.

Parameters:

  • atom_features (int) – Embedding size of the atom embedding.

  • edge_features (int) – Embedding size of the input edge embedding.

  • out_features (int) – Embedding size after the dense layer.

  • activation (str) – Activation function used in the dense layer.

in_features#
dense#
forward(h: torch.Tensor, m: torch.Tensor, edge_index) → torch.Tensor#
Parameters:

  • h (torch.Tensor, shape (num_atoms, atom_features)) – Atom embeddings.

  • m (torch.Tensor, shape (num_edges, edge_features)) – Radial basis in embedding block, edge embedding in interaction block.

Returns:

m_st – Edge embeddings.

Return type:

torch.Tensor, shape=(nEdges, emb_size)

Contents