TensorMACE¶

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

model:
   +TensorMACE:
      elements: [Si, O]
      channels: 32
      hidden_irreps: "0e+1o+2e+3o"
      props: tensor
      target_method: "tensor_product"
      target_tensor_irreps: "0e+1e+2e"
      irrep_tp: "3o"
      number_of_tps: 8

Definition¶

class graph_pes.models.TensorMACE( elements, cutoff=5.0, n_radial=8, radial_expansion='Bessel', weights_mlp={'activation': 'SiLU', 'hidden_depth': 3, 'hidden_features': 64}, channels=128, hidden_irreps='0e + 1o', l_max=3, layers=2, correlation=3, aggregation='constant_fixed', self_connection=True, readout_width=16, props='tensor', target_method='tensor_product', number_of_tps=None, target_tensor_irreps=None, irrep_tp='1o', )[source]¶

The MACE architecture, targeting arbitrary rank<https://arziv.org/abs/2412.15063> atomic tensors

One-hot encodings of the atomic numbers are used to condition the TensorProduct update in the residual connection of the message passing layers, as well as the contractions in the message passing layers.

Following the notation used in ACEsuite/mace, the first layer in this model is a RealAgnosticInteractionBlock. Subsequent layers are then RealAgnosticResidualInteractionBlocks

Please cite the following if you use this model in your research:

@misc{Batatia2022MACE,
    title = {
        MACE: Higher Order Equivariant Message Passing
        Neural Networks for Fast and Accurate Force Fields
    },
    author = {
        Batatia, Ilyes and Kov{\'a}cs, D{\'a}vid P{\'e}ter and
        Simm, Gregor N. C. and Ortner, Christoph and Cs{\'a}nyi, G{\'a}bor
    },
    year = {2022},
    doi = {10.48550/arXiv.2206.07697},
}

@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},
}

Parameters:

elements (list[str]) – list of elements that this MACE model will be able to handle.
cutoff (float) – radial cutoff (in Å) for the radial expansion (and message passing)
n_radial (int) – number of bases to expand the radial distances into
radial_expansion (type[DistanceExpansion] | str) – type of radial expansion to use. See DistanceExpansion for available options
weights_mlp (MLPConfig) – configuration for the MLPs that map the radial basis functions to the weights of the interactions’ tensor products
channels (int) – the multiplicity of the node features corresponding to each irrep specified in hidden_irreps
hidden_irreps (str | list[str]) – string representations of the e3nn.o3.Irreps to use for representing the node features between each message passing layer
l_max (int) – the highest order to consider in: * the spherical harmonics expansion of the neighbour vectors * the irreps of node features used within each message passing layer
layers (int) – number of message passing layers
correlation (int) – maximum correlation (body-order) of the messages
aggregation (NeighbourAggregationMode) – the type of aggregation to use when creating total messages from neigbour messages \(m_{j \rightarrow i}\)
self_connection (bool) – whether to use self-connections in the message passing layers
readout_width (int) – the width of the MLP used to read out the per-atom energies after the final message passing layer
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

Basic usage:

>>> from graph_pes.models import MACE
>>> model = TensorMACE(
...     elements=["H", "C", "N", "O"],
...     cutoff=5.0,
...     channels=16,
...     radial_expansion="Bessel",
... )

Specification in a YAML file:

model:
    +TensorMACE:
        elements: [H, C, N, O]
        cutoff: 5.0
        radial_expansion: Bessel
        target_method: tensor_product
        target_tensor_irreps: 0e + 1e + 2e
        number_of_tps: 128
        irrep_tp: 3o

        # change from the default MLP config:
        weights_mlp:
            hidden_depth: 2
            hidden_features: 16
            activation: SiLU

class graph_pes.models.ZEmbeddingTensorMACE( z_embed_dim=4, cutoff=5.0, n_radial=8, radial_expansion='Bessel', weights_mlp={'activation': 'SiLU', 'hidden_depth': 3, 'hidden_features': 64}, channels=128, hidden_irreps='0e + 1o', l_max=3, layers=2, correlation=3, aggregation='constant_fixed', self_connection=True, readout_width=16, props='tensor', target_method='tensor_product', number_of_tps=None, target_tensor_irreps=None, irrep_tp='1o', )[source]¶

A variant of TensorMACE that uses a fixed-size (z_embed_dim) per-element embedding of the atomic numbers to condition the TensorProduct update in the residual connection of the message passing layers, as well as the contractions in the message passing layers.

Please cite the following if you use this model in your research:

@misc{Batatia2022MACE,
    title = {
        MACE: Higher Order Equivariant Message Passing
        Neural Networks for Fast and Accurate Force Fields
    },
    author = {
        Batatia, Ilyes and Kov{'a}cs, D{'a}vid P{'e}ter and
        Simm, Gregor N. C. and Ortner, Christoph and Cs{'a}nyi, G{'a}bor
    },
    year = {2022},
    doi = {10.48550/arXiv.2206.07697},
}
@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},
}

All paramters are identical to MACE, except for the following:

elements is not required or used here
z_embed_dim controls size of the per-element embedding