MACEPack

MACEPack is an extension package for SchNetPack that provides MACE-based equivariant representations within the modular and highly customizable SchNetPack framework.

The main motivation behind this project is to bring the accuracy of MACE representations to users who want to retain the flexibility of SchNetPack for building custom models, training workflows, and prediction heads. Rather than being a standalone machine learning framework, MACEPack is designed to integrate seamlessly SchNetPack.

The package focuses on providing reusable representation modules that can be combined with SchNetPack's data handling, training infrastructure, and task-specific outputs, enabling users to build models for a wide range of molecular machine learning applications.

Installation

Clone the repository:

git clone https://github.com/smausenberger/MACEPack
cd MACEPack

Install the package:

pip install .

If you use MACEPack, you need compatible versions of MACE and SchNetPack. PyTorch, e3nn, ASE, and other required packages are installed as dependencies of those packages.

Usage

MACEPack provides SchNetPack-compatible representation modules based on MACE. These representations construct atom-wise scalar and equivariant features from molecular structures and expose them through the SchNetPack model interface, allowing them to be combined with task-specific prediction heads and training workflows.

Example usage from Python:

from mace.modules import PolynomialCutoff
from macepack.representation import MACERepresentation
from schnetpack.nn.radial import BesselRBF

cutoff = 5.0

representation = MACERepresentation(
    atomic_numbers=[1, 6, 7, 8],
    n_atom_basis=128,
    max_L=2,
    num_interactions=2,
    correlation=3,
    max_ell=2,
    avg_num_neighbors=20.0,
    radial_mlp=[64, 64, 64],
    radial_basis=BesselRBF(
        cutoff=cutoff,
        n_rbf=15,
    ),
    cutoff_fn=PolynomialCutoff(
        r_max=cutoff,
        p=8,
    ),
    scalar_key="scalar_representation",
    vector_key="vector_representation",
    full_key="mace_representation",
    eps=1.0e-8,
)

The exact constructor arguments depend on the representation variant used.

Important: Currently, atomic_numbers and avg_num_neighbors must be provided explicitly when constructing a MACERepresentation. See the note below for details and the rationale behind this design choice.

At the moment, atomic_numbers and avg_num_neighbors must be set manually. Both values are available from the dataset, but passing them automatically to the representation would currently require changes to SchNetPack or MACE internals. Since MACEPack aims to stay compatible with the upstream projects without custom patches, this has not been implemented yet. Future versions may provide functionality to include these values automatically.

Example Configuration

Below is a minimal example configuration showing the intended structure. The names and parameters may need to be adapted to the exact modules used in a given project.

globals:
  cutoff: 5.0

model:
  representation:
    _target_: macepack.representation.MACERepresentation
    atomic_numbers: [1, 6, 7, 8]
    n_atom_basis: 128
    max_L: 2
    num_interactions: 2
    correlation: 3
    max_ell: 2
    avg_num_neighbors: 20.0
    radial_mlp: [64, 64, 64]
    radial_basis:
      _target_: schnetpack.nn.radial.BesselRBF
      cutoff: ${globals.cutoff}
      n_rbf: 15
    cutoff_fn:
      _target_: mace.modules.PolynomialCutoff
      r_max: ${globals.cutoff}
      p: 8
    scalar_key: scalar_representation
    vector_key: vector_representation
    full_key: mace_representation
    eps: 1.0e-8

A MACE representation can also be configured directly from a SchNetPack YAML configuration file (see example above):

spktrain experiment=qm_atomwise run.data_dir=<path> model/representation=mace

This uses the predefined MACEPack representation configuration and integrates it into the standard SchNetPack training workflow.

Q&A

What is MACEPack?

MACEPack is an extension package for SchNetPack that provides MACE-based equivariant representations within the modular and highly customizable SchNetPack framework.

Who should use MACEPack?

MACEPack is intended for developers of atomistic machine learning models who want full control over the model architecture and training workflow. It is especially useful for experimenting with custom changes without digging through a large and convoluted code base.

Is MACEPack identical to MACE?

No. MACEPack tries to reproduce the MACE representation based on the original MACE package, but it is not a one-to-one copy of MACE. The implementation is adapted for use as a SchNetPack-compatible representation module. As a result, numerical results may vary slightly from the original MACE implementation.

Is this a replacement for MACE?

No. MACEPack is not intended to replace the original MACE package. It reuses or adapts MACE-style ideas and components for workflows where a separate representation module is needed.

Is this a replacement for SchNetPack?

No. The goal is to integrate with SchNetPack-style workflows, not to replace SchNetPack. MACEPack is intended to provide additional representation options.

Does MACEPack support periodic systems?

The current focus is on non-periodic molecular systems. Periodic systems are not the primary design target.

Can I use or fine-tune MACE foundation models?

No. MACE foundation models are not compatible with MACEPack models.

Can I use SPaiNN with MACEPack?

Yes. SPaiNN is fully compatible with MACEPack.

Can I include point charges?

Yes. MACEPack has a module for including point charges inspired by FieldMACE.

Is the package stable?

MACEPack was tested with MACE 0.3.16 and SchNetPack 2.2.0. It should also work with newer versions, as long as the relevant MACE and SchNetPack APIs remain compatible.