Engine Models

This guide covers the backend models available in IANN, including their architectures, features, and use cases.

Overview

IANN provides several state-of-the-art foundation models for interatomic potentials. All of them are graph-based equivariant neural networks:

  • FastPot (Fast Potential with high-order tensor features and equivariant message passing)

  • PaiNN (Polarizable atom interaction Neural Network)

  • NequIP (Neural equivariant Interatomic Potentials)

  • MACE (Message-passing Atomic Cluster Expansion)

  • EquiformerV2 (Equivariant Transformer V2)

Each model has its own strengths and is suitable for different applications.

Note

NVIDIA Python library cuEquivariance is an optional dependency for the NequIP and MACE models. If you want to use the optimized operations, you need to install it, and the models will automatically use it. Please check the official website cuEquivariance for more details.

FastPot

FastPot is a high-performance neural network model that combines the power of high-order tensor features and equivariant message passing for fast and accurate potential energy surface prediction.

Key features:

  • Very high computational efficiency

  • Good balance of accuracy and speed

  • Suitable for general-purpose applications

Example usage:

from iann.models.fastpot import FastPot

model = FastPot(
    num_layers=3,
    num_channels=128,
    cutoff=5.5,
    compute_forces=True
)

There are more adjustable parameters for FastPot model to setup the training process, please check the source code iann.models.fastpot.FastPot for more details (all adjustable parameters are passed as kwargs.get in the model class).

PaiNN

PaiNN is a message-passing neural network that considers features including:

  • Scalar: atomic number, distance

  • Vector: coordinate difference

Key features:

  • High computational efficiency

  • Good balance of accuracy and speed

  • Suitable for general-purpose applications

Example usage:

from iann.models.painn import PaiNN

model = PaiNN(
    num_layers=3,
    num_channels=128,
    cutoff=5.5,
    compute_forces=True
)

There are more adjustable parameters for PaiNN model to setup the training process, please check the source code iann.models.painn.PaiNN for more details (all adjustable parameters are passed as kwargs.get in the model class).

NequIP

NequIP is an equivariant neural network that considers features including:

  • Scalar: atomic number, distance

  • Vector: coordinate difference

  • Higher-order tensor: high order rotation

Key features:

  • Excellent accuracy

  • Uses spherical harmonics symmetries

  • Good for high-precision applications

Example usage:

from iann.models.nequip import NequIP

model = NequIP(
    num_layers=3,
    num_channels=64,
    cutoff=5.5,
    compute_forces=True
)

There are more adjustable parameters for NequIP model to setup the training process, please check the source code iann.models.nequip.NequIP for more details (all adjustable parameters are passed as kwargs.get in the model class).

MACE

MACE combines message-passing architecture and multi-body expansion, which considers features including:

  • Scalar: atomic number, distance

  • Vector: coordinate difference

  • Higher-order tensor: high order rotation

  • multi-body expansion

Key features:

  • Fast training and inference

  • Good scaling properties

  • Suitable for large-scale applications

Example usage:

from iann.models.mace import MACE

model = MACE(
    num_layers=3,
    num_channels=64,
    cutoff=5.5,
    compute_forces=True
)

There are more adjustable parameters for MACE model to setup the training process, please check the source code iann.models.mace.MACE for more details (all adjustable parameters are passed as kwargs.get in the model class).

EquiformerV2

EquiformerV2 is a transformer-based model that:

  • Uses attention mechanisms

  • Preserves physical symmetries

Key features:

  • State-of-the-art accuracy

  • Good for complex systems

Example usage:

from iann.models.equiformerV2 import EquiformerV2

model = EquiformerV2(
    num_layers=3,
    num_channels=32,
    cutoff=5.5,
    compute_forces=True
)

There are more adjustable parameters for EquiformerV2 model to setup the training process, please check the source code iann.models.equiformerV2.EquiformerV2 for more details (all adjustable parameters are passed as kwargs.get in the model class).

Model Selection

When choosing a model, consider:

  1. Accuracy Requirements

    • NequIP, MACE and EquiformerV2 for highest accuracy

    • FastPot, PaiNN for balanced performance

  2. Computational Resources

    • FastPot and PaiNN for fast training/inference

    • NequIP, MACE and EquiformerV2 for most complex systems

  3. System Size

    • FastPot, PaiNN for super large systems

    • NequIP, MACE for large systems

    • EquiformerV2 is not the best choice for small systems

Note

The description may be not reliable and the reliable data will be provided in future.

For detailed API documentation of each model, see the API Reference reference.