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)")

On a larger molecule (valinomycin with 168 atoms)

   D3-dispersion correction with zero damping:

   Zero-damping parameters: s6 = 1.0; rs6 = 1.261; s8 = 1.703
   3-body term will not be calculated
         17065233 function calls in 11.750 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000   11.750   11.750 <string>:1(<module>)
        1    0.003    0.003   11.750   11.750 dftd3-pre-jax-profile.py:7(main)
        1    5.212    5.212   11.748   11.748 dftd3.py:68(d3)
  4812106    2.880    0.000    5.980    0.000 utils.py:109(lin)
    14196    0.137    0.000    0.157    0.000 utils.py:119(getc6)
        1    0.000    0.000    0.000    0.000 utils.py:161(check_inputs)
        1    0.085    0.085    0.093    0.093 utils.py:67(ncoord)
        1    0.000    0.000    0.000    0.000 utils.py:84(<listcomp>)
        1    0.000    0.000   11.750   11.750 {built-in method builtins.exec}
    65810    0.008    0.000    0.008    0.000 {built-in method builtins.isinstance}
        4    0.000    0.000    0.000    0.000 {built-in method builtins.len}
  4812106    1.614    0.000    1.614    0.000 {built-in method builtins.max}
  4812106    1.487    0.000    1.487    0.000 {built-in method builtins.min}
        3    0.000    0.000    0.000    0.000 {built-in method builtins.print}
    93396    0.015    0.000    0.015    0.000 {built-in method math.exp}
    70476    0.013    0.000    0.013    0.000 {built-in method math.pow}
  2384760    0.297    0.000    0.297    0.000 {built-in method math.sqrt}
      262    0.000    0.000    0.000    0.000 {method 'append' of 'list' objects}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

There seems to be a lot of

def lin(i1, i2):
    """Linear interpolation."""

    idum1 = max(i1, i2)
    idum2 = min(i1, i2)
    lin = idum2 + idum1 * (idum1 - 1) / 2

    return lin

Possibly whatever calls this does so redundantly and storing things can save time, or maybe the lin routine itself can be optimized - it seems like there should be some way to do so.

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.