feat: add MACE interface by FelyCZ · Pull Request #224 · PHOTOX/ABIN

FelyCZ · 2026-03-07T22:41:28Z

Add MACE Interface via MPI

Description

This PR adds a new interface for the MACE (Machine Learning Atomic Cluster Expansion) potential to ABIN. MACE is integrated via an MPI-based client-server architecture: a Python server loads and evaluates the MACE ML model, while the Fortran side communicates coordinates, energies, and forces over MPI. This follows the same pattern used for the existing TeraChem (_tera_) interface.

The interface supports configurable model paths, device selection (CPU/CUDA), precision settings, and various MACE-specific options -- all controlled through a new &mace namelist in the ABIN input file.

File Changes

Fortran Core

File	Description
`src/mace_mpi_api.F90`	[NEW] Core MPI API module (`mod_mace_mpi`). Handles MPI port connection to the Python server, sends atom types/config/coordinates, and defines all MACE namelist variables (model path, device, dtype, batch size, etc.).
`src/force_mace.F90`	[NEW] Force driver module (`mod_force_mace`). Orchestrates per-bead force evaluation -- sends coordinates via `send_mace`, receives energy and forces via `receive_mace`. Supports OpenMP.
`src/forces.F90`	Added `_mace_` case to the `force_wrapper` potential selector, routing to `force_mace()`.
`src/init.F90`	Added `&mace` namelist reading and MACE interface/server initialization when `pot='_mace_'`.
`src/mpi_wrapper.F90`	Added `_mace_` to the MPI-required potential check so non-MPI builds raise a clear error.
`src/Makefile`	Added `mace_mpi_api.o` and `force_mace.o` to the build.

Python Server

File	Description
`interfaces/MACE/mace_server.py`	[NEW] MACE MPI server. Opens MPI port, accepts connection from ABIN, receives configuration and atom data, loads the MACE model, and returns energies/forces each step. Handles unit conversions (Bohr <-> Å, Hartree <-> eV) and supports foundation model auto-download.
`interfaces/MACE/README.md`	[NEW] Documentation covering requirements, manual/automatic launch instructions, and a full configuration reference table for all `&mace` namelist parameters. Includes MACE-OFF licensing notice.
`interfaces/MACE/MACE-OFF23_medium.model`	[NEW] Bundled pre-trained MACE-OFF23 medium model (~18 MB) for out-of-the-box usage. Licensing notice is included in the README.

Tests

File	Description
`tests/MACE/test.sh`	[NEW] Integration test script. Launches the mock MACE server and ABIN via MPI, runs a short MD simulation, and validates output against reference files.
`tests/MACE/mock_mace_server.py`	[NEW] Mock MACE server for testing. Implements the same MPI protocol as the real server but returns forces from a simple harmonic potential, avoiding the need for `mace-torch`/PyTorch.
`tests/MACE/input.in`	[NEW] Test input file with `pot='_mace_'` and `&mace` namelist.
`tests/MACE/mini.xyz`	[NEW] Test geometry (water molecule).
*`tests/MACE/.ref`**	[NEW] Reference output files (`energies.dat.ref`, `forces.xyz.ref`, `movie.xyz.ref`, `velocities.xyz.ref`, `temper.dat.ref`).
`tests/test.sh`	Registered `MACE` in the MPI test folder list.
`unit_tests/test_mace.pf`	[NEW] pFUnit tests verifying default namelist values, MPI tag constants, and non-MPI error handling.
`unit_tests/Makefile`	Added `mace` target to the unit test build, test runner, and clean targets.

Utilities & Docs

File	Description
`utils/run.mace_mpi_abin.sh`	[NEW] Launch script that starts both the MACE Python server and ABIN, connects them via MPI, and monitors both processes with automatic cleanup.
`dev_scripts/install_mace.sh`	[NEW] Convenience script to clone the MACE repo and install dependencies (`mace-torch`, `mpi4py`).
`sample_inputs/input.in.mace`	[NEW] Sample input file demonstrating `&mace` namelist configuration.
`README.md`	Added MACE to the list of optional libraries with a brief description and install reference.

modified: sample_inputs/input.in.mace modified: src/mace_mpi_api.F90

new file: interfaces/MACE/README.md

danielhollas · 2026-03-17T13:51:25Z

@FelyCZ this is super exciting, thanks! Is this PR ready for review?
cc @JanosJiri probably you could do first rounds of review in case you haven't already?

FelyCZ · 2026-03-17T14:12:40Z

Hi @danielhollas,
Yes, it's ready for review. I just don't have the necessary permissions to request this in the PR.

I have tested the code on a local machine (WSL Ubuntu 22.04.5 LTS) and it yielded no problems.

danielhollas · 2026-03-17T15:04:16Z

Perfect, thanks! I just gave you Triage permissions to the repo which should give you enough rights to ask for review and run tests on GitHub (I think). I'll leave it up to @JanosJiri regarding the level of permissions that you should have (perhaps you should be a member of the PHOTOX Dev Team which has write permissions to this repo?).

Thanks for adding unit tests! One of them is failing though, can you take a look?
https://github.com/PHOTOX/ABIN/actions/runs/22808780618/job/67410107594?pr=224

The build failure of GFortran 9 is unrelated to these changes so don't worry about it, I'll fix that separately before we merge this.

Can you also please run the autoformatter? https://github.com/PHOTOX/ABIN/blob/master/CONTRIBUTING.md#code-formatting

Once the format and test pass on GitHub please ping me for review, thanks!

FelyCZ · 2026-03-17T15:51:24Z

Thanks for the Triage role!
I think this is enough for me, I don't need to have write permissions. :)

I've fixed the unit test and run the autoformat script. The actions are passing now.

JanosJiri · 2026-03-17T16:03:02Z

cc @JanosJiri probably you could do first rounds of review in case you haven't already?

I could have a look, but I won't be able to get to it for quite some time. Sorry, I'm just overwhelmed with my thesis and finishing all papers before I leave the group. Also, I have very limited knowledge of MPI, so I won't be able to review it reliably (at least now).

Perfect, thanks! I just gave you Triage permissions to the repo, which should give you enough rights to ask for review and run tests on GitHub (I think). I'll leave it up to @JanosJiri regarding the level of permissions that you should have (perhaps you should be a member of the PHOTOX Dev Team which has write permissions to this repo?).

@FelyCZ, I can add you to the PHOTOX Dev Team](https://github.com/orgs/PHOTOX/teams/devs) if you want and if you plan to work on ABIN or anything else in the repo. There is no official policy for that.

danielhollas · 2026-03-17T16:09:01Z

@FelyCZ if you merge in master that should fix the gcc9 build.

danielhollas · 2026-03-17T16:41:55Z

Just a quick comment, I don't think we should be including the MACE-OFF model in the repo, it's a huge binary blob, and it will likely become obsolete soon as new models come out.
We can definitely provide instructions and link about where it can be downloaded though!

Is this file needed for tests? If so, we should download it as part of the MACE installation script perhaps.

I agree that it's not entirely necessary. The main idea was to provide users with a quick way to test the interface without any further configuration.

The file is not necessary for the tests; however, the default configuration (&mace namelist) uses the file as fallback.

Hi, @danielhollas,
Are you still of the opinion that we should remove the model from the repo?

Hi, yes I think we should remove it. Please add instructions somewhere (perhaps a README file) where we can tell the users how to download it.

Sorry for the slow review, this is a big PR and I want to have a careful look and we also plan to test it soon.

danielhollas

Did a very quick initial review, I hope these comments are useful.

danielhollas · 2026-03-30T14:20:10Z

+
+git clone "${REPO_URL}" "${REPO_DIR}" && cd "$REPO_DIR"
+
+pip install --upgrade pip


This will install the package into the global python environment (unless the user is already inside a venv or conda. I think we should create a fresh venv for the user.

danielhollas · 2026-03-30T14:20:44Z

+git clone "${REPO_URL}" "${REPO_DIR}" && cd "$REPO_DIR"
+
+pip install --upgrade pip
+pip install .


Why are we installing MACE from the repo? We should I think install the package from PyPI.

danielhollas · 2026-03-30T14:22:47Z

+
+pip install --upgrade pip
+pip install .
+pip install mpi4py


To be clear, I think we should do something like this:

python -m venv .mace_venv .mace_venv/bin/python3 -m pip install mace-torch mpi4py

danielhollas · 2026-03-30T14:25:24Z

+        self.device = str(torch_tools.init_device(device))
+
+        import os
+        if os.path.isfile(model_path):


This interface is a bit confusing. Should we instead have two diferrent (mutually exclusive) parameters, model_path and model_name?

danielhollas · 2026-03-30T14:26:33Z

+        if os.path.isfile(model_path):
+            # Load from local file
+            try:
+                self.model = torch.jit.load(f=model_path, map_location=device).to(device)


I am a little confused as to why we're importing the model directly to torch, and not using MACE specific functions / interface for it?

We didn't use the default pre-trained model; we used a custom, fine-tuned model based on MACE-OFF-23. However, the MACE functions do not support the use of custom models (afaik). Therefore the interface is based on eval_configs.py file.

Can you say a bit more about your fine-tuned model? Is it still in a *.model format?
The mace.calculators.MACECalculator definitely supports arbitrary models, not just the pre-trained ones. See for example the code here:

https://github.com/duartegroup/mlp-train/blob/64d2b0010d40e63b733febfb57679ad2f2e90601/mlptrain/potentials/mace/mace.py#L242

danielhollas · 2026-03-30T14:37:45Z

+    port_name = MPI.Open_port()
+    log(f"MPI port opened: {port_name}")
+
+    port_file = "mace_port.txt.1"


Let's define this as a "constant" at the top of the script, something like:

MACE_PORT_FILE = "mace_port.txt.1"

danielhollas · 2026-03-30T14:40:50Z

+When using OpenMPI, you need an `ompi-server` running for the connection handshake:
+
+```console
+# 1. Start ompi-server


Is this really the case? We don't seem to need something like this in the TeraChem interface?

danielhollas · 2026-03-30T14:43:50Z

+
+## Configuration
+
+MACE settings are controlled via the `&mace` namelist in the ABIN input file (`input.in`). Below is the full list of available parameters:


So I am not totally opposed to this, but note that generally we don't put interface-specific parameters in the abin input file since ABIN doesn't really need to know about these.

I feel that in this case it would be better (and more general!) to hava a YAML configuration file in the interface folder which we can than load directly into MACE (I believe MACE supports this natively).

danielhollas · 2026-03-30T14:46:15Z


      namelist /qmmm/ natqm, natmm, q, LJ_rmin, LJ_eps, mm_types

+      namelist /mace/ mace_model, mace_device, mace_default_dtype, &


Yeah, I don't really think we should be parsing MACE specific keywords here, they are irrelevant from ABIN's point of view (and note that since MACE excosystem is still evolving a lot, we need to have the flexibility to add new options easily, without having to add them here and recompile ABIN

danielhollas · 2026-03-30T14:50:51Z

+      end if
+   end subroutine finalize_mace
+
+   subroutine wait_for_mace(comm)


This code was added to TeraChem interface because a typing TC calculation takes (much) more than a 1 second so we could save a whole CPU by instructing ABIN to sleep a bit while waiting for TC.

In MACE case, the hope is that it is actually very quick, so I don't think this case makes sense. I'd just remove it for now to keep things simple? (happy to hear otherwise though!)

danielhollas · 2026-03-30T15:34:04Z

Thinking about this more, I don't think we should provide the install_mace.sh script, there's simply too many variables about how to handle the python environment (and how to handle Torch installation in particular).

I'd rather have a comprehensive installation section in the README. Given the (unfortunate) fact that the mace package is not published on conda, I think we should encourage users to use the uv package installer, in particular because it has an advanced handling of pytorch.

https://docs.astral.sh/uv/guides/integration/pytorch/

i.e. something like this

uv venv uv pip install mace-torch mpi4py --torch-backend=auto

To install CPU-only version

uv pip install mace-torch mpi4py --torch-backend=cpu

I am happy to help write this down before we merge this PR.

danielhollas · 2026-04-26T17:38:28Z

Hi @FelyCZ, just wanted to ask if you'll have time to work on this further in the near future?(No worries if not)

I am hoping to do some testing on this branch next week, as we will need this interface soon in our own MACE project.

FelyCZ · 2026-04-26T18:20:01Z

Hi @danielhollas ,

I am very sorry but I have so many things with deadlines piling up now. I probably won't be able to work on this at least till the end of May.

danielhollas · 2026-04-27T09:18:32Z

@FelyCZ no problem, thanks for letting me know!

I obviously don't want to take over your great work, but also since we need to press on with our project, I am wondering what would be best for now, I see couple of options:

We merge this PR more or less as is, to get it out of your TODO list and to have you have the full credit for the feature. I could then continue building on top of it.
If you'd like to go through the review yourself and have full ownership, I could in the meantime open a separate branch based on this one so I can do any necessary changes to unblock our work. And then once you have time we'd continue working together on this PR and merge it once it's fully polished.
Something else?

Let me know what would you prefer, I am fine either way. Also let me know if there's anybody else in Petr's group that would be interested in working on this (I know Jirka is busy, perhaps somebody else?)

FelyCZ added 13 commits February 22, 2026 18:56

feat: add MACE potential support with MPI interface

13fc591

fix: add test for MACE potential with MPI interface

547d35f

chore: cleanup

d37987d

fix: add MACE install script

b78376a

fix: add option for custom Python interpretter

02bd2dc

fix: MACE model loading

dd59254

fix: default MACE model path

789cdb7

fix: default MACE model path

8ca192a

modified: sample_inputs/input.in.mace modified: src/mace_mpi_api.F90

docs: add information about MACE to README

5f5828a

docs: add quick guide for MACE interface

1f8ee25

new file: interfaces/MACE/README.md

docs: add MACE licensing

c8dfc47

chore: cleanup

683d2ae

style: reformat & comment cleanup

d829cfc

FelyCZ added 2 commits March 17, 2026 16:12

style: autoformat

9c7bc19

fix: mace unit test default model path

b01d25d

FelyCZ requested a review from danielhollas March 17, 2026 15:51

danielhollas reviewed Mar 17, 2026

View reviewed changes

danielhollas requested changes Mar 30, 2026

View reviewed changes

danielhollas reviewed Mar 30, 2026

View reviewed changes


		git clone "${REPO_URL}" "${REPO_DIR}" && cd "$REPO_DIR"

		pip install --upgrade pip


		## Configuration

		MACE settings are controlled via the `&mace` namelist in the ABIN input file (`input.in`). Below is the full list of available parameters:


		namelist /qmmm/ natqm, natmm, q, LJ_rmin, LJ_eps, mm_types

		namelist /mace/ mace_model, mace_device, mace_default_dtype, &

Conversation

FelyCZ commented Mar 7, 2026

Add MACE Interface via MPI

Description

File Changes

Fortran Core

Python Server

Tests

Utilities & Docs

Uh oh!

danielhollas commented Mar 17, 2026

Uh oh!

FelyCZ commented Mar 17, 2026

Uh oh!

danielhollas commented Mar 17, 2026

Uh oh!

FelyCZ commented Mar 17, 2026

Uh oh!

JanosJiri commented Mar 17, 2026

Uh oh!

danielhollas commented Mar 17, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

danielhollas left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

FelyCZ Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

danielhollas commented Apr 26, 2026

Uh oh!

FelyCZ commented Apr 26, 2026

Uh oh!

danielhollas commented Apr 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

FelyCZ Mar 30, 2026 •

edited

Loading