WIP: MACE foundation-model integration (small / medium / mpa-0 / matpes parity)

[PythonFZ]

Warning

Claude Code against mace-torch and mace-jax. I have not written this code!

Summary

Integrates torch-mace foundation models into apax via a new apax convert-mace path. End-to-end parity (within rtol 1e-4 energy / rtol 1e-3 force) on:

• MACE-MP-0 small (Residual, scalar-only)
• MACE-MP-0 medium (Residual, multi-irrep)
• MACE-MPA-0 medium (Density variants + ZBL + AgnesiTransform)
• MACE-matpes-r2scan-omat-ft (same architecture as mpa-0)

55 commits ahead of main. Marked draft / WIP for review while the broader feature branch settles; the most recent commit (34188618) closes the mpa-0 / matpes parity gap.

What's in the latest commit

Two design specs landed in tandem:

• docs/superpowers/specs/2026-05-04-mace-density-zbl-design.md
• docs/superpowers/specs/2026-05-04-mace-distance-transform.md

Architecture

• _interaction_scaffold helper + three interaction blocks (Residual / Density / DensityResidual) with per-layer dispatch in MaceRepresentation.
• MaceZBLPairRepulsion empirical correction (faithful port of mace.modules.radial.ZBLBasis); new output_scale field reproduces torch's "ZBL inside scale_shift" semantics without mutating buffers.
• AgnesiTransform Linen module + MaceRadialEmbedding refactor that threads Z/idx through and applies the transform between cutoff and bessel.
• MaceModelConfig.interaction_cls widened to a Literal union (single str OR per-layer list); new distance_transform discriminated union; MaceZBLPairRepulsion correction config.

Converter

• Per-layer interaction-class detection; emits a single str when uniform, list otherwise.
• Density density_fn weight scatter; size-aware _scatter_o3_linear_blocks for non-uniform layer-1 linear shapes; per-path skip_tp scatter.
• _map_pair_repulsion and _map_distance_transform helpers respect both params (trainable) and buffers (fixed) collections.
• _map_scale_shift collapses 2-D atomic_energies tables (matpes ships (1, 89) even single-head).

Bug fixes worth calling out

• check_for_ensemble crashed on 0-d scalar buffers — treats ndim==0 as size 1 (preserves ensemble semantics; jnp.stack lifts every leaf to >=1-D).
• AgnesiTransform clips r and r_0 to 0.02 so masked / padding edges with dr=0 don't trigger inf gradients via x^(-3.66).
• Periodic SiO₂ test fixture bumped to [14, 14, 14] to avoid an unrelated apax PBC bug (NaN forces from self-image edges) that exists for every foundation; documented in fixture docstring.

Test plan

• uv run pytest tests/unit_tests/ --ignore=tests/unit_tests/md — 93 pass, 0 fail.
• uv run pytest tests/integration_tests/mace -m mace_parity — 17 pass (small / medium / mpa-0 water; matpes periodic SiO₂; ZBL dimer parity at rtol 1e-12).
• Reviewer: confirm new ZBL output_scale semantics make sense for fresh-trained apax models (default 1.0 = no scaling).
• Reviewer: cross-check _map_scale_shift 2-D atomic_energies handling vs. multi-head matpes when an explicit head is passed.
• Address the orthogonal apax PBC bug (NaN forces with self-image edges in tight cells) in a follow-up.

Out of scope

• RealAgnostic (non-residual non-density) interaction variant — not used by foundations in scope; additive when a model needs it.
• SoftTransform (the tanh-based distance transform) — additive one-class extension.
• Multi-head foundation training. Single-head loading (matpes / mpa-0) works.
• Fresh end-to-end training of a Density-variant model — covered by P4 work, not parity-tested here.

🤖 Generated with Claude Code