TensorNequIP¶

Train this architecture on your own data using the graph-pes-train CLI, using e.g. the following config:

model:
   +TensorNequIP:
      elements: [Si, O]
      features:
         channels: [64, 64, 64]
         l_max: 2
         use_odd_parity: true
      props: tensor
      target_tensor_irreps: 0e + 1e + 2e
      prune_last_layer: false
      target_method: "direct"

Definition¶

class graph_pes.models.TensorNequIP( elements, cutoff=5.0, layers=3, features={'channels': 16, 'l_max': 2, 'use_odd_parity': True}, self_interaction='tensor_product', prune_last_layer=True, neighbour_aggregation='sum', radial_features=8, props='tensor', target_method='direct', number_of_tps=None, target_tensor_irreps=None, irrep_tp='1o', )[source]¶

NequIP architecture from E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials targeting arbitrary rank<https://arziv.org/abs/2412.15063> atomic tensors

If you use this model in your research, please cite the original work:

@article{Batzner-22-05,
    title = {
              E(3)-Equivariant Graph Neural Networks for
              Data-Efficient and Accurate Interatomic Potentials
            },
    author = {
              Batzner, Simon and Musaelian, Albert and Sun, Lixin
              and Geiger, Mario and Mailoa, Jonathan P.
              and Kornbluth, Mordechai and Molinari, Nicola
              and Smidt, Tess E. and Kozinsky, Boris
            },
    year = {2022},
    journal = {Nature Communications},
    volume = {13},
    number = {1},
    pages = {2453},
    doi = {10.1038/s41467-022-29939-5},
    copyright = {2022 The Author(s)}
}

@misc{BenMahmoud2025NMR,
    title = {
        Graph-neural-network predictions of solid-state NMR parameters
          in silica from spherical tensor decomposition
        },
    author = {Ben Mahmoud, Chiheb and Rosset, Louise and Yates, Jonatha and
              Deringer, Volker
        },
    year = {2025},
    doi = {10.1063/5.0274240},
}

Parameters:

elements (list[str]) – The elements that the model will encounter in the training data. This is used to create the atomic one-hot embedding. If you intend to fine-tune this model on additional data, you must ensure that all elements you will encounter in both pre-training and fine-tuning are present in this list. (See ZEmbeddingNequIP for an alternative that allows for arbitrary atomic numbers.)
direct_force_predictions – Whether to predict forces directly. If True, the model will output forces (rather than infer them from the energy) using a LinearReadOut to map the final layer node embedding to a set of force predictions.
cutoff (float) – The cutoff radius for the model.
features (SimpleIrrepSpec | CompleteIrrepSpec) – A specification of the irreps to use for the node and edge embeddings. Can be either a SimpleIrrepSpec or a CompleteIrrepSpec.
layers (int) – The number of layers for the message passing.
self_interaction (Literal['linear', 'tensor_product'] | None) – The kind of self-interaction to use. If None, no self-interaction is applied. If "linear", a linear self-interaction is applied to update the node embeddings along the residual path. If "tensor_product", a tensor product combining the old node embedding with an embedding of the atomic number is applied. As first noticed by the authors of SevenNet, using linear self-interactions greatly reduces the number of parameters in the model, and helps to prevent overfitting.
prune_last_layer (bool) – Whether to prune irrep communication pathways in the final layer that do not contribute to the 0e output embedding.
neighbour_aggregation (NeighbourAggregationMode) – The neighbour aggregation mode. See NeighbourAggregationMode for more details. Note that "mean" or "sqrt" aggregations lead to un-physical discontinuities in the energy function as atoms enter and leave the cutoff radius of the model.
radial_features (int) – The number of features to expand the radial distances into. These features are then passed through an MLP to generate distance-conditioned weights for the message tensor product.
props (str) – the property targeted by the model, set to “tensor”
target_method (str) – determine how to reconstruct the target tensor, either by tensor product if speherical tensor contains “0o”, “1e”, “2o”,.. etc, or direct otherwise
number_of_tps – the number of tensors involved in the tensor product
target_tensor_irreps – the irreps of the target tensor
irrep_tp – the irrep of the tensors involved in the tensor product to reconstruct the target

Examples

Configure a TensorNequIP model for use with graph-pes-train:

model:
  +TensorNequIP:
    elements: [C, H, O]
    cutoff: 5.0

    # use 2 message passing layers
    layers: 2

    # using SimpleIrrepSpec
    features:
      channels: [64, 32, 8]
      l_max: 2
      use_odd_parity: true

    # scale the aggregation by the avg. number of
    # neighbours in the training set
    neighbour_aggregation: constant_fixed

    target_method: tensor_product
    target_tensor_irreps: 0e + 1e + 2e
    number_of_tps: 128
    irrep_tp: 3o

class graph_pes.models.ZEmbeddingTensorNequIP( cutoff=5.0, Z_embed_dim=8, features={'channels': 16, 'l_max': 2, 'use_odd_parity': True}, layers=3, self_interaction='tensor_product', prune_last_layer=True, neighbour_aggregation='sum', radial_features=8, props='tensor', target_method='direct', number_of_tps=None, target_tensor_irreps=None, irrep_tp='1o', )[source]¶

A modified version of the TensorNequIP architecture that embeds atomic numbers into a learnable embedding rather than using an atomic one-hot encoding.

This circumvents the need to know in advance all the elements that the model will encounter in any pre-training or fine-tuning datasets.

Relevant differences from TensorNequIP:

The elements argument is removed.
The Z_embed_dim (int) argument controls the size of the atomic number embedding (default: 8).

For all other options, see TensorNequIP.

If you use this model in your research, please cite the original work:

@article{Batzner-22-05,
    title = {
              E(3)-Equivariant Graph Neural Networks for
              Data-Efficient and Accurate Interatomic Potentials
            },
    author = {
              Batzner, Simon and Musaelian, Albert and Sun, Lixin
              and Geiger, Mario and Mailoa, Jonathan P.
              and Kornbluth, Mordechai and Molinari, Nicola
              and Smidt, Tess E. and Kozinsky, Boris
            },
    year = {2022},
    journal = {Nature Communications},
    volume = {13},
    number = {1},
    pages = {2453},
    doi = {10.1038/s41467-022-29939-5},
    copyright = {2022 The Author(s)}
}

@misc{BenMahmoud2025NMR,
    title = {
        Graph-neural-network predictions of solid-state NMR parameters
          in silica from spherical tensor decomposition
        },
    author = {Ben Mahmoud, Chiheb and Rosset, Louise and
              Yates, Jonathan and Deringer, Volker
        },
    year = {2025},
    doi = {10.1063/5.0274240},
}