API Reference

This document provides a reference for the main functions in the module.

Enums

Layers

class Layers

An enumeration of supported layer types.

ELEM_EMBEDDING = 0

EDGE_EMBEDDING = 1

RADIAL_BASIS = 2

LINEAR_E3NN = 3

NEQUIP_CONV = 4

NEQUIP_CONV_BLOCK = 5

ElemEmbedding

class ElemEmbedding

An enumeration of supported element embedding types.

ONE_HOT = 0

BINARY = 1

ELECTRON = 2

get_embed_type_from_str(embed_str: str) → ElemEmbedding

Get element embedding type from string.

Parameters:

embed_str – Element embedding type as a string.

Returns:

Corresponding ElemEmbedding enum value.

Functions

summary()

Print a summary of supported layers and their required parameters.

get_element_embedding(embedding_type: str, n_elems: int = 118) → torch.nn.Module

Get torch module for element embedding.

Parameters:

• embedding_type – Element embedding type.

• n_elems – Number of elements (only for one_hot).

Returns:

Element embedding module.

get_edge_embedding(lmax: int, normalize: bool = True, normalization: str = 'component', parity: bool = True) → torch.nn.Module

Returns edge embedding module.

Parameters:

• lmax – Maximum l value for spherical harmonics.

• normalize – Whether to normalize the spherical harmonics.

• normalization – Normalization scheme to use.

• parity – Whether to use parity.

Returns:

Edge embedding module.

get_radial_basis(r_max: float | torch.Tensor, num_basis: int = 8, trainable: bool = True, power: int = 6) → torch.nn.Module

Returns radial basis module.

Parameters:

• r_max – Cutoff radius.

• num_basis – Number of basis functions.

• trainable – Whether the basis functions are trainable.

• power – Power used in envelope function.

Returns:

Radial basis module.

get_linear_e3nn(irreps_in, irreps_out) → torch.nn.Module

Get linear e3nn module.

Parameters:

• irreps_in – Input irreps.

• irreps_out – Output irreps.

Returns:

Linear e3nn module.

get_nequip_conv(parity: bool, lmax: int, conv_feature_size: int, node_embedding_irrep_in, node_attr_irrep, edge_attr_irrep, edge_embedding_irrep, avg_neigh=1, nonlinearity_type='gate', resnet=False, nonlinearity_scalars={'e': 'ssp', 'o': 'tanh'}, nonlinearity_gates={'e': 'ssp', 'o': 'abs'}, radial_network_hidden_dim=64, radial_network_layers=2) → torch.nn.Module

Get NequIP convolution layer.

Parameters:

• parity – Whether to use parity.

• lmax – Maximum l value for spherical harmonics.

• conv_feature_size – Convolution feature size.

• node_embedding_irrep_in – Input node embedding irreps.

• node_attr_irrep – Node attribute irreps.

• edge_attr_irrep – Edge attribute irreps.

• edge_embedding_irrep – Edge embedding irreps.

• avg_neigh – Average number of neighbors.

• nonlinearity_type – Nonlinearity type.

• resnet – Whether to use resnet.

• nonlinearity_scalars – Nonlinearity scalars.

• nonlinearity_gates – Nonlinearity gates.

• radial_network_hidden_dim – Radial network hidden dimension.

• radial_network_layers – Radial network layers.

Returns:

NequIP convolution layer.

get_nequip_conv_block(n_conv_layers: int, parity: bool, lmax: int, conv_feature_size: int, node_embedding_irrep_in, node_attr_irrep, edge_attr_irrep, edge_embedding_irrep, avg_neigh=1, nonlinearity_type='gate', resnet=False, nonlinearity_scalars={'e': 'ssp', 'o': 'tanh'}, nonlinearity_gates={'e': 'ssp', 'o': 'abs'}, radial_network_hidden_dim=64, radial_network_layers=2) → torch.nn.Module

Returns NequIP convolution block, with multiple convolution layers.

Parameters:

• n_conv_layers – Number of conv layers.

• parity – Whether to use parity.

• lmax – Maximum l value for spherical harmonics.

• conv_feature_size – Convolution feature size.

• node_embedding_irrep_in – Input node embedding irreps.

• node_attr_irrep – Node attribute irreps.

• edge_attr_irrep – Edge attribute irreps.

• edge_embedding_irrep – Edge embedding irreps.

• avg_neigh – Average number of neighbors.

• nonlinearity_type – Nonlinearity type.

• resnet – Whether to use resnet.

• nonlinearity_scalars – Nonlinearity scalars.

• nonlinearity_gates – Nonlinearity gates.

• radial_network_hidden_dim – Radial network hidden dimension.

• radial_network_layers – Radial network layers.

Returns:

NequIP convolution block.

get_model_from_yaml(yaml_file) → torch.nn.Sequential

Generate model sequentially from yaml file.

Parameters:

yaml_file – Path to yaml file.

Returns:

Sequential model.