core.models.gemnet_oc.layers.efficient

core.models.gemnet_oc.layers.efficient#

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#

BasisEmbedding

Embed a basis (CBF, SBF), optionally using the efficient reformulation.

EfficientInteractionBilinear

Efficient reformulation of the bilinear layer and subsequent summation.

Module Contents#

class core.models.gemnet_oc.layers.efficient.BasisEmbedding(num_radial: int, emb_size_interm: int, num_spherical: int | None = None)#

Bases: torch.nn.Module

Embed a basis (CBF, SBF), optionally using the efficient reformulation.

Parameters:

  • num_radial (int) – Number of radial basis functions.

  • emb_size_interm (int) – Intermediate embedding size of triplets/quadruplets.

  • num_spherical (int) – Number of circular/spherical basis functions. Only required if there is a circular/spherical basis.

weight: torch.nn.Parameter#
num_radial#
num_spherical#
reset_parameters() → None#
forward(rad_basis, sph_basis=None, idx_rad_outer=None, idx_rad_inner=None, idx_sph_outer=None, idx_sph_inner=None, num_atoms=None)#
Parameters:

  • rad_basis (torch.Tensor, shape=(num_edges, num_radial or num_orders * num_radial)) – Raw radial basis.

  • sph_basis (torch.Tensor, shape=(num_triplets or num_quadruplets, num_spherical)) – Raw spherical or circular basis.

  • idx_rad_outer (torch.Tensor, shape=(num_edges)) – Atom associated with each radial basis value. Optional, used for efficient edge aggregation.

  • idx_rad_inner (torch.Tensor, shape=(num_edges)) – Enumerates radial basis values per atom. Optional, used for efficient edge aggregation.

  • idx_sph_outer (torch.Tensor, shape=(num_triplets or num_quadruplets)) – Edge associated with each circular/spherical basis value. Optional, used for efficient triplet/quadruplet aggregation.

  • idx_sph_inner (torch.Tensor, shape=(num_triplets or num_quadruplets)) – Enumerates circular/spherical basis values per edge. Optional, used for efficient triplet/quadruplet aggregation.

  • num_atoms (int) – Total number of atoms. Optional, used for efficient edge aggregation.

Returns:

  • rad_W1 (torch.Tensor, shape=(num_edges, emb_size_interm, num_spherical))

  • sph (torch.Tensor, shape=(num_edges, Kmax, num_spherical)) – Kmax = maximum number of neighbors of the edges

class core.models.gemnet_oc.layers.efficient.EfficientInteractionBilinear(emb_size_in: int, emb_size_interm: int, emb_size_out: int)#

Bases: torch.nn.Module

Efficient reformulation of the bilinear layer and subsequent summation.

Parameters:

  • emb_size_in (int) – Embedding size of input triplets/quadruplets.

  • emb_size_interm (int) – Intermediate embedding size of the basis transformation.

  • emb_size_out (int) – Embedding size of output triplets/quadruplets.

emb_size_in#
emb_size_interm#
emb_size_out#
bilinear#
forward(basis, m, idx_agg_outer, idx_agg_inner, idx_agg2_outer=None, idx_agg2_inner=None, agg2_out_size=None)#
Parameters:

  • basis (Tuple (torch.Tensor, torch.Tensor),) –

    shapes=((num_edges, emb_size_interm, num_spherical),

    (num_edges, num_spherical, Kmax))

    First element: Radial basis multiplied with weight matrix Second element: Circular/spherical basis

  • m (torch.Tensor, shape=(num_edges, emb_size_in)) – Input edge embeddings

  • idx_agg_outer (torch.Tensor, shape=(num_triplets or num_quadruplets)) – Output edge aggregating this intermediate triplet/quadruplet edge.

  • idx_agg_inner (torch.Tensor, shape=(num_triplets or num_quadruplets)) – Enumerates intermediate edges per output edge.

  • idx_agg2_outer (torch.Tensor, shape=(num_edges)) – Output atom aggregating this edge.

  • idx_agg2_inner (torch.Tensor, shape=(num_edges)) – Enumerates edges per output atom.

  • agg2_out_size (int) – Number of output embeddings when aggregating twice. Typically the number of atoms.

Returns:

m_ca – Aggregated edge/atom embeddings.

Return type:

torch.Tensor, shape=(num_edges, emb_size)

Contents