Surface-Aware MACE (Fork of ACEsuit/mace)

This repository is a fork of the official MACE project:

• Upstream code: https://github.com/ACEsuit/mace
• Upstream docs: https://mace-docs.readthedocs.io
• Upstream paper (MACE): https://openreview.net/forum?id=YPpSngE-ZU

This fork focuses on surface-aware design for slab/catalysis systems (SlabNet idea), while preserving energy-force consistency.

What This Fork Adds

Compared with upstream MACE, this fork adds:

• Surface descriptors in the model:
  • local surface normal n_i
  • local surface degree u_i
• Surface-aware anisotropic message gating
• Adaptive cutoff that expands neighborhood for undercoordinated atoms
• Surface-related training losses (vacuum / surface energy / NEB barrier terms)
• CLI/config options for all the above

Core implementation files:

• mace/modules/models.py
• mace/modules/loss.py
• mace/tools/arg_parser.py
• mace/cli/run_train.py
• mace/tools/model_script_utils.py
• mace/tools/scripts_utils.py

Surface-Aware Design (SlabNet Core Idea)

The design follows four coupled modifications.

1) Surface Descriptors in Embedding

For each atom, build differentiable local descriptors from neighbor geometry:

• Surface normal n_i: from weighted local covariance; use the smallest-eigenvalue eigenvector.
• Surface degree u_i in [0,1]: quantifies undercoordination/exposure.

A smooth radial kernel is used:

w_ij = exp(-(r_ij / r0)^p)

This keeps descriptors continuous w.r.t. coordinates and compatible with autodiff.

2) Anisotropic Message Gating

Edge features are modulated by surface exposure and orientation:

• exposure term via u_i, u_j
• anisotropy term via c_ij = uhat_ij · n_i (used in sign-invariant form, e.g. c_ij^2)

Conceptually:

e_ij -> e_ij * (1 + lambda_u * u_i u_j + lambda_c * c_ij^2)

This introduces slab-sensitive directional effects while preserving rotational behavior.

3) Adaptive Receptive Field

Instead of fixed cutoff for all atoms, use exposure-aware cutoff:

r_cut(i,j) = r0 * (1 + alpha_cut * min(u_i, u_j))

Bulk-like regions keep base cutoff; undercoordinated surface regions get larger context.

4) Physics-Enhanced Surface Losses

Total objective extends base energy/force loss with optional surface regularizers:

• Vacuum invariance term L_vac
• Surface energy consistency term L_gamma
• NEB barrier anchoring term L_NEB

L_total = L_base + beta_vac L_vac + beta_gamma L_gamma + beta_NEB L_NEB

This helps align slab energetics and reaction barriers.

Training Controls (Surface-Aware)

Main options exposed in training:

• --surface_aware
• --surface_kernel_r0
• --surface_kernel_p
• --surface_coord_ref
• --surface_ui_alpha
• --surface_lambda_ui
• --surface_lambda_c2
• --surface_adaptive_cutoff_alpha
• --surface_graph_cutoff_scale
• --beta_vac
• --beta_gamma
• --beta_neb

Latest Checkpoint in This Repo

Included checkpoint:

• checkpoints/oc22_mh1_surface_stage2_run-42.model

This checkpoint is trained with the surface-aware fork using staged OC22 S2EF training.

Quick Evaluation Example

mace_eval_configs \
  --configs your_configs.xyz \
  --model checkpoints/oc22_mh1_surface_stage2_run-42.model \
  --output predictions.xyz

Notes

• For general MACE usage, installation, and upstream features, please use upstream docs/code links above.
• This README intentionally focuses on the fork-specific surface-aware design.