External machine-learned force fields in QuantumATK¶

The field of universal machine-learned force fields (MLFFs) is advancing rapidly, with new models continuously being developed and published. While QuantumATK offers native support for several of these MLFFs (see the list here), it is not feasible to include every available model within the package due to the fast pace of innovation and the growing number of MLFFs.

To accommodate this, QuantumATK provides a flexible framework for integrating external MLFFs through the ExternalMachineLearnedFFCalculator. This functionality enables the use of any MLFF that implements an ASE-compatible calculator, including models not yet supported natively in QuantumATK. It also allows users to incorporate their own open-source trained, or fine-tuned MLFFs into their QuantumATK simulation workflows.

Note

Synopsys QuantumATK is not affiliated with, nor does it endorse or promote, any of the external packages mentioned in this guide.

Synopsys QuantumATK assumes no responsibility for the accuracy, safety, or reliability of any third-party packages or files, nor for any issues that may arise from their use. Users are strongly advised to install and use only packages and files obtained from trusted sources.

While the packages listed below have been tested, support is not guaranteed for all versions. Other versions may function correctly, but compatibility and performance cannot be assured.

Please ensure that you comply with the license terms of any external packages you choose to use.

Note

Some MLFFs may have dependencies that are incompatible with the Python environment provided by QuantumATK. This can occur, for example, when a MLFF requires a specific version of a library that differs from the one bundled with QuantumATK. Additionally, the set of supported external MLFFs may change over time as both QuantumATK and the MLFFs themselves are updated.

Note

All arguments passed to the underlying ASE calculator class should be basic python types, such as str, int, float, etc. This will ensure that the ExternalMachineLearnedFFCalculator can properly restore the calculator when reading it from disc and in complex workflows, such as SurfaceProcessSimulation If any arguments types are not not known to QuantumATK, a placeholder calculator is saved instead, which provides basic functionality, e.g. in the GUI, but cannot be used for further calculations. If needed, this can also be achieved by implementing a custom wrapper function or class that sets up the calculator using only basic types as input arguments arguments, as shown with the Orb example below.

Note

It is recommended to use Linux when running QuantumATK with external MLFFs, as support for Windows is limited for many of these models. Consequently, this guide focuses on installation and usage instructions for Linux environments.

Setting up the virtual environment¶

Although external MLFF modules can be installed directly into QuantumATK, is always recommended that external MLFF modules are installed in a separate Python virtual environment to be able to completely remove possible conflicting installations without having to re-install QuantumATK. For a detailed guide on how to set up a virtual environment, please refer to the virtual environment section of the QuantumATK manual.

The following is a brief example of how to set up a new virtual environment suitable for installing an external MLFF:

sh$ QATK_DIR="/path/to/your/QuantumATK/installation"
sh$ VENV_DIR="/path/to/your/custom/venv"
sh$ source "${QATK_DIR}/atkpython/bin/activate_env"
(qatk) sh$ python -m venv --system-site-packages --upgrade-deps "$VENV_DIR"
(qatk) sh$ source "${VENV_DIR}/bin/activate"

After running these commands, your command prompt should display the prefix (venv) (qatk), indicating that both environments are active. You are now ready to install the required packages for your external MLFF.

SevenNet¶

Note

The following guide applies to SevenNet version 0.11.1.

To install the required packages for using SevenNet models, run:

(venv) (qatk) sh$ pip install sevenn==0.11.1

This command may raise an error related to the version of the e3nn library. SevenNet requires e3nn>=0.5.0, whereas QuantumATK includes e3nn==0.4.4. As a result, you will have e3nn==0.4.4 in the QuantumATK shell environment and e3nn>=0.5.0 in the virtual environment.

This version mismatch is generally not an issue, as Python prioritizes packages installed in the active virtual environment. Therefore, the correct version of e3nn will be used when running SevenNet. However, this setup does mean that MACE training functionality will not work within the QuantumATK shell environment while the virtual environment is active.

Below is a simple example demonstrating how to use SevenNet. This snippet sets up a calculator using the 7net-mf-ompa model and runs it on a GPU:

from sevenn.calculator import SevenNetCalculator
calculator = ExternalMachineLearnedFFCalculator(
    SevenNetCalculator,
    model='7net-mf-ompa',
    modal='mpa',
    device='cuda'
)

Orb models¶

Note

The following guide applies to Orb models version 0.4.2.

To install the required packages for using Orb models, run the following commands:

(venv) (qatk) sh$ pip install orb-models==0.4.2
(venv) (qatk) sh$ pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"

Orb models require scipy>=1.13.1, while QuantumATK includes scipy==1.10.1. Despite this version mismatch, Orb models appear to function correctly with the older version. To avoid conflicts, you can simply uninstall the newer version installed as a dependency:

(venv) (qatk) sh$ pip uninstall scipy

To fully support Orb models with QuantumATK, you need to define a custom wrapper function, as the underlying ASE calculator expects an object as an argument (see note above). This wrapper function must be placed in a separate file in the same working directory, for example, wrap_orb.py

# wrap_orb.py
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

def wrap_orb_v2(device):
    orbff = pretrained.orb_v2(device=device)
    return ORBCalculator(orbff, device=device)

You can then use this wrapper in your workflow script to initialize the ExternalMachineLearnedFFCalculator:

from wrap_orb import wrap_orb_v2
calculator = ExternalMachineLearnedFFCalculator(wrap_orb_v2, device='cuda')

DeePMD-kit¶

Note

The following guide applies to DeePMD-kit version 3.0.2.

To install DeePMD-kit, use the following command. Note that due to incompatibilities in the precompiled binary packages, the library must be built from source using the --no-binary deepmd-kit flag:

(venv) (qatk) sh$ pip install deepmd-kit==3.0.2 --no-binary deepmd-kit

To use a DeePMD-kit model, you first need to download it from AIS Square and follow the specific instructions provided for that model. In this example, we use the model DPA-2.3.1-v3.0.0rc0 as an example. We will extract and freeze the Domains_SemiCond using the following command:

(venv) (qatk) sh$ dp --pt freeze -c DPA2_medium_28_10M_rc0.pt -o DPA2_SemiCond.pth --head Domains_SemiCond

Once the model is prepared, it can be used with the ExternalMachineLearnedFFCalculator as shown below:

from deepmd.calculator import DP
calculator = ExternalMachineLearnedFFCalculator(
    DP,
    model='DPA2_SemiCond.pth'
)

MatterSim¶

Note

The following guide applies to MatterSim version 1.1.2.

MatterSim lists atomate2 as a dependency, but this is not required for using the force field itself (as of version 1.1.2). Installing atomate2 introduces numerous additional dependencies, many of which conflict with packages included in QuantumATK. To avoid these issues, we recommend installing only the necessary dependencies manually:

(venv) (qatk) sh$ pip install "e3nn>=0.5.0" "torch_geometric>=2.5.3" "torch_runstats>=0.2.0" deprecated loguru
(venv) (qatk) sh$ pip install --no-deps mattersim==1.1.2

This command may raise an error related to the version of the e3nn library. MatterSim requires e3nn>=0.5.0, whereas QuantumATK includes e3nn==0.4.4. As a result, you will have e3nn==0.4.4 in the QuantumATK shell environment and e3nn>=0.5.0 in the virtual environment.

This version mismatch is generally not an issue, as Python prioritizes packages installed in the active virtual environment. Therefore, the correct version of e3nn will be used when running MatterSim. However, this setup does mean that MACE training functionality will not work within the QuantumATK shell environment while the virtual environment is active.

Below is a simple example demonstrating how to use MatterSim. This snippet sets up a calculator using the default MatterSim model and runs it on a GPU:

from mattersim.forcefield import MatterSimCalculator
calculator = ExternalMachineLearnedFFCalculator(
    MatterSimCalculator,
    device='cuda'
)

CHGnet¶

Note

The following guide applies to CHGNet version 0.3.8.

To install CHGNet, run the following command:

(venv) (qatk) sh$ pip install chgnet==0.3.8

Below is a simple example demonstrating how to use CHGNet. This snippet sets up a calculator using the default CHGNet model and runs it on a GPU:

from chgnet.model.dynamics import CHGNetCalculator
calculator = ExternalMachineLearnedFFCalculator(
    CHGNetCalculator,
    use_device='cuda'
)

FAIRChem UMA¶

Note

The following guide applies to FAIRChem version 2.2.0.

FAIRChem lists numpy>2.0 as a dependency, but QuantumATK includes numpy==1.26.4 as of the X-2025.06 release. Despite this version mismatch, it appears that FAIRChem models function correctly with the older version of NumPy when just using the force field itself. To install FAIRChem, run the following command:

(venv) (qatk) sh$ pip install "e3nn>=0.5.0" "numba>=0.61.2" "huggingface_hub>=0.27.1" wandb lmdb tensorboard orjson submitit hydra-core torchtnt
(venv) (qatk) sh$ pip install --no-deps fairchem-core==2.2.0

Similar to MatterSim and SevenNet, this command may raise an error related to the version of the e3nn library, which can be ignored. FAIRChem UMA models require a custom wrapper function to be used with the ExternalMachineLearnedFFCalculator, as the underlying ASE calculator expects an object as an argument (see note above). This wrapper function must be placed in a separate file in the same working directory - for example, wrap_uma.py

# wrap_uma.py
from fairchem.core import FAIRChemCalculator
from fairchem.core.units.mlip_unit import load_predict_unit

def wrap_uma(ckpt="uma-s-1.pt", task_name="omat", device="cuda"):
    predictor = load_predict_unit(ckpt, device=device)
    calc = FAIRChemCalculator(predictor, task_name=task_name)
    return calc

The FAIRChem UMA model is available for download on the Hugging Face Hub, but access must be requested beforehand.

After downloading the model checkpoint, you can then use the wrapper in your workflow script to initialize the calculator:

from wrap_uma import wrap_uma
calculator = ExternalMachineLearnedFFCalculator(
    wrap_uma,
    ckpt='path/to/uma_model.pt',
    task_name='omat',
    device='cuda'
)

Other Models¶

The above list presents a selection of tested models. Other models may work as well, similar to the examples shown above. In many cases they can be installed using pip and the installation instructions that come with the open-source module. Incompatibilities during the installation process may in some cases require some extra technical workarounds, which can be tested safely in a virtual environment.