core.models.gemnet_gp.gemnet

core.models.gemnet_gp.gemnet#

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#

GraphParallelGemNetT

GemNet-T, triplets-only variant of GemNet

Module Contents#

class core.models.gemnet_gp.gemnet.GraphParallelGemNetT(num_spherical: int, num_radial: int, num_blocks: int, 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, regress_forces: bool = True, direct_forces: bool = False, cutoff: float = 6.0, max_neighbors: int = 50, rbf: dict | None = None, envelope: dict | None = None, cbf: dict | None = None, extensive: bool = True, otf_graph: bool = False, use_pbc: bool = True, use_pbc_single: bool = False, output_init: str = 'HeOrthogonal', activation: str = 'swish', scale_num_blocks: bool = False, scatter_atoms: bool = True, scale_file: str | None = None)#

Bases: torch.nn.Module, fairchem.core.models.base.GraphModelMixin

GemNet-T, triplets-only variant of GemNet

Parameters:

  • num_spherical (int) – Controls maximum frequency.

  • num_radial (int) – Controls maximum frequency.

  • num_blocks (int) – Number of building blocks to be stacked.

  • 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.

  • regress_forces (bool) – Whether to predict forces. Default: True

  • direct_forces (bool) – If True predict forces based on aggregation of interatomic directions. If False predict forces based on negative gradient of energy potential.

  • cutoff (float) – Embedding cutoff for interactomic directions in Angstrom.

  • rbf (dict) – Name and hyperparameters of the radial basis function.

  • envelope (dict) – Name and hyperparameters of the envelope function.

  • cbf (dict) – Name and hyperparameters of the cosine basis function.

  • extensive (bool) – Whether the output should be extensive (proportional to the number of atoms)

  • output_init (str) – Initialization method for the final dense layer.

  • activation (str) – Name of the activation function.

  • scale_file (str) – Path to the json file containing the scaling factors.

num_blocks#
extensive#
scale_num_blocks#
scatter_atoms#
cutoff#
max_neighbors#
regress_forces#
otf_graph#
use_pbc#
use_pbc_single#
direct_forces#
radial_basis#
cbf_basis3#
mlp_rbf3#
mlp_cbf3#
mlp_rbf_h#
mlp_rbf_out#
atom_emb#
edge_emb#
out_blocks#
int_blocks#
get_triplets(edge_index, num_atoms)#

Get all b->a for each edge c->a. It is possible that b=c, as long as the edges are distinct.

Returns:

  • id3_ba (torch.Tensor, shape (num_triplets,)) – Indices of input edge b->a of each triplet b->a<-c

  • id3_ca (torch.Tensor, shape (num_triplets,)) – Indices of output edge c->a of each triplet b->a<-c

  • id3_ragged_idx (torch.Tensor, shape (num_triplets,)) – Indices enumerating the copies of id3_ca for creating a padded matrix

select_symmetric_edges(tensor: torch.Tensor, mask, reorder_idx, inverse_neg) → torch.Tensor#
reorder_symmetric_edges(edge_index, cell_offsets, neighbors, edge_dist, edge_vector)#

Reorder edges to make finding counter-directional edges easier.

Some edges are only present in one direction in the data, since every atom has a maximum number of neighbors. Since we only use i->j edges here, we lose some j->i edges and add others by making it symmetric. We could fix this by merging edge_index with its counter-edges, including the cell_offsets, and then running torch.unique. But this does not seem worth it.

select_edges(data, edge_index, cell_offsets, neighbors, edge_dist, edge_vector, cutoff=None)#
generate_interaction_graph(data)#
forward(data)#
property num_params#