MolecularDiffusion.modules.layers.esen.nn.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 MolecularDiffusion.modules.layers.esen.nn.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

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

edge_level = False
emb_size
extensive = False
r2tensor_MLP
scalar_nonlinearity
class MolecularDiffusion.modules.layers.esen.nn.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

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

change_mat
edge_level = False
emb_size
extensive = False
irrep2_MLP
scalar_MLP
scalar_nonlinearity
class MolecularDiffusion.modules.layers.esen.nn.rank2.Rank2SymmetricTensorHead(backbone: MolecularDiffusion.modules.layers.esen.nn.utils.BackboneInterface, output_name: str = 'stress', 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, MolecularDiffusion.modules.layers.esen.nn.utils.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

Parameters:

  • backbone – Backbone model that the head is attached to

  • decompose – Whether to decompose the rank2 tensor into isotropic and anisotropic components

  • edge_level_mlp – If true apply MLP at edge level before pooling, otherwise use MLP at nodes after pooling

  • num_mlp_layers – number of MLP layers

  • use_source_target_embedding – Whether to use both source and target atom embeddings

  • extensive – Whether to do sum-pooling (extensive) vs mean pooling (intensive).

  • avg_num_nodes – Used only if extensive to divide prediction by avg num nodes.

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}

avg_num_nodes = 1.0
decompose = False
lmax
output_name = 'stress'
sphharm_norm
use_source_target_embedding = False
xedge_layer_norm