MACE interface, vol. 2

.github/workflows/gfortran.yml

@@ -332,6 +332,15 @@ jobs:
         if [ $GCC_V -ge 10 ];then export FFLAGS="-fallow-argument-mismatch";export FCFLAGS="-fallow-argument-mismatch";fi && \
         ./dev_scripts/install_mpich.sh ${HOME}/mpich ${MPICH_V}
 
+    - name: Set up uv
+      uses: astral-sh/setup-uv@v7
+      with:
+        python-version: 3.12
+        activate-environment: true
+
+    - name: Install MACE
+      run: uv pip install --torch-backend=cpu -r interfaces/MACE/mace_server.py
+
     - name: build ABIN
       run: |
         if [ $GCC_V -ge 10 ];then export FFLAGS="-fallow-argument-mismatch $FFLAGS";fi && \
@@ -341,6 +350,7 @@ jobs:
         LDLIBS: ${{ env.ABIN_LDLIBS }}
     - name: test ABIN
       run: make test
+
     - name: Codecov upload
       uses: codecov/codecov-action@v7
       with:
@@ -374,6 +384,15 @@ jobs:
       if: steps.openmpi-cache.outputs.cache-hit != 'true'
       run: ./dev_scripts/install_openmpi.sh ${HOME}/openmpi ${OPENMPI_V} ${OPENMPI_PATCH}
 
+    - name: Set up uv
+      uses: astral-sh/setup-uv@v7
+      with:
+        python-version: 3.12
+        activate-environment: true
+
+    - name: Install MACE
+      run: uv pip install --torch-backend=cpu -r interfaces/MACE/mace_server.py
+
     - name: build ABIN
       run: ./configure --mpi "${HOME}/openmpi/${OPENMPI_V}/install" && make
       env:

README.md

@@ -7,9 +7,9 @@
 ABIN is a program for performing ab initio molecular dynamics.
 It is a general purpose program that was initially designed to model nuclear quantum effects (NQE).
 NQE can be most rigirously captured with path integral MD (PIMD), but also within the Quantum Thermostat based on General Langevin Equation framework developed by Michele Cerriotti.
-ABIN can also simulate non-adiabatic events using Surface-hoping algorithm, using either the classical fewest-switches algorithm (FSSH) or simpler Landau-Zener approach which does not require non-adiabatic couplings. The LZ approach can also capture singlet-triplet transitions.
+ABIN can also simulate non-adiabatic events using Surface-hoping algorithm, using either the classical fewest-switches algorithm (FSSH) or simpler Landau-Zener (LZ) approach which does not require non-adiabatic couplings. The LZ approach can also capture singlet-triplet transitions.
 
-The basic philosophy of ABIN program is simple — 
+The basic philosophy of ABIN program is simple —
 while the program itself handles the propagation of the system according to the equations of motion,
 the forces and energies are taken from an external electronic structure program such as ORCA or TeraChem.
 The call to the chosen external program is handled via a simple shell script interface.
@@ -72,11 +72,41 @@ To install the libraries, you can use the install scripts in `dev_scripts/`.
 We use these in our Continuous Integration testing suite on Github using the Ubuntu 18.04 image.
 
 The optional libraries are:
- - [MPICH](https://www.mpich.org/): An MPI implementation used for Replica Exchange MD and MPI interface with TeraChem.
+ - [MPICH](https://www.mpich.org/): An MPI implementation used for Replica Exchange MD and MPI interface with TeraChem and MACE.
       - If you just need REMD you can also use other MPI libraries such as OpenMPI or IntelMPI.
  - [FFTW](http://www.fftw.org/): Fast Fourier Transform library used for normal mode transformation in Path Integral MD.
  - [PLUMED](https://www.plumed.org/): A collection of very useful tools for free energy calculations (MetaDynamics, Umbrella Sampling etc).
  - [TCPB-CPP](https://github.com/mtzgroup/tcpb-cpp): [EXPERIMENTAL] TCPB interface to TeraChem
+ - [MACE](https://github.com/ACEsuit/mace): Machine Learning Atomic Cluster Expansion potential.
+      - The interface to MACE requires an MPICH installation (see below) and MACE python environment.
+        See `interfaces/MACE/README.md` for details
+
+
+### Installing with MPICH
+
+1. Install MPICH (adjust MPICH_PATH as needed)
+
+```bash
+MPICH_PATH=/home/$USER/software/mpich
+./dev_scripts/install_mpich.sh $MPICH_PATH
+```
+
+2. Compile ABIN and run tests
+
+(Substitute the `<mpich_version>` for the actual version that was installed)
+```bash
+MPICH_PATH=$MPICH_PATH/<mpich_version>/install
+./configure --mpi $MPICH_PATH
+make clean && make
+make test
+```
+
+Before running ABIN, you might need to adjust your PATH and LD_LIBRARY_PATH
+
+```bash
+export LD_LIBRARY_PATH=$MPICH_PATH/lib:$LD_LIBRARY_PATH
+export PATH=$MPICH_PATH/bin:$PATH
+```
 
 
 ## Structure of the repository

configure

@@ -123,7 +123,7 @@ done
 
 # Set to default if FFLAGS is not defined yet, see
 # https://stackoverflow.com/questions/11362250/in-bash-how-do-i-test-if-a-variable-is-defined-in-u-mode
-: ${FFLAGS:="-O2 -fopenmp -fimplicit-none $warning_flags"}
+: ${FFLAGS:="-O2 -g -fopenmp -fimplicit-none $warning_flags"}
 
 if [[ -z ${FC-} ]];then
    if [[ $MPI = "TRUE" ]];then

dev_scripts/install_mpich.sh

@@ -7,19 +7,26 @@
 # Exit script immediately upon error
 set -euo pipefail
 
+CC=${CC:-gcc}
+FC=${FC:-gfortran}
+
+if [[ -z ${1-} ]]; then
+  echo "ERROR: Provide prefix path where install MPICH as first parameter"
+  exit 1
+fi
 # Path as an optional first parameter
-MPICH_DIR="${1-$HOME/mpich}"
+MPICH_DIR="$(realpath "$1")"
 # We take current stable version as default
 # (as of 06 Nov 2020).
-MPICH_VERSION="${2-"4.0.2"}"
+MPICH_VERSION="${2-"4.3.2"}"
+
+# Detect number of CPUs
+NCPUS=$(nproc)
 
 TAR_FILE="mpich-${MPICH_VERSION}.tar.gz"
 DOWNLOAD_URL="https://www.mpich.org/static/downloads/${MPICH_VERSION}/${TAR_FILE}"
 INSTALL_DIR="$MPICH_DIR/$MPICH_VERSION/install"
 
-# Github Actions machines have two CPUs, per:
-# https://docs.github.com/en/free-pro-team@latest/actions/reference/specifications-for-github-hosted-runners#supported-runners-and-hardware-resources
-NCPUS=2
 
 if [[ -d "$INSTALL_DIR" ]];then
   echo "Found existing MPICH installation in $INSTALL_DIR"
@@ -56,9 +63,10 @@ cd "$MPICH_DIR/$MPICH_VERSION/src" && tar -xzf "../pkg/${TAR_FILE}" && cd "mpich
 # export FCFLAGS=-fallow-argument-mismatch
 ./configure FC=gfortran CC=gcc \
   --enable-fortran=all \
-  --with-pm=hydra --with-device=ch3:nemesis \
+  --with-pm=hydra \
+  --with-device=ch3:nemesis \
   --with-namepublisher=pmi \
-  --enable-static --disable-shared \
+  --enable-static \
   --prefix="${INSTALL_DIR}" 2>&1 |\
   tee configure.log
 make -j $NCPUS 2>&1 | tee make.log

interfaces/MACE/README.md

@@ -0,0 +1,64 @@
+# MACE MPI Interface for ABIN
+
+This directory contains the Python-based server for the MACE (Machine Learning Atomic Cluster Expansion) potential.
+
+## Requirements
+
+- Python >= 3.8
+- PyTorch >= 1.12
+- mace-torch
+- ase
+- mpi4py
+- numpy
+
+The exact dependencies are specified as inline metadata in `mace_server.py`.
+We highly recommend installing the dependencies in a fresh virtual environment
+using the [uv package manager](https://github.com/astral-sh/uv):
+
+```console
+# Install uv first, https://github.com/astral-sh/uv#installation
+uv venv  # Creates a new virtual environment in .venv/ folder
+uv pip install -r interfaces/MACE/mace_server.py --torch-backend=auto
+source .venv/bin/activate   # Activates the environment
+```
+
+> [!IMPORTANT]
+> The tricky part is to install the correct binary version of PyTorch, as it needs to match your CUDA version.
+> The `--torch-backend=auto` should autodetect your environment and install the correct version.
+> In HPE clusters, make sure you load your CUDA environment before running the installation.
+> Please read the [uv PyTorch documentation](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) for more information.
+> If your machine doesn't have a GPU accelerator, the installation will automatically pick up the CPU-only build.
+
+ABIN itself must be compiled using the MPICH compiler, see top-level [README.md](../../README.md#installing-with-mpich) for instructions.
+
+After installation, run the MACE tests to make sure the basic communication works.
+```console
+$ make test TEST="MACE MACE_ERROR MACE_ERROR2"
+Running tests in directories:
+MACE MACE_ERROR MACE_ERROR2
+MACE	PASSED
+=======================
+MACE_ERROR	PASSED
+=======================
+MACE_ERROR2	PASSED
+=======================
+
+3 tests PASSED.
+```
+
+These tests use a fake potential, not the actual MACE implementation.
+Make sure you test the energy conservation using a short NVE dynamics
+before running production calculations!
+
+## Usage
+
+The MACE server must be launched alongside ABIN using mpirun.
+It is recommended to use the provided launch script `utils/run.mace_mpi_abin.sh`.
+
+Rough steps:
+
+1. Specify `pot='_mace_'` in the ABIN input file, everything else is the same.
+2. Copy `interfaces/MACE/mace_server.py` and `utils/run.mace_mpi_abin.sh` into your working directory.
+3. In the launch script, specify the ABIN input files, and paths to MPICH and Python installation.
+4. In the launch script, specify path to your MACE model file, and whether to evaluate the model on CPU or CUDA GPU.
+5. Profit!