enerzyme.models.layers.electron_embedding.NonlinearElectronicEmbedding

class enerzyme.models.layers.electron_embedding.NonlinearElectronicEmbedding(dim_embedding: int, num_residual: int, activation_fn: str = 'swish', attribute: Literal['charge', 'spin'] = 'charge')

Bases: BaseElectronEmbedding

Block for updating atomic features through nonlocal interactions with the electrons.

Arguments:

num_features (int):

Dimensions of feature space.

num_basis_functions (int):

Number of radial basis functions.

num_residual_pre_i (int):

Number of residual blocks applied to atomic features in i branch (central atoms) before computing the interaction.

num_residual_pre_j (int):

Number of residual blocks applied to atomic features in j branch (neighbouring atoms) before computing the interaction.

num_residual_post (int):

Number of residual blocks applied to interaction features.

activation (str):

Kind of activation function. Possible values: ‘swish’: Swish activation function. ‘ssp’: Shifted softplus activation function.

__init__(dim_embedding: int, num_residual: int, activation_fn: str = 'swish', attribute: Literal['charge', 'spin'] = 'charge') → None

Initializes the NonlinearElectronicEmbedding class.

get_electron_embedding(atom_embedding: Tensor, Q: Tensor | None = None, S: Tensor | None = None, batch_seg: Tensor | None = None, mask: Tensor | None = None, eps: float = 1e-08) → Tensor

Evaluate interaction block. N: Number of atoms.

x (FloatTensor [N, num_features]):

Atomic feature vectors.

reset_parameters() → None

Initialize parameters.