core.models.uma.escn_md

Contents

core.models.uma.escn_md#

Copyright (c) Meta Platforms, Inc. and affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Attributes#

ESCNMD_DEFAULT_EDGE_CHUNK_SIZE

Classes#

eSCNMDBackbone

Base class for all neural network modules.

MLP_EFS_Head

Base class for all neural network modules.

MLP_Energy_Head

Base class for all neural network modules.

Linear_Energy_Head

Base class for all neural network modules.

Linear_Force_Head

Base class for all neural network modules.

MLP_Stress_Head

Base class for all neural network modules.

Functions#

compose_tensor(→ torch.Tensor)

Re-compose a tensor from its decomposition

Module Contents#

core.models.uma.escn_md.ESCNMD_DEFAULT_EDGE_CHUNK_SIZE#
class core.models.uma.escn_md.eSCNMDBackbone(max_num_elements: int = 100, sphere_channels: int = 128, lmax: int = 2, mmax: int = 2, grid_resolution: int | None = None, num_sphere_samples: int = 128, otf_graph: bool = False, max_neighbors: int = 300, use_pbc: bool = True, use_pbc_single: bool = True, cutoff: float = 5.0, edge_channels: int = 128, distance_function: Literal['gaussian'] = 'gaussian', num_distance_basis: int = 512, direct_forces: bool = True, regress_forces: bool = True, regress_stress: bool = False, num_layers: int = 2, hidden_channels: int = 128, norm_type: str = 'rms_norm_sh', act_type: str = 'gate', ff_type: str = 'grid', activation_checkpointing: bool = False, chg_spin_emb_type: Literal['pos_emb', 'lin_emb', 'rand_emb'] = 'pos_emb', cs_emb_grad: bool = False, dataset_emb_grad: bool = False, dataset_list: list[str] | None = None, use_dataset_embedding: bool = True, use_cuda_graph_wigner: bool = False, radius_pbc_version: int = 1, always_use_pbc: bool = True)#

Bases: torch.nn.Module, fairchem.core.models.uma.nn.mole_utils.MOLEInterface

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested 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) -> None:
        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 also have their parameters converted 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.

max_num_elements#
lmax#
mmax#
sphere_channels#
grid_resolution#
num_sphere_samples#
always_use_pbc#
regress_forces#
direct_forces#
regress_stress#
otf_graph#
max_neighbors#
radius_pbc_version#
enforce_max_neighbors_strictly = False#
chg_spin_emb_type#
cs_emb_grad#
dataset_emb_grad#
dataset_list#
use_dataset_embedding#
sph_feature_size#
mappingReduced#
SO3_grid#
sphere_embedding#
charge_embedding#
spin_embedding#
cutoff#
edge_channels#
distance_function#
num_distance_basis#
source_embedding#
target_embedding#
edge_channels_list#
edge_degree_embedding#
envelope#
num_layers#
hidden_channels#
norm_type#
act_type#
ff_type#
blocks#
norm#
_get_rotmat_and_wigner(edge_distance_vecs: torch.Tensor) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#
_get_displacement_and_cell(data_dict: fairchem.core.datasets.atomic_data.AtomicData) → tuple[torch.Tensor | None, torch.Tensor | None]#
csd_embedding(charge, spin, dataset)#
_generate_graph(data_dict)#
forward(data_dict: fairchem.core.datasets.atomic_data.AtomicData) → dict[str, torch.Tensor]#
_init_gp_partitions(graph_dict, atomic_numbers_full)#

Graph Parallel This creates the required partial tensors for each rank given the full tensors. The tensors are split on the dimension along the node index using node_partition.

property num_params: int#
no_weight_decay() → set#
class core.models.uma.escn_md.MLP_EFS_Head(backbone: eSCNMDBackbone, prefix: str | None = None, wrap_property: 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 them to be nested 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) -> None:
        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 also have their parameters converted 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.

regress_stress#
regress_forces#
prefix#
wrap_property#
sphere_channels#
hidden_channels#
energy_block#
forward(data: fairchem.core.datasets.atomic_data.AtomicData, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#

Head forward.

Parameters:

  • data (AtomicData) – 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.uma.escn_md.MLP_Energy_Head(backbone: eSCNMDBackbone, reduce: str = 'sum')#

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 them to be nested 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) -> None:
        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 also have their parameters converted 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#
sphere_channels#
hidden_channels#
energy_block#
forward(data_dict: fairchem.core.datasets.atomic_data.AtomicData, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#

Head forward.

Parameters:

  • data (AtomicData) – 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.uma.escn_md.Linear_Energy_Head(backbone: eSCNMDBackbone, reduce: str = 'sum')#

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 them to be nested 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) -> None:
        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 also have their parameters converted 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#
energy_block#
forward(data_dict: fairchem.core.datasets.atomic_data.AtomicData, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#

Head forward.

Parameters:

  • data (AtomicData) – 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.uma.escn_md.Linear_Force_Head(backbone: eSCNMDBackbone)#

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 them to be nested 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) -> None:
        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 also have their parameters converted 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.

linear#
forward(data_dict: fairchem.core.datasets.atomic_data.AtomicData, emb: dict[str, torch.Tensor])#

Head forward.

Parameters:

  • data (AtomicData) – 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]

core.models.uma.escn_md.compose_tensor(trace: torch.Tensor, l2_symmetric: torch.Tensor) → torch.Tensor#

Re-compose a tensor from its decomposition

Parameters:

  • trace – a tensor with scalar part of the decomposition of r2 tensors in the batch

  • l2_symmetric – tensor with the symmetric/traceless part of decomposition

Returns:

rank 2 tensor

Return type:

tensor

class core.models.uma.escn_md.MLP_Stress_Head(backbone: eSCNMDBackbone, reduce: str = 'mean')#

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 them to be nested 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) -> None:
        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 also have their parameters converted 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#
sphere_channels#
hidden_channels#
scalar_block#
l2_linear#
forward(data_dict: fairchem.core.datasets.atomic_data.AtomicData, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#

Head forward.

Parameters:

  • data (AtomicData) – 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]

Contents