core.models.equiformer_v2.equiformer_v2_dens

core.models.equiformer_v2.equiformer_v2_dens#

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#

`EqV2DeNSBackbone`	DeNS extra Args:
`DeNSScalarHead`	Base class for all neural network modules.
`DeNSVectorHead`	Base class for all neural network modules.
`DeNSRank2Head`	A rank 2 symmetric tensor prediction head.

Module Contents#

class core.models.equiformer_v2.equiformer_v2_dens.EqV2DeNSBackbone(use_pbc: bool = True, use_pbc_single: bool = False, regress_forces: bool = True, otf_graph: bool = True, max_neighbors: int = 500, max_radius: float = 5.0, max_num_elements: int = 90, num_layers: int = 12, sphere_channels: int = 128, attn_hidden_channels: int = 128, num_heads: int = 8, attn_alpha_channels: int = 32, attn_value_channels: int = 16, ffn_hidden_channels: int = 512, norm_type: str = 'rms_norm_sh', lmax_list: list[int] | None = None, mmax_list: list[int] | None = None, grid_resolution: int | None = None, num_sphere_samples: int = 128, edge_channels: int = 128, use_atom_edge_embedding: bool = True, share_atom_edge_embedding: bool = False, use_m_share_rad: bool = False, distance_function: str = 'gaussian', num_distance_basis: int = 512, attn_activation: str = 'scaled_silu', use_s2_act_attn: bool = False, use_attn_renorm: bool = True, ffn_activation: str = 'scaled_silu', use_gate_act: bool = False, use_grid_mlp: bool = False, use_sep_s2_act: bool = True, alpha_drop: float = 0.1, drop_path_rate: float = 0.05, proj_drop: float = 0.0, weight_init: str = 'normal', enforce_max_neighbors_strictly: bool = True, avg_num_nodes: float | None = None, avg_degree: float | None = None, use_energy_lin_ref: bool | None = False, load_energy_lin_ref: bool | None = False, activation_checkpoint: bool | None = False, use_force_encoding=True, use_noise_schedule_sigma_encoding: bool = False)#

Bases: fairchem.core.models.equiformer_v2.equiformer_v2.EquiformerV2Backbone

DeNS extra Args:: use_force_encoding (bool): For ablation study, whether to encode forces during denoising positions. Default: True. use_noise_schedule_sigma_encoding (bool): For ablation study, whether to encode the sigma (sampled std of Gaussian noises) during

denoising positions when fixed_noise_std = False in config files. Default: False.

use_force_encoding#

use_noise_schedule_sigma_encoding#

irreps_sh#

force_embedding#

forward(data) → dict[str, torch.Tensor]#

class core.models.equiformer_v2.equiformer_v2_dens.DeNSScalarHead(backbone: fairchem.core.models.base.BackboneInterface, output_name: str = 'energy', reduce: Literal['sum', 'mean'] = 'sum', use_denoising: bool = True)#

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

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

reduce#

avg_num_nodes#

scalar_block#

output_name#

use_denoising#

forward(data: torch_geometric.data.batch.Batch, emb: dict[str, torch.Tensor | fairchem.core.models.base.GraphData]) → dict[str, torch.Tensor]#

Head forward.

Parameters:

data (DataBatch) – Atomic systems as input
emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

class core.models.equiformer_v2.equiformer_v2_dens.DeNSVectorHead(backbone: fairchem.core.models.base.BackboneInterface, output_name: str = 'forces')#

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

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

output_name#

activation_checkpoint#

vector_block#

denoising_pos_block#

forward(data: torch_geometric.data.batch.Batch, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#

Head forward.

Parameters:

data (DataBatch) – Atomic systems as input
emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

class core.models.equiformer_v2.equiformer_v2_dens.DeNSRank2Head(backbone: fairchem.core.models.base.BackboneInterface, *args, use_denoising: bool = True, **kwargs)#

Bases: fairchem.core.models.equiformer_v2.heads.rank2.Rank2SymmetricTensorHead

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

use_denoising#

forward(data: torch_geometric.data.batch.Batch, 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}