python dftd3 can be as fast as the binary s-dftd3?

[jjhjm17]

Dear whom it may concern,

We thank you very much always for your help.

Maybe my installation needs to be changed.
We installed s-dftd3 binary using conda install simple-dftd3 dftd3-python
The system is 700 atom Si, O, H, Al in POSCAR format, under the periodic boundary condition.
s-dftd3 takes 0.5 sec and decreases to 0.3 sec when OMP_NUM_THREADS=2.

The python interface from dftd3.ase import DFTD3 takes 140 sec to get the energy and the force.
140 sec is much faster than DFT, but we want to combine it with a machine learning potential GRACE-2L-OAM, and run relaxation. The machine learning potential running in GPU takes 0.1 sec per step, and if we use ase SumCalculator to do the relaxation, the dominant part is dftd3..

Is this time diffference an expected behavior? Is there a python wrapper of this binary? (previous fortran dftd3 does not support openmp, this version s-dftd3 supports openmp.)
Below is a structure and the python code, and the s-dftd3 command.

from ase.io import read
from dftd3.ase import DFTD3
import time

if __name__ == '__main__':
    d3_calc = DFTD3(method='PBE', damping='d3zero')
    cfg = read('POSCAR')

    start = time.process_time()
    calc = d3_calc
    cfg.calc = calc

    print(f'{cfg.pbc =}')  # True for POSCAR
    calc.calculate(cfg, properties=['energy', 'forces'])
    energy = calc.results['energy']
    forces = calc.results['forces']
    end = time.process_time()
    print(f"Time taken for D3 calculation: {end - start:.2f} s")

s-dftd3 command:
time s-dftd3 POSCAR --zero pbe

POSCAR Structure

POSCAR.zip

[jjhjm17]

Here is a workaround, to use a python wrapper to the s-dftd3 binary, which gives 0.06 sec per step. It gives around 1 sec per step for 1 openmp thread, and with the parallelization (OMP_NUM_THREADS=48), it became 0.06 sec per step. Was my installation method of the python interface correct?

Here is the workaround.
A test result follows:

export OMP_NUM_THREADS=1; python c_test_d3_wrapper.py
    gives Average time: 0.991 s
export OMP_NUM_THREADS=48; python c_test_d3_wrapper.py
    Average time: 0.058 s

export OMP_NUM_THREADS=1; time  s-dftd3 POSCAR --zero pbe --grad --json
real    0m1.006s
export OMP_NUM_THREADS=48; time  s-dftd3 POSCAR --zero pbe --grad --json
real    0m0.106s
export OMP_NUM_THREADS=48; time  s-dftd3 POSCAR --zero pbe --grad --json > /dev/null 2>&1
real    0m0.093s

dftd3_wrapper.py

import subprocess
import os
import json
from numpy import array
from ase.calculators.calculator import Calculator, all_changes
from ase import units

class SDftd3Calculator(Calculator):
    """We use a compiled fortran binary s-dftd3, which is faster than the provided
    python interface."""
    implemented_properties = ["energy", "forces"]

    def __init__(self, workdir=None, **kwargs):
        super().__init__(**kwargs)

        # Use $SCRATCH if available for fast I/O
        if workdir is None:
            # shm_path = "/dev/shm"
            shm_path = os.environ['SCRATCH']
            if os.path.isdir(shm_path) and os.access(shm_path, os.W_OK):
                workdir = os.path.join(shm_path, "dftd3_tmp")
            else:
                workdir = "dftd3_tmp"  # fallback

        self.workdir = workdir
        os.makedirs(self.workdir, exist_ok=True)

    def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes):
        super().calculate(atoms, properties, system_changes)

        structure_file = os.path.join(self.workdir, "POSCAR")
        atoms.write(structure_file)

        # call s-dftd3
        subprocess.run(
            ["s-dftd3", "POSCAR", "--zero", "pbe", "--grad", "--json"],
            cwd=self.workdir,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            text=True
        )

        # parse output (you need to implement this)
        with open(f"{self.workdir}/dftd3.json") as f:
            json_out = json.load(f)
        energy = json_out["energy"]
        forces = - array(json_out['gradient'], dtype=float).reshape(-1, 3)

        # unit conversion from the atomic units to eV, eV/Ang
        energy *= units.Hartree
        forces *= units.Hartree / units.Bohr

        self.results = {
            "energy": energy,
            "forces": forces
        }

c_test_d3_wrapper.py

"""This script tests the speed of d3 part for a test backbone structure."""

# ── Import ase.db BEFORE tensorflow/tensorpotential ──────────────────────────
from ase.io import read
from dftd3_wrapper import SDftd3Calculator
import time

if __name__ == '__main__':
    cfg = read('POSCAR')

    start = time.perf_counter()  # wall clock time
    n_repeats = 5
    # n_repeats = 1
    calc = SDftd3Calculator()  # Use the CLI wrapper instead of the direct DFTD3 calculator
    for i in range(n_repeats):
        cfg.positions[0, 0] += 0.001 * i  # Perturb the first atom slightly to avoid caching effects
        cfg.calc = calc

        print(f'{cfg.pbc =}')  # True for POSCAR
        calc.calculate(cfg, properties=['energy', 'forces'])
        energy = calc.results['energy']
        forces = calc.results['forces']
    end = time.perf_counter()
    print(f"Time taken for D3 calculation: {end - start:.3f} s")
    print(f"Average time: {(end - start) / n_repeats:.3f} s")
    print('Energy:', energy)
    print('Forces:', forces)

The workaround was suggested, and the codes were written with the help of chat gpt web version.

[awvwgk]

The issue might be that the Python API is using a different default for the ATM contribution. In the CLI this contribution is opt-in with the --atm flag, in the Python API it is opt-out via s9=0.

[jjhjm17]

Thank you very much for the above explanation!

[awvwgk]

I am adding some optimizations in library with #172 and also trying to reduce the overhead from the ASE calculator by additional caching. In general, the overhead comes from ASE's way of handling calculators, you should not see this when directly using the Python API via dftd3.interface. If you the overhead from using two ASE calculators inside a SumCalculator is too high, I recommend to reimplement the evaluation logic for D3 directly in your calculator from the dftd3.interface functions.