Merge branch 'develop' into multi-head-interface

.github/workflows/pre-commit.yaml

@@ -0,0 +1,28 @@
+name: Pre-Commit Checks
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+      - name: Install requirements
+        run: |
+          pip install -U pip
+          pip install pylint
+          pip install -U black
+          pip install .[dev]
+          pip install wandb
+          pip install tqdm
+      - name: Run black
+        run: |
+          python -m black .
+      - uses: pre-commit/[email protected]

.github/workflows/release.yaml

@@ -0,0 +1,36 @@
+name: Release on tag
+
+on:
+  push:
+    tags:
+      # After vMajor.Minor.Patch _anything_ is allowed (without "/") !
+      - v[0-9]+.[0-9]+.[0-9]+*
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    if: github.repository == 'ACEsuit/mace' && startsWith(github.ref, 'refs/tags/v')
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+        cache: "pip"
+
+    - name: Build project for distribution
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install build
+        python -m build
+
+    - name: Create Release
+      uses: ncipollo/release-action@v1
+      with:
+        artifacts: "dist/*"
+        token: ${{ secrets.GITHUB_TOKEN }}
+        draft: false
+        skipIfReleaseExists: true

.github/workflow/unittest.yaml → .github/workflows/unittest.yaml

@@ -9,8 +9,8 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: "3.10"
           cache: "pip"

.pre-commit-config.yaml

@@ -1,50 +1,58 @@
+exclude: &exclude_files >
+    (?x)^(
+        docs/.*|
+        tests/.*|
+        .github/.*|
+        LICENSE.md|
+        README.md|
+    )$
+
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v2.5.0
     hooks:
-      - id: end-of-file-fixer
-      - id: trailing-whitespace
+    -   id: mixed-line-ending
+    -   id: trailing-whitespace
+        exclude: *exclude_files
 
   - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 23.12.1
+    rev: 24.4.0
     hooks:
       - id: black
         name: Black Formating
+        exclude: *exclude_files
 
   - repo: https://github.com/pycqa/isort
     rev: 5.13.2
     hooks:
       - id: isort
         name: Sort imports
+        exclude: *exclude_files
 
-  # Failing - to be investigated separately
-  # - repo: local
-  #   hooks:
-  #     - id: pylint
-  #       name: Pylint Checks
-  #       entry: pylint
-  #       language: system
-  #       types: [python]
-  #       args:
-  #         [
-  #           "--rcfile=pyproject.toml", 
-  #           "mace", 
-  #           "tests", 
-  #           "scripts"
-  #         ]  
-
-  # - repo: local
-  #   hooks:
-  #     - id: mypy
-  #       name: mypy type checks
-  #       entry: mypy
-  #       language: system
-  #       types: [python]
-  #       args:
-  #         [
-  #           --config-file=.mypy.ini,
-  #           mace,
-  #           tests,
-  #           scripts
-  #         ]
-
+  - repo: https://github.com/PyCQA/pylint
+    rev: pylint-2.5.2
+    hooks:
+    -   id: pylint
+        language: system
+        args: [
+          '--disable=line-too-long',
+          '--disable=no-member',
+          '--disable=missing-module-docstring',
+          '--disable=missing-class-docstring',
+          '--disable=missing-function-docstring',
+          '--disable=too-many-arguments',
+          '--disable=too-many-locals',
+          '--disable=not-callable',
+          '--disable=logging-fstring-interpolation',
+          '--disable=logging-not-lazy',
+          '--disable=invalid-name',
+          '--disable=too-few-public-methods',
+          '--disable=too-many-instance-attributes',
+          '--disable=too-many-statements',
+          '--disable=too-many-branches',
+          '--disable=import-outside-toplevel',
+          '--disable=cell-var-from-loop',
+          '--disable=duplicate-code',
+          '--disable=use-dict-literal',
+        ]
+        exclude: *exclude_files

README.md

@@ -26,6 +26,7 @@
       - [Example usage in ASE](#example-usage-in-ase)
     - [MACE-OFF: Transferable Organic Force Fields](#mace-off-transferable-organic-force-fields)
       - [Example usage in ASE](#example-usage-in-ase-1)
+    - [Finetuning foundation models](#finetuning-foundation-models)
   - [Development](#development)
   - [References](#references)
   - [Contact](#contact)
@@ -36,7 +37,7 @@
 MACE provides fast and accurate machine learning interatomic potentials with higher order equivariant message passing.
 
 This repository contains the MACE reference implementation developed by
-Ilyes Batatia, Gregor Simm, and David Kovacs.
+Ilyes Batatia, Gregor Simm, David Kovacs, and the group of Gabor Csanyi, and friends (see Contributors).
 
 Also available:
 
@@ -49,63 +50,33 @@ A partial documentation is available at: https://mace-docs.readthedocs.io
 
 ## Installation
 
-Requirements:
+### 1. Requirements:
 
-- Python >= 3.7
-- [PyTorch](https://pytorch.org/) >= 1.12 **(training with float64 is not supported with PyTorch 2.1)**.
+- Python >= 3.7  (for openMM, use Python = 3.9)
+- [PyTorch](https://pytorch.org/) >= 1.12 **(training with float64 is not supported with PyTorch 2.1 but is supported with 2.2 and later)**
 
-(for openMM, use Python = 3.9)
+**Make sure to install PyTorch.** Please refer to the [official PyTorch installation](https://pytorch.org/get-started/locally/) for the installation instructions. Select the appropriate options for your system.
 
-### pip installation
-
-To install via `pip`, follow the steps below:
+### 2a. Installation from PyPI
+This is the recommended way to install MACE. 
 
 ```sh
 pip install --upgrade pip
-pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
 pip install mace-torch
 ```
+**Note:** The homonymous package on [PyPI](https://pypi.org/project/MACE/) has nothing to do with this one.
 
-For CPU or MPS (Apple Silicon) installation, use `pip install torch torchvision torchaudio` instead.
 
-### conda installation
+### 2b. Installation from source
 
-If you do not have CUDA pre-installed, it is **recommended** to follow the conda installation process:
 
 ```sh
-# Create a virtual environment and activate it
-conda create --name mace_env
-conda activate mace_env
-
-# Install PyTorch
-conda install pytorch torchvision torchaudio pytorch-cuda=11.6 -c pytorch -c nvidia
-
-# (optional) Install MACE's dependencies from Conda as well
-conda install numpy scipy matplotlib ase opt_einsum prettytable pandas e3nn
-
-# Clone and install MACE (and all required packages)
 git clone https://github.com/ACEsuit/mace.git
 pip install ./mace
 ```
 
-### pip installation from source
-
-To install via `pip`, follow the steps below:
 
-```sh
-# Create a virtual environment and activate it
-python -m venv mace-venv
-source mace-venv/bin/activate
 
-# Install PyTorch (for example, for CUDA 11.6 [cu116])
-pip3 install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
-
-# Clone and install MACE (and all required packages)
-git clone https://github.com/ACEsuit/mace.git
-pip install ./mace
-```
-
-**Note:** The homonymous package on [PyPI](https://pypi.org/project/MACE/) has nothing to do with this one.
 
 ## Usage
 
@@ -215,7 +186,7 @@ python ./mace/scripts/preprocess_data.py \
     --seed=123 \
 ```
 
-To see all options and a little description of them run `python ./mace/scripts/preprocess_data.py --help` . The script will create a number of HDF5 files in the `processed_data` folder which can be used for training. There wiull be one file for trainin, one for validation and a separate one for each `config_type` in the test set. To train the model use the `run_train.py` script as follows:
+To see all options and a little description of them run `python ./mace/scripts/preprocess_data.py --help` . The script will create a number of HDF5 files in the `processed_data` folder which can be used for training. There will be one folder for training, one for validation and a separate one for each `config_type` in the test set. To train the model use the `run_train.py` script as follows:
 
 ```sh
 python ./mace/scripts/run_train.py \

mace/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.5"
+__version__ = "0.3.6"

mace/calculators/foundations_models.py

@@ -1,7 +1,7 @@
 import os
 import urllib.request
 from pathlib import Path
-from typing import Literal, Union
+from typing import Union
 
 import torch
 from ase import units
@@ -20,7 +20,7 @@ def mace_mp(
     device: str = "",
     default_dtype: str = "float32",
     dispersion: bool = False,
-    damping: Literal["zero", "bj", "zerom", "bjm"] = "bj",
+    damping: str = "bj",  # choices: ["zero", "bj", "zerom", "bjm"]
     dispersion_xc: str = "pbe",
     dispersion_cutoff: float = 40.0 * units.Bohr,
     **kwargs,
@@ -108,6 +108,7 @@ def mace_mp(
     mace_calc = MACECalculator(
         model_paths=model, device=device, default_dtype=default_dtype, **kwargs
     )
+    d3_calc = None
     if dispersion:
         gh_url = "https://github.com/pfnet-research/torch-dftd"
         try: