MACE - Machine Learning Force Fields

https://mace-docs.readthedocs.io/en/latest/
Admin
MACE is a machine learning software for predicting many-body atomic interactions and generating...
Tokens:28,467
Snippets:366
Trust Score:8.3
Update:2 months ago
Show doc for...
Tokens:
Raw
### Test installation

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

Verify the installation using a test model and example files.

```console
mace-md -f ejm_31.sdf --model_path MACE-OFF23_small.model --output_dir md_test  --steps 1000 --unwrap --minimiser openmm
```

--------------------------------

### Install MACE from source

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/installation.rst.txt

Manual installation method by cloning the repository and installing via pip.

```bash
git clone https://github.com/ACEsuit/mace.git
pip install ./mace
```

--------------------------------

### Install mamba

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

Install the mamba package manager for environment creation.

```console
conda install mamba -c conda-forge
```

--------------------------------

### Set Hyperparameters for Multi-head Finetuning

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/troubleshooting.rst.txt

Recommended starting hyperparameters for performing multi-head finetuning in MACE.

```bash
--ema_decay=0.99999
--lr=0.0001
--num_samples_pt=100000
--forces_weight=10
--energy_weight=1
--stress_weight=1
```

--------------------------------

### Install MACE via PyPI

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/installation.rst.txt

Recommended method for installing the MACE package using pip.

```bash
pip install --upgrade pip
pip install mace-torch
```

--------------------------------

### Initialize MACE-MP calculator and get descriptors

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/foundation_models.rst.txt

Demonstrates how to initialize the MACE-MP calculator and compute descriptors for a water molecule.

```python
from mace.calculators import mace_mp
from ase import build

macemp = mace_mp() # return ASE calculator
atoms = build.molecule('H2O')
descriptors_mp = macemp.get_descriptors(atoms)
```

--------------------------------

### Install LAMMPS with MACE support on CPU

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/lammps.rst.txt

Commands to set up the environment, download LibTorch, and compile LAMMPS for CPU execution.

```bash
module purge
module load rhel7/default-ccl
module load gcc/9
```

```bash
wget https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-1.13.0%2Bcpu.zip
unzip libtorch-shared-with-deps-1.13.0+cpu.zip
rm libtorch-shared-with-deps-1.13.0+cpu.zip
```

```bash
git clone --branch mace --depth=1 https://github.com/ACEsuit/lammps
cd lammps; mkdir build; cd build
cmake -DCMAKE_INSTALL_PREFIX=$(pwd) \
      -D CMAKE_CXX_STANDARD=17 \
      -D CMAKE_CXX_STANDARD_REQUIRED=ON \
      -D BUILD_MPI=ON \
      -D BUILD_OMP=ON \
      -D PKG_OPENMP=ON \
      -D PKG_ML-MACE=ON \
      -D CMAKE_PREFIX_PATH=$(pwd)/../../libtorch \
      ../cmake
make -j 4
make install
```

--------------------------------

### Verify Data Loading Output

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/troubleshooting.rst.txt

Example of the expected console output when training data is successfully loaded.

```text
Training set [100 configs, 100 energy, 30000 forces] loaded from 'training.xyz'
```

--------------------------------

### Build and install LAMMPS

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/lammps_mliap.rst.txt

Compile the LAMMPS binary and install the Python package.

```bash
make -j 8
```

```bash
make install-python
```

--------------------------------

### Install cuEquivariance for MACE 0.3.11 or higher

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/cuda_acceleration.rst.txt

Use this command for current MACE versions. Ensure the CUDA version suffix matches your environment.

```bash
pip install cuequivariance cuequivariance-torch cuequivariance-ops-torch-cu12
# Replace "cu12" with "cu11" if you are using CUDA 11.
```

--------------------------------

### Install cuEquivariance for MACE 0.3.10 or lower

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/cuda_acceleration.rst.txt

Use this command for older MACE versions. Ensure the CUDA version suffix matches your environment.

```bash
pip install cuequivariance==0.1.0 cuequivariance-torch==0.1.0 cuequivariance-ops-torch-cu12==0.1.0
# Replace "cu12" with "cu11" if you are using CUDA 11.
```

--------------------------------

### Install MACE-POLAR-1 dependencies

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/polar_mace.rst.txt

Install the required mace-torch library and the graph_electrostatics module.

```bash
pip install mace-torch
pip install git+https://github.com/WillBaldwin0/graph_electrostatics.git@v0.4.0
```

--------------------------------

### Install mace-md environment

Source: https://mace-docs.readthedocs.io/en/latest/guide/openmm.html

Commands to set up the conda environment, including CUDA override for headnode installations.

```bash
conda install mamba -c conda-forge
export CONDA_OVERRIDE_CUDA=12.4
mamba env create -f env.yml
```

--------------------------------

### Install LAMMPS with MACE support on GPU

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/lammps.rst.txt

Commands to clone the repository, download LibTorch, and compile LAMMPS with Kokkos and MACE support for GPU acceleration.

```bash
git clone --branch=mace --depth=1 https://github.com/ACEsuit/lammps
wget https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.2.0%2Bcu121.zip
unzip libtorch-shared-with-deps-2.2.0+cu121.zip
rm libtorch-shared-with-deps-2.2.0+cu121.zip
mv libtorch libtorch-gpu
```

```bash
sintr -A YOUR-ACCOUNT-GPU -p ampere -N 1 --gres=gpu:1 -t 1:00:00
```

```bash
module purge
module load intel-mkl-2017.4-gcc-5.4.0-2tzpyn7
module load rhel8/slurm rhel8/global gcc/9 openmpi/gcc/9.3/4.0.4 cuda/12.1 cudnn
```

```bash
cd lammps
mkdir build-ampere
cd build-ampere
cmake \
    -D CMAKE_BUILD_TYPE=Release \
    -D CMAKE_INSTALL_PREFIX=$(pwd) \
    -D CMAKE_CXX_STANDARD=17 \
    -D CMAKE_CXX_STANDARD_REQUIRED=ON \
    -D BUILD_MPI=ON \
    -D BUILD_SHARED_LIBS=ON \
    -D PKG_KOKKOS=ON \
    -D Kokkos_ENABLE_CUDA=ON \
    -D CMAKE_CXX_COMPILER=$(pwd)/../lib/kokkos/bin/nvcc_wrapper \
    -D Kokkos_ARCH_AMDAVX=ON \
    -D Kokkos_ARCH_AMPERE100=ON \
    -D CMAKE_PREFIX_PATH=$(pwd)/../../libtorch-gpu \
    -D PKG_ML-MACE=ON \
    ../cmake
make -j 20
make install
```

--------------------------------

### Configure LMDB Data Format

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multipreprocessing.rst.txt

Example YAML configuration for loading LMDB files, available from version 0.3.11.

```yaml
heads:
    omat_pbe_refit:
        train_file: "/omat/train/aimd-from-PBE-1000-npt:/omat/train/aimd-from-PBE-3000-nvt:/omat/train/rattled-1000-subsampled:/omat/train/rattled-300:/omat/train/aimd-from-PBE-3000-npt:/omat/train/rattled-300-subsampled:/omat/train/rattled-500:/omat/train/aimd-from-PBE-1000-nvt:/omat/train/rattled-1000:/omat/train/rattled-500-subsampled:/omat/train/rattled-relax"
        valid_file: "/omat/val/aimd-from-PBE-1000-npt:/omat/val/aimd-from-PBE-3000-nvt:/omat/val/rattled-1000-subsampled:/omat/val/rattled-300:/omat/val/aimd-from-PBE-3000-npt:/omat/val/rattled-300-subsampled:/omat/val/rattled-500:/omat/val/aimd-from-PBE-1000-nvt:/omat/val/rattled-1000:/omat/val/rattled-500-subsampled:/omat/val/rattled-relax"
```

--------------------------------

### Train on multiple GPUs with Slurm

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multigpu.rst.txt

Example Slurm script for distributed training across multiple nodes and GPUs using the --distributed flag.

```bash
#!/bin/bash
#SBATCH --partition=gpu
#SBATCH --job-name=train
#SBATCH --output=train.out
#SBATCH --nodes=2
#SBATCH --ntasks=20
#SBATCH --ntasks-per-node=10
#SBATCH --gpus-per-node=10
#SBATCH --cpus-per-task=8
#SBATCH --exclusive
#SBATCH --time=1:00:00

srun python <mace_repo_dir>/mace/cli/run_train.py \
    --name='model' \
    --model='MACE' \
    --num_interactions=2 \
    --num_channels=128 \
    --max_L=2 \
    --correlation=3 \
    --E0s='average' \
    --r_max=5.0 \
    --train_file='./h5_data/train.h5' \
    --valid_file='./h5_data/valid.h5' \
    --statistics_file='./h5_data/statistics.json' \
    --num_workers=8 \
    --batch_size=20 \
    --valid_batch_size=80 \
    --max_num_epochs=100 \
    --loss='weighted' \
    --error_table='PerAtomRMSE' \
    --default_dtype='float32' \
    --device='cuda' \
    --distributed \
    --seed=2222 \
```

--------------------------------

### Display help

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

View all available command line arguments.

```console
mace-md -h
```

--------------------------------

### Install LAMMPS Python package

Source: https://mace-docs.readthedocs.io/en/latest/guide/lammps_mliap.html

Generate and install the Python interface for LAMMPS.

```bash
make install-python
```

--------------------------------

### Run NVT simulation with MACE-MP-0

Source: https://mace-docs.readthedocs.io/en/latest/examples/foundation_examples.html

Initializes the MACE-MP-0 calculator, sets up an H2O molecule, and runs Langevin dynamics for 200 steps.

```python
from mace.calculators import mace_mp
from ase import build
from ase.md import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase import units

macemp = mace_mp() # return the default medium ASE calculator equivalent to mace_mp(model="medium")
#macemp = mace_mp(model="large") # return a larger model
#macemp = mace_mp(model="https://tinyurl.com/y7uhwpje") # downlaod the model at the given url
#macemp = mace_mp(dispersion=True) # return a model with D3 dispersion correction
atoms = build.molecule('H2O')
atoms.calc = macemp

# Initialize velocities.
T_init = 300  # Initial temperature in K
MaxwellBoltzmannDistribution(atoms, T_init * units.kB)

# Set up the Langevin dynamics engine for NVT ensemble.
dyn = Langevin(atoms, 0.5 * units.fs, T_init * units.kB, 0.001)
n_steps = 200 # Number of steps to run
dyn.run(n_steps)
```

--------------------------------

### Run NVT simulation with MACE-MP-0

Source: https://mace-docs.readthedocs.io/en/latest/_sources/examples/foundation_examples.rst.txt

Initializes the MACE-MP-0 calculator, sets up initial velocities, and runs Langevin dynamics for an NVT ensemble.

```python
from mace.calculators import mace_mp 
from ase import build
from ase.md import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase import units

macemp = mace_mp() # return the default medium ASE calculator equivalent to mace_mp(model="medium")
#macemp = mace_mp(model="large") # return a larger model
#macemp = mace_mp(model="https://tinyurl.com/y7uhwpje") # downlaod the model at the given url
#macemp = mace_mp(dispersion=True) # return a model with D3 dispersion correction
atoms = build.molecule('H2O')
atoms.calc = macemp

# Initialize velocities.
T_init = 300  # Initial temperature in K
MaxwellBoltzmannDistribution(atoms, T_init * units.kB)

# Set up the Langevin dynamics engine for NVT ensemble.
dyn = Langevin(atoms, 0.5 * units.fs, T_init * units.kB, 0.001)
n_steps = 200 # Number of steps to run
dyn.run(n_steps)
```

--------------------------------

### Run multihead finetuning via CLI

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multihead_finetuning.rst.txt

Executes the training process with multihead finetuning enabled using specific dataset and model paths.

```bash
python -m mace.cli.run_train \
  --name mymodel_finetuned \
  --train_file path/to/your_dataset.xyz \
  --foundation_model path/to/foundation_model.model \
  --pt_train_file path/to/replay_dataset.xyz \
  --num_samples_pt 10000 \
  --filter_type_pt combinations \
  --subselect_pt fps \
  --weight_pt 1.0 \
  --atomic_numbers "[1, 6, 7, 8]" \
  --multiheads_finetuning True \
  --force_mh_ft_lr False
```

--------------------------------

### Initialize MACE-MDP Calculator

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/polarizability.rst.txt

Load the pretrained MACE-MDP model and attach it to an ASE atoms object.

```python
from mace.calculators import mace_mdp
from ase import build

calc = mace_mdp(device="cuda", default_dtype="float64")
atoms = build.molecule("H2O")
atoms.calc = calc
```

--------------------------------

### Initialize MACE-MP calculator with cuEquivariance

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/cuda_acceleration.rst.txt

Enable acceleration by setting the enable_cueq flag to True when initializing the MACE-MP calculator.

```python
from mace.calculators import mace_mp
from ase import build

macemp = mace_mp(enable_cueq=True)  # Return ASE calculator
atoms = build.molecule('H2O')
descriptors_mp = macemp.get_descriptors(atoms)
```

--------------------------------

### Use MACE-OFF23 as an ASE calculator

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/foundation_models.rst.txt

Example of using the MACE-OFF23 transferable organic force field to calculate potential energy.

```python
from mace.calculators import mace_off
from ase import build

atoms = build.molecule('H2O')
calc = mace_off(model="medium", device='cuda')
atoms.set_calculator(calc)
print(atoms.get_potential_energy())
```

--------------------------------

### Train using MP shortcut

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multihead_finetuning.rst.txt

Automatically download and use a compatible replay dataset by setting the pt_train_file parameter to 'mp'.

```bash
python -m mace.cli.run_train \
  --train_file path/to/your_dataset.xyz \
  --foundation_model medium \
  --pt_train_file mp \
  --atomic_numbers "[1, 6, 7, 8]" \
  --multiheads_finetuning True
```

--------------------------------

### LAMMPS MACE ML-IAP Script

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/lammps_mliap.rst.txt

A complete LAMMPS input script demonstrating the setup of the mliap pair style with a MACE model.

```lammps
# MACE ML-IAP example
units         metal
atom_style    atomic
newton        on

# Read structure
read_data     structure.data

# Set up MACE potential
pair_style    mliap unified model-mliap_lammps.pt 0
pair_coeff    * * C H O N

# Run settings
timestep      0.0001
thermo        100

# MD run
fix           1 all nvt temp 300 300 100
run           1000
```

--------------------------------

### Configure Weights and Biases logging

Source: https://mace-docs.readthedocs.io/en/latest/guide/openmm.html

Optional arguments to track replica exchange progress in wandb.

```bash
--wandb \
--wandb_project test_repex_calculations \
--wandb_name "methane"
```

--------------------------------

### Configure Kokkos settings

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/lammps_mliap.rst.txt

Copy the preset configuration for Kokkos CUDA support.

```bash
cp ../cmake/presets/kokkos-cuda.cmake ./
```

--------------------------------

### Select configurations for replay via CLI

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multihead_finetuning.rst.txt

Use the fine_tuning_select.py script to filter and select configurations from a replay dataset based on a target dataset.

```bash
python -m mace.cli.fine_tuning_select \
  --configs_pt path/to/replay_dataset.xyz \
  --configs_ft path/to/your_dataset.xyz \
  --num_samples 10000 \
  --subselect fps \
  --model path/to/foundation_model.model \
  --output selected_configs.xyz \
  --filtering_type combinations \
  --head_pt pt_head \
  --head_ft target_head \
  --weight_pt 1.0 \
  --weight_ft 10.0
```

--------------------------------

### Initialize custom MACE calculator with cuEquivariance

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/cuda_acceleration.rst.txt

Enable acceleration for custom models by passing the enable_cueq flag to the MACECalculator constructor.

```python
from mace.calculators import MACECalculator
from ase import build

mace_calc = MACECalculator(model_paths="mace_agnesi_small.model", enable_cueq=True)
atoms = build.molecule('H2O')
atoms.calc = mace_calc
```

--------------------------------

### Train with combined dataset

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multihead_finetuning.rst.txt

Execute the training process using the previously generated selected configurations alongside the target dataset.

```bash
python -m mace.cli.run_train \
  --name mymodel_finetuned \
  --pt_train_file selected_configs.xyz \
  --train_file path/to/your_dataset.xyz \
  --valid_fraction 0.05 \
  --foundation_model path/to/foundation_model.model \
  --energy_weight 1.0 \
  --forces_weight 100.0 \
  --swa \
  --swa_energy_weight 10.0 \
  --swa_forces_weight 100.0
```

--------------------------------

### Build LAMMPS

Source: https://mace-docs.readthedocs.io/en/latest/guide/lammps_mliap.html

Compile the LAMMPS binary using make.

```bash
make -j 8
```

--------------------------------

### Create build directory

Source: https://mace-docs.readthedocs.io/en/latest/guide/lammps_mliap.html

Prepare a dedicated directory for the LAMMPS build process.

```bash
mkdir build-mliap
cd build-mliap
```

--------------------------------

### Configure multiple training files in YAML

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/heterogeneous_data.rst.txt

Provide a list of file paths to the train_file, valid_file, or test_file parameters in the configuration.

```yaml
heads:
  DFT:
    train_file: ["path/to/training1.xyz", "path/to/training2.xyz"]
    valid_file: ["path/to/validation1.xyz", "path/to/validation2/"]
    test_file: "path/to/test.xyz"
```

--------------------------------

### Preprocess HDF5 Dataset

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multipreprocessing.rst.txt

Run the preprocess_data.py script to convert an xyz file into HDF5 format for efficient training.

```bash
mkdir processed_data
python <mace_repo_dir>/mace/cli/preprocess_data.py \
    --train_file="/path/to/train_large.xyz" \
    --valid_fraction=0.05 \
    --test_file="/path/to/test_large.xyz" \
    --atomic_numbers="[1, 6, 7, 8, 9, 15, 16, 17, 35, 53]" \
    --r_max=4.5 \
    --h5_prefix="processed_data/" \
    --compute_statistics \
    --E0s="average" \
    --seed=123 \
```

--------------------------------

### Configure wandb logging

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

Optional arguments to log replica exchange progress to Weights & Biases.

```console
--wandb \
  --wandb_project test_repex_calculations \
  --wandb_name "methane"
```

--------------------------------

### Configure MACE Training Parameters

Source: https://mace-docs.readthedocs.io/en/latest/guide/troubleshooting.html

Command-line arguments for optimizing memory usage and model performance.

```bash
--enable_cueq=True
```

```bash
--batch_size=4
```

```bash
--num_channels=64
```

```bash
--max_L=0
```

--------------------------------

### Run Langevin dynamics with MACE-MP

Source: https://mace-docs.readthedocs.io/en/latest/guide/foundation_models.html

Configures an ASE calculator for molecular dynamics, including velocity initialization and Langevin integration.

```python
from mace.calculators import mace_mp
from ase import build
from ase.md import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase import units

macemp = mace_mp() # return the default medium ASE calculator equivalent to mace_mp(model="medium") in MACE < 0.3.10 and mace_mp(model="medium-mpa-0") in MACE >= 0.3.10
#macemp = mace_mp(model="large") # return a larger model
#macemp = mace_mp(model="https://tinyurl.com/y7uhwpje") # downlaod the model at the given url
#macemp = mace_mp(dispersion=True) # return a model with D3 dispersion correction
atoms = build.molecule('H2O')
atoms.calc = macemp

# Initialize velocities.
T_init = 300  # Initial temperature in K
MaxwellBoltzmannDistribution(atoms, T_init * units.kB)

# Set up the Langevin dynamics engine for NVT ensemble.
dyn = Langevin(atoms, 0.5 * units.fs, T_init * units.kB, 0.001)
n_steps = 200 # Number of steps to run
dyn.run(n_steps)

```

--------------------------------

### Preprocess ANI-1x dataset for MACE training

Source: https://mace-docs.readthedocs.io/en/latest/_sources/examples/training_examples.rst.txt

Prepares ANI-1x data for on-the-fly loading by computing statistics and generating HDF5 files.

```shell
python <mace_repo_dir>/mace/cli/preprocess_data.py \
    --train_file="ani1x_cc_dft.xyz" \
    --valid_fraction=0.03 \
    --energy_key="DFT_energy" \
    --forces_key="DFT_forces" \
    --r_max=5.0 \
    --h5_prefix="ANI1x_cc_DFT_rc5_" \
    --compute_statistics \
    --E0s="{1: -13.62222753701504, 6: -1029.4130839658328, 7: -1484.8710358098756, 8: -2041.8396277138045}" \
    --seed=12345
```

--------------------------------

### Create conda environment

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

Set the CUDA override and create the environment from the provided YAML file.

```console
export CONDA_OVERRIDE_CUDA=12.4
mamba env create -f env.yml
```

--------------------------------

### Import MACE-MP calculator and ASE MD modules

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/foundation_models.rst.txt

Imports necessary modules for setting up MACE-MP calculators and performing molecular dynamics simulations with ASE.

```python
from mace.calculators import mace_mp 
from ase import build
from ase.md import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase import units
```

--------------------------------

### Initialize LAMMPS input file

Source: https://mace-docs.readthedocs.io/en/latest/guide/lammps.html

Set the required atomic style and unit system at the beginning of the LAMMPS input script.

```text
units         metal
atom_style    atomic
atom_modify   map yes
newton        on
```

--------------------------------

### Run LAMMPS with GPU acceleration

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/lammps_mliap.rst.txt

Command line options for running simulations on single or multiple GPUs.

```bash
lmp -k on g 1 -sf kk -pk kokkos newton on neigh half -in your_input.in
```

```bash
mpirun -np 2 lmp -k on g 2 -sf kk -pk kokkos newton on neigh half -in input.in
```

--------------------------------

### Configure Multihead Training in YAML

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multihead_training.rst.txt

Define multiple heads in the configuration file, each specifying its own training/validation files, E0s, and energy/force keys.

```yaml
heads:
    3bpa_wb97x_d3bj:
        train_file: "train_3bpa_wb97x_d3bj.xyz"
        valid_file: "val_3bpa_wb97x_d3bj.xyz"
        E0s: "e0s_wb97x_d3bj.json"
        energy_key: "dft_energy"
        forces_key: "dft_forces"

    3bpa_ccdt:
        train_file: "train_3bpa_ccdt.xyz"
        valid_file: "val_3bpa_ccdt.xyz"
        E0s: "e0s_ccdt.json"
        energy_key: "ccdt_energy"
        forces_key: "ccdt_forces"
```

--------------------------------

### Clone mace-md repository

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

Initial step to obtain the mace-md source code.

```console
git clone https://github.com/jharrymoore/mace-md.git
```

--------------------------------

### Download and prepare LAMMPS dependencies

Source: https://mace-docs.readthedocs.io/en/latest/guide/lammps.html

Clone the MACE-compatible LAMMPS branch and download the required LibTorch library for GPU support.

```bash
git clone --branch=mace --depth=1 https://github.com/ACEsuit/lammps
wget https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.2.0%2Bcu121.zip
unzip libtorch-shared-with-deps-2.2.0+cu121.zip
rm libtorch-shared-with-deps-2.2.0+cu121.zip
mv libtorch libtorch-gpu
```

--------------------------------

### Configure CSD3 Cascade Lake Environment

Source: https://mace-docs.readthedocs.io/en/latest/guide/lammps.html

Load necessary modules for the CSD3 Cascade Lake compute nodes.

```bash
module purge
module load rhel7/default-ccl
module load gcc/9
```

--------------------------------

### Configure advanced heads dictionary

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multihead_finetuning.rst.txt

Customizes individual heads by providing specific training files and E0 values via a dictionary string.

```bash
python -m mace.cli.run_train \
  --name mymodel_finetuned \
  --foundation_model path/to/foundation_model.model \
  --multiheads_finetuning True \
  --heads "{'target_head': {'train_file': 'path/to/your_dataset.xyz', 'E0s': 'path/to/e0s.json'}, 'pt_head': {'train_file': 'path/to/replay_dataset.xyz', 'E0s': 'foundation'}}" \
  --atomic_numbers "[1, 6, 7, 8]" \
  --energy_weight 1.0 \
  --forces_weight 100.0
```

--------------------------------

### Run unit tests

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

Execute the test suite for the mace-md package.

```console
pytest -s mace_md/tests/test_mace-md.py
```

--------------------------------

### Train on a single node with multiple GPUs

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multigpu.rst.txt

Command to initiate multi-GPU training on a single node using torchrun.

```bash
torchrun --standalone --nnodes=1 --nproc_per_node=4 <mace_repo_dir>/mace/cli/run_train.py \
    --name='model' \
    --model='MACE' \
    --num_interactions=2 \
    --num_channels=128 \
    --max_L=2 \
    --correlation=3 \
    --E0s='average' \
    --r_max=5.0 \
    --train_file='./h5_data/train.h5' \
    --valid_file='./h5_data/valid.h5' \
    --statistics_file='./h5_data/statistics.json' \
    --num_workers=8 \
    --batch_size=20 \
    --valid_batch_size=80 \
    --max_num_epochs=100 \
    --loss='weighted' \
    --error_table='PerAtomRMSE' \
    --default_dtype='float32' \
    --device='cuda' \
    --distributed \
    --seed=2222 \
```

--------------------------------

### Mix different file formats in configuration

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/heterogeneous_data.rst.txt

MACE automatically detects and loads different file formats when provided in a list within the configuration.

```yaml
heads:
  QM7:
    train_file: ["data/qm7/train.xyz", "data/qm7/train.h5", "data/qm7/database/"]
    valid_file: ["data/qm7/valid.xyz", "data/qm7_h5/valid/"]
    test_file:  ["data/qm7/test.xyz", "data/qm7/test.h5"]
```

--------------------------------

### Enable MACE Timing

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/lammps_mliap.rst.txt

Set this environment variable to output timing information for each calculation step.

```bash
export MACE_TIME=true
```

--------------------------------

### Access model results

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/polar_mace.rst.txt

Populate the calculator results by running a potential energy calculation.

```python
# populate calc.results
_ = atoms.get_potential_energy()
```

--------------------------------

### Train Model with Preprocessed Data

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/multipreprocessing.rst.txt

Use the run_train.py script to train a model using the directory of preprocessed HDF5 files.

```bash
python <mace_repo_dir>/mace/cli/run_train.py \
    --name="MACE_on_big_data" \
    --num_workers=16 \
    --train_file="./processed_data/train" \
    --valid_file="./processed_data/val" \
    --test_dir="./processed_data/test" \
    --statistics_file="./processed_data/statistics.json" \
    --model="ScaleShiftMACE" \
    --num_interactions=2 \
    --num_channels=128 \
    --max_L=1 \
    --correlation=3 \
    --batch_size=32 \
    --valid_batch_size=32 \
    --max_num_epochs=100 \
    --swa \
    --start_swa=60 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --error_table='PerAtomMAE' \
    --device=cuda \
    --seed=123 \
```

--------------------------------

### Load environment modules

Source: https://mace-docs.readthedocs.io/en/latest/guide/lammps.html

Load the necessary compiler, MPI, and CUDA modules required for building LAMMPS on CSD3.

```bash
module purge
module load intel-mkl-2017.4-gcc-5.4.0-2tzpyn7
module load rhel8/slurm rhel8/global gcc/9 openmpi/gcc/9.3/4.0.4 cuda/12.1 cudnn
```

--------------------------------

### Run pure MD simulation

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

Execute a standard MD simulation using a MACE potential.

```console
mace-md -f molecule.xyz \
    --model_path /path/to/my-mace.model \
    --steps 1000 \
    --timestep 1.0 \
    --interval 100 \
    --output_dir ./test_output
```

--------------------------------

### Run hybrid ML/MM replica exchange simulation

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/openmm.rst.txt

Perform a replica exchange simulation for mechanical embedding regimes.

```console
mace-md -f complex.pdb --ml_mol ligand.sdf --run_type repex --replicas 8 --output_dir ./repex_output --steps 1000 --model_path /path/to/my-mace.model
```

--------------------------------

### Train a MACE model via CLI

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/training.rst.txt

Executes the training process using the run_train.py script with specified training, validation, and test files along with model architecture and optimization settings.

```bash
python <mace_repo_dir>/mace/cli/run_train.py \
    --name="MACE_model" \
    --train_file="train.xyz" \
    --valid_fraction=0.05 \
    --test_file="test.xyz" \
    --config_type_weights='{"Default":1.0}' \
    --E0s='{1:-13.663181292231226, 6:-1029.2809654211628, 7:-1484.1187695035828, 8:-2042.0330099956639}' \
    --model="MACE" \
    --hidden_irreps='128x0e + 128x1o' \
    --r_max=5.0 \
    --batch_size=10 \
    --max_num_epochs=1500 \
    --swa \
    --start_swa=1200 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --restart_latest \
    --device=cuda \
```

--------------------------------

### Run MD simulation with MACE ASE calculator

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/ase.rst.txt

Configures a Langevin molecular dynamics simulation using a pre-trained MACE model and saves the trajectory to an XYZ file.

```python
from ase import units
from ase.md.langevin import Langevin
from ase.io import read, write
import numpy as np
import time

from mace.calculators import MACECalculator

calculator = MACECalculator(model_path='/content/checkpoints/MACE_model_run-123.model', device='cuda')
init_conf = read('BOTNet-datasets/dataset_3BPA/test_300K.xyz', '0')
init_conf.set_calculator(calculator)

dyn = Langevin(init_conf, 0.5*units.fs, temperature_K=310, friction=5e-3)
def write_frame():
        dyn.atoms.write('md_3bpa.xyz', append=True)
dyn.attach(write_frame, interval=50)
dyn.run(100)
print("MD finished!")
```

--------------------------------

### Train Model with Preprocessed Data

Source: https://mace-docs.readthedocs.io/en/latest/guide/multipreprocessing.html

Use the run_train.py script to train a model using the on-line dataloader with preprocessed HDF5 files.

```bash
python <mace_repo_dir>/mace/cli/run_train.py \
--name="MACE_on_big_data" \
--num_workers=16 \
--train_file="./processed_data/train" \
--valid_file="./processed_data/val" \
--test_dir="./processed_data/test" \
--statistics_file="./processed_data/statistics.json" \
--model="ScaleShiftMACE" \
--num_interactions=2 \
--num_channels=128 \
--max_L=1 \
--correlation=3 \
--batch_size=32 \
--valid_batch_size=32 \
--max_num_epochs=100 \
--swa \
--start_swa=60 \
--ema \
--ema_decay=0.99 \
--amsgrad \
--error_table='PerAtomMAE' \
--device=cuda \
--seed=123 \
```

--------------------------------

### Train AtomicDielectricMACE Model

Source: https://mace-docs.readthedocs.io/en/latest/_sources/guide/polarizability.rst.txt

Execute the training command for an AtomicDielectricMACE model with specific irreps and loss configurations.

```bash
python /../mace/mace/cli/run_train.py \
     --name="mace_mu_alpha" \
     --train_file="train.xyz" \
     --valid_file="val.xyz" \
     --test_dir="test.xyz" \
     --model="AtomicDielectricMACE" \
     --E0s="average" \
     --num_interactions=2 \
     --num_channels=128 \
     --max_L=2 \
     --correlation=3 \
     --MLP_irreps="16x0e+16x1o+16x2e" \
     --dipole_key="REF_dipole" \
     --polarizability_key="REF_polarizability" \
     --loss="dipole_polar" \
     --weight_decay=5e-10 \
     --polarizability_weight=2000 \
     --dipole_weight=1000 \
     --clip_grad=1.0 \
     --batch_size=128 \
     --valid_batch_size=128 \
     --max_num_epochs=40 \
     --scheduler_patience=15 \
     --patience=15 \
     --eval_interval=1 \
     --ema \
     --error_table="DipolePolarRMSE" \
     --default_dtype="float64" \
     --device=cuda \
     --seed=123 \
     --restart_latest \
     --save_cpu
```
Tokens:
Raw