nn.NuclearEmbedding

class nn.NuclearEmbedding(*args: Any, **kwargs: Any)

Embedding which maps scalar nuclear charges Z to vectors in a (num_features)-dimensional feature space. The embedding consists of a freely learnable parameter matrix [Zmax, num_features] and a learned linear mapping from the electron configuration to a (num_features)-dimensional vector. The latter part encourages alchemically meaningful representations without restricting the expressivity of learned embeddings. Using complexe nuclear embedding can have negative impact on the model performance, when spin charge embedding is activated Negative performance in regard of the duration until the model converges. The model will converge to a lower value, but the duration is longer.

Args: num_features: Dimensions of feature space. Zmax: Maximum nuclear charge of atoms. The default is 100, so all

elements up to Fermium (Fm) (Z=100) are supported. Can be kept at the default value (has minimal memory impact).

zero_init: If True, initialize the embedding with zeros. Otherwise, use

uniform initialization.