Running universal machine-learned force fields

Universal machine-learned force fields (uMLFF, or alternatively: uMLIP for "universal machine-learned interatomic potentials") offer a fast alternative to the computational demands of DFT. Pre-trained uMLFFs can be used as an alternative to VASP-native force fields to drive essentially every VASP simulation which uses the the prediction-only mode ML_MODE = run. This includes molecular dynamics simulations, ionic optimization (see IBRION) and advanced sampling techniques.

Mind: For GRACE force fields and VASP ≥ 6.6.0, please refer to the dedicated wiki page.

We utilize VASP's Plugins feature to call pre-trained uMLFF models via Python and have them calculate forces and stresses of our systems.

Important: When employing uMLFFs, the accuracy and reliability of your calculations and results may vary. Thorough validation is advised.

Model selection

A dedicated list of publically available models, including performance benchmarks, can be found, for example, at MaterialsProject's matbench-discovery. As a general rule of thumb, choose a model with good relevant metrics that was trained on a broad choice of datasets. Between models with similar metrics, choose a model with fewer parameters for faster inference.

Step-by-step instructions

We offer two examples for running different uMLFFs in VASP. Please be advised that implementations may vary wildly between models, and the first step is always to understand how to call any one particular model in Python directly before trying to have it called via VASP plugin. This also allows for much easier and more comprehensive debugging.

Step 1: Setting up vasp_plugin.py

Most models specify how to call them from Python. The instructions for how to install the corresponding packages and load their calculator instances will differ from model to model, but the vasp_plugin.py file should always contain (at least) the following:

calculator = ...                                # different for each model

from vasp.force_field import AseForceField      # VASP force field wrapper class
force_field = AseForceField(calculator)         # apply wrapper class

def force_and_stress(constants, additions):     # to compute force and stress via the uMLFF model instead of VASP's DFT routines
    force_field.force_and_stress(constants, additions)

The vasp_plugin.py file should be copied/moved to your VASP calculation folder.

We will now look at two example models and how to run them in VASP via vasp_plugin.py.

Mind: These examples are arbitrarily chosen with the express purpose of demonstration, and are not recommendations.

Example A: UPET Models

First, let's use the models of the UPET package (see also the documentation here). The documentation clarifies how to load different models, so we install the package into our environment, and then follow their instructions in our own vasp_plugin.py:

from upet.calculator import UPETCalculator
calculator = UPETCalculator(model="pet-mad-s", version="1.5.0", device="cuda")

from vasp.force_field import AseForceField      # VASP force field wrapper class
force_field = AseForceField(calculator)         # apply wrapper class

def force_and_stress(constants, additions):     # to compute force and stress via the uMLFF model instead of VASP's DFT routines
    force_field.force_and_stress(constants, additions)

Note that most models offer device choices depending your CPU/GPU setup and preferences.

Example B:

Next, let's instead use the DPA-3.1-3M-FT model. The model checkpoint needs to be downloaded separately from the package, as indicated and linked on its matbench-discovery page. The documentation links to the DeePMD-kit package (see pip installation and other instructions here). Notice the file checkpoint has a .pth file extension, which indicates the PyTorch workflow is the one we need.

The DeePMD-kit instructions are relatively straightforward, so we adapt them for our vasp_plugin.py.

from deepmd.calculator import DP
calculator = DP(model="/path/to/dpa-3.1-3m-ft.pth")

from vasp.force_field import AseForceField      # VASP force field wrapper class
force_field = AseForceField(calculator)         # apply wrapper class

def force_and_stress(constants, additions):     # to compute force and stress via the uMLFF model instead of VASP's DFT routines
    force_field.force_and_stress(constants, additions)

Step 2: Setting up your INCAR

We need to add a few specific tags to our INCAR file, but your prior setup can stay largely the same:

PLUGINS/FORCE_AND_STRESS = True
PLUGINS/ML_MODE = run

Note that this can also be written differently:

PLUGINS {
    FORCE_AND_STRESS = True
    ML_MODE = run
}

You may find one or the other style more intuitive to read; they are functionally identical.

Important: The ML_MODE and PLUGINS/ML_MODE tags differ in scope and purpose. The latter specifically controls plugins behaviour.

Step 3: Run calculation

Start your VASP calculation the same way you usually would. If everything works as expected, you should notice a significant speedup and no electronic steps showing up in your OUTCAR.

Recommendations and advice

All else being equal, try picking models with good benchmark metrics, broader training datasets, and fewer parameters. The number of parameters directly affects inference speed.

References