core.models.equiformer_v2.prediction_heads.rank2

core.models.equiformer_v2.prediction_heads.rank2#

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#

Rank2Block

Output block for predicting rank-2 tensors (stress, dielectric tensor).

Rank2DecompositionEdgeBlock

Output block for predicting rank-2 tensors (stress, dielectric tensor, etc).

Rank2SymmetricTensorHead

A rank 2 symmetric tensor prediction head.

Module Contents#

class core.models.equiformer_v2.prediction_heads.rank2.Rank2Block(emb_size: int, num_layers: int = 2, edge_level: bool = False, extensive: bool = False)#

Bases: torch.nn.Module

Output block for predicting rank-2 tensors (stress, dielectric tensor). Applies outer product between edges and computes node-wise or edge-wise MLP.

Parameters:

  • emb_size (int) – Size of edge embedding used to compute outer product

  • num_layers (int) – Number of layers of the MLP

  • edge_level (bool) – If true apply MLP at edge level before pooling, otherwise use MLP at nodes after pooling

  • extensive (bool) – Whether to sum or average the outer products

edge_level#
emb_size#
extensive#
scalar_nonlinearity#
r2tensor_MLP#
forward(edge_distance_vec, x_edge, edge_index, data)#
Parameters:

  • edge_distance_vec (torch.Tensor) – Tensor of shape (…, 3)

  • x_edge (torch.Tensor) – Tensor of shape (…, emb_size)

  • edge_index (torch.Tensor) – Tensor of shape (2, nEdges)

  • data – LMDBDataset sample

class core.models.equiformer_v2.prediction_heads.rank2.Rank2DecompositionEdgeBlock(emb_size: int, num_layers: int = 2, edge_level: bool = False, extensive: bool = False)#

Bases: torch.nn.Module

Output block for predicting rank-2 tensors (stress, dielectric tensor, etc). Decomposes a rank-2 symmetric tensor into irrep degree 0 and 2.

Parameters:

  • emb_size (int) – Size of edge embedding used to compute outer product

  • num_layers (int) – Number of layers of the MLP

  • edge_level (bool) – If true apply MLP at edge level before pooling, otherwise use MLP at nodes after pooling

  • extensive (bool) – Whether to sum or average the outer products

emb_size#
edge_level#
extensive#
scalar_nonlinearity#
scalar_MLP#
irrep2_MLP#
change_mat#
forward(edge_distance_vec, x_edge, edge_index, data)#
Parameters:

  • edge_distance_vec (torch.Tensor) – Tensor of shape (…, 3)

  • x_edge (torch.Tensor) – Tensor of shape (…, emb_size)

  • edge_index (torch.Tensor) – Tensor of shape (2, nEdges)

  • data – LMDBDataset sample

class core.models.equiformer_v2.prediction_heads.rank2.Rank2SymmetricTensorHead(backbone: fairchem.core.models.base.BackboneInterface, output_name: str, decompose: bool = False, edge_level_mlp: bool = False, num_mlp_layers: int = 2, use_source_target_embedding: bool = False, extensive: bool = False, avg_num_nodes: int = 1.0, default_norm_type: str = 'layer_norm_sh')#

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

A rank 2 symmetric tensor prediction head.

ouput_name#

name of output prediction property (ie, stress)

sphharm_norm#

layer normalization for spherical harmonic edge weights

xedge_layer_norm#

embedding layer norm

block#

rank 2 equivariant symmetric tensor block

output_name#
decompose#
use_source_target_embedding#
avg_num_nodes#
sphharm_norm#
xedge_layer_norm#
forward(data: dict[str, torch.Tensor] | torch.Tensor, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#
Parameters:

  • data – data batch

  • emb – dictionary with embedding object and graph data

Returns: dict of {output property name: predicted value}

Contents