enerzyme.models.layers.atom_embedding.NuclearEmbedding

class enerzyme.models.layers.atom_embedding.NuclearEmbedding(max_Za: int, dim_embedding: int, zero_init: bool = True, use_electron_config: bool = True)

Bases: BaseAtomEmbedding

__init__(max_Za: int, dim_embedding: int, zero_init: bool = True, use_electron_config: bool = True) → None

get_embedding(Za: Tensor) → Tensor

Assign corresponding embeddings to nuclear charges. N: Number of atoms. num_features: Dimensions of feature space.

Arguments:

Za (LongTensor [N]):

Nuclear charges (atomic numbers) of atoms.

Returns:

x (FloatTensor [N, num_features]):

Embeddings of all atoms.

reset_parameters(zero_init: bool = True) → None

Initialize parameters.

train(mode: bool = True) → None

Switch between training and evaluation mode.