Node Embedding

Among the node embedding methods provided below, we recommend using LinearNodeEmbedding in most cases. Other embedding strategies may not consistently achieve the best performance across different datasets. The linear embedding, which relies solely on element, is the most conservative and stable choice.

In principle, node embeddings can also incorporate information from the local environment (including scalar and tensor features). However, based on our current experiments, the linear approach remains the most reliable.

class tace.models._e3nn.node.LinearNodeEmbedding(num_elements: int, num_radial_basis: int, num_channel: int, Lmax: int, lmax: int, avg_num_neighbors: float, bias: bool = False, so2_angular_basis: SO3Rotation | None = None)

Bases: NodeEmbedding

A simple node embedding module based on a linear transformation.

This class projects discrete node attributes (e.g., element types) into a continuous feature space using a single linear layer, without introducing nonlinearity or structural information.

class tace.models._e3nn.node.NonLinearNodeEmbedding(num_elements: int, num_radial_basis: int, num_channel: int, Lmax: int, lmax: int, avg_num_neighbors: float, bias: bool = False, so2_angular_basis: SO3Rotation | None = None)

Bases: NodeEmbedding

A node embedding module with nonlinear transformation.

This class applies a nonlinear activation function after a linear projection.

class tace.models._e3nn.node.TensorNodeEmbedding(num_elements: int, num_radial_basis: int, num_channel: int, Lmax: int, lmax: int, avg_num_neighbors: float, bias: bool = False, so2_angular_basis: SO3Rotation | None = None)

Bases: NodeEmbedding

class tace.models._e3nn.node.SO2TensorNodeEmbedding(num_elements: int, num_radial_basis: int, num_channel: int, Lmax: int, lmax: int, avg_num_neighbors: float, bias: bool = False, so2_angular_basis: SO3Rotation | None = None)

Bases: NodeEmbedding