Geometry optimization

Relax a geometry by descending the analytic forces with an optax optimizer. The forces come from one reverse-mode pass at each step's converged density; this is the envelope-theorem gradient, so no differentiation through the SCF loop is needed.

import jax; jax.config.update("jax_enable_x64", True)
import numpy as np
import optax
from dftax import KS, Molecule, becke, forces, scf
from dftax.energy.xc import PBE

symbols = ["O", "H", "H"]
coords = np.array([[0.0, 0.0, 0.0], [1.9, 0.0, 0.6], [1.9, 0.0, -0.6]])  # Bohr, stretched
grid = becke(n_radial=50, lebedev=110)

opt = optax.adam(2e-2)
state = opt.init(coords)
for step in range(30):
    mol = Molecule(symbols, coords, "sto-3g")
    res = scf(KS(mol, PBE(), grid=grid), e_tol=1e-9, d_tol=1e-7)
    F = np.asarray(forces(mol, PBE(), res, grid=grid))
    gmax = float(np.abs(F).max())
    if step % 5 == 0 or gmax < 3e-4:
        print(f"step {step:2d}  E = {res.e_tot:.8f} Ha   |F|max = {gmax:.2e}")
    if gmax < 3e-4:
        break
    updates, state = opt.update(-F, state)     # descend -F = dE/dR
    coords = np.asarray(optax.apply_updates(coords, updates))

step  0  E = -74.99765708 Ha   |F|max = 3.99e-01

step  5  E = -75.06025318 Ha   |F|max = 2.22e-01

step 10  E = -75.10128179 Ha   |F|max = 1.58e-01

step 15  E = -75.14397177 Ha   |F|max = 1.32e-01

step 20  E = -75.18135042 Ha   |F|max = 1.15e-01

step 25  E = -75.20709571 Ha   |F|max = 9.39e-02

d_oh = np.linalg.norm(coords[1] - coords[0])
print(f"relaxed O-H distance: {d_oh:.3f} Bohr ({d_oh*0.529177:.3f} Angstrom)")

relaxed O-H distance: 1.889 Bohr (0.999 Angstrom)