Example profiling script for pyDFTD3 v2.0a0

dftd3-pre-jax-profile.py

import cProfile

from dftd3.dftd3 import d3


def main(charges, coordinates, functional, damping):
    _, _, _ = d3(charges, coordinates, functional=functional, damp=damping)


coordinates = [
    4.01376695,
    -0.42094027,
    0.00413283,
    2.37659705,
    -2.06988695,
    -0.07118787,
    3.23930049,
    2.02346393,
    0.03394893,
    6.09063455,
    -0.5931548,
    0.05363799,
    1.40721479,
    2.12118734,
    -0.0108338,
    -4.04214497,
    0.34532288,
    -0.03309288,
    -2.55773675,
    2.13450991,
    -0.01370996,
    -3.05560778,
    -2.01961834,
    0.06482139,
    -6.12485749,
    0.3304205,
    -0.12008265,
    -1.22314224,
    -1.95250094,
    0.14005516,
]

charges = [6, 8, 8, 1, 1, 6, 8, 8, 1, 1]

functional = "B3LYP"

damping = "zero"
#damping = "bj"

cProfile.run("main(charges, coordinates, functional, damping)")

lin seems to only be used to compute the repulsive terms: it obtains some sort of compound index. Note that repulsion is not in the default definition of D3: a correctly placed if should avoid doing unnecessary steps.
The repulsion is in a triple loop and maybe we can arrange it so that the computation of the compound index can be done without resorting to calls to max and min inside lin.