MACEPOL-EF-OFF

MACE-POLAR + External Field — Organic Force Field

A fork of MACE and MACE-POLAR-1 targeting an external-field-responsive organic force field.

What this fork adds

• PolarHead — independent E_polar branch on top of a frozen PolarMACE foundation. Physics-aware multipole expansion of the polarisation energy:

  E_polar(E) = m_l0          (0th)
             + <m_l1, E>     (1st: dipole · field)
             + m_l0 · |E|²   (2nd isotropic: trace α)
             + <m_l2, Y_2(E)> (2nd anisotropic: traceless α)
             → MLP → E_polar_i
  

  The output Linear is zero-initialised so E_polar ≡ 0 at warm start (bit-wise equivalent to the frozen backbone). Enabled by default.
• Per-atom external field [N, 3] (upstream supports only [num_graphs, 3] uniform fields).
• Split E_app formula — q couples to the per-graph mean field E_uniform, μ couples to the per-atom field E_i:

  E_app = Σ_i q_i (E_uniform · R_i)  +  Σ_i μ_i · E_i
  

  Restores translation invariance for neutral molecules under MLMM-style spatially varying embedding fields.
• Sign convention unified — external_field is the physical field E_phys. Internal blocks expect E_input = -E_phys; both MACECalculator and HDF5Dataset (training) negate at load via AtomicData.from_config(negate_external_field=True) (default).
• TorchScript export — scripts/convert_polar_to_pt.py produces libtorch-loadable .pt for downstream C++ MLMM (sander).

Installation

conda create -n macepol python=3.9 -y && conda activate macepol
pip install torch --index-url https://download.pytorch.org/whl/cu118
cd mace-polar && pip install -e .
pip install git+https://github.com/WillBaldwin0/graph_electrostatics.git

Pins: e3nn==0.4.4, PyTorch>=2.2, ase. See pyproject.toml.

Quick start

import os
os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"

from ase import Atoms
from mace.calculators import MACECalculator

WEIGHTS = os.environ.get("MACE_POLAR_WEIGHTS", "../mace-polar-mlip")
calc = MACECalculator(
    model_paths=f"{WEIGHTS}/MACE-POLAR-1-S.model",
    device="cpu", default_dtype="float64",
)

water = Atoms(
    "OH2",
    positions=[[0.0, 0.0,  0.1173],
               [0.0, 0.7572, -0.4692],
               [0.0,-0.7572, -0.4692]],
)
water.info["charge"] = 0
water.info["spin"] = 0
water.info["external_field"] = [0.0, 0.0, 0.1]    # V/Å, physical (E_phys)
# Or per-atom:  [[Ex, Ey, Ez], ...] of shape (N, 3)
water.calc = calc

energy = water.get_potential_energy()              # eV
forces = water.get_forces()                        # eV/Å
charges = calc.results["charges"]                  # (N,)

MACECalculator.check_state() skips atoms.info ndarrays — changing external_field alone won't trigger recalculation. Workaround: calc.results = {} before each call.

Reference checkpoint

Model	Notes
MACE-POLAR-1-{S,M,L}.model	upstream OMol25 baselines (Batatia et al., 2026, ωB97M-V)
macepol-ef	this fork's production checkpoint (PolarHead-equipped, fine-tuned for external-field response) — recommended for MLMM/sander

Foundation weights live in the sibling ../mace-polar-mlip/ directory (not in git); override via MACE_POLAR_WEIGHTS=/path/to/weights.

E_app formula — gauge invariance

Under MLMM-style per-atom field E_i, the legacy formula Σ_i E_i · (q_i R_i + μ_i) breaks translation invariance: shifting a neutral molecule by d leaves a residual d · Σ_i E_i q_i ≠ 0 because Σ_i E_i q_i need not vanish when fields differ per atom. This produces origin-dependent forces in periodic MLMM (e.g. sander vlimit violations).

The split formula puts q onto the per-graph mean E_uniform:

E_app = E_uniform · Σ_i q_i R_i  +  Σ_i μ_i · E_i
      = E_uniform · D_charge      +  Σ_i μ_i · E_i

• Neutral molecules → translation invariant (Σ q = 0 ⇒ d·Σq term vanishes).
• Charged molecules → ΔE = Q · E_uniform · d on translation, which is the correct work done by the field on a net charge (physical, not bug). Forces are well-defined and translation-invariant in both cases.
• Uniform field → reduces to textbook −p · E, identical to the legacy formula. So the split is a no-op for upstream training data and a stability fix for downstream MLMM inference.

Training & Fine-tuning

# From scratch
python scripts/run_train.py \
    --name=PolarMACE_run \
    --train_file=train.xyz \
    --model=PolarMACE \
    --r_max=5.0 --device=cuda

# Fine-tune from a foundation checkpoint
python scripts/run_train.py \
    --name=ft_run \
    --foundation_model=../mace-polar-mlip/MACE-POLAR-1-S.model \
    --train_file=train.xyz --valid_file=valid.xyz \
    --model=PolarMACE --device=cuda

Optional WandB: pip install wandb && export WANDB_PROJECT=macepol-ft.

TorchScript export

python scripts/convert_polar_to_pt.py                    # all S/M/L
python scripts/convert_polar_to_pt.py --model path/to/X.model

Verifies energy / charge / force diff vs eager model = 0 on a water test geometry; bails out otherwise.

Differences vs upstream MACE

Aspect	upstream	this fork
External-field shape	[num_graphs, 3] uniform	[N, 3] per-atom
E_app formula	Σ E_i · (q_i R_i + μ_i)	split: q→E_uniform, μ→E_i
external_field sign	E_input = -E_phys (implicit)	E_phys (negated at load)
PolarHead (E_polar)	❌	✅ default on
TorchScript	partial	scatter / irreps_tools made TS-safe

Development

bash ./scripts/run_checks.sh                        # black + isort + pylint
python -m pytest tests/test_polar_models.py -v      # polar tests

Citation

@inproceedings{batatia2022mace,
  title={{MACE}: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields},
  author={Ilyes Batatia and David P Kovacs and Gregor N. C. Simm and Christoph Ortner and Gabor Csanyi},
  booktitle={Advances in Neural Information Processing Systems},
  year={2022},
}

@misc{batatia2026macepolar,
  title={MACE-POLAR-1: long-range electrostatics and field response in universal ML potentials},
  author={Ilyes Batatia et al.},
  year={2026},
  eprint={2602.19411},
  archivePrefix={arXiv},
}

License

MIT (same as upstream). See LICENSE.md.