Node encoder: `periodic_table_embedding`

Role

The node encoder produces the initial state $\bm h_i^{(0)}$ of every atom from its chemical species. Because the species is a label, the initial state is invariant (0e channels only); directional information enters later through edges.

Built-in members of this kind: species_irrep_embedding (one learned vector per species, the MANDALA default), embedding and one_hot followed by irrep_linear (DeepH-E3, MACE-H presets).

The new layer

A plain embedding table learns every element independently. Chemistry is not independent: elements in the same group share valence configurations and elements in the same period share core size. The periodic-table embedding factorizes the initial state accordingly,

\[\bm h_i^{(0)}=\bm E_{\text{species}}[z_i]+\bm E_{\text{period}}\big[\mathrm{period}(z_i)\big] +\bm E_{\text{group}}\big[\mathrm{group}(z_i)\big]\ \in\ \mathbb R^{d}\ (d\times\texttt{0e}).\]

In a multicomponent material such as ZnCu₂Sn(SeS)₂, S and Se share the group-16 vector and Zn and Cu share the period-4 vector, so information learned from one element transfers to its neighbours in the table. The species-specific table keeps the full flexibility of a plain embedding.

The species list is resolved from the training data and passed through the recipe as elements: ${species.order}, so the layer always agrees with the index convention of node_type_idx.

Symmetry. The output depends only on species labels: it is invariant under every element of E(3) and equivariant under permutations of atoms.

Implementation

def period_and_group(atomic_number: int) -> tuple[int, int]:
    """Period (1-7) and IUPAC group (1-18) of an element.

    Lanthanides and actinides are assigned to group 3.
    """
    period_ends = (2, 10, 18, 36, 54, 86, 118)
    period = next(i + 1 for i, end in enumerate(period_ends) if atomic_number <= end)
    position = atomic_number - (period_ends[period - 2] if period > 1 else 0)
    if period == 1:
        return period, 1 if position == 1 else 18
    if period in (2, 3):
        return period, position if position <= 2 else position + 10
    if period in (4, 5):
        return period, position
    if position <= 2:
        return period, position
    if position <= 17:
        return period, 3
    return period, position - 14

class PeriodicTableEmbedding(nn.Module):
    r"""``h_i = E_species[z_i] + E_period[period(z_i)] + E_group[group(z_i)]``.

    Elements sharing a row or column of the periodic table share part of their
    embedding, so chemically related species start from related features.
    """

    def __init__(self, elements, width):
        super().__init__()
        from ase.data import atomic_numbers

        rows = [period_and_group(atomic_numbers[symbol]) for symbol in elements]
        self.register_buffer("period", torch.tensor([row[0] - 1 for row in rows]))
        self.register_buffer("group", torch.tensor([row[1] - 1 for row in rows]))
        self.species = nn.Embedding(len(elements), width)
        self.periods = nn.Embedding(7, width)
        self.groups = nn.Embedding(18, width)
        for table in (self.species, self.periods, self.groups):
            nn.init.normal_(table.weight, std=0.2)

    def forward(self, species_index):
        return (
            self.species(species_index)
            + self.periods(self.period[species_index])
            + self.groups(self.group[species_index])
        )

The lookup tables of period and group are buffers: they move with the model between devices and are stored in checkpoints, but are not trained.

@register_transform("periodic_table_embedding")
def _periodic_table_embedding(options, inputs, mapper):
    (species,) = inputs
    require_index(species, "node")
    module = PeriodicTableEmbedding(**options)
    width = options["width"]
    return module, [FeatureSpec("node", irreps=f"{width}x0e", width=width)]

Usage

- id: node_encoder
  role: encoder
  transform: periodic_table_embedding
  inputs: [node_type_idx]
  outputs: [node_0]
  options: {elements: '${species.order}', width: 32}

Variations

Add embeddings of tabulated invariant descriptors (electronegativity, covalent radius, valence-electron count) through a small MLP.
Encode the basis set of each species (its shell list ${species.shells}), which distinguishes the same element computed with different orbital sets.

Node encoder: periodic_table_embedding