nn.ElectronicEmbedding

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

Single Head self attention like block for updating atomic features through nonlocal interactions with the electrons. The embeddings are used to map the total molecular charge or molecular spin to a feature vector. Since those properties are not localized on a specific atom they have to be delocalized over the whole molecule. The delocalization is achieved by using a self attention like mechanism.

Parameters:

• property_key – key of electronic property in the spk ‘inputs’ dictionary

• num_features – Dimensions of feature space aka the number of features to describe atomic environments. This determines the size of each embedding vector

• num_residual – Number of residual blocks applied to atomic features

• activation – activation function.

• is_charged – True corresponds to building embedding for molecular charge and separate weights are used for positive and negative charges. False corresponds to building embedding for spin values, no seperate weights are used

• epsilon – numerical stability parameter