Implicit differentiation

Gradients of converged quantities: implicit_density exposes the SCF fixed point as a differentiable function of the energy functional (CPHF as a custom_vjp), so a polarizability is one jax.jacobian through the converged density, with no field stepping.

import jax; jax.config.update("jax_enable_x64", True)
import equinox as eqx
import jax.numpy as jnp
import numpy as np
from dftax import KS, Molecule, becke, implicit_density, polarizability
from dftax.integrals.multipole import dipole_matrices
from dftax.ks.properties import nuclear_dipole
from dftax.energy.xc import PBE

mol = Molecule.from_xyz("O 0 0 0; H 0.7586 0 0.5043; H 0.7586 0 -0.5043", "sto-3g")
ks0 = KS(mol, PBE(), grid=becke(50, 110))
D = dipole_matrices(ks0.basis)
nuc = nuclear_dipole(mol)

def mu(field):
    ksf = eqx.tree_at(lambda k: k.hcore, ks0, ks0.hcore + jnp.einsum("i,ipq->pq", field, D))
    return nuc - jnp.einsum("ipq,pq->i", D, implicit_density(ksf))

alpha = jax.jacobian(mu)(jnp.zeros(3))
alpha = 0.5 * (alpha + alpha.T)
print("analytic polarizability (a.u.):\n", np.asarray(alpha))

analytic polarizability (a.u.):
 [[ 3.63690551e+00 -4.55636293e-16 -4.14868558e-13]
 [-4.55636293e-16  1.24682744e-02 -5.52033694e-16]
 [-4.14868558e-13 -5.52033694e-16  2.74677412e+00]]

alpha_fd = np.asarray(polarizability(mol, PBE(), method="fd", grid=becke(50, 110)))
print("finite-field reference:\n", alpha_fd)
print("max |diff| =", float(np.abs(np.asarray(alpha) - alpha_fd).max()))

finite-field reference:
 [[ 3.63690617e+00 -1.92407201e-09  5.34711164e-10]
 [-1.92407201e-09  1.24688180e-02 -2.12468931e-11]
 [ 5.34711164e-10 -2.12468931e-11  2.74675871e+00]]
max |diff| = 1.5413626930360635e-05