DOCSEARCHbyAlgolia
Submit Docs

MACE

repository·develop·Indexed 23 days ago

https://github.com/acesuit/mace

MACE provides fast and accurate machine learning interatomic potentials using higher-order equivariant message passing for molecular dynamics and materials science. It includes tools for training via mace_run_train, evaluation with mace_eval_configs, and integration with ASE using MACE-MP and MACE-OFF foundation models. The library supports active learning molecular dynamics, Weights and Biases experiment tracking, and model conversion between cuEquivariance and e3nn.

Tokens
10.4K
Snippets
11
Records
54
Agent score
79%

What's inside MACE

  2. Understanding the MACE implementation of torch_geometric

    develop
    MACE does not depend on the full pytorch_geometric library. Instead, it includes a trimmed-down, modified subset of torch_geometric code directly within the repository. This subset provides only the most basic graph data structures required by MACE. This approach is used to simplify installation and avoid unnecessary second-degree dependencies.

  3. Fine-tune MACE foundation models

    develop

    You can fine-tune MACE foundation models (like mace-mp-0) using the mace_run_train script. Use the --foundation_model argument to specify the model type or a path to a model file.

    If you provide a path via --foundation_model=$path_model, the script will automatically extract the hyperparameters from that model.

    Supported model types for mace-mp-0 include small, medium, and large.

    mace_run_train \
  --name="MACE" \
  --foundation_model="small" \
  --train_file="train.xyz" \
  --valid_fraction=0.05 \
  --test_file="test.xyz" \
  --energy_weight=1.0 \
  --forces_weight=1.0 \
  --E0s="average" \
  --lr=0.01 \
  --scaling="rms_forces_scaling" \
  --batch_size=2 \
  --max_num_epochs=6 \
  --ema \
  --ema_decay=0.99 \
  --amsgrad \
  --default_dtype="float32" \
  --device=cuda \
  --seed=3

  5. Set up a development environment for MACE

    develop

    To develop on MACE, install the project in editable mode with the dev extra dependencies and initialize pre-commit to automate code formatting and linting (using black, isort, pylint, and mypy).

    pip install -e ".[dev]"
pre-commit install

  6. Install MACE via PyPI

    develop

    The recommended way to install MACE is using pip. Note that there is a homonymous package on PyPI that is unrelated to this project; ensure you install mace-torch.

    Requirements

    • Python >= 3.10
    • PyTorch >= 1.12
      • Note on PyTorch versions: Training with float64 is not supported with PyTorch 2.1, but is supported with 2.2 and later. PyTorch 2.4.1 is not supported.
    pip install --upgrade pip
pip install mace-torch

  7. Install MACE with Weights and Biases support

    develop

    To use MACE with Weights and Biases for experiment tracking, install the package with the [wandb] extra. You can then use the following keyword arguments to configure logging:

    • --wandb
    • --wandb_project
    • --wandb_entity
    • --wandb_name
    • --wandb_log_hypers
    pip install ./mace[wandb]

  8. Use YAML configuration for training

    develop

    Instead of long command-line arguments, you can provide a YAML file to mace_run_train using the --config flag. Command-line arguments will overwrite values specified in the YAML file.

    Example your_configs.yaml:

    name: nacl
seed: 2024
train_file: train.xyz
stage_two: yes
start_stage_two: 1200
max_num_epochs: 1500
device: cpu
test_file: test.xyz
E0s:
  41: -1029.2809654211628
  38: -1484.1187695035828
  8: -2042.0330099956639
config_type_weights:
  Default: 1.0

    Usage:

    mace_run_train --config="your_configs.yaml"

  9. Train a MACE model using mace_run_train

    develop

    Use the mace_run_train CLI script to train a model. This script is typically available in your PATH after pip installation. Alternatively, you can run it directly via python3 <path_to_cloned_dir>/mace/cli/run_train.py.

    Key Arguments

    • --name: Name of the model.
    • --train_file: Path to the training XYZ file.
    • --valid_fraction: Fraction of the training set to use for validation.
    • --test_file: Path to the test XYZ file.
    • --config_type_weights: JSON string mapping config types to weights (e.g., '{"Default":1.0}').
    • --E0s: Atomic reference energies. Can be a JSON mapping of atomic numbers to energies, or special keywords:
      • "average": Estimates atomic energies using least squares regression (useful if energies are unknown, but may be less stable for MD).
      • "estimated": (Recommended for finetuning) Estimates reference energies by solving a linear system that corrects foundation model predictions.
    • --model: Model architecture (e.g., "MACE" or "ScaleShiftMACE").
    • --hidden_irreps: Controls model size. Recommended defaults are '256x0e' or '128x0e + 128x1o'. For higher accuracy, add higher order features like 128x2e or increase channels to 256.
    • --r_max: Cutoff radius.
    • --batch_size: Training batch size.
    • --max_num_epochs: Total training epochs.
    • --stage_two: Enables increasing the energy loss weight for the last ~20% of epochs (starting from --start_stage_two). Helps lower energy errors.
    • --ema / --ema_decay: Enables Exponential Moving Average.
    • --device: Computation device (cuda, cpu, or mps for Apple Silicon).
    • --default_dtype: Precision (default is float64; float32 provides ~2x speedup).
    • --plot: Enable/disable training progress plots (default is enabled).
    mace_run_train \
    --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 \
    --stage_two \
    --start_stage_two=1200 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --restart_latest \
    --device=cuda \

  10. Run MACE tests

    develop

    The test suite is organized by requirement type. You can use pytest with specific markers to control which tests run based on available hardware or network access.

    Test Directories

    • tests/unit: Fast, CPU-only unit tests.
    • tests/workflows: End-to-end CLI training tests.
    • tests/backends: Parity and converters (requires cueq / oeq extras and GPU).
    • tests/extensions/<x>: Extension-specific tests (requires corresponding dependencies).
    • tests/foundations: Foundation-model loaders (requires network access).
    • tests/integrations/<x>: Integration tests (e.g., LAMMPS).
    • tests/benchmarks: Performance measurements (requires GPU).

    Running Tests via Pytest

    Use the following commands to run specific subsets of the test suite:

    pytest tests/unit                             # quick check
pytest tests -m "not slow and not network"    # smoke test over everything runnable

    Environment Variables for CI Reproduction

    • MACE_REQUIRE_CAPS: When set in CI, turns skipped tests (due to missing requirements like GPU) into failures.
    • MACE_CI_ALLOW_NETWORK=1: Enables network access for tests that require downloading models or data.

  11. Configure MACE model and data caching

    develop

    By default, automatically downloaded models (such as mace_mp and mace_off) and fine-tuning data are stored in ~/.cache/mace.

    You can change this location by setting the XDG_CACHE_HOME environment variable. The new path will be $XDG_CACHE_HOME/.cache/mace.

  12. Manage checkpoints and training restarts

    develop

    The script uses a CheckpointHandler to manage model states.

    • Resuming Training: Use --restart_latest to attempt to load the most recent checkpoint (including SWA states) to resume training from the last saved epoch.
    • Checkpoint Storage: Specify the directory with --checkpoints_dir and control how many are kept with --keep_checkpoints.
    • SWA Checkpoints: The handler can specifically load stage-two (SWA) models using the swa_start parameter.