Shape of Hessian in batched calculations

[blaschma]

Hi,
I'm trying to calculate the Hessian in batched calculations (on commit 455854a). For this, I have adapted the provided example batch-2.py:

import tad_mctc as mctc
import torch

import dxtb

torch.set_printoptions(precision=10)


# S22 system 4: formamide dimer
numbers = mctc.batch.pack(
    (
        mctc.convert.symbol_to_number("C C N N H H H H H H O O".split()),
        mctc.convert.symbol_to_number("C O N H H H".split()),
    )
)

# coordinates in Bohr
positions = mctc.batch.pack(
    (
        torch.tensor(
            [
                [-3.81469488143921, +0.09993441402912, 0.00000000000000],
                [+3.81469488143921, -0.09993441402912, 0.00000000000000],
                [-2.66030049324036, -2.15898251533508, 0.00000000000000],
                [+2.66030049324036, +2.15898251533508, 0.00000000000000],
                [-0.73178529739380, -2.28237795829773, 0.00000000000000],
                [-5.89039325714111, -0.02589114569128, 0.00000000000000],
                [-3.71254944801331, -3.73605775833130, 0.00000000000000],
                [+3.71254944801331, +3.73605775833130, 0.00000000000000],
                [+0.73178529739380, +2.28237795829773, 0.00000000000000],
                [+5.89039325714111, +0.02589114569128, 0.00000000000000],
                [-2.74426102638245, +2.16115570068359, 0.00000000000000],
                [+2.74426102638245, -2.16115570068359, 0.00000000000000],
            ],
            dtype=torch.double,
        ),
        torch.tensor(
            [
                [-0.55569743203406, +1.09030425468557, 0.00000000000000],
                [+0.51473634678469, +3.15152550263611, 0.00000000000000],
                [+0.59869690244446, -1.16861263789477, 0.00000000000000],
                [-0.45355203669134, -2.74568780438064, 0.00000000000000],
                [+2.52721209544999, -1.29200800956867, 0.00000000000000],
                [-2.63139587595376, +0.96447869452240, 0.00000000000000],
            ],
            dtype=torch.double,
        ),
    )
)

# total charge of both system
charge = torch.tensor([0.0, 0.0], dtype=torch.double)
positions = positions.requires_grad_(True)

# instantiate calculator and calculate GFN1 energy in Hartree
calc = dxtb.calculators.GFN1Calculator(numbers, dtype=torch.double)
energy = calc.get_energy(positions, charge)

hessian = calc.get_hessian(positions, charge, derived_quantity = "energy")
print(f"hessian shape", hessian.shape)

print("Testing dimer interaction energy:")
print(f"Calculated: {energy}")
print(
    "Expected:   tensor([-23.2835232516, -11.6302093800], dtype=torch.float64)"
)
# tensor([-23.2835232516, -11.6302093800], dtype=torch.float64)

e = energy[0] - 2 * energy[1]
# tensor(-0.0231044917, dtype=torch.float64)

print("\nInteraction energy in kcal/mol:")
print(f"Calculated: {e * mctc.units.AU2KCALMOL}")
print("Expected:   -14.4982874136")
# tensor(-14.4982874136, dtype=torch.float64)`

The Hessian computed this way has shape: torch.Size([2, 12, 3, 2, 12, 3]), which is unexpected.
With this, the option matrix=True leads to an error. Do you know what could be the issue?

[marvinfriede]

It seems like batched Hessians are actually not supported yet, and I missed implementing an error message. Sorry for that.

Calling torch.func.hessian does not work for batched calculations, as this function does not have any function of "batch". This is, however, what essentially happens under the hood in dxtb and why the additional dimensions appear in the final result.

The best way to implement batched Hessians is with torch.func.vmap, but unfortunately, vmap has stricter code requirements (data flow, in-place operations, custom backward functions) that are not compatible with dxtb at the moment.
I just fixed the issues in the tad-dftd3 dependency (dftd3/tad-dftd3#89), but fixing tad-dftd4 and dxtb itself will take some more time, because the fix was unfortunately not trivial.

[marvinfriede]

I did some more testing. Making dxtb ready for using vmap will be more work than I initially anticipated, because we need to refactor some internal logic. Some function have no vmap compliant implementation (e.g. torch.unique).

[blaschma]

Thank you for the feedback and the additional testing! I will stick to unbatched calculations for now. However, having this feature available might be very interesting in the future.

[marvinfriede]

Definitely, I will try to fix this when I find the time.