Tutorial: Predicting the Optimal Mixing Parameter for aPBE0

Authors: Alastair J. A. Price, Danish Khan, O. Anatole von Lilienfeld
Date: 2025-03-10

What Does This Tutorial Cover?

In this tutorial, we demonstrate how to use an ASE-based Python script to predict the optimal fraction of exact exchange for the aPBE0 hybrid functional. The workflow takes in a molecular geometry (from an XYZ file) and reference training data, uses a machine-learning model to predict the fraction "a", and then shows you how to update your FHIaims input file (control.in) with the resulting value.

1. Background

Adaptive hybrid functionals like aPBE0 improve upon traditional global hybrids (e.g., PBE0) by adjusting the fraction of exact exchange on a per-system basis. In this approach:

Input: The molecular geometry and training data.
Processing: A machine-learning model (e.g., Kernel Ridge Regression) predicts the optimal fraction "a".
Output: The predicted fraction is used to modify the DFT calculation in FHIaims, ensuring that the tailored exchange fraction is employed.

2. ASE Code Example: A Simple Demonstration

Below is an example Python script that uses ASE to extract basic features from a molecule and trains a simple Kernel Ridge Regression model to predict the optimal fraction "a". This example uses dummy data for illustration.

# Import necessary libraries
from ase import Atoms
from sklearn.kernel_ridge import KernelRidge
import numpy as np

def extract_features(atoms):
    """
    Extract features from the ASE Atoms object.
    In a realistic setting, this function would compute descriptors 
    based on the molecular geometry (e.g., bond lengths, angles).
    Here, we use a simplified example: the number of atoms and 
    the x-coordinate of the center of mass.
    """
    num_atoms = len(atoms)
    com_x = atoms.get_center_of_mass()[0]
    return np.array([num_atoms, com_x])

# Example geometry: a water molecule (H2O)
atoms = Atoms('H2O', positions=[(0, 0, 0), (0, 0, 1), (0, 1, 0)])

# Dummy training data (replace with actual descriptors and optimal "a" values)
X_train = np.array([
    [3, 0.0],
    [3, 0.2],
    [4, 0.1]
])
y_train = np.array([0.25, 0.30, 0.28])

# Extract features from the current molecule
features = extract_features(atoms).reshape(1, -1)

# Train a simple Kernel Ridge Regression model
model = KernelRidge(kernel='linear')
model.fit(X_train, y_train)

# Predict the optimal fraction "a" for the current geometry
a_pred = model.predict(features)[0]
print("Predicted fraction of exact exchange:", a_pred)

Running this script will output a predicted fraction (for example, 0.27). This simple demonstration is meant to illustrate the basic idea before moving on to the full implementation.

3. Full Implementation with ase_wrapper.py

For a production-level prediction, we provide the ase_wrapper.py module that uses a trained machine-learning model stored in model.npz along with additional libraries. Make sure to include the model.npz file in your package.

Below is the content of ase_wrapper.py:

import numpy as np
from ase.io import read
from cMBDF import generate_mbdf
from numba import njit
from scipy.linalg import cho_solve
from joblib import Parallel, delayed

# Load the trained model parameters from model.npz
data = np.load('model.npz', allow_pickle=True)
convs = (data['rconvs'], data['aconvs'])
xtrain, qtrain, alpha, sigma = data['xtrain'], data['qtrain'], data['alpha'], data['sigma']
L = data['L']

@njit(parallel=True)
def local_kernel(A, B, Q1, Q2, sigma):
    n1, n2 = len(Q1), len(Q2)
    K = 0
    for i in range(n1):
        k = 0
        for j in range(n2):
            q1, q2 = Q1[i], Q2[j]
            if q1 == q2:
                dist = np.linalg.norm(A[i] - B[j])**2
                k += np.exp(-dist / (2 * sigma**2))
        K += k
    return K

def get_kernel(X1, X2, q1, q2, sigma):
    K = Parallel(n_jobs=-1)(delayed(local_kernel)(X1[i], X2, q1[i], q2, sigma) for i in range(len(q1)))
    return np.array(K)

def return_prediction(xyz, x0=0.7):
    """
    Returns the predicted exact-exchange admixture fraction for PBE0.
    """
    Atoms = read(xyz)
    charges = Atoms.get_atomic_numbers()
    coords = Atoms.get_positions()

    with open(xyz, 'r') as f:
        lines = f.readlines()
        comment_line = lines[1].strip()
        charge, multiplicity = map(int, comment_line.split()[:2])

    if len(charges) > 1:
        rep = generate_mbdf(np.array([charges]), np.array([coords]), convs, n_atm=2.0, pad=100)[0]
        Ne = np.sum(charges) - charge

        k = get_kernel(xtrain, rep, qtrain, charges, sigma)
        sk = local_kernel(rep, rep, charges, charges, sigma)

        alpha2 = cho_solve((L, True), k)
        variance = sk - np.dot(k, alpha2)
        unc = 1 - (variance / sk)
        cutoff = 1 - (0.5 * (1 + np.tanh(5000 * (unc - x0))))

        pred = (np.dot(k, alpha) / Ne) / 100

        return 0.25 + (cutoff * pred)

4. Integrating with FHIaims

Once you have obtained the predicted fraction (aopt), update your FHIaims input file (control.in) accordingly:

# FHIaims input: Set the hybrid mixing parameter
xc pbe0
hybrid_xc_coeff 0.27  # Replace 0.27 with the value of aopt from your prediction

Note this model has only been trained to work with PBE0 currently, other functionals are not guaranteed to work.

5. Requirements

Make sure that the following Python libraries are installed: - numpy - numba - scipy - ase - joblib

Download Required Files

To follow this tutorial, download the following required files:

ase_wrapper.py - Python module for predicting the mixing parameter.
model.npz - Pre-trained model file.
example.xyz - Sample XYZ geometry file.
cMBDF.py - Generalized convolutional many body distribution functional representation.

Right-click and choose "Save As" if the files open in your browser instead of downloading.

References

Khan et al., Sci. Adv. 11, eadt7769 (2025).
Additional relevant literature...